Add GraniteSWA (granite_swa) with a shared SinkAttention component

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)

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",

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."""

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",

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()
 

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 (