[Bugfix] Fix NemotronH MTP + Chunked Prefill

[benchislett]

Purpose

Functional

Test Plan

This branch adds a reproducer which causes garbage outputs with NemotronH MTP + Chunked Prefill.

It does not seem to happen with Qwen3-Next due to differences in how they separate out the spec decodes from non-spec decodes (GDN checks num_draft_tokens_cpu and dynamically splits the decodes into spec and non-spec).

Test Result

The diff in this branch fixes the reproducer giving results consistent with baseline. Further evaluation required.

[gemini-code-assist]

Code Review

This pull request addresses a bug in NemotronH models when using Multi-Token Prediction (MTP) with chunked prefill. The core fix in vllm/v1/attention/backends/mamba_attn.py correctly handles cases where small prefill chunks are misclassified as decodes, preventing an assertion failure. While this fix is sound, the pull request includes some changes that require attention. There's a block of dead code in vllm/v1/worker/gpu_model_runner.py that seems to be a work-in-progress and should be cleaned up before merging. Additionally, an unrelated change in vllm/model_executor/layers/layernorm.py has been identified, which should be reverted to avoid potential side effects.

[mergify]

Hi @benchislett, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mergify]

Hi @benchislett, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @benchislett.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

Hi @benchislett, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mergify]

Hi @benchislett, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[benchislett]

Not sure why pre-commit is breaking. Having a hard time figuring out what the discrepancy is

[tdoublep]

Do we really want to have these prefill chunks processed by the decode path?

[benchislett]

This is how it's done in all other attention backends. It's only GDN attention that does it differently, by manually splitting out the spec-decodes and non-spec-decodes. In my opinion, that strategy is inefficient and more difficult to maintain.

[tdoublep]

I guess we'd prefer not to have mamba2 specific logic in GPU model runner if possible

[benchislett]

Makes sense. I can try to refactor the behaviour out of gpu model runner and into a helper somewhere

[benchislett]

Moved the core logic into mamba_utils and renamed this toggle. But I'm still not sure what's the cleanest way to set this flag so that we can dispatch based on mamba2 vs gdn attention

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @benchislett.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[tdoublep]

Generally LGTM but would like @asafgardin's eyes on it from the Mamba1 perspective

[tdoublep]

Seems like there is a different (?) corner case being explored here: #32716

[Josephasafg]

Generally LGTM but would like @asafgardin's eyes on it from the Mamba1 perspective

Thanks @tdoublep
Since Mamba1 does not yet have support for speculative decoding, this change should be ok

[benchislett]

No idea why the hybrid test is hanging in CI. It passes locally.

[benchislett]

Looking into it.

[Josephasafg]

is this assertion still necessary?

[benchislett]

nope :p

[benchislett]

Tagging @vadiklyutiy to help with the CI issue, I cannot reproduce it

[tdoublep]

I can reproduce the hang locally on an L4 GPU (but not on H100)

[tdoublep]

It looks like for the test that's hanging we only have 38 blocks available on L4, and the request requires 100+ blocks (we need one block per speculative token in MTP for hybrid models) so it just sits there waiting to be scheduled. Can we change the test to require less blocks?

[benchislett]

added marks to limit the test to >= H100, and added coverage of nemotron-h + MTP