Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,153 @@
+---
+license: mit
+tags:
+- coreml
+- ANE
+- LLaMA
+- Qwen
+- DeepSeek
+- Gemma
+- Apple
+- Apple Neural Engine
+- DeepHermes
+---
+# ANEMLL
+
+**ANEMLL** (pronounced like "animal") is an open-source project focused on accelerating the porting of Large Language Models (LLMs) to tensor processors, starting with the Apple Neural Engine (ANE).
+
+The goal is to provide a fully open-source pipeline from model conversion to inference for common LLM architectures running on ANE.
+
+This enables seamless integration and on-device inference for low-power applications on edge devices, ensuring maximum privacy and security.
+
+This is critical for autonomous applications, where models run directly on the device without requiring an internet connection.
+
+For more information, visit the [ANEMLL GitHub repository](https://github.com/anemll/anemll).
+
+
+---
+
+## License
+
+ANEMLL is licensed under the [MIT License](https://opensource.org/license/mit).
+The original model may require a separate license depending on the architecture:
+- LLaMA models: Based on Meta's LLaMA and may require Meta's license
+- Qwen models: Based on Alibaba's Qwen and may require Alibaba's license
+- Gemma models: Based on Google's Gemma and subject to Gemma Terms of Use
+- Other models: Check respective original model licenses
+
+This model is converted for CoreML using ANEMLL's open-source conversion pipeline. It supports multiple LLM architectures including LLaMA, Qwen, Gemma, and DeepSeek variants.
+
+---
+
+## Requirements
+
+- **macOS 15 (Sequoia)** or later with Apple Neural Engine and 8GB RAM or more
+- **CoreML Tools 8.x+** and **HuggingFace Transformers** libraries
+- **Python 3.9+**
+
+`chat.py` provides a sample inference script.  
+`chat_full.py` provides a sample inference script with history and conversation management.  
+
+**Installation**
+
+1. Download the model from Hugging Face:
+```bash
+# Install required tools
+pip install huggingface_hub
+
+# Install Git LFS (Large File Support)
+# macOS with Homebrew:
+brew install git-lfs
+# Or Ubuntu/Debian:
+# sudo apt-get install git-lfs
+
+# Initialize Git LFS
+git lfs install
+
+# Clone the repository with model files
+git clone https://huggingface.co/anemll/anemll-google-gemma-3-4b-it-qat-int4-unquantized-ctx1024_0.3.5
+```
+
+2. Extract model files:
+```bash
+# Navigate to cloned directory
+cd anemll-google-gemma-3-4b-it-qat-int4-unquantized-ctx1024_0.3.5
+
+# Pull LFS files (model weights)
+git lfs pull
+
+# Extract CoreML model files
+find . -type f -name "*.zip" -exec unzip {} \;
+```
+
+3. Install dependencies:
+```bash
+pip install coremltools transformers
+```
+
+**Coremltools:**
+
+See coremltools installation guide at https://apple.github.io/coremltools/docs-guides/source/installing-coremltools.html 
+
+**How to Run**
+
+1. Basic chat interface:
+```bash
+python chat.py --meta ./meta.yaml
+```
+
+2. Full conversation mode with history:
+```bash
+python chat_full.py --meta ./meta.yaml
+```
+
+> Note: The first time the model loads, macOS will take some time to place it on the device.
+> Subsequent loads will be instantaneous.
+> Use Ctrl-D to exit, Ctrl-C to interrupt inference.
+
+**More Info**
+Please check following links for later updates:
+
+* [GitHub](https://github.com/anemll)
+* [Hugging Face Models](https://huggingface.co/anemll)
+* [Twitter/X](https://x.com/anemll)
+* [Website](https://anemll.com)
+
+
+realanemll@gmail.com
+
+# anemll-google-gemma-3-4b-it-qat-int4-unquantized-ctx1024_0.3.5
+
+This is a CoreML model converted using ANEMLL for Apple Neural Engine inference.
+
+## Available Distributions
+
+### Standard Distribution
+- Contains zipped MLMODELC files
+- Suitable for macOS and development
+
+### iOS Distribution
+- Contains unzipped MLMODELC files
+- Ready for iOS deployment
+- Includes offline tokenizer support
+
+## Model Information
+- Context Length: 1024
+- Batch Size: 64
+- Number of Chunks: 2
+- LUT Quantization: 4
+
+## Quick Start
+
+### Test in iOS/macOS App
+Try our sample Chat-Bot app on TestFlight:
+1. Install TestFlight from App Store
+2. Join beta test: [TestFlight Link](https://testflight.apple.com/join/jrQq1D1C)
+3. App includes a small demo model pre-installed
+4. You can add custom models via HuggingFace URLs
+
+> [!Note]
+> - The TestFlight app works on both iOS and macOS
+> - Demonstrates proper model integration and provides a reference implementation
+> - iOS requires unzipped MLMODELC files and config.json for offline tokenizer
+> - macOS supports both zipped and unzipped model formats

chat_full.py

@@ -0,0 +1,1993 @@
+# chat.py
+#!/usr/bin/env python3
+# chat.py
+# Copyright (c) 2025 Anemll
+# Licensed under the MIT License
+
+import argparse
+import os
+import re
+import glob
+from pathlib import Path
+import coremltools as ct
+from transformers import LlamaTokenizer, AutoTokenizer
+import torch
+import torch.nn.functional as F
+import numpy as np
+import queue
+import threading
+import time
+import yaml
+import sys
+
+# ANSI color codes
+LIGHT_BLUE = "\033[94m"
+DARK_BLUE = "\033[34m"
+LIGHT_GREEN = "\033[92m"
+SYSTEM_COLOR = "\033[93m"
+RESET_COLOR = "\033[0m"
+
+# Add at the top with other constants
+WARMUP_TOKEN_LIMIT = 10  # Maximum tokens to generate during warmup
+THINKING_MODE = False
+THINKING_PROMPT = """You are a deep thinking AI, you may use extremely long chains of thought to deeply consider the problem and deliberate with yourself via systematic reasoning processes to help come to a correct solution prior to answering. You should enclose your thoughts and internal monologue inside <think> </think> tags, and then provide your solution or response to the problem."""
+DEBUG_LEVEL = 0  # Default debug level
+
+def print_system(msg: str) -> None:
+    print(f"{SYSTEM_COLOR}[SYSTEM] {msg}{RESET_COLOR}")
+
+class TokenPrinter:
+    """Handles background printing of generated tokens."""
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+        self.token_queue = queue.Queue()
+        self.stop_event = threading.Event()
+        self.thread = None
+        self.buffer = ""
+        self.lock = threading.Lock()
+        self.thinking = True  # Track if we're still in thinking mode
+        self.decoding_buffer = []  # Buffer for token IDs
+        # Timing and stats tracking
+        self.start_time = time.time()
+        self.token_count = 0
+        self.prefill_time = 0
+        self.inference_time = 0
+        self.context_pos = 0
+        self.start()
+
+    def start(self):
+        """Start the printer thread."""
+        if self.thread is None:
+            self.thread = threading.Thread(target=self._print_worker)
+            self.thread.daemon = True
+            self.thread.start()
+
+    def add_token(self, token_id):
+        """Add a token to the print queue."""
+        if not self.stop_event.is_set():
+            self.token_queue.put(token_id)
+            self.token_count += 1
+
+    def drain_buffer(self):
+        """Decode token IDs from decoding_buffer in the main thread."""
+        if not self.decoding_buffer:
+            return
+
+        # Decode all tokens at once in the main thread
+        token_str = self.tokenizer.decode(self.decoding_buffer)
+        self.decoding_buffer.clear()
+
+        # Save to buffer for conversation history
+        self.buffer += token_str
+
+        # Color-handling logic
+        if self.thinking and "</think>" in token_str:
+            self.thinking = False
+            parts = token_str.split("</think>")
+            if len(parts) > 0:
+                print(parts[0] + "</think>", end='', flush=True)
+                if len(parts) > 1:
+                    print(LIGHT_BLUE + parts[1], end='', flush=True)
+        else:
+            if not self.thinking:
+                print(LIGHT_BLUE + token_str, end='', flush=True)
+            else:
+                print(token_str, end='', flush=True)
+
+    def _print_worker(self):
+        """Worker thread that takes token_ids from the queue."""
+        while not self.stop_event.is_set():
+            try:
+                token_id = self.token_queue.get(timeout=0.01)
+                with self.lock:
+                    self.decoding_buffer.append(token_id)
+                self.token_queue.task_done()
+            except queue.Empty:
+                continue
+            except Exception as e:
+                print(f"\nError: Token printer error: {str(e)}")
+                break
+
+    def stop(self):
+        """Stop the printer thread."""
+        if self.thread and self.thread.is_alive():
+            self.stop_event.set()
+            try:
+                self.thread.join(timeout=1.0)
+            except Exception:
+                pass
+            print(RESET_COLOR)  # Reset color at the end
+        return self.buffer
+
+    def set_timing(self, prefill_time, inference_time, context_pos):
+        """Set timing information."""
+        self.prefill_time = prefill_time
+        self.inference_time = inference_time
+        self.context_pos = context_pos
+
+def parse_model_path(path):
+    """Parse model path and return full path with .mlmodelc or .mlpackage extension."""
+    path = Path(path)
+    
+    # If path exists exactly as specified, return it
+    if path.exists():
+        return str(path)
+        
+    # Try with both extensions
+    candidates = [
+        path,  # Original path
+        path.with_suffix('.mlmodelc'),  # With .mlmodelc
+        path.with_suffix('.mlpackage'),  # With .mlpackage
+        Path(str(path) + '.mlmodelc'),  # Handle case where extension is included
+        Path(str(path) + '.mlpackage')
+    ]
+    
+    # Try all possible paths
+    for candidate in candidates:
+        if candidate.exists():
+            print(f"Found model at: {candidate}")
+            return str(candidate)
+    
+    # If embeddings with LUT suffix not found, try without LUT suffix
+    if "_lut" in str(path) and "embeddings" in str(path):
+        print(f"Failed to find {path}, trying without LUT suffix...")
+        # Remove LUT suffix
+        path_no_lut = str(path).split("_lut")[0]
+        path_no_lut = Path(path_no_lut)
+        
+        # Try candidates without LUT suffix
+        candidates_no_lut = [
+            path_no_lut,
+            path_no_lut.with_suffix('.mlmodelc'),
+            path_no_lut.with_suffix('.mlpackage'),
+            Path(str(path_no_lut) + '.mlmodelc'),
+            Path(str(path_no_lut) + '.mlpackage')
+        ]
+        
+        for candidate in candidates_no_lut:
+            if candidate.exists():
+                print(f"Found model at: {candidate}")
+                return str(candidate)
+        
+        # Add no-LUT candidates to the list for error reporting
+        candidates.extend(candidates_no_lut)
+
+    # If FFN path isn't chunked, try to find chunked variants.
+    path_str = str(path)
+    base_str = str(path.with_suffix('')) if path.suffix in ('.mlmodelc', '.mlpackage') else path_str
+    if "_chunk_" not in base_str:
+        chunk_pattern = f"{base_str}_chunk_*of*"
+        chunk_candidates = sorted(glob.glob(chunk_pattern + ".mlmodelc"))
+        if not chunk_candidates:
+            chunk_candidates = sorted(glob.glob(chunk_pattern + ".mlpackage"))
+        if chunk_candidates:
+            print(f"Found model at: {chunk_candidates[0]}")
+            return str(Path(chunk_candidates[0]))
+        candidates.extend([Path(p) for p in sorted(glob.glob(chunk_pattern + ".mlmodelc"))])
+        candidates.extend([Path(p) for p in sorted(glob.glob(chunk_pattern + ".mlpackage"))])
+            
+    # If we get here, no valid path was found
+    print("\nError: Model not found. Tried following paths:")
+    for candidate in candidates:
+        print(f"  {candidate}")
+    raise FileNotFoundError(f"Model not found: {path}")
+
+def build_stop_token_ids(tokenizer):
+    """Collect token IDs that should stop generation."""
+    def _get_token_id_if_present(token_str):
+        if not token_str:
+            return None
+        # Avoid calling get_vocab() as it can segfault on some tokenizers (e.g., Gemma)
+        # Use convert_tokens_to_ids() directly instead
+        try:
+            token_id = tokenizer.convert_tokens_to_ids(token_str)
+            if isinstance(token_id, list):
+                if len(token_id) == 1:
+                    token_id = token_id[0]
+                else:
+                    return None
+            if token_id is None:
+                return None
+            if tokenizer.unk_token_id is not None and token_id == tokenizer.unk_token_id:
+                return None
+            return token_id
+        except Exception:
+            return None
+
+    stop_ids = set()
+    eos_token_ids = tokenizer.eos_token_id
+    if isinstance(eos_token_ids, list):
+        stop_ids.update(eos_token_ids)
+    elif eos_token_ids is not None:
+        stop_ids.add(eos_token_ids)
+
+    for token_str in ("<|endoftext|>", "<end_of_turn>", "<|eot_id|>"):
+        token_id = _get_token_id_if_present(token_str)
+        if token_id is not None:
+            stop_ids.add(token_id)
+
+    return stop_ids
+
+def format_manual_prompt(messages):
+    """Format a plain text prompt when no chat template is available."""
+    system = None
+    turns = []
+    pending_user = None
+    for message in messages:
+        role = message.get("role")
+        content = message.get("content", "")
+        if role == "system":
+            system = content
+        elif role == "user":
+            pending_user = content
+        elif role == "assistant":
+            if pending_user is not None:
+                turns.append((pending_user, content))
+                pending_user = None
+
+    def _format_inst(user_text, system_text):
+        if system_text:
+            return f"[INST] <<SYS>>\n{system_text}\n<</SYS>>\n\n{user_text} [/INST]"
+        return f"[INST] {user_text} [/INST]"
+
+    blocks = []
+    for user_text, assistant_text in turns:
+        blocks.append(f"{_format_inst(user_text, system)} {assistant_text}")
+        system = None  # Only apply system prompt once.
+    if pending_user is not None:
+        blocks.append(_format_inst(pending_user, system))
+    return "\n".join(blocks)
+
+def parse_ffn_filename(path):
+    """Parse FFN model filename to extract chunk information."""
+    path = Path(path)
+    pattern = r'FFN_PF.*_chunk_(\d+)of(\d+)'
+    match = re.search(pattern, path.name)
+    
+    if match:
+        current_chunk = int(match.group(1))
+        total_chunks = int(match.group(2))
+        return current_chunk, total_chunks
+    return None, None
+
+def find_all_chunks(base_path):
+    """Find all chunk files matching the base FFN path pattern."""
+    path = Path(base_path)
+    pattern = re.sub(r'_chunk_\d+of\d+', '_chunk_*', str(path))
+    return sorted(glob.glob(pattern))
+
+def load_model(path, function_name=None, compute_unit=None):
+    """Load a CoreML model, handling both .mlmodelc and .mlpackage formats."""
+    path = Path(path)
+    if compute_unit is None:
+        compute_unit = ct.ComputeUnit.CPU_AND_NE
+    
+    try:
+        if path.suffix == '.mlmodelc':
+            # For compiled models (.mlmodelc), use CompiledMLModel
+            if function_name:
+                return ct.models.CompiledMLModel(str(path), compute_unit, function_name=function_name)
+            else:
+                return ct.models.CompiledMLModel(str(path), compute_unit)
+        else:
+            # For packages (.mlpackage)
+            if function_name:
+                return ct.models.MLModel(str(path), function_name=function_name)
+            else:
+                return ct.models.MLModel(str(path))
+                
+    except RuntimeError as e:
+        if "valid manifest does not exist" in str(e):
+            print(f"\nError: Could not load compiled model at {path}")
+            print("This might be because:")
+            print("1. The model is not properly compiled")
+            print("2. The model was compiled for a different OS version")
+            print("3. The model needs to be recompiled")
+            print("\nTry using the .mlpackage version instead, or recompile the model.")
+        raise
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Full Chat with CoreML LLaMA with context window shifting, gil resolved (c) 2025 Anemll')
+
+    # Add meta.yaml option
+    parser.add_argument('--meta', type=str, help='Path to meta.yaml to load all parameters')
+
+    # Add existing arguments
+    parser.add_argument('--d', '--dir', type=str, default='.',
+                       help='Directory containing model files (default: current directory)')
+    parser.add_argument('--embed', type=str, required=False,
+                       help='Path to embeddings model (relative to --dir)')
+    parser.add_argument('--ffn', type=str, required=False,
+                       help='Path to FFN model (can be chunked, relative to --dir)')
+    parser.add_argument('--lmhead', type=str, required=False,
+                       help='Path to LM head model (relative to --dir)')
+    parser.add_argument('--tokenizer', type=str, required=False,
+                       help='Path to tokenizer')
+
+    # Add new argument for auto-generation
+    parser.add_argument('--prompt', type=str,
+                       help='If specified, run once with this prompt and exit')
+    parser.add_argument('--max-tokens', type=int,
+                       help='Maximum number of tokens to generate')
+
+    # Add no-warmup flag
+    parser.add_argument('--nw', action='store_true',
+                       help='Skip warmup phase')
+
+    # Add debug level
+    parser.add_argument('--debug-level', type=int, default=0,
+                       help='Debug level (0=none, 1=print prompts, 2=more verbose)')
+
+    # Add CPU-only mode
+    parser.add_argument('--cpu', action='store_true',
+                       help='Run on CPU only (no ANE/GPU)')
+
+    # Model configuration
+    parser.add_argument('--context-length', type=int,
+                       help='Context length for the model (default: 512), if not provided, it will be detected from the model directory name ctxNUMBER')
+    parser.add_argument('--batch-size', type=int,
+                       help='Batch size for prefill (default: 64)')
+    parser.add_argument('--split-lm-head', type=int,
+                       help='Number of logits splits from LM head (default: 8 for llama, 16 for qwen)')
+
+    args = parser.parse_args()
+
+    # If meta.yaml is provided, load parameters from it
+    if args.meta:
+        try:
+            with open(args.meta, 'r') as f:
+                meta = yaml.safe_load(f)
+            params = meta['model_info']['parameters']
+
+            # Set model directory to meta.yaml directory if not specified
+            if not args.d or args.d == '.':
+                args.d = str(Path(args.meta).parent)
+
+            # Check if this is a monolithic model
+            model_type = meta['model_info'].get('model_type', 'chunked')
+            args.is_monolithic = (model_type == 'monolithic')
+
+            if args.is_monolithic:
+                # Monolithic model configuration
+                prefix = params.get('model_prefix', 'qwen')
+                lut_bits = params.get('lut_bits', 'none')
+                lut_suffix = f"_lut{lut_bits}" if lut_bits != 'none' else ''
+
+                # Set monolithic model path
+                args.monolithic_model = params.get('monolithic_model', f'{prefix}_monolithic_full{lut_suffix}.mlmodelc')
+
+                # Set other parameters
+                if args.context_length is None:
+                    args.context_length = int(params['context_length'])
+                if args.batch_size is None:
+                    args.batch_size = int(params['batch_size'])
+                args.num_chunks = 1  # Monolithic has no chunks
+
+                # state_length for split cache models (defaults to context_length if not specified)
+                args.state_length = int(params.get('state_length', args.context_length))
+
+                # Check for argmax_in_model flag (model outputs argmax instead of logits)
+                args.argmax_in_model = params.get('argmax_in_model', False)
+
+                # Set split_lm_head, but allow CLI override
+                if args.split_lm_head is None:
+                    if 'split_lm_head' in params:
+                        args.split_lm_head = int(params['split_lm_head'])
+                    else:
+                        args.split_lm_head = 16 if 'qwen' in prefix.lower() else 8
+
+                # Set tokenizer path
+                if not args.tokenizer:
+                    if 'tokenizer_path' in params:
+                        args.tokenizer = params['tokenizer_path']
+                    else:
+                        args.tokenizer = args.d
+
+                print(f"\nLoaded MONOLITHIC model from {args.meta}:")
+                print(f"  Model: {args.monolithic_model}")
+                print(f"  Context Length: {args.context_length}")
+                print(f"  State Length: {args.state_length}")
+                print(f"  Batch Size: {args.batch_size}")
+                print(f"  Split LM Head: {args.split_lm_head}")
+                print(f"  Argmax in Model: {args.argmax_in_model}")
+                print(f"  Models Directory: {args.d}")
+            else:
+                # Standard chunked model configuration
+                args.is_monolithic = False
+                # Build model paths based on parameters
+                prefix = params.get('model_prefix', 'llama')  # Default to 'llama' if not specified
+                lut_ffn = f"_lut{params['lut_ffn']}" if params['lut_ffn'] != 'none' else ''
+                lut_lmhead = f"_lut{params['lut_lmhead']}" if params['lut_lmhead'] != 'none' else ''
+                lut_embeddings = f"_lut{params['lut_embeddings']}" if params['lut_embeddings'] != 'none' else ''
+                num_chunks = int(params['num_chunks'])
+
+                # Set model paths if not specified
+                if not args.lmhead:
+                    args.lmhead = f'{prefix}_lm_head{lut_lmhead}'
+                if not args.embed:
+                    args.embed = f'{prefix}_embeddings{lut_embeddings}'  # Changed from lm_head to embeddings
+                if not args.ffn:
+                    args.ffn = f'{prefix}_FFN_PF{lut_ffn}_chunk_01of{num_chunks:02d}'
+                if not args.tokenizer:
+                    args.tokenizer = args.d
+
+                # Set other parameters if not overridden by command line
+                if args.context_length is None:
+                    args.context_length = int(params['context_length'])
+                if args.batch_size is None:
+                    args.batch_size = int(params['batch_size'])
+                args.num_chunks = num_chunks
+
+                # Parse split_lm_head parameter from meta.yaml, but allow CLI override
+                if args.split_lm_head is None:
+                    if 'split_lm_head' in params:
+                        args.split_lm_head = int(params['split_lm_head'])
+                    else:
+                        args.split_lm_head = 8  # Default value
+
+                # Check for argmax_in_model flag (for chunked models)
+                args.argmax_in_model = params.get('argmax_in_model', False)
+
+                # sliding_window for Gemma3 rotation support (default 512 for Gemma3)
+                # Only set if the model has a sliding window configured or if prefix is gemma3
+                if 'sliding_window' in params:
+                    args.sliding_window = int(params['sliding_window'])
+                elif prefix.lower().startswith('gemma3'):
+                    args.sliding_window = 512  # Default Gemma3 sliding window
+                else:
+                    args.sliding_window = None  # No rotation for other models
+
+                print(f"\nLoaded parameters from {args.meta}:")
+                print(f"  Context Length: {args.context_length}")
+                print(f"  Batch Size: {args.batch_size}")
+                print(f"  Num Chunks: {args.num_chunks}")
+                print(f"  Split LM Head: {args.split_lm_head}")
+                print(f"  Argmax in Model: {args.argmax_in_model}")
+                print(f"  Models Directory: {args.d}")
+                print(f"  Embeddings: {args.embed}")
+                print(f"  LM Head: {args.lmhead}")
+                print(f"  FFN: {args.ffn}")
+
+        except Exception as e:
+            print(f"\nError loading meta.yaml: {str(e)}")
+            sys.exit(1)
+    else:
+        # If no meta.yaml, set defaults
+        args.is_monolithic = False
+
+    return args
+
+def load_metadata(model,args):
+    # Extract metadata and config parameters
+    metadata = {}
+    if hasattr(model, 'user_defined_metadata'):
+        meta = model.user_defined_metadata
+        
+        # Extract key parameters with defaults
+        metadata['context_length'] = int(meta.get('com.anemll.context_length', 512))
+        metadata['state_length'] = int(meta.get('com.anemll.state_length', metadata['context_length']))  # Added state_length
+        metadata['batch_size'] = int(meta.get('com.anemll.batch_size', 64))
+        metadata['lut_bits'] = int(meta.get('com.anemll.lut_bits', 0))
+        metadata['num_chunks'] = int(meta.get('com.anemll.num_chunks', 1))
+
+        # If meta.yaml/args provide overrides, prefer those for reporting/usage
+        if getattr(args, 'context_length', None) is not None:
+            metadata['context_length'] = int(args.context_length)
+        if getattr(args, 'state_length', None) is not None:
+            metadata['state_length'] = int(args.state_length)
+        
+        print("\nExtracted Parameters:")
+        print(f"  Context Length: {metadata['context_length']}")
+        print(f"  State Length: {metadata['state_length']}")
+        print(f"  Prefill Batch Size: {metadata['batch_size']}")
+        print(f"  LUT Bits: {metadata['lut_bits']}")
+        print(f"  Number of Chunks: {metadata['num_chunks']}")
+        
+        # Print model info
+        print("\nModel Info:")
+        if 'com.anemll.info' in meta:
+            print(f"  {meta['com.anemll.info']}")
+        if 'com.github.apple.coremltools.version' in meta:
+            print(f"  CoreML Tools: {meta['com.github.apple.coremltools.version']}")
+        
+        # Print model input/output shapes
+        print("\nModel Shapes:")
+        if hasattr(model, 'input_description'):
+            print("  Inputs:")
+            try:
+                if hasattr(model.input_description, 'items'):
+                    for name, desc in model.input_description.items():
+                        print(f"    {name}: {desc}")
+                else:
+                    print(f"    {model.input_description}")
+            except:
+                print(f"    Input description: {type(model.input_description)}")
+        if hasattr(model, 'output_description'):
+            print("  Outputs:")
+            try:
+                if hasattr(model.output_description, 'items'):
+                    for name, desc in model.output_description.items():
+                        print(f"    {name}: {desc}")
+                else:
+                    print(f"    {model.output_description}")
+            except:
+                print(f"    Output description: {type(model.output_description)}")
+    else:
+        print("\nWarning: No metadata found in model")
+
+        # Check if model directory name contains context length pattern (ctxXXX)
+        ctx_len = 512
+        if args.context_length is  None:
+            import re
+            ctx_match = re.search(r'ctx(\d+)', str(args.d))
+            if ctx_match:
+                ctx_len0 = int(ctx_match.group(1))
+                if 512 <= ctx_len0 <= 8096:
+                    ctx_len = ctx_len0
+                    print(f"\nDetected context length {ctx_len} from directory name")
+            else:
+                print(f"\nWarning: No context length found in directory  {ctx_len} from directory name {args.d}")
+        else:
+            ctx_len = args.context_length
+
+        # Use defaults or values from args
+        metadata['context_length'] = ctx_len
+        metadata['state_length'] = ctx_len
+        # Get batch size from args or use default
+        metadata['batch_size'] = getattr(args, 'batch_size', 64)
+        metadata['lut_bits'] = 4
+        metadata['num_chunks'] = getattr(args, 'num_chunks', 4)
+        print("\nUsing parameters:")
+        print(f"  Context Length: {metadata['context_length']}")
+        print(f"  State Length: {metadata['state_length']}")
+        print(f"  Prefill Batch Size: {metadata['batch_size']}")
+        print(f"  LUT Bits: {metadata['lut_bits']}")
+        print(f"  Number of Chunks: {metadata['num_chunks']}")
+
+    # Override with values from args if they exist
+    if hasattr(args, 'batch_size') and args.batch_size is not None:
+        metadata['batch_size'] = args.batch_size
+        print(f"\nOverriding batch size from args: {args.batch_size}")
+    if hasattr(args, 'num_chunks') and args.num_chunks is not None:
+        metadata['num_chunks'] = args.num_chunks
+        print(f"\nOverriding num chunks from args: {args.num_chunks}")
+    
+    return metadata
+    
+def load_models(args,metadata):
+    """Load all required models and extract metadata."""
+    print("\nLoading models...")
+
+    # Determine compute unit
+    compute_unit = ct.ComputeUnit.CPU_ONLY if getattr(args, 'cpu', False) else ct.ComputeUnit.CPU_AND_NE
+    if getattr(args, 'cpu', False):
+        print("Running in CPU-only mode")
+
+    try:
+        # Load embeddings model
+        print("\nLoading embeddings model...")
+        embed_path = parse_model_path(args.embed)
+        print(f"Loading from: {embed_path}")
+        embed_model = load_model(embed_path, compute_unit=compute_unit)
+        print("Embeddings model loaded successfully")
+        metadata = load_metadata(embed_model,args)
+        
+
+        
+        # Load LM head model
+        print("\nLoading LM head model...")
+        lmhead_path = parse_model_path(args.lmhead)
+        print(f"Loading from: {lmhead_path}")
+        lmhead_model = load_model(lmhead_path, compute_unit=compute_unit)
+        print("LM head model loaded successfully")
+
+        # Parse FFN path and find chunks if needed
+        print("\nLoading FFN+PREFILL model(s)...")
+        ffn_path = parse_model_path(args.ffn)
+        chunk_no, total_chunks = parse_ffn_filename(ffn_path)
+
+        ffn_models = []
+        if chunk_no and total_chunks:
+            print(f"\nDetected chunked FFN+PREFILL model ({total_chunks} chunks)")
+            # Find and load all chunks
+            chunk_paths = find_all_chunks(ffn_path)
+            if len(chunk_paths) != total_chunks:
+                raise ValueError(f"Found {len(chunk_paths)} chunks but filename indicates {total_chunks} chunks")
+
+            for chunk_path in chunk_paths:
+                print(f"\nLoading FFN+PREFILL chunk: {Path(chunk_path).name}")
+                try:
+                    # For chunked models, we need both infer and prefill functions
+                    chunk_dict = {
+                        'infer': load_model(chunk_path, function_name='infer', compute_unit=compute_unit),
+                        'prefill': load_model(chunk_path, function_name='prefill', compute_unit=compute_unit)
+                    }
+                    # Try to load rotation functions only if context > sliding_window
+                    # If context_length <= sliding_window, rotation is never needed
+                    sliding_window = getattr(args, 'sliding_window', None)
+                    context_length = getattr(args, 'context_length', None)
+                    needs_rotation = (sliding_window is not None and
+                                     context_length is not None and
+                                     context_length > sliding_window)
+
+                    if needs_rotation:
+                        try:
+                            chunk_dict['infer_rotate'] = load_model(chunk_path, function_name='infer_rotate', compute_unit=compute_unit)
+                            chunk_dict['prefill_rotate'] = load_model(chunk_path, function_name='prefill_rotate', compute_unit=compute_unit)
+                            print("  Rotation functions loaded (4-function model)")
+                        except Exception:
+                            # Rotation functions not available - standard 2-function model
+                            pass
+                    elif sliding_window is not None:
+                        print(f"  Skipping rotation functions (context {context_length} <= sliding_window {sliding_window})")
+                    ffn_models.append(chunk_dict)
+                    print("Chunk loaded successfully")
+                except Exception as e:
+                    print(f"Error loading chunk {chunk_path}: {str(e)}")
+                    raise
+            metadata = load_metadata(ffn_models[0],args)
+
+        else:
+            print("\nLoading single FFN model...")
+            ffn_models.append(load_model(ffn_path, compute_unit=compute_unit))
+            print("FFN model loaded successfully")
+        
+        return embed_model, ffn_models, lmhead_model, metadata
+        
+    except Exception as e:
+        print(f"\nError loading models: {str(e)}")
+        print("\nPlease ensure all model files exist and are accessible.")
+        print("Expected files:")
+        print(f"  Embeddings: {args.embed}")
+        print(f"  LM Head: {args.lmhead}")
+        print(f"  FFN: {args.ffn}")
+        raise
+
+# At the top of the file, make this a default path
+
+def initialize_tokenizer(model_path=None):
+    """Initialize and configure the tokenizer."""
+    try:
+
+        
+        tokenizer = AutoTokenizer.from_pretrained(
+            str(model_path), 
+            use_fast=False,
+            trust_remote_code=True
+        )
+        
+        print("\nTokenizer Configuration:")
+        print(f"Tokenizer type: {type(tokenizer)}")
+        print(f"Tokenizer name: {tokenizer.__class__.__name__}")
+        print(f"Vocabulary size: {len(tokenizer)}")
+        print(f"Model max length: {tokenizer.model_max_length}")
+
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+            print("Set PAD token to EOS token")
+        
+        tokenizer.padding_side = "left"
+        
+        print(f"\nSpecial Tokens:")
+        print(f"PAD token: '{tokenizer.pad_token}' (ID: {tokenizer.pad_token_id})")
+        print(f"EOS token: '{tokenizer.eos_token}' (ID: {tokenizer.eos_token_id})")
+        print(f"BOS token: '{tokenizer.bos_token}' (ID: {tokenizer.bos_token_id})")
+        print(f"UNK token: '{tokenizer.unk_token}' (ID: {tokenizer.unk_token_id})")
+
+        return tokenizer
+        
+    except Exception as e:
+        print(f"\nError: Failed to load tokenizer from {model_path}")
+        print(f"Error details: {str(e)}")
+        print(f"Error type: {type(e)}")
+        print("\nThis code requires a Llama 3.2 model for chat template functionality.")
+        print("Please provide the path to a Llama 3.2 model directory.")
+        import traceback
+        traceback.print_exc()
+        raise
+
+
+
+def make_causal_mask(length, start):
+    """Create causal attention mask."""
+    mask = np.full((1, 1, length, length), -np.inf, dtype=np.float16)
+    row_indices = np.arange(length).reshape(length, 1)
+    col_indices = np.arange(length).reshape(1, length)
+    mask[:, :, col_indices <= (row_indices + start)] = 0
+    return mask
+
+def run_prefill(embed_model, ffn_models, input_ids, current_pos, context_length, batch_size, state, causal_mask, sliding_window=None):
+    """Run prefill on the input sequence.
+
+    For Gemma3 with 4-function models:
+    - Uses 'prefill' for positions < sliding_window
+    - Uses 'prefill_rotate' for positions >= sliding_window (if available)
+    """
+    #print(f"[DEBUG] Running prefill from 0 to {current_pos}")
+
+    # Check if rotation functions are available
+    has_rotation = isinstance(ffn_models[0], dict) and 'prefill_rotate' in ffn_models[0]
+
+    # If no rotation or no sliding_window, use standard prefill
+    if not has_rotation or sliding_window is None:
+        sliding_window = context_length  # Effectively disables rotation mode
+
+    # Process in batches
+    batch_pos = 0
+    while batch_pos < current_pos:
+        batch_end = min(batch_pos + batch_size, current_pos)
+        current_batch_size = batch_end - batch_pos
+
+        #print(f"[DEBUG] Prefill batch {batch_pos}-{batch_end} (size={current_batch_size})")
+
+        # Get current batch
+        batch_input = input_ids[:, batch_pos:batch_end]
+
+        # Pad to full batch size
+        batch_input = F.pad(
+            batch_input,
+            (0, batch_size - current_batch_size),
+            value=0
+        )
+
+        # Generate position IDs for this batch
+        position_ids = torch.arange(batch_pos, batch_pos + batch_size, dtype=torch.int32)
+
+        # Use the pre-initialized causal mask and extract the batch portion
+        batch_causal_mask = causal_mask[:, :, batch_pos:batch_pos + batch_size, :]
+
+        # Run embeddings
+        hidden_states = torch.from_numpy(
+            embed_model.predict({'input_ids': batch_input.numpy()})['hidden_states']
+        )
+
+        # Determine which prefill function to use based on position
+        # Use prefill_rotate for positions >= sliding_window
+        prefill_func_name = 'prefill_rotate' if batch_pos >= sliding_window and has_rotation else 'prefill'
+
+        # Run through FFN chunks
+        for ffn_model in ffn_models:
+            if isinstance(ffn_model, dict):
+                inputs = {
+                    'hidden_states': hidden_states.numpy(),
+                    'position_ids': position_ids.numpy(),
+                    'causal_mask': batch_causal_mask.numpy(),
+                    'current_pos': np.array([batch_pos], dtype=np.int32)
+                }
+                output = ffn_model[prefill_func_name].predict(inputs, state)
+                hidden_states = torch.from_numpy(output['output_hidden_states'])
+
+        batch_pos = batch_end
+
+    return torch.tensor([current_pos], dtype=torch.int32)
+
+def generate_next_token(embed_model, ffn_models, lmhead_model, input_ids, pos, context_length, state, causal_mask, metadata=None, temperature=0.0):
+    """Generate the next token.
+
+    For Gemma3 with 4-function models:
+    - Uses 'infer' for positions < sliding_window
+    - Uses 'infer_rotate' for positions >= sliding_window (if available)
+    """
+    sliding_window = metadata.get('sliding_window', None) if metadata else None
+
+    # Check if rotation functions are available
+    has_rotation = isinstance(ffn_models[0], dict) and 'infer_rotate' in ffn_models[0]
+
+    # Determine which infer function to use
+    # Use infer_rotate for positions >= sliding_window (0-indexed, so pos-1 is the actual position)
+    use_rotation = has_rotation and sliding_window is not None and (pos - 1) >= sliding_window
+    infer_func_name = 'infer_rotate' if use_rotation else 'infer'
+
+    # Get current token
+    current_token = input_ids[:, pos-1:pos]
+
+    # Run embeddings
+    hidden_states = torch.from_numpy(
+        embed_model.predict({'input_ids': current_token.numpy()})['hidden_states']
+    )
+
+    # Create masks
+    update_mask = torch.zeros((1, 1, context_length, 1), dtype=torch.float16)
+    update_mask[0, 0, pos-1, 0] = 1.0
+    position_ids = torch.tensor([pos-1], dtype=torch.int32)
+
+    # Use the pre-initialized causal mask and extract the single position portion
+    single_causal_mask = causal_mask[:, :, pos-1:pos, :]
+
+    # Run through FFN chunks
+    for ffn_model in ffn_models:
+        if isinstance(ffn_model, dict):
+            inputs = {
+                'hidden_states': hidden_states.numpy(),
+                'position_ids': position_ids.numpy(),
+                'causal_mask': single_causal_mask.numpy(),
+                'current_pos': position_ids.numpy()
+            }
+            # Add update_mask only if model expects it (older models)
+            try:
+                model_inputs = {inp.name for inp in ffn_model[infer_func_name].get_spec().description.input}
+            except Exception:
+                model_inputs = set()
+            if 'update_mask' in model_inputs:
+                inputs['update_mask'] = update_mask.numpy()
+            output = ffn_model[infer_func_name].predict(inputs, state)
+            hidden_states = torch.from_numpy(output['output_hidden_states'])
+    
+    # Run LM head and get next token
+    lm_output = lmhead_model.predict({'hidden_states': hidden_states.numpy()})
+
+    # Check if model uses argmax_in_model mode (outputs argmax_idx/argmax_val instead of logits)
+    argmax_in_model = metadata.get('argmax_in_model', False) if metadata else False
+
+    if argmax_in_model and 'argmax_idx' in lm_output:
+        # Model outputs argmax_idx and argmax_val (split across num_chunks chunks)
+        argmax_idx = lm_output['argmax_idx']  # shape: [num_chunks], LOCAL indices within chunk
+        argmax_val = lm_output['argmax_val']  # shape: [num_chunks], max logit values
+
+        # Flatten in case of extra dimensions
+        argmax_idx_flat = argmax_idx.flatten()
+        argmax_val_flat = argmax_val.flatten()
+
+        # Find the chunk with the highest value
+        best_chunk = int(np.argmax(argmax_val_flat))
+        local_idx = int(argmax_idx_flat[best_chunk])
+
+        # Calculate global token index: local_idx + chunk_offset
+        num_chunks = len(argmax_idx_flat)
+        vocab_size = 262144  # Standard for Gemma3
+        chunk_size = vocab_size // num_chunks
+        next_token = local_idx + (best_chunk * chunk_size)
+
+        return next_token
+
+    # Warn if argmax expected but not found
+    if argmax_in_model and 'argmax_idx' not in lm_output:
+        print(f"\n[WARNING] argmax_in_model=True but model outputs: {list(lm_output.keys())}")
+        print("Model may need reconversion with --argmax flag")
+
+    if 'logits1' in lm_output:
+        logit_indices = [
+            int(k[6:]) for k in lm_output.keys()
+            if k.startswith("logits") and k[6:].isdigit()
+        ]
+        max_available = max(logit_indices) if logit_indices else 0
+        num_logits = (
+            metadata.get('split_lm_head', metadata.get('num_logits', max_available or 8))
+            if metadata
+            else (max_available or 8)
+        )
+        if max_available and num_logits > max_available:
+            num_logits = max_available
+        logits_parts = []
+        for i in range(1, num_logits + 1):
+            key = f'logits{i}'
+            if key in lm_output:
+                logits_parts.append(torch.from_numpy(lm_output[key]))
+        logits = torch.cat(logits_parts, dim=-1)
+    else:
+        logits = torch.from_numpy(lm_output['output_logits'])
+    
+    if temperature > 0:
+        logits = logits / temperature
+        probs = F.softmax(logits[0, -1, :], dim=-1)
+        next_token = torch.multinomial(probs, num_samples=1).item()
+    else:
+        next_token = torch.argmax(logits[0, -1, :]).item()
+    
+    return next_token
+
+def create_unified_state(ffn_models, context_length):
+    """Create unified KV cache state for transformer."""
+    if isinstance(ffn_models[0], dict):
+        # Use first FFN model's prefill function to create state
+        state = ffn_models[0]['prefill'].make_state()
+        print(f"\nCreated unified transformer state for {len(ffn_models)} chunks")
+        return state
+    else:
+        state = ffn_models[0].make_state()
+        print("\nCreated unified transformer state")
+        return state
+
+def initialize_causal_mask(context_length):
+    """Initialize causal mask for transformer attention."""
+    causal_mask = make_causal_mask(context_length, 0)
+    causal_mask = torch.tensor(causal_mask, dtype=torch.float16)
+    print(f"\nInitialized causal mask for context length {context_length}")
+    return causal_mask
+
+
+def load_monolithic_model(args, metadata):
+    """Load monolithic model with infer, infer_rotate, prefill, and prefill_rotate functions."""
+    print("\nLoading monolithic model...")
+
+    # Determine compute unit
+    compute_unit = ct.ComputeUnit.CPU_ONLY if getattr(args, 'cpu', False) else ct.ComputeUnit.CPU_AND_NE
+    if getattr(args, 'cpu', False):
+        print("Running in CPU-only mode")
+
+    model_path = str(Path(args.d) / args.monolithic_model)
+    model_path = parse_model_path(model_path)
+
+    print(f"Loading from: {model_path}")
+
+    # Load all functions
+    infer_model = load_model(model_path, function_name='infer', compute_unit=compute_unit)
+    prefill_model = load_model(model_path, function_name='prefill', compute_unit=compute_unit)
+
+    # Try to load infer_rotate (optional, for models with split cache rotation)
+    infer_rotate_model = None
+    try:
+        infer_rotate_model = load_model(model_path, function_name='infer_rotate', compute_unit=compute_unit)
+    except Exception as e:
+        print(f"  Note: infer_rotate not available - using infer for all positions")
+
+    # Try to load prefill_rotate (optional, for long context prefill with rotation)
+    prefill_rotate_model = None
+    try:
+        prefill_rotate_model = load_model(model_path, function_name='prefill_rotate', compute_unit=compute_unit)
+    except Exception as e:
+        pass  # prefill_rotate is optional
+
+    # Report loaded functions
+    functions = ["infer", "prefill"]
+    if infer_rotate_model:
+        functions.insert(1, "infer_rotate")
+    if prefill_rotate_model:
+        functions.append("prefill_rotate")
+    print(f"Monolithic model loaded successfully ({' + '.join(functions)} functions)")
+
+    # Extract metadata from model
+    metadata = load_metadata(infer_model, args)
+
+    return infer_model, infer_rotate_model, prefill_model, prefill_rotate_model, metadata
+
+
+def run_monolithic_prefill(model, input_ids, context_pos, context_length, batch_size, state, causal_mask):
+    """Run prefill on monolithic model."""
+    batch_pos = 0
+    while batch_pos < context_pos:
+        batch_end = min(batch_pos + batch_size, context_pos)
+        current_batch_size = batch_end - batch_pos
+
+        # Get current batch
+        batch_input = input_ids[:, batch_pos:batch_end]
+
+        # Pad to full batch size
+        batch_input = F.pad(batch_input, (0, batch_size - current_batch_size), value=0)
+
+        # Generate position IDs for full batch size
+        position_ids = torch.arange(batch_pos, batch_pos + batch_size, dtype=torch.int32)
+        batch_causal_mask = causal_mask[:, :, batch_pos:batch_pos + batch_size, :]
+
+        # Run monolithic prefill (input_ids -> logits directly)
+        inputs = {
+            'input_ids': batch_input.numpy().astype(np.int32),
+            'position_ids': position_ids.numpy().astype(np.int32),
+            'causal_mask': batch_causal_mask.numpy().astype(np.float16),
+            'current_pos': np.array([batch_pos], dtype=np.int32)
+        }
+        output = model.predict(inputs, state)
+        # We don't need the output logits for prefill, just updating KV cache
+
+        batch_pos = batch_end
+
+    return torch.tensor([context_pos], dtype=torch.int32)
+
+
+def run_monolithic_prefill_with_rotation(prefill_model, prefill_rotate_model, input_ids, context_pos,
+                                         context_length, batch_size, state, causal_mask, sliding_window,
+                                         infer_rotate_model=None):
+    """Run prefill with rotation support for long contexts.
+
+    When context_pos > sliding_window, this splits the prefill into two phases:
+    - Phase 1: Fill mode (prefill_model) for positions 0 to sliding_window-1
+    - Phase 2: Rotation mode (prefill_rotate_model) for positions sliding_window to context_pos-1
+
+    If prefill_rotate_model is None or context_pos <= sliding_window, falls back to standard prefill.
+    """
+    # If no rotation model or short context, use standard prefill
+    if prefill_rotate_model is None or context_pos <= sliding_window:
+        return run_monolithic_prefill(prefill_model, input_ids, context_pos, context_length,
+                                      batch_size, state, causal_mask)
+
+    # Phase 1: Fill mode for positions 0 to sliding_window-1
+    print_system(f"Prefill Phase 1: Fill mode (0 to {sliding_window-1})")
+    batch_pos = 0
+    while batch_pos < sliding_window:
+        batch_end = min(batch_pos + batch_size, sliding_window)
+        current_batch_size = batch_end - batch_pos
+
+        batch_input = input_ids[:, batch_pos:batch_end]
+        batch_input = F.pad(batch_input, (0, batch_size - current_batch_size), value=0)
+
+        position_ids = torch.arange(batch_pos, batch_pos + batch_size, dtype=torch.int32)
+        batch_causal_mask = causal_mask[:, :, batch_pos:batch_pos + batch_size, :]
+
+        inputs = {
+            'input_ids': batch_input.numpy().astype(np.int32),
+            'position_ids': position_ids.numpy().astype(np.int32),
+            'causal_mask': batch_causal_mask.numpy().astype(np.float16),
+            'current_pos': np.array([batch_pos], dtype=np.int32)
+        }
+        prefill_model.predict(inputs, state)
+        batch_pos = batch_end
+
+    # Phase 2: Rotation mode for positions sliding_window to context_pos-1
+    print_system(f"Prefill Phase 2: Rotation mode ({sliding_window} to {context_pos-1})")
+    batch_pos = sliding_window
+    # Process full batches with prefill_rotate
+    while batch_pos + batch_size <= context_pos:
+        batch_end = batch_pos + batch_size
+
+        batch_input = input_ids[:, batch_pos:batch_end]
+        position_ids = torch.arange(batch_pos, batch_end, dtype=torch.int32)
+        batch_causal_mask = causal_mask[:, :, batch_pos:batch_end, :]
+
+        inputs = {
+            'input_ids': batch_input.numpy().astype(np.int32),
+            'position_ids': position_ids.numpy().astype(np.int32),
+            'causal_mask': batch_causal_mask.numpy().astype(np.float16),
+            'current_pos': np.array([batch_pos], dtype=np.int32)
+        }
+        prefill_rotate_model.predict(inputs, state)
+        batch_pos = batch_end
+
+    # Handle remainder tokens without padding (token-by-token rotation)
+    if batch_pos < context_pos:
+        if infer_rotate_model is not None:
+            print_system(f"Prefill Phase 2b: Rotation single-token fill ({batch_pos} to {context_pos-1})")
+            while batch_pos < context_pos:
+                token = input_ids[:, batch_pos:batch_pos + 1]
+                position_ids = torch.tensor([batch_pos], dtype=torch.int32)
+                single_causal_mask = causal_mask[:, :, batch_pos:batch_pos + 1, :]
+
+                inputs = {
+                    'input_ids': token.numpy().astype(np.int32),
+                    'position_ids': position_ids.numpy().astype(np.int32),
+                    'causal_mask': single_causal_mask.numpy().astype(np.float16),
+                    'current_pos': position_ids.numpy().astype(np.int32)
+                }
+                infer_rotate_model.predict(inputs, state)
+                batch_pos += 1
+        else:
+            # Fallback to padded batch if infer_rotate is unavailable
+            batch_end = context_pos
+            current_batch_size = batch_end - batch_pos
+            batch_input = input_ids[:, batch_pos:batch_end]
+            batch_input = F.pad(batch_input, (0, batch_size - current_batch_size), value=0)
+
+            position_ids = torch.arange(batch_pos, batch_pos + batch_size, dtype=torch.int32)
+            batch_causal_mask = causal_mask[:, :, batch_pos:batch_pos + batch_size, :]
+
+            inputs = {
+                'input_ids': batch_input.numpy().astype(np.int32),
+                'position_ids': position_ids.numpy().astype(np.int32),
+                'causal_mask': batch_causal_mask.numpy().astype(np.float16),
+                'current_pos': np.array([batch_pos], dtype=np.int32)
+            }
+            prefill_rotate_model.predict(inputs, state)
+
+    return torch.tensor([context_pos], dtype=torch.int32)
+
+
+def generate_next_token_monolithic(model, input_ids, pos, context_length, metadata, state, causal_mask, temperature=0.0):
+    """Generate next token using monolithic model."""
+    # Get current token
+    current_token = input_ids[:, pos-1:pos]  # [1, 1]
+
+    # Create inputs
+    position_ids = torch.tensor([pos-1], dtype=torch.int32)
+    single_causal_mask = causal_mask[:, :, pos-1:pos, :]
+
+    # Run monolithic infer
+    inputs = {
+        'input_ids': current_token.numpy().astype(np.int32),
+        'position_ids': position_ids.numpy().astype(np.int32),
+        'causal_mask': single_causal_mask.numpy().astype(np.float16),
+        'current_pos': position_ids.numpy().astype(np.int32)
+    }
+    output = model.predict(inputs, state)
+
+    # Check if model uses argmax_in_model mode (outputs 2 tensors instead of logits)
+    argmax_in_model = metadata.get('argmax_in_model', False)
+
+    if argmax_in_model and 'argmax_idx' in output:
+        # Model outputs argmax_idx and argmax_val (split across num_chunks chunks)
+        # Each chunk covers vocab_size / num_chunks tokens
+        argmax_idx = output['argmax_idx']  # shape: [num_chunks], LOCAL indices within chunk
+        argmax_val = output['argmax_val']  # shape: [num_chunks], max logit values
+
+        # Flatten in case of extra dimensions
+        argmax_idx_flat = argmax_idx.flatten()
+        argmax_val_flat = argmax_val.flatten()
+
+        # Find the chunk with the highest value
+        best_chunk = int(np.argmax(argmax_val_flat))
+        local_idx = int(argmax_idx_flat[best_chunk])
+
+        # Calculate global token index: local_idx + chunk_offset
+        # Each chunk covers vocab_size / num_chunks tokens (e.g., 16384 for 262k vocab / 16 chunks)
+        num_chunks = len(argmax_idx_flat)
+        vocab_size = 262144  # Standard for Gemma3
+        chunk_size = vocab_size // num_chunks
+        next_token = local_idx + (best_chunk * chunk_size)
+
+        return next_token
+
+    # Get number of logits from metadata
+    num_logits = metadata.get('split_lm_head', metadata.get('num_logits', 8))
+
+    # Combine logits1-N if they exist
+    if 'logits1' in output:
+        logit_indices = [
+            int(k[6:]) for k in output.keys()
+            if k.startswith("logits") and k[6:].isdigit()
+        ]
+        max_available = max(logit_indices) if logit_indices else 0
+        if max_available and num_logits > max_available:
+            num_logits = max_available
+        logits_parts = []
+        for i in range(1, num_logits + 1):
+            key = f'logits{i}'
+            if key in output:
+                logits_parts.append(torch.from_numpy(output[key]))
+        logits = torch.cat(logits_parts, dim=-1)
+    elif 'logits' in output:
+        logits = torch.from_numpy(output['logits'])
+    else:
+        # Try other common output names
+        for key in output.keys():
+            if 'logit' in key.lower():
+                logits = torch.from_numpy(output[key])
+                break
+
+    # Apply temperature and sample
+    if temperature > 0:
+        logits = logits / temperature
+        probs = F.softmax(logits[0, -1, :], dim=-1)
+        next_token = torch.multinomial(probs, num_samples=1).item()
+    else:
+        next_token = torch.argmax(logits[0, -1, :]).item()
+
+    return next_token
+
+
+def chat_loop_monolithic(infer_model, prefill_model, tokenizer, metadata, state, causal_mask, auto_prompt=None, warmup=False, max_tokens=None, infer_rotate_model=None, prefill_rotate_model=None):
+    """Chat loop for monolithic models with full conversation history.
+
+    Args:
+        infer_model: Model for single-token inference (fill mode, pos < sliding_window)
+        prefill_model: Model for batch prefill (fill mode, for positions 0 to sliding_window-1)
+        tokenizer: Tokenizer
+        metadata: Model metadata dict
+        state: CoreML state object
+        causal_mask: Causal mask tensor
+        auto_prompt: Optional auto-prompt string
+        warmup: If True, skip output
+        max_tokens: Maximum tokens to generate
+        infer_rotate_model: Optional model for single-token inference with cache rotation
+                           (rotation mode, pos >= sliding_window). If None, uses infer_model.
+        prefill_rotate_model: Optional model for batch prefill with cache rotation
+                              (rotation mode, for positions >= sliding_window). If None,
+                              uses prefill_model for all positions (legacy behavior).
+    """
+    global THINKING_MODE
+    global DEBUG_LEVEL
+    context_length = metadata.get('context_length')
+    state_length = metadata.get('state_length', context_length)
+    sliding_window = metadata.get('sliding_window', 512)  # For switching between infer modes
+    batch_size = metadata.get('batch_size', 64)
+
+    # For split cache models, sliding window is typically 512 (local attention)
+    # Global attention layers can see up to state_length tokens
+    total_tokens_in_memory = 0  # Track total tokens processed in conversation
+    cumulative_tokens = 0  # Track all tokens ever processed (including trimmed)
+    turn_number = 0  # Track conversation turns
+
+    if not warmup:
+        print(f"\nUsing context length: {context_length}")
+        print(f"State length (global attention): {state_length}")
+        print(f"Sliding window (local attention): {sliding_window}")
+        if infer_rotate_model is not None:
+            print(f"Cache rotation: ENABLED (infer_rotate function available)")
+            print(f"  - pos < {sliding_window}: infer (fill mode)")
+            print(f"  - pos >= {sliding_window}: infer_rotate (rotation mode)")
+        else:
+            print(f"Cache rotation: NOT AVAILABLE (using infer for all positions)")
+        print("\nStarting chat session. Press Ctrl+D to exit.")
+        print("Type your message and press Enter to chat. Use /t to toggle thinking mode.")
+        print(f"Thinking mode is {'ON' if THINKING_MODE else 'OFF'}")
+
+    # Keep track of conversation history
+    conversation = []
+    stop_token_ids = build_stop_token_ids(tokenizer)
+    use_chat_template = False
+    try:
+        tokenizer.apply_chat_template([{"role": "user", "content": "test"}], return_tensors="pt")
+        use_chat_template = True
+        if not warmup:
+            print("\nUsing chat template for prompts")
+    except Exception:
+        if not warmup:
+            print("\nUsing manual formatting for prompts")
+
+    def _build_base_input_ids(messages, show_debug):
+        if use_chat_template:
+            base_input_ids = tokenizer.apply_chat_template(
+                messages,
+                return_tensors="pt",
+                add_generation_prompt=True
+            ).to(torch.int32)
+            if show_debug and DEBUG_LEVEL >= 1 and not warmup:
+                label = "Full prompt with thinking" if THINKING_MODE else "Full prompt"
+                print(f"\n{DARK_BLUE}Debug: {label}:{RESET_COLOR}")
+                print(tokenizer.decode(base_input_ids[0]))
+            return base_input_ids
+
+        prompt_text = format_manual_prompt(messages)
+        base_input_ids = tokenizer(
+            prompt_text, return_tensors="pt", add_special_tokens=True
+        ).input_ids.to(torch.int32)
+        if show_debug and DEBUG_LEVEL >= 1 and not warmup:
+            label = "Full prompt with thinking" if THINKING_MODE else "Full prompt"
+            print(f"\n{DARK_BLUE}Debug: {label}:{RESET_COLOR}")
+            print(prompt_text)
+        return base_input_ids
+    use_chat_template = False
+    try:
+        tokenizer.apply_chat_template([{"role": "user", "content": "test"}], return_tensors="pt")
+        use_chat_template = True
+        if not warmup:
+            print("\nUsing chat template for prompts")
+    except Exception:
+        if not warmup:
+            print("\nUsing manual formatting for prompts")
+
+    def _build_base_input_ids(messages, show_debug):
+        if use_chat_template:
+            base_input_ids = tokenizer.apply_chat_template(
+                messages,
+                return_tensors="pt",
+                add_generation_prompt=True
+            ).to(torch.int32)
+            if show_debug and DEBUG_LEVEL >= 1 and not warmup:
+                label = "Full prompt with thinking" if THINKING_MODE else "Full prompt"
+                print(f"\n{DARK_BLUE}Debug: {label}:{RESET_COLOR}")
+                print(tokenizer.decode(base_input_ids[0]))
+            return base_input_ids
+
+        prompt_text = format_manual_prompt(messages)
+        base_input_ids = tokenizer(
+            prompt_text, return_tensors="pt", add_special_tokens=True
+        ).input_ids.to(torch.int32)
+        if show_debug and DEBUG_LEVEL >= 1 and not warmup:
+            label = "Full prompt with thinking" if THINKING_MODE else "Full prompt"
+            print(f"\n{DARK_BLUE}Debug: {label}:{RESET_COLOR}")
+            print(prompt_text)
+        return base_input_ids
+
+    try:
+        while True:
+            try:
+                if not warmup:
+                    print(f"\n{LIGHT_GREEN}You{' (thinking)' if THINKING_MODE else ''}:{RESET_COLOR}", end=' ', flush=True)
+                if auto_prompt is not None:
+                    user_input = auto_prompt
+                    if not warmup:
+                        print(user_input)
+                else:
+                    user_input = input().strip()
+            except EOFError:
+                if not warmup:
+                    print("\nExiting chat...")
+                break
+
+            if not user_input:
+                continue
+
+            # Handle /t command
+            if user_input == "/t":
+                THINKING_MODE = not THINKING_MODE
+                print(f"Thinking mode {'ON' if THINKING_MODE else 'OFF'}")
+                continue
+
+            # Add user message to conversation
+            conversation.append({"role": "user", "content": user_input})
+
+            messages = conversation
+            if THINKING_MODE:
+                messages = [{"role": "system", "content": THINKING_PROMPT}] + conversation
+            base_input_ids = _build_base_input_ids(messages, show_debug=True)
+
+            # Check if we need to trim history
+            # Use state_length (global context) for split cache models, context_length otherwise
+            history_trimmed = False
+            original_size = base_input_ids.size(1)
+            while base_input_ids.size(1) > state_length - 100:  # Leave room for response
+                history_trimmed = True
+                # Remove oldest message pair (user + assistant)
+                if len(conversation) > 2:
+                    conversation = conversation[2:]  # Remove oldest pair
+                    messages = conversation
+                    if THINKING_MODE:
+                        messages = [{"role": "system", "content": THINKING_PROMPT}] + conversation
+                    base_input_ids = _build_base_input_ids(messages, show_debug=False)
+                else:
+                    # If only current message remains and still too long, truncate
+                    base_input_ids = base_input_ids[:, -state_length//2:]
+                    break
+
+            context_pos = base_input_ids.size(1)
+            turn_number += 1
+
+            if history_trimmed and not warmup:
+                print_system(f"History trimmed: {original_size} → {context_pos} tokens, {len(conversation)} msgs remaining")
+                # Note: KV cache state should be re-prefilled with trimmed context
+                # The prefill that runs next will update the cache appropriately
+
+            # Debug: show conversation state
+            if DEBUG_LEVEL >= 2 and not warmup:
+                print(f"{DARK_BLUE}[Debug] Turn {turn_number}: context_pos={context_pos}, conversation={len(conversation)} msgs{RESET_COLOR}")
+
+            # Pad sequence to context_size
+            input_ids = F.pad(
+                base_input_ids,
+                (0, context_length - context_pos),
+                value=0
+            )
+
+            # Initialize token printer and collect response
+            token_printer = TokenPrinter(tokenizer)
+            response_tokens = []
+            generation_start_time = time.time()
+
+            try:
+                # Run prefill on entire context (uses rotation for pos >= sliding_window if available)
+                current_pos = run_monolithic_prefill_with_rotation(
+                    prefill_model,
+                    prefill_rotate_model,
+                    input_ids,
+                    context_pos,
+                    context_length,
+                    batch_size,
+                    state,
+                    causal_mask,
+                    sliding_window,
+                    infer_rotate_model
+                )
+
+                if not warmup:
+                    print(f"\n{LIGHT_BLUE}Assistant:{RESET_COLOR}", end=' ', flush=True)
+
+                # Generation loop
+                pos = context_pos
+                tokens_generated = 0
+                max_tokens_this_turn = (
+                    max_tokens
+                    if max_tokens is not None
+                    else max(0, context_length - context_pos)
+                )
+                inference_start = time.time()  # Start inference timing
+
+                while True:
+                    # Check if we need to shift window
+                    if pos >= context_length - 2:
+                        if DEBUG_LEVEL >= 1:
+                            print_system(f"Context window reached {context_length} tokens; shifting context to continue.")
+                        # Calculate shift to maintain full batches
+                        batch_size = metadata.get('batch_size', 64)
+                        # Calculate max batches that fit in context
+                        max_batches = context_length // batch_size
+                        desired_batches = max(1, max_batches - 2)  # Leave room for new tokens
+                        new_size = min(desired_batches * batch_size, context_length - batch_size)
+
+                        # Create shifted input_ids
+                        tmp = torch.zeros((1, context_length), dtype=torch.int32)
+                        tmp[:,0:new_size] = input_ids[:,pos-new_size:pos]
+                        input_ids = tmp
+
+                        # Reset state and run prefill (uses rotation for pos >= sliding_window if available)
+                        current_pos = run_monolithic_prefill_with_rotation(
+                            prefill_model,
+                            prefill_rotate_model,
+                            input_ids,
+                            new_size,  # Prefill the entire shifted content
+                            context_length,
+                            batch_size,
+                            state,
+                            causal_mask,
+                            sliding_window,
+                            infer_rotate_model
+                        )
+
+                        # Start generating from the next position
+                        pos = new_size  # Don't back up, continue from where we left off
+
+                        window_shifted = True
+
+                    # Generate next token
+                    # Select the appropriate model based on position:
+                    # - pos < sliding_window: use infer_model (fill mode)
+                    # - pos >= sliding_window: use infer_rotate_model (rotation mode) if available
+                    if pos >= sliding_window and infer_rotate_model is not None:
+                        current_infer_model = infer_rotate_model
+                    else:
+                        current_infer_model = infer_model
+
+                    next_token = generate_next_token_monolithic(
+                        current_infer_model,
+                        input_ids,
+                        pos,
+                        context_length,
+                        metadata,
+                        state,
+                        causal_mask
+                    )
+
+                    # Add token
+                    input_ids[0, pos] = next_token
+                    if not warmup:
+                        token_printer.add_token(next_token)
+                        token_printer.drain_buffer()
+                    response_tokens.append(next_token)
+
+                    pos += 1
+                    tokens_generated += 1
+
+                    # In warmup mode, limit tokens
+                    if warmup and tokens_generated >= WARMUP_TOKEN_LIMIT:
+                        break
+                    if not warmup and max_tokens_this_turn is not None and tokens_generated >= max_tokens_this_turn:
+                        break
+
+                    if next_token in stop_token_ids:
+                        break
+
+                inference_time = time.time() - inference_start  # Calculate inference time
+
+                # Add assistant response to conversation
+                response_text = token_printer.stop()
+                conversation.append({"role": "assistant", "content": response_text})
+
+                # Update total tokens in memory (prompt + response)
+                total_tokens_in_memory = context_pos + len(response_tokens)
+                cumulative_tokens += context_pos + len(response_tokens)
+
+                # Print stats only if not in warmup
+                if not warmup:
+                    total_time = time.time() - generation_start_time
+                    prefill_time = total_time - inference_time
+                    inference_tokens_per_sec = len(response_tokens) / inference_time if inference_time > 0 else 0
+                    prefill_ms = prefill_time * 1000
+                    prefill_tokens_per_sec = context_pos / prefill_time if prefill_time > 0 else 0
+
+                    # Show context status for split cache debugging
+                    # Final position after generation
+                    final_pos = context_pos + len(response_tokens)
+                    rotation_mode = "ROTATE" if (final_pos >= sliding_window and infer_rotate_model is not None) else "FILL"
+                    if total_tokens_in_memory > sliding_window:
+                        context_status = f"[Turn {turn_number} | GLOBAL+{rotation_mode}: {total_tokens_in_memory}/{state_length} ctx, {len(conversation)} msgs]"
+                    else:
+                        context_status = f"[Turn {turn_number} | LOCAL+{rotation_mode}: {total_tokens_in_memory}/{sliding_window} ctx, {len(conversation)} msgs]"
+
+                    print(f"{DARK_BLUE}{inference_tokens_per_sec:.1f} t/s, "
+                          f"TTFT: {prefill_ms:.1f}ms ({prefill_tokens_per_sec:.1f} t/s, {context_pos} tokens), "
+                          f"{len(response_tokens)} tokens {context_status}{RESET_COLOR}")
+
+                if auto_prompt is not None:
+                    break
+
+            except KeyboardInterrupt:
+                if not warmup:
+                    print("\nGeneration interrupted")
+                token_printer.stop()
+                continue
+
+    except Exception as e:
+        if not warmup:
+            print(f"\nError in chat loop: {str(e)}")
+            import traceback
+            traceback.print_exc()
+
+
+def get_user_input():
+    """Get input from user, handling special key combinations."""
+    global THINKING_MODE
+    try:
+        import termios
+        import tty
+        import sys
+
+        def _getch():
+            fd = sys.stdin.fileno()
+            old_settings = termios.tcgetattr(fd)
+            try:
+                tty.setraw(sys.stdin.fileno())
+                ch = sys.stdin.read(1)
+            finally:
+                termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
+            return ch
+
+        buffer = []
+        while True:
+            char = _getch()
+            
+            # Debug: print the character code
+            print(f"\nKey pressed: {repr(char)} (hex: {hex(ord(char))})")
+            
+            # Check for Enter key
+            if char == '\r' or char == '\n':
+                print()  # Move to next line
+                input_text = ''.join(buffer)
+                # Check if the command is /t
+                if input_text == '/t':
+                    THINKING_MODE = not THINKING_MODE
+                    print(f"Thinking mode {'ON' if THINKING_MODE else 'OFF'}")
+                    buffer = []  # Clear buffer
+                    print(f"\n{LIGHT_GREEN}You{' (thinking)' if THINKING_MODE else ''}:{RESET_COLOR}", end=' ', flush=True)
+                    continue
+                return input_text
+                
+            # Handle backspace
+            if char == '\x7f':  # backspace
+                if buffer:
+                    buffer.pop()
+                    sys.stdout.write('\b \b')  # Erase character
+                    sys.stdout.flush()
+                continue
+                
+            # Handle Ctrl-C
+            if char == '\x03':  # Ctrl-C
+                print("^C")
+                raise KeyboardInterrupt
+                
+            # Print character and add to buffer
+            sys.stdout.write(char)
+            sys.stdout.flush()
+            buffer.append(char)
+            
+    except ImportError:
+        # Fallback for systems without termios
+        return input("> ")
+
+def chat_loop(embed_model, ffn_models, lmhead_model, tokenizer, metadata, state, causal_mask, auto_prompt=None, warmup=False, max_tokens=None):
+    """Interactive chat loop."""
+    global THINKING_MODE
+    global DEBUG_LEVEL
+    context_length = metadata.get('context_length')
+    state_length = metadata.get('state_length', context_length)  # For split cache models
+    batch_size = metadata.get('batch_size', 64)
+    
+    if not warmup:
+        print(f"\nUsing context length: {context_length}")
+        print("\nStarting chat session. Press Ctrl+D to exit.")
+        print("Type your message and press Enter to chat. Use /t to toggle thinking mode.")
+        print(f"Thinking mode is {'ON' if THINKING_MODE else 'OFF'}")
+    
+    # Keep track of conversation history
+    conversation = []
+    stop_token_ids = build_stop_token_ids(tokenizer)
+    use_chat_template = False
+    try:
+        tokenizer.apply_chat_template([{"role": "user", "content": "test"}], return_tensors="pt")
+        use_chat_template = True
+        if not warmup:
+            print("\nUsing chat template for prompts")
+    except Exception:
+        if not warmup:
+            print("\nUsing manual formatting for prompts")
+
+    def _build_base_input_ids(messages, show_debug):
+        if use_chat_template:
+            base_input_ids = tokenizer.apply_chat_template(
+                messages,
+                return_tensors="pt",
+                add_generation_prompt=True
+            ).to(torch.int32)
+            if show_debug and DEBUG_LEVEL >= 1 and not warmup:
+                label = "Full prompt with thinking" if THINKING_MODE else "Full prompt"
+                print(f"\n{DARK_BLUE}Debug: {label}:{RESET_COLOR}")
+                print(tokenizer.decode(base_input_ids[0]))
+            return base_input_ids
+
+        prompt_text = format_manual_prompt(messages)
+        base_input_ids = tokenizer(
+            prompt_text, return_tensors="pt", add_special_tokens=True
+        ).input_ids.to(torch.int32)
+        if show_debug and DEBUG_LEVEL >= 1 and not warmup:
+            label = "Full prompt with thinking" if THINKING_MODE else "Full prompt"
+            print(f"\n{DARK_BLUE}Debug: {label}:{RESET_COLOR}")
+            print(prompt_text)
+        return base_input_ids
+    
+    try:
+        while True:
+            try:
+                if not warmup:
+                    print(f"\n{LIGHT_GREEN}You{' (thinking)' if THINKING_MODE else ''}:{RESET_COLOR}", end=' ', flush=True)
+                if auto_prompt is not None:
+                    user_input = auto_prompt
+                    if not warmup:
+                        print(user_input)
+                else:
+                    user_input = input().strip()
+            except EOFError:
+                if not warmup:
+                    print("\nExiting chat...")
+                break
+            
+            if not user_input:
+                continue
+
+            # Handle /t command
+            if user_input == "/t":
+                THINKING_MODE = not THINKING_MODE
+                print(f"Thinking mode {'ON' if THINKING_MODE else 'OFF'}")
+                continue
+            
+            # Add user message to conversation
+            conversation.append({"role": "user", "content": user_input})
+            
+            messages = conversation
+            if THINKING_MODE:
+                messages = [{"role": "system", "content": THINKING_PROMPT}] + conversation
+            base_input_ids = _build_base_input_ids(messages, show_debug=True)
+            
+            # Check if we need to trim history
+            # Use state_length (global context) for split cache models, context_length otherwise
+            while base_input_ids.size(1) > state_length - 100:  # Leave room for response
+                # Remove oldest message pair (user + assistant)
+                if len(conversation) > 2:
+                    conversation = conversation[2:]  # Remove oldest pair
+                    messages = conversation
+                    if THINKING_MODE:
+                        messages = [{"role": "system", "content": THINKING_PROMPT}] + conversation
+                    base_input_ids = _build_base_input_ids(messages, show_debug=False)
+                else:
+                    # If only current message remains and still too long, truncate
+                    base_input_ids = base_input_ids[:, -state_length//2:]
+                    break
+
+            context_pos = base_input_ids.size(1)
+            
+            # Pad sequence to context_size
+            input_ids = F.pad(
+                base_input_ids,
+                (0, context_length - context_pos),
+                value=0
+            )
+            
+            # split_lm_head should already be in metadata from caller
+
+            # Initialize token printer and collect response
+            token_printer = TokenPrinter(tokenizer)
+            response_tokens = []
+            generation_start_time = time.time()
+            
+            try:
+                # Get sliding_window for rotation support (Gemma3)
+                sliding_window = metadata.get('sliding_window', None)
+
+                # Run prefill on entire context
+                current_pos = run_prefill(
+                    embed_model,
+                    ffn_models,
+                    input_ids,
+                    context_pos,
+                    context_length,
+                    batch_size,
+                    state,
+                    causal_mask,
+                    sliding_window
+                )
+                #print(f"\n[DEBUG] After initial prefill - current_pos: {current_pos}")
+
+                if not warmup:
+                    print(f"\n{LIGHT_BLUE}Assistant:{RESET_COLOR}", end=' ', flush=True)
+                
+                # Generation loop
+                pos = context_pos
+                tokens_generated = 0
+                max_tokens_this_turn = (
+                    max_tokens
+                    if max_tokens is not None
+                    else max(0, context_length - context_pos)
+                )
+                inference_start = time.time()  # Start inference timing
+                
+                while True:
+                    # Check if we need to shift window
+                    if pos >= context_length - 2:
+                        if DEBUG_LEVEL >= 1:
+                            print_system(f"Context window reached {context_length} tokens; shifting context to continue.")
+                        # Calculate shift to maintain full batches
+                        batch_size = metadata.get('batch_size', 64)
+                        # Calculate max batches that fit in context
+                        max_batches = context_length // batch_size
+                        desired_batches = max(1, max_batches - 2)  # Leave room for new tokens
+                        new_size = min(desired_batches * batch_size, context_length - batch_size)
+
+                        # Create shifted input_ids
+                        tmp = torch.zeros((1, context_length), dtype=torch.int32)
+                        tmp[:,0:new_size] = input_ids[:,pos-new_size:pos]
+                        input_ids = tmp
+
+                        # Reset state and run prefill
+                        # keep the same state
+                        #state = create_unified_state(ffn_models, context_length)
+                        current_pos = run_prefill(
+                            embed_model,
+                            ffn_models,
+                            input_ids,
+                            new_size,  # Prefill the entire shifted content
+                            context_length,
+                            batch_size,
+                            state,
+                            causal_mask,
+                            sliding_window
+                        )
+
+                        # Start generating from the next position
+                        pos = new_size  # Don't back up, continue from where we left off
+
+                        #print(f"\n[DEBUG] After shift - next token will be at pos {pos}")
+                        #print(f"[DEBUG] Context before next token: {tokenizer.decode(input_ids[0, pos-40:pos])}")
+
+                        window_shifted = True
+                    
+                    # Generate next token
+                    next_token = generate_next_token(
+                        embed_model,
+                        ffn_models,
+                        lmhead_model,
+                        input_ids,
+                        pos,
+                        context_length,
+                        state,
+                        causal_mask,
+                        metadata
+                    )
+                    
+                    # Add token
+                    input_ids[0, pos] = next_token
+                    if not warmup:
+                        token_printer.add_token(next_token)
+                        token_printer.drain_buffer()
+                    response_tokens.append(next_token)
+                    
+                    pos += 1
+                    tokens_generated += 1
+                    
+                    # In warmup mode, limit tokens
+                    if warmup and tokens_generated >= WARMUP_TOKEN_LIMIT:
+                        break
+                    if not warmup and max_tokens_this_turn is not None and tokens_generated >= max_tokens_this_turn:
+                        break
+                    
+                    if next_token in stop_token_ids:
+                        break
+                inference_time = time.time() - inference_start  # Calculate inference time
+                
+                # Add assistant response to conversation
+                response_text = token_printer.stop()
+                conversation.append({"role": "assistant", "content": response_text})
+                
+                # Print stats only if not in warmup
+                if not warmup:
+                    total_time = time.time() - generation_start_time
+                    prefill_time = total_time - inference_time
+                    inference_tokens_per_sec = len(response_tokens) / inference_time if inference_time > 0 else 0
+                    prefill_ms = prefill_time * 1000
+                    prefill_tokens_per_sec = context_pos / prefill_time if prefill_time > 0 else 0
+                    print(f"{DARK_BLUE}{inference_tokens_per_sec:.1f} t/s, "
+                          f"TTFT: {prefill_ms:.1f}ms ({prefill_tokens_per_sec:.1f} t/s, {context_pos} tokens), "
+                          f"{len(response_tokens)} tokens{RESET_COLOR}")
+                
+                if auto_prompt is not None:
+                    break
+                
+            except KeyboardInterrupt:
+                if not warmup:
+                    print("\nGeneration interrupted")
+                token_printer.stop()
+                continue
+                
+    except Exception as e:
+        if not warmup:
+            print(f"\nError in chat loop: {str(e)}")
+            import traceback
+            traceback.print_exc()
+
+def main():
+    args = parse_args()
+    global DEBUG_LEVEL
+    DEBUG_LEVEL = args.debug_level
+
+    # Convert directory to absolute path
+    model_dir = Path(args.d).resolve()
+    if not model_dir.exists():
+        print(f"\nError: Model directory not found: {model_dir}")
+        return 1
+
+    print(f"\nUsing model directory: {model_dir}")
+    print(f"Context length: {args.context_length}")
+
+    try:
+        # Handle tokenizer path
+        if args.tokenizer is None:
+            args.tokenizer = str(model_dir)
+
+        if not Path(args.tokenizer).exists():
+            print(f"\nError: Tokenizer directory not found: {args.tokenizer}")
+            return 1
+
+        args.tokenizer = str(Path(args.tokenizer).resolve())  # Convert to absolute path
+        print(f"Using tokenizer path: {args.tokenizer}")
+
+        # Load tokenizer with resolved path
+        tokenizer = initialize_tokenizer(args.tokenizer)
+        if tokenizer is None:
+            raise RuntimeError("Failed to initialize tokenizer")
+
+        metadata = {}
+
+        # Branch based on model type
+        if getattr(args, 'is_monolithic', False):
+            # MONOLITHIC MODEL PATH
+            infer_model, infer_rotate_model, prefill_model, prefill_rotate_model, metadata = load_monolithic_model(args, metadata)
+
+            # Override context length from command line if provided
+            if args.context_length is not None:
+                metadata['context_length'] = args.context_length
+
+            # Use state_length from args (parsed from YAML) or default to context_length
+            metadata['state_length'] = getattr(args, 'state_length', metadata['context_length'])
+
+            # Set metadata values
+            metadata['batch_size'] = getattr(args, 'batch_size', 64)
+            metadata['split_lm_head'] = getattr(args, 'split_lm_head', 16)
+            metadata['argmax_in_model'] = getattr(args, 'argmax_in_model', False)
+            metadata['sliding_window'] = 512  # Local attention window for Gemma3
+
+            print(f"\nMonolithic metadata: {metadata}")
+
+            # Create state from infer model
+            state = infer_model.make_state()
+            print("\nCreated unified transformer state for monolithic model")
+
+            # Initialize causal mask - use state_length for split cache models
+            causal_mask = initialize_causal_mask(metadata['state_length'])
+
+            # Warmup runs
+            if not args.nw:
+                for _ in range(2):
+                    chat_loop_monolithic(
+                        infer_model=infer_model,
+                        infer_rotate_model=infer_rotate_model,
+                        prefill_model=prefill_model,
+                        prefill_rotate_model=prefill_rotate_model,
+                        tokenizer=tokenizer,
+                        metadata=metadata,
+                        state=state,
+                        causal_mask=causal_mask,
+                        warmup=True,
+                        auto_prompt="who are you?"
+                    )
+
+            # Main run
+            chat_loop_monolithic(
+                infer_model=infer_model,
+                infer_rotate_model=infer_rotate_model,
+                prefill_model=prefill_model,
+                prefill_rotate_model=prefill_rotate_model,
+                tokenizer=tokenizer,
+                metadata=metadata,
+                state=state,
+                causal_mask=causal_mask,
+                warmup=False,
+                auto_prompt=args.prompt,
+                max_tokens=args.max_tokens
+            )
+
+        else:
+            # CHUNKED MODEL PATH (original code)
+            # Update paths to be relative to model directory
+            args.embed = str(model_dir / args.embed)
+            args.ffn = str(model_dir / args.ffn)
+            args.lmhead = str(model_dir / args.lmhead)
+
+            # Load models and extract metadata
+            embed_model, ffn_models, lmhead_model, metadata = load_models(args, metadata)
+
+            print(f"\nMetadata befor args.context_length: {metadata}")
+
+            # Override context length from command line if provided
+            if args.context_length is not None:
+                metadata['context_length'] = args.context_length
+                metadata['state_length'] = args.context_length  # Also update state_length
+                print(f"\nOverriding context length from command line: {args.context_length}")
+
+            print(f"\nMetadata after load_models: {metadata}")
+
+            # Create unified state once
+            state = create_unified_state(ffn_models, metadata['context_length'])
+
+            # Initialize causal mask once
+            causal_mask = initialize_causal_mask(metadata['context_length'])
+
+            # Add split_lm_head to metadata for generate_next_token
+            metadata['split_lm_head'] = getattr(args, 'split_lm_head', 8)
+
+            # Add argmax_in_model flag for chunked models
+            metadata['argmax_in_model'] = getattr(args, 'argmax_in_model', False)
+
+            # Add sliding_window for Gemma3 rotation support
+            sliding_window = getattr(args, 'sliding_window', None)
+            metadata['sliding_window'] = sliding_window
+            if sliding_window is not None:
+                context_len = metadata['context_length']
+                if context_len > sliding_window:
+                    print(f"Sliding window: {sliding_window} (rotation enabled for pos >= {sliding_window})")
+                else:
+                    print(f"Sliding window: {sliding_window} (rotation disabled - context {context_len} <= sliding_window)")
+
+            # Warmup runs to prevent Python GIL issues with CoreML !
+            if not args.nw:
+                for i in range(2):
+                    chat_loop(
+                        embed_model=embed_model,
+                        ffn_models=ffn_models,
+                        lmhead_model=lmhead_model,
+                        tokenizer=tokenizer,
+                        metadata=metadata,
+                        state=state,  # Pass the state
+                        causal_mask=causal_mask,  # Pass the causal mask
+                        warmup=True,
+                        auto_prompt="who are you?"
+                    )
+
+            # Main run
+            chat_loop(
+                embed_model=embed_model,
+                ffn_models=ffn_models,
+                lmhead_model=lmhead_model,
+                tokenizer=tokenizer,
+                metadata=metadata,
+                state=state,  # Pass the state
+                causal_mask=causal_mask,  # Pass the causal mask
+                warmup=False,
+                auto_prompt=args.prompt,
+                max_tokens=args.max_tokens
+            )
+
+    except Exception as e:
+        print(f"\nError: {str(e)}")
+        import traceback
+        traceback.print_exc()
+        return 1
+
+    return 0
+
+if __name__ == "__main__":
+    exit(main()) 

config.json

@@ -0,0 +1,4 @@
+{
+  "tokenizer_class": "GemmaTokenizer",
+  "model_type": "gemma"
+}

gemma3_FFN_PF_lut4_chunk_01of02.mlmodelc/analytics/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e23843ead3a55dd604fe5f6237628cc34225064f6e9a7019c399fa7bd74dcc6
+size 243

gemma3_FFN_PF_lut4_chunk_01of02.mlmodelc/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:260c117c785e874f502b6eea34e7df49f980a72f5bd2be79b55e50ae8db6ff4a
+size 1087

gemma3_FFN_PF_lut4_chunk_01of02.mlmodelc/metadata.json

@@ -0,0 +1,349 @@
+[
+  {
+    "metadataOutputVersion" : "3.0",
+    "userDefinedMetadata" : {
+      "com.github.apple.coremltools.version" : "9.0",
+      "com.github.apple.coremltools.source_dialect" : "TorchScript",
+      "com.anemll.lut_bits" : "4",
+      "com.anemll.context_length" : "1024",
+      "com.github.apple.coremltools.source" : "torch==2.5.0",
+      "com.anemll.num_chunks" : "2",
+      "com.anemll.batch_size" : "64",
+      "com.anemll.info" : "Converted with Anemll v0.1.1",
+      "com.anemll.chunk_no" : "1"
+    },
+    "availability" : {
+      "macOS" : "15.0",
+      "tvOS" : "18.0",
+      "visionOS" : "2.0",
+      "watchOS" : "11.0",
+      "iOS" : "18.0",
+      "macCatalyst" : "18.0"
+    },
+    "inputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 2560]",
+        "name" : "hidden_states",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Int32",
+        "formattedType" : "MultiArray (Int32 1)",
+        "shortDescription" : "",
+        "shape" : "[1]",
+        "name" : "position_ids",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 1 × 1024)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 1, 1024]",
+        "name" : "causal_mask",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Int32",
+        "formattedType" : "MultiArray (Int32 1)",
+        "shortDescription" : "",
+        "shape" : "[1]",
+        "name" : "current_pos",
+        "type" : "MultiArray"
+      }
+    ],
+    "outputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 2560]",
+        "name" : "output_hidden_states",
+        "type" : "MultiArray"
+      }
+    ],
+    "modelParameters" : [
+
+    ],
+    "storagePrecision" : "Mixed (Float16, Palettized (12 bits), Palettized (13 bits), Palettized (14 bits), Palettized (16 bits), UInt4)",
+    "method" : "predict",
+    "functions" : [
+      {
+        "inputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 2560]",
+            "name" : "hidden_states",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 1)",
+            "shortDescription" : "",
+            "shape" : "[1]",
+            "name" : "position_ids",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 1 × 1024)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 1, 1024]",
+            "name" : "causal_mask",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 1)",
+            "shortDescription" : "",
+            "shape" : "[1]",
+            "name" : "current_pos",
+            "type" : "MultiArray"
+          }
+        ],
+        "computePrecision" : "Mixed (Float16, Int16, Int32, UInt16)",
+        "storagePrecision" : "Mixed (Float16, Palettized (12 bits), Palettized (13 bits), Palettized (14 bits), Palettized (16 bits), UInt4)",
+        "stateSchema" : [
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 58 × 4 × 1024 × 256)",
+            "shortDescription" : "",
+            "shape" : "[58, 4, 1024, 256]",
+            "name" : "model_model_kv_cache_local",
+            "type" : "State"
+          },
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 10 × 4 × 1024 × 256)",
+            "shortDescription" : "",
+            "shape" : "[10, 4, 1024, 256]",
+            "name" : "model_model_kv_cache_global",
+            "type" : "State"
+          }
+        ],
+        "outputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 2560]",
+            "name" : "output_hidden_states",
+            "type" : "MultiArray"
+          }
+        ],
+        "name" : "infer",
+        "mlProgramOperationTypeHistogram" : {
+          "Ios18.expandDims" : 68,
+          "Ios18.mul" : 342,
+          "Ios18.softmax" : 17,
+          "Ios18.matmul" : 34,
+          "Identity" : 1,
+          "Ios18.greaterEqual" : 2,
+          "Select" : 2,
+          "Ios18.readState" : 36,
+          "Tile" : 34,
+          "Ios18.gather" : 4,
+          "Ios18.add" : 90,
+          "Ios18.layerNorm" : 102,
+          "Ios18.sliceUpdate" : 34,
+          "Ios18.writeState" : 34,
+          "Ios18.reshape" : 107,
+          "Ios18.constexprLutToDense" : 119,
+          "Ios18.conv" : 119,
+          "Ios18.concat" : 196,
+          "Ios18.transpose" : 102,
+          "Ios18.cast" : 5,
+          "Ios18.gelu" : 17,
+          "Ios18.sliceByIndex" : 212,
+          "Ios18.squeeze" : 51
+        }
+      },
+      {
+        "inputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 64 × 2560)",
+            "shortDescription" : "",
+            "shape" : "[1, 64, 2560]",
+            "name" : "hidden_states",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 64)",
+            "shortDescription" : "",
+            "shape" : "[64]",
+            "name" : "position_ids",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 64 × 1024)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 64, 1024]",
+            "name" : "causal_mask",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 1)",
+            "shortDescription" : "",
+            "shape" : "[1]",
+            "name" : "current_pos",
+            "type" : "MultiArray"
+          }
+        ],
+        "computePrecision" : "Mixed (Float16, Int16, Int32, UInt16)",
+        "storagePrecision" : "Mixed (Float16, Palettized (12 bits), Palettized (13 bits), Palettized (14 bits), Palettized (16 bits), UInt4)",
+        "stateSchema" : [
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 58 × 4 × 1024 × 256)",
+            "shortDescription" : "",
+            "shape" : "[58, 4, 1024, 256]",
+            "name" : "model_model_kv_cache_local",
+            "type" : "State"
+          },
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 10 × 4 × 1024 × 256)",
+            "shortDescription" : "",
+            "shape" : "[10, 4, 1024, 256]",
+            "name" : "model_model_kv_cache_global",
+            "type" : "State"
+          }
+        ],
+        "outputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 64 × 2560)",
+            "shortDescription" : "",
+            "shape" : "[1, 64, 2560]",
+            "name" : "output_hidden_states",
+            "type" : "MultiArray"
+          }
+        ],
+        "name" : "prefill",
+        "mlProgramOperationTypeHistogram" : {
+          "Ios18.expandDims" : 68,
+          "Ios18.mul" : 342,
+          "Ios18.softmax" : 17,
+          "Ios18.matmul" : 34,
+          "Identity" : 1,
+          "Ios18.greaterEqual" : 2,
+          "Select" : 2,
+          "Ios18.readState" : 36,
+          "Tile" : 34,
+          "Ios18.gather" : 4,
+          "Ios18.add" : 89,
+          "Ios18.layerNorm" : 102,
+          "Ios18.sliceUpdate" : 34,
+          "Ios18.writeState" : 34,
+          "Ios18.reshape" : 141,
+          "Ios18.constexprLutToDense" : 119,
+          "Ios18.conv" : 119,
+          "Ios18.concat" : 136,
+          "Ios18.transpose" : 157,
+          "Ios18.cast" : 5,
+          "Ios18.gelu" : 17,
+          "Ios18.sliceByIndex" : 212,
+          "Ios18.squeeze" : 51
+        }
+      }
+    ],
+    "version" : "0.1.1",
+    "isUpdatable" : "0",
+    "defaultFunctionName" : "infer",
+    "specificationVersion" : 9,
+    "stateSchema" : [
+      {
+        "dataType" : "Float16",
+        "isOptional" : "0",
+        "formattedType" : "State (Float16 58 × 4 × 1024 × 256)",
+        "shortDescription" : "",
+        "shape" : "[58, 4, 1024, 256]",
+        "name" : "model_model_kv_cache_local",
+        "type" : "State"
+      },
+      {
+        "dataType" : "Float16",
+        "isOptional" : "0",
+        "formattedType" : "State (Float16 10 × 4 × 1024 × 256)",
+        "shortDescription" : "",
+        "shape" : "[10, 4, 1024, 256]",
+        "name" : "model_model_kv_cache_global",
+        "type" : "State"
+      }
+    ],
+    "computePrecision" : "Mixed (Float16, Int16, Int32, UInt16)",
+    "mlProgramOperationTypeHistogram" : {
+      "Ios18.expandDims" : 68,
+      "Ios18.mul" : 342,
+      "Ios18.softmax" : 17,
+      "Ios18.matmul" : 34,
+      "Identity" : 1,
+      "Ios18.greaterEqual" : 2,
+      "Select" : 2,
+      "Ios18.readState" : 36,
+      "Tile" : 34,
+      "Ios18.gather" : 4,
+      "Ios18.add" : 90,
+      "Ios18.layerNorm" : 102,
+      "Ios18.sliceUpdate" : 34,
+      "Ios18.writeState" : 34,
+      "Ios18.reshape" : 107,
+      "Ios18.constexprLutToDense" : 119,
+      "Ios18.conv" : 119,
+      "Ios18.concat" : 196,
+      "Ios18.transpose" : 102,
+      "Ios18.cast" : 5,
+      "Ios18.gelu" : 17,
+      "Ios18.sliceByIndex" : 212,
+      "Ios18.squeeze" : 51
+    },
+    "shortDescription" : "Anemll Model: Multifunction FFN+Prefill",
+    "generatedClassName" : "gemma3_FFN_PF_lut4_chunk_01of02",
+    "author" : "Converted with Anemll v0.1.1",
+    "modelType" : {
+      "name" : "MLModelType_mlProgram"
+    }
+  }
+]

gemma3_FFN_PF_lut4_chunk_01of02.mlmodelc/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7d3facd57dd7cb071f1a3bc535a6cc19dd633428c4aa0540c497034a64046417
+size 810781248

gemma3_FFN_PF_lut4_chunk_02of02.mlmodelc/analytics/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b528c0e95e7d7c5f99895ac6bbabdc60252aa3fb2e5527c9beab09a31f18a33c
+size 243

gemma3_FFN_PF_lut4_chunk_02of02.mlmodelc/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c95d0f76a8034da6b00e59ebb7144319517ab721ad5db2e209d3c534351e086
+size 1087

gemma3_FFN_PF_lut4_chunk_02of02.mlmodelc/metadata.json

@@ -0,0 +1,349 @@
+[
+  {
+    "metadataOutputVersion" : "3.0",
+    "userDefinedMetadata" : {
+      "com.github.apple.coremltools.source" : "torch==2.5.0",
+      "com.github.apple.coremltools.version" : "9.0",
+      "com.anemll.context_length" : "1024",
+      "com.github.apple.coremltools.source_dialect" : "TorchScript",
+      "com.anemll.chunk_no" : "2",
+      "com.anemll.num_chunks" : "2",
+      "com.anemll.info" : "Converted with Anemll v0.1.1",
+      "com.anemll.batch_size" : "64",
+      "com.anemll.lut_bits" : "4"
+    },
+    "availability" : {
+      "macOS" : "15.0",
+      "tvOS" : "18.0",
+      "visionOS" : "2.0",
+      "watchOS" : "11.0",
+      "iOS" : "18.0",
+      "macCatalyst" : "18.0"
+    },
+    "inputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 2560]",
+        "name" : "hidden_states",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Int32",
+        "formattedType" : "MultiArray (Int32 1)",
+        "shortDescription" : "",
+        "shape" : "[1]",
+        "name" : "position_ids",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 1 × 1024)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 1, 1024]",
+        "name" : "causal_mask",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Int32",
+        "formattedType" : "MultiArray (Int32 1)",
+        "shortDescription" : "",
+        "shape" : "[1]",
+        "name" : "current_pos",
+        "type" : "MultiArray"
+      }
+    ],
+    "outputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 2560]",
+        "name" : "output_hidden_states",
+        "type" : "MultiArray"
+      }
+    ],
+    "modelParameters" : [
+
+    ],
+    "storagePrecision" : "Mixed (Float16, Palettized (12 bits), Palettized (13 bits), Palettized (14 bits), Palettized (16 bits), UInt4)",
+    "method" : "predict",
+    "functions" : [
+      {
+        "inputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 2560]",
+            "name" : "hidden_states",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 1)",
+            "shortDescription" : "",
+            "shape" : "[1]",
+            "name" : "position_ids",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 1 × 1024)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 1, 1024]",
+            "name" : "causal_mask",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 1)",
+            "shortDescription" : "",
+            "shape" : "[1]",
+            "name" : "current_pos",
+            "type" : "MultiArray"
+          }
+        ],
+        "computePrecision" : "Mixed (Float16, Int16, Int32, UInt16)",
+        "storagePrecision" : "Mixed (Float16, Palettized (12 bits), Palettized (13 bits), Palettized (14 bits), Palettized (16 bits), UInt4)",
+        "stateSchema" : [
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 58 × 4 × 1024 × 256)",
+            "shortDescription" : "",
+            "shape" : "[58, 4, 1024, 256]",
+            "name" : "model_model_kv_cache_local",
+            "type" : "State"
+          },
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 10 × 4 × 1024 × 256)",
+            "shortDescription" : "",
+            "shape" : "[10, 4, 1024, 256]",
+            "name" : "model_model_kv_cache_global",
+            "type" : "State"
+          }
+        ],
+        "outputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 2560]",
+            "name" : "output_hidden_states",
+            "type" : "MultiArray"
+          }
+        ],
+        "name" : "infer",
+        "mlProgramOperationTypeHistogram" : {
+          "Ios18.expandDims" : 68,
+          "Ios18.mul" : 344,
+          "Ios18.softmax" : 17,
+          "Ios18.matmul" : 34,
+          "Identity" : 1,
+          "Ios18.greaterEqual" : 2,
+          "Select" : 2,
+          "Ios18.readState" : 36,
+          "Tile" : 34,
+          "Ios18.gather" : 4,
+          "Ios18.add" : 90,
+          "Ios18.layerNorm" : 103,
+          "Ios18.sliceUpdate" : 34,
+          "Ios18.writeState" : 34,
+          "Ios18.reshape" : 107,
+          "Ios18.constexprLutToDense" : 119,
+          "Ios18.conv" : 119,
+          "Ios18.concat" : 205,
+          "Ios18.transpose" : 102,
+          "Ios18.cast" : 5,
+          "Ios18.gelu" : 17,
+          "Ios18.sliceByIndex" : 213,
+          "Ios18.squeeze" : 51
+        }
+      },
+      {
+        "inputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 64 × 2560)",
+            "shortDescription" : "",
+            "shape" : "[1, 64, 2560]",
+            "name" : "hidden_states",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 64)",
+            "shortDescription" : "",
+            "shape" : "[64]",
+            "name" : "position_ids",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 64 × 1024)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 64, 1024]",
+            "name" : "causal_mask",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 1)",
+            "shortDescription" : "",
+            "shape" : "[1]",
+            "name" : "current_pos",
+            "type" : "MultiArray"
+          }
+        ],
+        "computePrecision" : "Mixed (Float16, Int16, Int32, UInt16)",
+        "storagePrecision" : "Mixed (Float16, Palettized (12 bits), Palettized (13 bits), Palettized (14 bits), Palettized (16 bits), UInt4)",
+        "stateSchema" : [
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 58 × 4 × 1024 × 256)",
+            "shortDescription" : "",
+            "shape" : "[58, 4, 1024, 256]",
+            "name" : "model_model_kv_cache_local",
+            "type" : "State"
+          },
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 10 × 4 × 1024 × 256)",
+            "shortDescription" : "",
+            "shape" : "[10, 4, 1024, 256]",
+            "name" : "model_model_kv_cache_global",
+            "type" : "State"
+          }
+        ],
+        "outputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 2560]",
+            "name" : "output_hidden_states",
+            "type" : "MultiArray"
+          }
+        ],
+        "name" : "prefill",
+        "mlProgramOperationTypeHistogram" : {
+          "Ios18.expandDims" : 68,
+          "Ios18.mul" : 342,
+          "Ios18.softmax" : 17,
+          "Ios18.matmul" : 34,
+          "Identity" : 1,
+          "Ios18.greaterEqual" : 2,
+          "Select" : 2,
+          "Ios18.readState" : 36,
+          "Tile" : 34,
+          "Ios18.gather" : 4,
+          "Ios18.add" : 89,
+          "Ios18.layerNorm" : 102,
+          "Ios18.sliceUpdate" : 34,
+          "Ios18.writeState" : 34,
+          "Ios18.reshape" : 141,
+          "Ios18.constexprLutToDense" : 119,
+          "Ios18.conv" : 119,
+          "Ios18.concat" : 136,
+          "Ios18.transpose" : 157,
+          "Ios18.cast" : 5,
+          "Ios18.gelu" : 17,
+          "Ios18.sliceByIndex" : 213,
+          "Ios18.squeeze" : 51
+        }
+      }
+    ],
+    "version" : "0.1.1",
+    "isUpdatable" : "0",
+    "defaultFunctionName" : "infer",
+    "specificationVersion" : 9,
+    "stateSchema" : [
+      {
+        "dataType" : "Float16",
+        "isOptional" : "0",
+        "formattedType" : "State (Float16 58 × 4 × 1024 × 256)",
+        "shortDescription" : "",
+        "shape" : "[58, 4, 1024, 256]",
+        "name" : "model_model_kv_cache_local",
+        "type" : "State"
+      },
+      {
+        "dataType" : "Float16",
+        "isOptional" : "0",
+        "formattedType" : "State (Float16 10 × 4 × 1024 × 256)",
+        "shortDescription" : "",
+        "shape" : "[10, 4, 1024, 256]",
+        "name" : "model_model_kv_cache_global",
+        "type" : "State"
+      }
+    ],
+    "computePrecision" : "Mixed (Float16, Int16, Int32, UInt16)",
+    "mlProgramOperationTypeHistogram" : {
+      "Ios18.expandDims" : 68,
+      "Ios18.mul" : 344,
+      "Ios18.softmax" : 17,
+      "Ios18.matmul" : 34,
+      "Identity" : 1,
+      "Ios18.greaterEqual" : 2,
+      "Select" : 2,
+      "Ios18.readState" : 36,
+      "Tile" : 34,
+      "Ios18.gather" : 4,
+      "Ios18.add" : 90,
+      "Ios18.layerNorm" : 103,
+      "Ios18.sliceUpdate" : 34,
+      "Ios18.writeState" : 34,
+      "Ios18.reshape" : 107,
+      "Ios18.constexprLutToDense" : 119,
+      "Ios18.conv" : 119,
+      "Ios18.concat" : 205,
+      "Ios18.transpose" : 102,
+      "Ios18.cast" : 5,
+      "Ios18.gelu" : 17,
+      "Ios18.sliceByIndex" : 213,
+      "Ios18.squeeze" : 51
+    },
+    "shortDescription" : "Anemll Model: Multifunction FFN+Prefill",
+    "generatedClassName" : "gemma3_FFN_PF_lut4_chunk_02of02",
+    "author" : "Converted with Anemll v0.1.1",
+    "modelType" : {
+      "name" : "MLModelType_mlProgram"
+    }
+  }
+]

gemma3_FFN_PF_lut4_chunk_02of02.mlmodelc/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0eedb2a3f38b818d22a49908f59bf7e2bef70ac30d27fa375b5b43d73072d3ee
+size 810786432

gemma3_embeddings.mlmodelc/analytics/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b4f85141f904b9284fd7ef474649035cdd7be02f0ba4c314531bb6560ffa023a
+size 243

gemma3_embeddings.mlmodelc/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:264ac6761ac3bc746937f0a99b9b1b57df8807e2c24a33f1b0fe1c515427835e
+size 560

gemma3_embeddings.mlmodelc/metadata.json

@@ -0,0 +1,72 @@
+[
+  {
+    "shortDescription" : "Anemll Model (Embeddings) converted to CoreML",
+    "metadataOutputVersion" : "3.0",
+    "outputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16)",
+        "shortDescription" : "",
+        "shape" : "[]",
+        "name" : "hidden_states",
+        "type" : "MultiArray"
+      }
+    ],
+    "version" : "0.1.1",
+    "modelParameters" : [
+
+    ],
+    "author" : "Converted with Anemll v0.1.1",
+    "specificationVersion" : 9,
+    "storagePrecision" : "Float16",
+    "mlProgramOperationTypeHistogram" : {
+      "Ios18.greaterEqual" : 1,
+      "Ios18.add" : 1,
+      "Select" : 1,
+      "Ios18.gather" : 1,
+      "Ios18.mul" : 1
+    },
+    "computePrecision" : "Mixed (Float16, Int32)",
+    "stateSchema" : [
+
+    ],
+    "isUpdatable" : "0",
+    "availability" : {
+      "macOS" : "15.0",
+      "tvOS" : "18.0",
+      "visionOS" : "2.0",
+      "watchOS" : "11.0",
+      "iOS" : "18.0",
+      "macCatalyst" : "18.0"
+    },
+    "modelType" : {
+      "name" : "MLModelType_mlProgram"
+    },
+    "inputSchema" : [
+      {
+        "shortDescription" : "",
+        "dataType" : "Int32",
+        "hasShapeFlexibility" : "1",
+        "isOptional" : "0",
+        "shapeFlexibility" : "1 × 1 | 1 × 64",
+        "formattedType" : "MultiArray (Int32 1 × 1)",
+        "type" : "MultiArray",
+        "shape" : "[1, 1]",
+        "name" : "input_ids",
+        "enumeratedShapes" : "[[1, 1], [1, 64]]"
+      }
+    ],
+    "userDefinedMetadata" : {
+      "com.github.apple.coremltools.conversion_date" : "2026-01-30",
+      "com.github.apple.coremltools.source_dialect" : "TorchScript",
+      "com.anemll.context_length" : "1024",
+      "com.github.apple.coremltools.source" : "torch==2.5.0",
+      "com.github.apple.coremltools.version" : "9.0",
+      "com.anemll.info" : "Converted with Anemll v0.1.1"
+    },
+    "generatedClassName" : "gemma3_embeddings",
+    "method" : "predict"
+  }
+]

gemma3_embeddings.mlmodelc/model.mil

@@ -0,0 +1,18 @@
+program(1.3)
+[buildInfo = dict<string, string>({{"coremlc-component-MIL", "3510.2.1"}, {"coremlc-version", "3500.32.1"}, {"coremltools-component-torch", "2.5.0"}, {"coremltools-source-dialect", "TorchScript"}, {"coremltools-version", "9.0"}})]
+{
+    func main<ios18>(tensor<int32, [1, ?]> input_ids) [FlexibleShapeInformation = tuple<tuple<string, dict<string, tensor<int32, [?]>>>, tuple<string, dict<string, dict<string, tensor<int32, [?]>>>>>((("DefaultShapes", {{"input_ids", [1, 1]}}), ("EnumeratedShapes", {{"79ae981e", {{"input_ids", [1, 1]}}}, {"ed9b58c8", {{"input_ids", [1, 64]}}}})))] {
+            tensor<fp16, [262208, 2560]> embed_tokens_weight = const()[name = string("embed_tokens_weight"), val = tensor<fp16, [262208, 2560]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(64)))];
+            int32 hidden_states_1_batch_dims_0 = const()[name = string("hidden_states_1_batch_dims_0"), val = int32(0)];
+            bool hidden_states_1_validate_indices_0 = const()[name = string("hidden_states_1_validate_indices_0"), val = bool(false)];
+            int32 greater_equal_0_y_0 = const()[name = string("greater_equal_0_y_0"), val = int32(0)];
+            tensor<bool, [1, ?]> greater_equal_0 = greater_equal(x = input_ids, y = greater_equal_0_y_0)[name = string("greater_equal_0")];
+            int32 slice_by_index_0 = const()[name = string("slice_by_index_0"), val = int32(262208)];
+            tensor<int32, [1, ?]> add_0 = add(x = input_ids, y = slice_by_index_0)[name = string("add_0")];
+            tensor<int32, [1, ?]> select_0 = select(a = input_ids, b = add_0, cond = greater_equal_0)[name = string("select_0")];
+            int32 hidden_states_1_axis_1 = const()[name = string("hidden_states_1_axis_1"), val = int32(0)];
+            tensor<fp16, [1, ?, 2560]> hidden_states_1 = gather(axis = hidden_states_1_axis_1, batch_dims = hidden_states_1_batch_dims_0, indices = select_0, validate_indices = hidden_states_1_validate_indices_0, x = embed_tokens_weight)[name = string("hidden_states_1")];
+            fp16 var_7_to_fp16 = const()[name = string("op_7_to_fp16"), val = fp16(0x1.94cp+5)];
+            tensor<fp16, [1, ?, 2560]> hidden_states = mul(x = hidden_states_1, y = var_7_to_fp16)[name = string("hidden_states_cast_fp16")];
+        } -> (hidden_states);
+}

gemma3_embeddings.mlmodelc/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f1b931a4f2d2ea11483fdc02fc827ae1a0f96f11626993f37cea53293762b219
+size 1342505088

gemma3_lm_head_lut6.mlmodelc/analytics/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e5bfc6a9a1696895465dee71ddfdba69f0249a96b6b16274510b8cc6e5eadc47
+size 243

gemma3_lm_head_lut6.mlmodelc/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c7faf2025bd907e13c83406271713bbafa76a879e3a9cef9e3f00b7521d7948
+size 973

gemma3_lm_head_lut6.mlmodelc/metadata.json

@@ -0,0 +1,221 @@
+[
+  {
+    "shortDescription" : "Anemll Model (LM Head) converted to CoreML",
+    "metadataOutputVersion" : "3.0",
+    "outputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 16388)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 16388]",
+        "name" : "logits1",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 16388)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 16388]",
+        "name" : "logits2",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 16388)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 16388]",
+        "name" : "logits3",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 16388)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 16388]",
+        "name" : "logits4",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 16388)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 16388]",
+        "name" : "logits5",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 16388)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 16388]",
+        "name" : "logits6",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 16388)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 16388]",
+        "name" : "logits7",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 16388)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 16388]",
+        "name" : "logits8",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 16388)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 16388]",
+        "name" : "logits9",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 16388)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 16388]",
+        "name" : "logits10",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 16388)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 16388]",
+        "name" : "logits11",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 16388)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 16388]",
+        "name" : "logits12",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 16388)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 16388]",
+        "name" : "logits13",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 16388)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 16388]",
+        "name" : "logits14",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 16388)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 16388]",
+        "name" : "logits15",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 16388)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 16388]",
+        "name" : "logits16",
+        "type" : "MultiArray"
+      }
+    ],
+    "version" : "0.1.1",
+    "modelParameters" : [
+
+    ],
+    "author" : "Converted with Anemll v0.1.1",
+    "specificationVersion" : 9,
+    "storagePrecision" : "Mixed (Float16, Palettized (19 bits), UInt6)",
+    "mlProgramOperationTypeHistogram" : {
+      "Ios18.transpose" : 17,
+      "Ios18.constexprLutToDense" : 16,
+      "Ios18.expandDims" : 1,
+      "Ios18.conv" : 16,
+      "Ios18.squeeze" : 16
+    },
+    "computePrecision" : "Mixed (Float16, Int32)",
+    "stateSchema" : [
+
+    ],
+    "isUpdatable" : "0",
+    "availability" : {
+      "macOS" : "15.0",
+      "tvOS" : "18.0",
+      "visionOS" : "2.0",
+      "watchOS" : "11.0",
+      "iOS" : "18.0",
+      "macCatalyst" : "18.0"
+    },
+    "modelType" : {
+      "name" : "MLModelType_mlProgram"
+    },
+    "inputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 2560]",
+        "name" : "hidden_states",
+        "type" : "MultiArray"
+      }
+    ],
+    "userDefinedMetadata" : {
+      "com.github.apple.coremltools.version" : "9.0",
+      "com.github.apple.coremltools.source_dialect" : "TorchScript",
+      "com.github.apple.coremltools.conversion_date" : "2026-01-30",
+      "com.github.apple.coremltools.source" : "torch==2.5.0",
+      "com.anemll.context_length" : "1024",
+      "com.anemll.info" : "Converted with Anemll v0.1.1",
+      "com.anemll.lut_bits" : "6"
+    },
+    "generatedClassName" : "gemma3_lm_head_lut6",
+    "method" : "predict"
+  }
+]

gemma3_lm_head_lut6.mlmodelc/model.mil

@@ -0,0 +1,186 @@
+program(1.3)
+[buildInfo = dict<string, string>({{"coremlc-component-MIL", "3510.2.1"}, {"coremlc-version", "3500.32.1"}})]
+{
+    func main<ios18>(tensor<fp16, [1, 1, 2560]> hidden_states) {
+            tensor<int32, [3]> var_5 = const()[name = string("op_5"), val = tensor<int32, [3]>([0, 2, 1])];
+            tensor<int32, [1]> input_axes_0 = const()[name = string("input_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 2560, 1]> var_6_cast_fp16 = transpose(perm = var_5, x = hidden_states)[name = string("transpose_16")];
+            tensor<fp16, [1, 2560, 1, 1]> input_cast_fp16 = expand_dims(axes = input_axes_0, x = var_6_cast_fp16)[name = string("input_cast_fp16")];
+            string var_29_pad_type_0 = const()[name = string("op_29_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_29_strides_0 = const()[name = string("op_29_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_29_pad_0 = const()[name = string("op_29_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_29_dilations_0 = const()[name = string("op_29_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_29_groups_0 = const()[name = string("op_29_groups_0"), val = int32(1)];
+            tensor<fp16, [16388, 2560, 1, 1]> op_9_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16388, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(64))), lut = tensor<fp16, [4097, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(31465088))))[name = string("op_9_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16388, 1, 1]> var_29_cast_fp16 = conv(dilations = var_29_dilations_0, groups = var_29_groups_0, pad = var_29_pad_0, pad_type = var_29_pad_type_0, strides = var_29_strides_0, weight = op_9_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_29_cast_fp16")];
+            tensor<int32, [1]> var_31_axes_0 = const()[name = string("op_31_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 16388, 1]> var_31_cast_fp16 = squeeze(axes = var_31_axes_0, x = var_29_cast_fp16)[name = string("op_31_cast_fp16")];
+            tensor<int32, [3]> var_34_perm_0 = const()[name = string("op_34_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_55_pad_type_0 = const()[name = string("op_55_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_55_strides_0 = const()[name = string("op_55_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_55_pad_0 = const()[name = string("op_55_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_55_dilations_0 = const()[name = string("op_55_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_55_groups_0 = const()[name = string("op_55_groups_0"), val = int32(1)];
+            tensor<fp16, [16388, 2560, 1, 1]> op_35_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16388, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(31989568))), lut = tensor<fp16, [4097, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(63454592))))[name = string("op_35_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16388, 1, 1]> var_55_cast_fp16 = conv(dilations = var_55_dilations_0, groups = var_55_groups_0, pad = var_55_pad_0, pad_type = var_55_pad_type_0, strides = var_55_strides_0, weight = op_35_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_55_cast_fp16")];
+            tensor<int32, [1]> var_57_axes_0 = const()[name = string("op_57_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 16388, 1]> var_57_cast_fp16 = squeeze(axes = var_57_axes_0, x = var_55_cast_fp16)[name = string("op_57_cast_fp16")];
+            tensor<int32, [3]> var_60_perm_0 = const()[name = string("op_60_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_81_pad_type_0 = const()[name = string("op_81_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_81_strides_0 = const()[name = string("op_81_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_81_pad_0 = const()[name = string("op_81_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_81_dilations_0 = const()[name = string("op_81_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_81_groups_0 = const()[name = string("op_81_groups_0"), val = int32(1)];
+            tensor<fp16, [16388, 2560, 1, 1]> op_61_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16388, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(63979072))), lut = tensor<fp16, [4097, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(95444096))))[name = string("op_61_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16388, 1, 1]> var_81_cast_fp16 = conv(dilations = var_81_dilations_0, groups = var_81_groups_0, pad = var_81_pad_0, pad_type = var_81_pad_type_0, strides = var_81_strides_0, weight = op_61_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_81_cast_fp16")];
+            tensor<int32, [1]> var_83_axes_0 = const()[name = string("op_83_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 16388, 1]> var_83_cast_fp16 = squeeze(axes = var_83_axes_0, x = var_81_cast_fp16)[name = string("op_83_cast_fp16")];
+            tensor<int32, [3]> var_86_perm_0 = const()[name = string("op_86_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_107_pad_type_0 = const()[name = string("op_107_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_107_strides_0 = const()[name = string("op_107_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_107_pad_0 = const()[name = string("op_107_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_107_dilations_0 = const()[name = string("op_107_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_107_groups_0 = const()[name = string("op_107_groups_0"), val = int32(1)];
+            tensor<fp16, [16388, 2560, 1, 1]> op_87_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16388, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(95968576))), lut = tensor<fp16, [4097, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(127433600))))[name = string("op_87_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16388, 1, 1]> var_107_cast_fp16 = conv(dilations = var_107_dilations_0, groups = var_107_groups_0, pad = var_107_pad_0, pad_type = var_107_pad_type_0, strides = var_107_strides_0, weight = op_87_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_107_cast_fp16")];
+            tensor<int32, [1]> var_109_axes_0 = const()[name = string("op_109_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 16388, 1]> var_109_cast_fp16 = squeeze(axes = var_109_axes_0, x = var_107_cast_fp16)[name = string("op_109_cast_fp16")];
+            tensor<int32, [3]> var_112_perm_0 = const()[name = string("op_112_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_133_pad_type_0 = const()[name = string("op_133_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_133_strides_0 = const()[name = string("op_133_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_133_pad_0 = const()[name = string("op_133_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_133_dilations_0 = const()[name = string("op_133_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_133_groups_0 = const()[name = string("op_133_groups_0"), val = int32(1)];
+            tensor<fp16, [16388, 2560, 1, 1]> op_113_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16388, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(127958080))), lut = tensor<fp16, [4097, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(159423104))))[name = string("op_113_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16388, 1, 1]> var_133_cast_fp16 = conv(dilations = var_133_dilations_0, groups = var_133_groups_0, pad = var_133_pad_0, pad_type = var_133_pad_type_0, strides = var_133_strides_0, weight = op_113_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_133_cast_fp16")];
+            tensor<int32, [1]> var_135_axes_0 = const()[name = string("op_135_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 16388, 1]> var_135_cast_fp16 = squeeze(axes = var_135_axes_0, x = var_133_cast_fp16)[name = string("op_135_cast_fp16")];
+            tensor<int32, [3]> var_138_perm_0 = const()[name = string("op_138_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_159_pad_type_0 = const()[name = string("op_159_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_159_strides_0 = const()[name = string("op_159_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_159_pad_0 = const()[name = string("op_159_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_159_dilations_0 = const()[name = string("op_159_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_159_groups_0 = const()[name = string("op_159_groups_0"), val = int32(1)];
+            tensor<fp16, [16388, 2560, 1, 1]> op_139_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16388, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(159947584))), lut = tensor<fp16, [4097, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(191412608))))[name = string("op_139_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16388, 1, 1]> var_159_cast_fp16 = conv(dilations = var_159_dilations_0, groups = var_159_groups_0, pad = var_159_pad_0, pad_type = var_159_pad_type_0, strides = var_159_strides_0, weight = op_139_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_159_cast_fp16")];
+            tensor<int32, [1]> var_161_axes_0 = const()[name = string("op_161_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 16388, 1]> var_161_cast_fp16 = squeeze(axes = var_161_axes_0, x = var_159_cast_fp16)[name = string("op_161_cast_fp16")];
+            tensor<int32, [3]> var_164_perm_0 = const()[name = string("op_164_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_185_pad_type_0 = const()[name = string("op_185_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_185_strides_0 = const()[name = string("op_185_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_185_pad_0 = const()[name = string("op_185_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_185_dilations_0 = const()[name = string("op_185_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_185_groups_0 = const()[name = string("op_185_groups_0"), val = int32(1)];
+            tensor<fp16, [16388, 2560, 1, 1]> op_165_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16388, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(191937088))), lut = tensor<fp16, [4097, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(223402112))))[name = string("op_165_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16388, 1, 1]> var_185_cast_fp16 = conv(dilations = var_185_dilations_0, groups = var_185_groups_0, pad = var_185_pad_0, pad_type = var_185_pad_type_0, strides = var_185_strides_0, weight = op_165_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_185_cast_fp16")];
+            tensor<int32, [1]> var_187_axes_0 = const()[name = string("op_187_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 16388, 1]> var_187_cast_fp16 = squeeze(axes = var_187_axes_0, x = var_185_cast_fp16)[name = string("op_187_cast_fp16")];
+            tensor<int32, [3]> var_190_perm_0 = const()[name = string("op_190_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_211_pad_type_0 = const()[name = string("op_211_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_211_strides_0 = const()[name = string("op_211_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_211_pad_0 = const()[name = string("op_211_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_211_dilations_0 = const()[name = string("op_211_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_211_groups_0 = const()[name = string("op_211_groups_0"), val = int32(1)];
+            tensor<fp16, [16388, 2560, 1, 1]> op_191_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16388, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(223926592))), lut = tensor<fp16, [4097, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(255391616))))[name = string("op_191_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16388, 1, 1]> var_211_cast_fp16 = conv(dilations = var_211_dilations_0, groups = var_211_groups_0, pad = var_211_pad_0, pad_type = var_211_pad_type_0, strides = var_211_strides_0, weight = op_191_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_211_cast_fp16")];
+            tensor<int32, [1]> var_213_axes_0 = const()[name = string("op_213_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 16388, 1]> var_213_cast_fp16 = squeeze(axes = var_213_axes_0, x = var_211_cast_fp16)[name = string("op_213_cast_fp16")];
+            tensor<int32, [3]> var_216_perm_0 = const()[name = string("op_216_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_237_pad_type_0 = const()[name = string("op_237_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_237_strides_0 = const()[name = string("op_237_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_237_pad_0 = const()[name = string("op_237_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_237_dilations_0 = const()[name = string("op_237_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_237_groups_0 = const()[name = string("op_237_groups_0"), val = int32(1)];
+            tensor<fp16, [16388, 2560, 1, 1]> op_217_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16388, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(255916096))), lut = tensor<fp16, [4097, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(287381120))))[name = string("op_217_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16388, 1, 1]> var_237_cast_fp16 = conv(dilations = var_237_dilations_0, groups = var_237_groups_0, pad = var_237_pad_0, pad_type = var_237_pad_type_0, strides = var_237_strides_0, weight = op_217_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_237_cast_fp16")];
+            tensor<int32, [1]> var_239_axes_0 = const()[name = string("op_239_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 16388, 1]> var_239_cast_fp16 = squeeze(axes = var_239_axes_0, x = var_237_cast_fp16)[name = string("op_239_cast_fp16")];
+            tensor<int32, [3]> var_242_perm_0 = const()[name = string("op_242_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_263_pad_type_0 = const()[name = string("op_263_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_263_strides_0 = const()[name = string("op_263_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_263_pad_0 = const()[name = string("op_263_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_263_dilations_0 = const()[name = string("op_263_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_263_groups_0 = const()[name = string("op_263_groups_0"), val = int32(1)];
+            tensor<fp16, [16388, 2560, 1, 1]> op_243_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16388, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(287905600))), lut = tensor<fp16, [4097, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(319370624))))[name = string("op_243_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16388, 1, 1]> var_263_cast_fp16 = conv(dilations = var_263_dilations_0, groups = var_263_groups_0, pad = var_263_pad_0, pad_type = var_263_pad_type_0, strides = var_263_strides_0, weight = op_243_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_263_cast_fp16")];
+            tensor<int32, [1]> var_265_axes_0 = const()[name = string("op_265_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 16388, 1]> var_265_cast_fp16 = squeeze(axes = var_265_axes_0, x = var_263_cast_fp16)[name = string("op_265_cast_fp16")];
+            tensor<int32, [3]> var_268_perm_0 = const()[name = string("op_268_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_289_pad_type_0 = const()[name = string("op_289_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_289_strides_0 = const()[name = string("op_289_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_289_pad_0 = const()[name = string("op_289_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_289_dilations_0 = const()[name = string("op_289_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_289_groups_0 = const()[name = string("op_289_groups_0"), val = int32(1)];
+            tensor<fp16, [16388, 2560, 1, 1]> op_269_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16388, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(319895104))), lut = tensor<fp16, [4097, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(351360128))))[name = string("op_269_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16388, 1, 1]> var_289_cast_fp16 = conv(dilations = var_289_dilations_0, groups = var_289_groups_0, pad = var_289_pad_0, pad_type = var_289_pad_type_0, strides = var_289_strides_0, weight = op_269_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_289_cast_fp16")];
+            tensor<int32, [1]> var_291_axes_0 = const()[name = string("op_291_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 16388, 1]> var_291_cast_fp16 = squeeze(axes = var_291_axes_0, x = var_289_cast_fp16)[name = string("op_291_cast_fp16")];
+            tensor<int32, [3]> var_294_perm_0 = const()[name = string("op_294_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_315_pad_type_0 = const()[name = string("op_315_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_315_strides_0 = const()[name = string("op_315_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_315_pad_0 = const()[name = string("op_315_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_315_dilations_0 = const()[name = string("op_315_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_315_groups_0 = const()[name = string("op_315_groups_0"), val = int32(1)];
+            tensor<fp16, [16388, 2560, 1, 1]> op_295_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16388, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(351884608))), lut = tensor<fp16, [4097, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(383349632))))[name = string("op_295_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16388, 1, 1]> var_315_cast_fp16 = conv(dilations = var_315_dilations_0, groups = var_315_groups_0, pad = var_315_pad_0, pad_type = var_315_pad_type_0, strides = var_315_strides_0, weight = op_295_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_315_cast_fp16")];
+            tensor<int32, [1]> var_317_axes_0 = const()[name = string("op_317_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 16388, 1]> var_317_cast_fp16 = squeeze(axes = var_317_axes_0, x = var_315_cast_fp16)[name = string("op_317_cast_fp16")];
+            tensor<int32, [3]> var_320_perm_0 = const()[name = string("op_320_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_341_pad_type_0 = const()[name = string("op_341_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_341_strides_0 = const()[name = string("op_341_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_341_pad_0 = const()[name = string("op_341_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_341_dilations_0 = const()[name = string("op_341_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_341_groups_0 = const()[name = string("op_341_groups_0"), val = int32(1)];
+            tensor<fp16, [16388, 2560, 1, 1]> op_321_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16388, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(383874112))), lut = tensor<fp16, [4097, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(415339136))))[name = string("op_321_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16388, 1, 1]> var_341_cast_fp16 = conv(dilations = var_341_dilations_0, groups = var_341_groups_0, pad = var_341_pad_0, pad_type = var_341_pad_type_0, strides = var_341_strides_0, weight = op_321_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_341_cast_fp16")];
+            tensor<int32, [1]> var_343_axes_0 = const()[name = string("op_343_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 16388, 1]> var_343_cast_fp16 = squeeze(axes = var_343_axes_0, x = var_341_cast_fp16)[name = string("op_343_cast_fp16")];
+            tensor<int32, [3]> var_346_perm_0 = const()[name = string("op_346_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_367_pad_type_0 = const()[name = string("op_367_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_367_strides_0 = const()[name = string("op_367_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_367_pad_0 = const()[name = string("op_367_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_367_dilations_0 = const()[name = string("op_367_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_367_groups_0 = const()[name = string("op_367_groups_0"), val = int32(1)];
+            tensor<fp16, [16388, 2560, 1, 1]> op_347_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16388, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(415863616))), lut = tensor<fp16, [4097, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(447328640))))[name = string("op_347_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16388, 1, 1]> var_367_cast_fp16 = conv(dilations = var_367_dilations_0, groups = var_367_groups_0, pad = var_367_pad_0, pad_type = var_367_pad_type_0, strides = var_367_strides_0, weight = op_347_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_367_cast_fp16")];
+            tensor<int32, [1]> var_369_axes_0 = const()[name = string("op_369_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 16388, 1]> var_369_cast_fp16 = squeeze(axes = var_369_axes_0, x = var_367_cast_fp16)[name = string("op_369_cast_fp16")];
+            tensor<int32, [3]> var_372_perm_0 = const()[name = string("op_372_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_393_pad_type_0 = const()[name = string("op_393_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_393_strides_0 = const()[name = string("op_393_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_393_pad_0 = const()[name = string("op_393_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_393_dilations_0 = const()[name = string("op_393_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_393_groups_0 = const()[name = string("op_393_groups_0"), val = int32(1)];
+            tensor<fp16, [16388, 2560, 1, 1]> op_373_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16388, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(447853120))), lut = tensor<fp16, [4097, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(479318144))))[name = string("op_373_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16388, 1, 1]> var_393_cast_fp16 = conv(dilations = var_393_dilations_0, groups = var_393_groups_0, pad = var_393_pad_0, pad_type = var_393_pad_type_0, strides = var_393_strides_0, weight = op_373_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_393_cast_fp16")];
+            tensor<int32, [1]> var_395_axes_0 = const()[name = string("op_395_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 16388, 1]> var_395_cast_fp16 = squeeze(axes = var_395_axes_0, x = var_393_cast_fp16)[name = string("op_395_cast_fp16")];
+            tensor<int32, [3]> var_398_perm_0 = const()[name = string("op_398_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_419_pad_type_0 = const()[name = string("op_419_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_419_strides_0 = const()[name = string("op_419_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_419_pad_0 = const()[name = string("op_419_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_419_dilations_0 = const()[name = string("op_419_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_419_groups_0 = const()[name = string("op_419_groups_0"), val = int32(1)];
+            tensor<fp16, [16388, 2560, 1, 1]> op_399_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16388, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(479842624))), lut = tensor<fp16, [4097, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(511307648))))[name = string("op_399_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16388, 1, 1]> var_419_cast_fp16 = conv(dilations = var_419_dilations_0, groups = var_419_groups_0, pad = var_419_pad_0, pad_type = var_419_pad_type_0, strides = var_419_strides_0, weight = op_399_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_419_cast_fp16")];
+            tensor<int32, [1]> var_421_axes_0 = const()[name = string("op_421_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 16388, 1]> var_421_cast_fp16 = squeeze(axes = var_421_axes_0, x = var_419_cast_fp16)[name = string("op_421_cast_fp16")];
+            tensor<int32, [3]> var_424_perm_0 = const()[name = string("op_424_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            tensor<fp16, [1, 1, 16388]> logits1 = transpose(perm = var_34_perm_0, x = var_31_cast_fp16)[name = string("transpose_0")];
+            tensor<fp16, [1, 1, 16388]> logits2 = transpose(perm = var_60_perm_0, x = var_57_cast_fp16)[name = string("transpose_1")];
+            tensor<fp16, [1, 1, 16388]> logits3 = transpose(perm = var_86_perm_0, x = var_83_cast_fp16)[name = string("transpose_2")];
+            tensor<fp16, [1, 1, 16388]> logits4 = transpose(perm = var_112_perm_0, x = var_109_cast_fp16)[name = string("transpose_3")];
+            tensor<fp16, [1, 1, 16388]> logits5 = transpose(perm = var_138_perm_0, x = var_135_cast_fp16)[name = string("transpose_4")];
+            tensor<fp16, [1, 1, 16388]> logits6 = transpose(perm = var_164_perm_0, x = var_161_cast_fp16)[name = string("transpose_5")];
+            tensor<fp16, [1, 1, 16388]> logits7 = transpose(perm = var_190_perm_0, x = var_187_cast_fp16)[name = string("transpose_6")];
+            tensor<fp16, [1, 1, 16388]> logits8 = transpose(perm = var_216_perm_0, x = var_213_cast_fp16)[name = string("transpose_7")];
+            tensor<fp16, [1, 1, 16388]> logits9 = transpose(perm = var_242_perm_0, x = var_239_cast_fp16)[name = string("transpose_8")];
+            tensor<fp16, [1, 1, 16388]> logits10 = transpose(perm = var_268_perm_0, x = var_265_cast_fp16)[name = string("transpose_9")];
+            tensor<fp16, [1, 1, 16388]> logits11 = transpose(perm = var_294_perm_0, x = var_291_cast_fp16)[name = string("transpose_10")];
+            tensor<fp16, [1, 1, 16388]> logits12 = transpose(perm = var_320_perm_0, x = var_317_cast_fp16)[name = string("transpose_11")];
+            tensor<fp16, [1, 1, 16388]> logits13 = transpose(perm = var_346_perm_0, x = var_343_cast_fp16)[name = string("transpose_12")];
+            tensor<fp16, [1, 1, 16388]> logits14 = transpose(perm = var_372_perm_0, x = var_369_cast_fp16)[name = string("transpose_13")];
+            tensor<fp16, [1, 1, 16388]> logits15 = transpose(perm = var_398_perm_0, x = var_395_cast_fp16)[name = string("transpose_14")];
+            tensor<fp16, [1, 1, 16388]> logits16 = transpose(perm = var_424_perm_0, x = var_421_cast_fp16)[name = string("transpose_15")];
+        } -> (logits1, logits2, logits3, logits4, logits5, logits6, logits7, logits8, logits9, logits10, logits11, logits12, logits13, logits14, logits15, logits16);
+}

gemma3_lm_head_lut6.mlmodelc/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b278ded6af56b5a6162fcadbfb0d9b3964339408cc02647e317593886012fc44
+size 511832128

meta.yaml

@@ -0,0 +1,28 @@
+model_info:
+  name: anemll-google-gemma-3-4b-it-qat-int4-unquantized-ctx1024
+  version: 0.3.5
+  description: |
+    Demonstarates running google-gemma-3-4b-it-qat-int4-unquantized on Apple Neural Engine
+    Context length: 1024
+    Batch size: 64
+    Chunks: 2
+  license: MIT
+  author: Anemll
+  framework: Core ML
+  language: Python
+  architecture: gemma3
+  parameters:
+    context_length: 1024
+    batch_size: 64
+    lut_embeddings: none
+    lut_ffn: 4
+    lut_ffn_per_channel: 4
+    lut_lmhead: 6
+    lut_lmhead_per_channel: 4
+    num_chunks: 2
+    model_prefix: gemma3
+    embeddings: gemma3_embeddings.mlmodelc
+    lm_head: gemma3_lm_head_lut6.mlmodelc
+    ffn: gemma3_FFN_PF_lut4_chunk_01of02.mlmodelc
+    split_lm_head: 16
+    sliding_window: 1024

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7d4046bf0505a327dd5a0abbb427ecd4fc82f99c2ceaa170bc61ecde12809b0c
+size 33384570

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1299c11d7cf632ef3b4e11937501358ada021bbdf7c47638d13c0ee982f2e79c
+size 4689074