Fix use_cache with seq_len > 1 ( #46032) by Ramshankar07 · Pull Request #46084 · huggingface/transformers

Ramshankar07 · 2026-05-19T20:01:34Z

What does this PR do?

Add a third dispatch path to Mamba2Mixer for the case where a prior cached state exists and seq_len > 1 (chunked prefill / speculative decode verification).

Previously:

cuda_kernels_forward: .squeeze(1) crashed or corrupted tensors for seq_len > 1 when has_previous_state was True.
torch_forward: dt[:, 0, :] silently dropped all tokens after index 0, and previous_states was always zero-initialised ignoring the cache.

Both paths now:

Gate the single-step kernel on seq_len == 1.
For seq_len > 1, prepend the cached conv buffer (last K-1 entries) to the chunk input, run a full causal conv, drop the prepended prefix from the output, and pass the cached recurrent_state as initial_states to mamba_chunk_scan_combined / as previous_states in the naive SSD.

This PR follows the same pattern applied to GDN-based models in #45513, as discussed in the issues I added test_mamba2_chunked_prefill (CPU/torch path, always runs) and test_mamba2_chunked_prefill_cuda_path (skipped without mamba-ssm).

Fixes #46032

Code Agent Policy

The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by
code agents. We are currently bottlenecked by our ability to review and respond to them. As a result,
we ask that new users do not submit pure code agent PRs at this time.
You may use code agents in drafting or to help you diagnose issues. We'd also ask autonomous "OpenClaw"-like agents
not to open any PRs or issues for the moment.

PRs that appear to be fully agent-written will probably be closed without review, and we may block users who do this
repeatedly or maliciously.

This is a rapidly-evolving situation that's causing significant shockwaves in the open-source community. As a result,
this policy is likely to be updated regularly in the near future. For more information, please read CONTRIBUTING.md.

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).
Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case. Mamba2Mixer: use_cache with seq_len > 1 silently produces incorrect results (both CPU and GPU paths) #46032
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@zucchini-nlp @Cyrilvallez 🤗@vasqu @clara2911

Ramshankar07 · 2026-05-19T20:02:38Z

I didn't run the cuda tests, can anybody test it?

vasqu

Some initial comments from my side, this is not about functionality (it seems correct to me) but more on style

Ideally we would propogate the changes to other mamba2 related models like bamba, falcon, zamba2 etc. But no worries let's focus on getting the base mamba2 right first

vasqu

Sorry that im so nitpicky but I really want to avoid multiple different versions of essentially the same. We should follow qwen3 next as close as possible where we use the precomputed states and seq len var where we need + more aligned forward

vasqu · 2026-06-01T17:36:58Z


        # Single step calculations via cache
-        if cache_params is not None and cache_params.has_previous_state(self.layer_idx):
+        if cache_params is not None and cache_params.has_previous_state(self.layer_idx) and seq_len == 1:


Please let's follow Bamba and make this a var use_precomputed_states = ...

transformers/src/transformers/models/bamba/modeling_bamba.py

Line 500 in 1aa495d

use_precomputed_states = (

but without the sequence length (we want this to be a separate check similar to qwen3 next)

vasqu · 2026-06-01T17:37:10Z

        else:
            A = -torch.exp(self.A_log.float())  # (num_heads) or (intermediate_size, state_size)
            dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit}
+            has_previous_state = cache_params is not None and cache_params.has_previous_state(self.layer_idx)


Then we dont need this

vasqu · 2026-06-01T17:40:38Z

                    return_final_states=True,
                    dt_bias=self.dt_bias,
                    dt_softplus=True,
+                    initial_states=cache_params.layers[self.layer_idx].recurrent_states if has_previous_state else None,


We can extract the prev states (conv and recurrent before) similar to qwen3 next

transformers/src/transformers/models/qwen3_next/modeling_qwen3_next.py

Lines 613 to 616 in 1aa495d

# getting projected states from cache if it exists

if use_precomputed_states:

conv_state = cache_params.layers[self.layer_idx].conv_states

recurrent_state = cache_params.layers[self.layer_idx].recurrent_states

vasqu · 2026-06-01T17:42:28Z

-                    )
-                else:
-                    hidden_states_B_C = causal_conv1d_fn(
+                if has_previous_state:


Would also be nice if we could follow qwen3 next more closely here

transformers/src/transformers/models/qwen3_next/modeling_qwen3_next.py

Lines 626 to 656 in 1aa495d

if use_precomputed_states and seq_len == 1:

# Single-token cached decode: the fused per-step kernel updates the conv state in-place.

mixed_qkv = self.causal_conv1d_update(

mixed_qkv,

conv_state,

self.conv1d.weight.squeeze(1),

self.conv1d.bias,

self.activation,

)

else:

# Multi-token forward (prefill, or chunked-tokens decode when the cache has prior state).

if use_precomputed_states:

# Cached chunked-tokens decode: prepend the cached conv context so the causal conv

# sees the correct left-context rather than zero-padding. Dropped from the output

# at the end of this branch.

mixed_qkv = torch.cat([conv_state, mixed_qkv], dim=-1)

if cache_params is not None:

new_conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0))

cache_params.update_conv_state(new_conv_state, self.layer_idx)

if self.causal_conv1d_fn is not None:

mixed_qkv = self.causal_conv1d_fn(

x=mixed_qkv,

weight=self.conv1d.weight.squeeze(1),

bias=self.conv1d.bias,

activation=self.activation,

seq_idx=kwargs.get("seq_idx"),

)

else:

mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, : mixed_qkv.shape[-1]])

if use_precomputed_states:

mixed_qkv = mixed_qkv[:, :, -seq_len:]

I just want to avoid having the same but in different ways, we should stay consistent

vasqu · 2026-06-01T17:42:53Z

+            cache_params is not None and cache_params.has_previous_state(self.layer_idx) and seq_len == 1
+        )
+        # Snapshot before any cache
+        has_previous_state = cache_params is not None and cache_params.has_previous_state(self.layer_idx)


Would avoid this extra var and just check with seq len where actually needed

Ramshankar07 · 2026-06-01T17:46:50Z

I was working to make it look closer like bamba script. I'll rework on it today

vasqu · 2026-06-01T17:51:39Z

Sorry I was to eager then 😬 feel free to ping me when ready, ignoring the thread for now then

Ramshankar07 · 2026-06-01T17:54:15Z

Sure thing, I'll let you know

Ramshankar07 · 2026-06-02T04:43:45Z

Should I change initial_states=cache_params.layers[self.layer_idx].conv_states[:, :, 1:], to initial_states=conv_state[:, :, 1:], to match the qwen's style?

Ramshankar07 · 2026-06-02T04:49:12Z

I changed for now, @vasqu lmk if that's right

vasqu

Smaller nits but would you be up to fixing this on any mamba2 related model? I think this is a nice opportunity to do so

See it as pretty much approved for the core structure

Ramshankar07 · 2026-06-13T18:06:17Z

Hey @vasqu , Got bit busy lately. I have modified based on your suggestion. Can you take look and lmk if I'm missing any. I couldn't run test_mamba2_chunked_prefill_cuda_path(), test_mamba2_slow_vs_fast_forward() , test_mamba2_slow_vs_fast_forward_grouped() cause of GPU. if every thing looks good. I'll try to rent a gpu to run.

vasqu

Heya looks pretty good, I made a suggestion on the cuda fast path because I want this to be closer to the qwen implementation. Lmk what you think

Other than that should be ready to merge then 🤗

Ramshankar07 · 2026-06-21T19:00:02Z

Made changes based on your suggestion, let me know if I miss any

Match mamba2 (huggingface#46084) and the qwen3-next reference by prepending the full cached conv_states instead of conv_states[:, :, 1:]. The two are numerically identical: the extra oldest column only adds one warm-up conv position that the trailing [-seq_len:] slice drops. This keeps the chunked-prefill left-context consistent with the reference path. Signed-off-by: Ting Sun <suntcrick@gmail.com>

vasqu

The forward looks good to me now. I would only adjust the tests a bit then we can merge 🤗

vasqu · 2026-06-22T14:43:34Z


-    def create_and_check_mamba2_slow_vs_fast_forward(self, config, input_ids, *args, gradient_checkpointing=False):
-        model = Mamba2Model(config)
+    def create_and_check_mamba2_chunked_prefill(self, config, *args, device="cpu"):


Hmm, should we not adapt

transformers/tests/models/qwen3_5/test_modeling_qwen3_5.py

Line 203 in f007617

def test_linear_attention_multi_token_cached_forward_matches_single_token(self):

for mamba2 instead?

just the self?

I mean that the test in qwen 3.5 is already testing what we need, we don't need to come up with something else can just reuse the patterns there

I changed the single token prefill and added other changes but not sure at this

…on-H cuda_kernels_forward had the same defect as the torch path: a multi-token forward with a primed cache (chunked prefill / speculative verification) was routed into the single-step causal_conv1d_update / selective_state_update kernels and crashed ("weight must have shape (dim, width)"). Gate the single-step branch on seq_len == 1; for the multi-token cached case run the full kernels seeded from the cache (prepend the cached conv left-context and drop it after the conv, pass the cached recurrent state as initial_states to mamba_chunk_scan_combined), mirroring huggingface#46084. Signed-off-by: Ting Sun <suntcrick@gmail.com>

Add a third dispatch path to Mamba2Mixer for the case where a prior cached state exists and seq_len > 1 (chunked prefill / speculative decode verification). Previously: - cuda_kernels_forward: .squeeze(1) crashed or corrupted tensors for seq_len > 1 when has_previous_state was True. - torch_forward: dt[:, 0, :] silently dropped all tokens after index 0, and previous_states was always zero-initialised ignoring the cache. Both paths now: - Gate the single-step kernel on seq_len == 1. - For seq_len > 1, prepend the cached conv buffer (last K-1 entries) to the chunk input, run a full causal conv, drop the prepended prefix from the output, and pass the cached recurrent_state as initial_states to mamba_chunk_scan_combined / as previous_states in the naive SSD. This PR follows the same pattern applied to GDN-based models in huggingface#45513, as discussed in the issues I added test_mamba2_chunked_prefill (CPU/torch path, always runs) and test_mamba2_chunked_prefill_cuda_path (skipped without mamba-ssm).

…lution

…proving cache usage

…ice execution paths

Ramshankar07 · 2026-06-23T14:11:34Z

I think it looks clean now! imo

@vasqu have look when you're free

github-actions · 2026-06-23T14:31:37Z

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

run-slow: mamba2

vasqu

Should be good to go now, just using slow CI in a sec to sanity check

vasqu · 2026-06-23T14:32:54Z

run-slow: mamba2

github-actions · 2026-06-23T14:34:27Z

Workflow Run ⚙️

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

models: ["models/mamba2"]
quantizations: []

github-actions · 2026-06-23T14:49:16Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	6155764d	workflow commit (merge commit)
PR	bb795032	branch commit (from PR)
main	5bbed98a	base commit (on `main`)

✅ No failing test specific to this PR 🎉 👏 !

HuggingFaceDocBuilderDev · 2026-06-23T14:49:49Z

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.

vasqu · 2026-06-23T15:16:59Z

Also checked locally and seems fine, merging now thanks a lot!

Ramshankar07 · 2026-06-23T15:24:05Z

Thanks for your suggestions! Got many lessons on mistakes that I make.

Reshape the regression test to match the form merged in huggingface#46084: a model-level first-token causal check (the first position of a multi-token cached forward must equal the single-token cached forward), with separate CPU and accelerator variants. Nemotron-H has no CPU variant because its MoE layers use grouped_mm, which has no CPU kernel. Signed-off-by: Ting Sun <suntcrick@gmail.com>

…eMoeHybrid These models carry their own Mamba2 mixer and gate the single-step branch on seq_len == 1, but their multi-token cached path still starts from a zero conv context and zero SSM state, so chunked-prefill / speculative verification is history-less. Seed the multi-token cached path from the cache (prepend the cached conv left-context, pass the cached recurrent state as the SSM initial state), mirroring the Zamba2/Nemotron-H fix and huggingface#46084. Add the chunked-prefill regression test to each. Signed-off-by: Ting Sun <suntcrick@gmail.com>

…n-H, Bamba, FalconH1 and GraniteMoeHybrid (huggingface#46741) * Fix Mamba2 mixer chunked-prefill/speculative decode for Zamba2 and Nemotron-H The Mamba2 mixer took the single-step decode branch whenever the cache held a previous state, ignoring seq_len. For seq_len > 1 with a non-empty cache (chunked prefill, speculative/assisted decode verification) it used only the first tokens dt plus one conv/SSM step, silently producing wrong logits for the rest of the chunk and corrupting the cached state. Gate the single-step branch on seq_len == 1; for the multi-token cached case run the full path seeded from the cache (prepend the cached conv left context, pass the cached recurrent state as the initial SSM state). Same root cause and fix as mamba2 (huggingface#46032); applied in modular_zamba2.py so the change propagates to nemotron_h, which inherits the mixer. Signed-off-by: Ting Sun <suntcrick@gmail.com> * Prepend the full cached conv_states in the Mamba2 chunked path Match mamba2 (huggingface#46084) and the qwen3-next reference by prepending the full cached conv_states instead of conv_states[:, :, 1:]. The two are numerically identical: the extra oldest column only adds one warm-up conv position that the trailing [-seq_len:] slice drops. This keeps the chunked-prefill left-context consistent with the reference path. Signed-off-by: Ting Sun <suntcrick@gmail.com> * Fix Mamba2 cuda_kernels_forward chunked-prefill for Zamba2 and Nemotron-H cuda_kernels_forward had the same defect as the torch path: a multi-token forward with a primed cache (chunked prefill / speculative verification) was routed into the single-step causal_conv1d_update / selective_state_update kernels and crashed ("weight must have shape (dim, width)"). Gate the single-step branch on seq_len == 1; for the multi-token cached case run the full kernels seeded from the cache (prepend the cached conv left-context and drop it after the conv, pass the cached recurrent state as initial_states to mamba_chunk_scan_combined), mirroring huggingface#46084. Signed-off-by: Ting Sun <suntcrick@gmail.com> * Align Zamba2/Nemotron-H chunked-prefill tests with merged mamba2 test Reshape the regression test to match the form merged in huggingface#46084: a model-level first-token causal check (the first position of a multi-token cached forward must equal the single-token cached forward), with separate CPU and accelerator variants. Nemotron-H has no CPU variant because its MoE layers use grouped_mm, which has no CPU kernel. Signed-off-by: Ting Sun <suntcrick@gmail.com> * Propagate Mamba2 chunked-prefill cache fix to Bamba, FalconH1, GraniteMoeHybrid These models carry their own Mamba2 mixer and gate the single-step branch on seq_len == 1, but their multi-token cached path still starts from a zero conv context and zero SSM state, so chunked-prefill / speculative verification is history-less. Seed the multi-token cached path from the cache (prepend the cached conv left-context, pass the cached recurrent state as the SSM initial state), mirroring the Zamba2/Nemotron-H fix and huggingface#46084. Add the chunked-prefill regression test to each. Signed-off-by: Ting Sun <suntcrick@gmail.com> * Extract cached conv/recurrent state into locals in the Mamba2 mixers Bind the cached conv_states / recurrent_states once after use_precomputed_states, matching mamba2, instead of re-reading cache_params.layers[self.layer_idx] inline in the cat / initial_states / single-step paths. Pure refactor, no behavior change. Signed-off-by: Ting SUN <suntcrick@gmail.com> * Add CPU variant for the nemotron_h chunked-prefill test NemotronH MoE experts dispatch through integrations.moe to torch._grouped_mm, which only gained a CPU kernel in torch 2.11; guard the new _cpu test with require_torch_greater_or_equal("2.11") so it runs on current CI and skips on older torch. Signed-off-by: Ting SUN <suntcrick@gmail.com> * rename a bit --------- Signed-off-by: Ting Sun <suntcrick@gmail.com> Signed-off-by: Ting SUN <suntcrick@gmail.com> Co-authored-by: vasqu <antonprogamer@gmail.com>

clara2911 reviewed May 20, 2026

View reviewed changes

Comment thread src/transformers/models/mamba2/modeling_mamba2.py Outdated

Rocketknight1 mentioned this pull request May 20, 2026

Fix Mamba2Mixer: use_cache with seq_len > 1 silently produces incorrect results (both CPU and GPU paths) (#46032) #46098

Closed

5 tasks

Ramshankar07 force-pushed the fix/46032-mamba2-chunked-prefill branch from 6bc871f to 03f48bf Compare May 20, 2026 18:50

vasqu reviewed May 25, 2026

View reviewed changes

Comment thread src/transformers/models/mamba2/modeling_mamba2.py Outdated

Comment thread src/transformers/models/mamba2/modeling_mamba2.py Outdated

Comment thread src/transformers/models/mamba2/modeling_mamba2.py Outdated

Comment thread src/transformers/models/mamba2/modeling_mamba2.py Outdated

Ramshankar07 force-pushed the fix/46032-mamba2-chunked-prefill branch from df45189 to aac20d1 Compare May 31, 2026 00:53

vasqu reviewed Jun 1, 2026

View reviewed changes

Ramshankar07 force-pushed the fix/46032-mamba2-chunked-prefill branch from 4e08cde to 9752bab Compare June 2, 2026 04:18

vasqu reviewed Jun 2, 2026

View reviewed changes

Comment thread src/transformers/models/mamba2/modeling_mamba2.py Outdated

Comment thread tests/models/mamba2/test_modeling_mamba2.py Outdated

Comment thread tests/models/mamba2/test_modeling_mamba2.py Outdated

Ramshankar07 force-pushed the fix/46032-mamba2-chunked-prefill branch from 4e51f2f to 720ca05 Compare June 13, 2026 17:36

vasqu reviewed Jun 15, 2026

View reviewed changes

Comment thread src/transformers/models/mamba2/modeling_mamba2.py

Comment thread tests/models/mamba2/test_modeling_mamba2.py Outdated

Sunt-ing mentioned this pull request Jun 18, 2026

Fix Mamba2 chunked-prefill / speculative decoding for Zamba2, Nemotron-H, Bamba, FalconH1 and GraniteMoeHybrid #46741

Merged

6 tasks

Ramshankar07 force-pushed the fix/46032-mamba2-chunked-prefill branch from c769ea0 to 13efaa4 Compare June 19, 2026 01:13

vasqu approved these changes Jun 22, 2026

View reviewed changes

Ramshankar07 force-pushed the fix/46032-mamba2-chunked-prefill branch from ffb3df6 to d805833 Compare June 23, 2026 13:43

Ramshankar07 added 6 commits June 23, 2026 10:10

Refactor convolution logic to utilize initial_states for causal convo…

d9e7fef

…lution

Code quality

78f8228

merging the chunked-prefill elif branch, renaming

a6d446e

making the flow cleaner

2c75d62

Refactor of has_previous_states and others.

d474fb1

Ramshankar07 added 4 commits June 23, 2026 10:10

adding in require_* decorator format

f48bec7

adding in __init__

fd5dcb0

Refactor Mamba2Mixer to make it close to qwen 3 implementation and im…

8e9fa80

…proving cache usage

Updated input handling and added separate tests for CPU and torch dev…

3424102

…ice execution paths

Ramshankar07 force-pushed the fix/46032-mamba2-chunked-prefill branch from e744224 to 3424102 Compare June 23, 2026 14:10

adjust tests slightly to require kernels instead

bb79503

vasqu approved these changes Jun 23, 2026

View reviewed changes

vasqu added this pull request to the merge queue Jun 23, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Jun 23, 2026

vasqu added this pull request to the merge queue Jun 23, 2026

Merged via the queue into huggingface:main with commit 892be11 Jun 23, 2026
104 checks passed

abcgco mentioned this pull request Jul 5, 2026

Hybrid models: padding is not masked out of linear-attention/conv state on continued forwards (chunked prefill, cache continuation) #47086

Open

4 tasks

	# getting projected states from cache if it exists
	if use_precomputed_states:
	conv_state = cache_params.layers[self.layer_idx].conv_states
	recurrent_state = cache_params.layers[self.layer_idx].recurrent_states

	if use_precomputed_states and seq_len == 1:
	# Single-token cached decode: the fused per-step kernel updates the conv state in-place.
	mixed_qkv = self.causal_conv1d_update(
	mixed_qkv,
	conv_state,
	self.conv1d.weight.squeeze(1),
	self.conv1d.bias,
	self.activation,
	)
	else:
	# Multi-token forward (prefill, or chunked-tokens decode when the cache has prior state).
	if use_precomputed_states:
	# Cached chunked-tokens decode: prepend the cached conv context so the causal conv
	# sees the correct left-context rather than zero-padding. Dropped from the output
	# at the end of this branch.
	mixed_qkv = torch.cat([conv_state, mixed_qkv], dim=-1)
	if cache_params is not None:
	new_conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0))
	cache_params.update_conv_state(new_conv_state, self.layer_idx)
	if self.causal_conv1d_fn is not None:
	mixed_qkv = self.causal_conv1d_fn(
	x=mixed_qkv,
	weight=self.conv1d.weight.squeeze(1),
	bias=self.conv1d.bias,
	activation=self.activation,
	seq_idx=kwargs.get("seq_idx"),
	)
	else:
	mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, : mixed_qkv.shape[-1]])
	if use_precomputed_states:
	mixed_qkv = mixed_qkv[:, :, -seq_len:]

Uh oh!

Conversation

Ramshankar07 commented May 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Code Agent Policy

Before submitting

Who can review?

Uh oh!

Ramshankar07 commented May 19, 2026

Uh oh!

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Ramshankar07 commented Jun 1, 2026

Uh oh!

vasqu commented Jun 1, 2026

Uh oh!

Ramshankar07 commented Jun 1, 2026

Uh oh!

Ramshankar07 commented Jun 2, 2026

Uh oh!

Ramshankar07 commented Jun 2, 2026

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Ramshankar07 commented Jun 13, 2026

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Ramshankar07 commented Jun 21, 2026

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Ramshankar07 commented Jun 23, 2026

Uh oh!

github-actions Bot commented Jun 23, 2026

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Ramshankar07 commented May 19, 2026 •

edited

Loading