Enable batched attention for LLM exports

python/src/coreai_models/_constants.py

@@ -71,3 +71,12 @@
 # the converter in RELEASE mode and embeds minimum debug information. The
 # --include-debug-info flag turns it on (DEBUG mode).
 DEFAULT_INCLUDE_DEBUG_INFO = False
+
+# Max of the dynamic batch dim for `--attention-mask batched` exports when
+# --dynamic-batch-size is not given.
+DEFAULT_DYNAMIC_MAX_BATCH_SIZE = 8
+
+# metadata.json schema version (default for all bundle kinds).
+METADATA_VERSION = "0.2"
+# Batched LLM bundles use 0.3 (adds the max_batch_size field).
+METADATA_VERSION_BATCHED = "0.3"

python/src/coreai_models/export/bundle.py

@@ -13,9 +13,9 @@
 
 from transformers import AutoTokenizer
 
-logger = logging.getLogger(__name__)
+from coreai_models._constants import METADATA_VERSION, METADATA_VERSION_BATCHED
 
-METADATA_VERSION = "0.2"
+logger = logging.getLogger(__name__)
 
 
 def bundle_llm_asset(
@@ -27,8 +27,9 @@ def bundle_llm_asset(
     tokenizer_model_id: str | None = None,
     drafter_name: str | None = None,
     speculative_config: dict[str, Any] | None = None,
+    max_batch_size: int = 1,
 ) -> None:
-    """Add tokenizer and metadata.json (0.2 schema) to an LLM bundle.
+    """Add tokenizer and metadata.json to an LLM bundle.
 
     Expects ``{name}.aimodel`` to already exist inside bundle_path.
     When *drafter_name* is set, ``{drafter_name}.aimodel`` must also exist;
@@ -44,6 +45,8 @@ def bundle_llm_asset(
         speculative_config: Runtime config for speculative decoding (e.g.
             ``{"num_draft_tokens": 5, "shared_embeddings": True}``).
             Written as the ``"speculative"`` key in metadata.json.
+        max_batch_size: Upper bound on the request batch the graph serves;
+            > 1 selects the 0.3 schema (adds the ``max_batch_size`` field).
     """
     tok_id = tokenizer_model_id or hf_model_id
     _write_tokenizer(bundle_path / "tokenizer", tok_id)
@@ -55,6 +58,7 @@ def bundle_llm_asset(
         name,
         drafter_name=drafter_name,
         speculative_config=speculative_config,
+        max_batch_size=max_batch_size,
     )
 
 
@@ -72,23 +76,29 @@ def _write_metadata(
     name: str,
     drafter_name: str | None = None,
     speculative_config: dict[str, Any] | None = None,
+    max_batch_size: int = 1,
 ) -> None:
     assets: dict[str, str] = {"main": f"{name}.aimodel"}
     if drafter_name is not None:
         assets["drafter"] = f"{drafter_name}.aimodel"
 
+    is_batched = max_batch_size > 1
+    language: dict[str, Any] = {
+        "tokenizer": hf_model_id,
+        "vocab_size": getattr(hf_config, "vocab_size", None),
+        "max_context_length": getattr(hf_config, "max_position_embeddings", None),
+        "embedded_tokenizer": True,
+        "function_map": {"main": ["main"]},
+    }
+    if is_batched:
+        language["max_batch_size"] = max_batch_size
+
     metadata: dict[str, Any] = {
-        "metadata_version": METADATA_VERSION,
+        "metadata_version": METADATA_VERSION_BATCHED if is_batched else METADATA_VERSION,
         "kind": "llm",
         "name": name,
         "assets": assets,
-        "language": {
-            "tokenizer": hf_model_id,
-            "vocab_size": getattr(hf_config, "vocab_size", None),
-            "max_context_length": getattr(hf_config, "max_position_embeddings", None),
-            "embedded_tokenizer": True,
-            "function_map": {"main": ["main"]},
-        },
+        "language": language,
         "source": {
             "model_definition": "torch",
             "hf_model_id": hf_model_id,

python/src/coreai_models/export/macos.py

@@ -45,22 +45,35 @@
 
 logger = logging.getLogger(__name__)
 
+#: Batch size a *dynamic* batch graph is traced at. Must be >= 2 (torch.export specializes a
+#: size-1 dim) and strictly below the declared max, so the batch axis stays symbolic.
+_DYNAMIC_BATCH_TRACE_SIZE = 2
+
 
 def _build_reference_inputs(
     model: BaseForCausalLM,
     config,
     target_dtype: torch.dtype,
     max_context_length: int,
+    dynamic_max_batch_size: int | None = None,
 ) -> tuple[dict[str, Any], dict]:
     """Reference inputs and dynamic shapes for macOS export.
 
     Thin wrapper over the model's export-contract hooks, where the per-model variation
     lives. Returns ``(reference_inputs, dynamic_shapes)``.
+
+    ``dynamic_max_batch_size``, when set, declares a dynamic ``batch`` Dim up to that max,
+    tracing at the minimum valid size (2) so torch.export does not specialize the batch
+    axis; otherwise the model traces single-batch (the single-sequence contract).
     """
+    # Trace at the minimum valid batch for a dynamic batch dim; else single-batch.
+    trace_batch = _DYNAMIC_BATCH_TRACE_SIZE if dynamic_max_batch_size is not None else 1
     # The trace cache length only bounds peak memory, so cap it at the context it serves.
     spec = TraceSpec(
         max_context_length=max_context_length,
         cache_seq_len=min(TRACE_KV_CACHE_SEQ_LEN, max_context_length),
+        batch_size=trace_batch,
+        max_batch_size=dynamic_max_batch_size,
     )
     reference_inputs = model.build_reference_inputs(config, target_dtype, spec)
     dynamic_shapes = model.build_dynamic_shapes(config, spec)
@@ -406,6 +419,8 @@ def export_macos_model(
     if max_context_length is None:
         max_context_length = getattr(config, "max_position_embeddings", 2048)
 
+    dynamic_max_batch_size = getattr(export_config, "dynamic_max_batch_size", None)
+
     # Graph-mode quantization flattens the model into a torch.fx.GraphModule, which
     # carries none of the export-contract hooks. `externalized_model` is the eager
     # module that graph was captured from, so query the contract there.
@@ -420,7 +435,11 @@ def export_macos_model(
     )
 
     reference_inputs, dynamic_shapes = _build_reference_inputs(
-        contract_model, config, target_dtype, max_context_length
+        contract_model,
+        config,
+        target_dtype,
+        max_context_length,
+        dynamic_max_batch_size,
     )
 
     reference_inputs = _retype_inputs_to_quantized_graph(reference_inputs, model)

python/src/coreai_models/export/pipeline.py

@@ -17,13 +17,15 @@
 import re
 import tempfile
 from dataclasses import dataclass, field
+from enum import Enum
 from pathlib import Path
 from typing import Any, Literal, NamedTuple
 
 import torch
 from transformers import AutoConfig, AutoTokenizer
 
 from coreai_models._constants import (
+    DEFAULT_DYNAMIC_MAX_BATCH_SIZE,
     DEFAULT_INCLUDE_DEBUG_INFO,
     IOS_DEFAULT_MAX_CONTEXT_LENGTH,
     TRACE_KV_CACHE_SEQ_LEN,
@@ -50,6 +52,17 @@
 logger = logging.getLogger(__name__)
 
 
+class AttentionMaskMode(str, Enum):
+    """How the exported macOS graph handles the attention mask / per-row KV writes."""
+
+    #: Implicit ``is_causal`` mask (the shipping contract).
+    DEFAULT = "default"
+    #: Additive ``attn_mask`` graph input; enables the dense batched / ragged decode path.
+    BATCHED = "batched"
+    #: Reserved for the hardware paged-attention backend; not yet supported at export.
+    PAGED = "paged"
+
+
 @dataclass
 class ExportConfig:
     """Everything needed to export a model."""
@@ -78,12 +91,49 @@ class ExportConfig:
     model_type_override: str | None = None
     # Speculative decoding: export the drafter model alongside the target.
     with_drafter: bool = False
+    # Experimental, macOS only. Attention-mask / KV-write mode of the exported graph;
+    # see ``AttentionMaskMode`` for the per-value contract. Default keeps the implicit causal mask.
+    attention_mask_mode: AttentionMaskMode = AttentionMaskMode.DEFAULT
+    # Experimental, macOS only. Max of the dynamic ``batch`` dim across input_ids/position_ids/
+    # attn_mask and the KV cache. Defaults to 8 for batched exports; None for the causal mode.
+    dynamic_max_batch_size: int | None = None
+
+    @property
+    def emits_attn_mask(self) -> bool:
+        """True when the exported graph declares an additive ``attn_mask`` input."""
+        return self.attention_mask_mode is not AttentionMaskMode.DEFAULT
 
     def __post_init__(self) -> None:
+        if self.attention_mask_mode is AttentionMaskMode.PAGED:
+            raise NotImplementedError(
+                "The 'paged' attention-mask option is not available: it is reserved for the "
+                "hardware paged-attention backend, which does not ship yet. "
+                "Use --attention-mask default or --attention-mask batched."
+            )
         if self.quantization_mode == "graph" and self.variant != "macOS":
             raise ValueError(
                 f"quantization_mode='graph' is macOS only (got variant '{self.variant}')."
             )
+        if self.emits_attn_mask and self.variant != "macOS":
+            raise ValueError(
+                f"attention_mask_mode={self.attention_mask_mode.value!r} is macOS only "
+                f"(got variant '{self.variant}'). iOS has its own causal-mask export contract."
+            )
+        # Batched export implies a dynamic batch dim; default its max so `--attention-mask batched`
+        # alone yields a batch-capable graph. An explicit --dynamic-batch-size overrides.
+        if self.emits_attn_mask and self.dynamic_max_batch_size is None:
+            self.dynamic_max_batch_size = DEFAULT_DYNAMIC_MAX_BATCH_SIZE
+        if self.dynamic_max_batch_size is not None:
+            # iOS is already rejected above: dynamic batch requires attn_mask, which is macOS only.
+            if self.dynamic_max_batch_size < 2:
+                raise ValueError(
+                    f"dynamic_max_batch_size must be >= 2 (got {self.dynamic_max_batch_size})."
+                )
+            if not self.emits_attn_mask:
+                raise ValueError(
+                    "dynamic batch export requires attention_mask_mode=BATCHED; a batched graph "
+                    "without an attn_mask input cannot be served."
+                )
 
     @classmethod
     def from_preset(
@@ -286,6 +336,11 @@ def build_model(
 
     logger.info(f"Loading {config.hf_model_id} ({config.variant}, dtype={target_dtype})...")
     use_memory_efficient = config.variant == "macOS"
+    # Only forward `use_attention_mask` when a batched export actually needs it. In default mode we
+    # omit it entirely so every model's `from_hf*` override keeps working unchanged (gpt_oss,
+    # muse_glimmer, gemma3n, ... override the factory with signatures that predate this flag); this
+    # keeps default-mode export byte-identical to before the attention-mask work.
+    attn_mask_kw = {"use_attention_mask": True} if config.emits_attn_mask else {}
     if use_memory_efficient:
         layer_mmap_dir = None
         if temp_dir is not None:
@@ -299,6 +354,7 @@ def build_model(
             num_layers=config.num_layers,
             hf_config_attr=entry.hf_config_attr,
             hf_state_dict_prefix=entry.hf_state_dict_prefix,
+            **attn_mask_kw,
         )
     else:
         model = model_class.from_hf(
@@ -307,6 +363,7 @@ def build_model(
             target_dtype=target_dtype,
             num_layers=config.num_layers,
             disable_embedding_quantization=config.disable_embedding_quantization,
+            **attn_mask_kw,
         )
     model = model.eval()
     torch_quantization_config = resolved.quantization_config
@@ -548,6 +605,7 @@ async def _async_export_model(config: ExportConfig) -> str:
             tokenizer_model_id=entry.tokenizer_model_id,
             drafter_name=drafter_name,
             speculative_config=entry.drafter_config if drafter_name else None,
+            max_batch_size=config.dynamic_max_batch_size or 1,
         )
 
     logger.info(f"Export complete: {bundle_path}")

python/src/coreai_models/llm/export.py

@@ -20,7 +20,7 @@
 from coreai_models.diffusion.presets import (
     PRESETS as DIFFUSION_PRESETS,
 )
-from coreai_models.export.pipeline import ExportConfig, export_model
+from coreai_models.export.pipeline import AttentionMaskMode, ExportConfig, export_model
 from coreai_models.export.presets import ALL_PRESET_NAMES, IOS_PRESETS, MACOS_PRESETS, list_presets
 from coreai_models.export.presets import (
     DEFAULT_IOS_COMPRESSION_PRESET as IOS_DEFAULT,
@@ -182,6 +182,28 @@ def build_parser() -> argparse.ArgumentParser:
             "The drafter is looked up from the model registry; not all models have one."
         ),
     )
+    parser.add_argument(
+        "--attention-mask",
+        choices=["default", "batched", "paged"],
+        default="default",
+        help=(
+            "Experimental, macOS only. Attention-mask / KV-write mode of the exported graph. "
+            "'default' keeps the implicit causal mask (shipping contract); 'batched' emits an "
+            "explicit additive attn_mask input (is_causal off) for the dense batched/ragged "
+            "decoder; 'paged' is reserved for the hardware paged backend (not yet supported). "
+            "Default: default. Non-default is rejected when --platform is iOS."
+        ),
+    )
+    parser.add_argument(
+        "--dynamic-batch-size",
+        type=int,
+        default=None,
+        help=(
+            "Experimental, macOS only. Declare a DYNAMIC batch dim up to this max (traced at 2) "
+            "across input_ids/position_ids/attn_mask and the KV cache. Must be >= 2. Pairs with "
+            "--attention-mask batched."
+        ),
+    )
     return parser
 
 
@@ -400,6 +422,8 @@ def _resolve_export_config(args: argparse.Namespace) -> ExportConfig:
         registry_compression_config=registry_compression_config,
     )
 
+    attention_mask_mode = AttentionMaskMode(args.attention_mask)
+
     return ExportConfig(
         hf_model_id=hf_model_id,
         variant=variant,
@@ -416,6 +440,8 @@ def _resolve_export_config(args: argparse.Namespace) -> ExportConfig:
         include_debug_info=args.include_debug_info,
         model_type_override=getattr(preset, "_model_type_override", None) if preset else None,
         with_drafter=args.with_drafter,
+        attention_mask_mode=attention_mask_mode,
+        dynamic_max_batch_size=args.dynamic_batch_size,
     )