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-1b-it-ctx4096_0.3.4
+```
+
+2. Extract model files:
+```bash
+# Navigate to cloned directory
+cd anemll-google-gemma-3-1b-it-ctx4096_0.3.4
+
+# 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-1b-it-ctx4096_0.3.4
+
+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: 4096
+- Batch Size: 64
+- Number of Chunks: 1
+- LUT Quantization: 6
+
+## 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.py

@@ -0,0 +1,1950 @@
+# 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 json
+import sys
+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
+import resource
+
+
+def _get_rss_mb() -> float:
+    """Best-effort RSS in MB (macOS/Linux)."""
+    try:
+        # On macOS ru_maxrss is bytes; on Linux it is kilobytes.
+        rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
+        if rss > 10_000_000:  # heuristic: likely bytes
+            return rss / (1024 * 1024)
+        return rss / 1024  # kB -> MB
+    except Exception:
+        return -1.0
+
+
+def _maybe_report_mem(label: str, enabled: bool):
+    if not enabled:
+        return
+    rss_mb = _get_rss_mb()
+    if rss_mb >= 0:
+        print(f"[mem] {label}: rss≈{rss_mb:.1f} MB")
+
+# ANSI color codes
+LIGHT_BLUE = "\033[94m"
+DARK_BLUE = "\033[34m"
+LIGHT_GREEN = "\033[92m"
+RESET_COLOR = "\033[0m"
+
+# Add at top with other constants
+WARMUP_TOKEN_LIMIT = 10  # Maximum tokens to generate during warmup
+
+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
+        # Add token counting and timing
+        self.start_time = time.time()
+        self.token_count = 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, eval_mode=False):
+        """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()
+        
+        # Store the text in buffer for later saving to file
+        with self.lock:
+            self.buffer += token_str
+
+        # Skip printing in eval mode
+        if eval_mode:
+            return
+
+        # 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, eval_mode=False):
+        """Stop the printer thread."""
+        if self.thread and self.thread.is_alive():
+            # Ensure any remaining tokens are processed
+            self.drain_buffer()
+            self.stop_event.set()
+            try:
+                self.thread.join(timeout=1.0)
+            except Exception:
+                pass
+            # Calculate and print tokens/s with shorter format in blue (unless in eval mode)
+            if not eval_mode:
+                elapsed = time.time() - self.start_time
+                if elapsed > 0 and self.token_count > 0:
+                    tokens_per_sec = self.token_count / elapsed
+                    print(f"\n{DARK_BLUE}{tokens_per_sec:.1f} t/s{RESET_COLOR}")
+                else:
+                    print(RESET_COLOR)  # Reset color at the end
+        return self.buffer
+
+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():
+            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():
+                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:
+            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
+        if hasattr(tokenizer, "get_vocab"):
+            vocab = tokenizer.get_vocab()
+            if token_str in vocab:
+                return vocab[token_str]
+        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
+
+    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 parse_ffn_filename(path):
+    """Parse FFN model filename to extract chunk information."""
+    path = Path(path)
+    # Support multiple naming conventions:
+    # - FFN_PF_lut6_chunk_01of04 (legacy/prefill style)
+    # - gemma3_1b_FFN_lut6_chunk_01of04 (new Gemma3 style)
+    # - any_prefix_FFN_*_chunk_NNofNN
+    pattern = r'FFN[^/]*_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 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 not args.eval:
+            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:
+        if not args.eval:
+            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)
+        if not args.eval:
+            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
+        if not args.eval:
+            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
+        if not args.eval:
+            print(f"\nOverriding num chunks from args: {args.num_chunks}")
+    
+    return metadata
+    
+def load_models(args,metadata):
+    """Load all required models and extract metadata."""
+    if not args.eval:
+        print("\nLoading models...")
+    _maybe_report_mem("start load_models", getattr(args, "mem_report", False))
+
+    # Determine compute unit
+    compute_unit = ct.ComputeUnit.CPU_ONLY if getattr(args, 'cpu', False) else ct.ComputeUnit.CPU_AND_NE
+    if not args.eval and getattr(args, 'cpu', False):
+        print("Running in CPU-only mode")
+
+    try:
+        # Load embeddings model
+        if not args.eval:
+            print("\nLoading embeddings model...")
+        embed_path = parse_model_path(args.embed)
+        if not args.eval:
+            print(f"Loading from: {embed_path}")
+        embed_model = load_model(embed_path, compute_unit=compute_unit)
+        if not args.eval:
+            print("Embeddings model loaded successfully")
+        _maybe_report_mem("after embeddings load", getattr(args, "mem_report", False))
+        metadata = load_metadata(embed_model,args)
+        
+
+        
+        # Load LM head model
+        if not args.eval:
+            print("\nLoading LM head model...")
+        lmhead_path = parse_model_path(args.lmhead)
+        if not args.eval:
+            print(f"Loading from: {lmhead_path}")
+        lmhead_model = load_model(lmhead_path, compute_unit=compute_unit)
+        if not args.eval:
+            print("LM head model loaded successfully")
+        _maybe_report_mem("after lmhead load", getattr(args, "mem_report", False))
+        
+        # Parse FFN path and find chunks if needed
+        if not args.eval:
+            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:
+            if not args.eval:
+                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:
+                if not args.eval:
+                    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 (Gemma3 with context > 512)
+                    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)
+                        if not args.eval:
+                            print("  Rotation functions loaded (4-function model)")
+                    except Exception:
+                        # Rotation functions not available - standard 2-function model
+                        pass
+                    ffn_models.append(chunk_dict)
+                    if not args.eval:
+                        print("Chunk loaded successfully")
+                    _maybe_report_mem(f"after FFN chunk load {Path(chunk_path).name}", getattr(args, "mem_report", False))
+                except Exception as e:
+                    if not args.eval:
+                        print(f"Error loading chunk {chunk_path}: {str(e)}")
+                    raise
+            metadata = load_metadata(ffn_models[0],args)
+
+        else:
+            if not args.eval:
+                print("\nLoading single FFN model...")
+            ffn_models.append(load_model(ffn_path, compute_unit=compute_unit))
+            if not args.eval:
+                print("FFN model loaded successfully")
+            _maybe_report_mem("after FFN load", getattr(args, "mem_report", False))
+        
+        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, eval_mode=False):
+    """Initialize and configure the tokenizer."""
+    try:
+
+        
+        tokenizer = AutoTokenizer.from_pretrained(
+            str(model_path), 
+            use_fast=False,
+            trust_remote_code=True
+        )
+
+        # Try to load a chat template if the tokenizer doesn't have one.
+        if getattr(tokenizer, "chat_template", None) in (None, "") and model_path:
+            template = None
+            config_path = Path(model_path) / "tokenizer_config.json"
+            if config_path.exists():
+                try:
+                    config_data = json.loads(config_path.read_text())
+                    template = config_data.get("chat_template")
+                except Exception as e:
+                    if not eval_mode:
+                        print(f"Warning: Failed to read tokenizer_config.json chat_template: {e}")
+            if template is None:
+                jinja_path = Path(model_path) / "chat_template.jinja"
+                if jinja_path.exists():
+                    template = jinja_path.read_text()
+            if template:
+                tokenizer.chat_template = template
+                if not eval_mode:
+                    print("Loaded chat_template from model files")
+        
+        if not eval_mode:
+            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
+            if not eval_mode:
+                print("Set PAD token to EOS token")
+        
+        tokenizer.padding_side = "left"
+        
+        if not eval_mode:
+            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 appears to be a tokenizer loading issue.")
+        
+        # Check if it's the specific Qwen tokenizer file issue
+        if "expected str, bytes or os.PathLike object, not NoneType" in str(e):
+            print("\nThis error suggests the tokenizer files are missing or incomplete.")
+            print("For Qwen models, you need the original model directory with tokenizer files.")
+            print("Try using: --tokenizer ~/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/YOUR_SNAPSHOT_ID")
+        else:
+            print("Please provide the path to a compatible model directory with tokenizer files.")
+        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 initialize_causal_mask(context_length, eval_mode=False):
+    """Initialize causal mask for transformer attention."""
+    causal_mask = make_causal_mask(context_length, 0)
+    causal_mask = torch.tensor(causal_mask, dtype=torch.float16)
+    if not eval_mode:
+        print(f"\nInitialized causal mask for context length {context_length}")
+    return causal_mask
+
+def run_prefill(embed_model, ffn_models, input_ids, context_pos, context_length, batch_size=64, state=None, causal_mask=None, 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)
+    """
+    # Use provided causal mask or create one if not provided
+    if causal_mask is None:
+        causal_mask = make_causal_mask(context_length, 0)
+        causal_mask = torch.tensor(causal_mask, dtype=torch.float16)
+
+    # 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 < 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]
+
+        # Always pad to full batch size for prefill
+        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)  # Changed: Always use full batch size
+        batch_causal_mask = causal_mask[:, :, batch_pos:batch_pos+batch_size, :]  # Changed: Use full batch size
+
+        # Run embeddings
+        hidden_states = torch.from_numpy(
+            embed_model.predict({
+                'input_ids': batch_input.numpy().astype(np.int32)
+            })['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 with state
+        for ffn_model in ffn_models:
+            if isinstance(ffn_model, dict):
+                inputs = {
+                    'hidden_states': hidden_states.numpy().astype(np.float16),  # [1, 64, hidden_size]
+                    'position_ids': position_ids.numpy().astype(np.int32),    # [64]
+                    'causal_mask': batch_causal_mask.numpy().astype(np.float16), # [1, 1, 64, context_length]
+                    'current_pos': np.array([batch_pos], dtype=np.int32)  # [1]
+                }
+                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([context_pos], dtype=torch.int32)
+
+def generate_next_token(embed_model, ffn_models, lmhead_model, input_ids, pos, context_length, metadata, state=None, causal_mask=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)
+    """
+    debug = metadata.get('debug', False)
+    attention_size = metadata.get('attention_size', context_length)
+    sliding_window = metadata.get('sliding_window', 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]  # [1, 1]
+
+    # Ensure proper data type for CoreML
+    current_token_array = current_token.numpy().astype(np.int32)
+
+    # Run embeddings
+    hidden_states = torch.from_numpy(
+        embed_model.predict({'input_ids': current_token_array})['hidden_states']
+    )  # [1, 1, hidden_size]
+
+    # 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)  # [1]
+
+    # Use provided causal mask or create one if not provided
+    if causal_mask is None:
+        causal_mask_data = make_causal_mask(context_length, 0)
+        single_causal_mask = torch.tensor(causal_mask_data[:, :, pos-1:pos, :], dtype=torch.float16)  # [1, 1, 1, context_length]
+    else:
+        single_causal_mask = causal_mask[:, :, pos-1:pos, :]
+
+    if debug:
+        print(f"\n[DEBUG] generate_next_token: pos={pos}, context_length={context_length}, attention_size={attention_size}")
+        print(f"[DEBUG]   position_ids={position_ids.item()}, current_token={current_token.item()}")
+        print(f"[DEBUG]   causal_mask shape={single_causal_mask.shape}")
+        print(f"[DEBUG]   hidden_states shape={hidden_states.shape}")
+
+    # Run through FFN chunks with state
+    for ffn_model in ffn_models:
+        if isinstance(ffn_model, dict):
+            # Build inputs dict - only include inputs that the model expects
+            inputs = {
+                'hidden_states': hidden_states.numpy().astype(np.float16),
+                '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)
+            }
+            # Add update_mask only if model expects it (older models)
+            # Get model input names from the spec
+            try:
+                model_inputs = {inp.name for inp in ffn_model[infer_func_name].get_spec().description.input}
+            except:
+                model_inputs = set()
+            if 'update_mask' in model_inputs:
+                inputs['update_mask'] = update_mask.numpy().astype(np.float16)
+            if debug:
+                print(f"[DEBUG]   FFN {infer_func_name} inputs: position_ids={inputs['position_ids']}, current_pos={inputs['current_pos']}")
+                print(f"[DEBUG]   FFN {infer_func_name} causal_mask shape={inputs['causal_mask'].shape}")
+            try:
+                output = ffn_model[infer_func_name].predict(inputs, state)
+            except Exception as e:
+                print(f"\n[ERROR] FFN {infer_func_name} failed at pos={pos}, position_ids={position_ids.item()}")
+                print(f"[ERROR]   context_length={context_length}, attention_size={attention_size}")
+                print(f"[ERROR]   causal_mask shape={single_causal_mask.shape}")
+                print(f"[ERROR]   Exception: {e}")
+                raise
+            hidden_states = torch.from_numpy(output['output_hidden_states'])
+    
+    # Run LM head
+    lm_output = lmhead_model.predict({'hidden_states': hidden_states.numpy().astype(np.float16)})
+
+    # 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)
+
+    # Debug: show LM head output keys if argmax mode 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("[WARNING] Model may need to be reconverted with --argmax flag")
+        # Fall through to logits processing
+
+    if argmax_in_model and 'argmax_idx' in lm_output:
+        # Argmax-in-model mode: find the chunk with highest value and compute global index
+        # Model outputs LOCAL indices (0 to chunk_size-1) for each chunk
+        # We compute: global_idx = local_idx + (best_chunk * chunk_size)
+        argmax_idx = lm_output['argmax_idx']  # shape: [num_chunks], LOCAL indices
+        argmax_val = lm_output['argmax_val']  # shape: [num_chunks]
+
+        # Flatten arrays
+        argmax_idx_flat = argmax_idx.flatten()
+        argmax_val_flat = argmax_val.flatten()
+
+        # Find best chunk (highest value)
+        best_chunk = int(np.argmax(argmax_val_flat))
+        local_idx = int(argmax_idx_flat[best_chunk])
+
+        # Compute global index: local_idx + (best_chunk * chunk_size)
+        num_chunks = len(argmax_idx_flat)
+        chunk_size = 262144 // num_chunks  # Gemma3 vocab = 262144
+        global_idx = local_idx + (best_chunk * chunk_size)
+
+        if metadata.get('debug_argmax', False):
+            print(f"\nLM head argmax mode (chunked):")
+            print(f"  argmax_idx shape: {argmax_idx.shape}, dtype: {argmax_idx.dtype}")
+            print(f"  argmax_val shape: {argmax_val.shape}, dtype: {argmax_val.dtype}")
+            print(f"  best_chunk={best_chunk}, local_idx={local_idx}, global_idx={global_idx}")
+            print(f"  best_val={argmax_val_flat[best_chunk]:.4f}")
+
+        return global_idx
+
+    # Get number of logits from metadata, using split_lm_head if available
+    # First check for split_lm_head (new), then num_logits (legacy), default to 8
+    num_logits = metadata.get('split_lm_head', metadata.get('num_logits', 8))
+
+    # Combine logits1-N if they exist
+    if 'logits1' in lm_output:
+        # Concatenate all logits parts
+        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)  # Concatenate along vocab dimension
+    else:
+        # Try output_logits as fallback
+        logits = torch.from_numpy(lm_output['output_logits'])
+
+    # 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 create_unified_state(ffn_models, context_length, eval_mode=False):
+    """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()
+        if not eval_mode:
+            print(f"\nCreated unified transformer state for {len(ffn_models)} chunks")
+        return state
+    else:
+        state = ffn_models[0].make_state()
+        if not eval_mode:
+            print("\nCreated unified transformer state")
+        return state
+
+def chat_loop(embed_model, ffn_models, lmhead_model, tokenizer, metadata, state, causal_mask=None, auto_prompt=None, warmup=False, save_file=None, max_tokens=None, no_template=False, eval_mode=False):
+    """Interactive chat loop."""
+    context_length = metadata.get('context_length')
+    batch_size = metadata.get('batch_size', 64)
+    
+    if not warmup and not eval_mode:
+        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.")
+    
+    # Check if tokenizer has chat template and if it works
+    has_chat_template = False
+    try:
+        # Test if chat template works
+        test_messages = [{"role": "user", "content": "test"}]
+        tokenizer.apply_chat_template(test_messages, return_tensors="pt")
+        has_chat_template = True
+        if not warmup and not eval_mode:
+            print("\nUsing chat template for prompts")
+    except:
+        if not warmup and not eval_mode:
+            print("\nUsing manual formatting for prompts")
+
+    stop_token_ids = build_stop_token_ids(tokenizer)
+    
+    conversation = []
+    
+    try:
+        while True:
+            try:
+                if not warmup and not eval_mode:
+                    print(f"\n{LIGHT_GREEN}You:{RESET_COLOR}", end=' ', flush=True)
+                if auto_prompt is not None:
+                    user_input = auto_prompt
+                    if not warmup and not eval_mode:
+                        print(user_input)
+                else:
+                    user_input = input().strip()
+            except EOFError:
+                if not warmup and not eval_mode:
+                    print("\nExiting chat...")
+                break
+                
+            if not user_input:
+                continue
+            
+            # Format prompt based on no_template flag and tokenizer capabilities
+            if no_template:
+                # Use raw input without any chat template formatting
+                input_ids = tokenizer(
+                    user_input,
+                    return_tensors="pt",
+                    add_special_tokens=True
+                ).input_ids.to(torch.int32)
+                if not warmup and not eval_mode:
+                    print("Using raw input without chat template")
+            elif has_chat_template:
+                messages = [{"role": "user", "content": user_input}]
+                input_ids = tokenizer.apply_chat_template(
+                    messages,
+                    return_tensors="pt",
+                    add_generation_prompt=True
+                ).to(torch.int32)
+            else:
+                # Manual formatting for Llama models without chat template
+                formatted_prompt = f"[INST] {user_input} [/INST]"
+                input_ids = tokenizer(
+                    formatted_prompt,
+                    return_tensors="pt",
+                    add_special_tokens=True
+                ).input_ids.to(torch.int32)
+            
+            context_pos = input_ids.size(1)
+            
+            if not warmup and not eval_mode:
+                print(f"\n{LIGHT_BLUE}Assistant:{RESET_COLOR}", end=' ', flush=True)
+            
+            # Initialize token printer
+            token_printer = TokenPrinter(tokenizer)
+            tokens_generated = 0  # Track number of tokens
+            
+            try:
+                # Start prefill timing
+                prefill_start = time.time()
+                
+                # Run prefill with state and causal mask
+                # Ensure batch_size is not None
+                if batch_size is None:
+                    batch_size = 64
+                    if not eval_mode:
+                        print(f"Warning: batch_size was None, using default: {batch_size}")
+                
+                # Get sliding_window for rotation support (Gemma3)
+                sliding_window = metadata.get('sliding_window', None)
+
+                _ = run_prefill(
+                    embed_model,
+                    ffn_models,
+                    input_ids,
+                    context_pos,
+                    context_length,
+                    batch_size,
+                    state,
+                    causal_mask,
+                    sliding_window
+                )
+                
+                # Calculate prefill timing
+                prefill_time = time.time() - prefill_start
+                prefill_tokens = context_pos  # Number of tokens in input
+                prefill_tokens_per_sec = prefill_tokens / prefill_time if prefill_time > 0 else 0
+                
+                # Generation loop with state
+                input_ids = input_ids
+                pos = context_pos
+                inference_start = time.time()
+                inference_tokens = 0
+                
+                while pos < context_length - 1:
+                    # Generate next token with causal mask
+                    next_token = generate_next_token(
+                        embed_model,
+                        ffn_models,
+                        lmhead_model,
+                        input_ids,
+                        pos,
+                        context_length,
+                        metadata,
+                        state,
+                        causal_mask
+                    )
+                    
+                    # Add token to sequence
+                    if pos < input_ids.size(1):
+                        input_ids[0, pos] = next_token
+                    else:
+                        input_ids = torch.cat([
+                            input_ids,
+                            torch.tensor([[next_token]], dtype=torch.int32)
+                        ], dim=1)
+                    
+                    # Add to printer only if not in warmup
+                    if not warmup:
+                        token_printer.add_token(next_token)
+                        token_printer.drain_buffer(eval_mode)
+                    
+                    pos += 1
+                    tokens_generated += 1
+                    inference_tokens += 1
+                    
+                    # Check limits
+                    if warmup and tokens_generated >= WARMUP_TOKEN_LIMIT:
+                        break
+                    
+                    # Check max_tokens limit
+                    if max_tokens is not None and tokens_generated >= max_tokens:
+                        break
+                        
+                    if next_token in stop_token_ids:
+                        break
+                
+                # Calculate inference timing
+                inference_time = time.time() - inference_start
+                inference_tokens_per_sec = inference_tokens / inference_time if inference_time > 0 else 0
+                
+                # Get final response and add to conversation
+                if not warmup:
+                    response = token_printer.stop(eval_mode)
+                    if eval_mode:
+                        # In eval mode, only print the model response
+                        print(response, end='')
+                    else:
+                        # Print timing stats
+                        prefill_ms = prefill_time * 1000  # Convert to milliseconds
+                        print(f"\nPrefill: {prefill_ms:.1f}ms ({prefill_tokens_per_sec:.1f} t/s)")
+                        print(f"Inference: {inference_tokens_per_sec:.1f} t/s")
+                        print(f"Total: Generated {tokens_generated} tokens in {prefill_time + inference_time:.2f}s")
+                    conversation.append({"role": "assistant", "content": response})
+                    
+                    # Save response to file if requested
+                    if save_file and not eval_mode:
+                        try:
+                            # Add small delay to ensure all tokens are processed
+                            time.sleep(0.5)
+                            
+                            # Make sure response ends with EOS token if it's supposed to
+                            if response and not response.endswith("<|eot_id|>") and not response.endswith("</s>") and not response.endswith("<end_of_turn>"):
+                                if tokenizer.eos_token:
+                                    eos_text = tokenizer.decode([tokenizer.eos_token_id])
+                                    if not response.endswith(eos_text):
+                                        print(f"\n{DARK_BLUE}Adding missing EOS token for consistency{RESET_COLOR}")
+                                        response += eos_text
+                            
+                            with open(save_file, 'w') as f:
+                                f.write(response)
+                            print(f"\n{DARK_BLUE}Response saved to file: {save_file}{RESET_COLOR}")
+                        except Exception as e:
+                            print(f"\n{DARK_BLUE}Error saving to file: {str(e)}{RESET_COLOR}")
+                else:
+                    token_printer.stop(eval_mode)  # Clean up without printing stats
+                
+                # Exit after one response in auto_prompt mode
+                if auto_prompt is not None:
+                    break
+                
+            except KeyboardInterrupt:
+                if not eval_mode:
+                    print("\nGeneration interrupted")
+                token_printer.stop(eval_mode)
+                continue
+                
+    except Exception as e:
+        print(f"\nError in chat loop: {str(e)}")
+        import traceback
+        traceback.print_exc()
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Chat with CoreML LLaMA, gil resolved  (c) 2025 Anemll')
+    
+    # Add meta.yaml option
+    parser.add_argument('--meta', type=str, help='Path to meta.yaml to load all parameters')
+    
+    # Model paths
+    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')
+    
+    # Add save option
+    parser.add_argument('--save', type=str,
+                       help='Save assistant\'s response to specified file')
+    
+    # Add max-tokens option
+    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 no-template flag  
+    parser.add_argument('--no-template', action='store_true',
+                       help='Prefill the question itself and start inference directly without chat template')
+    
+    # Add eval mode flag
+    parser.add_argument('--eval', action='store_true',
+                       help='Evaluation mode: suppress all output except model response')
+
+    # Add CPU-only mode
+    parser.add_argument('--cpu', action='store_true',
+                       help='Run on CPU only (no ANE/GPU)')
+    parser.add_argument('--mem-report', action='store_true',
+                       help='Print approximate RSS after large steps (debugging)')
+
+    # 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('--num-logits', type=int, default=8,
+                       help='Number of logits outputs from LM head (default: 8, legacy)')
+    parser.add_argument('--split-lm-head', type=int,
+                       help='Number of logits splits from LM head (default: 8 for llama, 16 for qwen)')
+    parser.add_argument('--debug-argmax', action='store_true',
+                       help='Enable debug output for argmax mode (print indices and values)')
+    parser.add_argument('--debug', action='store_true',
+                       help='Enable debug output (position, state, shapes)')
+
+    args = parser.parse_args()
+
+    def _strip_model_ext(value):
+        if value is None:
+            return None
+        return value.replace('.mlmodelc', '').replace('.mlpackage', '')
+    
+    # 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'])
+                # state_length for split cache models (defaults to context_length if not specified)
+                args.state_length = int(params.get('state_length', args.context_length))
+                if args.batch_size is None:
+                    args.batch_size = int(params['batch_size'])
+                args.num_chunks = 1  # Monolithic has no chunks
+
+                # 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
+
+                # Check for argmax_in_model flag
+                args.argmax_in_model = params.get('argmax_in_model', False)
+
+                # Set tokenizer path
+                if not args.tokenizer:
+                    if 'tokenizer_path' in params:
+                        args.tokenizer = params['tokenizer_path']
+                    else:
+                        args.tokenizer = args.d
+
+                if not args.eval:
+                    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
+                prefix = params.get('model_prefix', 'llama')
+                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:
+                    if 'lm_head' in params:
+                        args.lmhead = _strip_model_ext(params['lm_head'])
+                    else:
+                        args.lmhead = f'{prefix}_lm_head{lut_lmhead}'
+                if not args.embed:
+                    if 'embeddings' in params:
+                        args.embed = _strip_model_ext(params['embeddings'])
+                    else:
+                        args.embed = f'{prefix}_embeddings{lut_embeddings}'
+                if not args.ffn:
+                    if 'ffn' in params:
+                        ffn_candidate = _strip_model_ext(params['ffn'])
+                        ffn_path = Path(ffn_candidate)
+                        if "_chunk_" not in ffn_candidate:
+                            default_ffn = f'{prefix}_FFN_PF{lut_ffn}_chunk_01of{num_chunks:02d}'
+                            base_dir = ffn_path.parent if ffn_path.is_absolute() else Path(args.d)
+                            if (base_dir / f"{default_ffn}.mlmodelc").exists() or (base_dir / f"{default_ffn}.mlpackage").exists():
+                                args.ffn = str(base_dir / default_ffn) if ffn_path.is_absolute() else default_ffn
+                            else:
+                                args.ffn = ffn_candidate
+                        else:
+                            args.ffn = ffn_candidate
+                    else:
+                        args.ffn = f'{prefix}_FFN_PF{lut_ffn}_chunk_01of{num_chunks:02d}'
+                if not args.tokenizer:
+                    if 'tokenizer_path' in params:
+                        args.tokenizer = params['tokenizer_path']
+                    else:
+                        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
+                if 'num_logits' in params:
+                    args.num_logits = int(params['num_logits'])
+                # attention_size is used for causal mask (sliding window for Gemma3)
+                args.attention_size = int(params.get('attention_size', args.context_length))
+
+                # 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
+
+                # 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 = 8
+
+                # Check for argmax_in_model flag (for chunked models)
+                args.argmax_in_model = params.get('argmax_in_model', False)
+
+                if not args.eval:
+                    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"  Num Logits: {args.num_logits}")
+                    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
+        if not hasattr(args, 'split_lm_head') or args.split_lm_head is None:
+            args.split_lm_head = args.num_logits  # Use num_logits as fallback
+
+    return args
+
+
+def load_monolithic_model(args, metadata):
+    """Load monolithic model with infer, infer_rotate, and prefill functions."""
+    if not args.eval:
+        print("\nLoading monolithic model...")
+    _maybe_report_mem("start load_monolithic_model", getattr(args, "mem_report", False))
+
+    # Determine compute unit
+    compute_unit = ct.ComputeUnit.CPU_ONLY if getattr(args, 'cpu', False) else ct.ComputeUnit.CPU_AND_NE
+    if not args.eval and 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)
+
+    if not args.eval:
+        print(f"Loading from: {model_path}")
+
+    def _progress_bar(done, total, label, width=18):
+        if total <= 0:
+            total = 1
+        filled = int(width * done / total)
+        bar = "[" + ("#" * filled) + ("." * (width - filled)) + "]"
+        sys.stdout.write(f"\r{bar} {done}/{total} {label}")
+        sys.stdout.flush()
+        if done == total:
+            sys.stdout.write("\n")
+
+    # Decide whether to attempt rotate functions
+    attempt_rotate = True
+    if getattr(args, "context_length", None) is not None and args.context_length <= 512:
+        attempt_rotate = False
+
+    functions_to_load = [("infer", True), ("prefill", True)]
+    if attempt_rotate:
+        functions_to_load += [("infer_rotate", False), ("prefill_rotate", False)]
+
+    infer_model = None
+    prefill_model = None
+    infer_rotate_model = None
+    prefill_rotate_model = None
+    loaded = []
+    missing = []
+
+    total = len(functions_to_load)
+    for idx, (name, required) in enumerate(functions_to_load, start=1):
+        if not args.eval:
+            _progress_bar(idx, total, name)
+        try:
+            model = load_model(model_path, function_name=name, compute_unit=compute_unit)
+            loaded.append(name)
+            if name == "infer":
+                infer_model = model
+            elif name == "prefill":
+                prefill_model = model
+            elif name == "infer_rotate":
+                infer_rotate_model = model
+            elif name == "prefill_rotate":
+                prefill_rotate_model = model
+        except Exception:
+            if required:
+                raise
+            missing.append(name)
+
+    _maybe_report_mem("after load monolithic functions", getattr(args, "mem_report", False))
+
+    if not args.eval:
+        summary = "Monolithic model loaded (" + ", ".join(loaded) + ")"
+        if missing:
+            summary += f" [missing: {', '.join(missing)}]"
+        print(summary)
+
+    # 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
+    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
+    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:
+            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:
+        # Argmax-in-model mode: find the chunk with highest value and compute global index
+        # Model outputs LOCAL indices (0 to chunk_size-1) for each chunk
+        # We compute: global_idx = local_idx + (best_chunk * chunk_size)
+        argmax_idx = output['argmax_idx']  # shape: [num_chunks], LOCAL indices
+        argmax_val = output['argmax_val']  # shape: [num_chunks]
+
+        # Flatten in case of extra dimensions
+        argmax_idx_flat = argmax_idx.flatten()
+        argmax_val_flat = argmax_val.flatten()
+
+        # Find chunk with highest value
+        best_chunk = int(np.argmax(argmax_val_flat))
+        local_idx = int(argmax_idx_flat[best_chunk])
+
+        # Compute global token ID: local_idx + (chunk * chunk_size)
+        chunk_size = 16384  # 262144 / 16
+        global_idx = local_idx + (best_chunk * chunk_size)
+
+        # Debug: print shapes and values
+        if metadata.get('debug_argmax', False):
+            print(f"\n=== Argmax Debug (pos={pos}) ===")
+            print(f"argmax_idx shape: {argmax_idx.shape}, dtype: {argmax_idx.dtype}")
+            print(f"argmax_val shape: {argmax_val.shape}, dtype: {argmax_val.dtype}")
+            print(f"Per-chunk results (LOCAL indices, chunk_size={chunk_size}):")
+
+            # Find top 3 chunks by value for comparison
+            sorted_indices = np.argsort(argmax_val_flat)[::-1][:3]
+
+            for i in range(min(16, len(argmax_idx_flat))):
+                local = int(argmax_idx_flat[i])
+                val = float(argmax_val_flat[i])
+                computed_global = local + (i * chunk_size)
+                in_range = 0 <= local < chunk_size
+                marker = " <-- SELECTED" if i == best_chunk else ""
+                if i in sorted_indices and i != best_chunk:
+                    marker += f" (top-{list(sorted_indices).index(i)+1})"
+                range_ok = "✓" if in_range else f"✗ (expected 0-{chunk_size-1})"
+                print(f"  Chunk {i:2d}: local={local:5d}, global={computed_global:6d}, val={val:8.4f}, range={range_ok}{marker}")
+
+            print(f"Result: best_chunk={best_chunk}, local_idx={local_idx}, global_idx={global_idx}, best_val={argmax_val_flat[best_chunk]:.4f}")
+
+            # Value comparison: show if there are close competing values
+            top_values = [float(argmax_val_flat[i]) for i in sorted_indices]
+            if len(top_values) >= 2:
+                val_diff = abs(top_values[0] - top_values[1])
+                print(f"Value comparison: top-1={top_values[0]:.6f}, top-2={top_values[1]:.6f}, diff={val_diff:.6f}")
+                if val_diff < 0.01:
+                    print(f"  WARNING: Values are very close - possible precision issue!")
+
+        return global_idx
+
+    # 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:
+        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=None,
+                         auto_prompt=None, warmup=False, save_file=None, max_tokens=None,
+                         no_template=False, eval_mode=False, infer_rotate_model=None, prefill_rotate_model=None):
+    """Chat loop for monolithic models.
+
+    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
+        save_file: Optional file to save conversation
+        max_tokens: Maximum tokens to generate
+        no_template: If True, don't use chat template
+        eval_mode: If True, minimal output for evaluation
+        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).
+    """
+    context_length = metadata.get('context_length')
+    batch_size = metadata.get('batch_size', 64)
+    sliding_window = metadata.get('sliding_window', 512)  # For switching between infer modes
+
+    if not warmup and not eval_mode:
+        print(f"\nUsing context length: {context_length}")
+        if infer_rotate_model is not None:
+            print(f"Cache rotation: ENABLED (infer_rotate function available)")
+        else:
+            print(f"Cache rotation: NOT AVAILABLE (using infer for all positions)")
+        print("\nStarting chat session. Press Ctrl+D to exit.")
+
+    # Check chat template
+    has_chat_template = False
+    try:
+        test_messages = [{"role": "user", "content": "test"}]
+        tokenizer.apply_chat_template(test_messages, return_tensors="pt")
+        has_chat_template = True
+        if not warmup and not eval_mode:
+            print("\nUsing chat template for prompts")
+    except:
+        if not warmup and not eval_mode:
+            print("\nUsing manual formatting for prompts")
+
+    stop_token_ids = build_stop_token_ids(tokenizer)
+
+    try:
+        while True:
+            try:
+                if not warmup and not eval_mode:
+                    print(f"\n{LIGHT_GREEN}You:{RESET_COLOR}", end=' ', flush=True)
+                if auto_prompt is not None:
+                    user_input = auto_prompt
+                    if not warmup and not eval_mode:
+                        print(user_input)
+                else:
+                    user_input = input().strip()
+            except EOFError:
+                if not warmup and not eval_mode:
+                    print("\nExiting chat...")
+                break
+
+            if not user_input:
+                continue
+
+            # Format prompt
+            if no_template:
+                input_ids = tokenizer(user_input, return_tensors="pt", add_special_tokens=True).input_ids.to(torch.int32)
+            elif has_chat_template:
+                messages = [{"role": "user", "content": user_input}]
+                input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True).to(torch.int32)
+            else:
+                formatted_prompt = f"[INST] {user_input} [/INST]"
+                input_ids = tokenizer(formatted_prompt, return_tensors="pt", add_special_tokens=True).input_ids.to(torch.int32)
+
+            context_pos = input_ids.size(1)
+
+            if not warmup and not eval_mode:
+                print(f"\n{LIGHT_BLUE}Assistant:{RESET_COLOR}", end=' ', flush=True)
+
+            token_printer = TokenPrinter(tokenizer)
+            tokens_generated = 0
+
+            try:
+                prefill_start = time.time()
+
+                # Run prefill with monolithic model (uses rotation for pos >= sliding_window if available)
+                _ = 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
+                )
+
+                prefill_time = time.time() - prefill_start
+                prefill_tokens_per_sec = context_pos / prefill_time if prefill_time > 0 else 0
+
+                # Generation loop
+                pos = context_pos
+                inference_start = time.time()
+                inference_tokens = 0
+
+                while pos < context_length - 1:
+                    # 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
+                    )
+
+                    if pos < input_ids.size(1):
+                        input_ids[0, pos] = next_token
+                    else:
+                        input_ids = torch.cat([input_ids, torch.tensor([[next_token]], dtype=torch.int32)], dim=1)
+
+                    if not warmup:
+                        token_printer.add_token(next_token)
+                        token_printer.drain_buffer(eval_mode)
+
+                    pos += 1
+                    tokens_generated += 1
+                    inference_tokens += 1
+
+                    if warmup and tokens_generated >= WARMUP_TOKEN_LIMIT:
+                        break
+                    if max_tokens is not None and tokens_generated >= max_tokens:
+                        break
+
+                    if next_token in stop_token_ids:
+                        break
+
+                inference_time = time.time() - inference_start
+                inference_tokens_per_sec = inference_tokens / inference_time if inference_time > 0 else 0
+
+                if not warmup:
+                    response = token_printer.stop(eval_mode)
+                    if eval_mode:
+                        print(response, end='')
+                    else:
+                        prefill_ms = prefill_time * 1000
+                        print(f"\nPrefill: {prefill_ms:.1f}ms ({prefill_tokens_per_sec:.1f} t/s)")
+                        print(f"Inference: {inference_tokens_per_sec:.1f} t/s")
+                        print(f"Total: Generated {tokens_generated} tokens in {prefill_time + inference_time:.2f}s")
+
+                    if save_file and not eval_mode:
+                        try:
+                            time.sleep(0.5)
+                            with open(save_file, 'w') as f:
+                                f.write(response)
+                            print(f"\n{DARK_BLUE}Response saved to file: {save_file}{RESET_COLOR}")
+                        except Exception as e:
+                            print(f"\n{DARK_BLUE}Error saving to file: {str(e)}{RESET_COLOR}")
+                else:
+                    token_printer.stop(eval_mode)
+
+                if auto_prompt is not None:
+                    break
+
+            except KeyboardInterrupt:
+                if not eval_mode:
+                    print("\nGeneration interrupted")
+                token_printer.stop(eval_mode)
+                continue
+
+    except Exception as e:
+        print(f"\nError in chat loop: {str(e)}")
+        import traceback
+        traceback.print_exc()
+
+
+def main():
+    args = parse_args()
+    
+    # Convert directory to absolute path
+    model_dir = Path(args.d).resolve()
+    if not model_dir.exists():
+        if not args.eval:
+            print(f"\nError: Model directory not found: {model_dir}")
+        return 1
+        
+    if not args.eval:
+        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)
+
+        # Check if tokenizer directory exists and has required files
+        tokenizer_path = Path(args.tokenizer)
+        if not tokenizer_path.exists():
+            if not args.eval:
+                print(f"\nError: Tokenizer directory not found: {args.tokenizer}")
+            return 1
+
+        required_files = ['tokenizer.json', 'tokenizer_config.json']
+        missing_files = [f for f in required_files if not (tokenizer_path / f).exists()]
+
+        if missing_files and not args.eval:
+            print(f"\nWarning: Tokenizer directory missing required files: {missing_files}")
+            print(f"Current tokenizer path: {args.tokenizer}")
+            print("\nFor Qwen models, you may need to specify the original model directory:")
+            print("  python chat.py --meta /tmp/qwen/meta.yaml --tokenizer ~/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/YOUR_SNAPSHOT_ID")
+
+        args.tokenizer = str(Path(args.tokenizer).resolve())
+        if not args.eval:
+            print(f"Using tokenizer path: {args.tokenizer}")
+
+        # Load tokenizer
+        tokenizer = initialize_tokenizer(args.tokenizer, args.eval)
+        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['debug_argmax'] = getattr(args, 'debug_argmax', False)
+            metadata['debug'] = getattr(args, 'debug', False)
+            metadata['sliding_window'] = 512  # Local attention window for Gemma3
+
+            if not args.eval:
+                print(f"\nMonolithic metadata: {metadata}")
+
+            # Create state from infer model
+            state = infer_model.make_state()
+            if not args.eval:
+                print("\nCreated unified transformer state for monolithic model")
+            _maybe_report_mem("after monolithic make_state", getattr(args, "mem_report", False))
+
+            # Initialize causal mask - use state_length for split cache models
+            causal_mask = initialize_causal_mask(metadata['state_length'], args.eval)
+
+            # Warmup runs
+            if not args.nw and not args.eval:
+                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?",
+                        no_template=args.no_template,
+                        eval_mode=args.eval
+                    )
+
+            # 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,
+                save_file=args.save,
+                max_tokens=args.max_tokens,
+                no_template=args.no_template,
+                eval_mode=args.eval
+            )
+
+        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)
+
+            if not args.eval:
+                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
+                if not args.eval:
+                    print(f"\nOverriding context length from command line: {args.context_length}")
+
+            # Add num_logits to metadata (legacy support)
+            metadata['num_logits'] = getattr(args, 'num_logits', 8)
+
+            # Add split_lm_head to metadata (preferred)
+            metadata['split_lm_head'] = getattr(args, 'split_lm_head', getattr(args, 'num_logits', 8))
+
+            # Add debug flag
+            metadata['debug'] = getattr(args, 'debug', False)
+
+            # Add argmax_in_model flag for chunked models
+            metadata['argmax_in_model'] = getattr(args, 'argmax_in_model', False)
+            metadata['debug_argmax'] = getattr(args, 'debug_argmax', False)
+
+            if not args.eval:
+                print(f"\nMetadata after load_models: {metadata}")
+                print(f"Using split_lm_head value: {metadata.get('split_lm_head', 8)}")
+                if metadata.get('argmax_in_model'):
+                    print("Argmax mode enabled for LM head")
+
+            # Create unified state once
+            state = create_unified_state(ffn_models, metadata['context_length'], args.eval)
+            _maybe_report_mem("after chunked make_state", getattr(args, "mem_report", False))
+
+            # Initialize causal mask once
+            # For Gemma3 with split cache, use attention_size (sliding window) for causal mask
+            attention_size = getattr(args, 'attention_size', metadata['context_length'])
+            metadata['attention_size'] = attention_size
+
+            # 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 and not args.eval:
+                print(f"Sliding window: {sliding_window} (rotation enabled for pos >= {sliding_window})")
+
+            causal_mask = initialize_causal_mask(attention_size, args.eval)
+
+            # Warmup runs to prevent Python GIL issues with CoreML
+            if not args.nw and not args.eval:
+                for _ in range(2):
+                    chat_loop(
+                        embed_model=embed_model,
+                        ffn_models=ffn_models,
+                        lmhead_model=lmhead_model,
+                        tokenizer=tokenizer,
+                        metadata=metadata,
+                        state=state,
+                        causal_mask=causal_mask,
+                        warmup=True,
+                        auto_prompt="who are you?",
+                        no_template=args.no_template,
+                        eval_mode=args.eval
+                    )
+
+            # Main run
+            chat_loop(
+                embed_model=embed_model,
+                ffn_models=ffn_models,
+                lmhead_model=lmhead_model,
+                tokenizer=tokenizer,
+                metadata=metadata,
+                state=state,
+                causal_mask=causal_mask,
+                warmup=False,
+                auto_prompt=args.prompt,
+                save_file=args.save,
+                max_tokens=args.max_tokens,
+                no_template=args.no_template,
+                eval_mode=args.eval
+            )
+        
+    except Exception as e:
+        if not args.eval:
+            print(f"\nError: {str(e)}")
+            import traceback
+            traceback.print_exc()
+        return 1
+    
+    return 0
+
+if __name__ == "__main__":
+    exit(main()) 

chat_full.py

@@ -0,0 +1,1978 @@
+# 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
+        if hasattr(tokenizer, "get_vocab"):
+            vocab = tokenizer.get_vocab()
+            if token_str in vocab:
+                return vocab[token_str]
+        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
+
+    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 (Gemma3 with context > 512)
+                    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
+                    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:
+                print(f"Sliding window: {sliding_window} (rotation enabled for pos >= {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_lut6_chunk_01of01.mlmodelc/analytics/coremldata.bin

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

gemma3_FFN_PF_lut6_chunk_01of01.mlmodelc/coremldata.bin

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

gemma3_FFN_PF_lut6_chunk_01of01.mlmodelc/metadata.json

@@ -0,0 +1,560 @@
+[
+  {
+    "metadataOutputVersion" : "3.0",
+    "userDefinedMetadata" : {
+      "com.github.apple.coremltools.source" : "torch==2.5.0",
+      "com.github.apple.coremltools.source_dialect" : "TorchScript",
+      "com.anemll.context_length" : "4096",
+      "com.github.apple.coremltools.version" : "9.0",
+      "com.anemll.lut_bits" : "6",
+      "com.anemll.num_chunks" : "1",
+      "com.anemll.info" : "Converted with Anemll v0.1.1",
+      "com.anemll.batch_size" : "64",
+      "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 × 1152)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 1152]",
+        "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 × 4096)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 1, 4096]",
+        "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 × 1152)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 1152]",
+        "name" : "output_hidden_states",
+        "type" : "MultiArray"
+      }
+    ],
+    "modelParameters" : [
+
+    ],
+    "storagePrecision" : "Mixed (Float16, Palettized (11 bits), Palettized (13 bits), Palettized (14 bits), Palettized (16 bits), UInt6)",
+    "method" : "predict",
+    "functions" : [
+      {
+        "inputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 1152)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 1152]",
+            "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 × 4096)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 1, 4096]",
+            "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 (11 bits), Palettized (13 bits), Palettized (14 bits), Palettized (16 bits), UInt6)",
+        "stateSchema" : [
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 44 × 1 × 512 × 256)",
+            "shortDescription" : "",
+            "shape" : "[44, 1, 512, 256]",
+            "name" : "model_model_kv_cache_local",
+            "type" : "State"
+          },
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 8 × 1 × 4096 × 256)",
+            "shortDescription" : "",
+            "shape" : "[8, 1, 4096, 256]",
+            "name" : "model_model_kv_cache_global",
+            "type" : "State"
+          }
+        ],
+        "outputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 1152)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 1152]",
+            "name" : "output_hidden_states",
+            "type" : "MultiArray"
+          }
+        ],
+        "name" : "infer",
+        "mlProgramOperationTypeHistogram" : {
+          "Ios18.expandDims" : 52,
+          "Ios18.softmax" : 26,
+          "Ios18.mul" : 522,
+          "Ios18.matmul" : 52,
+          "Identity" : 1,
+          "Ios18.greaterEqual" : 2,
+          "Select" : 2,
+          "Ios18.readState" : 54,
+          "Tile" : 52,
+          "Ios18.gather" : 4,
+          "Ios18.add" : 133,
+          "Ios18.layerNorm" : 157,
+          "Ios18.sliceUpdate" : 52,
+          "Ios18.writeState" : 52,
+          "Ios18.reshape" : 108,
+          "Ios18.constexprLutToDense" : 182,
+          "Ios18.conv" : 182,
+          "Ios18.concat" : 297,
+          "Ios18.transpose" : 156,
+          "Ios18.cast" : 5,
+          "Ios18.clip" : 52,
+          "Ios18.gelu" : 26,
+          "Ios18.sliceByIndex" : 314,
+          "Ios18.squeeze" : 26
+        }
+      },
+      {
+        "inputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 1152)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 1152]",
+            "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 × 4096)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 1, 4096]",
+            "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 (11 bits), Palettized (13 bits), Palettized (14 bits), Palettized (16 bits), UInt6)",
+        "stateSchema" : [
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 44 × 1 × 512 × 256)",
+            "shortDescription" : "",
+            "shape" : "[44, 1, 512, 256]",
+            "name" : "model_model_kv_cache_local",
+            "type" : "State"
+          },
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 8 × 1 × 4096 × 256)",
+            "shortDescription" : "",
+            "shape" : "[8, 1, 4096, 256]",
+            "name" : "model_model_kv_cache_global",
+            "type" : "State"
+          }
+        ],
+        "outputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 1152)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 1152]",
+            "name" : "output_hidden_states",
+            "type" : "MultiArray"
+          }
+        ],
+        "name" : "infer_rotate",
+        "mlProgramOperationTypeHistogram" : {
+          "Ios18.expandDims" : 52,
+          "Ios18.softmax" : 26,
+          "Ios18.mul" : 522,
+          "Ios18.matmul" : 52,
+          "Identity" : 1,
+          "Ios18.greaterEqual" : 2,
+          "Select" : 2,
+          "Ios18.readState" : 54,
+          "Tile" : 52,
+          "Ios18.gather" : 4,
+          "Ios18.add" : 133,
+          "Ios18.layerNorm" : 157,
+          "Ios18.sliceUpdate" : 52,
+          "Ios18.writeState" : 52,
+          "Ios18.reshape" : 108,
+          "Ios18.constexprLutToDense" : 182,
+          "Ios18.conv" : 182,
+          "Ios18.concat" : 269,
+          "Ios18.transpose" : 156,
+          "Ios18.cast" : 5,
+          "Ios18.clip" : 52,
+          "Ios18.gelu" : 26,
+          "Ios18.sliceByIndex" : 402,
+          "Ios18.squeeze" : 26
+        }
+      },
+      {
+        "inputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 64 × 1152)",
+            "shortDescription" : "",
+            "shape" : "[1, 64, 1152]",
+            "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 × 4096)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 64, 4096]",
+            "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 (11 bits), Palettized (13 bits), Palettized (14 bits), Palettized (16 bits), UInt6)",
+        "stateSchema" : [
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 44 × 1 × 512 × 256)",
+            "shortDescription" : "",
+            "shape" : "[44, 1, 512, 256]",
+            "name" : "model_model_kv_cache_local",
+            "type" : "State"
+          },
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 8 × 1 × 4096 × 256)",
+            "shortDescription" : "",
+            "shape" : "[8, 1, 4096, 256]",
+            "name" : "model_model_kv_cache_global",
+            "type" : "State"
+          }
+        ],
+        "outputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 1152)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 1152]",
+            "name" : "output_hidden_states",
+            "type" : "MultiArray"
+          }
+        ],
+        "name" : "prefill",
+        "mlProgramOperationTypeHistogram" : {
+          "Ios18.expandDims" : 52,
+          "Ios18.softmax" : 26,
+          "Ios18.mul" : 520,
+          "Ios18.matmul" : 52,
+          "Ios18.greaterEqual" : 2,
+          "Select" : 2,
+          "Ios18.readState" : 54,
+          "Tile" : 52,
+          "Ios18.gather" : 4,
+          "Ios18.add" : 133,
+          "Ios18.layerNorm" : 156,
+          "Ios18.sliceUpdate" : 52,
+          "Ios18.writeState" : 52,
+          "Ios18.reshape" : 186,
+          "Ios18.constexprLutToDense" : 182,
+          "Ios18.conv" : 182,
+          "Ios18.concat" : 296,
+          "Ios18.transpose" : 238,
+          "Ios18.cast" : 5,
+          "Ios18.clip" : 52,
+          "Ios18.gelu" : 26,
+          "Ios18.sliceByIndex" : 314,
+          "Ios18.squeeze" : 26
+        }
+      },
+      {
+        "inputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 64 × 1152)",
+            "shortDescription" : "",
+            "shape" : "[1, 64, 1152]",
+            "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 × 4096)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 64, 4096]",
+            "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 (11 bits), Palettized (13 bits), Palettized (14 bits), Palettized (16 bits), UInt6)",
+        "stateSchema" : [
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 44 × 1 × 512 × 256)",
+            "shortDescription" : "",
+            "shape" : "[44, 1, 512, 256]",
+            "name" : "model_model_kv_cache_local",
+            "type" : "State"
+          },
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 8 × 1 × 4096 × 256)",
+            "shortDescription" : "",
+            "shape" : "[8, 1, 4096, 256]",
+            "name" : "model_model_kv_cache_global",
+            "type" : "State"
+          }
+        ],
+        "outputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 1152)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 1152]",
+            "name" : "output_hidden_states",
+            "type" : "MultiArray"
+          }
+        ],
+        "name" : "prefill_rotate",
+        "mlProgramOperationTypeHistogram" : {
+          "Ios18.expandDims" : 52,
+          "Ios18.softmax" : 26,
+          "Ios18.mul" : 520,
+          "Ios18.matmul" : 52,
+          "Ios18.greaterEqual" : 2,
+          "Select" : 2,
+          "Ios18.readState" : 54,
+          "Tile" : 52,
+          "Ios18.gather" : 4,
+          "Ios18.add" : 133,
+          "Ios18.layerNorm" : 156,
+          "Ios18.sliceUpdate" : 52,
+          "Ios18.writeState" : 52,
+          "Ios18.reshape" : 186,
+          "Ios18.constexprLutToDense" : 182,
+          "Ios18.conv" : 182,
+          "Ios18.concat" : 268,
+          "Ios18.transpose" : 238,
+          "Ios18.cast" : 5,
+          "Ios18.clip" : 52,
+          "Ios18.gelu" : 26,
+          "Ios18.sliceByIndex" : 402,
+          "Ios18.squeeze" : 26
+        }
+      }
+    ],
+    "version" : "0.1.1",
+    "isUpdatable" : "0",
+    "defaultFunctionName" : "infer",
+    "specificationVersion" : 9,
+    "stateSchema" : [
+      {
+        "dataType" : "Float16",
+        "isOptional" : "0",
+        "formattedType" : "State (Float16 44 × 1 × 512 × 256)",
+        "shortDescription" : "",
+        "shape" : "[44, 1, 512, 256]",
+        "name" : "model_model_kv_cache_local",
+        "type" : "State"
+      },
+      {
+        "dataType" : "Float16",
+        "isOptional" : "0",
+        "formattedType" : "State (Float16 8 × 1 × 4096 × 256)",
+        "shortDescription" : "",
+        "shape" : "[8, 1, 4096, 256]",
+        "name" : "model_model_kv_cache_global",
+        "type" : "State"
+      }
+    ],
+    "computePrecision" : "Mixed (Float16, Int16, Int32, UInt16)",
+    "mlProgramOperationTypeHistogram" : {
+      "Ios18.expandDims" : 52,
+      "Ios18.softmax" : 26,
+      "Ios18.mul" : 522,
+      "Ios18.matmul" : 52,
+      "Identity" : 1,
+      "Ios18.greaterEqual" : 2,
+      "Select" : 2,
+      "Ios18.readState" : 54,
+      "Tile" : 52,
+      "Ios18.gather" : 4,
+      "Ios18.add" : 133,
+      "Ios18.layerNorm" : 157,
+      "Ios18.sliceUpdate" : 52,
+      "Ios18.writeState" : 52,
+      "Ios18.reshape" : 108,
+      "Ios18.constexprLutToDense" : 182,
+      "Ios18.conv" : 182,
+      "Ios18.concat" : 297,
+      "Ios18.transpose" : 156,
+      "Ios18.cast" : 5,
+      "Ios18.clip" : 52,
+      "Ios18.gelu" : 26,
+      "Ios18.sliceByIndex" : 314,
+      "Ios18.squeeze" : 26
+    },
+    "shortDescription" : "Anemll Model: Multifunction FFN+Prefill",
+    "generatedClassName" : "gemma3_FFN_PF_lut6_chunk_01of01",
+    "author" : "Converted with Anemll v0.1.1",
+    "modelType" : {
+      "name" : "MLModelType_mlProgram"
+    }
+  }
+]

gemma3_FFN_PF_lut6_chunk_01of01.mlmodelc/weights/weight.bin

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

gemma3_embeddings_lut6.mlmodelc/analytics/coremldata.bin

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

gemma3_embeddings_lut6.mlmodelc/coremldata.bin

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

gemma3_embeddings_lut6.mlmodelc/metadata.json

@@ -0,0 +1,74 @@
+[
+  {
+    "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" : "Mixed (Float16, Palettized (21 bits), UInt6)",
+    "mlProgramOperationTypeHistogram" : {
+      "Ios18.greaterEqual" : 2,
+      "Ios18.constexprLutToDense" : 1,
+      "Ios18.add" : 2,
+      "Select" : 2,
+      "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.version" : "9.0",
+      "com.github.apple.coremltools.source" : "torch==2.5.0",
+      "com.anemll.context_length" : "4096",
+      "com.github.apple.coremltools.conversion_date" : "2026-01-28",
+      "com.github.apple.coremltools.source_dialect" : "TorchScript",
+      "com.anemll.info" : "Converted with Anemll v0.1.1",
+      "com.anemll.lut_bits" : "6"
+    },
+    "generatedClassName" : "gemma3_embeddings_lut6",
+    "method" : "predict"
+  }
+]

gemma3_embeddings_lut6.mlmodelc/model.mil

@@ -0,0 +1,23 @@
+program(1.3)
+[buildInfo = dict<string, string>({{"coremlc-component-MIL", "3510.2.1"}, {"coremlc-version", "3500.32.1"}})]
+{
+    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, [262144, 1152]> embed_tokens_weight_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [262144, 1152]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(64))), lut = tensor<fp16, [32768, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(226492544))))[name = string("embed_tokens_weight_palettized")];
+            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(262144)];
+            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 greater_equal_0_y_0_1 = const()[name = string("greater_equal_0_y_0_1"), val = int32(0)];
+            tensor<bool, [1, ?]> greater_equal_0_1 = greater_equal(x = select_0, y = greater_equal_0_y_0_1)[name = string("greater_equal_0_1")];
+            int32 slice_by_index_0_1 = const()[name = string("slice_by_index_0_1"), val = int32(262144)];
+            tensor<int32, [1, ?]> add_0_1 = add(x = select_0, y = slice_by_index_0_1)[name = string("add_0_1")];
+            tensor<int32, [1, ?]> select_0_1 = select(a = select_0, b = add_0_1, cond = greater_equal_0_1)[name = string("select_0_1")];
+            int32 hidden_states_1_axis_0 = const()[name = string("hidden_states_1_axis_0"), val = int32(0)];
+            tensor<fp16, [1, ?, 1152]> hidden_states_1 = gather(axis = hidden_states_1_axis_0, batch_dims = hidden_states_1_batch_dims_0, indices = select_0_1, validate_indices = hidden_states_1_validate_indices_0, x = embed_tokens_weight_palettized)[name = string("hidden_states_1")];
+            fp16 var_7_to_fp16 = const()[name = string("op_7_to_fp16"), val = fp16(0x1.0f8p+5)];
+            tensor<fp16, [1, ?, 1152]> hidden_states = mul(x = hidden_states_1, y = var_7_to_fp16)[name = string("hidden_states_cast_fp16")];
+        } -> (hidden_states);
+}

gemma3_embeddings_lut6.mlmodelc/weights/weight.bin

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

gemma3_lm_head_lut6.mlmodelc/analytics/coremldata.bin

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

gemma3_lm_head_lut6.mlmodelc/coremldata.bin

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

gemma3_lm_head_lut6.mlmodelc/metadata.json

@@ -0,0 +1,85 @@
+[
+  {
+    "shortDescription" : "Anemll Model (LM Head) converted to CoreML",
+    "metadataOutputVersion" : "3.0",
+    "outputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Int32",
+        "formattedType" : "MultiArray (Int32 16)",
+        "shortDescription" : "",
+        "shape" : "[16]",
+        "name" : "argmax_idx",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 16)",
+        "shortDescription" : "",
+        "shape" : "[16]",
+        "name" : "argmax_val",
+        "type" : "MultiArray"
+      }
+    ],
+    "version" : "0.1.1",
+    "modelParameters" : [
+
+    ],
+    "author" : "Converted with Anemll v0.1.1",
+    "specificationVersion" : 9,
+    "storagePrecision" : "Mixed (Float16, Palettized (17 bits), UInt6)",
+    "mlProgramOperationTypeHistogram" : {
+      "Ios18.squeeze" : 20,
+      "Ios18.gatherAlongAxis" : 16,
+      "Ios18.reduceArgmax" : 16,
+      "Ios18.concat" : 2,
+      "Ios18.transpose" : 17,
+      "Ios18.constexprLutToDense" : 16,
+      "Ios18.expandDims" : 1,
+      "Ios18.conv" : 16,
+      "Ios18.cast" : 18
+    },
+    "computePrecision" : "Mixed (Float16, Int16, Int32, UInt16)",
+    "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 × 1152)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 1152]",
+        "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-28",
+      "com.github.apple.coremltools.source" : "torch==2.5.0",
+      "com.anemll.context_length" : "4096",
+      "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,348 @@
+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, 1152]> 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, 1152, 1]> var_6_cast_fp16 = transpose(perm = var_5, x = hidden_states)[name = string("transpose_16")];
+            tensor<fp16, [1, 1152, 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, [16384, 1152, 1, 1]> op_9_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16384, 1152, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(64))), lut = tensor<fp16, [2048, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(14155904))))[name = string("op_9_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16384, 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, 16384, 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]> logits_1_perm_0 = const()[name = string("logits_1_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, [16384, 1152, 1, 1]> op_35_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16384, 1152, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(14418112))), lut = tensor<fp16, [2048, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(28573952))))[name = string("op_35_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16384, 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, 16384, 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]> logits_3_perm_0 = const()[name = string("logits_3_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, [16384, 1152, 1, 1]> op_61_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16384, 1152, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(28836160))), lut = tensor<fp16, [2048, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(42992000))))[name = string("op_61_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16384, 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, 16384, 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]> logits_5_perm_0 = const()[name = string("logits_5_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, [16384, 1152, 1, 1]> op_87_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16384, 1152, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(43254208))), lut = tensor<fp16, [2048, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(57410048))))[name = string("op_87_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16384, 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, 16384, 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]> logits_7_perm_0 = const()[name = string("logits_7_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, [16384, 1152, 1, 1]> op_113_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16384, 1152, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(57672256))), lut = tensor<fp16, [2048, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(71828096))))[name = string("op_113_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16384, 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, 16384, 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]> logits_9_perm_0 = const()[name = string("logits_9_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, [16384, 1152, 1, 1]> op_139_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16384, 1152, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(72090304))), lut = tensor<fp16, [2048, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(86246144))))[name = string("op_139_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16384, 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, 16384, 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]> logits_11_perm_0 = const()[name = string("logits_11_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, [16384, 1152, 1, 1]> op_165_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16384, 1152, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(86508352))), lut = tensor<fp16, [2048, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(100664192))))[name = string("op_165_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16384, 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, 16384, 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]> logits_13_perm_0 = const()[name = string("logits_13_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, [16384, 1152, 1, 1]> op_191_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16384, 1152, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(100926400))), lut = tensor<fp16, [2048, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(115082240))))[name = string("op_191_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16384, 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, 16384, 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]> logits_15_perm_0 = const()[name = string("logits_15_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, [16384, 1152, 1, 1]> op_217_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16384, 1152, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(115344448))), lut = tensor<fp16, [2048, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(129500288))))[name = string("op_217_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16384, 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, 16384, 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]> logits_17_perm_0 = const()[name = string("logits_17_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, [16384, 1152, 1, 1]> op_243_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16384, 1152, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(129762496))), lut = tensor<fp16, [2048, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(143918336))))[name = string("op_243_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16384, 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, 16384, 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]> logits_19_perm_0 = const()[name = string("logits_19_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, [16384, 1152, 1, 1]> op_269_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16384, 1152, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(144180544))), lut = tensor<fp16, [2048, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(158336384))))[name = string("op_269_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16384, 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, 16384, 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]> logits_21_perm_0 = const()[name = string("logits_21_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, [16384, 1152, 1, 1]> op_295_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16384, 1152, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(158598592))), lut = tensor<fp16, [2048, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(172754432))))[name = string("op_295_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16384, 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, 16384, 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]> logits_23_perm_0 = const()[name = string("logits_23_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, [16384, 1152, 1, 1]> op_321_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16384, 1152, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(173016640))), lut = tensor<fp16, [2048, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(187172480))))[name = string("op_321_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16384, 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, 16384, 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]> logits_25_perm_0 = const()[name = string("logits_25_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, [16384, 1152, 1, 1]> op_347_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16384, 1152, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(187434688))), lut = tensor<fp16, [2048, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(201590528))))[name = string("op_347_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16384, 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, 16384, 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]> logits_27_perm_0 = const()[name = string("logits_27_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, [16384, 1152, 1, 1]> op_373_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16384, 1152, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(201852736))), lut = tensor<fp16, [2048, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(216008576))))[name = string("op_373_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16384, 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, 16384, 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]> logits_29_perm_0 = const()[name = string("logits_29_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, [16384, 1152, 1, 1]> op_399_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint6, [16384, 1152, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(216270784))), lut = tensor<fp16, [2048, 1, 1, 1, 64, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(230426624))))[name = string("op_399_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 16384, 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, 16384, 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]> logits_perm_0 = const()[name = string("logits_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            int32 chunk_argmax_1_axis_0 = const()[name = string("chunk_argmax_1_axis_0"), val = int32(-1)];
+            bool chunk_argmax_1_keep_dims_0 = const()[name = string("chunk_argmax_1_keep_dims_0"), val = bool(true)];
+            string chunk_argmax_1_output_dtype_0 = const()[name = string("chunk_argmax_1_output_dtype_0"), val = string("int32")];
+            tensor<fp16, [1, 1, 16384]> logits_1_cast_fp16 = transpose(perm = logits_1_perm_0, x = var_31_cast_fp16)[name = string("transpose_15")];
+            tensor<int32, [1, 1, 1]> chunk_argmax_1_cast_fp16 = reduce_argmax(axis = chunk_argmax_1_axis_0, keep_dims = chunk_argmax_1_keep_dims_0, output_dtype = chunk_argmax_1_output_dtype_0, x = logits_1_cast_fp16)[name = string("chunk_argmax_1_cast_fp16")];
+            int32 var_428 = const()[name = string("op_428"), val = int32(-1)];
+            bool var_430_validate_indices_0 = const()[name = string("op_430_validate_indices_0"), val = bool(false)];
+            string chunk_argmax_1_cast_fp16_to_uint16_dtype_0 = const()[name = string("chunk_argmax_1_cast_fp16_to_uint16_dtype_0"), val = string("uint16")];
+            tensor<uint16, [1, 1, 1]> chunk_argmax_1_cast_fp16_to_uint16 = cast(dtype = chunk_argmax_1_cast_fp16_to_uint16_dtype_0, x = chunk_argmax_1_cast_fp16)[name = string("cast_17")];
+            tensor<fp16, [1, 1, 1]> var_430_cast_fp16_cast_int16 = gather_along_axis(axis = var_428, indices = chunk_argmax_1_cast_fp16_to_uint16, validate_indices = var_430_validate_indices_0, x = logits_1_cast_fp16)[name = string("op_430_cast_fp16_cast_int16")];
+            int32 chunk_argmax_3_axis_0 = const()[name = string("chunk_argmax_3_axis_0"), val = int32(-1)];
+            bool chunk_argmax_3_keep_dims_0 = const()[name = string("chunk_argmax_3_keep_dims_0"), val = bool(true)];
+            string chunk_argmax_3_output_dtype_0 = const()[name = string("chunk_argmax_3_output_dtype_0"), val = string("int32")];
+            tensor<fp16, [1, 1, 16384]> logits_3_cast_fp16 = transpose(perm = logits_3_perm_0, x = var_57_cast_fp16)[name = string("transpose_14")];
+            tensor<int32, [1, 1, 1]> chunk_argmax_3_cast_fp16 = reduce_argmax(axis = chunk_argmax_3_axis_0, keep_dims = chunk_argmax_3_keep_dims_0, output_dtype = chunk_argmax_3_output_dtype_0, x = logits_3_cast_fp16)[name = string("chunk_argmax_3_cast_fp16")];
+            int32 var_439 = const()[name = string("op_439"), val = int32(-1)];
+            bool var_441_validate_indices_0 = const()[name = string("op_441_validate_indices_0"), val = bool(false)];
+            string chunk_argmax_3_cast_fp16_to_int16_dtype_0 = const()[name = string("chunk_argmax_3_cast_fp16_to_int16_dtype_0"), val = string("int16")];
+            tensor<int16, [1, 1, 1]> chunk_argmax_3_cast_fp16_to_int16 = cast(dtype = chunk_argmax_3_cast_fp16_to_int16_dtype_0, x = chunk_argmax_3_cast_fp16)[name = string("cast_16")];
+            tensor<fp16, [1, 1, 1]> var_441_cast_fp16_cast_int16 = gather_along_axis(axis = var_439, indices = chunk_argmax_3_cast_fp16_to_int16, validate_indices = var_441_validate_indices_0, x = logits_3_cast_fp16)[name = string("op_441_cast_fp16_cast_int16")];
+            int32 chunk_argmax_5_axis_0 = const()[name = string("chunk_argmax_5_axis_0"), val = int32(-1)];
+            bool chunk_argmax_5_keep_dims_0 = const()[name = string("chunk_argmax_5_keep_dims_0"), val = bool(true)];
+            string chunk_argmax_5_output_dtype_0 = const()[name = string("chunk_argmax_5_output_dtype_0"), val = string("int32")];
+            tensor<fp16, [1, 1, 16384]> logits_5_cast_fp16 = transpose(perm = logits_5_perm_0, x = var_83_cast_fp16)[name = string("transpose_13")];
+            tensor<int32, [1, 1, 1]> chunk_argmax_5_cast_fp16 = reduce_argmax(axis = chunk_argmax_5_axis_0, keep_dims = chunk_argmax_5_keep_dims_0, output_dtype = chunk_argmax_5_output_dtype_0, x = logits_5_cast_fp16)[name = string("chunk_argmax_5_cast_fp16")];
+            int32 var_450 = const()[name = string("op_450"), val = int32(-1)];
+            bool var_452_validate_indices_0 = const()[name = string("op_452_validate_indices_0"), val = bool(false)];
+            string chunk_argmax_5_cast_fp16_to_int16_dtype_0 = const()[name = string("chunk_argmax_5_cast_fp16_to_int16_dtype_0"), val = string("int16")];
+            tensor<int16, [1, 1, 1]> chunk_argmax_5_cast_fp16_to_int16 = cast(dtype = chunk_argmax_5_cast_fp16_to_int16_dtype_0, x = chunk_argmax_5_cast_fp16)[name = string("cast_15")];
+            tensor<fp16, [1, 1, 1]> var_452_cast_fp16_cast_int16 = gather_along_axis(axis = var_450, indices = chunk_argmax_5_cast_fp16_to_int16, validate_indices = var_452_validate_indices_0, x = logits_5_cast_fp16)[name = string("op_452_cast_fp16_cast_int16")];
+            int32 chunk_argmax_7_axis_0 = const()[name = string("chunk_argmax_7_axis_0"), val = int32(-1)];
+            bool chunk_argmax_7_keep_dims_0 = const()[name = string("chunk_argmax_7_keep_dims_0"), val = bool(true)];
+            string chunk_argmax_7_output_dtype_0 = const()[name = string("chunk_argmax_7_output_dtype_0"), val = string("int32")];
+            tensor<fp16, [1, 1, 16384]> logits_7_cast_fp16 = transpose(perm = logits_7_perm_0, x = var_109_cast_fp16)[name = string("transpose_12")];
+            tensor<int32, [1, 1, 1]> chunk_argmax_7_cast_fp16 = reduce_argmax(axis = chunk_argmax_7_axis_0, keep_dims = chunk_argmax_7_keep_dims_0, output_dtype = chunk_argmax_7_output_dtype_0, x = logits_7_cast_fp16)[name = string("chunk_argmax_7_cast_fp16")];
+            int32 var_461 = const()[name = string("op_461"), val = int32(-1)];
+            bool var_463_validate_indices_0 = const()[name = string("op_463_validate_indices_0"), val = bool(false)];
+            string chunk_argmax_7_cast_fp16_to_int16_dtype_0 = const()[name = string("chunk_argmax_7_cast_fp16_to_int16_dtype_0"), val = string("int16")];
+            tensor<int16, [1, 1, 1]> chunk_argmax_7_cast_fp16_to_int16 = cast(dtype = chunk_argmax_7_cast_fp16_to_int16_dtype_0, x = chunk_argmax_7_cast_fp16)[name = string("cast_14")];
+            tensor<fp16, [1, 1, 1]> var_463_cast_fp16_cast_int16 = gather_along_axis(axis = var_461, indices = chunk_argmax_7_cast_fp16_to_int16, validate_indices = var_463_validate_indices_0, x = logits_7_cast_fp16)[name = string("op_463_cast_fp16_cast_int16")];
+            int32 chunk_argmax_9_axis_0 = const()[name = string("chunk_argmax_9_axis_0"), val = int32(-1)];
+            bool chunk_argmax_9_keep_dims_0 = const()[name = string("chunk_argmax_9_keep_dims_0"), val = bool(true)];
+            string chunk_argmax_9_output_dtype_0 = const()[name = string("chunk_argmax_9_output_dtype_0"), val = string("int32")];
+            tensor<fp16, [1, 1, 16384]> logits_9_cast_fp16 = transpose(perm = logits_9_perm_0, x = var_135_cast_fp16)[name = string("transpose_11")];
+            tensor<int32, [1, 1, 1]> chunk_argmax_9_cast_fp16 = reduce_argmax(axis = chunk_argmax_9_axis_0, keep_dims = chunk_argmax_9_keep_dims_0, output_dtype = chunk_argmax_9_output_dtype_0, x = logits_9_cast_fp16)[name = string("chunk_argmax_9_cast_fp16")];
+            int32 var_472 = const()[name = string("op_472"), val = int32(-1)];
+            bool var_474_validate_indices_0 = const()[name = string("op_474_validate_indices_0"), val = bool(false)];
+            string chunk_argmax_9_cast_fp16_to_int16_dtype_0 = const()[name = string("chunk_argmax_9_cast_fp16_to_int16_dtype_0"), val = string("int16")];
+            tensor<int16, [1, 1, 1]> chunk_argmax_9_cast_fp16_to_int16 = cast(dtype = chunk_argmax_9_cast_fp16_to_int16_dtype_0, x = chunk_argmax_9_cast_fp16)[name = string("cast_13")];
+            tensor<fp16, [1, 1, 1]> var_474_cast_fp16_cast_int16 = gather_along_axis(axis = var_472, indices = chunk_argmax_9_cast_fp16_to_int16, validate_indices = var_474_validate_indices_0, x = logits_9_cast_fp16)[name = string("op_474_cast_fp16_cast_int16")];
+            int32 chunk_argmax_11_axis_0 = const()[name = string("chunk_argmax_11_axis_0"), val = int32(-1)];
+            bool chunk_argmax_11_keep_dims_0 = const()[name = string("chunk_argmax_11_keep_dims_0"), val = bool(true)];
+            string chunk_argmax_11_output_dtype_0 = const()[name = string("chunk_argmax_11_output_dtype_0"), val = string("int32")];
+            tensor<fp16, [1, 1, 16384]> logits_11_cast_fp16 = transpose(perm = logits_11_perm_0, x = var_161_cast_fp16)[name = string("transpose_10")];
+            tensor<int32, [1, 1, 1]> chunk_argmax_11_cast_fp16 = reduce_argmax(axis = chunk_argmax_11_axis_0, keep_dims = chunk_argmax_11_keep_dims_0, output_dtype = chunk_argmax_11_output_dtype_0, x = logits_11_cast_fp16)[name = string("chunk_argmax_11_cast_fp16")];
+            int32 var_483 = const()[name = string("op_483"), val = int32(-1)];
+            bool var_485_validate_indices_0 = const()[name = string("op_485_validate_indices_0"), val = bool(false)];
+            string chunk_argmax_11_cast_fp16_to_int16_dtype_0 = const()[name = string("chunk_argmax_11_cast_fp16_to_int16_dtype_0"), val = string("int16")];
+            tensor<int16, [1, 1, 1]> chunk_argmax_11_cast_fp16_to_int16 = cast(dtype = chunk_argmax_11_cast_fp16_to_int16_dtype_0, x = chunk_argmax_11_cast_fp16)[name = string("cast_12")];
+            tensor<fp16, [1, 1, 1]> var_485_cast_fp16_cast_int16 = gather_along_axis(axis = var_483, indices = chunk_argmax_11_cast_fp16_to_int16, validate_indices = var_485_validate_indices_0, x = logits_11_cast_fp16)[name = string("op_485_cast_fp16_cast_int16")];
+            int32 chunk_argmax_13_axis_0 = const()[name = string("chunk_argmax_13_axis_0"), val = int32(-1)];
+            bool chunk_argmax_13_keep_dims_0 = const()[name = string("chunk_argmax_13_keep_dims_0"), val = bool(true)];
+            string chunk_argmax_13_output_dtype_0 = const()[name = string("chunk_argmax_13_output_dtype_0"), val = string("int32")];
+            tensor<fp16, [1, 1, 16384]> logits_13_cast_fp16 = transpose(perm = logits_13_perm_0, x = var_187_cast_fp16)[name = string("transpose_9")];
+            tensor<int32, [1, 1, 1]> chunk_argmax_13_cast_fp16 = reduce_argmax(axis = chunk_argmax_13_axis_0, keep_dims = chunk_argmax_13_keep_dims_0, output_dtype = chunk_argmax_13_output_dtype_0, x = logits_13_cast_fp16)[name = string("chunk_argmax_13_cast_fp16")];
+            int32 var_494 = const()[name = string("op_494"), val = int32(-1)];
+            bool var_496_validate_indices_0 = const()[name = string("op_496_validate_indices_0"), val = bool(false)];
+            string chunk_argmax_13_cast_fp16_to_int16_dtype_0 = const()[name = string("chunk_argmax_13_cast_fp16_to_int16_dtype_0"), val = string("int16")];
+            tensor<int16, [1, 1, 1]> chunk_argmax_13_cast_fp16_to_int16 = cast(dtype = chunk_argmax_13_cast_fp16_to_int16_dtype_0, x = chunk_argmax_13_cast_fp16)[name = string("cast_11")];
+            tensor<fp16, [1, 1, 1]> var_496_cast_fp16_cast_int16 = gather_along_axis(axis = var_494, indices = chunk_argmax_13_cast_fp16_to_int16, validate_indices = var_496_validate_indices_0, x = logits_13_cast_fp16)[name = string("op_496_cast_fp16_cast_int16")];
+            int32 chunk_argmax_15_axis_0 = const()[name = string("chunk_argmax_15_axis_0"), val = int32(-1)];
+            bool chunk_argmax_15_keep_dims_0 = const()[name = string("chunk_argmax_15_keep_dims_0"), val = bool(true)];
+            string chunk_argmax_15_output_dtype_0 = const()[name = string("chunk_argmax_15_output_dtype_0"), val = string("int32")];
+            tensor<fp16, [1, 1, 16384]> logits_15_cast_fp16 = transpose(perm = logits_15_perm_0, x = var_213_cast_fp16)[name = string("transpose_8")];
+            tensor<int32, [1, 1, 1]> chunk_argmax_15_cast_fp16 = reduce_argmax(axis = chunk_argmax_15_axis_0, keep_dims = chunk_argmax_15_keep_dims_0, output_dtype = chunk_argmax_15_output_dtype_0, x = logits_15_cast_fp16)[name = string("chunk_argmax_15_cast_fp16")];
+            int32 var_505 = const()[name = string("op_505"), val = int32(-1)];
+            bool var_507_validate_indices_0 = const()[name = string("op_507_validate_indices_0"), val = bool(false)];
+            string chunk_argmax_15_cast_fp16_to_int16_dtype_0 = const()[name = string("chunk_argmax_15_cast_fp16_to_int16_dtype_0"), val = string("int16")];
+            tensor<int16, [1, 1, 1]> chunk_argmax_15_cast_fp16_to_int16 = cast(dtype = chunk_argmax_15_cast_fp16_to_int16_dtype_0, x = chunk_argmax_15_cast_fp16)[name = string("cast_10")];
+            tensor<fp16, [1, 1, 1]> var_507_cast_fp16_cast_int16 = gather_along_axis(axis = var_505, indices = chunk_argmax_15_cast_fp16_to_int16, validate_indices = var_507_validate_indices_0, x = logits_15_cast_fp16)[name = string("op_507_cast_fp16_cast_int16")];
+            int32 chunk_argmax_17_axis_0 = const()[name = string("chunk_argmax_17_axis_0"), val = int32(-1)];
+            bool chunk_argmax_17_keep_dims_0 = const()[name = string("chunk_argmax_17_keep_dims_0"), val = bool(true)];
+            string chunk_argmax_17_output_dtype_0 = const()[name = string("chunk_argmax_17_output_dtype_0"), val = string("int32")];
+            tensor<fp16, [1, 1, 16384]> logits_17_cast_fp16 = transpose(perm = logits_17_perm_0, x = var_239_cast_fp16)[name = string("transpose_7")];
+            tensor<int32, [1, 1, 1]> chunk_argmax_17_cast_fp16 = reduce_argmax(axis = chunk_argmax_17_axis_0, keep_dims = chunk_argmax_17_keep_dims_0, output_dtype = chunk_argmax_17_output_dtype_0, x = logits_17_cast_fp16)[name = string("chunk_argmax_17_cast_fp16")];
+            int32 var_516 = const()[name = string("op_516"), val = int32(-1)];
+            bool var_518_validate_indices_0 = const()[name = string("op_518_validate_indices_0"), val = bool(false)];
+            string chunk_argmax_17_cast_fp16_to_int16_dtype_0 = const()[name = string("chunk_argmax_17_cast_fp16_to_int16_dtype_0"), val = string("int16")];
+            tensor<int16, [1, 1, 1]> chunk_argmax_17_cast_fp16_to_int16 = cast(dtype = chunk_argmax_17_cast_fp16_to_int16_dtype_0, x = chunk_argmax_17_cast_fp16)[name = string("cast_9")];
+            tensor<fp16, [1, 1, 1]> var_518_cast_fp16_cast_int16 = gather_along_axis(axis = var_516, indices = chunk_argmax_17_cast_fp16_to_int16, validate_indices = var_518_validate_indices_0, x = logits_17_cast_fp16)[name = string("op_518_cast_fp16_cast_int16")];
+            int32 chunk_argmax_19_axis_0 = const()[name = string("chunk_argmax_19_axis_0"), val = int32(-1)];
+            bool chunk_argmax_19_keep_dims_0 = const()[name = string("chunk_argmax_19_keep_dims_0"), val = bool(true)];
+            string chunk_argmax_19_output_dtype_0 = const()[name = string("chunk_argmax_19_output_dtype_0"), val = string("int32")];
+            tensor<fp16, [1, 1, 16384]> logits_19_cast_fp16 = transpose(perm = logits_19_perm_0, x = var_265_cast_fp16)[name = string("transpose_6")];
+            tensor<int32, [1, 1, 1]> chunk_argmax_19_cast_fp16 = reduce_argmax(axis = chunk_argmax_19_axis_0, keep_dims = chunk_argmax_19_keep_dims_0, output_dtype = chunk_argmax_19_output_dtype_0, x = logits_19_cast_fp16)[name = string("chunk_argmax_19_cast_fp16")];
+            int32 var_527 = const()[name = string("op_527"), val = int32(-1)];
+            bool var_529_validate_indices_0 = const()[name = string("op_529_validate_indices_0"), val = bool(false)];
+            string chunk_argmax_19_cast_fp16_to_int16_dtype_0 = const()[name = string("chunk_argmax_19_cast_fp16_to_int16_dtype_0"), val = string("int16")];
+            tensor<int16, [1, 1, 1]> chunk_argmax_19_cast_fp16_to_int16 = cast(dtype = chunk_argmax_19_cast_fp16_to_int16_dtype_0, x = chunk_argmax_19_cast_fp16)[name = string("cast_8")];
+            tensor<fp16, [1, 1, 1]> var_529_cast_fp16_cast_int16 = gather_along_axis(axis = var_527, indices = chunk_argmax_19_cast_fp16_to_int16, validate_indices = var_529_validate_indices_0, x = logits_19_cast_fp16)[name = string("op_529_cast_fp16_cast_int16")];
+            int32 chunk_argmax_21_axis_0 = const()[name = string("chunk_argmax_21_axis_0"), val = int32(-1)];
+            bool chunk_argmax_21_keep_dims_0 = const()[name = string("chunk_argmax_21_keep_dims_0"), val = bool(true)];
+            string chunk_argmax_21_output_dtype_0 = const()[name = string("chunk_argmax_21_output_dtype_0"), val = string("int32")];
+            tensor<fp16, [1, 1, 16384]> logits_21_cast_fp16 = transpose(perm = logits_21_perm_0, x = var_291_cast_fp16)[name = string("transpose_5")];
+            tensor<int32, [1, 1, 1]> chunk_argmax_21_cast_fp16 = reduce_argmax(axis = chunk_argmax_21_axis_0, keep_dims = chunk_argmax_21_keep_dims_0, output_dtype = chunk_argmax_21_output_dtype_0, x = logits_21_cast_fp16)[name = string("chunk_argmax_21_cast_fp16")];
+            int32 var_538 = const()[name = string("op_538"), val = int32(-1)];
+            bool var_540_validate_indices_0 = const()[name = string("op_540_validate_indices_0"), val = bool(false)];
+            string chunk_argmax_21_cast_fp16_to_int16_dtype_0 = const()[name = string("chunk_argmax_21_cast_fp16_to_int16_dtype_0"), val = string("int16")];
+            tensor<int16, [1, 1, 1]> chunk_argmax_21_cast_fp16_to_int16 = cast(dtype = chunk_argmax_21_cast_fp16_to_int16_dtype_0, x = chunk_argmax_21_cast_fp16)[name = string("cast_7")];
+            tensor<fp16, [1, 1, 1]> var_540_cast_fp16_cast_int16 = gather_along_axis(axis = var_538, indices = chunk_argmax_21_cast_fp16_to_int16, validate_indices = var_540_validate_indices_0, x = logits_21_cast_fp16)[name = string("op_540_cast_fp16_cast_int16")];
+            int32 chunk_argmax_23_axis_0 = const()[name = string("chunk_argmax_23_axis_0"), val = int32(-1)];
+            bool chunk_argmax_23_keep_dims_0 = const()[name = string("chunk_argmax_23_keep_dims_0"), val = bool(true)];
+            string chunk_argmax_23_output_dtype_0 = const()[name = string("chunk_argmax_23_output_dtype_0"), val = string("int32")];
+            tensor<fp16, [1, 1, 16384]> logits_23_cast_fp16 = transpose(perm = logits_23_perm_0, x = var_317_cast_fp16)[name = string("transpose_4")];
+            tensor<int32, [1, 1, 1]> chunk_argmax_23_cast_fp16 = reduce_argmax(axis = chunk_argmax_23_axis_0, keep_dims = chunk_argmax_23_keep_dims_0, output_dtype = chunk_argmax_23_output_dtype_0, x = logits_23_cast_fp16)[name = string("chunk_argmax_23_cast_fp16")];
+            int32 var_549 = const()[name = string("op_549"), val = int32(-1)];
+            bool var_551_validate_indices_0 = const()[name = string("op_551_validate_indices_0"), val = bool(false)];
+            string chunk_argmax_23_cast_fp16_to_int16_dtype_0 = const()[name = string("chunk_argmax_23_cast_fp16_to_int16_dtype_0"), val = string("int16")];
+            tensor<int16, [1, 1, 1]> chunk_argmax_23_cast_fp16_to_int16 = cast(dtype = chunk_argmax_23_cast_fp16_to_int16_dtype_0, x = chunk_argmax_23_cast_fp16)[name = string("cast_6")];
+            tensor<fp16, [1, 1, 1]> var_551_cast_fp16_cast_int16 = gather_along_axis(axis = var_549, indices = chunk_argmax_23_cast_fp16_to_int16, validate_indices = var_551_validate_indices_0, x = logits_23_cast_fp16)[name = string("op_551_cast_fp16_cast_int16")];
+            int32 chunk_argmax_25_axis_0 = const()[name = string("chunk_argmax_25_axis_0"), val = int32(-1)];
+            bool chunk_argmax_25_keep_dims_0 = const()[name = string("chunk_argmax_25_keep_dims_0"), val = bool(true)];
+            string chunk_argmax_25_output_dtype_0 = const()[name = string("chunk_argmax_25_output_dtype_0"), val = string("int32")];
+            tensor<fp16, [1, 1, 16384]> logits_25_cast_fp16 = transpose(perm = logits_25_perm_0, x = var_343_cast_fp16)[name = string("transpose_3")];
+            tensor<int32, [1, 1, 1]> chunk_argmax_25_cast_fp16 = reduce_argmax(axis = chunk_argmax_25_axis_0, keep_dims = chunk_argmax_25_keep_dims_0, output_dtype = chunk_argmax_25_output_dtype_0, x = logits_25_cast_fp16)[name = string("chunk_argmax_25_cast_fp16")];
+            int32 var_560 = const()[name = string("op_560"), val = int32(-1)];
+            bool var_562_validate_indices_0 = const()[name = string("op_562_validate_indices_0"), val = bool(false)];
+            string chunk_argmax_25_cast_fp16_to_int16_dtype_0 = const()[name = string("chunk_argmax_25_cast_fp16_to_int16_dtype_0"), val = string("int16")];
+            tensor<int16, [1, 1, 1]> chunk_argmax_25_cast_fp16_to_int16 = cast(dtype = chunk_argmax_25_cast_fp16_to_int16_dtype_0, x = chunk_argmax_25_cast_fp16)[name = string("cast_5")];
+            tensor<fp16, [1, 1, 1]> var_562_cast_fp16_cast_int16 = gather_along_axis(axis = var_560, indices = chunk_argmax_25_cast_fp16_to_int16, validate_indices = var_562_validate_indices_0, x = logits_25_cast_fp16)[name = string("op_562_cast_fp16_cast_int16")];
+            int32 chunk_argmax_27_axis_0 = const()[name = string("chunk_argmax_27_axis_0"), val = int32(-1)];
+            bool chunk_argmax_27_keep_dims_0 = const()[name = string("chunk_argmax_27_keep_dims_0"), val = bool(true)];
+            string chunk_argmax_27_output_dtype_0 = const()[name = string("chunk_argmax_27_output_dtype_0"), val = string("int32")];
+            tensor<fp16, [1, 1, 16384]> logits_27_cast_fp16 = transpose(perm = logits_27_perm_0, x = var_369_cast_fp16)[name = string("transpose_2")];
+            tensor<int32, [1, 1, 1]> chunk_argmax_27_cast_fp16 = reduce_argmax(axis = chunk_argmax_27_axis_0, keep_dims = chunk_argmax_27_keep_dims_0, output_dtype = chunk_argmax_27_output_dtype_0, x = logits_27_cast_fp16)[name = string("chunk_argmax_27_cast_fp16")];
+            int32 var_571 = const()[name = string("op_571"), val = int32(-1)];
+            bool var_573_validate_indices_0 = const()[name = string("op_573_validate_indices_0"), val = bool(false)];
+            string chunk_argmax_27_cast_fp16_to_int16_dtype_0 = const()[name = string("chunk_argmax_27_cast_fp16_to_int16_dtype_0"), val = string("int16")];
+            tensor<int16, [1, 1, 1]> chunk_argmax_27_cast_fp16_to_int16 = cast(dtype = chunk_argmax_27_cast_fp16_to_int16_dtype_0, x = chunk_argmax_27_cast_fp16)[name = string("cast_4")];
+            tensor<fp16, [1, 1, 1]> var_573_cast_fp16_cast_int16 = gather_along_axis(axis = var_571, indices = chunk_argmax_27_cast_fp16_to_int16, validate_indices = var_573_validate_indices_0, x = logits_27_cast_fp16)[name = string("op_573_cast_fp16_cast_int16")];
+            int32 chunk_argmax_29_axis_0 = const()[name = string("chunk_argmax_29_axis_0"), val = int32(-1)];
+            bool chunk_argmax_29_keep_dims_0 = const()[name = string("chunk_argmax_29_keep_dims_0"), val = bool(true)];
+            string chunk_argmax_29_output_dtype_0 = const()[name = string("chunk_argmax_29_output_dtype_0"), val = string("int32")];
+            tensor<fp16, [1, 1, 16384]> logits_29_cast_fp16 = transpose(perm = logits_29_perm_0, x = var_395_cast_fp16)[name = string("transpose_1")];
+            tensor<int32, [1, 1, 1]> chunk_argmax_29_cast_fp16 = reduce_argmax(axis = chunk_argmax_29_axis_0, keep_dims = chunk_argmax_29_keep_dims_0, output_dtype = chunk_argmax_29_output_dtype_0, x = logits_29_cast_fp16)[name = string("chunk_argmax_29_cast_fp16")];
+            int32 var_582 = const()[name = string("op_582"), val = int32(-1)];
+            bool var_584_validate_indices_0 = const()[name = string("op_584_validate_indices_0"), val = bool(false)];
+            string chunk_argmax_29_cast_fp16_to_int16_dtype_0 = const()[name = string("chunk_argmax_29_cast_fp16_to_int16_dtype_0"), val = string("int16")];
+            tensor<int16, [1, 1, 1]> chunk_argmax_29_cast_fp16_to_int16 = cast(dtype = chunk_argmax_29_cast_fp16_to_int16_dtype_0, x = chunk_argmax_29_cast_fp16)[name = string("cast_3")];
+            tensor<fp16, [1, 1, 1]> var_584_cast_fp16_cast_int16 = gather_along_axis(axis = var_582, indices = chunk_argmax_29_cast_fp16_to_int16, validate_indices = var_584_validate_indices_0, x = logits_29_cast_fp16)[name = string("op_584_cast_fp16_cast_int16")];
+            int32 chunk_argmax_axis_0 = const()[name = string("chunk_argmax_axis_0"), val = int32(-1)];
+            bool chunk_argmax_keep_dims_0 = const()[name = string("chunk_argmax_keep_dims_0"), val = bool(true)];
+            string chunk_argmax_output_dtype_0 = const()[name = string("chunk_argmax_output_dtype_0"), val = string("int32")];
+            tensor<fp16, [1, 1, 16384]> logits_cast_fp16 = transpose(perm = logits_perm_0, x = var_421_cast_fp16)[name = string("transpose_0")];
+            tensor<int32, [1, 1, 1]> chunk_argmax_cast_fp16 = reduce_argmax(axis = chunk_argmax_axis_0, keep_dims = chunk_argmax_keep_dims_0, output_dtype = chunk_argmax_output_dtype_0, x = logits_cast_fp16)[name = string("chunk_argmax_cast_fp16")];
+            int32 var_593 = const()[name = string("op_593"), val = int32(-1)];
+            bool chunk_max_val_validate_indices_0 = const()[name = string("chunk_max_val_validate_indices_0"), val = bool(false)];
+            string chunk_argmax_cast_fp16_to_int16_dtype_0 = const()[name = string("chunk_argmax_cast_fp16_to_int16_dtype_0"), val = string("int16")];
+            tensor<int16, [1, 1, 1]> chunk_argmax_cast_fp16_to_int16 = cast(dtype = chunk_argmax_cast_fp16_to_int16_dtype_0, x = chunk_argmax_cast_fp16)[name = string("cast_2")];
+            tensor<fp16, [1, 1, 1]> chunk_max_val_cast_fp16_cast_int16 = gather_along_axis(axis = var_593, indices = chunk_argmax_cast_fp16_to_int16, validate_indices = chunk_max_val_validate_indices_0, x = logits_cast_fp16)[name = string("chunk_max_val_cast_fp16_cast_int16")];
+            int32 var_602 = const()[name = string("op_602"), val = int32(-1)];
+            bool var_603_interleave_0 = const()[name = string("op_603_interleave_0"), val = bool(false)];
+            tensor<int32, [1, 1, 16]> var_603 = concat(axis = var_602, interleave = var_603_interleave_0, values = (chunk_argmax_1_cast_fp16, chunk_argmax_3_cast_fp16, chunk_argmax_5_cast_fp16, chunk_argmax_7_cast_fp16, chunk_argmax_9_cast_fp16, chunk_argmax_11_cast_fp16, chunk_argmax_13_cast_fp16, chunk_argmax_15_cast_fp16, chunk_argmax_17_cast_fp16, chunk_argmax_19_cast_fp16, chunk_argmax_21_cast_fp16, chunk_argmax_23_cast_fp16, chunk_argmax_25_cast_fp16, chunk_argmax_27_cast_fp16, chunk_argmax_29_cast_fp16, chunk_argmax_cast_fp16))[name = string("op_603")];
+            tensor<int32, [1]> var_605_axes_0 = const()[name = string("op_605_axes_0"), val = tensor<int32, [1]>([0])];
+            string var_603_to_int16_dtype_0 = const()[name = string("op_603_to_int16_dtype_0"), val = string("int16")];
+            tensor<int16, [1, 1, 16]> var_603_to_int16 = cast(dtype = var_603_to_int16_dtype_0, x = var_603)[name = string("cast_1")];
+            tensor<int16, [1, 16]> var_605_cast_uint16 = squeeze(axes = var_605_axes_0, x = var_603_to_int16)[name = string("op_605_cast_uint16")];
+            tensor<int32, [1]> var_607_axes_0 = const()[name = string("op_607_axes_0"), val = tensor<int32, [1]>([0])];
+            tensor<int16, [16]> var_607_cast_uint16 = squeeze(axes = var_607_axes_0, x = var_605_cast_uint16)[name = string("op_607_cast_uint16")];
+            string var_607_cast_uint16_to_int32_dtype_0 = const()[name = string("op_607_cast_uint16_to_int32_dtype_0"), val = string("int32")];
+            int32 var_609 = const()[name = string("op_609"), val = int32(-1)];
+            bool var_610_interleave_0 = const()[name = string("op_610_interleave_0"), val = bool(false)];
+            tensor<fp16, [1, 1, 16]> var_610_cast_fp16 = concat(axis = var_609, interleave = var_610_interleave_0, values = (var_430_cast_fp16_cast_int16, var_441_cast_fp16_cast_int16, var_452_cast_fp16_cast_int16, var_463_cast_fp16_cast_int16, var_474_cast_fp16_cast_int16, var_485_cast_fp16_cast_int16, var_496_cast_fp16_cast_int16, var_507_cast_fp16_cast_int16, var_518_cast_fp16_cast_int16, var_529_cast_fp16_cast_int16, var_540_cast_fp16_cast_int16, var_551_cast_fp16_cast_int16, var_562_cast_fp16_cast_int16, var_573_cast_fp16_cast_int16, var_584_cast_fp16_cast_int16, chunk_max_val_cast_fp16_cast_int16))[name = string("op_610_cast_fp16")];
+            tensor<int32, [1]> var_612_axes_0 = const()[name = string("op_612_axes_0"), val = tensor<int32, [1]>([0])];
+            tensor<fp16, [1, 16]> var_612_cast_fp16 = squeeze(axes = var_612_axes_0, x = var_610_cast_fp16)[name = string("op_612_cast_fp16")];
+            tensor<int32, [1]> var_614_axes_0 = const()[name = string("op_614_axes_0"), val = tensor<int32, [1]>([0])];
+            tensor<fp16, [16]> argmax_val = squeeze(axes = var_614_axes_0, x = var_612_cast_fp16)[name = string("op_614_cast_fp16")];
+            tensor<int32, [16]> argmax_idx = cast(dtype = var_607_cast_uint16_to_int32_dtype_0, x = var_607_cast_uint16)[name = string("cast_0")];
+        } -> (argmax_idx, argmax_val);
+}

gemma3_lm_head_lut6.mlmodelc/weights/weight.bin

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

meta.yaml

@@ -0,0 +1,26 @@
+model_info:
+  name: anemll-google-gemma-3-1b-it-ctx4096
+  version: 0.3.4
+  description: |
+    Demonstarates running google-gemma-3-1b-it on Apple Neural Engine
+    Context length: 4096
+    Batch size: 64
+    Chunks: 1
+  license: MIT
+  author: Anemll
+  framework: Core ML
+  language: Python
+  architecture: gemma3_text
+  parameters:
+    context_length: 4096
+    batch_size: 64
+    lut_embeddings: 6
+    lut_ffn: 6
+    lut_lmhead: 6
+    num_chunks: 1
+    model_prefix: gemma3
+    embeddings: gemma3_embeddings_lut6.mlmodelc
+    lm_head: gemma3_lm_head_lut6.mlmodelc
+    ffn: gemma3_FFN_PF_lut6_chunk_01of01.mlmodelc
+    split_lm_head: 16
+    argmax_in_model: true

tokenizer.json

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

tokenizer.model

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