[Model][Spec Decode] Nemotron-H MTP and Mamba Speculative Decoding Support

[benchislett]

Purpose

This PR adds support for MTP for the Nemotron-H model family, which will be introduced with Nemotron V3 Super.

To facilitate this, we also implement speculative decoding support for the Mamba attention backends. Previously, mamba-style speculative decoding support was limited to Qwen3-Next. This PR attempts to implement the attention metadata in a simple and unified manner that does not introduce too much complexity to the backend.

Co-authored with @shaharmor98

Design

The core change to mamba attention is to separate prefill and decode state indices, since decode state indices will have spec tokens on the second axis. When using specdec we also add num_accepted_tokens and query_start_loc_d (decode) for indexing into the decode batch and selecting the right ssm cache states. Kernel support already exists for this.

Test Plan

Testing is currently all local with checkpoints of Nemotron V3 Super MTP. GSM8k is consistent between MTP on/off.

Open to suggestions for a test plan for the mamba specdec pathway. I'm not sure if there are any good test model/draft combinations we could use. Maybe n-gram with nemotron nano v3?

Example run command:

vllm serve $MODEL --no-enable-chunked-prefill --no-enable-prefix-caching --max-model-len 32768 --max-num-batched-tokens 32768 --trust-remote-code -tp 2 --speculative-config '{"method": "nemotron_h_mtp", "num_speculative_tokens": 7}'

Limitations

Most of the feature incompatibility of this PR is due to feature gaps in the way vLLM handles speculative decoding for linear attention models. Specifically, the following features are confirmed to cause bugs and/or crashes when used with speculative decoding and mamba2:

• Chunked Prefill
• Prefix Caching (old style, have not tested 'align' mode)
• supports_update_block_table (will be disabled for mamba2 when specdec is on)

Also, asynchronous scheduling is ineffective with linear+specdec since _update_states_after_model_execute causes a synchronization. I alleviated issues here by moving it to after the drafting. Full overlap can be achieved by reusing logic from _copy_valid_sampled_token_count, but my preliminary attempt crashed due to a suspected race condition, so I think full overlap can be left as a future performance optimization.

[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]

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