Fix recurrent-layer padding mask being skipped on continued forwards (chunked prefill, cache continuation)
Author: abcgco
Pull request: yes

<!-- ci-dashboard-badge:start -->
[![CI](https://transformers-ci.lor-e.huggingface.cool/badge/pr?pr=47087)](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=47087)
<!-- ci-dashboard-badge:end -->

# What does this PR do?

Fixes #47086

`create_recurrent_attention_mask` returned `None` as soon as the cache had a previous recurrent state, so multi-token forwards continuing from a cache (chunked prefill via `prefill_chunk_size`, cache continuation) never zeroed the padding of the new segment out of the linear-attention/conv path of hybrid models (`bamba`, `falcon_h1`, `granitemoehybrid`, `jamba`, `lfm2`, `lfm2_moe`, `nemotron_h`, `olmo_hybrid`, `qwen3_5`, `qwen3_5_moe`, `qwen3_next`, `zaya`). The padding of the new segment leaked into the recurrent state and silently corrupted all subsequent real tokens (full analysis, repros and real-checkpoint numbers in #47086).

The fix replaces the state-based skip with a single-token (decode) skip:

- a generated token is never padding, so decode still does not need the mask;
- the check is shape-based, so compiled decode keeps receiving `None` exactly as today (the state check existed to keep the growing 2D mask out of the compiled decode graph; that property is preserved);
- multi-token continued forwards now receive the mask sliced to the current segment, so `apply_mask_to_padding_states` zeroes the new segment's padding again.

With this fix on `Qwen/Qwen3.5-0.8B` (float32, CPU): chunked prefill divergence vs the unpadded baseline goes from `5.661` to `3.3e-05` (the fp32 floor), cache-continuation vs the single forward of the same batch from `4.445` to `2.2e-05`, and greedy `generate` outputs with and without `prefill_chunk_size` become identical; unpadded control rows are unaffected.

Not a duplicate of #46773: that PR fixes the batch-size-1 gate inside `apply_mask_to_padding_states` for the first (prefill) forward and does not change `create_recurrent_attention_mask`; its alignment gate is intentionally false on continued forwards. The two changes compose (the sliced mask restored here has exactly the aligned length).

Tests: a unit test in `tests/utils/test_masking_utils.py` (mask kept on continued multi-token forwards, still `None` on single-token decode) and a general regression test in `tests/generation/test_utils.py` (`GenerationTesterMixin.test_recurrent_layers_mask_padding_on_continued_forward`, added after review) that runs for every model whose modeling module uses `create_recurrent_attention_mask`: ten hybrid models fail on `main` and pass with the fix, NemotronH and Jamba skip it with a documented reason (a residual downstream of the mask, in their Mamba mixer), and purely recurrent models (Mamba/RWKV/...) skip as out of scope.

Test commands run locally (CPU, torch 2.12.1):

```
python -m pytest tests/utils/test_masking_utils.py                       # 11 passed
python -m pytest tests/models/qwen3_next/test_modeling_qwen3_next.py    # fast suite green; tp/ep and training_kwargs failures are pre-existing (fail identically on main, no distributed env locally)
python -m pytest tests/models/lfm2/test_modeling_lfm2.py                # 143 passed
# the general test, swept over 20+ models: the 10 hybrids using create_recurrent_attention_mask
# pass (and fail without the fix); NemotronH/Jamba documented skip; pure-recurrent and
# non-recurrent models skip cleanly
python -m pytest tests/models/<model> -k test_recurrent_layers_mask_padding_on_continued_forward
ruff check / ruff format                                                 # clean (one pre-existing FURB171 on an untouched line)
```

AI assistance was used for the analysis and drafting; I reviewed every changed line, ran the tests, and reproduced both the bug and the fix on my machine (see #47086).

## Code Agent Policy

- [x] I confirm that this is not a pure code agent PR.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://huggingface.co/docs/transformers/contributing) and the
      [Pull Request](https://huggingface.co/docs/transformers/pr_checks) checks?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case. -> #47086
- [x] Did you make sure to update the documentation with your changes according to the [guidelines](https://github.com/huggingface/transformers/tree/main/docs)? (the docstring of `create_recurrent_attention_mask` is updated to match the new behavior)
- [x] Did you write any new necessary [tests](https://huggingface.co/docs/transformers/testing)?

## Who can review?

@Cyrilvallez (masking utils)


## Round 3 (review follow-up)

Per review: the NemotronH mixer bug is now fixed rather than skipped (7faa96d — the shared `Zamba2MambaMixer.torch_forward` masks padding before `in_proj` whenever the mask is the 2D current-tokens slice; decode and Zamba2's full-history plumbing verified bit-for-bit unchanged), the test seed is dropped, and the Jamba/Zamba skip reasons are corrected to the Mamba1 conv-continuation contract (4623260).

Commands run this round (Python 3.12, torch CPU):

```bash
python -m pytest tests/models/nemotron_h/test_modeling_nemotron_h.py::NemotronHModelTest::test_recurrent_layers_mask_padding_on_continued_forward
# PASSED (skip removed this round; RED on the pre-fix mixer, 2 failed)

python -m pytest tests/models/{bamba,falcon_h1,granitemoehybrid,jamba,lfm2,lfm2_moe,minimax,nemotron_h,olmo_hybrid,qwen3_5,qwen3_5_moe,qwen3_next,zamba,zamba2,zaya} \
  -k test_recurrent_layers_mask_padding_on_continued_forward
# 13 passed, 4 skipped (jamba, minimax, zamba, zamba2 — documented skips); identical across 3 consecutive seedless runs

python -m pytest tests/models/nemotron_h/ tests/models/zamba2/ tests/models/zamba/ tests/models/jamba/ tests/utils/test_masking_utils.py
# 560 passed, 0 failed (5614 subtests)

python utils/check_modular_conversion.py --files src/transformers/models/zamba2/modular_zamba2.py src/transformers/models/nemotron_h/modular_nemotron_h.py
# exit 0

ruff==0.14.10 check / format --check on the touched files
# All checks passed / already formatted
```


Comments:
Round 3 pushed: 7faa96d fixes the mixer per the nemotron_h thread (its skip is removed and the common test passes; details inline), 4623260 drops the seed and corrects the Jamba/Zamba skip reasons, plus a merge of current main. Skip inventory after this round: Jamba/Zamba (Mamba1 conv-continuation contract), Zamba2 (mask plumbing, with a pre-existing shape crash on masked multi-token continuations that the follow-up slice fix would cure), MiniMax (lightning-attention inline path) — NemotronH is fixed and unskipped.

On the two red shards of the previous run — both unrelated, from the job logs: in `tests_generate [4/8]`, a 120s `torch.compile` timeout on ministral3 leaked its fx `nn.Module.__call__` patch on that xdist worker, and the following 33 tests on gw4 failed instantly with `NameError: module is not installed as a submodule` (the same fullgraph test passed for 15+ models on the other workers); in `tests_torch [7/8]`, `xcodec::test_can_init_all_missing_weights` (`ConvTranspose1d` missing from `_init_weights`) fails on clean main as well. Our regression test was green in both shards.


Thank you for your contribution 🤗!

## CI Security Gate — automatic approval blocked

This PR was **not** automatically approved for CI because the security gate failed.

**Possible reasons:**
- The PR touches 50 or more files — only PRs with fewer than 50 changed files are automatically approved
- A changed file is outside the allowed directories (`src/`, `tests/`, `docs/`, `utils/`), has a disallowed extension (only `.py`, `.txt`, `.md` permitted outside `tests/` and `docs/`), or is not `.md`/`.yml` inside `docs/`
- A new high-severity security issue was detected in the changed Python files (Bandit check)

See the [workflow run](https://github.com/huggingface/transformers/actions/runs/29518088431) for the exact violations.

A maintainer can review and manually approve CI if a finding is a false positive.

run-slow: falcon_mamba, jamba, mamba, mamba2, minimax, nemotron_h, zamba, zamba2

[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/29580090178)

This comment contains `run-slow`, running the specified jobs: 

models: ["models/falcon_mamba", "models/jamba", "models/mamba", "models/mamba2", "models/minimax", "models/nemotron_h", "models/zamba", "models/zamba2"]
quantizations: []

## CI Results
[Workflow Run ⚙️](https://github.com/huggingface/transformers/actions/runs/29580090178)

### Commit Info
| Context | Commit | Description |
|---------|--------|-------------|
| RUN | [c92c6886](https://github.com/huggingface/transformers/commit/c92c68866ff5ba05ad70f5e74c9d1ce135bf173e) | workflow commit (merge commit) |
| PR | [12a808c8](https://github.com/abcgco/transformers/commit/12a808c87bd3e03938a72048df03e8b50d43b637) | branch commit (from PR) |
| main | [13337f29](https://github.com/huggingface/transformers/commit/13337f2978a2016ea5e71ea6ae42201e760e927d) | base commit (on `main`) |

✅ No failing test specific to this PR 🎉 👏 !


**[For maintainers]** Suggested jobs to run (before merge)

run-slow: falcon_mamba, jamba, mamba, mamba2, minimax, nemotron_h, zamba, zamba2

<!-- ci-dashboard-recap:start -->

---

### CI recap

**Dashboard:** [View test results in Grafana](https://transformers-ci.lor-e.huggingface.cool/d/pytest-observability-by-pr/pytest-observability-branch?var-pr=47087)
**Latest run:** [28869791091:2](https://github.com/huggingface/transformers/actions/runs/29600333464)
**Result:** `success` | **Jobs:** 15 | **Tests:** 172,892 | **Failures:** 34 | **Duration:** 8h 5m

<!-- ci-dashboard-recap:end -->