diff --git a/scripts/generate_golden.py b/scripts/generate_golden.py index e05617ba7..0ae3ea5cc 100644 --- a/scripts/generate_golden.py +++ b/scripts/generate_golden.py @@ -241,6 +241,22 @@ def _get_model_device(model: object, device: str): # transformers) are deferred to avoid import cost when --dry-run. +# HuggingFace ``model_type`` values whose reference must run the eager attention +# kernel. GraniteSWA's learnable per-head sink is an extra logit inside the +# softmax denominator, which SDPA cannot express — upstream sets +# ``GraniteSWAPreTrainedModel._supports_sdpa = False`` for exactly this reason. +# Pin it explicitly so the golden reference can never drift onto a kernel that +# silently drops the sink. +_EAGER_ATTENTION_MODEL_TYPES = frozenset({"granite_swa"}) + + +def _forced_attn_implementation(case: TestCase) -> str | None: + """Return the attention backend to pin for this case, or ``None``.""" + if case.model_type in _EAGER_ATTENTION_MODEL_TYPES: + return "eager" + return None + + def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None: """Generate golden data for a causal-lm (text-generation) model.""" from mobius._testing.golden import save_generation_json, save_golden_ref @@ -267,7 +283,11 @@ def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None: ) else: model, tokenizer = load_torch_model( - case.model_id, device=device, trust_remote_code=case.trust_remote_code + case.model_id, + device=device, + trust_remote_code=case.trust_remote_code, + revision=case.revision, + attn_implementation=_forced_attn_implementation(case), ) encoded = tokenizer(case.prompts[0], return_tensors="np", padding=False) diff --git a/src/mobius/_configs/__init__.py b/src/mobius/_configs/__init__.py index 2fbc36c8f..b52856dda 100644 --- a/src/mobius/_configs/__init__.py +++ b/src/mobius/_configs/__init__.py @@ -38,6 +38,7 @@ Gemma4Config, GlmAsrConfig, GraniteMoeHybridConfig, + GraniteSwaConfig, JambaConfig, JetMoeConfig, Lfm2Config, @@ -114,6 +115,7 @@ "Gemma4Config", "GlmAsrConfig", "GraniteMoeHybridConfig", + "GraniteSwaConfig", "JambaConfig", "JetMoeConfig", "Lfm2Config", diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index cafe91f3c..5ead16b20 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -1494,6 +1494,74 @@ def from_transformers(cls, config, parent_config=None) -> MuseGlimmerConfig: ) +@dataclasses.dataclass +class GraniteSwaConfig(CausalLMConfig): + """Configuration for GraniteSWA (``GraniteSWAForCausalLM``). + + Adds the two fields that distinguish GraniteSWA from plain Granite, on top + of the shared Granite scaling multipliers (``embedding_multiplier``, + ``attention_multiplier``, ``logits_scaling``, ``residual_multiplier``) that + :class:`ArchitectureConfig` already extracts: + + * ``layer_types`` — per-layer ``"full_attention"`` / ``"sliding_attention"``. + HuggingFace defaults this to ``full_attention`` on every fourth layer + (``i % 4 == 0``) when the checkpoint omits it. + * ``layer_rope_theta`` — per-layer RoPE base frequency, where ``0`` marks a + NoPE layer. HuggingFace defaults it to the global ``rope_theta`` for + every layer. + + Both defaults are re-applied here so that a raw ``config.json`` mapping + (which never runs ``GraniteSWAConfig.__post_init__``) yields the same + architecture as a materialised HuggingFace config object. + """ + + layer_rope_theta: list[float | int] | None = None + + @classmethod + def from_transformers(cls, config, parent_config=None) -> GraniteSwaConfig: + base = ArchitectureConfig.from_transformers(config, parent_config) + + # GraniteSWA is definitionally a RoPE architecture: NoPE is expressed + # per layer via ``layer_rope_theta[i] == 0``, never globally. A raw + # ``config.json`` mapping carries only the flat ``rope_theta: 10000`` + # (the HF default) with no ``rope_parameters``, which the generic + # extractor reads as "no RoPE signal at all". Re-assert RoPE here so + # the raw-JSON path builds the same graph as a materialised HF config. + if base.rope_type is None: + base = dataclasses.replace( + base, + rope_type="default", + rope_theta=float(getattr(config, "rope_theta", None) or 10_000.0), + partial_rotary_factor=1.0, + ) + + # HF ``__post_init__``: every fourth layer is full attention, rest slide. + layer_types = base.layer_types + if not layer_types: + layer_types = [ + "full_attention" if index % 4 == 0 else "sliding_attention" + for index in range(base.num_hidden_layers) + ] + + # HF ``__post_init__``: default to the global rope_theta on every layer. + # ``0`` is a real, meaningful value here (NoPE), so only a missing or + # empty list falls back — never a list that legitimately contains 0. + raw_layer_rope_theta = getattr(config, "layer_rope_theta", None) + if raw_layer_rope_theta: + layer_rope_theta: list[float | int] = list(raw_layer_rope_theta) + else: + layer_rope_theta = [base.rope_theta or 0.0] * base.num_hidden_layers + + base = dataclasses.replace( + base, + layer_types=layer_types, + no_rope_layers=[ + index for index, theta in enumerate(layer_rope_theta) if not theta + ], + ) + return cls(**_shallow_fields(base), layer_rope_theta=layer_rope_theta) + + @dataclasses.dataclass class Lfm2Config(CausalLMConfig): """Configuration for LFM2's automatically adjusted feed-forward width.""" diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index d626c1f7b..3b7ee1fa7 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -28,6 +28,7 @@ Gemma3nMultiModalConfig, Gemma4AssistantConfig, Gemma4Config, + GraniteSwaConfig, Lfm2Config, Lfm2VlConfig, MMSConfig, @@ -140,6 +141,7 @@ from mobius.models.gpt2 import GPT2CausalLMModel from mobius.models.gpt_neox import GPTNeoXCausalLMModel, GPTNeoXJapaneseCausalLMModel from mobius.models.gptj_codegen import CodeGenCausalLMModel, GPTJCausalLMModel +from mobius.models.granite_swa import GraniteSwaCausalLMModel from mobius.models.granitemoehybrid import GraniteMoeHybridCausalLMModel from mobius.models.internvl import InternVL2Model from mobius.models.jamba import JambaCausalLMModel @@ -492,6 +494,7 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: "gpt_oss": ModelRegistration(GPTOSSCausalLMModel), "gptj": ModelRegistration(GPTJCausalLMModel), "granite": ModelRegistration(GraniteCausalLMModel), + "granite_swa": ModelRegistration(GraniteSwaCausalLMModel, config_class=GraniteSwaConfig), "hunyuan_v1_dense": ModelRegistration(HunYuanV1DenseCausalLMModel), "internlm2": ModelRegistration(InternLM2CausalLMModel), "llama4_text": ModelRegistration(Llama4CausalLMModel), @@ -986,6 +989,7 @@ def _create_default_registry() -> ModelRegistry: "gemma3n_text": "google/gemma-3n-E2B-it", "gemma4_text": "google/gemma-4-E2B-it", "granite": "ibm-granite/granite-3.3-2b-instruct", + "granite_swa": "ibm-granite/granite-swash-2b", "internlm2": "internlm/internlm2_5-7b-chat", "nemotron": "nvidia/Nemotron-Mini-4B-Instruct", "olmo": "allenai/OLMo-1B-hf", diff --git a/src/mobius/_testing/torch_reference.py b/src/mobius/_testing/torch_reference.py index 271ec2693..de2eda4aa 100644 --- a/src/mobius/_testing/torch_reference.py +++ b/src/mobius/_testing/torch_reference.py @@ -198,6 +198,8 @@ def load_torch_model( dtype: torch.dtype = torch.float32, device: str = "cpu", trust_remote_code: bool = True, + revision: str | None = None, + attn_implementation: str | None = None, ): """Load a HuggingFace causal LM model for reference inference. @@ -211,6 +213,14 @@ def load_torch_model( natively supported by the installed transformers, so the transformers-5.x-compatible implementation is used instead of an older bundled ``modeling_*.py`` that relies on removed cache APIs. + revision: Optional immutable HuggingFace revision (commit SHA) applied + to the tokenizer, config, and weight downloads alike, so a single + pin covers every Hub artifact the reference depends on. + attn_implementation: Optional HuggingFace attention backend to force + (e.g. ``"eager"``). Needed for architectures whose attention is + not SDPA-expressible — GraniteSWA's learnable sink is an extra + logit inside the softmax denominator, so only the eager kernel + reproduces the published semantics. Returns: Tuple of (model, tokenizer). @@ -219,26 +229,37 @@ def load_torch_model( _install_dynamic_cache_legacy_shims() - tokenizer = transformers.AutoTokenizer.from_pretrained( - model_id, trust_remote_code=trust_remote_code - ) + hub_kwargs: dict = {"trust_remote_code": trust_remote_code} + if revision is not None: + hub_kwargs["revision"] = revision + + tokenizer = transformers.AutoTokenizer.from_pretrained(model_id, **hub_kwargs) # NemotronH: disable rescale_prenorm_residual before loading to # prevent _init_weights from corrupting out_proj.weight with # random kaiming_uniform_ initialization after checkpoint loading. - config = transformers.AutoConfig.from_pretrained( - model_id, trust_remote_code=trust_remote_code - ) + config = transformers.AutoConfig.from_pretrained(model_id, **hub_kwargs) if getattr(config, "model_type", None) == "nemotron_h": config.rescale_prenorm_residual = False + model_kwargs: dict = dict(hub_kwargs) + if attn_implementation is not None: + model_kwargs["attn_implementation"] = attn_implementation + model = transformers.AutoModelForCausalLM.from_pretrained( model_id, config=config, dtype=dtype, device_map=device, - trust_remote_code=trust_remote_code, + **model_kwargs, ) + if attn_implementation is not None: + actual = model.config._attn_implementation + if actual != attn_implementation: + raise RuntimeError( + f"Requested attn_implementation={attn_implementation!r} for " + f"{model_id} but transformers resolved {actual!r}." + ) _fix_nemotron_h_init_weights(model, model_id) model.eval() diff --git a/src/mobius/components/__init__.py b/src/mobius/components/__init__.py index e0afec8cc..c7d48a908 100644 --- a/src/mobius/components/__init__.py +++ b/src/mobius/components/__init__.py @@ -29,6 +29,7 @@ "EncoderDecoderAttention", "EncoderLayer", "FCMLP", + "Float32SinkAttention", "FusedGateUpMLP", "GatedDeltaNet", "GatedMLP", @@ -73,6 +74,7 @@ "Siglip2NaFlexVisionEmbeddings", "Siglip2NaFlexVisionModel", "SigmoidTopKGate", + "SinkAttention", "SnakeBeta", "SoftmaxTopKGate", "SparseMixerGate", @@ -103,7 +105,9 @@ from mobius.components._activations import SiLU, get_activation from mobius.components._attention import ( Attention, + Float32SinkAttention, GQAContext, + SinkAttention, StaticCacheState, ) from mobius.components._attention import ( diff --git a/src/mobius/components/_attention.py b/src/mobius/components/_attention.py index f6724e841..ae2647c2a 100644 --- a/src/mobius/components/_attention.py +++ b/src/mobius/components/_attention.py @@ -313,6 +313,33 @@ def __init__( self.q_norm = None self.k_norm = None + def _apply_qk_norm( + self, + op: OpBuilder, + query_states: ir.Value, + key_states: ir.Value, + ) -> tuple[ir.Value, ir.Value]: + """Apply optional Q/K normalization, returning the (possibly) normed pair. + + No-op when ``config.attn_qk_norm`` is False. Shared by every + :class:`Attention` forward path (fused, GQA, and the manual + :class:`SinkAttention` path) so the normalization semantics can + never drift between them. + """ + if self.q_norm is None or self.k_norm is None: + return query_states, key_states + if self._qk_norm_full: + # Apply norm on 3D tensor (across all heads) + return self.q_norm(op, query_states), self.k_norm(op, key_states) + # Apply norm per-head on 4D tensor + query_states = op.Reshape(query_states, [0, 0, -1, self.head_dim]) + key_states = op.Reshape(key_states, [0, 0, -1, self.head_dim]) + query_states = self.q_norm(op, query_states) + key_states = self.k_norm(op, key_states) + query_states = op.Reshape(query_states, [0, 0, -1]) + key_states = op.Reshape(key_states, [0, 0, -1]) + return query_states, key_states + def forward( self, op: OpBuilder, @@ -326,19 +353,7 @@ def forward( key_states = self.k_proj(op, hidden_states) value_states = self.v_proj(op, hidden_states) - if self.q_norm is not None and self.k_norm is not None: - if self._qk_norm_full: - # Apply norm on 3D tensor (across all heads) - query_states = self.q_norm(op, query_states) - key_states = self.k_norm(op, key_states) - else: - # Apply norm per-head on 4D tensor - query_states = op.Reshape(query_states, [0, 0, -1, self.head_dim]) - key_states = op.Reshape(key_states, [0, 0, -1, self.head_dim]) - query_states = self.q_norm(op, query_states) - key_states = self.k_norm(op, key_states) - query_states = op.Reshape(query_states, [0, 0, -1]) - key_states = op.Reshape(key_states, [0, 0, -1]) + query_states, key_states = self._apply_qk_norm(op, query_states, key_states) # Direct GroupQueryAttention path: skip external RoPE, fuse everything. if isinstance(attention_bias, GQAContext): @@ -459,6 +474,284 @@ def _forward_gqa( return attn_out, (present_key, present_value) +class SinkAttention(Attention): + """Attention with a learnable per-head sink logit in the softmax denominator. + + An *attention sink* is one extra learnable logit per head that participates + in the softmax denominator but carries no value vector. It lets a head + "discard" probability mass into a virtual null position, so the attention + output shrinks instead of being forced to sum to one over real tokens. + + Implemented by appending the per-head sink as one extra column to the + attention scores before the softmax and then dropping that column:: + + combined = concat([scores, sinks], dim=-1) # [B, H, S_q, S_kv + 1] + combined = combined - max(combined, dim=-1) # numerical stability + probs = softmax(combined, dim=-1)[..., :-1] + out = probs @ V + + This is algebraically identical to the two documented upstream forms: + + * GPT-OSS ``eager_attention_forward`` concatenates the sink column exactly + as above. + * GraniteSWA ``eager_attention_forward`` instead computes an ordinary + softmax and rescales the output by + ``sigmoid(logsumexp(scores) - sink)``. Writing ``Z = sum(exp(scores))``, + that factor is ``Z / (Z + exp(sink))``, which is precisely the mass this + implementation removes by keeping the sink column in the denominator. + + Because the sink lives inside the softmax, the fused ONNX ``Attention`` and + ``GroupQueryAttention`` ops cannot be used: the score matrix is built + explicitly with MatMul/Softmax. Consequently ``attention_bias`` must be a + float additive bias that already bakes in causality, any sliding window, + and padding (see :func:`~mobius.components._common.create_attention_bias`). + + This class reproduces the GPT-OSS precision contract: the stabilised + softmax runs in the *compute* dtype, matching upstream's explicit + ``F.softmax(combined_logits, dim=-1, dtype=combined_logits.dtype)``. Models + whose upstream kernel forces float32 instead must use + :class:`Float32SinkAttention`. + + Args: + config: Architecture configuration. + rms_norm_class: Norm class for optional Q/K normalization. + scale: Custom attention scale factor (default: ``1/sqrt(head_dim)``). + Granite-family models pass ``config.attention_multiplier``. + linear_class: Factory callable for the projection layers. + """ + + #: Whether the sink softmax is evaluated in float32 regardless of the + #: model compute dtype. ``False`` keeps the graph in the compute dtype + #: (GPT-OSS); :class:`Float32SinkAttention` flips it to ``True``. + upcast_sink_softmax: bool = False + + def __init__( + self, + config: ArchitectureConfig, + rms_norm_class: type[nn.Module] | None = None, + scale: float | None = None, + linear_class: type | None = None, + ): + super().__init__( + config, + rms_norm_class=rms_norm_class, + scale=scale, + linear_class=linear_class, + ) + if self._softcap: + # Softcapping would have to be applied to the score matrix *before* + # the sink column is concatenated (the sink itself is never capped + # upstream). No sink model uses it, so refuse rather than guess. + raise ValueError( + "SinkAttention does not implement attention-logit softcapping " + f"(attn_logit_softcapping={self._softcap})." + ) + self.num_kv_groups = self.num_attention_heads // self.num_key_value_heads + self._dtype = config.dtype + + # Learnable sink logit: one scalar per attention head [num_heads]. + self.sinks = nn.Parameter([self.num_attention_heads]) + + def _expand_kv_for_gqa( + self, + op: OpBuilder, + kv: ir.Value, + batch_1d: ir.Value, + kv_len_1d: ir.Value, + ) -> ir.Value: + """Expand KV from [B, kv_heads, S, d] to [B, q_heads, S, d] for GQA. + + Uses unsqueeze+expand+reshape to replicate each KV head + ``num_kv_groups`` times consecutively: ``[kv0]*g, [kv1]*g, ...``, + which is what HuggingFace's ``repeat_kv`` does. + """ + # [B, kv_heads, S, d] → [B, kv_heads, 1, S, d] + kv_5d = op.Unsqueeze(kv, [2]) + # Expand to [B, kv_heads, num_kv_groups, S, d] + expand_shape = op.Concat( + batch_1d, + [self.num_key_value_heads, self.num_kv_groups], + kv_len_1d, + [self.head_dim], + axis=0, + ) + kv_exp = op.Expand(kv_5d, expand_shape) + # Flatten to [B, q_heads, S, d] + flat_shape = op.Concat( + batch_1d, + [self.num_attention_heads], + kv_len_1d, + [self.head_dim], + axis=0, + ) + return op.Reshape(kv_exp, flat_shape) + + def forward( + self, + op: OpBuilder, + hidden_states: ir.Value, + attention_bias: ir.Value | GQAContext | None, + position_embeddings: tuple | None = None, + past_key_value: tuple | None = None, + static_cache: StaticCacheState | None = None, + ): + if isinstance(attention_bias, GQAContext): + raise TypeError( + "SinkAttention cannot emit GroupQueryAttention: the sink logit " + "must take part in the softmax denominator. Build this model " + "with a float additive attention bias instead." + ) + if static_cache is not None: + raise NotImplementedError( + "SinkAttention does not support the opset-24 static KV cache." + ) + + # hidden_states: [B, S, H] + batch_1d = op.Shape(hidden_states, start=0, end=1) # 1-D tensor holding B + seq_1d = op.Shape(hidden_states, start=1, end=2) # 1-D tensor holding S_q + + # QKV projections: [B, S, heads * d] + query = self.q_proj(op, hidden_states) + key = self.k_proj(op, hidden_states) + value = self.v_proj(op, hidden_states) + + query, key = self._apply_qk_norm(op, query, key) + + # Apply RoPE on the 3D packed format [B, S, heads * d]. + # ``position_embeddings is None`` marks a NoPE layer. + if position_embeddings is not None: + query = apply_rotary_pos_emb( + op, + x=query, + position_embeddings=position_embeddings, + num_heads=self.num_attention_heads, + rotary_embedding_dim=self.rotary_embedding_dim, + interleaved=self._rope_interleave, + ) + key = apply_rotary_pos_emb( + op, + x=key, + position_embeddings=position_embeddings, + num_heads=self.num_key_value_heads, + rotary_embedding_dim=self.rotary_embedding_dim, + interleaved=self._rope_interleave, + ) + + # Reshape to 4D and transpose: [B, S, heads, d] → [B, heads, S, d] + query = op.Transpose( + op.Reshape(query, [0, 0, self.num_attention_heads, self.head_dim]), + perm=[0, 2, 1, 3], + ) # [B, q_heads, S_q, d] + key = op.Transpose( + op.Reshape(key, [0, 0, self.num_key_value_heads, self.head_dim]), + perm=[0, 2, 1, 3], + ) # [B, kv_heads, S_q, d] + value = op.Transpose( + op.Reshape(value, [0, 0, self.num_key_value_heads, self.head_dim]), + perm=[0, 2, 1, 3], + ) # [B, kv_heads, S_q, d] + + # KV cache: prepend past tokens + if past_key_value is not None: + key = op.Concat(past_key_value[0], key, axis=2) # [B, kv_heads, past+S, d] + value = op.Concat(past_key_value[1], value, axis=2) # [B, kv_heads, past+S, d] + present_key_value = (key, value) + + # Total KV sequence length (after cache concatenation) + kv_len_1d = op.Shape(key, start=2, end=3) # [S_kv] + + # GQA: expand key/value from kv_heads to q_heads + if self.num_kv_groups > 1: + key_exp = self._expand_kv_for_gqa(op, key, batch_1d, kv_len_1d) + value_exp = self._expand_kv_for_gqa(op, value, batch_1d, kv_len_1d) + else: + key_exp = key + value_exp = value + + # Attention scores: [B, q_heads, S_q, d] @ [B, q_heads, d, S_kv] + key_t = op.Transpose(key_exp, perm=[0, 1, 3, 2]) # [B, q_heads, d, S_kv] + attn_scores = op.MatMul(query, key_t) # [B, q_heads, S_q, S_kv] + attn_scores = op.Mul(attn_scores, self.scaling) + + # Add causal + sliding-window + padding mask (float additive bias) + if attention_bias is not None: + # attention_bias: [B, 1, S_q, S_kv] — broadcasts over q_heads + attn_scores = op.Add(attn_scores, attention_bias) + + # Append the sink column: [q_heads] → [B, q_heads, S_q, 1] + sinks_4d = op.Reshape(self.sinks, [1, self.num_attention_heads, 1, 1]) + expand_shape = op.Concat( + batch_1d, + [self.num_attention_heads], + seq_1d, + [1], + axis=0, + ) + sinks_expanded = op.Expand(sinks_4d, expand_shape) # [B, q_heads, S_q, 1] + # combined: [B, q_heads, S_q, S_kv + 1] + combined = op.Concat(attn_scores, sinks_expanded, axis=-1) + + # HuggingFace's precision contract differs per model: GPT-OSS softmaxes + # in the compute dtype, GraniteSWA forces float32. ``upcast_sink_softmax`` + # selects between them; for float32 builds both are identical and the + # graph carries no Cast nodes at all. + upcast = self.upcast_sink_softmax and self._dtype != ir.DataType.FLOAT + if upcast: + combined = op.Cast(combined, to=ir.DataType.FLOAT) + + # Numerical stability: subtract the per-row max before the softmax + row_max = op.ReduceMax(combined, [-1], keepdims=True) # [B, q_heads, S_q, 1] + combined = op.Sub(combined, row_max) + + # Softmax over the extended (S_kv + 1) axis + probs = op.Softmax(combined, axis=-1) # [B, q_heads, S_q, S_kv + 1] + + # Drop the sink column: slice axis 3 from 0 to -1 (all but the last). + # The removed mass is exactly the sink's share of the denominator, so + # the remaining probabilities sum to sigmoid(logsumexp(scores) - sink). + scores = op.Slice(probs, [0], [-1], [3]) # [B, q_heads, S_q, S_kv] + if upcast: + scores = op.Cast(scores, to=self._dtype) + + # Weighted sum with value: [B, q_heads, S_q, d] + attn_out = op.MatMul(scores, value_exp) + + # Transpose and flatten heads: [B, q_heads, S_q, d] → [B, S_q, q_heads*d] + attn_out = op.Transpose(attn_out, perm=[0, 2, 1, 3]) # [B, S_q, q_heads, d] + attn_out = op.Reshape(attn_out, [0, 0, -1]) # [B, S_q, q_heads * d] + + # Output projection + attn_out = self.o_proj(op, attn_out) + return attn_out, present_key_value + + +class Float32SinkAttention(SinkAttention): + """Sink attention whose sink scaling and softmax are forced to float32. + + Some upstream sink kernels deliberately leave the compute dtype for the + softmax. GraniteSWA's ``eager_attention_forward`` is one of them:: + + lse = torch.logsumexp(attn_weights, dim=-1) + sink_scale = (lse - sinks).to(torch.float32).sigmoid() # forced fp32 + attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32) + + In the extra-column formulation, the ``logsumexp``/``sigmoid`` pair and the + softmax are the same reduction, so forcing that one softmax to float32 + reproduces both upstream upcasts at once. + + This is deliberately a separate class rather than a constructor flag: it is + a precision *contract* tied to the architecture, not a tuning knob, and + keeping it in the type means a model cannot silently acquire (or lose) the + upcast. GPT-OSS, whose upstream kernel softmaxes in + ``combined_logits.dtype``, must keep using :class:`SinkAttention`. + + For float32 builds this class emits exactly the same graph as + :class:`SinkAttention` — the upcast would be a no-op, so no Cast is added. + """ + + upcast_sink_softmax: bool = True + + class Qwen35Attention(nn.Module): """Multi-head attention with output gating for Qwen3.5. diff --git a/src/mobius/components/_decoder.py b/src/mobius/components/_decoder.py index 67ca345e4..63a36e471 100644 --- a/src/mobius/components/_decoder.py +++ b/src/mobius/components/_decoder.py @@ -44,6 +44,12 @@ class DecoderLayer(nn.Module): Pass :class:`~mobius.components.FusedGateUpMLP` for models that store the gate and up projections as a single fused ``gate_up_proj`` weight (e.g. Phi-3, Phi-4, GLM). + attention_class: Attention module class to use (default: + :class:`~mobius.components.Attention`). Must accept the same + ``(config, rms_norm_class=..., scale=..., linear_class=...)`` + constructor signature. Pass + :class:`~mobius.components.SinkAttention` for models with learnable + per-head attention sinks (GraniteSWA, GPT-OSS). """ def __init__( @@ -56,17 +62,20 @@ def __init__( post_norm: bool = False, linear_class: type | None = None, mlp_class: type | None = None, + attention_class: type[nn.Module] | None = None, ): super().__init__() if norm_class is None: norm_class = RMSNorm if mlp_class is None: mlp_class = MLP + if attention_class is None: + attention_class = Attention self._post_norm = post_norm self._residual_multiplier = residual_multiplier - self.self_attn = Attention( + self.self_attn = attention_class( config, rms_norm_class=norm_class, scale=attention_scale, @@ -192,6 +201,7 @@ def create_decoder_layer( post_norm: bool = False, linear_class: type | None = None, mlp_class: type | None = None, + attention_class: type[nn.Module] | None = None, ) -> DecoderLayer: """Config-driven factory for creating decoder layers. @@ -209,6 +219,8 @@ def create_decoder_layer( factory for LoRA-adapted layers. mlp_class: MLP module class override (default: MLP). Pass FusedGateUpMLP for models with fused gate_up_proj weights. + attention_class: Attention module class override (default: Attention). + Pass SinkAttention for models with learnable attention sinks. Returns: A configured DecoderLayer instance. @@ -224,6 +236,7 @@ def create_decoder_layer( post_norm=post_norm, linear_class=linear_class, mlp_class=mlp_class, + attention_class=attention_class, ) diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 57c617fac..9839faf72 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -70,6 +70,7 @@ "GraniteCausalLMModel", "GraniteMoECausalLMModel", "GraniteMoeHybridCausalLMModel", + "GraniteSwaCausalLMModel", "HunYuanMoEV1CausalLMModel", "HunYuanV1DenseCausalLMModel", "HunYuanVLMoTModel", @@ -243,6 +244,7 @@ from mobius.models.gptj_codegen import CodeGenCausalLMModel, GPTJCausalLMModel from mobius.models.gptoss import GPTOSSCausalLMModel from mobius.models.granite import GraniteCausalLMModel, GraniteMoECausalLMModel +from mobius.models.granite_swa import GraniteSwaCausalLMModel from mobius.models.granitemoehybrid import GraniteMoeHybridCausalLMModel from mobius.models.hunyuan_dit import HunyuanDiT2DModel from mobius.models.hunyuan_v1 import HunYuanV1DenseCausalLMModel diff --git a/src/mobius/models/gptoss.py b/src/mobius/models/gptoss.py index 5dc0cff59..134f53a5c 100644 --- a/src/mobius/models/gptoss.py +++ b/src/mobius/models/gptoss.py @@ -18,7 +18,6 @@ from __future__ import annotations -import math from typing import TYPE_CHECKING import torch @@ -29,10 +28,10 @@ Embedding, Linear, RMSNorm, + SinkAttention, create_attention_bias, initialize_rope, ) -from mobius.components._rotary_embedding import apply_rotary_pos_emb from mobius.models.base import CausalLMModel if TYPE_CHECKING: @@ -146,209 +145,24 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value): return result -class _GptOssAttention(nn.Module): +class _GptOssAttention(SinkAttention): """GQA attention with learned per-head sinks for GPT-OSS. - HF ``eager_attention_forward`` appends one extra logit per token per head - (the learnable ``sinks`` value) to the attention scores before softmax. - This lets each head "discard" a token's weight into a virtual null position: - - combined = cat([attn_weights, sinks_expanded], dim=-1) # [B, H, S, S_kv+1] - combined = combined - max(combined) # numerical stability - probs = softmax(combined, dim=-1)[..., :-1] # drop sink, [B, H, S, S_kv] - out = probs @ V - - Implements this manually (cannot use fused op.Attention with sinks). + HF ``GptOssAttention``'s ``eager_attention_forward`` appends one extra logit + per token per head (the learnable ``sinks`` value) to the attention scores + before the softmax, letting each head "discard" weight into a virtual null + position. That is exactly the shared + :class:`~mobius.components.SinkAttention` behaviour; GPT-OSS keeps the + default ``1/sqrt(head_dim)`` scale. + + It also keeps the *base* precision contract on purpose. Upstream softmaxes + in the compute dtype (``F.softmax(combined_logits, dim=-1, + dtype=combined_logits.dtype)``), unlike GraniteSWA which forces float32, so + GPT-OSS must NOT use :class:`~mobius.components.Float32SinkAttention` — + doing so would change f16/bf16 numerics and double the size of the largest + score tensor. """ - def __init__(self, config: ArchitectureConfig): - super().__init__() - self.hidden_size = config.hidden_size - self.head_dim = config.head_dim - self.num_attention_heads = config.num_attention_heads - self.num_key_value_heads = config.num_key_value_heads - self.num_kv_groups = config.num_attention_heads // config.num_key_value_heads - self.scale = config.head_dim**-0.5 - self._rotary_embedding_dim = ( - 0 - if math.isclose(config.partial_rotary_factor, 1.0) - else int(self.head_dim * config.partial_rotary_factor) - ) - self._rope_interleave = config.rope_interleave - - # QKV projections with bias (attention_bias=True for GPT-OSS) - self.q_proj = Linear( - config.hidden_size, - config.num_attention_heads * config.head_dim, - bias=config.attn_qkv_bias, - ) - self.k_proj = Linear( - config.hidden_size, - config.num_key_value_heads * config.head_dim, - bias=config.attn_qkv_bias, - ) - self.v_proj = Linear( - config.hidden_size, - config.num_key_value_heads * config.head_dim, - bias=config.attn_qkv_bias, - ) - self.o_proj = Linear( - config.num_attention_heads * config.head_dim, - config.hidden_size, - bias=config.attn_o_bias, - ) - - # Learnable sink logit: one scalar per attention head [num_heads] - self.sinks = nn.Parameter([config.num_attention_heads]) - - def _expand_kv_for_gqa( - self, - op: OpBuilder, - kv: ir.Value, - batch_1d: ir.Value, - kv_len_1d: ir.Value, - ) -> ir.Value: - """Expand KV from [B, kv_heads, S, d] to [B, q_heads, S, d] for GQA. - - Uses unsqueeze+expand+reshape to replicate each KV head ``num_kv_groups`` - times consecutively: [kv0]*g, [kv1]*g, ..., which is what ``repeat_kv`` does. - """ - # [B, kv_heads, S, d] → [B, kv_heads, 1, S, d] - kv_5d = op.Unsqueeze(kv, [2]) - # Expand to [B, kv_heads, num_kv_groups, S, d] - expand_shape = op.Concat( - batch_1d, - [self.num_key_value_heads, self.num_kv_groups], - kv_len_1d, - [self.head_dim], - axis=0, - ) - kv_exp = op.Expand(kv_5d, expand_shape) - # Flatten to [B, q_heads, S, d] - flat_shape = op.Concat( - batch_1d, - [self.num_attention_heads], - kv_len_1d, - [self.head_dim], - axis=0, - ) - return op.Reshape(kv_exp, flat_shape) - - def forward( - self, - op: OpBuilder, - hidden_states: ir.Value, - attention_bias: ir.Value | None, - position_embeddings: tuple | None = None, - past_key_value: tuple | None = None, - ): - # hidden_states: [B, S, H] - batch_1d = op.Shape(hidden_states, start=0, end=1) # 1D tensor containing B - seq_1d = op.Shape(hidden_states, start=1, end=2) # 1D tensor containing S - - # QKV projections: [B, S, heads * d] - query = self.q_proj(op, hidden_states) - key = self.k_proj(op, hidden_states) - value = self.v_proj(op, hidden_states) - - # Apply RoPE on 3D packed format [B, S, heads * d] - if position_embeddings is not None: - query = apply_rotary_pos_emb( - op, - x=query, - position_embeddings=position_embeddings, - num_heads=self.num_attention_heads, - rotary_embedding_dim=self._rotary_embedding_dim, - interleaved=self._rope_interleave, - ) - key = apply_rotary_pos_emb( - op, - x=key, - position_embeddings=position_embeddings, - num_heads=self.num_key_value_heads, - rotary_embedding_dim=self._rotary_embedding_dim, - interleaved=self._rope_interleave, - ) - - # Reshape to 4D and transpose: [B, S, heads, d] → [B, heads, S, d] - query = op.Transpose( - op.Reshape(query, [0, 0, self.num_attention_heads, self.head_dim]), - perm=[0, 2, 1, 3], - ) # [B, q_heads, S, d] - key = op.Transpose( - op.Reshape(key, [0, 0, self.num_key_value_heads, self.head_dim]), - perm=[0, 2, 1, 3], - ) # [B, kv_heads, S, d] - value = op.Transpose( - op.Reshape(value, [0, 0, self.num_key_value_heads, self.head_dim]), - perm=[0, 2, 1, 3], - ) # [B, kv_heads, S, d] - - # KV cache: prepend past tokens - if past_key_value is not None: - key = op.Concat(past_key_value[0], key, axis=2) # [B, kv_heads, past+S, d] - value = op.Concat(past_key_value[1], value, axis=2) # [B, kv_heads, past+S, d] - present_key_value = (key, value) - - # Total KV sequence length (after cache concatenation) - kv_len_1d = op.Shape(key, start=2, end=3) # [total_S] - - # GQA: expand key/value from kv_heads to q_heads - if self.num_kv_groups > 1: - key_exp = self._expand_kv_for_gqa(op, key, batch_1d, kv_len_1d) - value_exp = self._expand_kv_for_gqa(op, value, batch_1d, kv_len_1d) - else: - key_exp = key - value_exp = value - - # Attention scores: [B, q_heads, S_q, S_kv] - # query @ key.T: [B, q_heads, S_q, d] @ [B, q_heads, d, S_kv] - key_t = op.Transpose(key_exp, perm=[0, 1, 3, 2]) # [B, q_heads, d, S_kv] - attn_scores = op.MatMul(query, key_t) # [B, q_heads, S_q, S_kv] - attn_scores = op.Mul(attn_scores, self.scale) - - # Add causal+sliding_window+padding mask (float additive bias) - if attention_bias is not None: - # attention_bias: [B, 1, S_q, S_kv] — broadcasts over q_heads - attn_scores = op.Add(attn_scores, attention_bias) - - # Append sinks column: [q_heads] → [B, q_heads, S_q, 1] - sinks_4d = op.Reshape( - self.sinks, - [1, self.num_attention_heads, 1, 1], - ) - expand_shape = op.Concat( - batch_1d, - [self.num_attention_heads], - seq_1d, - [1], - axis=0, - ) - sinks_expanded = op.Expand(sinks_4d, expand_shape) # [B, q_heads, S_q, 1] - # combined: [B, q_heads, S_q, S_kv+1] - combined = op.Concat(attn_scores, sinks_expanded, axis=-1) - - # Numerical stability: subtract per-row max before softmax - row_max = op.ReduceMax(combined, [-1], keepdims=True) # [B, q_heads, S_q, 1] - combined = op.Sub(combined, row_max) - - # Softmax over the extended sequence (S_kv + 1) dimension - probs = op.Softmax(combined, axis=-1) # [B, q_heads, S_q, S_kv+1] - - # Drop sink column: slice axis=3 from 0 to -1 (all but last) - scores = op.Slice(probs, [0], [-1], [3]) # [B, q_heads, S_q, S_kv] - - # Weighted sum with value: [B, q_heads, S_q, d] - attn_out = op.MatMul(scores, value_exp) - - # Transpose and flatten heads: [B, q_heads, S_q, d] → [B, S_q, q_heads*d] - attn_out = op.Transpose(attn_out, perm=[0, 2, 1, 3]) # [B, S_q, q_heads, d] - attn_out = op.Reshape(attn_out, [0, 0, -1]) # [B, S_q, q_heads*d] - - # Output projection - attn_out = self.o_proj(op, attn_out) - return attn_out, present_key_value - class _GptOssDecoderLayer(nn.Module): """GPT-OSS decoder layer: pre-norm attention (with sinks) + pre-norm MoE FFN.""" diff --git a/src/mobius/models/gptoss_test.py b/src/mobius/models/gptoss_test.py index 111655230..fd54b4b82 100644 --- a/src/mobius/models/gptoss_test.py +++ b/src/mobius/models/gptoss_test.py @@ -7,10 +7,14 @@ from unittest.mock import patch +import onnx_ir as ir +import pytest import torch from mobius._testing import make_config +from mobius.components import Float32SinkAttention, SinkAttention from mobius.models.gptoss import GPTOSSCausalLMModel +from mobius.tasks import get_task class TestGPTOSSPreprocessWeightsMXFP4: @@ -90,3 +94,71 @@ def test_no_blocks_passes_through(self): result = model.preprocess_weights(state_dict) assert "model.layers.0.self_attn.q_proj.weight" in result + + +class TestGPTOSSSinkSoftmaxPrecision: + """GPT-OSS must softmax the sink in the compute dtype, not float32. + + Upstream ``gpt_oss.eager_attention_forward`` is explicit about this:: + + probs = F.softmax(combined_logits, dim=-1, dtype=combined_logits.dtype) + + GraniteSWA forces float32 instead, so the two models use different + ``SinkAttention`` subclasses. Upcasting GPT-OSS would change f16/bf16 + numerics and double the size of the largest score tensor, so this is a + regression guard, not a preference. + """ + + @staticmethod + def _build(dtype: ir.DataType): + config = make_config( + num_local_experts=2, + num_experts_per_tok=1, + layer_types=["sliding_attention", "full_attention"], + sliding_window=256, + partial_rotary_factor=1.0, + rope_interleave=False, + attn_qkv_bias=True, + attn_o_bias=True, + dtype=dtype, + ) + module = GPTOSSCausalLMModel(config) + package = get_task("text-generation").build(module, config) + return config, module, package["model"].graph + + def test_uses_the_compute_dtype_sink_attention(self): + _, module, _ = self._build(ir.DataType.FLOAT16) + for layer in module.model.layers: + assert isinstance(layer.self_attn, SinkAttention) + assert not isinstance(layer.self_attn, Float32SinkAttention) + assert layer.self_attn.upcast_sink_softmax is False + + @pytest.mark.parametrize( + "dtype", + [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16], + ids=["f32", "f16", "bf16"], + ) + def test_no_cast_on_the_sink_softmax_path(self, dtype): + config, _, graph = self._build(dtype) + # GPT-OSS also has a Softmax in each MoE router; select only the + # attention sink softmaxes. + softmax_nodes = [ + node + for node in graph + if node.op_type == "Softmax" and "self_attn" in (node.name or "") + ] + assert len(softmax_nodes) == config.num_hidden_layers + for softmax in softmax_nodes: + # Softmax <- Sub(row-max stabilise) <- Concat(scores, sinks): + # no Cast may sit between the sink Concat and the Softmax. + stabilise = softmax.inputs[0].producer() + assert stabilise is not None and stabilise.op_type == "Sub" + producer = stabilise.inputs[0].producer() + assert producer is not None and producer.op_type == "Concat", ( + f"unexpected {producer.op_type if producer else None} before the " + f"sink softmax for dtype={dtype}; GPT-OSS must not upcast" + ) + # ... and none after the sink column is dropped either. + (consumer, _), *_ = softmax.outputs[0].uses() + assert consumer.op_type == "Slice" + assert all(use[0].op_type != "Cast" for use in consumer.outputs[0].uses()) diff --git a/src/mobius/models/granite_swa.py b/src/mobius/models/granite_swa.py new file mode 100644 index 000000000..2a33cba53 --- /dev/null +++ b/src/mobius/models/granite_swa.py @@ -0,0 +1,227 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""GraniteSWA: Granite with sliding-window attention and learnable attention sinks. + +GraniteSWA (``ibm-granite/granite-swash-2b``) is the Granite decoder-only +architecture plus three changes: + +* **Mixed attention span.** ``config.layer_types`` selects, per layer, either + full causal attention or a ``config.sliding_window``-wide local window. +* **Learnable per-head attention sinks.** Every layer owns a + ``sinks[num_attention_heads]`` parameter that adds one extra logit to the + softmax denominator, letting a head shed probability mass instead of being + forced to distribute it over real tokens. +* **Per-layer RoPE base.** ``config.layer_rope_theta[i]`` gives layer ``i`` its + own RoPE base frequency; a value of ``0`` marks a NoPE layer that receives no + positional rotation at all. + +Everything else follows Granite, including the four scaling multipliers +(``embedding_multiplier``, ``attention_multiplier``, ``logits_scaling``, +``residual_multiplier``). + +HuggingFace reference: ``GraniteSWAForCausalLM``. +""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING + +from onnxscript import OpBuilder, nn + +from mobius._configs import ArchitectureConfig, GraniteSwaConfig +from mobius.components import ( + Embedding, + Float32SinkAttention, + RMSNorm, + create_attention_bias, + create_decoder_layer, + initialize_rope, +) +from mobius.models.base import CausalLMModel + +if TYPE_CHECKING: + import onnx_ir as ir + + +def resolve_layer_rope_theta(config: ArchitectureConfig) -> list[float | int]: + """Return the per-layer RoPE base frequency, with ``0`` marking NoPE. + + ``layer_rope_theta`` only exists on :class:`~mobius._configs.GraniteSwaConfig`, + and even there it is optional, so a plain :class:`ArchitectureConfig` (used + by the tiny graph-build fixtures) must still resolve to something sensible. + ``no_rope_layers`` names the same set of layers, so prefer it before falling + back to rotating every layer at the global ``rope_theta``. + """ + layer_rope_theta = getattr(config, "layer_rope_theta", None) + if layer_rope_theta: + return list(layer_rope_theta) + + no_rope_layers = set(config.no_rope_layers or ()) + return [ + 0 if index in no_rope_layers else (config.rope_theta or 0) + for index in range(config.num_hidden_layers) + ] + + +class GraniteSwaTextModel(nn.Module): + """GraniteSWA backbone: mixed full/sliding attention with per-layer RoPE. + + Replicates HuggingFace ``GraniteSWAModel``: scale the embeddings by + ``embedding_multiplier``, build one causal mask per attention span, build one + ``(cos, sin)`` table per distinct non-zero RoPE base, and dispatch both + per layer. + """ + + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.config = config + self._dtype = config.dtype + self._layer_types = config.layer_types or [ + "full_attention" if index % 4 == 0 else "sliding_attention" + for index in range(config.num_hidden_layers) + ] + self._layer_rope_theta = resolve_layer_rope_theta(config) + self._sliding_window = config.sliding_window + self.embedding_multiplier = config.embedding_multiplier + + self.embed_tokens = Embedding( + config.vocab_size, config.hidden_size, config.pad_token_id + ) + self.layers = nn.ModuleList( + [ + # Float32SinkAttention, not SinkAttention: GraniteSWA's eager + # kernel forces the sink scaling and the softmax to float32. + create_decoder_layer(config, attention_class=Float32SinkAttention) + for _ in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + # One rotary module per distinct non-zero base theta, mirroring HF's + # ``GraniteSWAModel.rotary_embs``. ``sorted`` keeps construction (and + # therefore initializer naming) deterministic across builds. + self._rope_thetas = sorted({theta for theta in self._layer_rope_theta if theta}) + self.rotary_embs = nn.ModuleList( + [ + initialize_rope(dataclasses.replace(config, rope_theta=float(theta))) + for theta in self._rope_thetas + ] + ) + # ``TextModel``-style single-rope attribute, kept so that generic + # tooling that expects ``model.rotary_emb`` (metadata emitters, GQA + # rewrite guards) sees the common case. ``None`` when every layer is + # NoPE. + self.rotary_emb = self.rotary_embs[0] if self.rotary_embs else None + + def forward( + self, + op: OpBuilder, + input_ids: ir.Value | None, + attention_mask: ir.Value, + position_ids: ir.Value, + past_key_values: list | None = None, + inputs_embeds: ir.Value | None = None, + ): + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_tokens(op, input_ids) + # Granite embedding multiplier, applied before the first decoder layer. + hidden_states = op.Mul(hidden_states, self.embedding_multiplier) + + # SinkAttention builds the score matrix explicitly, so the bias must + # carry the FULL mask (causal + window + padding); there is no + # ``is_causal`` flag to fall back on. + mask_source = position_ids if input_ids is None else input_ids + full_attention_bias = create_attention_bias( + op, + input_ids=mask_source, + attention_mask=attention_mask, + dtype=self._dtype, + ) # (B, 1, S_q, S_kv) + sliding_attention_bias = full_attention_bias + if self._sliding_window and "sliding_attention" in self._layer_types: + sliding_attention_bias = create_attention_bias( + op, + input_ids=mask_source, + attention_mask=attention_mask, + sliding_window=self._sliding_window, + dtype=self._dtype, + ) # (B, 1, S_q, S_kv) + + # (cos, sin) per distinct non-zero base theta; NoPE layers get None. + position_embeddings_by_theta = { + theta: rotary_emb(op, position_ids) + for theta, rotary_emb in zip(self._rope_thetas, self.rotary_embs) + } + + present_key_values = [] + past_kvs = past_key_values or [None] * len(self.layers) + for layer_idx, (layer, past_key_value) in enumerate(zip(self.layers, past_kvs)): + is_sliding = self._layer_types[layer_idx] == "sliding_attention" + theta = self._layer_rope_theta[layer_idx] + hidden_states, present_key_value = layer( + op, + hidden_states=hidden_states, + attention_bias=(sliding_attention_bias if is_sliding else full_attention_bias), + position_embeddings=position_embeddings_by_theta.get(theta), + past_key_value=past_key_value, + ) + present_key_values.append(present_key_value) + + hidden_states = self.norm(op, hidden_states) + return hidden_states, present_key_values + + +class GraniteSwaCausalLMModel(CausalLMModel): + """GraniteSWA causal LM: Granite scaling + sliding windows + attention sinks. + + Extends the Granite architecture with per-layer sliding-window attention, + a learnable per-head attention sink folded into the softmax denominator, + and a per-layer RoPE base frequency (``0`` = NoPE). The four Granite + scaling multipliers still apply: + + - ``embedding_multiplier``: scales embeddings after lookup + - ``attention_multiplier``: replaces ``1/sqrt(head_dim)`` as attention scale + - ``residual_multiplier``: scales attention/MLP outputs before residual add + - ``logits_scaling``: divides the final logits + + HuggingFace model_type: ``granite_swa`` (``GraniteSWAForCausalLM``). + """ + + # Declared here as well as on the registry entry: the registry-consistency + # check requires the two to agree, and inheriting CausalLMConfig would + # silently drop ``layer_rope_theta``. + config_class: type = GraniteSwaConfig + + def __init__(self, config: ArchitectureConfig): + super().__init__(config) + self.model = GraniteSwaTextModel(config) + # ``CausalLMModel.__init__`` tied ``lm_head.weight`` to the *base* + # ``TextModel``'s embedding table, which the assignment above just + # replaced. Re-tie against the real backbone so the graph keeps a + # single shared initializer instead of two identical copies. + quantization = getattr(config, "quantization", None) + embed_quantized = quantization is not None and getattr( + quantization, "quantize_embeddings", False + ) + if config.tie_word_embeddings and not embed_quantized: + self.lm_head.weight = self.model.embed_tokens.weight + self.logits_scaling = config.logits_scaling + + def forward( + self, + op: OpBuilder, + input_ids: ir.Value, + attention_mask: ir.Value, + position_ids: ir.Value, + past_key_values: list | None = None, + ): + logits, present_key_values = super().forward( + op, input_ids, attention_mask, position_ids, past_key_values + ) + # Granite logits scaling + logits = op.Div(logits, self.logits_scaling) + return logits, present_key_values diff --git a/src/mobius/models/granite_swa_test.py b/src/mobius/models/granite_swa_test.py new file mode 100644 index 000000000..79bff6ee7 --- /dev/null +++ b/src/mobius/models/granite_swa_test.py @@ -0,0 +1,464 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Unit tests for the GraniteSWA architecture (``granite_swa``). + +Covers the three ways GraniteSWA departs from Granite — mixed +full/sliding attention spans, learnable per-head attention sinks, and a +per-layer RoPE base with ``0`` meaning NoPE — plus the config extraction +that feeds them. +""" + +from __future__ import annotations + +import dataclasses +import types + +import numpy as np +import onnx_ir as ir +import pytest +import torch + +from mobius._configs import ArchitectureConfig, GraniteSwaConfig +from mobius._registry import registry +from mobius._testing import count_op_type, create_test_builder, create_test_input +from mobius._testing.ort_inference import OnnxModelSession +from mobius.components import Float32SinkAttention, SinkAttention +from mobius.components._attention import GQAContext +from mobius.models.granite_swa import ( + GraniteSwaCausalLMModel, + GraniteSwaTextModel, + resolve_layer_rope_theta, +) +from mobius.tasks import get_task + +# The pinned ``ibm-granite/granite-swash-2b`` config.json contents at revision +# af1e3227100b61088eead48389ab5409b5d0e39c, inlined so the extraction test +# needs no network access. +PINNED_REVISION = "af1e3227100b61088eead48389ab5409b5d0e39c" +_PINNED_CONFIG_JSON: dict = { + "model_type": "granite_swa", + "architectures": ["GraniteSWAForCausalLM"], + "vocab_size": 100352, + "hidden_size": 2560, + "intermediate_size": 8192, + "num_hidden_layers": 24, + "num_attention_heads": 20, + "num_key_value_heads": 4, + "hidden_act": "silu", + "max_position_embeddings": 8192, + "rms_norm_eps": 1e-05, + "attention_bias": False, + "mlp_bias": False, + "tie_word_embeddings": True, + "initializer_range": 0.1, + "bos_token_id": 100257, + "eos_token_id": 100257, + "pad_token_id": 100256, + "sliding_window": 128, + "layer_types": ["full_attention"] + + ["sliding_attention", "sliding_attention", "sliding_attention", "full_attention"] * 5 + + ["sliding_attention", "sliding_attention", "sliding_attention"], + "embedding_multiplier": 12, + "residual_multiplier": 0.28, + "logits_scaling": 10, + "attention_multiplier": 0.0078125, + "rope_theta": 10000, +} + + +def _tiny_config(_config_cls: type[GraniteSwaConfig] = GraniteSwaConfig, **overrides): + """A 4-layer GraniteSWA config exercising every per-layer dispatch.""" + fields = dict( + vocab_size=64, + max_position_embeddings=32, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + hidden_act="silu", + pad_token_id=0, + rope_type="default", + rope_theta=10_000.0, + sliding_window=4, + layer_types=[ + "full_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ], + layer_rope_theta=[10_000.0, 10_000.0, 0, 500_000.0], + embedding_multiplier=12.0, + attention_multiplier=0.0078125, + logits_scaling=10.0, + residual_multiplier=0.28, + dtype=ir.DataType.FLOAT, + ) + fields.update(overrides) + return _config_cls(**fields) + + +class TestConfigExtraction: + """``GraniteSwaConfig.from_transformers`` must match HF ``__post_init__``.""" + + def test_pinned_checkpoint_config(self): + hf_config = types.SimpleNamespace(**_PINNED_CONFIG_JSON) + config = GraniteSwaConfig.from_transformers(hf_config) + + assert config.num_hidden_layers == 24 + assert config.hidden_size == 2560 + assert config.num_attention_heads == 20 + assert config.num_key_value_heads == 4 + # head_dim is derived: 2560 / 20, matching HF LlamaAttention's default. + assert config.head_dim == 128 + assert config.sliding_window == 128 + assert config.tie_word_embeddings is True + assert config.rms_norm_eps == pytest.approx(1e-5) + # Granite scaling multipliers survive extraction verbatim. + assert config.embedding_multiplier == pytest.approx(12.0) + assert config.residual_multiplier == pytest.approx(0.28) + assert config.logits_scaling == pytest.approx(10.0) + assert config.attention_multiplier == pytest.approx(0.0078125) + # Every fourth layer is full attention in this checkpoint. + assert config.layer_types is not None + assert len(config.layer_types) == 24 + assert [i for i, t in enumerate(config.layer_types) if t == "full_attention"] == [ + 0, + 4, + 8, + 12, + 16, + 20, + ] + # No layer_rope_theta in the checkpoint → global theta on every layer. + assert config.layer_rope_theta == [10_000.0] * 24 + assert config.no_rope_layers == [] + + def test_layer_types_default_when_absent(self): + """HF defaults full attention on every fourth layer.""" + raw = dict(_PINNED_CONFIG_JSON) + raw.pop("layer_types") + raw["num_hidden_layers"] = 6 + config = GraniteSwaConfig.from_transformers(types.SimpleNamespace(**raw)) + + assert config.layer_types == [ + "full_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + ] + + def test_zero_theta_marks_a_nope_layer(self): + raw = dict(_PINNED_CONFIG_JSON) + raw["num_hidden_layers"] = 4 + raw["layer_types"] = ["full_attention"] * 4 + raw["layer_rope_theta"] = [10_000.0, 0, 500_000.0, 0] + config = GraniteSwaConfig.from_transformers(types.SimpleNamespace(**raw)) + + assert config.layer_rope_theta == [10_000.0, 0, 500_000.0, 0] + assert config.no_rope_layers == [1, 3] + + def test_registry_wiring(self): + assert registry.get("granite_swa") is GraniteSwaCausalLMModel + assert registry.get_config_class("granite_swa") is GraniteSwaConfig + # The registry entry and the model class must agree; inheriting + # CausalLMConfig would silently drop ``layer_rope_theta``. + assert GraniteSwaCausalLMModel.config_class is GraniteSwaConfig + registration = registry.get_registration("granite_swa") + assert registration.test_model_id == "ibm-granite/granite-swash-2b" + + +class TestResolveLayerRopeTheta: + def test_explicit_list_wins(self): + config = _tiny_config() + assert resolve_layer_rope_theta(config) == [10_000.0, 10_000.0, 0, 500_000.0] + + def test_falls_back_to_no_rope_layers(self): + """A plain ArchitectureConfig has no ``layer_rope_theta`` field at all.""" + config = ArchitectureConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + rope_type="default", + rope_theta=10_000.0, + no_rope_layers=[2], + ) + assert resolve_layer_rope_theta(config) == [10_000.0, 10_000.0, 0, 10_000.0] + + def test_rotates_every_layer_by_default(self): + config = ArchitectureConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=3, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + rope_type="default", + rope_theta=10_000.0, + ) + assert resolve_layer_rope_theta(config) == [10_000.0] * 3 + + +class TestModuleStructure: + def test_one_rotary_module_per_distinct_nonzero_theta(self): + model = GraniteSwaTextModel(_tiny_config()) + # thetas {10_000, 500_000}; the NoPE layer contributes none. + assert len(model.rotary_embs) == 2 + assert model._rope_thetas == [10_000.0, 500_000.0] + assert model.rotary_emb is model.rotary_embs[0] + + def test_sink_parameter_names_match_huggingface(self): + config = _tiny_config() + module = GraniteSwaCausalLMModel(config) + package = get_task("text-generation").build(module, config) + names = set(package["model"].graph.initializers) + + for layer in range(config.num_hidden_layers): + name = f"model.layers.{layer}.self_attn.sinks" + assert name in names + assert list(package["model"].graph.initializers[name].shape) == [ + config.num_attention_heads + ] + + def test_tied_embeddings_share_one_initializer(self): + config = _tiny_config(tie_word_embeddings=True) + module = GraniteSwaCausalLMModel(config) + package = get_task("text-generation").build(module, config) + names = set(package["model"].graph.initializers) + + assert "model.embed_tokens.weight" in names + assert "lm_head.weight" not in names + + def test_attention_uses_the_granite_multiplier_as_scale(self): + config = _tiny_config() + module = GraniteSwaCausalLMModel(config) + for layer in module.model.layers: + # Float32SinkAttention, not the base SinkAttention: GraniteSWA's + # eager kernel forces the sink scaling and softmax to float32, + # whereas GPT-OSS deliberately stays in the compute dtype. + assert isinstance(layer.self_attn, Float32SinkAttention) + assert layer.self_attn.upcast_sink_softmax is True + assert layer.self_attn.scaling == pytest.approx(config.attention_multiplier) + + def test_nope_layer_skips_rotary_embedding(self): + """Three RoPE layers, two RotaryEmbedding nodes each (q, k) = 6, not 8.""" + config = _tiny_config() + module = GraniteSwaCausalLMModel(config) + package = get_task("text-generation").build(module, config) + assert count_op_type(package["model"].graph, "RotaryEmbedding") == 6 + + def test_all_nope_config_builds_without_rotary_modules(self): + config = _tiny_config(layer_rope_theta=[0, 0, 0, 0]) + module = GraniteSwaCausalLMModel(config) + assert len(module.model.rotary_embs) == 0 + assert module.model.rotary_emb is None + package = get_task("text-generation").build(module, config) + assert count_op_type(package["model"].graph, "RotaryEmbedding") == 0 + + def test_sliding_layers_use_a_distinct_bias_from_full_layers(self): + config = _tiny_config() + module = GraniteSwaCausalLMModel(config) + package = get_task("text-generation").build(module, config) + graph = package["model"].graph + # ``create_attention_bias`` only emits a ``Less`` comparison for the + # window term, so exactly one of the two biases is windowed. + assert count_op_type(graph, "Less") == 1 + + def test_uniform_full_attention_builds_a_single_bias(self): + config = _tiny_config( + layer_types=["full_attention"] * 4, + ) + module = GraniteSwaCausalLMModel(config) + package = get_task("text-generation").build(module, config) + assert count_op_type(package["model"].graph, "Less") == 0 + + +class TestSinkAttention: + """The sink must behave exactly like HF's ``sigmoid(logsumexp - sink)``.""" + + def test_rejects_gqa_context(self): + config = _tiny_config() + attn = SinkAttention(config) + builder, op, _ = create_test_builder() + hidden = create_test_input(builder, "hidden", [1, 4, config.hidden_size]) + ctx = GQAContext( + seqlens_k=create_test_input(builder, "seqlens_k", [1], ir.DataType.INT32), + total_seq_len=create_test_input(builder, "tsl", [], ir.DataType.INT32), + cos_cache=create_test_input(builder, "cos", [32, 4]), + sin_cache=create_test_input(builder, "sin", [32, 4]), + ) + with pytest.raises(TypeError, match="cannot emit GroupQueryAttention"): + attn(op, hidden, attention_bias=ctx) + + def test_rejects_softcapping(self): + @dataclasses.dataclass + class _SoftcappedConfig(GraniteSwaConfig): + attn_logit_softcapping: float = 30.0 + + config = _tiny_config(_config_cls=_SoftcappedConfig) + assert config.attn_logit_softcapping == pytest.approx(30.0) + with pytest.raises(ValueError, match="softcapping"): + SinkAttention(config) + + def test_emits_no_fused_attention_op(self): + """The sink lives inside the softmax, so no fused op may appear.""" + config = _tiny_config() + module = GraniteSwaCausalLMModel(config) + package = get_task("text-generation").build(module, config) + graph = package["model"].graph + assert count_op_type(graph, "Attention") == 0 + assert count_op_type(graph, "GroupQueryAttention") == 0 + assert count_op_type(graph, "Softmax") == config.num_hidden_layers + + @pytest.mark.parametrize( + "dtype", [ir.DataType.FLOAT16, ir.DataType.BFLOAT16], ids=["f16", "bf16"] + ) + def test_reduced_precision_upcasts_the_sink_softmax(self, dtype): + """Mirror upstream's forced-fp32 sink softmax for f16/bf16 builds. + + HuggingFace computes ``logsumexp``/``sigmoid`` and the softmax in + float32 regardless of the compute dtype. The equivalent here is to + upcast the extended (scores + sink) logits before the softmax and cast + the probabilities back, which shows up as two extra Cast nodes per + layer around the Softmax. + """ + config = _tiny_config(dtype=dtype) + module = GraniteSwaCausalLMModel(config) + graph = get_task("text-generation").build(module, config)["model"].graph + + softmax_nodes = [node for node in graph if node.op_type == "Softmax"] + assert len(softmax_nodes) == config.num_hidden_layers + for softmax in softmax_nodes: + # Softmax <- Sub(row-max stabilise) <- Cast(to float32) + stabilise = softmax.inputs[0].producer() + assert stabilise is not None and stabilise.op_type == "Sub" + upcast = stabilise.inputs[0].producer() + assert upcast is not None and upcast.op_type == "Cast" + assert upcast.attributes["to"].as_int() == ir.DataType.FLOAT + # Softmax -> Slice(drop sink column) -> Cast(back to compute dtype) + (consumer, _), *_ = softmax.outputs[0].uses() + assert consumer.op_type == "Slice" + (downcast, _), *_ = consumer.outputs[0].uses() + assert downcast.op_type == "Cast" + assert downcast.attributes["to"].as_int() == dtype + + def test_float32_build_has_no_sink_softmax_casts(self): + """float32 builds must stay Cast-free around the sink softmax.""" + config = _tiny_config() + module = GraniteSwaCausalLMModel(config) + graph = get_task("text-generation").build(module, config)["model"].graph + + softmax_nodes = [node for node in graph if node.op_type == "Softmax"] + assert len(softmax_nodes) == config.num_hidden_layers + for softmax in softmax_nodes: + stabilise = softmax.inputs[0].producer() + assert stabilise is not None and stabilise.op_type == "Sub" + # The stabilised logits come straight from the sink Concat. + producer = stabilise.inputs[0].producer() + assert producer is not None and producer.op_type == "Concat" + (consumer, _), *_ = softmax.outputs[0].uses() + assert consumer.op_type == "Slice" + assert all(use[0].op_type != "Cast" for use in consumer.outputs[0].uses()) + + def test_matches_the_logsumexp_sigmoid_reference(self): + """Run one SinkAttention layer in ORT against the upstream formula. + + HuggingFace ``granite_swa.eager_attention_forward`` computes an + ordinary softmax and then rescales the output by + ``sigmoid(logsumexp(scores) - sink)``. This implementation instead + keeps the sink as an extra softmax column. Both must agree. + """ + rng = np.random.default_rng(0) + config = _tiny_config(num_hidden_layers=1) + num_heads = config.num_attention_heads + num_kv_heads = config.num_key_value_heads + head_dim = config.head_dim + hidden_size = config.hidden_size + batch, seq_len = 1, 5 + + attn = Float32SinkAttention(config, scale=config.attention_multiplier) + builder, op, graph = create_test_builder() + hidden = create_test_input(builder, "hidden", [batch, seq_len, hidden_size]) + bias = create_test_input(builder, "bias", [batch, 1, seq_len, seq_len]) + output, _ = attn(op, hidden, attention_bias=bias, position_embeddings=None) + builder._adapt_outputs([output], "") + output.name = "attn_out" + graph.outputs.append(output) + + weights = {} + for name, out_features in ( + ("q_proj", num_heads * head_dim), + ("k_proj", num_kv_heads * head_dim), + ("v_proj", num_kv_heads * head_dim), + ("o_proj", hidden_size), + ): + in_features = num_heads * head_dim if name == "o_proj" else hidden_size + weights[name] = rng.standard_normal((out_features, in_features)).astype(np.float32) + weights["sinks"] = rng.standard_normal(num_heads).astype(np.float32) + + for value in graph.initializers.values(): + # Skip the folded scalar/shape constants; only the four projection + # weights and the sink vector are unset parameters. + if value.const_value is not None: + continue + key = "sinks" if value.name.endswith("sinks") else value.name.split(".")[0] + value.const_value = ir.Tensor(weights[key]) + + model = ir.Model(graph, ir_version=10) + hidden_states = rng.standard_normal((batch, seq_len, hidden_size)).astype(np.float32) + # Causal float additive bias, matching create_attention_bias's output. + mask = np.triu(np.full((seq_len, seq_len), np.float32(-3.4e38)), k=1) + attention_bias = mask.reshape(1, 1, seq_len, seq_len) + + session = OnnxModelSession(model) + try: + actual = session.run({"hidden": hidden_states, "bias": attention_bias}) + finally: + session.close() + onnx_out = next(iter(actual.values())) + + # --- Upstream reference (torch, sigmoid-LSE form) --- + hs = torch.from_numpy(hidden_states) + query = ( + (hs @ torch.from_numpy(weights["q_proj"]).T) + .view(batch, seq_len, num_heads, head_dim) + .transpose(1, 2) + ) + key = ( + (hs @ torch.from_numpy(weights["k_proj"]).T) + .view(batch, seq_len, num_kv_heads, head_dim) + .transpose(1, 2) + ) + value = ( + (hs @ torch.from_numpy(weights["v_proj"]).T) + .view(batch, seq_len, num_kv_heads, head_dim) + .transpose(1, 2) + ) + groups = num_heads // num_kv_heads + key = key.repeat_interleave(groups, dim=1) + value = value.repeat_interleave(groups, dim=1) + + scores = (query @ key.transpose(2, 3)) * config.attention_multiplier + scores = scores + torch.from_numpy(attention_bias) + lse = torch.logsumexp(scores, dim=-1) # (B, H, S) + sink_scale = ( + (lse - torch.from_numpy(weights["sinks"]).view(1, -1, 1)) + .to(torch.float32) + .sigmoid() + ) + probs = torch.softmax(scores, dim=-1, dtype=torch.float32).to(scores.dtype) + expected = probs @ value + expected = expected * sink_scale.unsqueeze(-1) + expected = expected.transpose(1, 2).reshape(batch, seq_len, num_heads * head_dim) + expected = expected @ torch.from_numpy(weights["o_proj"]).T + + np.testing.assert_allclose(onnx_out, expected.numpy(), rtol=1e-4, atol=1e-4) diff --git a/testdata/cases/causal-lm/granite-swash-2b.yaml b/testdata/cases/causal-lm/granite-swash-2b.yaml new file mode 100644 index 000000000..6d24b5485 --- /dev/null +++ b/testdata/cases/causal-lm/granite-swash-2b.yaml @@ -0,0 +1,23 @@ +model_id: "ibm-granite/granite-swash-2b" +model_type: "granite_swa" +revision: "af1e3227100b61088eead48389ab5409b5d0e39c" +task_type: "text-generation" +dtype: "float32" + +inputs: + prompts: + - "Here is my poem:" + +level: "L4+L5" + +generation: + max_new_tokens: 20 + do_sample: false + +notes: > + Granite SWA 2B (GraniteSWAForCausalLM). Mixed full/sliding-window attention + layers, learnable per-head attention sinks, per-layer RoPE theta, and the + Granite embedding/attention/residual/logit scaling multipliers. The HF + reference is pinned to eager attention because the sink is an extra logit in + the softmax denominator, which SDPA cannot express + (upstream sets ``GraniteSWAPreTrainedModel._supports_sdpa = False``). diff --git a/testdata/golden/causal-lm/granite-swash-2b.json b/testdata/golden/causal-lm/granite-swash-2b.json new file mode 100644 index 000000000..db345f9d3 --- /dev/null +++ b/testdata/golden/causal-lm/granite-swash-2b.json @@ -0,0 +1,41 @@ +{ + "top1_id": 330, + "top2_id": 2355, + "top10_ids": [ + 330, + 2355, + 578, + 720, + 100257, + 358, + 4815, + 19124, + 362, + 1054 + ], + "top10_logits": [ + "0x1.b0257c0000000p+3", + "0x1.a61ae60000000p+3", + "0x1.9947ec0000000p+3", + "0x1.8e445c0000000p+3", + "0x1.89f2c00000000p+3", + "0x1.8584540000000p+3", + "0x1.8125140000000p+3", + "0x1.7d47da0000000p+3", + "0x1.7552ea0000000p+3", + "0x1.6fbab60000000p+3" + ], + "logits_summary": [ + "0x1.b0257c0000000p+3", + "-0x1.184c640000000p+3", + "-0x1.37d419ed04e18p-1", + "0x1.49de7a500aca8p+1" + ], + "input_ids": [ + 8586, + 374, + 856, + 33894, + 25 + ] +} diff --git a/testdata/golden/causal-lm/granite-swash-2b_generation.json b/testdata/golden/causal-lm/granite-swash-2b_generation.json new file mode 100644 index 000000000..1458298fd --- /dev/null +++ b/testdata/golden/causal-lm/granite-swash-2b_generation.json @@ -0,0 +1,27 @@ +{ + "model_id": "ibm-granite/granite-swash-2b", + "prompt": "Here is my poem:", + "generated_tokens": [ + 330, + 791, + 8155, + 315, + 279, + 20409, + 38473, + 1, + 555, + 7957, + 45406, + 22096, + 24421, + 13, + 1102, + 374, + 264, + 33894, + 922, + 279 + ], + "generated_text": " \"The Last of the Mohicans\" by James Fenimore Cooper. It is a poem about the" +} diff --git a/tests/_test_configs.py b/tests/_test_configs.py index cbcc7ad6b..c749b5514 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -31,6 +31,7 @@ Gemma4Config, GlmAsrConfig, GraniteMoeHybridConfig, + GraniteSwaConfig, JambaConfig, JetMoeConfig, Lfm2Config, @@ -443,6 +444,31 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: # (Gemma3nMultiModalModel), so its entry lives in VL_CONFIGS. The text # decoder is covered by the two "gemma3n_text" entries. ("granite", {}, True), + # granite_swa: mixed full/sliding layers, learnable per-head attention + # sinks, and a per-layer RoPE base (0 = NoPE) on top of Granite's scaling + # multipliers. 4 layers so the tiny config can carry two distinct + # non-zero thetas plus a NoPE layer and still alternate attention spans. + ( + "granite_swa", + { + "_config_cls": GraniteSwaConfig, + "num_hidden_layers": 4, + "layer_types": [ + "full_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ], + "layer_rope_theta": [10_000.0, 10_000.0, 0, 500_000.0], + "sliding_window": 8, + "tie_word_embeddings": True, + "embedding_multiplier": 12.0, + "attention_multiplier": 0.0625, + "logits_scaling": 10.0, + "residual_multiplier": 0.28, + }, + True, + ), ("olmo", {}, False), ("internlm2", {"attn_qkv_bias": True}, True), ( diff --git a/tests/integration_test.py b/tests/integration_test.py index 0c535dd28..a97ea8d76 100644 --- a/tests/integration_test.py +++ b/tests/integration_test.py @@ -4973,6 +4973,163 @@ def test_bamba_prefill_logits_match(): assert_logits_close(onnx_outputs["logits"], hf_logits, rtol=1e-3, atol=1e-3) +# --------------------------------------------------------------------------- +# GraniteSWA real-weight parity +# +# The whole point of GraniteSWA's sliding layers only shows up once the prompt +# is longer than ``sliding_window`` (128 for granite-swash-2b): below that, a +# windowed layer and a full-attention layer see the same keys. These tests +# therefore use a prompt of >128 tokens so a window bug cannot hide, and pin +# the HuggingFace reference to the eager kernel because the learnable sink is +# an extra logit inside the softmax denominator that SDPA cannot express. +# --------------------------------------------------------------------------- + +_GRANITE_SWA_MODEL_ID = "ibm-granite/granite-swash-2b" +_GRANITE_SWA_REVISION = "af1e3227100b61088eead48389ab5409b5d0e39c" + + +def _granite_swa_long_prompt(tokenizer, min_tokens: int) -> np.ndarray: + """Tokenize a prompt guaranteed to exceed ``min_tokens`` tokens.""" + sentence = ( + "The city archives recorded every harvest, every flood, and every " + "quiet year in between, so that later readers could trace how the " + "valley changed. " + ) + text = sentence + while True: + input_ids = tokenizer(text, return_tensors="np")["input_ids"].astype(np.int64) + if input_ids.shape[1] > min_tokens: + return input_ids + text += sentence + + +def _load_granite_swa_reference(): + """Load the pinned GraniteSWA checkpoint with the eager attention kernel.""" + model, tokenizer = load_torch_model( + _GRANITE_SWA_MODEL_ID, + revision=_GRANITE_SWA_REVISION, + trust_remote_code=False, + attn_implementation="eager", + ) + # load_torch_model raises on mismatch, but assert here too so the intent of + # this test is visible at the call site. + assert model.config._attn_implementation == "eager" + return model, tokenizer + + +@pytest.mark.integration +def test_granite_swa_prefill_logits_match(): + """Prefill past the sliding window matches HuggingFace eager attention.""" + onnx_model = build( + _GRANITE_SWA_MODEL_ID, + revision=_GRANITE_SWA_REVISION, + dtype="f32", + load_weights=True, + ) + torch_model, tokenizer = _load_granite_swa_reference() + config = _get_config(_GRANITE_SWA_MODEL_ID) + + sliding_window = config.sliding_window + assert sliding_window == 128 + assert config.layer_types is not None + assert "sliding_attention" in config.layer_types + assert "full_attention" in config.layer_types + + input_ids = _granite_swa_long_prompt(tokenizer, min_tokens=sliding_window + 32) + seq_len = input_ids.shape[1] + attention_mask = np.ones((1, seq_len), dtype=np.int64) + position_ids = np.arange(seq_len, dtype=np.int64)[np.newaxis, :] + + torch_logits, _ = torch_forward(torch_model, input_ids, attention_mask, position_ids) + + session = _make_session(onnx_model) + try: + onnx_outputs = session.run( + _make_prefill_feeds(config, input_ids, attention_mask, position_ids) + ) + finally: + session.close() + + assert_logits_close(onnx_outputs["logits"], torch_logits, rtol=1e-3, atol=1e-3) + + +@pytest.mark.integration +def test_granite_swa_decode_step_logits_match(): + """A cached decode step past the window still matches HuggingFace. + + HuggingFace and mobius represent the sliding-layer KV cache differently and + both are correct: HF's ``DynamicSlidingWindowLayer`` physically drops keys + that have fallen out of the window, while the exported graph keeps the full + cache and masks the same keys out through the sliding attention bias. The + two cannot be cross-fed, so each side runs prefill and decode against its + own cache and only the decode logits are compared. + """ + onnx_model = build( + _GRANITE_SWA_MODEL_ID, + revision=_GRANITE_SWA_REVISION, + dtype="f32", + load_weights=True, + ) + torch_model, tokenizer = _load_granite_swa_reference() + config = _get_config(_GRANITE_SWA_MODEL_ID) + + input_ids = _granite_swa_long_prompt(tokenizer, min_tokens=config.sliding_window + 32) + seq_len = input_ids.shape[1] + attention_mask = np.ones((1, seq_len), dtype=np.int64) + position_ids = np.arange(seq_len, dtype=np.int64)[np.newaxis, :] + + # --- HuggingFace: prefill, then decode against its own Cache object --- + with torch.no_grad(): + hf_prefill = torch_model( + input_ids=torch.from_numpy(input_ids), + attention_mask=torch.from_numpy(attention_mask), + position_ids=torch.from_numpy(position_ids), + use_cache=True, + ) + hf_cache = hf_prefill.past_key_values + + # Document the representation difference this test works around: sliding + # layers are physically trimmed by HF, full-attention layers are not. + sliding_idx = config.layer_types.index("sliding_attention") + full_idx = config.layer_types.index("full_attention") + assert hf_cache.layers[sliding_idx].keys.shape[2] < seq_len + assert hf_cache.layers[full_idx].keys.shape[2] == seq_len + + next_token = np.array([[int(np.argmax(hf_prefill.logits[0, -1].numpy()))]], dtype=np.int64) + decode_attention_mask = np.ones((1, seq_len + 1), dtype=np.int64) + decode_position_ids = np.array([[seq_len]], dtype=np.int64) + + with torch.no_grad(): + hf_decode = torch_model( + input_ids=torch.from_numpy(next_token), + attention_mask=torch.from_numpy(decode_attention_mask), + position_ids=torch.from_numpy(decode_position_ids), + past_key_values=hf_cache, + use_cache=True, + ) + torch_logits_2 = hf_decode.logits.numpy() + + # --- ONNX: prefill, then decode against the full exported cache --- + session = _make_session(onnx_model) + try: + onnx_out_1 = session.run( + _make_prefill_feeds(config, input_ids, attention_mask, position_ids) + ) + onnx_out_2 = session.run( + _make_decode_feeds( + config, + next_token, + decode_attention_mask, + decode_position_ids, + onnx_out_1, + ) + ) + finally: + session.close() + + assert_logits_close(onnx_out_2["logits"], torch_logits_2, rtol=1e-3, atol=1e-3) + + # --------------------------------------------------------------------------- # Encoder-only parity tests (BERT, DistilBERT, RoBERTa) # diff --git a/tests/synthetic_parity_test.py b/tests/synthetic_parity_test.py index 3b60e623a..1bb24dd86 100644 --- a/tests/synthetic_parity_test.py +++ b/tests/synthetic_parity_test.py @@ -459,6 +459,10 @@ "head_dim": TINY_HEAD_DIM, "layer_types": ["sliding_attention", "full_attention"], }, + # GraniteSWA: HF defaults bos/eos to 100257 (the granite-swash tokenizer), + # which is outside the tiny 256-token vocab. Pin them in range so the + # reference config is self-consistent. + "granite_swa": {"bos_token_id": 1, "eos_token_id": 2}, # VL MoE text sub-models: need same HF extras as their base model types. # qwen3_vl_moe, qwen3_omni_moe → qwen3_moe (needs head_dim + moe_intermediate_size) "qwen3_vl_moe": {"head_dim": TINY_HEAD_DIM, "moe_intermediate_size": TINY_INTERMEDIATE}, @@ -740,8 +744,27 @@ def forward(self, **kwargs): return _SoftcappedBackboneCausalLM() +def _create_eager_causal_lm(hf_config): + """Build an HF causal LM pinned to the eager attention kernel. + + Required for attention-sink models (GraniteSWA): the sink is an extra logit + inside the softmax denominator, which SDPA cannot express. Upstream marks + ``GraniteSWAPreTrainedModel._supports_sdpa = False`` for exactly this + reason, but pin it explicitly here so the reference can never silently + drift onto a non-sink kernel. + """ + from transformers import AutoModelForCausalLM + + model = AutoModelForCausalLM.from_config(hf_config, attn_implementation="eager") + assert model.config._attn_implementation == "eager", ( + f"expected eager attention, got {model.config._attn_implementation!r}" + ) + return model + + _HF_MODEL_FACTORIES = { "muse_glimmer_text": _create_softcapped_backbone_causal_lm, + "granite_swa": _create_eager_causal_lm, }