Support GraniteMoeHybrid dense variant (granite-4.0-1b) with L4/L5 tests

scripts/generate_golden.py

@@ -237,12 +237,23 @@ def _generate_causal_lm(case: TestCase, json_path: Path, device: str) -> None:
         _apply_nemotron_h_generate_patch(model)
 
         model_device = _get_model_device(model, device)
+        gen_ids = torch.from_numpy(input_ids).to(model_device)
+        max_new = case.generation_params.get("max_new_tokens", 20)
         with torch.no_grad():
-            gen_output = model.generate(
-                torch.from_numpy(input_ids).to(model_device),
-                max_new_tokens=case.generation_params.get("max_new_tokens", 20),
-                do_sample=False,
-            )
+            try:
+                gen_output = model.generate(gen_ids, max_new_tokens=max_new, do_sample=False)
+            except ValueError as e:
+                # All-attention GraniteMoeHybrid variants (e.g. granite-4.0-1b)
+                # trip transformers' hybrid Mamba/attention generation cache,
+                # which assumes at least one linear-attention (Mamba) layer:
+                # "`has_previous_state` can only be called on LinearAttention
+                # layers". Greedy output is cache-independent, so fall back to
+                # the (slower) cache-free path.
+                if "has_previous_state" not in str(e):
+                    raise
+                gen_output = model.generate(
+                    gen_ids, max_new_tokens=max_new, do_sample=False, use_cache=False
+                )
         generated_ids = gen_output[0, seq_len:].cpu().numpy()
 
     save_golden_ref(

src/mobius/_configs/_base.py

@@ -2226,10 +2226,39 @@ class GraniteMoeHybridConfig(BambaConfig):
 
     @classmethod
     def from_transformers(cls, config, parent_config=None) -> GraniteMoeHybridConfig:
-        # Reuse BambaConfig.from_transformers for mamba fields + layer_types conversion
-        # (converts HF "mamba"→"mamba2" and "attention"→"full_attention")
+        # Reuse BambaConfig.from_transformers for mamba fields, MoE/RoPE/multiplier
+        # extraction, then rebuild layer_types from GraniteMoeHybrid's own naming.
         bamba = BambaConfig.from_transformers(config, parent_config)
         bamba_fields = _shallow_fields(bamba)
+
+        # GraniteMoeHybrid names layers "full_attention" / "linear_attention"
+        # (linear_attention == Mamba2/SSD), unlike Bamba's "attention" / "mamba".
+        raw_layer_types = (
+            getattr(config, "layer_types", None)
+            or getattr(config, "layers_block_type", None)
+            or []
+        )
+        _attn = {"full_attention", "attention"}
+        _mamba = {"linear_attention", "mamba", "mamba2"}
+        layer_types: list[str] = []
+        for ltype in raw_layer_types:
+            if ltype in _attn:
+                layer_types.append("full_attention")
+            elif ltype in _mamba:
+                layer_types.append("mamba2")
+            else:
+                raise ValueError(f"Unknown GraniteMoeHybrid layer type: {ltype!r}")
+        if layer_types:
+            bamba_fields["layer_types"] = layer_types
+
+        # Respect position_embedding_type: GraniteMoeHybrid checkpoints ship
+        # default ``rope_parameters`` even for the NoPE variant
+        # (granite-4.0-tiny-preview: position_embedding_type='nope'). Only apply
+        # RoPE when explicitly requested; otherwise disable it so
+        # ``initialize_rope`` returns None and attention runs NoPE.
+        if getattr(config, "position_embedding_type", "rope") != "rope":
+            bamba_fields["rope_type"] = None
+
         return cls(
             **bamba_fields,
             shared_intermediate_size=getattr(config, "shared_intermediate_size", 1024),

src/mobius/_testing/torch_reference.py

@@ -334,7 +334,18 @@ def torch_forward(
             )
         kwargs["past_key_values"] = cache
 
-    outputs = model(**kwargs)
+    try:
+        outputs = model(**kwargs)
+    except ValueError as e:
+        # All-attention hybrid models (e.g. GraniteMoeHybrid granite-4.0-1b)
+        # trip transformers' recurrent-mask builder, which assumes the hybrid
+        # cache contains a linear-attention (Mamba) layer: "`has_previous_state`
+        # can only be called on LinearAttention layers". A single-pass forward's
+        # logits don't depend on caching, so retry without a cache.
+        if "has_previous_state" not in str(e) or "past_key_values" in kwargs:
+            raise
+        kwargs["use_cache"] = False
+        outputs = model(**kwargs)
     logits = outputs.logits.cpu().numpy()
 
     # Extract KV cache if available (Mamba models don't have it)

src/mobius/models/granitemoehybrid.py

@@ -1,22 +1,26 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-"""GraniteMoeHybrid: Mamba2/SSD + Attention hybrid with MoE FFN on all layers.
+"""GraniteMoeHybrid: Mamba2/SSD + Attention hybrid with optional MoE FFN.
 
-Every layer has both a routed MoE block (``block_sparse_moe``) and a dense
-shared MLP (``shared_mlp``). The layer type ("mamba2" or "full_attention")
-controls whether the attention sub-block is a Mamba2/SSD or standard GQA.
-Attention layers use NoPE (no rotary position embeddings).
+Each layer has a dense shared MLP (``shared_mlp``) and, when the config has
+routed experts (``num_local_experts > 0``), a routed MoE block
+(``block_sparse_moe``); variants with no experts (e.g. granite-4.0-1b) run
+only the shared MLP. The layer type ("mamba2" or "full_attention") controls
+whether the token-mixing sub-block is a Mamba2/SSD or standard GQA. Attention
+layers use RoPE when ``position_embedding_type == 'rope'`` (granite-4.0-1b)
+and NoPE otherwise (granite-4.0-tiny-preview). Granite scaling multipliers
+(embedding/attention/residual/logits) are applied throughout.
 
 Forward pass per layer::
 
     residual = x
     x = input_layernorm(x)
-    x = mamba(x)  OR  self_attn(x, position_embeddings=None)
+    x = mamba(x)  OR  self_attn(x, position_embeddings)
     x = residual + x * residual_multiplier
     residual = x
     x = post_attention_layernorm(x)
-    x = block_sparse_moe(x) + shared_mlp(x)
+    x = shared_mlp(x)  [+ block_sparse_moe(x) if experts]
     x = residual + x * residual_multiplier
 
 HuggingFace reference: ``GraniteMoeHybridForCausalLM``.
@@ -39,6 +43,7 @@
     TopKGate,
     create_attention_bias,
     get_activation,
+    initialize_rope,
 )
 
 if TYPE_CHECKING:
@@ -183,6 +188,24 @@ def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value:
 # ---------------------------------------------------------------------------
 
 
+def _feedforward(
+    op: OpBuilder,
+    hidden_states: ir.Value,
+    block_sparse_moe: nn.Module | None,
+    shared_mlp: nn.Module,
+) -> ir.Value:
+    """Combined routed-MoE + shared-MLP feedforward.
+
+    When routed experts exist, the output is ``moe(x) + shared_mlp(x)``; for
+    variants with ``num_local_experts == 0`` (e.g. granite-4.0-1b) only the
+    dense shared MLP runs. Mirrors HF ``GraniteMoeHybridDecoderLayer``.
+    """
+    shared_out = shared_mlp(op, hidden_states)
+    if block_sparse_moe is None:
+        return shared_out
+    return op.Add(block_sparse_moe(op, hidden_states), shared_out)
+
+
 class _GraniteMoeHybridMambaDecoderLayer(nn.Module):
     """GraniteMoeHybrid Mamba2 layer.
 
@@ -212,8 +235,12 @@ def __init__(self, config: GraniteMoeHybridConfig):
         )
         self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
-        # Routed MoE with fused 3D expert weights
-        self.block_sparse_moe = _FusedMoEBlock(config)
+        # Routed MoE with fused 3D expert weights. GraniteMoeHybrid variants with
+        # ``num_local_experts == 0`` (e.g. granite-4.0-1b) have no routed experts;
+        # only the dense shared MLP runs. Mirrors HF ``block_sparse_moe = MoE(...)
+        # if num_local_experts > 0 else None``.
+        self._has_experts = bool(config.num_local_experts)
+        self.block_sparse_moe = _FusedMoEBlock(config) if self._has_experts else None
 
         # Dense shared MLP with fused gate+up weight
         self.shared_mlp = _FusedSharedMLP(config)
@@ -225,7 +252,7 @@ def forward(
         op: OpBuilder,
         hidden_states: ir.Value,
         attention_bias: ir.Value,
-        position_embeddings: tuple,
+        position_embeddings: tuple | None,
         past_key_value: tuple | None,
     ):
         """Forward pass. Returns (hidden_states, (conv_state, ssm_state)).
@@ -250,11 +277,7 @@ def forward(
         # MoE + shared-MLP path with pre-norm
         residual = hidden_states
         hidden_states = self.post_attention_layernorm(op, hidden_states)
-        # Both routed MoE and shared MLP run on every layer; outputs are summed
-        hidden_states = op.Add(
-            self.block_sparse_moe(op, hidden_states),
-            self.shared_mlp(op, hidden_states),
-        )
+        hidden_states = _feedforward(op, hidden_states, self.block_sparse_moe, self.shared_mlp)
         rm = op.CastLike(op.Constant(value_float=self._residual_multiplier), hidden_states)
         hidden_states = op.Add(residual, op.Mul(hidden_states, rm))
 
@@ -274,12 +297,16 @@ class _GraniteMoeHybridAttentionDecoderLayer(nn.Module):
     def __init__(self, config: GraniteMoeHybridConfig):
         super().__init__()
         self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        # GraniteMoeHybrid attention uses NoPE (no position embeddings)
-        self.self_attn = Attention(config)
+        # Attention scale is Granite's ``attention_multiplier`` (a fixed value,
+        # not the default 1/sqrt(head_dim)); RoPE is applied only when the text
+        # model supplies ``position_embeddings`` (position_embedding_type='rope').
+        self.self_attn = Attention(config, scale=config.attention_multiplier)
         self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
-        # Routed MoE with fused 3D expert weights
-        self.block_sparse_moe = _FusedMoEBlock(config)
+        # Routed MoE with fused 3D expert weights. Absent when there are no
+        # routed experts (num_local_experts == 0), e.g. granite-4.0-1b.
+        self._has_experts = bool(config.num_local_experts)
+        self.block_sparse_moe = _FusedMoEBlock(config) if self._has_experts else None
 
         # Dense shared MLP with fused gate+up weight
         self.shared_mlp = _FusedSharedMLP(config)
@@ -291,20 +318,22 @@ def forward(
         op: OpBuilder,
         hidden_states: ir.Value,
         attention_bias: ir.Value,
-        position_embeddings: tuple,
+        position_embeddings: tuple | None,
         past_key_value: tuple | None,
     ):
-        """Forward pass. Returns (hidden_states, (key, value))."""
-        del position_embeddings  # GraniteMoeHybrid uses NoPE: no rotary embeddings
+        """Forward pass. Returns (hidden_states, (key, value)).
 
-        # GQA attention path (no RoPE)
+        ``position_embeddings`` is ``None`` for NoPE variants and a
+        ``(cos, sin)`` tuple when ``position_embedding_type == 'rope'``.
+        """
+        # GQA attention path (RoPE applied iff position_embeddings is not None)
         residual = hidden_states
         hidden_states = self.input_layernorm(op, hidden_states)
         attn_out, present_kv = self.self_attn(
             op,
             hidden_states=hidden_states,
             attention_bias=attention_bias,
-            position_embeddings=None,  # NoPE: skip rotary embedding application
+            position_embeddings=position_embeddings,
             past_key_value=past_key_value,
         )
         rm = op.CastLike(op.Constant(value_float=self._residual_multiplier), attn_out)
@@ -313,10 +342,7 @@ def forward(
         # MoE + shared-MLP path with pre-norm
         residual = hidden_states
         hidden_states = self.post_attention_layernorm(op, hidden_states)
-        hidden_states = op.Add(
-            self.block_sparse_moe(op, hidden_states),
-            self.shared_mlp(op, hidden_states),
-        )
+        hidden_states = _feedforward(op, hidden_states, self.block_sparse_moe, self.shared_mlp)
         rm = op.CastLike(op.Constant(value_float=self._residual_multiplier), hidden_states)
         hidden_states = op.Add(residual, op.Mul(hidden_states, rm))
 
@@ -332,7 +358,9 @@ class _GraniteMoeHybridTextModel(nn.Module):
     """GraniteMoeHybrid text backbone: embedding -> N x (Mamba2|Attention) layers -> norm.
 
     Layer type ("mamba2" or "full_attention") is read from ``config.layer_types``.
-    No rotary embeddings are used (NoPE for attention layers).
+    Attention layers use RoPE when ``config`` declares rotary parameters
+    (``position_embedding_type == 'rope'``, e.g. granite-4.0-1b) and NoPE
+    otherwise (e.g. granite-4.0-tiny-preview).
     """
 
     def __init__(self, config: GraniteMoeHybridConfig):
@@ -352,7 +380,10 @@ def __init__(self, config: GraniteMoeHybridConfig):
                 self.layers.append(_GraniteMoeHybridAttentionDecoderLayer(config))
 
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        # No rotary_emb: GraniteMoeHybrid uses NoPE (no positional encodings)
+        # RoPE when the config declares rotary params; None => NoPE (attention
+        # layers then receive position_embeddings=None).
+        self.rotary_emb = initialize_rope(config)
+        self.embedding_multiplier = config.embedding_multiplier
 
     def forward(
         self,
@@ -362,10 +393,15 @@ def forward(
         position_ids: ir.Value,
         past_key_values: list | None = None,
     ):
-        del position_ids  # unused: NoPE architecture has no positional embeddings
-
         # (batch, seq, hidden)
         hidden_states = self.embed_tokens(op, input_ids)
+        # Granite scales embeddings by embedding_multiplier after lookup.
+        hidden_states = op.Mul(hidden_states, self.embedding_multiplier)
+
+        # (cos, sin) tuple for RoPE, or None for NoPE variants.
+        position_embeddings = (
+            self.rotary_emb(op, position_ids) if self.rotary_emb is not None else None
+        )
         attention_bias = create_attention_bias(
             op,
             input_ids=input_ids,
@@ -380,7 +416,7 @@ def forward(
                 op,
                 hidden_states=hidden_states,
                 attention_bias=attention_bias,
-                position_embeddings=None,  # NoPE: no RoPE
+                position_embeddings=position_embeddings,
                 past_key_value=past_kv,
             )
             present_key_values.append(present_kv)
@@ -413,6 +449,7 @@ def __init__(self, config: GraniteMoeHybridConfig):
         self.lm_head = Linear(config.hidden_size, config.vocab_size, bias=False)
         if config.tie_word_embeddings:
             self.lm_head.weight = self.model.embed_tokens.weight
+        self.logits_scaling = config.logits_scaling
 
     def forward(
         self,
@@ -430,6 +467,8 @@ def forward(
             past_key_values=past_key_values,
         )
         logits = self.lm_head(op, hidden_states)
+        # Granite divides final logits by logits_scaling.
+        logits = op.Div(logits, self.logits_scaling)
         return logits, present_key_values
 
     def preprocess_weights(

testdata/cases/causal-lm/granite-4_0-1b.yaml

@@ -0,0 +1,23 @@
+model_id: "ibm-granite/granite-4.0-1b"
+model_type: "granitemoehybrid"
+revision: "main"
+task_type: "text-generation"
+dtype: "float32"
+
+inputs:
+  prompts:
+    - "Here is my poem:"
+
+level: "L4+L5"
+
+generation:
+  max_new_tokens: 20
+  do_sample: false
+
+notes: >-
+  GraniteMoeHybrid dense variant: all full_attention layers (no Mamba), RoPE
+  (position_embedding_type='rope'), and no routed experts (num_local_experts=0)
+  so only the shared MLP runs. Exercises the Granite scaling multipliers
+  (embedding/attention/residual/logits). Unlike granite-4.0-tiny-preview this
+  runs natively in transformers (no trust_remote_code, no mamba-ssm), so a
+  golden can be generated in CI.

testdata/golden/causal-lm/granite-4_0-1b.json

@@ -0,0 +1,41 @@
+{
+  "top1_id": 330,
+  "top2_id": 4815,
+  "top10_ids": [
+    330,
+    4815,
+    720,
+    57277,
+    578,
+    763,
+    3146,
+    2355,
+    1035,
+    362
+  ],
+  "top10_logits": [
+    "0x1.2c68300000000p+4",
+    "0x1.2894c40000000p+4",
+    "0x1.2880740000000p+4",
+    "0x1.215f900000000p+4",
+    "0x1.1de6be0000000p+4",
+    "0x1.1bb2a80000000p+4",
+    "0x1.1aa89c0000000p+4",
+    "0x1.17b8080000000p+4",
+    "0x1.12ce680000000p+4",
+    "0x1.0c2f880000000p+4"
+  ],
+  "logits_summary": [
+    "0x1.2c68300000000p+4",
+    "-0x1.80fc560000000p+3",
+    "-0x1.01df734ac7e0ap-2",
+    "0x1.f3583dce9b92bp+1"
+  ],
+  "input_ids": [
+    8586,
+    374,
+    856,
+    33894,
+    25
+  ]
+}