[Sunt-ing]

What does this PR do?

Follow-up to #41790. apply_mask_to_padding_states (and the inline mixer copies in Bamba and FalconH1) zeros padding positions before the depthwise conv and the SSM / gated-delta recurrence. It was gated on attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1, so a batch-size-1 padded sequence was never masked: its padding flowed through the conv and recurrence into the real tokens and corrupted their logits. Batch size > 1 (the mask ran) and right padding (it follows the real tokens, which a causal model ignores) were unaffected.

The attention_mask.shape[0] > 1 gate cannot just be dropped. During cached multi-token decode (for example assisted / speculative decoding) attention_mask spans the whole sequence while hidden_states holds only the new tokens, so the multiply would shape-mismatch and raise; the batch-size gate happened to hide that for batch size 1. This replaces it with an alignment check, attention_mask.shape[1] == hidden_states.shape[1], which is true exactly in the prefill / full forward where the padding needs zeroing and false during cached decode. #41790 already noted the intent was to keep only the prefill guard ("we only make sure it is applied in prefill stage"); this makes that guard robust.

On Qwen/Qwen3.5-0.8B in float32, the last-token logits of a left-padded batch-size-1 sequence diverge from the unpadded sequence by ~5.7 on main and ~1e-5 with this change. Added Qwen3_5TextModelTest::test_left_padding_matches_unpadded_batch_size_one (fails before, passes after); the existing Lfm2ModelTest::test_assisted_decoding_* tests crash on a naive removal of the gate and pass with the alignment check.

All ten linear-attention models sharing this code are updated: Bamba, FalconH1, GraniteMoeHybrid, LFM2, LFM2-MoE, Mamba2, Olmo-hybrid, Qwen3.5, Qwen3.5-MoE, Qwen3-Next. The modeling files are regenerated from the modular sources, so make fix-repo and check_modular_conversion stay clean.

Reproduction, evidence, and environment

Real-last-token logits, unpadded vs the same sequence padded, max abs diff (right padding compares the last real token, which precedes the pad; left padding compares position -1):

                       float32         bfloat16
right pad (b=1)        ~1e-5           ~0        # real tokens precede pad, causal -> unaffected
left  pad (b>1)        ~2e-5           ~0.19     # mask already runs; bf16 has a left-pad precision floor
left  pad (b=1) before  5.7            5.7       # mask skipped -> padding leaks into the real tokens
left  pad (b=1) after  ~9e-6           ~0.13     # fixed: back to each dtype's b>1 baseline

float32 isolates the bug from precision noise: batch-size-1 left padding jumps from ~1e-5 to 5.7 and the fix restores it. In bfloat16 the corruption is the same magnitude (5.7) and the fix brings batch size 1 back to the ~0.19 floor that batch size > 1 already sits at. Per layer the divergence originates at the first gated-delta-net layer.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("Qwen/Qwen3.5-0.8B")
m = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3.5-0.8B", dtype=torch.float32).eval()
seq = tok("The capital of France is one of the most", return_tensors="pt").input_ids

ref = m(seq, use_cache=False).logits[0, -1]
pad = torch.full((1, 4), tok.pad_token_id)
ids = torch.cat([pad, seq], 1)
mask = torch.cat([torch.zeros(1, 4, dtype=torch.long), torch.ones_like(seq)], 1)
out = m(ids, attention_mask=mask, use_cache=False).logits[0, -1]
print((ref - out).abs().max())   # ~5.7 before, ~1e-5 after

Environment: transformers main (5.13.0.dev0), torch 2.8.0+cu128, Python 3.12.3, Linux. The float32 numbers reproduce on CPU (torch fallback); the bfloat16 column was measured on an RTX 4090.

Code Agent Policy

• I confirm that this is not a pure code agent PR.

Who can review?

@zucchini-nlp

[github-actions]

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

run-slow: bamba, falcon_h1, granitemoehybrid, lfm2, lfm2_moe, mamba2, olmo_hybrid, qwen3_5, qwen3_5_moe, qwen3_next

[github-actions]

CI Dashboard: View test results in Grafana

[zucchini-nlp]

I guess the intuition was to skip non-padded inputs, and usually an input of bs==1 or seq-len==1 are not padded

Interesting that you have padding in a single input

[Sunt-ing]

Right, the original attention_mask.shape[0] > 1 was a proxy for "is this input padded", assuming a single sequence is never padded. But left-padding a single prompt is a legitimate, common case (e.g. tokenizer(..., padding="max_length"), or a batch=1 request in a serving setup), and there it silently corrupts the real tokens because the pad positions are never zeroed before the conv/SSM recurrence.

I replaced the batch-size proxy with attention_mask.shape[1] == hidden_states.shape[1]: apply the mask exactly when the 2D mask aligns with the current hidden states (the padded prefill), which also naturally skips the cached decode step where the mask spans past+new tokens and is longer than the single-step hidden states.

[zucchini-nlp]

hmm, padded single sequences feel like a weird edge case though I get the reason ppl might be using it (static-shaped traced model ig?)

Let me trigger slow integration CI

[zucchini-nlp]

run-slow: bamba, falcon_h1, granitemoehybrid, lfm2, lfm2_moe, mamba2, olmo_hybrid, qwen3_5, qwen3_5_moe, qwen3_next

[github-actions]

Workflow Run ⚙️

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

models: ["models/bamba", "models/falcon_h1", "models/granitemoehybrid", "models/lfm2", "models/lfm2_moe", "models/mamba2", "models/olmo_hybrid", "models/qwen3_5", "models/qwen3_5_moe", "models/qwen3_next"]
quantizations: []

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[github-actions]

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	846d0422	workflow commit (merge commit)
PR	c1b2cb4c	branch commit (from PR)
main	decfe6a7	base commit (on main)

Model CI Report

❌ 1 new failed tests from this PR 😭

• mamba2:
  tests/models/mamba2/test_modeling_mamba2.py::Mamba2IntegrationTest::test_simple_generate (❌ ⟹ ❌)

[Sunt-ing]

@zucchini-nlp Slow CI is clean: every failure also fails on base (decfe6a7) per the bot's diff (mostly OOM and stale string asserts). The one flagged "new", mamba2 test_simple_generate, also fails on base; it's an unpadded single sequence, where the new guard is just an all-ones no-op, so this PR can't change it.

On the edge case: it's not only traced/compiled models. A batch-of-1 request through a padded collator and tokenizer(..., padding="max_length") both yield a left-padded single sequence at inference. The point is it fails silently, no error, just wrong logits on the real tokens (~5.7 in fp32), corrupting sampling/logprobs/scoring. The guard strictly narrows when the mask runs, so batch>1, right padding, and cached decode stay byte-identical.