[BUG] GraniteMoeHybrid drops cache lengths when calling ssm_update.

[kvcache670]

The mask change proposed in PR #1736 is not sufficient to make padded batch generation agree with standalone generation for granitemoehybrid. Its _ssm method also reads cache.lengths without passing it to ssm_update.

In a tiny randomly initialized GraniteMoeHybrid model, fixing only the mask leaves a maximum absolute log-probability difference of 0.2981 for the shorter prompt.

To Reproduce

No model weights or tokenizer downloads are needed. The float32 model has one Mamba-2 layer and one attention layer. The script compares the first-step log probabilities of the shorter prompt in a batch with those from running it alone.

Two in-process changes isolate the causes:

• A: combine the left_padding and lengths masks in ArraysCache.make_mask, as proposed in PR Fix batched generation for hybrid-cache models (silent corruption + broadcast crash) #1736.
• B: pass cache.lengths to ssm_update in GraniteMoeHybridMamba2Mixer._ssm.

The script sets A_log to -4 and multiplies convolution weights by 30 to make recurrent-state differences visible with random initialization. It uses mamba_d_state=32. Save it as repro.py.

python3 -m venv .venv
source .venv/bin/activate
python -m pip install "mlx==0.32.2" "mlx-lm==0.31.3"
python repro.py

# No weights needed: a tiny randomly initialised granitemoehybrid (1 Mamba-2 layer + 1 attention layer),
# float32, greedy. The same prompt alone and as the shorter member of a batch, under four variants:
#   A = ArraysCache.make_mask combining left_padding and lengths (the change proposed in PR #1736)
#   B = GraniteMoeHybridMamba2Mixer._ssm passing cache.lengths on to ssm_update (falcon_h1 / mamba2 / plamo2 do)
import mlx.core as mx
import mlx_lm
from mlx_lm.generate import BatchGenerator, generate_step
from mlx_lm.models import cache as cache_mod, granitemoehybrid as G
from mlx_lm.models.ssm import ssm_update

mx.random.seed(0)
model = G.Model(G.ModelArgs(
    model_type="granitemoehybrid", vocab_size=100, hidden_size=64, intermediate_size=128,
    shared_intermediate_size=128, num_hidden_layers=2, layer_types=["mamba", "attention"],
    max_position_embeddings=512, num_attention_heads=4, num_key_value_heads=2, attention_bias=False,
    embedding_multiplier=1.0, attention_multiplier=0.25, logits_scaling=1.0, residual_multiplier=1.0,
    rms_norm_eps=1e-5, rope_theta=10000.0, position_embedding_type="nope",
    mamba_n_heads=4, mamba_d_head=32, mamba_d_state=32, mamba_d_conv=4, mamba_n_groups=1,
    mamba_proj_bias=False, mamba_conv_bias=True))
for layer in model.layers:          # random init gives the recurrent state almost no weight; make it matter,
    if hasattr(layer, "mamba"):     # as it does in a trained checkpoint (slow decay, larger conv output)
        layer.mamba.A_log = mx.full(layer.mamba.A_log.shape, -4.0)
        layer.mamba.conv1d.weight = layer.mamba.conv1d.weight * 30
mx.eval(model.parameters())
long_prompt, short_prompt = list(range(1, 13)), [7, 8, 9]


def padded_member_error():
    _, alone = next(iter(generate_step(mx.array(short_prompt), model, max_tokens=1)))
    gen = BatchGenerator(model, max_tokens=1, stop_tokens=None)
    uids = gen.insert([long_prompt, short_prompt])
    got = {}
    while len(got) < 2:
        for r in gen.next_generated():
            got.setdefault(r.uid, r.logprobs)
    gen.close()
    return mx.abs(got[uids[1]] - alone).max().item()


def make_mask_A(self, N):
    mask = None
    if self.left_padding is not None:
        mask = mx.arange(N) >= self.left_padding[:, None]
    if self.lengths is not None:
        right = mx.arange(N) < self.lengths[:, None]
        mask = right if mask is None else (mask & right)
    return mask


def _ssm_B(self, hidden_states, B, C, dt, cache, mask):
    b, t, _ = hidden_states.shape
    hidden_states = hidden_states.reshape(b, t, self.num_heads, self.head_dim)
    B = B.reshape(b, t, self.n_groups, self.ssm_state_size)
    C = C.reshape(b, t, self.n_groups, self.ssm_state_size)
    state, lengths = (cache[1], cache.lengths) if cache else (None, None)
    y, state = ssm_update(hidden_states, self.A_log, B, C, self.D.astype(hidden_states.dtype), dt,
                          self.dt_bias, state, self.time_step_limit, mask, lengths)   # <- lengths
    if cache:
        cache[1] = state
    return y.reshape(b, t, self.intermediate_size)


mixer = next(v for v in vars(G).values() if isinstance(v, type) and hasattr(v, "_ssm"))
released = (cache_mod.ArraysCache.make_mask, mixer._ssm)
print("mlx", mx.__version__, "| mlx-lm", mlx_lm.__version__, "| max |logprobs in batch - alone| for the padded prompt")
err = {}
for name, a, b in (("as released", 0, 0), ("A only", 1, 0), ("B only", 0, 1), ("A and B", 1, 1)):
    cache_mod.ArraysCache.make_mask = make_mask_A if a else released[0]
    mixer._ssm = _ssm_B if b else released[1]
    err[name] = padded_member_error()
    print(f"  {name:12s} {err[name]:.3e}")
assert err["A only"] < 1e-3, (f"with make_mask fixed (A) the padded prompt is still off by {err['A only']:.3g}; "
                              f"it takes B as well ({err['A and B']:.1e})")

Observed behavior

Recorded output (exit code 1; only the local traceback path has been shortened):

ProductName:  macOS ProductVersion:  26.6.1 BuildVersion:  25G76 
Apple M1 Ultra
mlx 0.32.2 | mlx-lm 0.31.3 | max |logprobs in batch - alone| for the padded prompt
  as released  1.004e+00
  A only       2.981e-01
  B only       1.143e+00
  A and B      4.768e-07
Traceback (most recent call last):
  File "repro.py", line 72, in <module>
    assert err["A only"] < 1e-3, (f"with make_mask fixed (A) the padded prompt is still off by {err['A only']:.3g}; "
           ^^^^^^^^^^^^^^^^^^^^
AssertionError: with make_mask fixed (A) the padded prompt is still off by 0.298; it takes B as well (4.8e-07)
exit code: 1

Expected behavior

After both the padding mask and sequence lengths are respected, the shorter prompt's first-step log probabilities should agree between standalone and batch generation within floating-point tolerance.

Mask behavior	Original _ssm call	Pass lengths (B)
Original make_mask	1.004e+00	1.143e+00
Combined mask (A)	2.981e-01	4.768e-07

Values are maximum absolute log-probability differences from the recorded run. The final assertion checks the A-only case deliberately: it demonstrates the remaining problem after the mask fix.

Desktop (please complete the following information):

• OS Version: macOS 26.6.1 (25G76)
• Hardware: Apple M1 Ultra
• MLX: 0.32.2
• MLX-LM: 0.31.3
• Device: Metal GPU