From 4f67e7c1b52fc94034a2c952a271ee26cb6ca77d Mon Sep 17 00:00:00 2001 From: Sukru Tikves Date: Thu, 8 Oct 2026 14:46:20 -0700 Subject: [PATCH] Add opt-in batched / attention-mask LLM export with conditional 0.3 metadata Introduce opt-in macOS attention-mask / dynamic-batch export for LLMs: - Add the --attention-mask and --dynamic-batch-size CLI flags, the AttentionMaskMode export mode, and the use_attention_mask model plumbing (default off, so existing exports are byte-identical; masked-vs-causal parity is covered by test_qwen3_attention_mask). - Make the LLM metadata schema conditional: a batched bundle (max_batch_size > 1, exported with --attention-mask batched) stamps metadata_version 0.3 and carries the max_batch_size field, while every single-sequence bundle stays on 0.2 and omits the field. Other bundle kinds (diffusion/segmentation/VLM) are untouched and remain on 0.2, so current runtimes keep running existing models. First stage of the batched model export and serving work; the runtime that consumes 0.3 bundles lands in a later stage. --- python/src/coreai_models/_constants.py | 5 + python/src/coreai_models/export/bundle.py | 32 ++- python/src/coreai_models/export/macos.py | 21 +- python/src/coreai_models/export/pipeline.py | 62 ++++++ python/src/coreai_models/llm/export.py | 28 ++- python/src/coreai_models/models/base.py | 109 ++++++++-- .../src/coreai_models/models/macos/qwen3.py | 35 +++- .../coreai_models/primitives/macos/cache.py | 13 +- .../test_export/test_bundle_metadata.py | 59 ++++++ .../test_export/test_export_contract.py | 190 ++++++++++++++++++ .../test_export/test_use_attention_mask.py | 86 ++++++++ .../test_qwen3_attention_mask.py | 85 ++++++++ 12 files changed, 684 insertions(+), 41 deletions(-) create mode 100644 python/tests/test_model_units/test_export/test_bundle_metadata.py create mode 100644 python/tests/test_model_units/test_export/test_use_attention_mask.py create mode 100644 python/tests/test_model_units/test_models/test_macos_layers/test_qwen3_attention_mask.py diff --git a/python/src/coreai_models/_constants.py b/python/src/coreai_models/_constants.py index b515b73a..00f5c84c 100644 --- a/python/src/coreai_models/_constants.py +++ b/python/src/coreai_models/_constants.py @@ -67,3 +67,8 @@ # 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 + +# 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" diff --git a/python/src/coreai_models/export/bundle.py b/python/src/coreai_models/export/bundle.py index 65571f56..fc811b4d 100644 --- a/python/src/coreai_models/export/bundle.py +++ b/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, diff --git a/python/src/coreai_models/export/macos.py b/python/src/coreai_models/export/macos.py index d9f25ce5..07ffbb1b 100644 --- a/python/src/coreai_models/export/macos.py +++ b/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) diff --git a/python/src/coreai_models/export/pipeline.py b/python/src/coreai_models/export/pipeline.py index f1f43ba9..c78cd303 100644 --- a/python/src/coreai_models/export/pipeline.py +++ b/python/src/coreai_models/export/pipeline.py @@ -17,6 +17,7 @@ import re import tempfile from dataclasses import dataclass, field +from enum import Enum from pathlib import Path from typing import Any, Literal @@ -49,6 +50,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.""" @@ -77,12 +89,53 @@ 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 to 8 so `--attention-mask + # batched` alone yields a batch-8-capable graph (serves 1..8). An explicit + # --dynamic-batch-size overrides. + if self.emits_attn_mask and self.dynamic_max_batch_size is None: + self.dynamic_max_batch_size = 8 + if self.dynamic_max_batch_size is not None: + if self.variant != "macOS": + raise ValueError( + f"dynamic_max_batch_size is macOS only (got variant '{self.variant}')." + ) + 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." + ) def _generate_output_name(config: ExportConfig) -> str: @@ -216,6 +269,12 @@ async def _async_export_model(config: ExportConfig) -> str: else contextlib.nullcontext(None) ) + # 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 {} + 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 @@ -229,6 +288,7 @@ async def _async_export_model(config: ExportConfig) -> str: 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( @@ -237,6 +297,7 @@ async def _async_export_model(config: ExportConfig) -> str: target_dtype=target_dtype, num_layers=config.num_layers, disable_embedding_quantization=config.disable_embedding_quantization, + **attn_mask_kw, ) model = model.eval() # ---- 3. Resolve compression preset ---- @@ -434,6 +495,7 @@ def get_calibration_data(): # type: ignore[no-untyped-def] 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}") diff --git a/python/src/coreai_models/llm/export.py b/python/src/coreai_models/llm/export.py index 0d82a966..7640f4d0 100644 --- a/python/src/coreai_models/llm/export.py +++ b/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 @@ -385,6 +407,8 @@ def _resolve_export_config(args: argparse.Namespace) -> ExportConfig: f"preset. Available: {list_presets()}" ) + attention_mask_mode = AttentionMaskMode(args.attention_mask) + return ExportConfig( hf_model_id=hf_model_id, variant=variant, @@ -401,6 +425,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, ) diff --git a/python/src/coreai_models/models/base.py b/python/src/coreai_models/models/base.py index 2e064405..995aa4a8 100644 --- a/python/src/coreai_models/models/base.py +++ b/python/src/coreai_models/models/base.py @@ -78,6 +78,14 @@ class TraceSpec: cache_seq_len: int = TRACE_KV_CACHE_SEQ_LEN query_len: int = QUANT_TRACE_QUERY_LEN offset: int = QUANT_TRACE_OFFSET + #: Request-batch size the graph is traced at. When ``max_batch_size`` is None this also + #: pins dim0 (the single-batch / static-batch contract). + batch_size: int = 1 + #: Upper bound for a *dynamic* ``batch`` dim across input_ids/position_ids/attn_mask and the + #: cache's leading batch dim (index 1). None pins the batch to ``batch_size`` (the contract). + #: When set, the graph is traced at ``batch_size`` (which must be >= 2 and strictly below this + #: max) so torch.export does not specialize the batch axis to the trace size. + max_batch_size: int | None = None def __post_init__(self) -> None: # Below this there is no legal `Dim(min=query_len, max=max_context_length - 1)`. @@ -92,6 +100,13 @@ def __post_init__(self) -> None: f"cache_seq_len = {self.cache_seq_len}, " f"max_context_length = {self.max_context_length}" ) + if self.max_batch_size is not None and not 2 <= self.batch_size < self.max_batch_size: + # torch.export specializes a size-1 dim and rejects a dim whose only legal value is + # the trace size, so the trace batch must be >= 2 and strictly below the dynamic max. + raise ValueError( + "max_batch_size requires a trace batch_size in [2, max_batch_size); received " + f"batch_size={self.batch_size}, max_batch_size={self.max_batch_size}." + ) @property def caches_are_static(self) -> bool: @@ -314,6 +329,11 @@ class BaseForCausalLM(torch.nn.Module): #: Set while the exporter traces the prefill graph. See :meth:`set_prefill_mode`. prefill_mode: bool = False + #: Opt in to exporting with an explicit additive ``attn_mask`` graph input + #: (``is_causal`` off) instead of the implicit causal mask. Default off keeps the + #: standard causal export contract unchanged. + use_attention_mask: bool = False + def set_prefill_mode(self, prefill_mode: bool) -> None: """Toggle prefill mode for the next trace. @@ -360,7 +380,9 @@ def wrapper(*args, **kwargs): return wrapper - def __init__(self: Self, config, model_device: str = "cpu") -> None: + def __init__( + self: Self, config, model_device: str = "cpu", use_attention_mask: bool = False + ) -> None: """Initialize the model using template method pattern. Initializing the model on the meta device allows us to avoid @@ -370,9 +392,13 @@ def __init__(self: Self, config, model_device: str = "cpu") -> None: config: Model configuration object model_device: Device to use for initializing model components (e.g., "cpu" or "meta") + use_attention_mask: Build attention with an explicit additive mask + input (``is_causal`` off) rather than the implicit causal + mask. Default off preserves the standard causal contract. """ super().__init__() self.config = config + self.use_attention_mask = use_attention_mask with torch.device(model_device): self._init_model(config) @@ -411,10 +437,15 @@ def _mutate_state_dict(self: Self, state_dict: dict[str, torch.Tensor]) -> None: # value_cache). # ------------------------------------------------------------------ - @classmethod - def export_input_names(cls) -> dict[str, tuple[str, ...]]: - """Graph input names per graph, in relative order among the non-state args.""" - return {MAIN_GRAPH_NAME: ("input_ids", "position_ids")} + def export_input_names(self) -> dict[str, tuple[str, ...]]: + """Graph input names per graph, in relative order among the non-state args. + + Opting into :attr:`use_attention_mask` appends an additive ``attn_mask`` input. + """ + names = ("input_ids", "position_ids") + if self.use_attention_mask: + names = (*names, "attn_mask") + return {MAIN_GRAPH_NAME: names} @classmethod def export_state_names(cls) -> dict[str, tuple[str, ...]]: @@ -455,16 +486,19 @@ def build_reference_inputs( parameters in *exact* signature order. Pass ``spec`` to :meth:`build_dynamic_shapes` too, so the tensors and their dims cannot disagree. """ - input_ids = torch.randint(1, config.vocab_size, (1, spec.query_len), dtype=torch.int32) + input_ids = torch.randint( + 1, config.vocab_size, (spec.batch_size, spec.query_len), dtype=torch.int32 + ) position_ids = ( torch.arange(spec.query_len + spec.offset, dtype=torch.int32) .unsqueeze(0) - .expand(1, spec.query_len + spec.offset) + .expand(spec.batch_size, spec.query_len + spec.offset) + .contiguous() ) k_cache, v_cache = KVCache.create_cache_tensors( - config, dtype=target_dtype, seq_len=spec.cache_seq_len + config, dtype=target_dtype, seq_len=spec.cache_seq_len, batch_size=spec.batch_size ) - return { + inputs: dict[str, dict[str, Any]] = { MAIN_GRAPH_NAME: { "input_ids": input_ids, "position_ids": position_ids, @@ -472,6 +506,16 @@ def build_reference_inputs( "v_cache": v_cache, } } + if self.use_attention_mask: + # Additive mask broadcast over heads onto the per-row score matrix + # [B, heads, query_len, key_len]. The key length is the fetched cache + # prefix, which equals the position_ids length (query_len + offset). + # Zeros for the trace (shape only); LAST so the dict stays a contiguous + # in-order prefix of ``forward``. + inputs[MAIN_GRAPH_NAME]["attn_mask"] = torch.zeros( + spec.batch_size, 1, spec.query_len, spec.query_len + spec.offset, dtype=target_dtype + ) + return inputs def build_dynamic_shapes(self, config, spec: TraceSpec) -> dict[str, Any]: """``dynamic_shapes`` per graph, matching :meth:`build_reference_inputs`. @@ -479,9 +523,11 @@ def build_dynamic_shapes(self, config, spec: TraceSpec) -> dict[str, Any]: Keyed like the reference inputs; ``None`` pins that input to its traced shape. """ max_ctx = spec.max_context_length + seq_ids_dim = torch.export.Dim("seq_ids", max=max_ctx - 2) + seq_pos_dim = torch.export.Dim("seq_pos", min=spec.query_len, max=max_ctx - 1) shapes: dict[str, Any] = { - "input_ids": {1: torch.export.Dim("seq_ids", max=max_ctx - 2)}, - "position_ids": {1: torch.export.Dim("seq_pos", min=spec.query_len, max=max_ctx - 1)}, + "input_ids": {1: seq_ids_dim}, + "position_ids": {1: seq_pos_dim}, } seq_dim = KVCache.seq_len_dim() if spec.caches_are_static: @@ -494,6 +540,29 @@ def build_dynamic_shapes(self, config, spec: TraceSpec) -> dict[str, Any]: shapes["v_cache"] = { seq_dim: torch.export.Dim("v_seq_len", min=spec.cache_seq_len, max=max_ctx) } + if self.use_attention_mask: + # Mask query axis == input_ids; key axis == the attention key length, + # which is the fetched cache prefix tied to position_ids (always + # dynamic, even when the cache tensor dims are pinned). + shapes["attn_mask"] = {2: seq_ids_dim, 3: seq_pos_dim} + if spec.max_batch_size is None: + # No dynamic batch requested: dim0 (ids/pos/mask) and cache dim1 stay PINNED to + # spec.batch_size (the single-batch / static-batch contract). + return {MAIN_GRAPH_NAME: shapes} + # One shared ``batch`` Dim across every batch axis so torch.export ties them together: + # dim0 of ids/pos/mask and dim1 of the caches. Traced at spec.batch_size (>= 2, strictly + # below max), so the batch axis is not specialized to the trace size. + batch_dim = torch.export.Dim("batch", min=1, max=spec.max_batch_size) + shapes["input_ids"][0] = batch_dim + shapes["position_ids"][0] = batch_dim + if "attn_mask" in shapes: + shapes["attn_mask"][0] = batch_dim + # Caches carry the batch Dim on dim1 even when the seq dim is pinned (static cache): + # replace the None pin with a batch-only shape dict. + for cache_name in ("k_cache", "v_cache"): + cache_shape = shapes[cache_name] or {} + cache_shape[1] = batch_dim + shapes[cache_name] = cache_shape return {MAIN_GRAPH_NAME: shapes} def validate_export_contract( @@ -609,6 +678,7 @@ def from_hf( mmap_path: str | None = None, num_layers: int | None = None, disable_embedding_quantization: bool = False, + use_attention_mask: bool = False, ) -> T: """Load model from HuggingFace model hub. @@ -625,6 +695,9 @@ def from_hf( disable_embedding_quantization: iOS only. When True, the embedding table is not quantized to int8. Ignored for macOS model classes. + use_attention_mask: macOS only. When True, build the model with an + explicit additive ``attn_mask`` input (``is_causal`` off). + Ignored for iOS model classes. Returns: Instance of the model class loaded with HuggingFace weights @@ -643,10 +716,13 @@ def from_hf( ) # Create our model instance and load the state dict. - # disable_embedding_quantization is only accepted by the iOS base class. + # disable_embedding_quantization is only accepted by the iOS base class; + # use_attention_mask only by the macOS base class. init_kwargs: dict = {"config": config, "model_device": "meta"} if issubclass(cls, BaseForCausalLMForiOS): init_kwargs["disable_embedding_quantization"] = disable_embedding_quantization + else: + init_kwargs["use_attention_mask"] = use_attention_mask model = cls(**init_kwargs) model.to(dtype=target_dtype) state_dict = hf_model.state_dict() @@ -691,6 +767,7 @@ def from_hf_memory_efficient( hf_config_attr: str | None = None, hf_state_dict_prefix: str = "", disable_embedding_quantization: bool = False, + use_attention_mask: bool = False, ) -> T: """Load model from HuggingFace with layer-by-layer memory offloading. @@ -719,6 +796,9 @@ def from_hf_memory_efficient( disable_embedding_quantization: iOS only. When True, the embedding table is not quantized to int8. Ignored for non-iOS model classes. + use_attention_mask: macOS only. When True, build the model with an + explicit additive ``attn_mask`` input (``is_causal`` off). + Ignored for iOS model classes. """ model_dir = snapshot_download( huggingface_model_id, @@ -730,10 +810,13 @@ def from_hf_memory_efficient( config = cls._get_reauthored_config(hf_config, max_context_length, num_layers=num_layers) - # disable_embedding_quantization is only accepted by the iOS base class. + # disable_embedding_quantization is only accepted by the iOS base class; + # use_attention_mask only by the macOS base class. init_kwargs: dict = {"config": config, "model_device": "meta"} if issubclass(cls, BaseForCausalLMForiOS): init_kwargs["disable_embedding_quantization"] = disable_embedding_quantization + else: + init_kwargs["use_attention_mask"] = use_attention_mask model = cls(**init_kwargs) model.to(dtype=target_dtype) diff --git a/python/src/coreai_models/models/macos/qwen3.py b/python/src/coreai_models/models/macos/qwen3.py index 96cfc732..9c01c32b 100644 --- a/python/src/coreai_models/models/macos/qwen3.py +++ b/python/src/coreai_models/models/macos/qwen3.py @@ -25,7 +25,9 @@ class Attention(nn.Module): - def __init__(self, config: Qwen3Config, layer_idx: int) -> None: + def __init__( + self, config: Qwen3Config, layer_idx: int, use_attention_mask: bool = False + ) -> None: super().__init__() self.layer_idx = layer_idx @@ -47,7 +49,7 @@ def __init__(self, config: Qwen3Config, layer_idx: int) -> None: self.q_norm = RMSNorm(head_dim, eps=config.rms_norm_eps) self.k_norm = RMSNorm(head_dim, eps=config.rms_norm_eps) - self.sdpa = SDPA(is_causal=True) + self.sdpa = SDPA(is_causal=not use_attention_mask) assert is_default_rope_scaling(config), f"unsupported rope_scaling: {config.rope_scaling}" self.rope = initialize_rope(base=resolve_rope_theta(config)) @@ -56,6 +58,7 @@ def forward( x: torch.Tensor, position_ids: torch.IntTensor, cache: KVCache | None = None, + attn_mask: torch.Tensor | None = None, ) -> torch.Tensor: batch_size, query_len, _ = x.shape n_heads, n_kv_heads = self.n_heads, self.n_kv_heads @@ -101,7 +104,7 @@ def forward( ) output = ( - self.sdpa(query, key, value) + self.sdpa(query, key, value, attn_mask=attn_mask) .permute(0, 2, 1, 3) .reshape(batch_size, query_len, self.n_heads * self.head_dim) ) @@ -109,10 +112,14 @@ def forward( class TransformerBlock(nn.Module): - def __init__(self, config: Qwen3Config, layer_idx: int) -> None: + def __init__( + self, config: Qwen3Config, layer_idx: int, use_attention_mask: bool = False + ) -> None: super().__init__() hidden_size = config.hidden_size - self.self_attn = Attention(config, layer_idx=layer_idx) + self.self_attn = Attention( + config, layer_idx=layer_idx, use_attention_mask=use_attention_mask + ) self.mlp = MLP(hidden_size, config.intermediate_size) self.input_layernorm = RMSNorm(hidden_size, eps=config.rms_norm_eps) @@ -123,20 +130,24 @@ def forward( x: torch.Tensor, position_ids: torch.IntTensor, cache: KVCache | None = None, + attn_mask: torch.Tensor | None = None, ) -> torch.Tensor: - r = self.self_attn(self.input_layernorm(x), position_ids, cache) + r = self.self_attn(self.input_layernorm(x), position_ids, cache, attn_mask=attn_mask) h = x + r r = self.mlp(self.post_attention_layernorm(h)) return h + r class Qwen3Model(nn.Module): - def __init__(self, config: Qwen3Config) -> None: + def __init__(self, config: Qwen3Config, use_attention_mask: bool = False) -> None: super().__init__() hidden_size = config.hidden_size self.embed_tokens = nn.Embedding(config.vocab_size, hidden_size) self.layers = nn.ModuleList( - [TransformerBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + [ + TransformerBlock(config, layer_idx, use_attention_mask=use_attention_mask) + for layer_idx in range(config.num_hidden_layers) + ] ) self.norm = RMSNorm(hidden_size, eps=config.rms_norm_eps) @@ -145,10 +156,11 @@ def forward( input_ids: torch.Tensor, position_ids: torch.IntTensor = None, cache: KVCache | None = None, + attn_mask: torch.Tensor | None = None, ) -> torch.Tensor: h = self.embed_tokens(input_ids) for layer in self.layers: - h = layer(h, position_ids, cache) + h = layer(h, position_ids, cache, attn_mask=attn_mask) return self.norm(h) @@ -160,7 +172,7 @@ class Qwen3ForCausalLM(BaseForCausalLM): @override def _init_model(self, config: Qwen3Config) -> None: - self.model = Qwen3Model(config) + self.model = Qwen3Model(config, use_attention_mask=self.use_attention_mask) self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) if config.tie_word_embeddings: self.lm_head.weight = self.model.embed_tokens.weight @@ -172,9 +184,10 @@ def forward( position_ids: torch.IntTensor, k_cache: torch.Tensor, v_cache: torch.Tensor, + attn_mask: torch.Tensor | None = None, ) -> torch.Tensor | tuple: cache = KVCache(k_cache, v_cache) - out = self.model(input_ids, position_ids, cache) + out = self.model(input_ids, position_ids, cache, attn_mask=attn_mask) if self.prefill_mode: # A bare `return` causes torch export to trace a leaf node with value # `None` rather than having no leaf nodes whatsoever. Remedied with diff --git a/python/src/coreai_models/primitives/macos/cache.py b/python/src/coreai_models/primitives/macos/cache.py index ce1926d0..4b7fda83 100644 --- a/python/src/coreai_models/primitives/macos/cache.py +++ b/python/src/coreai_models/primitives/macos/cache.py @@ -37,6 +37,7 @@ def create_cache_tensors( config, dtype: torch.dtype = torch.float32, seq_len: int | None = None, + batch_size: int = 1, ) -> tuple[torch.Tensor, torch.Tensor]: """Create zero-initialized KV cache tensors from a model config. @@ -45,9 +46,10 @@ def create_cache_tensors( dtype: Cache dtype. seq_len: Sequence-dim length; defaults to ``config.max_position_embeddings``. Pass explicitly to build a trace-sized cache without mutating config. + batch_size: Leading request-batch dim (dim 1). Defaults to 1. Returns: - (k_cache, v_cache) of shape (n_layers, 1, n_kv_heads, max_seq_len, head_dim). + (k_cache, v_cache) of shape (n_layers, batch_size, n_kv_heads, max_seq_len, head_dim). """ n_kv_heads = config.num_key_value_heads n_layers = config.num_hidden_layers @@ -56,8 +58,8 @@ def create_cache_tensors( head_dim = config.head_dim else: head_dim = config.hidden_size // config.num_attention_heads - k_cache = torch.zeros(n_layers, 1, n_kv_heads, max_seq_len, head_dim, dtype=dtype) - v_cache = torch.zeros(n_layers, 1, n_kv_heads, max_seq_len, head_dim, dtype=dtype) + k_cache = torch.zeros(n_layers, batch_size, n_kv_heads, max_seq_len, head_dim, dtype=dtype) + v_cache = torch.zeros(n_layers, batch_size, n_kv_heads, max_seq_len, head_dim, dtype=dtype) return k_cache, v_cache @classmethod @@ -167,7 +169,10 @@ def update_and_fetch( end=torch.cat( [ layer_index_end, - torch.tensor((int(self._v_cache.size(1)),), dtype=torch.int32, device=device), + # Mirror the K path (no ``int()``): wrapping the batch dim in ``int()`` + # collapses a symbolic batch size to the trace constant, which breaks a + # dynamic-batch export. size(2)/size(4) are genuinely static. + torch.tensor((self._v_cache.size(1),), dtype=torch.int32, device=device), torch.tensor((int(self._v_cache.size(2)),), dtype=torch.int32, device=device), torch.tensor((offset + v.size(-2),), dtype=torch.int32, device=device), torch.tensor((int(self._v_cache.size(4)),), dtype=torch.int32, device=device), diff --git a/python/tests/test_model_units/test_export/test_bundle_metadata.py b/python/tests/test_model_units/test_export/test_bundle_metadata.py new file mode 100644 index 00000000..e4829323 --- /dev/null +++ b/python/tests/test_model_units/test_export/test_bundle_metadata.py @@ -0,0 +1,59 @@ +# 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 + +"""LLM bundle metadata_version contract: 0.3 (with max_batch_size) for batched bundles, +0.2 for everything else.""" + +import json +from pathlib import Path +from types import SimpleNamespace + +from coreai_models._constants import METADATA_VERSION, METADATA_VERSION_BATCHED +from coreai_models.export.bundle import _write_metadata + + +def test_llm_writer_defaults_to_0_2(tmp_path: Path) -> None: + """Omitted max_batch_size => single-sequence => 0.2 schema.""" + cfg = SimpleNamespace(vocab_size=256, max_position_embeddings=4096) + _write_metadata(tmp_path, "org/model", cfg, "4bit", "model") + meta = json.loads((tmp_path / "metadata.json").read_text()) + assert meta["metadata_version"] == "0.2" + + +def test_llm_writer_omits_max_batch_size_when_single_sequence(tmp_path: Path) -> None: + """A 0.2 bundle does not carry max_batch_size, keeping its output byte-identical.""" + cfg = SimpleNamespace(vocab_size=256, max_position_embeddings=4096) + _write_metadata(tmp_path, "org/model", cfg, "4bit", "model") + meta = json.loads((tmp_path / "metadata.json").read_text()) + assert "max_batch_size" not in meta["language"] + + +def test_llm_writer_emits_0_3_when_batched(tmp_path: Path) -> None: + """A dynamic-batch export (max_batch_size > 1) stamps the 0.3 schema.""" + cfg = SimpleNamespace(vocab_size=256, max_position_embeddings=4096) + _write_metadata(tmp_path, "org/model", cfg, "4bit", "model", max_batch_size=8) + meta = json.loads((tmp_path / "metadata.json").read_text()) + assert meta["metadata_version"] == "0.3" + + +def test_llm_writer_stamps_max_batch_size_when_batched(tmp_path: Path) -> None: + """A batched export records its max batch so the runtime can reject larger requests.""" + cfg = SimpleNamespace(vocab_size=256, max_position_embeddings=4096) + _write_metadata(tmp_path, "org/model", cfg, "4bit", "model", max_batch_size=8) + meta = json.loads((tmp_path / "metadata.json").read_text()) + assert meta["language"]["max_batch_size"] == 8 + + +def test_version_constants() -> None: + """The default schema is 0.2; batched LLM bundles use 0.3.""" + assert METADATA_VERSION == "0.2" + assert METADATA_VERSION_BATCHED == "0.3" + + +def test_other_bundle_kinds_stay_on_0_2() -> None: + """Conditional versioning leaves diffusion (and the other kinds) on 0.2.""" + from coreai_models.diffusion.pipeline import METADATA_VERSION as DIFFUSION_VERSION + + assert DIFFUSION_VERSION == "0.2" diff --git a/python/tests/test_model_units/test_export/test_export_contract.py b/python/tests/test_model_units/test_export/test_export_contract.py index e3a97098..5bcff0f2 100644 --- a/python/tests/test_model_units/test_export/test_export_contract.py +++ b/python/tests/test_model_units/test_export/test_export_contract.py @@ -300,6 +300,196 @@ def forward(self, input_ids, position_ids, k_cache, v_cache): assert len(m.reference_inputs_as_args(self._graph(m, config))) == 4 +class _MaskLM(_StandardLM): + """Default contract plus an opt-in additive ``attn_mask`` input. + + ``forward`` takes ``attn_mask`` after the KV pair; the extra graph input only + appears when the model is built with ``use_attention_mask=True``. + """ + + @override + def forward( + self, + input_ids: torch.Tensor, + position_ids: torch.IntTensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + attn_mask: torch.Tensor = None, + ) -> torch.Tensor: + raise NotImplementedError("shape contract only; never traced in these tests") + + +class TestMaskContract: + """The opt-in additive ``attn_mask`` graph input, gated on ``use_attention_mask``.""" + + def _built(self, model, config, spec=None): + spec = spec or TraceSpec(max_context_length=MAX_CONTEXT_LENGTH) + return ( + model.build_reference_inputs(config, torch.float16, spec), + model.build_dynamic_shapes(config, spec), + ) + + def test_mask_mode_declares_attn_mask_input(self, config) -> None: + m = _MaskLM(config, use_attention_mask=True) + assert m.export_input_names()[MAIN] == ("input_ids", "position_ids", "attn_mask") + + def test_mask_mode_appends_attn_mask_last(self, config) -> None: + m = _MaskLM(config, use_attention_mask=True) + refs, _ = self._built(m, config) + assert list(refs[MAIN]) == [ + "input_ids", + "position_ids", + "k_cache", + "v_cache", + "attn_mask", + ] + mask = refs[MAIN]["attn_mask"] + # Scores are [B, heads, query_len, key_len]; the key length is the fetched + # cache prefix, which equals the position_ids length (query_len + offset). + assert mask.shape == ( + 1, + 1, + QUANT_TRACE_QUERY_LEN, + QUANT_TRACE_QUERY_LEN + QUANT_TRACE_OFFSET, + ) + assert mask.dtype == torch.float16 + + def test_mask_mode_dynamic_shapes_reuse_dims(self, config) -> None: + m = _MaskLM(config, use_attention_mask=True) + _, shapes = self._built(m, config) + graph = shapes[MAIN] + assert "attn_mask" in graph + # Query axis is the same dim as input_ids; key axis is the attention key + # length, which is tied to position_ids (not the k_cache tensor's seq dim). + assert graph["attn_mask"][2] is graph["input_ids"][1] + assert graph["attn_mask"][3] is graph["position_ids"][1] + + def test_mask_mode_validates_and_converts_to_args(self, config) -> None: + m = _MaskLM(config, use_attention_mask=True) + refs, shapes = self._built(m, config) + m.validate_export_contract(refs, shapes) + args = m.reference_inputs_as_args(refs[MAIN]) + assert len(args) == 5 + + def test_mask_mode_static_cache_still_ties_key_to_position_ids(self, config) -> None: + config.max_position_embeddings = TRACE_KV_CACHE_SEQ_LEN + spec = TraceSpec( + max_context_length=TRACE_KV_CACHE_SEQ_LEN, cache_seq_len=TRACE_KV_CACHE_SEQ_LEN + ) + assert spec.caches_are_static + m = _MaskLM(config, use_attention_mask=True) + refs, shapes = self._built(m, config, spec) + # Caches are pinned, but position_ids stays dynamic, so the mask's key axis + # does too -- it tracks the attention key length, not the cache tensor. + assert shapes[MAIN]["k_cache"] is None + assert shapes[MAIN]["attn_mask"] == { + 2: shapes[MAIN]["input_ids"][1], + 3: shapes[MAIN]["position_ids"][1], + } + m.validate_export_contract(refs, shapes) + + def test_flag_off_leaves_the_contract_unchanged(self, config) -> None: + m = _MaskLM(config) # flag defaults off + refs, shapes = self._built(m, config) + assert m.export_input_names()[MAIN] == ("input_ids", "position_ids") + assert "attn_mask" not in refs[MAIN] + assert "attn_mask" not in shapes[MAIN] + m.validate_export_contract(refs, shapes) + + def test_batched_contract_traces_at_batch_size_with_dim0_pinned(self, config) -> None: + """batch_size>1 traces reference tensors at B with dim0/cache-dim1 PINNED (static batch).""" + spec = TraceSpec(max_context_length=MAX_CONTEXT_LENGTH, batch_size=2) + m = _MaskLM(config, use_attention_mask=True) + refs = m.build_reference_inputs(config, torch.float16, spec)[MAIN] + shapes = m.build_dynamic_shapes(config, spec)[MAIN] + + # Reference tensors carry the fixed batch on dim0 (ids/pos/mask) and cache dim1. + assert refs["input_ids"].shape[0] == 2 + assert refs["position_ids"].shape[0] == 2 + assert refs["attn_mask"].shape[0] == 2 + assert refs["k_cache"].shape[1] == 2 # cache batch dim is index 1 + assert refs["v_cache"].shape[1] == 2 + + # The batch dim is PINNED (static), not a dynamic Dim: dim0 / cache-dim1 absent from shapes. + assert shapes["input_ids"].get(0) is None + assert shapes["position_ids"].get(0) is None + assert shapes["attn_mask"].get(0) is None + assert (shapes["k_cache"] or {}).get(1) is None + assert (shapes["v_cache"] or {}).get(1) is None + m.validate_export_contract({MAIN: refs}, {MAIN: shapes}) + + def test_batch_size_one_traces_single_batch(self, config) -> None: + """Default batch_size=1 keeps the validated single-batch contract untouched.""" + spec = TraceSpec(max_context_length=MAX_CONTEXT_LENGTH) # batch_size defaults to 1 + m = _MaskLM(config, use_attention_mask=True) + refs = m.build_reference_inputs(config, torch.float16, spec)[MAIN] + shapes = m.build_dynamic_shapes(config, spec)[MAIN] + assert refs["input_ids"].shape[0] == 1 + assert shapes["input_ids"].get(0) is None # dim0 pinned, not dynamic + assert shapes["k_cache"].get(1) is None # cache batch dim pinned + + def test_dynamic_batch_declares_one_shared_batch_dim(self, config) -> None: + """max_batch_size declares ONE dynamic ``batch`` Dim on dim0 of ids/pos/mask and + dim1 of the caches, while reference tensors still trace at the smaller batch_size + (torch.export rejects a dim whose only legal value is the trace size).""" + spec = TraceSpec(max_context_length=MAX_CONTEXT_LENGTH, batch_size=2, max_batch_size=8) + m = _MaskLM(config, use_attention_mask=True) + refs = m.build_reference_inputs(config, torch.float16, spec)[MAIN] + shapes = m.build_dynamic_shapes(config, spec)[MAIN] + + # Reference tensors still trace at batch_size (the trace size), NOT the max. + assert refs["input_ids"].shape[0] == 2 + assert refs["k_cache"].shape[1] == 2 # cache batch dim is index 1 + + batch_dim = shapes["input_ids"][0] + assert batch_dim.min == 1 + assert batch_dim.max == 8 + # The SAME Dim object on every batch axis, so torch.export ties them together. + assert shapes["position_ids"][0] is batch_dim + assert shapes["attn_mask"][0] is batch_dim + assert shapes["k_cache"][1] is batch_dim + assert shapes["v_cache"][1] is batch_dim + m.validate_export_contract({MAIN: refs}, {MAIN: shapes}) + + def test_dynamic_batch_static_cache_still_declares_batch_dim(self, config) -> None: + """With seq-static caches, dim1 still carries the batch Dim (not pinned to None).""" + config.max_position_embeddings = TRACE_KV_CACHE_SEQ_LEN + spec = TraceSpec( + max_context_length=TRACE_KV_CACHE_SEQ_LEN, + cache_seq_len=TRACE_KV_CACHE_SEQ_LEN, + batch_size=2, + max_batch_size=8, + ) + assert spec.caches_are_static + m = _MaskLM(config, use_attention_mask=True) + refs = m.build_reference_inputs(config, torch.float16, spec)[MAIN] + shapes = m.build_dynamic_shapes(config, spec)[MAIN] + batch_dim = shapes["input_ids"][0] + # Seq dim pinned, batch dim dynamic: dim1 carries the batch Dim, seq axis absent. + assert shapes["k_cache"] == {1: batch_dim} + assert shapes["v_cache"] == {1: batch_dim} + m.validate_export_contract({MAIN: refs}, {MAIN: shapes}) + + def test_dynamic_batch_without_mask_still_declares_batch_dim(self, config) -> None: + """Batching is independent of the attn_mask flag: ids/pos/caches get the batch Dim + even with the mask off (there is simply no attn_mask axis to tie).""" + spec = TraceSpec(max_context_length=MAX_CONTEXT_LENGTH, batch_size=2, max_batch_size=8) + m = _MaskLM(config) # mask flag off + shapes = m.build_dynamic_shapes(config, spec)[MAIN] + batch_dim = shapes["input_ids"][0] + assert batch_dim.min == 1 and batch_dim.max == 8 + assert shapes["position_ids"][0] is batch_dim + assert shapes["k_cache"][1] is batch_dim + assert "attn_mask" not in shapes + + def test_max_batch_size_must_exceed_trace_batch(self) -> None: + """Trace size must be strictly below the max, and >= 2, or torch.export specializes it.""" + with pytest.raises(ValueError, match="max_batch_size"): + TraceSpec(max_context_length=MAX_CONTEXT_LENGTH, batch_size=2, max_batch_size=2) + with pytest.raises(ValueError, match="max_batch_size"): + TraceSpec(max_context_length=MAX_CONTEXT_LENGTH, batch_size=1, max_batch_size=8) + + class TestRegisteredModelsSatisfyTheContract: """Pins the real registry, so renaming a forward parameter fails here.""" diff --git a/python/tests/test_model_units/test_export/test_use_attention_mask.py b/python/tests/test_model_units/test_export/test_use_attention_mask.py new file mode 100644 index 00000000..e5880e3e --- /dev/null +++ b/python/tests/test_model_units/test_export/test_use_attention_mask.py @@ -0,0 +1,86 @@ +# 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 + +"""Tests for the ``--attention-mask [default|batched|paged]`` export flag. + +The flag threads CLI -> ``ExportConfig`` -> ``from_hf*`` -> model so the exported macOS +graph declares an additive ``attn_mask`` input instead of the implicit ``is_causal`` mask. +It is macOS-only and defaults to ``default``, leaving the standard causal export untouched. + +Flag plumbing only -- nothing here downloads weights or runs an export. +""" + +from __future__ import annotations + +import pytest + +from coreai_models.export.pipeline import AttentionMaskMode, ExportConfig +from coreai_models.llm.export import _resolve_export_config, build_parser + + +def test_defaults() -> None: + config = ExportConfig(hf_model_id="org/model") + assert config.attention_mask_mode is AttentionMaskMode.DEFAULT + assert config.emits_attn_mask is False + assert config.dynamic_max_batch_size is None + + +def test_batched_emits_mask_and_defaults_batch_to_8() -> None: + """`--attention-mask batched` alone emits the mask and yields a batch-8 dynamic graph.""" + config = ExportConfig( + hf_model_id="org/model", variant="macOS", attention_mask_mode=AttentionMaskMode.BATCHED + ) + assert config.emits_attn_mask is True + assert config.dynamic_max_batch_size == 8 + + +def test_rejects_batched_mask_on_ios() -> None: + """iOS has its own causal-mask contract; the additive mask is macOS-only.""" + with pytest.raises(ValueError, match="macOS"): + ExportConfig( + hf_model_id="org/model", variant="iOS", attention_mask_mode=AttentionMaskMode.BATCHED + ) + + +def test_rejects_paged() -> None: + """``paged`` is reserved for the hardware paged backend; it fails with an explicit error.""" + with pytest.raises(NotImplementedError, match="not available"): + ExportConfig(hf_model_id="org/model", attention_mask_mode=AttentionMaskMode.PAGED) + + +def test_rejects_dynamic_batch_on_ios() -> None: + with pytest.raises(ValueError, match="macOS"): + ExportConfig(hf_model_id="org/model", variant="iOS", dynamic_max_batch_size=8) + + +def test_rejects_dynamic_batch_below_two() -> None: + with pytest.raises(ValueError, match="dynamic_max_batch_size"): + ExportConfig(hf_model_id="org/model", dynamic_max_batch_size=1) + + +def test_rejects_dynamic_batch_without_attn_mask() -> None: + """A dynamic batch without a batched attn_mask would trace a graph the runtime can't serve.""" + with pytest.raises(ValueError, match="requires attention_mask_mode=BATCHED"): + ExportConfig(hf_model_id="org/model", variant="macOS", dynamic_max_batch_size=8) + + +def test_cli_resolves_attention_mask_onto_config() -> None: + assert ( + _resolve_export_config(build_parser().parse_args(["qwen3-0.6b"])).attention_mask_mode + is AttentionMaskMode.DEFAULT + ) + assert ( + _resolve_export_config( + build_parser().parse_args(["qwen3-0.6b", "--attention-mask", "batched"]) + ).attention_mask_mode + is AttentionMaskMode.BATCHED + ) + + +def test_cli_resolves_dynamic_batch_size_onto_config() -> None: + args = build_parser().parse_args( + ["qwen3-0.6b", "--dynamic-batch-size", "8", "--attention-mask", "batched"] + ) + assert _resolve_export_config(args).dynamic_max_batch_size == 8 diff --git a/python/tests/test_model_units/test_models/test_macos_layers/test_qwen3_attention_mask.py b/python/tests/test_model_units/test_models/test_macos_layers/test_qwen3_attention_mask.py new file mode 100644 index 00000000..fd5c024f --- /dev/null +++ b/python/tests/test_model_units/test_models/test_macos_layers/test_qwen3_attention_mask.py @@ -0,0 +1,85 @@ +# 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 + +"""Opt-in additive-attention-mask override for macOS Qwen3 attention. + +Proves, at batch=1, that an explicit additive ``attn_mask`` reproduces the default +``is_causal`` path. Default-off, so the existing export contract is untouched. +""" + +import torch +from transformers.models.qwen3.modeling_qwen3 import Qwen3Config + +from coreai_models.models.macos.qwen3 import Attention, Qwen3ForCausalLM +from coreai_models.primitives.macos.cache import KVCache + + +def _tiny_config() -> Qwen3Config: + config = Qwen3Config( + hidden_size=64, + num_attention_heads=4, + num_key_value_heads=2, + num_hidden_layers=1, + intermediate_size=128, + vocab_size=100, + max_position_embeddings=32, + head_dim=16, + ) + config.rope_scaling = None + config.rope_theta = 10000.0 + return config + + +def _causal_additive_mask(seq_len: int) -> torch.Tensor: + """Additive causal mask broadcastable to [B, heads, q, k].""" + return torch.triu(torch.full((seq_len, seq_len), float("-inf")), diagonal=1).view( + 1, 1, seq_len, seq_len + ) + + +def test_attention_mask_mode_matches_causal_single_batch() -> None: + """Mask-mode attention fed a causal additive mask == default is_causal path (B=1).""" + torch.manual_seed(0) + config = _tiny_config() + seq_len = 6 + + causal = Attention(config, layer_idx=0).eval() + masked = Attention(config, layer_idx=0, use_attention_mask=True).eval() + masked.load_state_dict(causal.state_dict()) + + x = torch.rand(1, seq_len, config.hidden_size) + position_ids = torch.arange(seq_len, dtype=torch.int32).unsqueeze(0) + + with torch.no_grad(): + out_causal = causal(x, position_ids) + out_masked = masked(x, position_ids, attn_mask=_causal_additive_mask(seq_len)) + + torch.testing.assert_close(out_masked, out_causal, atol=1e-5, rtol=1e-5) + + +def test_qwen3_for_causal_lm_mask_mode_matches_causal_single_batch() -> None: + """End-to-end: mask-mode Qwen3ForCausalLM threads attn_mask == default is_causal (B=1).""" + torch.manual_seed(0) + config = _tiny_config() + seq_len = 6 + + causal = Qwen3ForCausalLM(config, model_device="cpu") + causal.to(torch.float32).eval() + masked = Qwen3ForCausalLM(config, model_device="cpu", use_attention_mask=True) + masked.to(torch.float32).eval() + masked.load_state_dict(causal.state_dict()) + + input_ids = torch.randint(0, config.vocab_size, (1, seq_len)) + position_ids = torch.arange(seq_len, dtype=torch.int32).unsqueeze(0) + kc, vc = KVCache.create_cache_tensors(config, dtype=torch.float32, seq_len=seq_len) + kc2, vc2 = KVCache.create_cache_tensors(config, dtype=torch.float32, seq_len=seq_len) + + with torch.no_grad(): + out_causal = causal(input_ids, position_ids, kc, vc) + out_masked = masked( + input_ids, position_ids, kc2, vc2, attn_mask=_causal_additive_mask(seq_len) + ) + + torch.testing.assert_close(out_masked, out_causal, atol=1e-5, rtol=1e-5)