Right-padded prefill produces a corrupted cache in Mamba-family models (zero conv state, decayed SSM state)

System Info

• transformers version: 5.15.1
• Platform: Linux-7.0.0-30-generic-x86_64-with-glibc2.43
• Python version: 3.12.13
• Huggingface_hub version: 1.28.0
• Safetensors version: 0.8.0
• Accelerate version: 1.14.0
• Accelerate config: not found
• DeepSpeed version: not installed
• PyTorch version (accelerator?): 2.13.0+cu130 (CUDA)
• Using distributed or parallel set-up in script?: no
• Using GPU in script?: yes, single GPU; also reproduced CPU-only
• GPU type: NVIDIA GeForce RTX 5090

Kernel packages installed for the GPU run: causal-conv1d 1.7.0, mamba-ssm 2.3.2.post1 (built from source for sm_120). The bug reproduces identically without them (see Reproduction, part 2), so it is in the cache-update logic, not the kernels.

Who can help?

@vasqu @CyrilVallez (linear-attention cache layer, shared code path)

Information:

• My own modified scripts

Tasks:

• My own task or dataset (give details below)

Preparation work for a vision-language build on NVIDIA-Nemotron-Nano-9B-v2; systematic padding-correctness checks before training surfaced this.

Reproduction

No traceback: the failure is silent. Nothing raises; the model simply produces wrong continuations after a right-padded prefill.

A use_cache=True prefill whose rows are right-padded produces a cache that is wrong for the padded rows. Every subsequent decoded token is wrong (argmax flip, logit error > 12). Left-padded prefill produces a correct cache. generate() is unaffected because it left-pads.

What the cache holds after prefill, layer 0, padded row, vs the same row prefilled alone:

A single pad token is enough to corrupt (hidden-state error 0.85 at p=1); it saturates as the state finishes decaying.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

M = "nvidia/NVIDIA-Nemotron-Nano-9B-v2"
tok = AutoTokenizer.from_pretrained(M)
tok.pad_token = tok.eos_token
PAD = tok.pad_token_id
model = AutoModelForCausalLM.from_pretrained(M, dtype=torch.bfloat16).cuda().eval()

short = tok("The capital of France is Paris.", return_tensors=None)["input_ids"]
long_ = tok(
    "Selective state space models keep a fixed-size recurrent state per token, "
    "unlike attention whose cost grows with the sequence length, so hybrids "
    "interleave the two.", return_tensors=None)["input_ids"]
Ls, Ll = len(short), len(long_)

# clean reference: prefill short alone, decode one step
with torch.inference_mode():
    out = model(input_ids=torch.tensor([short], device="cuda"), use_cache=True)
    nxt = out.logits[0, -1].argmax().item()
    ref_state = out.past_key_values.layers[0].recurrent_states[0].clone()
    ref_conv = out.past_key_values.layers[0].conv_states[0].clone().float()
    ref = model(input_ids=torch.tensor([[nxt]], device="cuda"),
                past_key_values=out.past_key_values, use_cache=True).logits[0, -1].float()

for side in ("left", "right"):
    if side == "left":
        ids = torch.tensor([[PAD] * (Ll - Ls) + short, long_], device="cuda")
        mask = torch.ones_like(ids); mask[0, : Ll - Ls] = 0
    else:
        ids = torch.tensor([short + [PAD] * (Ll - Ls), long_], device="cuda")
        mask = torch.ones_like(ids); mask[0, Ls:] = 0
    with torch.inference_mode():
        out = model(input_ids=ids, attention_mask=mask, use_cache=True)
        # read the cache before the decode step mutates it in place
        st = out.past_key_values.layers[0].recurrent_states[0][0].clone()
        cv = out.past_key_values.layers[0].conv_states[0][0].clone().float()
        row1_next = out.logits[1, -1].argmax().item()
        acc_mask = torch.cat([mask, torch.ones_like(mask[:, :1])], dim=1)
        step = model(input_ids=torch.tensor([[nxt], [row1_next]], device="cuda"),
                     past_key_values=out.past_key_values, use_cache=True,
                     attention_mask=acc_mask)
        lg = step.logits[0, -1].float()
    print(f"{side:>5}-pad: logits max|diff|={float((lg - ref).abs().max()):7.3f} "
          f"top1 ref={ref.argmax().item()} got={lg.argmax().item()} "
          f"conv abs_mean={float(cv.abs().mean()):.6f} (clean {float(ref_conv.abs().mean()):.6f}) "
          f"ssm rel_err={float((st.float() - ref_state.float()).norm() / ref_state.float().norm()):.4f}")

Actual output (CPU, pure-torch fallback; the CUDA run with kernels gives 0.125 / 12.622 with the same argmax flip):

 left-pad: logits max|diff|=  0.094 top1 ref=8961 got=8961 conv abs_mean=0.536767 (clean 0.536767) ssm rel_err=0.0000
right-pad: logits max|diff|= 12.561 top1 ref=8961 got=1044 conv abs_mean=0.000000 (clean 0.536767) ssm rel_err=0.7776

Part 2, same failure without any kernel package: run the same script on CPU with causal-conv1d and mamba-ssm made unimportable (e.g. a PYTHONPATH stub that raises ImportError). The pure-torch fallback path reproduces the exact signature above, including the same wrong token (1044), which is why I think this belongs in the cache layer rather than with the kernel integrations. (Side note: with the kernel packages installed, a CPU run does not fall back — causal_conv1d_fn raises Expected x.is_cuda() to be true. The stub works around that.)

Expected behavior

Both pad directions should match the clean reference: the cached conv state should hold the row's last real pre-conv inputs, and the cached recurrent state should be the state at the row's last real token. Pads masked out by attention_mask should not change any cached state.

Mechanism

Two independent defects in the prefill path of the Mamba mixer (modeling_nemotron_h.py; line numbers from 5.15.1):

1. Conv state. update_conv_state (modeling_nemotron_h.py:489 → cache_utils.py:1066) caches full_conv_states[..., -conv_kernel_size:] — the last positions of the padded sequence, with no mask awareness. For a right-padded row those positions are pads, and the layer input at pads is zeroed by apply_mask_to_padding_states, so the cache holds zeros instead of the row's last real tokens.

2. SSM state. apply_mask_to_padding_states (modeling_nemotron_h.py:517) zeroes x/B/C at pad positions, but dt is never masked. The chunk scan still runs over the pad tail, where each step applies h ← exp(dt·A)·h with zero input, so the cached final_state is decayed across the pads (update_recurrent_state, modeling_nemotron_h.py:547). With zero input the state does not hold at the last real token — it keeps decaying for every pad step.

With left padding both problems disappear: the pads sit upstream (the state at the first real token is the fresh state) and the last positions of the sequence are real tokens.

Scope

update_conv_state and the scan-over-pads pattern live in the shared linear-attention cache layer and the shared padding helper, used by every model that mixes apply_mask_to_padding_states with a conv/recurrent cache: nemotron_h (confirmed here), jamba, zamba2, falcon_h1, qwen3_next, bamba, lfm2, mamba2, granitemoehybrid, olmo_hybrid, and others. I have only run the reproducer on nemotron_h; the code path is common.

This only bites code that right-pads a batched prefill and then continues from the cache (custom harnesses, assistant-chunked prefill, anything with padding_side="right" + use_cache). Prefill-time outputs and losses are unaffected — pads are masked correctly for the forward computation.

Possible fixes

• Conv state: select each row's last conv_kernel_size real positions via the attention mask (gather) instead of slicing the last positions of the padded sequence. This needs the mask to reach update_conv_state, which the cache layer does not see today.
• SSM state: the scan kernels have varlen support (seq_idx for mamba_chunk_scan_combined); packing rows to their real lengths and taking per-row final states would make padding direction irrelevant. mamba2_chunk_scan does not expose it.
• Short term, a warning (or error) when a mask with trailing zeros is combined with use_cache on these layers would turn silent wrong continuations into a visible one.

Workaround

Left-pad batched prefills (the generate() default) and thread the accumulated attention_mask through decode steps.