Hybrid models: padding is not masked out of linear-attention/conv state on continued forwards (chunked prefill, cache continuation)
Author: abcgco
Pull request: no
Labels: name: bug

### System Info

- `transformers` main (5.14.0.dev0, commit b70d02fc)
- torch 2.12.1, CPU, float32 (fla / causal_conv1d not installed, so all numbers below use the torch fallback path; the skipped masking happens before kernel dispatch, so the fast paths receive the same unmasked hidden states)
- Python 3.12

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

For models with recurrent layers (`bamba`, `falcon_h1`, `granitemoehybrid`, `jamba`, `lfm2`, `lfm2_moe`, `nemotron_h`, `olmo_hybrid`, `qwen3_5`, `qwen3_5_moe`, `qwen3_next`, `zaya`), padding tokens are zeroed out of the linear-attention/conv path via `apply_mask_to_padding_states`, using the 2D mask returned by `masking_utils.create_recurrent_attention_mask`. That function returns `None` as soon as the cache has a previous recurrent state:

https://github.com/huggingface/transformers/blob/b70d02fc724d04c916832ca4ead03ff05e8fb1ee/src/transformers/masking_utils.py#L1481-L1482

```python
if past_key_values is not None and past_key_values.has_previous_state():
    return None
```

So on any multi-token forward that *continues* from an existing cache, pad positions inside the new segment are never zeroed, and their hidden states flow into the depthwise conv and the gated-delta/SSM recurrence, corrupting the state for all subsequent real tokens. The full-attention layers are unaffected (their 4D mask handles padding on every forward), which makes the divergence silent.

Two public-API ways to hit this:

1. **Chunked prefill** (`generation_config.prefill_chunk_size`): with a left-padded batch, rows whose padding is longer than the chunk size have pad tokens inside chunk 2+, which run with `has_previous_state=True`. The manual loop below mimics the chunked-prefill loop of `GenerationMixin._prefill` (https://github.com/huggingface/transformers/blob/b70d02fc724d04c916832ca4ead03ff05e8fb1ee/src/transformers/generation/utils.py#L3919-L3931).
2. **Cache continuation / batched multi-turn**: turn 2 is left-padded to a common length and forwarded with the turn-1 cache; the pads of the new segment are inside the continued forward.

On `Qwen/Qwen3.5-0.8B` (18 `linear_attention` / 6 `full_attention` layers), float32, CPU, padding with its regular `pad_token_id` (`
`):

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL = "Qwen/Qwen3.5-0.8B"
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.float32).eval()

PAD = tok.pad_token_id
text_long = ("The library of Alexandria was one of the largest and most significant libraries "
             "of the ancient world, and its destruction remains a symbol of the loss of knowledge")
text_tgt = "The capital of France is one of the most"

ids_long = tok(text_long, return_tensors="pt").input_ids[0]   # 29 tokens
ids_tgt = tok(text_tgt, return_tensors="pt").input_ids[0]     # 9 tokens
pad_len = ids_long.shape[0] - ids_tgt.shape[0]                # 20
CHUNK = pad_len - 4                                           # 16 -> row 1 pads spill into chunk 2

input_ids = torch.stack([ids_long, torch.cat([torch.full((pad_len,), PAD), ids_tgt])])
attention_mask = torch.stack([
    torch.ones_like(ids_long),
    torch.cat([torch.zeros(pad_len, dtype=torch.long), torch.ones_like(ids_tgt)]),
])
position_ids = (attention_mask.cumsum(-1) - 1).clamp(min=0)

@torch.no_grad()
def chunked_prefill(chunk):
    cache, past = None, 0
    for ids in torch.split(input_ids, chunk, dim=-1):
        cur = past + ids.shape[-1]
        out = model(input_ids=ids, attention_mask=attention_mask[:, :cur],
                    position_ids=position_ids[:, past:cur], past_key_values=cache, use_cache=True)
        cache, past = out.past_key_values, cur
    return out.logits

with torch.no_grad():
    gt = model(input_ids=ids_tgt[None], attention_mask=torch.ones_like(ids_tgt)[None],
               position_ids=(torch.ones_like(ids_tgt)[None].cumsum(-1) - 1)).logits[0, -1]
    full = model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids).logits
chunked = chunked_prefill(CHUNK)

print(f"padded full prefill    vs unpadded: {(full[1, -1] - gt).abs().max():.3e}")
print(f"padded chunked prefill vs unpadded: {(chunked[1, -1] - gt).abs().max():.3e}")
print(f"control (unpadded row 0) full vs chunked: {(full[0, -1] - chunked[0, -1]).abs().max():.3e}")

# end-to-end: greedy generate, full vs chunked prefill
gen_kwargs = dict(attention_mask=attention_mask, max_new_tokens=12, do_sample=False)
with torch.no_grad():
    g_full = model.generate(input_ids, **gen_kwargs)
    g_chunk = model.generate(input_ids, prefill_chunk_size=CHUNK, **gen_kwargs)
print("full   :", repr(tok.decode(g_full[1, input_ids.shape[1]:], skip_special_tokens=True)))
print("chunked:", repr(tok.decode(g_chunk[1, input_ids.shape[1]:], skip_special_tokens=True)))
```

Output on main:

```
padded full prefill    vs unpadded: 3.147e-05     # fp32 floor: single-forward prefill masks pads correctly
padded chunked prefill vs unpadded: 5.661e+00     # pads of chunk 2 leak into the recurrent state
control (unpadded row 0) full vs chunked: 3.099e-05   # chunking itself is exact
full   : ' important cities in the world.\nThe capital of France is'
chunked: ' important cities in the world. It is also one of the'
```

The last-token logits of the padded row diverge by ~5.7 (vs the ~3e-5 float32 floor), and greedy `generate` output for the padded row changes.

The cache-continuation path fails the same way. Here the correct baseline is the single full forward of the *same padded batch*, not the unpadded sequence: zeroing pad hidden states is itself an approximation for mid-sequence pads (a zeroed step contributes nothing additively but still applies its decay), and that inherent cost is shared by the single forward and the split one. Splitting the same batch into a turn-1 forward plus a left-padded turn-2 forward diverges from its own single-forward baseline by `4.445e+00` (snippet below) — with masking working, splitting should be a no-op.

<details>
<summary>Cache-continuation snippet (continues from the script above)</summary>

```python
t1 = tok("The weather in Paris today is", return_tensors="pt").input_ids[0]
n0 = tok(" absolutely wonderful and the streets are full of people", return_tensors="pt").input_ids[0]
n1 = tok(" quite rainy", return_tensors="pt").input_ids[0]
pad2 = n0.shape[0] - n1.shape[0]

turn1 = torch.stack([t1, t1])
mask1 = torch.ones(2, t1.shape[0], dtype=torch.long)
turn2 = torch.stack([n0, torch.cat([torch.full((pad2,), PAD), n1])])
mask2 = torch.stack([torch.ones_like(n0),
                     torch.cat([torch.zeros(pad2, dtype=torch.long), torch.ones_like(n1)])])
full_ids = torch.cat([turn1, turn2], dim=-1)
full_mask = torch.cat([mask1, mask2], dim=-1)
pos = (full_mask.cumsum(-1) - 1).clamp(min=0)

with torch.no_grad():
    single = model(input_ids=full_ids, attention_mask=full_mask, position_ids=pos).logits
    out1 = model(input_ids=turn1, attention_mask=mask1, position_ids=pos[:, :t1.shape[0]], use_cache=True)
    out2 = model(input_ids=turn2, attention_mask=full_mask, position_ids=pos[:, t1.shape[0]:],
                 past_key_values=out1.past_key_values, use_cache=True)
print((out2.logits[1, -1] - single[1, -1]).abs().max())  # ~4.4 on main, ~2e-5 with the fix below
```

</details>

<details>
<summary>Self-contained deterministic repro (tiny random Qwen3Next, no download)</summary>

Padding uses an ordinary token id — as with `pad = eos` setups — because randomly initialized models zero the `padding_idx` embedding row, which would mask the bug; padding semantics come from `attention_mask`.

```python
import torch
from transformers import Qwen3NextConfig, Qwen3NextForCausalLM

torch.manual_seed(0)
cfg = Qwen3NextConfig(
    vocab_size=128, hidden_size=32, intermediate_size=64, num_hidden_layers=2,
    layer_types=["linear_attention", "full_attention"],
    num_attention_heads=2, num_key_value_heads=1, head_dim=16,
    linear_conv_kernel_dim=4, linear_key_head_dim=8, linear_value_head_dim=8,
    linear_num_key_heads=2, linear_num_value_heads=4,
    moe_intermediate_size=32, shared_expert_intermediate_size=32,
    num_experts_per_tok=2, num_experts=4, pad_token_id=0, attn_implementation="eager",
)
model = Qwen3NextForCausalLM(cfg).eval()

PAD = 7
torch.manual_seed(1)
real0 = torch.randint(1, cfg.vocab_size, (18,))
real1 = torch.randint(1, cfg.vocab_size, (8,))
input_ids = torch.stack([real0, torch.cat([torch.full((10,), PAD), real1])])
attention_mask = torch.stack([torch.ones(18, dtype=torch.long),
                              torch.cat([torch.zeros(10, dtype=torch.long), torch.ones(8, dtype=torch.long)])])
position_ids = (attention_mask.cumsum(-1) - 1).clamp(min=0)

@torch.no_grad()
def chunked_prefill(chunk):
    cache, past = None, 0
    for ids in torch.split(input_ids, chunk, dim=-1):
        cur = past + ids.shape[-1]
        out = model(input_ids=ids, attention_mask=attention_mask[:, :cur],
                    position_ids=position_ids[:, past:cur], past_key_values=cache, use_cache=True)
        cache, past = out.past_key_values, cur
    return out.logits

with torch.no_grad():
    logits_full = model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids).logits
    ids_gt = real1[None]
    logits_gt = model(input_ids=ids_gt, attention_mask=torch.ones_like(ids_gt),
                      position_ids=(torch.ones_like(ids_gt).cumsum(-1) - 1)).logits
logits_chunked = chunked_prefill(6)  # row1 pads (10) spill into chunk 2: [6 pads][4 pads + 2 real][6 real]

print((logits_chunked[0, -1] - logits_full[0, -1]).abs().max())  # unpadded row, control: ~4e-08
print((logits_full[1, -1]    - logits_gt[0, -1]).abs().max())    # padded row, single forward: ~1e-08
print((logits_chunked[1, -1] - logits_gt[0, -1]).abs().max())    # padded row, chunked: ~1.2e-03
```

</details>

Instrumenting `create_recurrent_attention_mask` confirms the mechanism: it returns the sliced mask on the first forward and `None` on every continued forward, so `apply_mask_to_padding_states` no-ops exactly on the chunks that still contain padding.


### Expected behavior

Continued forwards should zero pad positions of the current segment for the recurrent path, like the first forward does — the logits of real tokens should not depend on how the same padded batch was split into forwards.

### Possible fix

Instead of early-returning `None` when `has_previous_state()` is true, return the mask sliced to the current segment (the function already does this at the end: `attention_mask[:, -inputs_embeds.shape[1]:]`). If I read #46738 right, the state check exists so decode does not drag the growing 2D mask into the compiled graph — a `q_len == 1` skip preserves that (a generated token is never padding, and the check is shape-based, so compiled decode stays graph-stable) while not turning masking off for multi-token continued forwards. Note the current all-ones skip checks the full 2D mask, so it does not cover decode for batches whose prompt had padding — hence a dedicated `q_len == 1` skip rather than relying on it.

I have applied this fix locally and re-run everything on `Qwen/Qwen3.5-0.8B`: chunked prefill goes from `5.661` to `3.3e-05` (the fp32 floor of the single-forward baseline), the continuation-vs-single-forward divergence from `4.445` to `2.2e-05`, greedy `generate` outputs with and without `prefill_chunk_size` become identical, and the unpadded control rows are bit-for-bit unaffected.

Related work, for context (none of these cover the case above):
- #45513 and #46084 fixed the *state plumbing* of exactly these continued multi-token forwards (GDN and Mamba2), so the code path itself is supported — but padding on that path was left unhandled.
- #46773 (open) fixes the batch-size-1 gate inside `apply_mask_to_padding_states` for the *first* (prefill) forward. It does not change `create_recurrent_attention_mask`, and its alignment gate (`attention_mask.shape[1] == hidden_states.shape[1]`) is intentionally false on continued forwards, so the case above remains after it (the sliced mask this fix restores has exactly the aligned length, so the two changes compose).

Happy to open a PR with the fix and regression tests for both scenarios if the direction sounds right.

This analysis was AI-assisted; I reviewed and reproduced everything on my machine.



Comments:
cc @IlyasMoutawwakil since I think this might be a regression introduced by #46738, see also the agent PR at #47087

yes this looks like a regression, not sure how it was handled before since the masking was a bit scattered between modeling and cache classes. 
it's weird it wasn't caught by the continuing from an existing cache tests.
@vasqu @Cyrilvallez