@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f5972a872f

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 7, 19:00 UTC

[Model][Spec Decode] Nemotron-H MTP and Mamba Speculative Decoding Support (#33726)

Signed-off-by: Shahar Mor <smor@nvidia.com> Signed-off-by: Benjamin Chislett <bchislett@nvidia.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Shahar Mor <smor@nvidia.com> Co-authored-by: Roi Koren <roik@nvidia.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>

Benjamin Chislett committed Feb 24, 2026 at 12:49 UTC f5972a872fa3fabd94b7a6c6f031f4b5bcee2b2d

19 files changed +800 -158

tests/models/registry.py

@@ -1200,6 +1200,11 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {

1200 },

1201 is_available_online=False,

1202 ),

1203 + "NemotronHMTPModel": _HfExamplesInfo(

1204 + "nvidia/Nemotron-Super-Placeholder",

1205 + speculative_model="nvidia/Nemotron-Super-Placeholder",

1206 + is_available_online=False,

1207 + ),

1208 }

1209

1210 _TRANSFORMERS_BACKEND_MODELS = {

tests/v1/attention/test_mamba_update_block_table.py

+7 -1

@@ -41,6 +41,9 @@ def _make_vllm_config(block_size, max_model_len, max_num_seqs):

41 cudagraph_mode=CUDAGraphMode.FULL,

42 max_cudagraph_capture_size=None,

43 ),

44 + speculative_config=None,

45 + num_speculative_tokens=0,

46 + parallel_config=SimpleNamespace(decode_context_parallel_size=1),

47 scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs),

48 model_config=SimpleNamespace(max_model_len=max_model_len),

49 )

@@ -92,7 +95,10 @@ def test_update_block_table_copies_block_idx_to_persistent_buffers():

95 has_initial_states_p=None,

96 query_start_loc_p=None,

97 num_computed_tokens_p=None,

95 - state_indices_tensor=builder_a.state_indices_tensor[:num_reqs],

98 + state_indices_tensor_p=None,

99 + query_start_loc_d=None,

100 + num_accepted_tokens=None,

101 + state_indices_tensor_d=builder_a.state_indices_tensor_d[:num_reqs],

102 block_idx_last_scheduled_token=(

103 builder_a.block_idx_last_scheduled_token[:num_reqs]

104 ),

vllm/config/speculative.py

+17 -2

@@ -36,6 +36,7 @@ MTPModelTypes = Literal[

36 "glm4_moe_lite_mtp",

37 "glm_ocr_mtp",

38 "ernie_mtp",

39 + "nemotron_h_mtp",

40 "exaone_moe_mtp",

41 "qwen3_next_mtp",

42 "qwen3_5_mtp",

@@ -255,6 +256,19 @@ class SpeculativeConfig:

256 {"n_predict": n_predict, "architectures": ["ErnieMTPModel"]}

257 )

258

259 + if (

260 + hf_config.model_type == "nemotron_h"

261 + and hasattr(hf_config, "num_nextn_predict_layers")

262 + and hf_config.num_nextn_predict_layers > 0

263 + ):

264 + # Check if this is an MTP variant

265 + hf_config.model_type = "nemotron_h_mtp"

266 + if hf_config.model_type == "nemotron_h_mtp":

267 + n_predict = getattr(hf_config, "num_nextn_predict_layers", 1)

268 + hf_config.update(

269 + {"n_predict": n_predict, "architectures": ["NemotronHMTPModel"]}

270 + )

271 +

272 if hf_config.model_type == "qwen3_next":

273 hf_config.model_type = "qwen3_next_mtp"

274 if hf_config.model_type == "qwen3_next_mtp":

@@ -325,7 +339,7 @@ class SpeculativeConfig:

339 if self.target_model_config is None:

340 raise ValueError("target_model_config must be present for mtp")

341 if self.target_model_config.hf_text_config.model_type == "deepseek_v32":

328 - # FIXME(luccafong): cudgraph with v32 MTP is not supported,

342 + # FIXME(luccafong): cudagraph with v32 MTP is not supported,

343 # remove this when the issue is fixed.

344 self.enforce_eager = True

345 # use the draft model from the same model:

@@ -427,7 +441,7 @@ class SpeculativeConfig:

441 self.method = "mtp"

442 if self.num_speculative_tokens > 1:

443 logger.warning(

430 - "Enabling num_speculative_tokens > 1 will run"

444 + "Enabling num_speculative_tokens > 1 will run "

445 "multiple times of forward on same MTP layer"

446 ",which may result in lower acceptance rate"

447 )

@@ -712,6 +726,7 @@ class SpeculativeConfig:

726 "hunyuan_vl",

727 "hunyuan_v1_dense",

728 "afmoe",

729 + "nemotron_h",

730 ]

731 if (

732 self.method == "eagle3"

vllm/config/vllm.py

@@ -395,6 +395,15 @@ class VllmConfig:

395 ]

396 return hash_str

397

398 + @property

399 + def num_speculative_tokens(self) -> int:

400 + if (

401 + self.speculative_config is not None

402 + and self.speculative_config.num_speculative_tokens is not None

403 + ):

404 + return self.speculative_config.num_speculative_tokens

405 + return 0

406 +

407 @property

408 def needs_dp_coordinator(self) -> bool:

409 """

vllm/model_executor/layers/mamba/abstract.py

-8

@@ -41,14 +41,6 @@ class MambaBase(AttentionLayerBase):

41 pass

43 def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:

44 - if (

45 - vllm_config.speculative_config is not None

46 - and vllm_config.model_config.hf_config.model_type

47 - not in ["qwen3_next", "qwen3_5", "qwen3_5_moe"]

48 - ):

49 - raise NotImplementedError(

50 - "Mamba with speculative decoding is not supported yet."

51 - )

44 mamba_block_size = vllm_config.cache_config.mamba_block_size

45 page_size_padded = vllm_config.cache_config.mamba_page_size_padded

46 return MambaSpec(

vllm/model_executor/layers/mamba/mamba_mixer.py

+2 -18

@@ -265,7 +265,8 @@ class MambaMixer(MambaBase, PluggableLayer):

265 attn_metadata = attn_metadata[self.prefix]

266 assert isinstance(attn_metadata, Mamba1AttentionMetadata)

267 query_start_loc_p = attn_metadata.query_start_loc_p

268 - state_indices_tensor = attn_metadata.state_indices_tensor

268 + state_indices_tensor_p = attn_metadata.state_indices_tensor_p

269 + state_indices_tensor_d = attn_metadata.state_indices_tensor_d

270 self_kv_cache = self.kv_cache[forward_context.virtual_engine]

271 conv_state = self_kv_cache[0].transpose(-1, -2)

272 ssm_state = self_kv_cache[1]

@@ -295,17 +296,13 @@ class MambaMixer(MambaBase, PluggableLayer):

296 prefill_decode_split = split_batch_to_prefill_and_decode(

297 hidden_states_BC,

298 gate,

298 - state_indices_tensor,

299 num_prefill_tokens,

300 - num_prefills,

300 num_decode_tokens,

301 )

302 hidden_states_BC_p = prefill_decode_split.hidden_states_BC_p

303 hidden_states_BC_d = prefill_decode_split.hidden_states_BC_d

304 gate_p = prefill_decode_split.gate_p

305 gate_d = prefill_decode_split.gate_d

307 - state_indices_tensor_p = prefill_decode_split.state_indices_tensor_p

308 - state_indices_tensor_d = prefill_decode_split.state_indices_tensor_d

306

307 if is_mamba_cache_all:

308 block_idx_last_computed_token_d, block_idx_last_computed_token_p = (

@@ -477,16 +474,12 @@ class PrefillDecodeSplit(NamedTuple):

474 hidden_states_BC_d: torch.Tensor

475 gate_p: torch.Tensor

476 gate_d: torch.Tensor

480 - state_indices_tensor_p: torch.Tensor

481 - state_indices_tensor_d: torch.Tensor

477

478

479 def split_batch_to_prefill_and_decode(

480 hidden_states_BC: torch.Tensor,

481 gate: torch.Tensor,

487 - state_indices_tensor: torch.Tensor,

482 num_prefill_tokens: int,

489 - num_prefills: int,

483 num_decode_tokens: int,

484 ) -> PrefillDecodeSplit:

485 num_actual_tokens = num_prefill_tokens + num_decode_tokens

@@ -501,20 +494,11 @@ def split_batch_to_prefill_and_decode(

494 gate[..., :num_actual_tokens], [num_decode_tokens, num_prefill_tokens], dim=-1

495 )

496

504 - # num_decode_tokens accounts for CUDA graph padding when applicable

505 - state_indices_tensor_d, state_indices_tensor_p = torch.split(

506 - state_indices_tensor[: num_decode_tokens + num_prefills],

507 - [num_decode_tokens, num_prefills],

508 - dim=0,

509 - )

510 -

497 return PrefillDecodeSplit(

498 hidden_states_BC_p=hidden_states_BC_p,

499 hidden_states_BC_d=hidden_states_BC_d,

500 gate_p=gate_p,

501 gate_d=gate_d,

516 - state_indices_tensor_p=state_indices_tensor_p,

517 - state_indices_tensor_d=state_indices_tensor_d,

502 )

503

504

vllm/model_executor/layers/mamba/mamba_mixer2.py

+27 -19

@@ -477,7 +477,8 @@ class MambaMixer2(MambaBase, PluggableLayer):

477 dim=-1,

478 )

479

480 - compilation_config = get_current_vllm_config().compilation_config

480 + vllm_config = get_current_vllm_config()

481 + compilation_config = vllm_config.compilation_config

482 if prefix in compilation_config.static_forward_context:

483 raise ValueError(f"Duplicate layer name: {prefix}")

484 compilation_config.static_forward_context[prefix] = self

@@ -488,6 +489,8 @@ class MambaMixer2(MambaBase, PluggableLayer):

489 self.cache_config = cache_config

490 self.prefix = prefix

491

492 + self.num_spec = vllm_config.num_speculative_tokens

493 +

494 # Pre-compute sizes for forward pass

495 self.tped_intermediate_size = self.intermediate_size // self.tp_size

496 self.tped_conv_size = self.conv_dim // self.tp_size

@@ -576,7 +579,6 @@ class MambaMixer2(MambaBase, PluggableLayer):

579 # conv_state = (..., dim, width-1) yet contiguous along 'dim'

580 conv_state = self_kv_cache[0].transpose(-1, -2)

581 ssm_state = self_kv_cache[1]

579 - state_indices_tensor = attn_metadata.state_indices_tensor

582 has_initial_states_p = attn_metadata.has_initial_states_p

583 prep_initial_states = attn_metadata.prep_initial_states

584 chunk_size = attn_metadata.chunk_size

@@ -584,6 +586,12 @@ class MambaMixer2(MambaBase, PluggableLayer):

586 query_start_loc_p = attn_metadata.query_start_loc_p

587 cu_chunk_seqlen_p = attn_metadata.cu_chunk_seqlen_p

588 last_chunk_indices_p = attn_metadata.last_chunk_indices_p

589 + state_indices_tensor_p = attn_metadata.state_indices_tensor_p

590 + state_indices_tensor_d = attn_metadata.state_indices_tensor_d

591 + num_accepted_tokens = attn_metadata.num_accepted_tokens

592 + query_start_loc_d = attn_metadata.query_start_loc_d

593 + num_decodes = attn_metadata.num_decodes

594 + num_decode_tokens = attn_metadata.num_decode_tokens

595

596 if attn_metadata is None:

597 # profile run

@@ -593,29 +601,21 @@ class MambaMixer2(MambaBase, PluggableLayer):

601 hidden_states, _B, _C = self.split_hidden_states_B_C_fn(hidden_states_B_C)

602 return hidden_states

603

596 - num_prefills = attn_metadata.num_prefills # request count

597 - num_decodes = attn_metadata.num_decode_tokens # token count (=request)

598 - num_prefill_tokens = attn_metadata.num_prefill_tokens # token count

604 + num_prefills = attn_metadata.num_prefills

605 + num_prefill_tokens = attn_metadata.num_prefill_tokens

606 has_prefill = num_prefills > 0

607 has_decode = num_decodes > 0

601 - num_actual_tokens = num_prefill_tokens + num_decodes

608 + num_actual_tokens = num_prefill_tokens + num_decode_tokens

609

603 - # Separate prefill and decode by splitting varlen input

610 # Split along token dimension

611 hidden_states_B_C_d, hidden_states_B_C_p = torch.split(

612 hidden_states_B_C[:num_actual_tokens],

607 - [num_decodes, num_prefill_tokens],

613 + [num_decode_tokens, num_prefill_tokens],

614 dim=0,

615 )

616 dt_d, dt_p = torch.split(

617 dt[:num_actual_tokens],

612 - [num_decodes, num_prefill_tokens],

613 - dim=0,

614 - )

615 - # Split along batch dimension

616 - state_indices_tensor_d, state_indices_tensor_p = torch.split(

617 - state_indices_tensor[:num_actual_tokens],

618 - [num_decodes, num_prefills],

618 + [num_decode_tokens, num_prefill_tokens],

619 dim=0,

620 )

621

@@ -642,16 +642,16 @@ class MambaMixer2(MambaBase, PluggableLayer):

642 )

643 num_computed_tokens_p = attn_metadata.num_computed_tokens_p

644 else:

645 - block_idx_last_computed_token_d = None

645 block_idx_last_computed_token_p = None

647 - block_idx_last_scheduled_token_d = None

646 block_idx_last_scheduled_token_p = None

647 block_idx_first_scheduled_token_p = None

648 + block_idx_last_scheduled_token_d = None

649 + block_idx_last_computed_token_d = None

650 num_computed_tokens_p = None

651

652 preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split(

653 output[:num_actual_tokens],

654 - [num_decodes, num_prefill_tokens],

654 + [num_decode_tokens, num_prefill_tokens],

655 dim=0,

656 )

657

@@ -709,6 +709,7 @@ class MambaMixer2(MambaBase, PluggableLayer):

709 )

710

711 # NOTE: final output is an in-place update of out tensor

712 + assert preallocated_ssm_out_p is not None

713 varlen_states = mamba_chunk_scan_combined_varlen(

714 hidden_states_p.view(

715 num_prefill_tokens, self.num_heads // self.tp_size, self.head_dim

@@ -840,6 +841,9 @@ class MambaMixer2(MambaBase, PluggableLayer):

841 conv_state_indices=state_indices_tensor_d,

842 block_idx_last_scheduled_token=block_idx_last_scheduled_token_d,

843 initial_state_idx=block_idx_last_computed_token_d,

844 + num_accepted_tokens=num_accepted_tokens,

845 + query_start_loc=query_start_loc_d,

846 + max_query_len=state_indices_tensor_d.size(-1),

847 )

848

849 hidden_states_d, B_d, C_d = self.split_hidden_states_B_C_fn(

@@ -862,6 +866,7 @@ class MambaMixer2(MambaBase, PluggableLayer):

866 -1, self.num_heads // self.tp_size, self.head_dim

867 )

868

869 + assert preallocated_ssm_out_d is not None

870 # - the hidden is reshaped into (bs, num_heads, head_dim)

871 # - mamba_cache_params.ssm_state's slots will be selected

872 # using state_indices_tensor_d

@@ -879,7 +884,9 @@ class MambaMixer2(MambaBase, PluggableLayer):

884 dt_softplus=True,

885 state_batch_indices=state_indices_tensor_d_input,

886 dst_state_batch_indices=state_indices_tensor_d_output,

882 - out=preallocated_ssm_out_d.view(num_decodes, -1, self.head_dim),

887 + out=preallocated_ssm_out_d.view(num_decode_tokens, -1, self.head_dim),

888 + num_accepted_tokens=num_accepted_tokens,

889 + cu_seqlens=query_start_loc_d,

890 is_blackwell=self.is_blackwell,

891 )

892

@@ -901,6 +908,7 @@ class MambaMixer2(MambaBase, PluggableLayer):

908 head_dim=self.head_dim,

909 state_size=self.ssm_state_size,

910 conv_kernel=self.conv_kernel_size,

911 + num_spec=self.num_spec,

912 )

913

914 @property

vllm/model_executor/layers/mamba/mamba_utils.py

+2 -1

@@ -133,6 +133,7 @@ class MambaStateShapeCalculator:

133 head_dim: int,

134 state_size: int,

135 conv_kernel: int,

136 + num_spec: int = 0,

137 ) -> tuple[tuple[int, int], tuple[int, int, int]]:

138 # if n_groups is not divisible by world_size, need to extend the shards

139 # to ensure all groups needed by a head is sharded along with it

@@ -141,7 +142,7 @@ class MambaStateShapeCalculator:

142 conv_dim = intermediate_size + 2 * n_groups * state_size

143

144 # contiguous along 'dim' axis

144 - conv_state_shape = (conv_kernel - 1, divide(conv_dim, tp_world_size))

145 + conv_state_shape = (conv_kernel - 1 + num_spec, divide(conv_dim, tp_world_size))

146

147 # These are not TP-ed as they depend on A, dt_bias, D

148 # - they are typically small

vllm/model_executor/layers/mamba/ops/causal_conv1d.py

+3 -1

@@ -1155,7 +1155,9 @@ def causal_conv1d_update(

1155 if conv_state_indices is None:

1156 assert conv_state.size(0) >= batch

1157 else:

1158 - assert (batch,) == conv_state_indices.shape

1158 + assert batch == conv_state_indices.shape[0], (

1159 + f"ERROR: conv_state_indices should have shape ({batch},*) but got {conv_state_indices.shape}"

1160 + )

1161

1162 assert num_cache_lines >= batch

1163 assert weight.stride(1) == 1 # Need this

vllm/model_executor/layers/mamba/short_conv.py

+2 -8

@@ -119,7 +119,8 @@ class ShortConv(MambaBase, CustomOp):

119 assert isinstance(attn_metadata, ShortConvAttentionMetadata)

120 self_kv_cache = self.kv_cache[forward_context.virtual_engine]

121 conv_state = self_kv_cache[0].transpose(-1, -2)

122 - state_indices_tensor = attn_metadata.state_indices_tensor

122 + state_indices_tensor_p = attn_metadata.state_indices_tensor_p

123 + state_indices_tensor_d = attn_metadata.state_indices_tensor_d

124 has_initial_states_p = attn_metadata.has_initial_states_p

125 query_start_loc_p = attn_metadata.query_start_loc_p

126

@@ -163,13 +164,6 @@ class ShortConv(MambaBase, CustomOp):

164 [num_decodes, num_prefill_tokens],

165 dim=0,

166 )

166 - # Split along batch dimension

167 - state_indices_tensor_d, state_indices_tensor_p = torch.split(

168 - state_indices_tensor,

169 - [num_decodes, num_prefills],

170 - dim=0,

171 - )

172 -

167 conv_output_list = []

168

169 if has_prefill:

vllm/model_executor/models/mamba2.py

@@ -228,6 +228,7 @@ class Mamba2ForCausalLM(

228 head_dim=hf_config.head_dim,

229 state_size=hf_config.state_size,

230 conv_kernel=hf_config.conv_kernel,

231 + num_spec=vllm_config.num_speculative_tokens,

232 )

233

234 @classmethod

vllm/model_executor/models/nemotron_h.py

@@ -636,6 +636,9 @@ class NemotronHModel(nn.Module):

636 hidden_states, _ = self.norm_f(hidden_states, residual)

637 return hidden_states

638

639 + def is_spec_layer(self, config: NemotronHConfig, weight_name: str) -> bool:

640 + return weight_name.startswith("mtp.")

641 +

642 def _get_max_n_routed_experts(self) -> int:

643 """Get max n_routed_experts from config or block_configs for puzzle models.

644

@@ -702,6 +705,10 @@ class NemotronHModel(nn.Module):

705 if name is None:

706 continue

707

708 + # Skip MTP/spec decode layers early (before stacked params mapping)

709 + if name.startswith("mtp."):

710 + continue

711 +

712 # load stacked params

713 for param_name, weight_name, shard_id in stacked_params_mapping:

714 if weight_name not in name:

@@ -845,6 +852,7 @@ class NemotronHForCausalLM(

852 head_dim=hf_config.mamba_head_dim,

853 state_size=hf_config.ssm_state_size,

854 conv_kernel=hf_config.conv_kernel,

855 + num_spec=vllm_config.num_speculative_tokens,

856 )

857

858 @classmethod

vllm/model_executor/models/nemotron_h_mtp.py new

+503

@@ -0,0 +1,503 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""NemotronH-MTP model with attention layers."""

4 +

5 +import typing

6 +from collections.abc import Callable, Iterable

7 +

8 +import torch

9 +import torch.nn as nn

10 +

11 +from vllm.compilation.decorators import support_torch_compile

12 +from vllm.config import CacheConfig, ModelConfig, VllmConfig

13 +from vllm.config.parallel import ParallelConfig

14 +from vllm.model_executor.layers.fused_moe import FusedMoE

15 +from vllm.model_executor.layers.layernorm import RMSNorm

16 +from vllm.model_executor.layers.linear import ColumnParallelLinear

17 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

18 +from vllm.model_executor.layers.quantization import QuantizationConfig

19 +from vllm.model_executor.layers.vocab_parallel_embedding import (

20 + ParallelLMHead,

21 + VocabParallelEmbedding,

22 +)

23 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

24 +from vllm.model_executor.models.utils import (

25 + make_empty_intermediate_tensors_factory,

26 + maybe_prefix,

27 +)

28 +from vllm.sequence import IntermediateTensors

29 +from vllm.transformers_utils.configs import NemotronHConfig

30 +

31 +from .interfaces import SupportsPP

32 +from .nemotron_h import (

33 + NemotronHAttentionDecoderLayer,

34 + NemotronHMoEDecoderLayer,

35 +)

36 +

37 +

38 +class NemotronHMTPAttentionDecoderLayer(NemotronHAttentionDecoderLayer):

39 + def __init__(

40 + self,

41 + config: NemotronHConfig,

42 + layer_idx: int,

43 + model_config: ModelConfig | None = None,

44 + cache_config: CacheConfig | None = None,

45 + quant_config: QuantizationConfig | None = None,

46 + parallel_config: ParallelConfig | None = None,

47 + prefix: str = "",

48 + has_start_projections: bool = False,

49 + has_end_norm: bool = False,

50 + ) -> None:

51 + super().__init__(

52 + config=config,

53 + layer_idx=layer_idx,

54 + model_config=model_config,

55 + cache_config=cache_config,

56 + quant_config=quant_config,

57 + parallel_config=parallel_config,

58 + prefix=prefix,

59 + )

60 + self.has_start_projections = has_start_projections

61 + self.has_end_norm = has_end_norm

62 +

63 + if has_start_projections:

64 + self.enorm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)

65 + self.hnorm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)

66 +

67 + # Fusion layer to combine embeddings with target hidden states

68 + self.eh_proj = ColumnParallelLinear(

69 + input_size=config.hidden_size * 2,

70 + output_size=config.hidden_size,

71 + bias=False,

72 + gather_output=True,

73 + params_dtype=config.dtype

74 + if hasattr(config, "dtype")

75 + else torch.bfloat16,

76 + quant_config=quant_config,

77 + prefix=f"{prefix}.eh_proj",

78 + )

79 +

80 + if has_end_norm:

81 + self.final_layernorm = RMSNorm(

82 + config.hidden_size,

83 + eps=getattr(config, "layer_norm_epsilon", 1e-5),

84 + )

85 +

86 + def forward(

87 + self,

88 + inputs_embeds: torch.Tensor,

89 + positions: torch.Tensor,

90 + hidden_states: torch.Tensor,

91 + residual: torch.Tensor | None = None,

92 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

93 + # Start projections (Fusion)

94 + if self.has_start_projections:

95 + # Normalize both inputs before fusion

96 + assert inputs_embeds is not None

97 + inputs_embeds_normed = self.enorm(inputs_embeds)

98 + previous_hidden_states_normed = self.hnorm(hidden_states)

99 +

100 + # Fuse via concatenation and linear projection

101 + fused = torch.cat(

102 + [inputs_embeds_normed, previous_hidden_states_normed], dim=-1

103 + )

104 + hidden_states, _ = self.eh_proj(fused)

105 +

106 + # Call parent forward (Attention)

107 + # Parent forward expects: hidden_states, residual

108 + hidden_states, residual = super().forward(

109 + positions=positions,

110 + hidden_states=hidden_states,

111 + residual=residual,

112 + )

113 +

114 + # End norm

115 + if self.has_end_norm:

116 + if residual is not None:

117 + hidden_states = hidden_states + residual

118 + residual = None # Consumed residual

119 +

120 + hidden_states = self.final_layernorm(hidden_states)

121 +

122 + return hidden_states, residual

123 +

124 +

125 +class NemotronHMTPMoEDecoderLayer(NemotronHMoEDecoderLayer):

126 + def __init__(

127 + self,

128 + config: NemotronHConfig,

129 + layer_idx: int,

130 + model_config: ModelConfig | None = None,

131 + cache_config: CacheConfig | None = None,

132 + quant_config: QuantizationConfig | None = None,

133 + parallel_config: ParallelConfig | None = None,

134 + prefix: str = "",

135 + has_start_projections: bool = False,

136 + has_end_norm: bool = False,

137 + ) -> None:

138 + super().__init__(

139 + config=config,

140 + layer_idx=layer_idx,

141 + model_config=model_config,

142 + cache_config=cache_config,

143 + quant_config=quant_config,

144 + parallel_config=parallel_config,

145 + prefix=prefix,

146 + )

147 + self.has_start_projections = has_start_projections

148 + self.has_end_norm = has_end_norm

149 +

150 + if has_start_projections:

151 + self.enorm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)

152 + self.hnorm = RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)

153 +

154 + # Fusion layer to combine embeddings with target hidden states

155 + self.eh_proj = ColumnParallelLinear(

156 + input_size=config.hidden_size * 2,

157 + output_size=config.hidden_size,

158 + bias=False,

159 + gather_output=True,

160 + params_dtype=config.dtype

161 + if hasattr(config, "dtype")

162 + else torch.bfloat16,

163 + quant_config=quant_config,

164 + prefix=f"{prefix}.eh_proj",

165 + )

166 +

167 + if has_end_norm:

168 + self.final_layernorm = RMSNorm(

169 + config.hidden_size,

170 + eps=getattr(config, "layer_norm_epsilon", 1e-5),

171 + )

172 +

173 + def forward(

174 + self,

175 + inputs_embeds: torch.Tensor,

176 + positions: torch.Tensor,

177 + hidden_states: torch.Tensor,

178 + residual: torch.Tensor | None = None,

179 + ) -> tuple[torch.Tensor, torch.Tensor | None]:

180 + # Start projections (Fusion)

181 + if self.has_start_projections:

182 + # Normalize both inputs before fusion

183 + assert inputs_embeds is not None

184 + inputs_embeds_normed = self.enorm(inputs_embeds)

185 + previous_hidden_states_normed = self.hnorm(hidden_states)

186 +

187 + # Fuse via concatenation and linear projection

188 + fused = torch.cat(

189 + [inputs_embeds_normed, previous_hidden_states_normed], dim=-1

190 + )

191 + hidden_states, _ = self.eh_proj(fused)

192 +

193 + # Call parent forward (MoE)

194 + hidden_states, residual = super().forward(

195 + hidden_states=hidden_states,

196 + residual=residual,

197 + )

198 +

199 + # End norm

200 + if self.has_end_norm:

201 + if residual is not None:

202 + hidden_states = hidden_states + residual

203 + residual = None # Consumed residual

204 +

205 + hidden_states = self.final_layernorm(hidden_states)

206 +

207 + return hidden_states, residual

208 +

209 +

210 +@support_torch_compile

211 +class NemotronHMultiTokenPredictor(nn.Module):

212 + """MTP predictor with NemotronH layers."""

213 +

214 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

215 + super().__init__()

216 +

217 + config = vllm_config.model_config.hf_config

218 +

219 + self.config = config

220 + self.vocab_size = config.vocab_size

221 + self.org_vocab_size = config.vocab_size

222 +

223 + self.mtp_start_layer_idx = config.num_hidden_layers

224 + self.num_mtp_layers = getattr(config, "num_nextn_predict_layers", 1)

225 + assert self.num_mtp_layers == 1, (

226 + "Only one MTP layer is supported for NemotronH-MTP"

227 + )

228 +

229 + self.pattern_str = config.mtp_hybrid_override_pattern

230 + self.pattern_len = len(self.pattern_str)

231 + assert self.pattern_len > 0

232 +

233 + self.embed_tokens = VocabParallelEmbedding(

234 + self.vocab_size,

235 + config.hidden_size,

236 + )

237 +

238 + # Build flat list of layers

239 + self.layers = torch.nn.ModuleDict()

240 +

241 + # Total number of physical layers = num_steps * pattern_len

242 + total_layers = self.num_mtp_layers * self.pattern_len

243 + for i in range(total_layers):

244 + step_rel_idx = i % self.pattern_len

245 +

246 + char = self.pattern_str[step_rel_idx]

247 +

248 + is_start_of_step = step_rel_idx == 0

249 + is_end_of_step = step_rel_idx == self.pattern_len - 1

250 +

251 + layer_prefix = f"{prefix}.layers.{i}"

252 +

253 + # TODO smor- remove double layers formation

254 + common_kwargs = dict(

255 + config=config,

256 + layer_idx=self.mtp_start_layer_idx + i,

257 + model_config=vllm_config.model_config,

258 + cache_config=vllm_config.cache_config,

259 + quant_config=vllm_config.quant_config,

260 + parallel_config=vllm_config.parallel_config,

261 + prefix=layer_prefix,

262 + has_start_projections=is_start_of_step,

263 + has_end_norm=is_end_of_step,

264 + )

265 +

266 + if char == "*":

267 + self.layers[str(i)] = NemotronHMTPAttentionDecoderLayer(**common_kwargs)

268 + elif char == "E":

269 + self.layers[str(i)] = NemotronHMTPMoEDecoderLayer(**common_kwargs)

270 + else:

271 + raise NotImplementedError(

272 + f"Pattern char '{char}' in {self.pattern_str} not implemented"

273 + )

274 +

275 + self.make_empty_intermediate_tensors: Callable[..., IntermediateTensors] = (

276 + make_empty_intermediate_tensors_factory(

277 + ["hidden_states", "residual"], config.hidden_size

278 + )

279 + )

280 +

281 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

282 + assert self.embed_tokens is not None, (

283 + "embed_tokens not initialized - must be shared from target model"

284 + )

285 + return self.embed_tokens(input_ids)

286 +

287 + def forward(

288 + self,

289 + input_ids: torch.Tensor,

290 + positions: torch.Tensor,

291 + hidden_states: torch.Tensor,

292 + intermediate_tensors: IntermediateTensors | None = None,

293 + inputs_embeds: torch.Tensor | None = None,

294 + ) -> torch.Tensor | IntermediateTensors:

295 + if inputs_embeds is None:

296 + inputs_embeds = self.get_input_embeddings(input_ids)

297 +

298 + residual = None

299 +

300 + for i in range(self.pattern_len):

301 + hidden_states, residual = self.layers[str(i)](

302 + inputs_embeds=inputs_embeds,

303 + positions=positions,

304 + hidden_states=hidden_states,

305 + residual=residual,

306 + )

307 + return hidden_states

308 +

309 +

310 +class NemotronHMTP(nn.Module, SupportsPP):

311 + """NemotronH MTP model."""

312 +

313 + packed_modules_mapping = {

314 + "qkv_proj": [

315 + "q_proj",

316 + "k_proj",

317 + "v_proj",

318 + ],

319 + }

320 +

321 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

322 + super().__init__()

323 + config = vllm_config.model_config.hf_config

324 + self.vllm_config = vllm_config

325 + self.config = config

326 + self.quant_config = vllm_config.quant_config

327 +

328 + # Needed for load_weights mapping

329 + self.mtp_start_layer_idx = config.num_hidden_layers

330 +

331 + # EPLB config for experts

332 + self.num_redundant_experts = 0

333 + if vllm_config.parallel_config and vllm_config.parallel_config.eplb_config:

334 + self.num_redundant_experts = (

335 + vllm_config.parallel_config.eplb_config.num_redundant_experts

336 + )

337 +

338 + # MTP predictor

339 + self.model = NemotronHMultiTokenPredictor(

340 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "mtp")

341 + )

342 +

343 + # LM head for generating logits

344 + self.lm_head = ParallelLMHead(

345 + self.config.vocab_size,

346 + self.config.hidden_size,

347 + prefix=maybe_prefix(prefix, "lm_head"),

348 + )

349 +

350 + self.logits_processor = LogitsProcessor(self.config.vocab_size)

351 +

352 + self.make_empty_intermediate_tensors = (

353 + self.model.make_empty_intermediate_tensors

354 + )

355 +

356 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

357 + return self.model.get_input_embeddings(input_ids)

358 +

359 + def forward(

360 + self,

361 + input_ids: torch.Tensor,

362 + positions: torch.Tensor,

363 + hidden_states: torch.Tensor,

364 + intermediate_tensors: IntermediateTensors | None = None,

365 + inputs_embeds: torch.Tensor | None = None,

366 + **kwargs: object,

367 + ) -> torch.Tensor:

368 + """Forward - applies attention-based MTP."""

369 + hidden_states = self.model(

370 + input_ids,

371 + positions,

372 + hidden_states,

373 + intermediate_tensors,

374 + inputs_embeds,

375 + )

376 + return hidden_states

377 +

378 + def compute_logits(

379 + self,

380 + hidden_states: torch.Tensor,

381 + ) -> torch.Tensor | None:

382 + """Compute logits for DRAFT token generation."""

383 + assert self.lm_head is not None, (

384 + "lm_head not initialized - must be shared from target model"

385 + )

386 + return self.logits_processor(self.lm_head, hidden_states)

387 +

388 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

389 + """Load MTP weights with proper name remapping."""

390 + stacked_params_mapping = [

391 + # (param_name, shard_name, shard_id)

392 + ("qkv_proj", "q_proj", "q"),

393 + ("qkv_proj", "k_proj", "k"),

394 + ("qkv_proj", "v_proj", "v"),

395 + ]

396 +

397 + expert_params_mapping = []

398 + if hasattr(self.config, "n_routed_experts") and self.config.n_routed_experts:

399 + expert_params_mapping = FusedMoE.make_expert_params_mapping(

400 + self,

401 + ckpt_gate_proj_name="up_proj",

402 + ckpt_down_proj_name="down_proj",

403 + ckpt_up_proj_name="", # Empty - non-gated MoE

404 + num_experts=self.config.n_routed_experts,

405 + num_redundant_experts=self.num_redundant_experts,

406 + )

407 +

408 + params_dict = dict(self.named_parameters())

409 + loaded_params: set[str] = set()

410 +

411 + for name, loaded_weight in weights:

412 + # Only process MTP weights - skip all non-MTP weights

413 + if (

414 + not name.startswith("mtp.")

415 + and "embeddings" not in name

416 + and "lm_head" not in name

417 + ):

418 + continue

419 + # Skip rotary embeddings (computed, not loaded)

420 + if "rotary_emb.inv_freq" in name:

421 + continue

422 +

423 + name = name.replace("mtp.layers.", "model.layers.")

424 +

425 + if "embeddings" in name:

426 + name = name.replace("embeddings", "embed_tokens")

427 + if name.startswith("backbone."):

428 + name = name.replace("backbone.", "model.")

429 +

430 + # Handle stacked parameters (qkv_proj) for attention layers

431 + is_stacked = False

432 + for param_name, weight_name, shard_id in stacked_params_mapping:

433 + if weight_name not in name:

434 + continue

435 + # Must be in a mixer (attention layer)

436 + if ".mixer." not in name:

437 + continue

438 +

439 + is_stacked = True

440 + stacked_name = name.replace(weight_name, param_name)

441 +

442 + if stacked_name.endswith(".bias") and stacked_name not in params_dict:

443 + continue

444 +

445 + if stacked_name not in params_dict:

446 + # Might be that mapping failed or param doesn't exist

447 + continue

448 +

449 + param = params_dict[stacked_name]

450 + weight_loader = getattr(param, "weight_loader", None)

451 + if weight_loader is not None:

452 + weight_loader(param, loaded_weight, shard_id)

453 + loaded_params.add(stacked_name)

454 + break

455 +

456 + if is_stacked:

457 + continue

458 +

459 + is_expert_weight = False

460 + for mapping in expert_params_mapping:

461 + param_name, weight_name, expert_id, shard_id = mapping

462 + # weight_name is like "experts.0.up_proj."

463 + if weight_name not in name:

464 + continue

465 +

466 + is_expert_weight = True

467 +

468 + # Replace the expert-specific weight name with fused parameter name

469 + # e.g., "experts.0.up_proj." -> "experts.w13_"

470 + name_mapped = name.replace(weight_name, param_name)

471 +

472 + if name_mapped not in params_dict:

473 + continue

474 +

475 + param = params_dict[name_mapped]

476 + weight_loader = typing.cast(Callable[..., bool], param.weight_loader)

477 + success = weight_loader(

478 + param,

479 + loaded_weight,

480 + name_mapped,

481 + shard_id=shard_id,

482 + expert_id=expert_id,

483 + return_success=True,

484 + )

485 + if success:

486 + loaded_params.add(name_mapped)

487 + break

488 +

489 + if is_expert_weight:

490 + continue

491 +

492 + if name.endswith(".bias") and name not in params_dict:

493 + continue

494 +

495 + if name not in params_dict:

496 + continue

497 +

498 + param = params_dict[name]

499 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

500 + weight_loader(param, loaded_weight)

501 + loaded_params.add(name)

502 +

503 + return loaded_params

vllm/model_executor/models/plamo2.py

+3 -9

@@ -266,7 +266,8 @@ class Plamo2MambaMixer(MambaBase, PluggableLayer):

266 # conv_state = (..., dim, width-1) yet contiguous along 'dim'

267 conv_state = self_kv_cache[0].transpose(-1, -2)

268 ssm_state = self_kv_cache[1]

269 - state_indices_tensor = attn_metadata.state_indices_tensor

269 + state_indices_tensor_p = attn_metadata.state_indices_tensor_p

270 + state_indices_tensor_d = attn_metadata.state_indices_tensor_d

271 has_initial_states_p = attn_metadata.has_initial_states_p

272 prep_initial_states = attn_metadata.prep_initial_states

273 chunk_size = attn_metadata.chunk_size

@@ -309,13 +310,6 @@ class Plamo2MambaMixer(MambaBase, PluggableLayer):

310 gate_d, gate_p = torch.split(

311 gate[:num_actual_tokens], [num_decodes, num_prefill_tokens], dim=0

312 )

312 - # Split along batch dimension

313 - state_indices_tensor_d, state_indices_tensor_p = torch.split(

314 - state_indices_tensor,

315 - [num_decodes, num_prefills],

316 - dim=0,

317 - )

318 -

313 # Preallocate output tensor to avoid memcpy cost for merging prefill

314 # and decode outputs

315 preallocated_ssm_out = torch.empty(

@@ -336,7 +330,7 @@ class Plamo2MambaMixer(MambaBase, PluggableLayer):

330 if has_prefill:

331 # 2. Convolution sequence transformation

332 # - "cache_indices" updates the conv_state cache in positions

339 - # pointed to by "state_indices_tensor"

333 + # pointed to by "state_indices_tensor_p"

334 x = hidden_states_p.transpose(0, 1) # this is the form that causal-conv see

335 hidden_states_p = causal_conv1d_fn(

336 x,

vllm/model_executor/models/registry.py

@@ -522,6 +522,7 @@ _SPECULATIVE_DECODING_MODELS = {

522 "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),

523 "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),

524 "ExaoneMoeMTP": ("exaone_moe_mtp", "ExaoneMoeMTP"),

525 + "NemotronHMTPModel": ("nemotron_h_mtp", "NemotronHMTP"),

526 "LongCatFlashMTPModel": ("longcat_flash_mtp", "LongCatFlashMTP"),

527 "Glm4MoeMTPModel": ("glm4_moe_mtp", "Glm4MoeMTP"),

528 "Glm4MoeLiteMTPModel": ("glm4_moe_lite_mtp", "Glm4MoeLiteMTP"),

vllm/transformers_utils/configs/nemotron_h.py

+6 -3

@@ -51,6 +51,8 @@ class NemotronHConfig(PretrainedConfig):

51 The pattern of the hybrid model. The pattern is a string of

52 characters where each character represents

53 M: Mamba2, *: Attention, -: MLP

54 + mtp_hybrid_override_pattern (`str`, *optional*, defaults to `"*E"`):

55 + The pattern of the MTP layers.

56 num_attention_heads (`int`, *optional*, defaults to 32):

57 Number of attention heads for each attention layer in the

58 Transformer encoder.

@@ -150,6 +152,7 @@ class NemotronHConfig(PretrainedConfig):

152 intermediate_size=21504,

153 num_hidden_layers=52,

154 hybrid_override_pattern="M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-",

155 + mtp_hybrid_override_pattern="*E",

156 num_attention_heads=32,

157 head_dim=128,

158 num_key_value_heads=8, # nemo: num_query_groups

@@ -203,6 +206,7 @@ class NemotronHConfig(PretrainedConfig):

206 self.intermediate_size = intermediate_size

207 self.num_hidden_layers = num_hidden_layers

208 self.hybrid_override_pattern = hybrid_override_pattern

209 + self.mtp_hybrid_override_pattern = mtp_hybrid_override_pattern

210 self.num_attention_heads = num_attention_heads

211 self.head_dim = head_dim

212 self.sliding_window = sliding_window

@@ -215,10 +219,9 @@ class NemotronHConfig(PretrainedConfig):

219 assert len(self.hybrid_override_pattern) == self.num_hidden_layers, (

220 "hybrid_override_pattern must have same length as num_hidden_layers"

221 )

218 - assert re.match(r"^[*-M]+$", self.hybrid_override_pattern), (

219 - "hybrid_override_pattern must only contain characters 'M', '*', or '-'"

222 + assert re.match(r"^[*-ME]+$", self.hybrid_override_pattern), (

223 + "hybrid_override_pattern must only contain characters 'M', '*', '-', or 'E'"

224 )

221 -

225 # for backward compatibility

226 if num_key_value_heads is None:

227 num_key_value_heads = num_attention_heads

vllm/v1/attention/backends/mamba2_attn.py

+5 -1

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import itertools

4 from dataclasses import dataclass, replace

5 +from typing import Any

7 import torch

@@ -200,8 +201,11 @@ class Mamba2AttentionMetadataBuilder(

201 common_prefix_len: int,

202 common_attn_metadata: CommonAttentionMetadata,

203 fast_build: bool = False,

204 + **kwargs: Any,

205 ) -> Mamba2AttentionMetadata:

204 - common = self._compute_common_metadata(common_attn_metadata)

206 + common = self._compute_common_metadata(

207 + common_attn_metadata, num_accepted_tokens=kwargs.get("num_accepted_tokens")

208 + )

209

210 seq_idx_p = None

211 cu_chunk_seqlen_p = None

vllm/v1/attention/backends/mamba_attn.py

+193 -85

@@ -2,9 +2,8 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 import abc

5 -import copy

6 -from dataclasses import dataclass

7 -from typing import ClassVar, TypeVar

5 +from dataclasses import dataclass, replace

6 +from typing import Any, ClassVar, TypeVar

8 import torch

@@ -35,12 +34,21 @@ class BaseMambaAttentionMetadata:

34 num_reqs: int

36 # The following tensors only contain prefill requests and will be None if

38 - # the batch has no prefill request.

37 + # the batch has no prefill requests.

38 has_initial_states_p: torch.Tensor | None

39 query_start_loc_p: torch.Tensor | None

40 num_computed_tokens_p: torch.Tensor | None

41 + state_indices_tensor_p: torch.Tensor | None

43 - state_indices_tensor: torch.Tensor

43 + # The following tensors are used for decode requests and

44 + # speculative decoding compatibility, and will be None if the batch

45 + # has no decode requests.

46 + state_indices_tensor_d: torch.Tensor | None

47 + query_start_loc_d: torch.Tensor | None # shape: [num_decodes + 1,]

48 +

49 + # Number of accepted tokens for each spec sequence (for loading correct checkpoint)

50 + # Includes the bonus token (so minimum is 1)

51 + num_accepted_tokens: torch.Tensor | None # shape: [batch,]

53 # The following tensors are only used for prefix caching in all mode and

54 # are None if disabled

@@ -60,9 +68,9 @@ class BaseMambaAttentionMetadata:

68 class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

69 metadata_cls: type[M]

70 reorder_batch_threshold: int = 1

63 - _cudagraph_support: ClassVar[AttentionCGSupport] = (

64 - AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE

65 - )

71 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

72 +

73 + # Will be disabled if speculative decoding is used

74 supports_update_block_table: bool = True

76 def __init__(

@@ -74,6 +82,12 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

82 ):

83 super().__init__(kv_cache_spec, layer_names, vllm_config, device)

85 + # Enable speculative decoding support

86 + self.speculative_config = vllm_config.speculative_config

87 + self.compilation_config = vllm_config.compilation_config

88 + self.num_spec_tokens: int = vllm_config.num_speculative_tokens

89 + self.use_spec_decode = self.num_spec_tokens > 0

90 +

91 assert isinstance(kv_cache_spec, MambaSpec)

92 self.compilation_config = vllm_config.compilation_config

93 self.decode_cudagraph_max_bs = self.vllm_config.scheduler_config.max_num_seqs

@@ -84,13 +98,17 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

98 )

100 if self.vllm_config.cache_config.mamba_cache_mode == "all":

87 - self.state_indices_tensor = torch.empty(

101 + max_num_blocks = cdiv(

102 + self.vllm_config.model_config.max_model_len,

103 + self.kv_cache_spec.block_size,

104 + )

105 + # Speculative decoding not supported with prefix caching,

106 + # so keep shape consistent with prefill buffer

107 + # TODO: reduce this size as needed for decode-only cudagraph capture

108 + self.state_indices_tensor_d = torch.empty(

109 (

110 self.decode_cudagraph_max_bs,

90 - cdiv(

91 - self.vllm_config.model_config.max_model_len,

92 - self.kv_cache_spec.block_size,

93 - ),

111 + max_num_blocks,

112 ),

113 dtype=torch.int32,

114 device=device,

@@ -106,12 +124,25 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

124 device=device,

125 )

126 else:

109 - self.state_indices_tensor = torch.empty(

127 + self.state_indices_tensor_d = torch.empty(

128 + (self.decode_cudagraph_max_bs, 1 + self.num_spec_tokens),

129 + dtype=torch.int32,

130 + device=device,

131 + )

132 +

133 + # For speculative decoding, we need to store the following buffers

134 + # for CUDA graph capture during decode

135 + if self.num_spec_tokens > 0:

136 + self.decode_num_accepted_tokens = torch.empty(

137 (self.decode_cudagraph_max_bs,),

138 dtype=torch.int32,

139 device=device,

140 )

141

142 + self._init_reorder_batch_threshold(1, self.use_spec_decode)

143 + if self.use_spec_decode:

144 + self.supports_update_block_table = False

145 +

146 def build_for_cudagraph_capture(

147 self, common_attn_metadata: CommonAttentionMetadata

148 ) -> M:

@@ -121,26 +152,38 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

152 """

153 m = common_attn_metadata

154

124 - assert m.num_reqs == m.num_actual_tokens, (

155 + assert (

156 + m.max_query_len <= 1 + self.num_spec_tokens

157 + and m.num_reqs <= self.decode_cudagraph_max_bs

158 + ), (

159 "Mamba only supports decode-only full CUDAGraph capture. "

160 "Make sure all cudagraph capture sizes <= max_num_seq."

161 )

162

129 - m.max_query_len = 1 # decode-only

163 + assert m.max_query_len == 1 + self.num_spec_tokens # decode-only

164

131 - return self.build(0, m)

165 + num_accepted_tokens = None

166 + if self.num_spec_tokens > 0:

167 + num_accepted_tokens = torch.diff(m.query_start_loc)

168 +

169 + return self.build(0, m, num_accepted_tokens=num_accepted_tokens)

170

171 def build(

172 self,

173 common_prefix_len: int,

174 common_attn_metadata: CommonAttentionMetadata,

175 fast_build: bool = False,

176 + *,

177 + num_accepted_tokens: torch.Tensor | None = None,

178 + **kwargs: Any,

179 ) -> M:

180 """

181 Default build implementation for Mamba-like attention backends.

182 Subclasses (e.g., Mamba2) can override to add additional metadata.

183 """

143 - return self._compute_common_metadata(common_attn_metadata)

184 + return self._compute_common_metadata(

185 + common_attn_metadata, num_accepted_tokens=num_accepted_tokens

186 + )

187

188 def _compute_prefix_caching_block_indices(

189 self,

@@ -176,21 +219,32 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

219 def _compute_common_metadata(

220 self,

221 common_attn_metadata: CommonAttentionMetadata,

222 + *,

223 + num_accepted_tokens: torch.Tensor | None = None,

224 ) -> M:

225 """

226 Compute metadata common to both Mamba1 and Mamba2.

227 """

228 num_reqs = common_attn_metadata.num_reqs

229

230 + # Treat multi-token queries as decode requests when

231 + # speculative decoding is enabled. Otherwise, use the

232 + # default decode threshold to prevent misclassification

233 + # of prefill queries as decode requests.

234 + decode_threshold = (

235 + self.reorder_batch_threshold if num_accepted_tokens is not None else 1

236 + )

237 +

238 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

239 split_decodes_and_prefills(

187 - common_attn_metadata, decode_threshold=self.reorder_batch_threshold

240 + common_attn_metadata, decode_threshold=decode_threshold

241 )

242 )

243

244 # Need flags to indicate if there are initial states

245 has_initial_states_p = None

246 query_start_loc_p = None

247 + query_start_loc_d = None

248 num_computed_tokens = None

249 num_computed_tokens_p = None

250

@@ -218,13 +272,31 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

272 common_attn_metadata, mamba_block_size

273 )

274 else:

221 - # Always return just a single block per each request:

275 state_indices_tensor = mamba_get_block_table_tensor(

276 common_attn_metadata.block_table_tensor,

277 common_attn_metadata.seq_lens,

278 self.kv_cache_spec,

279 self.vllm_config.cache_config.mamba_cache_mode,

227 - )[:, 0]

280 + )

281 +

282 + if state_indices_tensor.dim() == 1:

283 + state_indices_tensor = state_indices_tensor.unsqueeze(-1)

284 +

285 + state_indices_tensor_d, state_indices_tensor_p = torch.split(

286 + state_indices_tensor,

287 + [num_decodes, num_prefills],

288 + dim=0,

289 + )

290 + if self.vllm_config.cache_config.mamba_cache_mode != "all":

291 + state_indices_tensor_d = state_indices_tensor_d[

292 + :, : 1 + self.num_spec_tokens

293 + ]

294 + state_indices_tensor_p = state_indices_tensor_p[:, 0]

295 +

296 + if num_decodes > 0 and self.use_spec_decode:

297 + assert num_accepted_tokens is not None

298 + query_start_loc_d = common_attn_metadata.query_start_loc[: num_decodes + 1]

299 + num_accepted_tokens = num_accepted_tokens[:num_decodes]

300

301 if num_prefills > 0:

302 if num_computed_tokens is None:

@@ -258,39 +330,18 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

330 block_idx_first_scheduled_token_p = block_idx_first_scheduled_token[

331 num_reqs - num_prefills : num_reqs

332 ]

261 - elif (

262 - num_decodes <= self.decode_cudagraph_max_bs

263 - and self.compilation_config.cudagraph_mode.has_full_cudagraphs()

264 - ):

265 - self.state_indices_tensor[:num_decodes].copy_(

266 - state_indices_tensor, non_blocking=True

267 - )

268 - state_indices_tensor = self.state_indices_tensor[:num_decode_tokens]

269 - state_indices_tensor[num_decodes:] = PAD_SLOT_ID

270 -

271 - if self.vllm_config.cache_config.mamba_cache_mode == "all":

272 - self.block_idx_last_scheduled_token[:num_decodes].copy_(

273 - block_idx_last_scheduled_token, non_blocking=True

274 - )

275 - block_idx_last_scheduled_token = self.block_idx_last_scheduled_token[

276 - :num_decode_tokens

277 - ]

333

279 - self.block_idx_last_computed_token[:num_decodes].copy_(

280 - block_idx_last_computed_token, non_blocking=True

281 - )

282 - block_idx_last_computed_token = self.block_idx_last_computed_token[

283 - :num_decode_tokens

284 - ]

285 -

286 - return self.metadata_cls(

334 + metadata = self.metadata_cls(

335 num_prefills=num_prefills,

336 num_prefill_tokens=num_prefill_tokens,

337 num_decodes=num_decodes,

338 num_decode_tokens=num_decode_tokens,

339 query_start_loc_p=query_start_loc_p,

340 has_initial_states_p=has_initial_states_p,

293 - state_indices_tensor=state_indices_tensor,

341 + state_indices_tensor_p=state_indices_tensor_p,

342 + state_indices_tensor_d=state_indices_tensor_d,

343 + num_accepted_tokens=num_accepted_tokens,

344 + query_start_loc_d=query_start_loc_d,

345 block_idx_last_scheduled_token=block_idx_last_scheduled_token,

346 block_idx_first_scheduled_token_p=block_idx_first_scheduled_token_p,

347 block_idx_last_computed_token=block_idx_last_computed_token,

@@ -302,55 +353,112 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

353 token_chunk_offset_ptr=token_chunk_offset_ptr,

354 )

355

305 - def update_block_table(

356 + return self._update_metadata_for_cudagraph_capture(metadata)

357 +

358 + def _update_metadata_for_cudagraph_capture(

359 self,

360 metadata: M,

308 - blk_table: torch.Tensor,

309 - slot_mapping: torch.Tensor,

361 ) -> M:

311 - new_metadata = copy.copy(metadata)

312 - state_indices_t = mamba_get_block_table_tensor(

313 - blk_table,

314 - metadata.seq_lens,

315 - self.kv_cache_spec,

316 - self.vllm_config.cache_config.mamba_cache_mode,

317 - )

318 - if self.vllm_config.cache_config.mamba_cache_mode in ("none", "align"):

319 - # Only needs the block that saves the running state

320 - state_indices_t = state_indices_t[:, 0]

321 -

322 - num_reqs = blk_table.shape[0]

323 -

324 - # For CUDA graphs, copy to persistent buffer

362 + """

363 + Update the metadata for cudagraph capture.

364 + Currently, only decode is supported for full cudagraphs with Mamba.

365 + """

366 + state_indices_tensor_d = metadata.state_indices_tensor_d

367 + query_start_loc_d = metadata.query_start_loc_d

368 + num_accepted_tokens = metadata.num_accepted_tokens

369 + block_idx_last_scheduled_token = metadata.block_idx_last_scheduled_token

370 + block_idx_last_computed_token = metadata.block_idx_last_computed_token

371 if (

372 metadata.num_prefills == 0

327 - and num_reqs <= self.decode_cudagraph_max_bs

373 + and metadata.num_decodes <= self.decode_cudagraph_max_bs

374 and self.compilation_config.cudagraph_mode.has_full_cudagraphs()

375 ):

330 - persistent_state_indices_t = self.state_indices_tensor[:num_reqs]

331 - persistent_state_indices_t.copy_(state_indices_t, non_blocking=True)

332 - state_indices_t = persistent_state_indices_t

376 + padded_bs = metadata.num_reqs

377 + self.state_indices_tensor_d[: metadata.num_decodes].copy_(

378 + state_indices_tensor_d, non_blocking=True

379 + )

380 + state_indices_tensor_d = self.state_indices_tensor_d[:padded_bs]

381 + state_indices_tensor_d[metadata.num_decodes :] = PAD_SLOT_ID

382 +

383 + if self.use_spec_decode:

384 + assert query_start_loc_d is not None

385 + assert num_accepted_tokens is not None

386 + query_start_loc_d = query_start_loc_d[: padded_bs + 1]

387 + self.decode_num_accepted_tokens[: metadata.num_decodes].copy_(

388 + num_accepted_tokens, non_blocking=True

389 + )

390 + num_accepted_tokens = self.decode_num_accepted_tokens[:padded_bs]

391 + num_accepted_tokens[metadata.num_decodes :] = (

392 + 1 # pad with 1st slot index

393 + )

394

334 - # For 'all' mode, also update prefix caching block indices

335 - # to use this builder's persistent buffers (required for CUDA

336 - # graph replay to read from the correct memory addresses).

395 if self.vllm_config.cache_config.mamba_cache_mode == "all":

338 - assert metadata.block_idx_last_scheduled_token is not None

339 - assert metadata.block_idx_last_computed_token is not None

340 - self.block_idx_last_scheduled_token[:num_reqs].copy_(

341 - metadata.block_idx_last_scheduled_token[:num_reqs],

396 + assert block_idx_last_scheduled_token is not None

397 + assert block_idx_last_computed_token is not None

398 + self.block_idx_last_scheduled_token[: metadata.num_decodes].copy_(

399 + block_idx_last_scheduled_token[: metadata.num_decodes],

400 non_blocking=True,

401 )

344 - new_metadata.block_idx_last_scheduled_token = (

345 - self.block_idx_last_scheduled_token[: metadata.num_decode_tokens]

346 - )

347 - self.block_idx_last_computed_token[:num_reqs].copy_(

348 - metadata.block_idx_last_computed_token[:num_reqs],

402 + block_idx_last_scheduled_token = self.block_idx_last_scheduled_token[

403 + : metadata.num_decode_tokens

404 + ]

405 +

406 + self.block_idx_last_computed_token[: metadata.num_decodes].copy_(

407 + block_idx_last_computed_token[: metadata.num_decodes],

408 non_blocking=True,

409 )

351 - new_metadata.block_idx_last_computed_token = (

352 - self.block_idx_last_computed_token[: metadata.num_decode_tokens]

353 - )

410 + block_idx_last_computed_token = self.block_idx_last_computed_token[

411 + : metadata.num_decode_tokens

412 + ]

413 +

414 + return replace(

415 + metadata,

416 + state_indices_tensor_d=state_indices_tensor_d,

417 + query_start_loc_d=query_start_loc_d,

418 + num_accepted_tokens=num_accepted_tokens,

419 + block_idx_last_scheduled_token=block_idx_last_scheduled_token,

420 + block_idx_last_computed_token=block_idx_last_computed_token,

421 + )

422 +

423 + def update_block_table(

424 + self,

425 + metadata: M,

426 + blk_table: torch.Tensor,

427 + slot_mapping: torch.Tensor,

428 + ) -> M:

429 + state_indices_tensor = mamba_get_block_table_tensor(

430 + blk_table,

431 + metadata.seq_lens,

432 + self.kv_cache_spec,

433 + self.vllm_config.cache_config.mamba_cache_mode,

434 + )

435 + if state_indices_tensor.dim() == 1:

436 + state_indices_tensor = state_indices_tensor.unsqueeze(-1)

437 +

438 + assert (

439 + metadata.num_prefills + metadata.num_decodes

440 + == state_indices_tensor.shape[0]

441 + ), (

442 + "Mismatch in number of requests when updating block table."

443 + f" Expected {metadata.num_prefills + metadata.num_decodes}, "

444 + f"got {state_indices_tensor.shape[0]}."

445 + )

446 +

447 + state_indices_tensor_d, state_indices_tensor_p = torch.split(

448 + state_indices_tensor,

449 + [metadata.num_decodes, metadata.num_prefills],

450 + dim=0,

451 + )

452 + if self.vllm_config.cache_config.mamba_cache_mode != "all":

453 + state_indices_tensor_d = state_indices_tensor_d[

454 + :, : 1 + self.num_spec_tokens

455 + ]

456 + state_indices_tensor_p = state_indices_tensor_p[:, 0]

457 +

458 + new_metadata = replace(

459 + metadata,

460 + state_indices_tensor_d=state_indices_tensor_d,

461 + state_indices_tensor_p=state_indices_tensor_p,

462 + )

463

355 - new_metadata.state_indices_tensor = state_indices_t

356 - return new_metadata

464 + return self._update_metadata_for_cudagraph_capture(new_metadata)

vllm/v1/worker/gpu_model_runner.py

+6 -2

@@ -113,6 +113,7 @@ from vllm.v1.attention.backend import (

113 MultipleOf,

114 )

115 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder

116 +from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder

117 from vllm.v1.attention.backends.utils import (

118 create_fast_prefill_custom_backend,

119 get_dcp_local_seq_lens,

@@ -1852,7 +1853,9 @@ class GPUModelRunner(

1853 )

1854

1855 extra_attn_metadata_args = {}

1855 - if use_spec_decode and isinstance(builder, GDNAttentionMetadataBuilder):

1856 + if use_spec_decode and isinstance(

1857 + builder, (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder)

1858 + ):

1859 assert ubid is None, "UBatching not supported with GDN yet"

1860 extra_attn_metadata_args = dict(

1861 num_accepted_tokens=self.num_accepted_tokens.gpu[:num_reqs_padded],

@@ -4725,7 +4728,7 @@ class GPUModelRunner(

4728 # Set num_scheduled_tokens based on num_tokens and max_num_seqs

4729 # for dummy run with LoRA so that the num_reqs collectively

4730 # has num_tokens in total.

4728 - assert num_tokens <= self.scheduler_config.max_num_batched_tokens

4731 + assert num_tokens <= self.max_num_tokens

4732 max_num_reqs = self.scheduler_config.max_num_seqs

4733 if create_mixed_batch:

4734 assert not uniform_decode

@@ -4849,6 +4852,7 @@ class GPUModelRunner(

4852 ubatch_slices=(ubatch_slices_padded if pad_attn else ubatch_slices),

4853 for_cudagraph_capture=is_graph_capturing,

4854 slot_mappings=slot_mappings_by_group,

4855 + use_spec_decode=self.speculative_config is not None,

4856 )

4857

4858 with self.maybe_dummy_run_with_lora(