Repository navigation
Pass kwargs to the Mamba2 mixer in Nemotron-H, Falcon-H1 and Mamba2 - #48490
Conversation
NemotronHBlock forwards **kwargs to the attention mixer but dropped them for the linear-attention branch, so kernel kwargs such as seq_idx never reached the mixer. The mixer already splats them into the fused conv1d+scan, the conv fallback and the chunk scan, and Bamba, GraniteMoEHybrid and Zamba2 all pass them through. Co-authored-by: Cursor <cursoragent@cursor.com>
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Just some small comments, is this issue on other models? E.g. the mamba2 related models could all have the same issue, could we cross check pls
| model(input_ids, seq_idx=seq_idx) | ||
|
|
||
| self.assertTrue(seen, "the linear-attention mixer was never called") | ||
| self.assertIn("seq_idx", seen[0]) |
There was a problem hiding this comment.
Can we change the test design as in test_mamba2_slow_vs_fast_forward so we actually do the heavy lifting in the tester
There was a problem hiding this comment.
Done — moved to create_and_check_kwargs_reach_mamba2_mixer on the tester, with the test method just calling it. Added the same pair to the falcon_h1 and mamba2 testers.
| def recording_forward(*args, **kwargs): | ||
| seen.append(set(kwargs)) | ||
| # The torch fallback kernels do not implement seq_idx, so only record it. | ||
| kwargs.pop("seq_idx", None) |
There was a problem hiding this comment.
Imo this pop is not needed no?
There was a problem hiding this comment.
Correct, removed. The three fallbacks (mamba2_split_conv1d_scan_combined, causal_conv1d_fn, mamba2_chunk_scan) all take **kwargs, so seq_idx passes through harmlessly on the torch path — I had added the pop defensively without checking.
Same gap as nemotron_h: the block accepts **kwargs but drops them at the mixer, so kernel kwargs such as seq_idx never reach the kernels. In these two the model also did not pass kwargs down to the block. Also move the nemotron_h test into the tester and drop the seq_idx pop, which is not needed since the torch fallback kernels take **kwargs. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Cross-checked all six models with a Mamba2 mixer, on both hops (model → block, block → mixer):
So you were right that it was not just Nemotron-H.
Two things I noticed but left alone, happy to fold in either way:
|
kfastino
changed the title
|
[For maintainers] Suggested jobs to run (before merge) run-slow: falcon_h1, mamba2, nemotron_h |
CI recapDashboard: View test results in Grafana |
|
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. |
…uggingface#48490) * Pass kwargs to the Nemotron-H Mamba2 mixer NemotronHBlock forwards **kwargs to the attention mixer but dropped them for the linear-attention branch, so kernel kwargs such as seq_idx never reached the mixer. The mixer already splats them into the fused conv1d+scan, the conv fallback and the chunk scan, and Bamba, GraniteMoEHybrid and Zamba2 all pass them through. Co-authored-by: Cursor <cursoragent@cursor.com> * Pass kwargs to the Mamba2 mixer in falcon_h1 and mamba2 too Same gap as nemotron_h: the block accepts **kwargs but drops them at the mixer, so kernel kwargs such as seq_idx never reach the kernels. In these two the model also did not pass kwargs down to the block. Also move the nemotron_h test into the tester and drop the seq_idx pop, which is not needed since the torch fallback kernels take **kwargs. Co-authored-by: Cursor <cursoragent@cursor.com> --------- Co-authored-by: Cursor <cursoragent@cursor.com>
What does this PR do?
Three models with a Mamba2 mixer never deliver kernel kwargs to it, because the block accepts
**kwargsand then drops them at the mixer call:The Mamba2 mixer is already written to receive them — it splats them into all three kernel calls (
mamba2_split_conv1d_scan_combinedviafused_kwargs,causal_conv1d_fn, andmamba2_chunk_scan) — so this is the missing hop.The main consequence is
seq_idx, which the Mamba2 kernels use to keep state from bleeding across concatenated sequences. Without it there is no way to run these models on packed / variable-length batches: the conv and scan treat a packed row as one long sequence. With this change,model(input_ids, seq_idx=seq_idx)reaches the kernels.Which models
Checked all six models with a Mamba2 mixer, on both hops:
bamba**kwargs**kwargsgranitemoehybrid**kwargs**kwargszamba2**kwargs**kwargsnemotron_h**kwargsfalcon_h1mamba2bamba,granitemoehybridandzamba2already forward kwargs in exactly this position, so this brings the other three in line rather than introducing a new pattern.falcon_h1andmamba2lost them at both hops.nemotron_handfalcon_h1are changed through their modular files and regenerated;mamba2has no modular file.Tests
test_kwargs_reach_mamba2_mixerplus acreate_and_check_kwargs_reach_mamba2_mixertester method for each of the three models, asserting a kwarg handed to the model arrives at the mixer. Each fails without its corresponding fix.utils/check_modular_conversion.pyandruff check/ruff formatare clean.Context
Found while adding Mamba2 packing support to Unsloth (unslothai/unsloth#9812). Without this, reaching the kernels from outside requires rebinding the kernel symbol in the modeling module's globals, which this makes unnecessary.
Who can review?
@ArthurZucker @vasqu