From d4a70276f4692cb3885864c7e67517e7080a2cdc Mon Sep 17 00:00:00 2001 From: Guru Deshpande <66385690+guru-desh@users.noreply.github.com> Date: Mon, 5 Oct 2026 20:57:01 -0700 Subject: [PATCH 1/7] Move the export's model loading and compression into resolve_model and build_model _async_export_model's resolution, loading and compression move, unchanged, into two functions: resolve_model reads the checkpoint's config and resolves it against the registry, and build_model loads the weights and compresses them. _async_export_model calls both. The moved lines are the original lines, re-indented where they leave the `with` block. The one exception is an assert that ruff re-wraps once it is less indented. Review with: git show --color-moved=dimmed-zebra --color-moved-ws=allow-indentation-change On a tiny Qwen3, the model, config and externalized model handed to the exporter are identical before and after, for macOS uncompressed, macOS 4-bit and iOS palettized. --- python/src/coreai_models/export/pipeline.py | 355 +++++++++++--------- 1 file changed, 203 insertions(+), 152 deletions(-) diff --git a/python/src/coreai_models/export/pipeline.py b/python/src/coreai_models/export/pipeline.py index f1f43ba9..02a1ca3b 100644 --- a/python/src/coreai_models/export/pipeline.py +++ b/python/src/coreai_models/export/pipeline.py @@ -18,7 +18,7 @@ import tempfile from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Literal +from typing import Any, Literal, NamedTuple import torch from transformers import AutoConfig, AutoTokenizer @@ -44,7 +44,7 @@ DEFAULT_MACOS_COMPRESSION_PRESET, get_preset, ) -from coreai_models.models.registry import get_model_entry +from coreai_models.models.registry import ModelEntry, get_model_entry logger = logging.getLogger(__name__) @@ -118,34 +118,28 @@ def _resolve_precision(precision_str: str) -> torch.dtype: return dtype -def export_model(config_or_model_id: ExportConfig | str) -> str: - """Export a HuggingFace model to Core AI format. +@dataclass(frozen=True) +class ResolvedModel: + """What an ``ExportConfig`` resolves to before any weight is read.""" - This is the main public API. It orchestrates: - 1. Resolve model class from HuggingFace config - 2. Load model with HF weights - 3. Apply pre-export compression (torch quantization) - 4. Variant-specific export (macOS or iOS) - 5. Save as .aimodel - - Args: - config_or_model_id: Either an ExportConfig or a HuggingFace model ID string. - If a string, uses default settings (macOS, 4bit compression). + config: ExportConfig + entry: ModelEntry + model_type: str + model_class: Any + hf_config: Any + target_dtype: torch.dtype + max_context_length: int | None - Returns: - Path to the exported .aimodel file. - """ - if isinstance(config_or_model_id, str): - config = ExportConfig(hf_model_id=config_or_model_id) - else: - config = config_or_model_id - return asyncio.run(_async_export_model(config)) +class BuiltModel(NamedTuple): + """The torch model a ``ResolvedModel`` describes, loaded and compressed.""" + model: Any + externalized_model: torch.nn.Module | None -async def _async_export_model(config: ExportConfig) -> str: - """Async implementation of export_model.""" +def resolve_model(config: ExportConfig) -> ResolvedModel: + """Resolve ``config`` against the registry and the checkpoint's config, reading no weights.""" # ---- 1. Resolve model class ---- hf_config = AutoConfig.from_pretrained(config.hf_model_id) model_type = config.model_type_override or getattr(hf_config, "model_type", None) @@ -203,7 +197,191 @@ async def _async_export_model(config: ExportConfig) -> str: if config.num_layers is not None: hf_config.num_hidden_layers = config.num_layers + return ResolvedModel( + config=config, + entry=entry, + model_type=model_type, + model_class=model_class, + hf_config=hf_config, + target_dtype=target_dtype, + max_context_length=max_context_length, + ) + + +def build_model(resolved: ResolvedModel, temp_dir: str | None) -> BuiltModel: + """Load ``resolved``'s weights and apply its compression. + + Args: + temp_dir: The macOS export's working directory, or None for iOS. + """ + config, entry, model_class = resolved.config, resolved.entry, resolved.model_class + hf_config, target_dtype = resolved.hf_config, resolved.target_dtype + max_context_length = resolved.max_context_length + logger.info(f"Loading {config.hf_model_id} ({config.variant}, dtype={target_dtype})...") + use_memory_efficient = config.variant == "macOS" + if use_memory_efficient: + assert temp_dir is not None # nullcontext yields None only when not memory-efficient + layer_mmap_dir = os.path.join(temp_dir, "layers") + os.makedirs(layer_mmap_dir, exist_ok=True) + model = model_class.from_hf_memory_efficient( + config.hf_model_id, + max_context_length=max_context_length, + target_dtype=target_dtype, + mmap_path=layer_mmap_dir, + num_layers=config.num_layers, + hf_config_attr=entry.hf_config_attr, + hf_state_dict_prefix=entry.hf_state_dict_prefix, + ) + else: + model = model_class.from_hf( + config.hf_model_id, + max_context_length=max_context_length, + target_dtype=target_dtype, + num_layers=config.num_layers, + disable_embedding_quantization=config.disable_embedding_quantization, + ) + model = model.eval() + # ---- 3. Resolve compression preset ---- + if config.compression_config_object is not None: + torch_quantization_config, torch_palettization_config = split_compression_config( + config.compression_config_object + ) + else: + preset = get_preset(config.compression) + torch_quantization_config = preset.get("torch_quantization_config") + torch_palettization_config = preset.get("torch_palettization_config") + + assert not (torch_quantization_config is not None and torch_palettization_config is not None), ( + "Both a quantization and a palettization config were provided, this should never happen." + ) + + # ---- 3a. Pre-export torch quantization (if configured) ---- + effective_max_ctx = max_context_length or getattr( + hf_config, "max_position_embeddings", TRACE_KV_CACHE_SEQ_LEN + ) + vocab_size = hf_config.vocab_size + batch_size = 1 + # Set when composite ops are marked for externalization before quantization. + externalized_model: torch.nn.Module | None = None + if torch_quantization_config is not None: + logger.info(f"Applying pre-export torch quantization (preset={config.compression})") + + def get_calibration_data(): # type: ignore[no-untyped-def] + tokenizer = AutoTokenizer.from_pretrained(config.hf_model_id) + return get_c4(tokenizer) + + # Copy so we don't mutate the shared preset. + quant_cfg = dict(torch_quantization_config) + + # The preset or YAML is the source of truth; `--quantization-mode` overrides + # it only when given. + if config.quantization_mode is not None: + logger.warning( + "Overriding execution_mode for `coreai-opt` compression with " + f"{config.quantization_mode}" + ) + quant_cfg["execution_mode"] = config.quantization_mode + elif "execution_mode" not in quant_cfg: + raise ValueError( + f"Compression config '{config.compression}' does not set " + "'execution_mode'. Set it there, or pass --quantization-mode " + "{eager,graph}." + ) + + graph_mode = is_compression_mode_graph(quant_cfg) + + quantizer_mmap_dir: str | None = None + # coreai-opt only supports mmap-backed finalization in eager mode. + if use_memory_efficient and not graph_mode: + assert temp_dir is not None + quantizer_mmap_dir = os.path.join(temp_dir, "quantized") + os.makedirs(quantizer_mmap_dir, exist_ok=True) + + if graph_mode: + patch_model_for_externalization(model) + # externalization patches live on the eager module's composite op + # submodules but quantize_for_export in graph-mode below returns a + # new GraphModule which get overwritten to `model`. + # So, keep a handle on the eager module the composites were patched on, + # for the sub-export later. + externalized_model = model + + model = quantize_for_export( + model, + hf_config, + target_dtype, + quant_cfg, + calibration_data_fn=get_calibration_data, + mmap_dir=quantizer_mmap_dir, + ) + + if torch_palettization_config is not None: + assert config.variant == "iOS", "palettization is only supported for iOS variant." + + query_len = 8 + input_ids = torch.randint(1, vocab_size, (batch_size, query_len), dtype=torch.int32) + position_ids = ( + torch.arange(query_len).to(torch.uint16).unsqueeze(0).expand(batch_size, query_len) + ) + in_step = torch.zeros((1,), dtype=torch.int32) + causal_mask = torch.zeros(1, effective_max_ctx, 1, query_len, dtype=torch.float16) + if hasattr(hf_config, "head_dim") and isinstance(hf_config.head_dim, int): + head_dim = hf_config.head_dim + else: + head_dim = hf_config.hidden_size // hf_config.num_attention_heads + key_cache = torch.zeros( + hf_config.num_hidden_layers, + 1, # batch_size + hf_config.num_key_value_heads * head_dim, + 1, + effective_max_ctx, + dtype=torch.float16, + ) + value_cache = key_cache.clone() + palettization_inputs = ( + input_ids, + position_ids, + in_step, + causal_mask, + key_cache, + value_cache, + ) + model = palettize_pytorch_model(model, palettization_inputs, torch_palettization_config) + + return BuiltModel(model, externalized_model) + + +def export_model(config_or_model_id: ExportConfig | str) -> str: + """Export a HuggingFace model to Core AI format. + + This is the main public API. It orchestrates: + 1. Resolve model class from HuggingFace config + 2. Load model with HF weights + 3. Apply pre-export compression (torch quantization) + 4. Variant-specific export (macOS or iOS) + 5. Save as .aimodel + + Args: + config_or_model_id: Either an ExportConfig or a HuggingFace model ID string. + If a string, uses default settings (macOS, 4bit compression). + + Returns: + Path to the exported .aimodel file. + """ + if isinstance(config_or_model_id, str): + config = ExportConfig(hf_model_id=config_or_model_id) + else: + config = config_or_model_id + + return asyncio.run(_async_export_model(config)) + + +async def _async_export_model(config: ExportConfig) -> str: + """Async implementation of export_model.""" + resolved = resolve_model(config) + entry, hf_config, model_type = resolved.entry, resolved.hf_config, resolved.model_type + target_dtype, max_context_length = resolved.target_dtype, resolved.max_context_length # Memory-efficient layer-by-layer loading + quantizer disk-checkpointing # is macOS-only for now. The iOS variant keeps the legacy full-RAM path @@ -217,134 +395,7 @@ async def _async_export_model(config: ExportConfig) -> str: ) with temp_dir_ctx as temp_dir: - if use_memory_efficient: - assert temp_dir is not None # nullcontext yields None only when not memory-efficient - layer_mmap_dir = os.path.join(temp_dir, "layers") - os.makedirs(layer_mmap_dir, exist_ok=True) - model = model_class.from_hf_memory_efficient( - config.hf_model_id, - max_context_length=max_context_length, - target_dtype=target_dtype, - mmap_path=layer_mmap_dir, - num_layers=config.num_layers, - hf_config_attr=entry.hf_config_attr, - hf_state_dict_prefix=entry.hf_state_dict_prefix, - ) - else: - model = model_class.from_hf( - config.hf_model_id, - max_context_length=max_context_length, - target_dtype=target_dtype, - num_layers=config.num_layers, - disable_embedding_quantization=config.disable_embedding_quantization, - ) - model = model.eval() - # ---- 3. Resolve compression preset ---- - if config.compression_config_object is not None: - torch_quantization_config, torch_palettization_config = split_compression_config( - config.compression_config_object - ) - else: - preset = get_preset(config.compression) - torch_quantization_config = preset.get("torch_quantization_config") - torch_palettization_config = preset.get("torch_palettization_config") - - assert not ( - torch_quantization_config is not None and torch_palettization_config is not None - ), "Both a quantization and a palettization config were provided, this should never happen." - - # ---- 3a. Pre-export torch quantization (if configured) ---- - effective_max_ctx = max_context_length or getattr( - hf_config, "max_position_embeddings", TRACE_KV_CACHE_SEQ_LEN - ) - vocab_size = hf_config.vocab_size - batch_size = 1 - # Set when composite ops are marked for externalization before quantization. - externalized_model: torch.nn.Module | None = None - if torch_quantization_config is not None: - logger.info(f"Applying pre-export torch quantization (preset={config.compression})") - - def get_calibration_data(): # type: ignore[no-untyped-def] - tokenizer = AutoTokenizer.from_pretrained(config.hf_model_id) - return get_c4(tokenizer) - - # Copy so we don't mutate the shared preset. - quant_cfg = dict(torch_quantization_config) - - # The preset or YAML is the source of truth; `--quantization-mode` overrides - # it only when given. - if config.quantization_mode is not None: - logger.warning( - "Overriding execution_mode for `coreai-opt` compression with " - f"{config.quantization_mode}" - ) - quant_cfg["execution_mode"] = config.quantization_mode - elif "execution_mode" not in quant_cfg: - raise ValueError( - f"Compression config '{config.compression}' does not set " - "'execution_mode'. Set it there, or pass --quantization-mode " - "{eager,graph}." - ) - - graph_mode = is_compression_mode_graph(quant_cfg) - - quantizer_mmap_dir: str | None = None - # coreai-opt only supports mmap-backed finalization in eager mode. - if use_memory_efficient and not graph_mode: - assert temp_dir is not None - quantizer_mmap_dir = os.path.join(temp_dir, "quantized") - os.makedirs(quantizer_mmap_dir, exist_ok=True) - - if graph_mode: - patch_model_for_externalization(model) - # externalization patches live on the eager module's composite op - # submodules but quantize_for_export in graph-mode below returns a - # new GraphModule which get overwritten to `model`. - # So, keep a handle on the eager module the composites were patched on, - # for the sub-export later. - externalized_model = model - - model = quantize_for_export( - model, - hf_config, - target_dtype, - quant_cfg, - calibration_data_fn=get_calibration_data, - mmap_dir=quantizer_mmap_dir, - ) - - if torch_palettization_config is not None: - assert config.variant == "iOS", "palettization is only supported for iOS variant." - - query_len = 8 - input_ids = torch.randint(1, vocab_size, (batch_size, query_len), dtype=torch.int32) - position_ids = ( - torch.arange(query_len).to(torch.uint16).unsqueeze(0).expand(batch_size, query_len) - ) - in_step = torch.zeros((1,), dtype=torch.int32) - causal_mask = torch.zeros(1, effective_max_ctx, 1, query_len, dtype=torch.float16) - if hasattr(hf_config, "head_dim") and isinstance(hf_config.head_dim, int): - head_dim = hf_config.head_dim - else: - head_dim = hf_config.hidden_size // hf_config.num_attention_heads - key_cache = torch.zeros( - hf_config.num_hidden_layers, - 1, # batch_size - hf_config.num_key_value_heads * head_dim, - 1, - effective_max_ctx, - dtype=torch.float16, - ) - value_cache = key_cache.clone() - palettization_inputs = ( - input_ids, - position_ids, - in_step, - causal_mask, - key_cache, - value_cache, - ) - model = palettize_pytorch_model(model, palettization_inputs, torch_palettization_config) + model, externalized_model = build_model(resolved, temp_dir) # ---- 4. Variant-specific export ---- if config.variant == "macOS": From d5c3af463827db594710f9f792ac6fb2a2e20529 Mon Sep 17 00:00:00 2001 From: Guru Deshpande <66385690+guru-desh@users.noreply.github.com> Date: Mon, 5 Oct 2026 20:57:02 -0700 Subject: [PATCH 2/7] Move the export CLI's compression resolution into resolve_compression The if/elif chain in _resolve_export_config moves, unchanged, into resolve_compression, so another entry point can resolve a compression exactly as the export CLI does. The only edit is `args.compression_config` becoming the `compression_config` parameter. --- python/src/coreai_models/llm/export.py | 65 ++++++++++++++++---------- 1 file changed, 40 insertions(+), 25 deletions(-) diff --git a/python/src/coreai_models/llm/export.py b/python/src/coreai_models/llm/export.py index 0d82a966..df8a3cc5 100644 --- a/python/src/coreai_models/llm/export.py +++ b/python/src/coreai_models/llm/export.py @@ -289,6 +289,40 @@ def _resolve_registry_compression_config(relpath: str, variant: str) -> tuple[ob return _load_compression_config_object(yaml_path, variant), yaml_path.stem +def resolve_compression( + variant: str, + *, + compression: str | None = None, + compression_config: Path | None = None, + registry_compression_config: str | None = None, +) -> tuple[str, object]: + """The compression name and coreai-opt config object an export of ``variant`` uses.""" + compression_config_object = None + if compression_config is not None: + if not compression_config.is_file(): + raise SystemExit(f"--compression-config: file not found: {compression_config}") + compression_config_object = _load_compression_config_object(compression_config, variant) + compression = compression_config.stem + elif registry_compression_config is not None: + compression_config_object, compression = _resolve_registry_compression_config( + registry_compression_config, variant + ) + elif compression == "none": + pass + elif not compression: + compression = MACOS_DEFAULT if variant == "macOS" else IOS_DEFAULT + elif compression in MACOS_PRESETS and variant == "iOS": + raise RuntimeError("macOS quantization preset provided, but platform is iOS.") + elif compression in IOS_PRESETS and variant == "macOS": + raise RuntimeError("iOS palettization preset provided, but platform is macOS.") + elif compression not in ALL_PRESET_NAMES and compression != "none": + raise RuntimeError( + f"Compression preset {compression} is not a valid compression " + f"preset. Available: {list_presets()}" + ) + return compression, compression_config_object + + def _resolve_export_config(args: argparse.Namespace) -> ExportConfig: """Resolve CLI args + registry preset into a final ExportConfig.""" hf_model_id = args.model @@ -297,7 +331,6 @@ def _resolve_export_config(args: argparse.Namespace) -> ExportConfig: compute_precision = args.compute_precision max_context_length = args.max_context_length output_dir = args.output_dir or _default_output_dir() - compression_config_object = None registry_compression_config: str | None = None # repo-root-relative path, e.g. models/qwen3/... preset = None @@ -360,30 +393,12 @@ def _resolve_export_config(args: argparse.Namespace) -> ExportConfig: if args.quantization_mode == "graph" and variant != "macOS": raise SystemExit(f"--quantization-mode graph requires --platform macOS (got '{variant}').") - if args.compression_config is not None: - if not args.compression_config.is_file(): - raise SystemExit(f"--compression-config: file not found: {args.compression_config}") - compression_config_object = _load_compression_config_object( - args.compression_config, variant - ) - compression = args.compression_config.stem - elif registry_compression_config is not None: - compression_config_object, compression = _resolve_registry_compression_config( - registry_compression_config, variant - ) - elif compression == "none": - pass - elif not compression: - compression = MACOS_DEFAULT if variant == "macOS" else IOS_DEFAULT - elif compression in MACOS_PRESETS and variant == "iOS": - raise RuntimeError("macOS quantization preset provided, but platform is iOS.") - elif compression in IOS_PRESETS and variant == "macOS": - raise RuntimeError("iOS palettization preset provided, but platform is macOS.") - elif compression not in ALL_PRESET_NAMES and compression != "none": - raise RuntimeError( - f"Compression preset {compression} is not a valid compression " - f"preset. Available: {list_presets()}" - ) + compression, compression_config_object = resolve_compression( + variant, + compression=compression, + compression_config=args.compression_config, + registry_compression_config=registry_compression_config, + ) return ExportConfig( hf_model_id=hf_model_id, From 4a9209715f150ddc0eec9db92a8cf6175041f353 Mon Sep 17 00:00:00 2001 From: Guru Deshpande <66385690+guru-desh@users.noreply.github.com> Date: Mon, 5 Oct 2026 20:58:15 -0700 Subject: [PATCH 3/7] Resolve the compression preset in resolve_model The "Resolve compression preset" block moves, unchanged, from build_model into resolve_model, and ResolvedModel carries the two configs it resolves. A caller can then see which compression a model will get before any weight is read. The one behaviour change: a config naming both a quantization and a palettization now fails before the weights load rather than after. --- python/src/coreai_models/export/pipeline.py | 33 +++++++++++++-------- 1 file changed, 20 insertions(+), 13 deletions(-) diff --git a/python/src/coreai_models/export/pipeline.py b/python/src/coreai_models/export/pipeline.py index 02a1ca3b..c1be5f82 100644 --- a/python/src/coreai_models/export/pipeline.py +++ b/python/src/coreai_models/export/pipeline.py @@ -129,6 +129,8 @@ class ResolvedModel: hf_config: Any target_dtype: torch.dtype max_context_length: int | None + quantization_config: dict | None + palettization_config: Any class BuiltModel(NamedTuple): @@ -197,6 +199,20 @@ def resolve_model(config: ExportConfig) -> ResolvedModel: if config.num_layers is not None: hf_config.num_hidden_layers = config.num_layers + # ---- 3. Resolve compression preset ---- + if config.compression_config_object is not None: + torch_quantization_config, torch_palettization_config = split_compression_config( + config.compression_config_object + ) + else: + preset = get_preset(config.compression) + torch_quantization_config = preset.get("torch_quantization_config") + torch_palettization_config = preset.get("torch_palettization_config") + + assert not (torch_quantization_config is not None and torch_palettization_config is not None), ( + "Both a quantization and a palettization config were provided, this should never happen." + ) + return ResolvedModel( config=config, entry=entry, @@ -205,6 +221,8 @@ def resolve_model(config: ExportConfig) -> ResolvedModel: hf_config=hf_config, target_dtype=target_dtype, max_context_length=max_context_length, + quantization_config=torch_quantization_config, + palettization_config=torch_palettization_config, ) @@ -242,19 +260,8 @@ def build_model(resolved: ResolvedModel, temp_dir: str | None) -> BuiltModel: disable_embedding_quantization=config.disable_embedding_quantization, ) model = model.eval() - # ---- 3. Resolve compression preset ---- - if config.compression_config_object is not None: - torch_quantization_config, torch_palettization_config = split_compression_config( - config.compression_config_object - ) - else: - preset = get_preset(config.compression) - torch_quantization_config = preset.get("torch_quantization_config") - torch_palettization_config = preset.get("torch_palettization_config") - - assert not (torch_quantization_config is not None and torch_palettization_config is not None), ( - "Both a quantization and a palettization config were provided, this should never happen." - ) + torch_quantization_config = resolved.quantization_config + torch_palettization_config = resolved.palettization_config # ---- 3a. Pre-export torch quantization (if configured) ---- effective_max_ctx = max_context_length or getattr( From d3c12ec0f4f04b53554d8e663ec859144dc23750 Mon Sep 17 00:00:00 2001 From: Guru Deshpande <66385690+guru-desh@users.noreply.github.com> Date: Mon, 5 Oct 2026 20:59:00 -0700 Subject: [PATCH 4/7] Pass build_model's export backend through to quantize_for_export quantize_for_export already takes the backend coreai-opt finalizes quantized weights for; build_model now passes one through, defaulting to Core AI as before. With another backend nothing is exported, so graph-mode quantization skips marking composite ops for externalization. --- python/src/coreai_models/export/pipeline.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/python/src/coreai_models/export/pipeline.py b/python/src/coreai_models/export/pipeline.py index c1be5f82..1da3a59a 100644 --- a/python/src/coreai_models/export/pipeline.py +++ b/python/src/coreai_models/export/pipeline.py @@ -226,11 +226,15 @@ def resolve_model(config: ExportConfig) -> ResolvedModel: ) -def build_model(resolved: ResolvedModel, temp_dir: str | None) -> BuiltModel: +def build_model( + resolved: ResolvedModel, temp_dir: str | None, export_backend: object | None = None +) -> BuiltModel: """Load ``resolved``'s weights and apply its compression. Args: temp_dir: The macOS export's working directory, or None for iOS. + export_backend: The backend coreai-opt finalizes quantized weights for. None is + Core AI. """ config, entry, model_class = resolved.config, resolved.entry, resolved.model_class hf_config, target_dtype = resolved.hf_config, resolved.target_dtype @@ -305,7 +309,7 @@ def get_calibration_data(): # type: ignore[no-untyped-def] quantizer_mmap_dir = os.path.join(temp_dir, "quantized") os.makedirs(quantizer_mmap_dir, exist_ok=True) - if graph_mode: + if graph_mode and export_backend is None: patch_model_for_externalization(model) # externalization patches live on the eager module's composite op # submodules but quantize_for_export in graph-mode below returns a @@ -321,6 +325,7 @@ def get_calibration_data(): # type: ignore[no-untyped-def] quant_cfg, calibration_data_fn=get_calibration_data, mmap_dir=quantizer_mmap_dir, + export_backend=export_backend, ) if torch_palettization_config is not None: From 82fbd1078d483d0fa72b336816e8b0fa3dd392eb Mon Sep 17 00:00:00 2001 From: Guru Deshpande <66385690+guru-desh@users.noreply.github.com> Date: Mon, 5 Oct 2026 21:00:24 -0700 Subject: [PATCH 5/7] Let build_model keep macOS weights in RAM when it gets no temp_dir macOS still loads through from_hf_memory_efficient, with mmap_path=None, and the quantizer skips its disk checkpoints. The export always passes a temp_dir on macOS, so it is unaffected. --- python/src/coreai_models/export/pipeline.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/python/src/coreai_models/export/pipeline.py b/python/src/coreai_models/export/pipeline.py index 1da3a59a..0eae0f9a 100644 --- a/python/src/coreai_models/export/pipeline.py +++ b/python/src/coreai_models/export/pipeline.py @@ -227,12 +227,13 @@ def resolve_model(config: ExportConfig) -> ResolvedModel: def build_model( - resolved: ResolvedModel, temp_dir: str | None, export_backend: object | None = None + resolved: ResolvedModel, temp_dir: str | None = None, export_backend: object | None = None ) -> BuiltModel: """Load ``resolved``'s weights and apply its compression. Args: - temp_dir: The macOS export's working directory, or None for iOS. + temp_dir: Where macOS weights and quantizer checkpoints are disk-backed. None keeps + them in RAM. export_backend: The backend coreai-opt finalizes quantized weights for. None is Core AI. """ @@ -243,9 +244,10 @@ def build_model( logger.info(f"Loading {config.hf_model_id} ({config.variant}, dtype={target_dtype})...") use_memory_efficient = config.variant == "macOS" if use_memory_efficient: - assert temp_dir is not None # nullcontext yields None only when not memory-efficient - layer_mmap_dir = os.path.join(temp_dir, "layers") - os.makedirs(layer_mmap_dir, exist_ok=True) + layer_mmap_dir = None + if temp_dir is not None: + layer_mmap_dir = os.path.join(temp_dir, "layers") + os.makedirs(layer_mmap_dir, exist_ok=True) model = model_class.from_hf_memory_efficient( config.hf_model_id, max_context_length=max_context_length, @@ -304,8 +306,7 @@ def get_calibration_data(): # type: ignore[no-untyped-def] quantizer_mmap_dir: str | None = None # coreai-opt only supports mmap-backed finalization in eager mode. - if use_memory_efficient and not graph_mode: - assert temp_dir is not None + if use_memory_efficient and temp_dir is not None and not graph_mode: quantizer_mmap_dir = os.path.join(temp_dir, "quantized") os.makedirs(quantizer_mmap_dir, exist_ok=True) From e4993d3c7943bd936e58ff03b2456c84976f325e Mon Sep 17 00:00:00 2001 From: Guru Deshpande <66385690+guru-desh@users.noreply.github.com> Date: Mon, 5 Oct 2026 21:03:21 -0700 Subject: [PATCH 6/7] Add ExportConfig.from_preset The config that exports a registry preset, with any value passed in winning over the preset's, for callers that start from a registry preset rather than from CLI arguments. The export CLI keeps its own resolution, unchanged. --- python/src/coreai_models/export/pipeline.py | 34 ++++++++++++ .../test_export/test_from_preset.py | 52 +++++++++++++++++++ 2 files changed, 86 insertions(+) create mode 100644 python/tests/test_model_units/test_export/test_from_preset.py diff --git a/python/src/coreai_models/export/pipeline.py b/python/src/coreai_models/export/pipeline.py index 0eae0f9a..9176eced 100644 --- a/python/src/coreai_models/export/pipeline.py +++ b/python/src/coreai_models/export/pipeline.py @@ -44,6 +44,7 @@ DEFAULT_MACOS_COMPRESSION_PRESET, get_preset, ) +from coreai_models.model_registry import ModelPreset from coreai_models.models.registry import ModelEntry, get_model_entry logger = logging.getLogger(__name__) @@ -84,6 +85,39 @@ def __post_init__(self) -> None: f"quantization_mode='graph' is macOS only (got variant '{self.variant}')." ) + @classmethod + def from_preset( + cls, + preset: ModelPreset, + *, + platform: str | None = None, + compression: str | None = None, + compression_config_object: Any = None, + compute_precision: str | None = None, + max_context_length: int | None = None, + **options: Any, + ) -> "ExportConfig": + """The config that exports ``preset``, with any value passed here winning over the preset's. + + Raises: + ValueError: If neither ``compute_precision`` nor the preset names a precision. + """ + precision = compute_precision or preset.compute_precision + if precision is None: + raise ValueError( + f"{preset.short_name!r} registers no compute precision; pass one explicitly." + ) + return cls( + hf_model_id=preset.hf_id, + variant=platform or preset.variant or "macOS", + compression=compression if compression is not None else preset.compression, + compression_config_object=compression_config_object, + compute_precision=precision, + max_context_length=max_context_length or preset.max_context_length, + model_type_override=preset._model_type_override, + **options, + ) + def _generate_output_name(config: ExportConfig) -> str: """Generate a filesystem-safe output name from config.""" diff --git a/python/tests/test_model_units/test_export/test_from_preset.py b/python/tests/test_model_units/test_export/test_from_preset.py new file mode 100644 index 00000000..733cac6b --- /dev/null +++ b/python/tests/test_model_units/test_export/test_from_preset.py @@ -0,0 +1,52 @@ +# Copyright 2026 Apple Inc. +# +# Use of this source code is governed by a BSD-3-clause license that can +# be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause + +"""``ExportConfig.from_preset`` and ``resolve_compression``.""" + +import dataclasses +from pathlib import Path + +import pytest + +from coreai_models.export.pipeline import ExportConfig +from coreai_models.llm.export import resolve_compression +from coreai_models.model_registry import LLM_PRESETS, lookup_preset + + +@pytest.mark.parametrize("preset", LLM_PRESETS, ids=lambda p: f"{p.short_name}-{p.variant}") +def test_every_registry_field_reaches_the_config(preset) -> None: + config = ExportConfig.from_preset( + preset, compute_precision=preset.compute_precision or "float16" + ) + assert config.hf_model_id == preset.hf_id + assert config.variant == preset.variant + assert config.compression == preset.compression + assert config.max_context_length == preset.max_context_length + assert config.model_type_override == preset._model_type_override + + +def test_stated_values_win_over_the_preset() -> None: + preset = lookup_preset("qwen3-0.6b", model_type="llm", variant="macOS") + config = ExportConfig.from_preset( + preset, compression="none", compute_precision="float32", max_context_length=512 + ) + assert (config.compression, config.compute_precision, config.max_context_length) == ( + "none", + "float32", + 512, + ) + + +def test_a_preset_without_a_precision_needs_one_stated() -> None: + preset = dataclasses.replace( + lookup_preset("qwen3-0.6b", model_type="llm", variant="macOS"), compute_precision=None + ) + with pytest.raises(ValueError, match="compute precision"): + ExportConfig.from_preset(preset) + + +def test_a_missing_compression_yaml_exits_with_the_flags_message() -> None: + with pytest.raises(SystemExit, match="--compression-config: file not found"): + resolve_compression("macOS", compression_config=Path("does/not/exist.yaml")) From ce7cec775f39bf22bfcf33c74e452cb83ac42647 Mon Sep 17 00:00:00 2001 From: Guru Deshpande <66385690+guru-desh@users.noreply.github.com> Date: Mon, 5 Oct 2026 21:04:35 -0700 Subject: [PATCH 7/7] Add an eager inference API to the authored models new_eager_state, eager_forward_args and eager_step call forward on real inputs, one sequence at a time. macOS takes positions over the whole history and the KV pair; iOS takes only the query's positions, the write offset, and an additive mask over the context, which causal_mask builds. Export's palettization keeps its own trace inputs, so what it traces is unchanged. --- python/src/coreai_models/models/base.py | 63 +++++++++++++++++++ .../src/coreai_models/primitives/ios/cache.py | 17 +++++ 2 files changed, 80 insertions(+) diff --git a/python/src/coreai_models/models/base.py b/python/src/coreai_models/models/base.py index 2e064405..9afb12f0 100644 --- a/python/src/coreai_models/models/base.py +++ b/python/src/coreai_models/models/base.py @@ -50,12 +50,21 @@ VALUE_CACHE_NAME, VALUE_CACHE_OUTPUT_NAME, ) +from coreai_models.primitives.ios.cache import KVCacheHandler, causal_mask from coreai_models.primitives.ios.embedding import GatherEmbeddings, LoadEmbeddings from coreai_models.primitives.macos.cache import KVCache T = TypeVar("T", bound="BaseForCausalLM") +@dataclass +class EagerState: + """The caches of one sequence being run eagerly, and how many positions they already hold.""" + + tensors: tuple[torch.Tensor, ...] + offset: int = 0 + + @dataclass(frozen=True) class TraceSpec: """Shapes for one ``torch.export`` trace of a causal LM forward. @@ -568,6 +577,30 @@ def reference_inputs_as_args(self, reference_inputs: dict[str, Any]) -> tuple[An ) return tuple(reference_inputs.values()) + def new_eager_state( + self, *, dtype: torch.dtype, device: torch.device | str, max_seq_len: int | None = None + ) -> EagerState: + """Empty caches for one sequence of at most ``max_seq_len`` tokens (default: the + context).""" + with torch.device(device): + k_cache, v_cache = KVCache.create_cache_tensors( + self.config, dtype=dtype, seq_len=max_seq_len + ) + return EagerState((k_cache, v_cache)) + + def eager_forward_args(self, input_ids: torch.Tensor, state: EagerState) -> tuple[Any, ...]: + """``forward``'s arguments for ``input_ids`` written at ``state.offset``.""" + position_ids = torch.arange( + state.offset + input_ids.shape[1], dtype=torch.int32, device=input_ids.device + ).unsqueeze(0) + return (input_ids, position_ids, *state.tensors) + + def eager_step(self, input_ids: torch.Tensor, state: EagerState) -> torch.Tensor: + """Logits ``[1, query_len, vocab]`` for ``input_ids``, advancing ``state`` past them.""" + logits = self(*self.eager_forward_args(input_ids, state)) + state.offset += input_ids.shape[1] + return logits.reshape(1, input_ids.shape[1], -1) + @classmethod def _get_reauthored_config( cls, @@ -877,6 +910,36 @@ def set_prefill_mode(self, prefill_mode: bool): # than on the top-level module. self.extend.prefill_mode = prefill_mode + @override + def new_eager_state( + self, *, dtype: torch.dtype, device: torch.device | str, max_seq_len: int | None = None + ) -> EagerState: + """iOS graphs index the cache and the mask over the whole context, so ``max_seq_len`` + is ignored.""" + with torch.device(device): + key_cache, value_cache = KVCacheHandler.get_kv_cache_from_hf(self.config, dtype=dtype) + return EagerState((key_cache, value_cache)) + + @override + def eager_forward_args(self, input_ids: torch.Tensor, state: EagerState) -> tuple[Any, ...]: + """iOS takes only the query's positions, the write offset, and an additive mask over + the context.""" + query_len, device = input_ids.shape[1], input_ids.device + position_ids = ( + torch.arange(state.offset, state.offset + query_len, device=device) + .to(torch.uint16) + .unsqueeze(0) + ) + in_step = torch.tensor([state.offset], dtype=torch.int32, device=device) + mask = causal_mask( + self.config.max_position_embeddings, + query_len, + state.offset, + dtype=state.tensors[0].dtype, + device=device, + ) + return (input_ids, position_ids, in_step, mask, *state.tensors) + # ------------------------------------------------------------------ # Export contract # diff --git a/python/src/coreai_models/primitives/ios/cache.py b/python/src/coreai_models/primitives/ios/cache.py index a80ae757..f50e50dd 100644 --- a/python/src/coreai_models/primitives/ios/cache.py +++ b/python/src/coreai_models/primitives/ios/cache.py @@ -186,3 +186,20 @@ def k_cache(self) -> torch.Tensor: @property def v_cache(self) -> torch.Tensor: return self._v_cache + + +def causal_mask( + max_context_length: int, + query_len: int, + offset: int, + *, + dtype: torch.dtype, + device: torch.device | str = "cpu", +) -> torch.Tensor: + """Additive mask in the iOS layout ``(1, key, 1, query)``: query ``i`` sees keys to + ``offset + i``. + """ + keys = torch.arange(max_context_length, device=device).view(1, -1, 1, 1) + queries = torch.arange(query_len, device=device).view(1, 1, 1, -1) + mask = torch.zeros(1, max_context_length, 1, query_len, dtype=dtype, device=device) + return mask.masked_fill(keys > offset + queries, float("-inf"))