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

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 2, 12:15 UTC

[Model]Support Step-3.7-Flash (#43859)

Signed-off-by: luotingdan <luotingdan@stepfun.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: luotingdan <luotingdan@stepfun.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Yu Huang <yuhuang@nvidia.com> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai>

ltd0924 committed May 29, 2026 at 08:01 UTC b690b2bb672b8a762b22a095cee1102082a58d42

15 files changed +696 -25

docs/models/supported_models.md

@@ -633,6 +633,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen

633 | `SmolVLMForConditionalGeneration` | SmolVLM2 | T + I | `SmolVLM2-2.2B-Instruct` | ✅︎ | |

634 | `Step3VLForConditionalGeneration` | Step3-VL | T + I+ | `stepfun-ai/step3` | | ✅︎ |

635 | `StepVLForConditionalGeneration` | Step3-VL-10B | T + I+ | `stepfun-ai/Step3-VL-10B` | | ✅︎ |

636 +| `Step3p7ForConditionalGeneration` | Step-3.7-Flash | T + I+ | `stepfun-ai/Step-3.7-Flash` | | ✅︎ |

637 | `TarsierForConditionalGeneration` | Tarsier | T + IE+ | `omni-search/Tarsier-7b`, `omni-search/Tarsier-34b` | | ✅︎ |

639 | `UltravoxModel` | Ultravox | T + AE+ | `fixie-ai/ultravox-v0_5-llama-3_2-1b` | ✅︎ | ✅︎ |

tests/models/registry.py

@@ -1368,6 +1368,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {

1368 "StepVLForConditionalGeneration": _HfExamplesInfo(

1369 "stepfun-ai/Step3-VL-10B", trust_remote_code=True

1370 ),

1371 + "Step3p7ForConditionalGeneration": _HfExamplesInfo(

1372 + "stepfun-ai/Step-3.7-Flash", is_available_online=False, trust_remote_code=True

1373 + ),

1374 "UltravoxModel": _HfExamplesInfo(

1375 "fixie-ai/ultravox-v0_5-llama-3_2-1b",

1376 trust_remote_code=True,

vllm/config/speculative.py

+23 -2

@@ -485,7 +485,16 @@ class SpeculativeConfig:

485 {"n_predict": n_predict, "architectures": ["LongCatFlashMTPModel"]}

486 )

487

488 - if hf_config.model_type == "step3p5":

488 + if hf_config.model_type in ("step3p5", "step3p7") or hf_config.architectures[

489 + 0

490 + ] in ("Step3p5ForCausalLM", "Step3p7ForConditionalGeneration"):

491 + quantization_config = getattr(hf_config, "quantization_config", None)

492 + hf_config = getattr(hf_config, "text_config", hf_config)

493 + if (

494 + quantization_config is not None

495 + and getattr(hf_config, "quantization_config", None) is None

496 + ):

497 + hf_config.update({"quantization_config": quantization_config})

498 hf_config.model_type = "step3p5_mtp"

499 n_predict = getattr(hf_config, "num_nextn_predict_layers", 1)

500 hf_config.update({"n_predict": n_predict, "architectures": ["Step3p5MTP"]})

@@ -705,7 +714,11 @@ class SpeculativeConfig:

714 MTPModelTypes

715 ):

716 self.method = "mtp"

708 - if self.num_speculative_tokens > 1:

717 + if (

718 + self.num_speculative_tokens > 1

719 + and self.draft_model_config.hf_config.model_type

720 + != "step3p5_mtp"

721 + ):

722 logger.warning(

723 "Enabling num_speculative_tokens > 1 will run "

724 "multiple times of forward on same MTP layer"

@@ -1056,6 +1069,14 @@ class SpeculativeConfig:

1069 == "gemma4_mtp"

1070 )

1071

1072 + def use_step3p5_mtp(self) -> bool:

1073 + return (

1074 + self.method == "mtp"

1075 + and self.draft_model_config is not None

1076 + and getattr(self.draft_model_config.hf_config, "model_type", None)

1077 + == "step3p5_mtp"

1078 + )

1079 +

1080 def use_eagle(self) -> bool:

1081 return self.method in ("eagle", "eagle3", "mtp", "dflash")

1082

vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py

+1 -1

@@ -260,7 +260,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit

260 router_logits_dtype: torch.dtype | None,

261 routing_method: RoutingMethodType,

262 ) -> bool:

263 - return True

263 + return router_logits_dtype != torch.float32

264

265 @staticmethod

266 def _supports_routing_method(

vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py

+1 -1

@@ -363,7 +363,7 @@ class TrtLlmNvFp4ExpertsMonolithic(

363 router_logits_dtype: torch.dtype | None,

364 routing_method: RoutingMethodType,

365 ) -> bool:

366 - return True

366 + return router_logits_dtype != torch.float32

367

368 def apply(

369 self,

vllm/model_executor/layers/fused_moe/layer.py

+13

@@ -974,6 +974,19 @@ class FusedMoE(PluggableLayer):

974 # this is needed for compressed-tensors only

975 loaded_weight = loaded_weight.to(param.data.device)

976

977 + # ModelOpt NVFP4 stores w13 input scales as two logical shards.

978 + # The generic assignment below would broadcast w1/w3 into the

979 + # whole expert row, so the second shard would overwrite the first.

980 + if (

981 + "ModelOpt" in quant_method_name

982 + and param.data.ndim == 2

983 + and shard_id in ("w1", "w3")

984 + ):

985 + scale_expert_id = global_expert_id if use_global_sf else expert_id

986 + scale_shard_id = 0 if shard_id == "w1" else 1

987 + param.data[scale_expert_id][scale_shard_id] = loaded_weight.reshape(())

988 + return True if return_success else None

989 +

990 if (

991 "compressed" in quant_method_name.lower()

992 and param.data[expert_id] != 1

vllm/model_executor/layers/quantization/modelopt.py

+13 -1

@@ -192,7 +192,11 @@ class ModelOptQuantConfigBase(QuantizationConfig):

192 # exclude_modules config. But need to keep them for loading quantized

193 # checkpoints generated by older versions. Then check substring matching

194 # for patterns not caught by exact match

195 - if "vision_tower" in prefix or "vision_model" in prefix:

195 + if (

196 + "vision_tower" in prefix

197 + or "vision_model" in prefix

198 + or "vit_large_projector" in prefix

199 + ):

200 return UnquantizedLinearMethod()

201

202 # now, the layer is quantized, handle it here

@@ -2340,6 +2344,14 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):

2344 if key.startswith(prefix_dot):

2345 return info["quant_algo"].upper()

2346

2347 + # FusedMoE expert prefix is e.g. "...moe.experts", while ModelOpt's

2348 + # quantized_layers entries use "...moe.gate_proj" / "...moe.up_proj".

2349 + if prefix.endswith(".experts"):

2350 + parent_dot = prefix.rsplit(".experts", 1)[0] + "."

2351 + for key, info in self.quantized_layers.items():

2352 + if key.startswith(parent_dot):

2353 + return info["quant_algo"].upper()

2354 +

2355 return None

2356

2357 @staticmethod

vllm/model_executor/models/registry.py

@@ -567,6 +567,7 @@ _MULTIMODAL_MODELS = {

567 "SmolVLMForConditionalGeneration": ("smolvlm", "SmolVLMForConditionalGeneration"),

568 "StepVLForConditionalGeneration": ("step_vl", "StepVLForConditionalGeneration"),

569 "Step3VLForConditionalGeneration": ("step3_vl", "Step3VLForConditionalGeneration"),

570 + "Step3p7ForConditionalGeneration": ("step3p7", "Step3p7ForConditionalGeneration"),

571 "TarsierForConditionalGeneration": ("tarsier", "TarsierForConditionalGeneration"),

572 "Tarsier2ForConditionalGeneration": (

573 "qwen2_vl",

vllm/model_executor/models/step3p5.py

+53 -1

@@ -637,6 +637,54 @@ class Step3p5Model(nn.Module):

637 (f".moe.experts.{base_layer}w13_weight", ".moe.gate_proj.weight", "w1"),

638 (f".moe.experts.{base_layer}w13_weight", ".moe.up_proj.weight", "w3"),

639 (f".moe.experts.{base_layer}w2_weight", ".moe.down_proj.weight", "w2"),

640 + (

641 + f".moe.experts.{base_layer}w13_weight_scale_2",

642 + ".moe.gate_proj.weight_scale_2",

643 + "w1",

644 + ),

645 + (

646 + f".moe.experts.{base_layer}w13_weight_scale_2",

647 + ".moe.up_proj.weight_scale_2",

648 + "w3",

649 + ),

650 + (

651 + f".moe.experts.{base_layer}w2_weight_scale_2",

652 + ".moe.down_proj.weight_scale_2",

653 + "w2",

654 + ),

655 + (

656 + f".moe.experts.{base_layer}w13_weight_scale",

657 + ".moe.gate_proj.weight_scale",

658 + "w1",

659 + ),

660 + (

661 + f".moe.experts.{base_layer}w13_weight_scale",

662 + ".moe.up_proj.weight_scale",

663 + "w3",

664 + ),

665 + (

666 + f".moe.experts.{base_layer}w2_weight_scale",

667 + ".moe.down_proj.weight_scale",

668 + "w2",

669 + ),

670 + # Required due to the Step3 HF model's packed expert format:

671 + # input scales are stored as moe.{gate,up,down}_proj.input_scale

672 + # rather than the standard per-expert format handled generically.

673 + (

674 + f".moe.experts.{base_layer}w13_input_scale",

675 + ".moe.gate_proj.input_scale",

676 + "w1",

677 + ),

678 + (

679 + f".moe.experts.{base_layer}w13_input_scale",

680 + ".moe.up_proj.input_scale",

681 + "w3",

682 + ),

683 + (

684 + f".moe.experts.{base_layer}w2_input_scale",

685 + ".moe.down_proj.input_scale",

686 + "w2",

687 + ),

688 ]

689

690 # New per-expert format: .moe.experts.E.gate_proj.weight_packed [out, in]

@@ -756,7 +804,11 @@ class Step3p5Model(nn.Module):

804 # Per-tensor global scales (e.g. weight_global_scale)

805 # have shape [1] in compressed-tensors NVFP4 checkpoints.

806 # Expand to per-expert before the iteration loop.

759 - if (

807 + if loaded_weight.ndim == 0:

808 + loaded_weight = loaded_weight.unsqueeze(0).expand(

809 + moe_expert_num

810 + )

811 + elif (

812 loaded_weight.shape[0] == 1

813 and loaded_weight.shape[0] != moe_expert_num

814 ):

vllm/model_executor/models/step3p7.py new

+90

@@ -0,0 +1,90 @@

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

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

3 +"""Inference-only Jurassic model."""

4 +

5 +import torch

6 +

7 +from vllm.config import VllmConfig

8 +from vllm.logger import init_logger

9 +from vllm.model_executor.layers.activation import get_act_fn

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

11 +

12 +from .step3_vl import Step3VLForConditionalGeneration

13 +from .step_vl import PerceptionEncoder

14 +from .utils import WeightsMapper, init_vllm_registered_model, maybe_prefix

15 +from .vision import run_dp_sharded_vision_model

16 +

17 +logger = init_logger(__name__)

18 +

19 +

20 +class Step3p7ForConditionalGeneration(Step3VLForConditionalGeneration):

21 + hf_to_vllm_mapper = WeightsMapper(

22 + orig_to_new_prefix={

23 + "model.vision_model.": "vision_model.",

24 + "model.vit_large_projector.": "vit_large_projector.",

25 + "model.vit_large_projector": "vit_large_projector",

26 + "model.language_model.": "language_model.model.",

27 + "model.language_model": "language_model.model",

28 + "model.": "language_model.model.",

29 + "lm_head.": "language_model.lm_head.",

30 + "lm_head": "language_model.lm_head",

31 + },

32 + orig_to_new_substr={

33 + ".attn.in_proj_weight": ".attn.qkv_proj.weight",

34 + ".attn.in_proj_bias": ".attn.qkv_proj.bias",

35 + ".mlp.c_fc": ".mlp.fc1",

36 + ".mlp.c_proj": ".mlp.fc2",

37 + },

38 + )

39 +

40 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:

41 + super(Step3VLForConditionalGeneration, self).__init__()

42 +

43 + config = vllm_config.model_config.hf_config

44 + multimodal_config = vllm_config.model_config.multimodal_config

45 + quant_config = vllm_config.quant_config

46 +

47 + self.config = config

48 + self.multimodal_config = multimodal_config

49 + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"

50 +

51 + with self._mark_tower_model(vllm_config, "image"):

52 + self.vision_model = PerceptionEncoder(

53 + config.vision_config,

54 + get_act_fn(config.vision_config.hidden_act),

55 + quant_config=quant_config,

56 + prefix=maybe_prefix(prefix, "vision_model"),

57 + )

58 + self.vit_large_projector = ColumnParallelLinear(

59 + config.vision_config.width * 4,

60 + config.text_config.hidden_size,

61 + bias=config.projector_bias,

62 + gather_output=True,

63 + quant_config=quant_config,

64 + prefix=maybe_prefix(prefix, "vit_large_projector"),

65 + disable_tp=self.use_data_parallel,

66 + )

67 +

68 + with self._mark_language_model(vllm_config):

69 + self.language_model = init_vllm_registered_model(

70 + vllm_config=vllm_config,

71 + hf_config=config.text_config,

72 + prefix=maybe_prefix(prefix, "language_model"),

73 + )

74 +

75 + self.make_empty_intermediate_tensors = (

76 + self.language_model.make_empty_intermediate_tensors

77 + )

78 +

79 + def _get_vision_model_output(

80 + self, input_tensor: torch.Tensor | None

81 + ) -> torch.Tensor | None:

82 + if input_tensor is None:

83 + return None

84 + if self.use_data_parallel:

85 + return run_dp_sharded_vision_model(input_tensor, self.vision_model)

86 + return self.vision_model(input_tensor)

87 +

88 + def _process_image_features(self, image_features: torch.Tensor) -> torch.Tensor:

89 + image_features, _ = self.vit_large_projector(image_features)

90 + return image_features

vllm/tokenizers/registry.py

+5 -1

@@ -38,7 +38,7 @@ logger = init_logger(__name__)

38 # temporary workaround and better long term solutions are:

39 # - Add model type to MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS in transformers (better)

40 # - Fix tokenizer_class on the hub for the affected models (best)

41 -_MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {"step3_vl"}

41 +_MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {"step3_vl", "step3p7"}

43 _VLLM_TOKENIZERS = {

44 "deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"),

@@ -249,6 +249,10 @@ def get_tokenizer(

249 tokenizer_cls_ = tokenizer_cls

250

251 tokenizer = tokenizer_cls_.from_pretrained(tokenizer_name, *args, **kwargs)

252 + if model_type in _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS:

253 + from vllm.tokenizers.hf import get_cached_tokenizer

254 +

255 + tokenizer = get_cached_tokenizer(tokenizer)

256 if not tokenizer.is_fast:

257 logger.warning(

258 "Using a slow tokenizer. This might cause a significant "

vllm/transformers_utils/model_arch_config_convertor.py

+23 -17

@@ -502,6 +502,11 @@ class Qwen3_5MTPModelArchConfigConvertor(ModelArchConfigConvertorBase):

502 return getattr(self.hf_text_config, "mtp_num_hidden_layers", 0)

503

504

505 +class Step3p5MTPModelArchConfigConvertor(ModelArchConfigConvertorBase):

506 + def get_num_hidden_layers(self) -> int:

507 + return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)

508 +

509 +

510 class PanguUltraMoeMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):

511 def get_num_hidden_layers(self) -> int:

512 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)

@@ -543,31 +548,32 @@ class Gemma4ModelArchConfigConvertor(ModelArchConfigConvertorBase):

548 # hf_config.model_type -> convertor class

549 MODEL_ARCH_CONFIG_CONVERTORS = {

550 "cohere_asr": CohereAsrModelArchConfigConvertor,

546 - "mamba": MambaModelArchConfigConvertor,

547 - "falcon_mamba": MambaModelArchConfigConvertor,

548 - "timm_wrapper": TerratorchModelArchConfigConvertor,

549 - "medusa": MedusaModelArchConfigConvertor,

550 - "zamba2": Zamba2ModelArchConfigConvertor,

551 - "mpt": MPTModelArchConfigConvertor,

551 "dbrx": DbrxModelArchConfigConvertor,

552 + "deepseek_mtp": DeepSeekMTPModelArchConfigConvertor,

553 + "ernie_mtp": ErnieMTPModelArchConfigConvertor,

554 "falcon": FalconModelArchConfigConvertor,

555 + "falcon_mamba": MambaModelArchConfigConvertor,

556 "gemma4": Gemma4ModelArchConfigConvertor,

555 - "gemma4_text": Gemma4ModelArchConfigConvertor,

557 "gemma4_mtp": Gemma4MTPModelArchConfigConvertor,

557 - "RefinedWeb": FalconModelArchConfigConvertor,

558 - "RefinedWebModel": FalconModelArchConfigConvertor,

559 - "nemotron-nas": NemotronNasModelArchConfigConvertor,

560 - "deepseek_mtp": DeepSeekMTPModelArchConfigConvertor,

561 - "qwen3_next_mtp": Qwen3NextMTPModelArchConfigConvertor,

562 - "qwen3_5_mtp": Qwen3_5MTPModelArchConfigConvertor,

558 + "gemma4_text": Gemma4ModelArchConfigConvertor,

559 + "glm4_moe_mtp": GLM4MoeMTPModelArchConfigConvertor,

560 + "glm_ocr_mtp": GLM4MoeMTPModelArchConfigConvertor,

561 + "longcat_flash_mtp": LongCatFlashMTPModelArchConfigConvertor,

562 + "mamba": MambaModelArchConfigConvertor,

563 + "medusa": MedusaModelArchConfigConvertor,

564 "mimo_mtp": MimoMTPModelArchConfigConvertor,

565 "mimo_v2": MimoV2ModelArchConfigConvertor,

566 "mimo_v2_flash": MimoV2ModelArchConfigConvertor,

567 "mimo_v2_mtp": MimoV2MTPModelArchConfigConvertor,

568 "mimo_v2_omni_mtp": MimoV2MTPModelArchConfigConvertor,

568 - "glm4_moe_mtp": GLM4MoeMTPModelArchConfigConvertor,

569 - "glm_ocr_mtp": GLM4MoeMTPModelArchConfigConvertor,

570 - "ernie_mtp": ErnieMTPModelArchConfigConvertor,

569 + "mpt": MPTModelArchConfigConvertor,

570 + "nemotron-nas": NemotronNasModelArchConfigConvertor,

571 "pangu_ultra_moe_mtp": PanguUltraMoeMTPModelArchConfigConvertor,

572 - "longcat_flash_mtp": LongCatFlashMTPModelArchConfigConvertor,

572 + "qwen3_5_mtp": Qwen3_5MTPModelArchConfigConvertor,

573 + "qwen3_next_mtp": Qwen3NextMTPModelArchConfigConvertor,

574 + "RefinedWeb": FalconModelArchConfigConvertor,

575 + "RefinedWebModel": FalconModelArchConfigConvertor,

576 + "step3p5_mtp": Step3p5MTPModelArchConfigConvertor,

577 + "timm_wrapper": TerratorchModelArchConfigConvertor,

578 + "zamba2": Zamba2ModelArchConfigConvertor,

579 }

vllm/v1/spec_decode/llm_base_proposer.py

@@ -1232,6 +1232,7 @@ class SpecDecodeBaseProposer:

1232 "Qwen3VLForConditionalGeneration",

1233 "Qwen3VLMoeForConditionalGeneration",

1234 "Gemma4ForConditionalGeneration",

1235 + "Step3p7ForConditionalGeneration",

1236 ]:

1237 self.model.config.image_token_index = target_model.config.image_token_id

1238 elif self.get_model_name(target_model) == "PixtralForConditionalGeneration":

vllm/v1/spec_decode/step3p5.py new

+459

@@ -0,0 +1,459 @@

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

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

3 +

4 +from copy import copy

5 +

6 +import torch

7 +

8 +from vllm.config import VllmConfig, get_layers_from_vllm_config, replace

9 +from vllm.forward_context import set_forward_context

10 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

11 +from vllm.model_executor.models.utils import get_draft_quant_config

12 +from vllm.v1.attention.backend import CommonAttentionMetadata

13 +from vllm.v1.kv_cache_interface import (

14 + KVCacheConfig,

15 + KVCacheSpec,

16 + UniformTypeKVCacheSpecs,

17 +)

18 +from vllm.v1.sample.metadata import SamplingMetadata

19 +from vllm.v1.spec_decode.eagle import EagleProposer

20 +from vllm.v1.spec_decode.utils import PADDING_SLOT_ID

21 +from vllm.v1.worker.utils import AttentionGroup

22 +

23 +

24 +class Step3p5MTPProposer(EagleProposer):

25 + """Step3.5 MTP proposer with per-layer draft-step selection."""

26 +

27 + def __init__(

28 + self,

29 + vllm_config: VllmConfig,

30 + device: torch.device,

31 + runner=None,

32 + ):

33 + super().__init__(vllm_config, device, runner)

34 + self._per_group_block_tables: dict[int, torch.Tensor] = {}

35 + self._per_group_slot_mappings: dict[int, torch.Tensor] = {}

36 + # Slot-mapping buffers for non-primary KV cache groups (the primary

37 + # group reuses self._slot_mapping_buffer from the base class).

38 + self._per_group_slot_mapping_buffers: dict[int, torch.Tensor] = {}

39 +

40 + def set_per_group_attn_metadata(

41 + self,

42 + gid: int,

43 + block_table: torch.Tensor,

44 + slot_mapping: torch.Tensor,

45 + ) -> None:

46 + self._per_group_block_tables[gid] = block_table

47 + self._per_group_slot_mappings[gid] = slot_mapping

48 +

49 + def _slot_mapping_buffer_for(self, gid: int) -> torch.Tensor:

50 + if gid == self.kv_cache_gid:

51 + return self._slot_mapping_buffer

52 + buf = self._per_group_slot_mapping_buffers.get(gid)

53 + if buf is None:

54 + buf = torch.zeros(self.max_positions, dtype=torch.int64, device=self.device)

55 + self._per_group_slot_mapping_buffers[gid] = buf

56 + return buf

57 +

58 + def _get_slot_mapping(

59 + self,

60 + num_tokens: int,

61 + slot_mapping: torch.Tensor | None = None,

62 + ) -> dict[str, torch.Tensor]:

63 + """Per-layer slot_mapping with one buffer per KV cache group."""

64 + per_layer: dict[str, torch.Tensor] = {}

65 + for attn_group in self.draft_attn_groups:

66 + gid = attn_group.kv_cache_group_id

67 + buf = self._slot_mapping_buffer_for(gid)

68 + source = self._per_group_slot_mappings.get(gid, slot_mapping)

69 + if source is not None and buf.data_ptr() != source.data_ptr():

70 + n = source.shape[0]

71 + buf[:n].copy_(source)

72 + if num_tokens > n:

73 + buf[n:num_tokens].fill_(PADDING_SLOT_ID)

74 + view = buf[:num_tokens]

75 + for layer_name in attn_group.layer_names:

76 + per_layer[layer_name] = view

77 + return per_layer

78 +

79 + def _update_positions_dependent_metadata(

80 + self,

81 + positions: torch.Tensor,

82 + common_attn_metadata: CommonAttentionMetadata,

83 + batch_size: int,

84 + input_batch_size: int,

85 + block_size: int,

86 + ) -> torch.Tensor:

87 + old_positions_1d = positions[0] if self.uses_mrope else positions

88 + positions = super()._update_positions_dependent_metadata(

89 + positions,

90 + common_attn_metadata,

91 + batch_size,

92 + input_batch_size,

93 + block_size,

94 + )

95 + # Parent already produced slot_mapping for the primary gid.

96 + self._per_group_slot_mappings[self.kv_cache_gid] = (

97 + common_attn_metadata.slot_mapping

98 + )

99 + # Recompute slot_mapping for the remaining gids using their own block tables.

100 + new_positions_1d = positions[0] if self.uses_mrope else positions

101 + exceeds = old_positions_1d + 1 >= self.max_model_len

102 + for attn_group in self.draft_attn_groups:

103 + gid = attn_group.kv_cache_group_id

104 + if gid == self.kv_cache_gid:

105 + continue

106 + block_table = self._per_group_block_tables.get(gid)

107 + if block_table is None:

108 + continue

109 + n_blocks = block_table.shape[1]

110 + bn = (new_positions_1d // block_size).clamp(max=n_blocks - 1).to(torch.long)

111 + block_ids = block_table[:batch_size].gather(1, bn.unsqueeze(1)).squeeze(1)

112 + sm = block_ids * block_size + (new_positions_1d % block_size)

113 + sm.masked_fill_(exceeds, PADDING_SLOT_ID)

114 + buf = self._slot_mapping_buffer_for(gid)

115 + buf[:batch_size].copy_(sm)

116 + if input_batch_size > batch_size:

117 + buf[batch_size:input_batch_size].fill_(PADDING_SLOT_ID)

118 + self._per_group_slot_mappings[gid] = buf[:batch_size]

119 + return positions

120 +

121 + def build_per_group_and_layer_attn_metadata(

122 + self,

123 + common_attn_metadata: CommonAttentionMetadata,

124 + draft_index: int = 0,

125 + ) -> tuple[list[object], dict[str, object]]:

126 + per_group_attn_metadata: list[object] = []

127 + per_layer_attn_metadata: dict[str, object] = {}

128 + # The proposer always works in unpadded shape. Per-group block tables

129 + # registered via set_per_group_attn_metadata are stored at the model

130 + # runner's padded shape; slice them to match cm's num_reqs.

131 + num_reqs = common_attn_metadata.num_reqs

132 + num_actual_tokens = common_attn_metadata.num_actual_tokens

133 + for attn_group in self.draft_attn_groups:

134 + gid = attn_group.kv_cache_group_id

135 + if gid in self._per_group_block_tables:

136 + cm = copy(common_attn_metadata)

137 + cm.block_table_tensor = self._per_group_block_tables[gid][:num_reqs]

138 + if gid in self._per_group_slot_mappings:

139 + sm = self._per_group_slot_mappings[gid]

140 + if sm.shape[0] >= num_actual_tokens:

141 + sm = sm[:num_actual_tokens]

142 + cm.slot_mapping = sm

143 + else:

144 + cm = common_attn_metadata

145 + attn_metadata = attn_group.get_metadata_builder().build_for_drafting(

146 + common_attn_metadata=cm,

147 + draft_index=draft_index,

148 + )

149 + per_group_attn_metadata.append(attn_metadata)

150 + for layer_name in attn_group.layer_names:

151 + per_layer_attn_metadata[layer_name] = attn_metadata

152 + return per_group_attn_metadata, per_layer_attn_metadata

153 +

154 + def _maybe_share_lm_head(self, target_language_model: torch.nn.Module) -> None:

155 + """Step3.5 MTP uses the lm_head stored in each MTP layer."""

156 +

157 + # The base MTP path shares target lm_head into shared_head.head.

158 + # Step3.5 checkpoints carry per-MTP-layer shared_head weights.

159 + return

160 +

161 + def _create_draft_vllm_config(self) -> VllmConfig:

162 + base = super()._create_draft_vllm_config()

163 + return replace(

164 + base,

165 + model_config=self.draft_model_config,

166 + quant_config=get_draft_quant_config(base),

167 + )

168 +

169 + def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:

170 + """Step3.5 MTP draft layers may span multiple KV cache groups."""

171 + return

172 +

173 + def initialize_attn_backend(

174 + self,

175 + kv_cache_config: KVCacheConfig,

176 + kernel_block_sizes: list[int] | None = None,

177 + ) -> None:

178 + all_attn_layers = get_layers_from_vllm_config(

179 + self.vllm_config,

180 + AttentionLayerBase, # type: ignore[type-abstract]

181 + )

182 +

183 + layer_to_gid: dict[str, int] = {}

184 + layer_to_spec: dict[str, KVCacheSpec] = {}

185 + for gid, group in enumerate(kv_cache_config.kv_cache_groups):

186 + group_spec = group.kv_cache_spec

187 + for layer_name in group.layer_names:

188 + layer_to_gid[layer_name] = gid

189 + if isinstance(group_spec, UniformTypeKVCacheSpecs):

190 + if layer_name in group_spec.kv_cache_specs:

191 + layer_to_spec[layer_name] = group_spec.kv_cache_specs[

192 + layer_name

193 + ]

194 + else:

195 + target_layer_name = getattr(

196 + all_attn_layers.get(layer_name),

197 + "kv_sharing_target_layer_name",

198 + None,

199 + )

200 + if (

201 + target_layer_name

202 + and target_layer_name in group_spec.kv_cache_specs

203 + ):

204 + layer_to_spec[layer_name] = group_spec.kv_cache_specs[

205 + target_layer_name

206 + ]

207 + else:

208 + layer_to_spec[layer_name] = group_spec

209 + else:

210 + layer_to_spec[layer_name] = group_spec

211 +

212 + attention_groups: dict[tuple[tuple[str, str], int], AttentionGroup] = {}

213 + for layer_name in sorted(self._draft_attn_layer_names):

214 + if layer_name not in layer_to_spec:

215 + continue

216 + attn_layer = all_attn_layers[layer_name]

217 + attn_backend = attn_layer.get_attn_backend()

218 + spec = layer_to_spec[layer_name]

219 + gid = layer_to_gid[layer_name]

220 + group_key = (attn_backend.full_cls_name(), gid)

221 +

222 + if group_key not in attention_groups:

223 + kernel_block_size = (

224 + kernel_block_sizes[gid]

225 + if kernel_block_sizes is not None and gid < len(kernel_block_sizes)

226 + else None

227 + )

228 + attn_group = AttentionGroup(

229 + backend=attn_backend,

230 + layer_names=[layer_name],

231 + kv_cache_spec=spec,

232 + kv_cache_group_id=gid,

233 + )

234 + attn_group.create_metadata_builders(

235 + self.vllm_config,

236 + self.device,

237 + kernel_block_size=kernel_block_size,

238 + )

239 + attention_groups[group_key] = attn_group

240 + else:

241 + attention_groups[group_key].layer_names.append(layer_name)

242 +

243 + self.draft_attn_groups = list(attention_groups.values())

244 + if self.draft_attn_groups:

245 + self.kv_cache_gid = self.draft_attn_groups[0].kv_cache_group_id

246 + self.block_size = (

247 + self.draft_attn_groups[0]

248 + .get_metadata_builder()

249 + .kv_cache_spec.block_size

250 + )

251 + else:

252 + self.kv_cache_gid = 0

253 + self.block_size = kv_cache_config.kv_cache_groups[

254 + 0

255 + ].kv_cache_spec.block_size

256 +

257 + def _sample_draft_tokens_for_step(

258 + self,

259 + hidden_states: torch.Tensor,

260 + sampling_metadata: SamplingMetadata,

261 + spec_step_idx: int,

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

263 + if not self._enable_probabilistic_draft_probs or sampling_metadata.all_greedy:

264 + if self.use_local_argmax_reduction:

265 + return self.model.get_top_tokens(hidden_states), None

266 + logits = self.model.compute_logits(

267 + hidden_states, spec_step_idx=spec_step_idx

268 + )

269 + return logits.argmax(dim=-1), None

270 +

271 + logits = self.model.compute_logits(hidden_states, spec_step_idx=spec_step_idx)

272 + return self._sample_from_logits(logits, sampling_metadata)

273 +

274 + def propose(

275 + self,

276 + target_token_ids: torch.Tensor,

277 + target_positions: torch.Tensor,

278 + target_hidden_states: torch.Tensor,

279 + next_token_ids: torch.Tensor,

280 + token_indices_to_sample: torch.Tensor | None,

281 + common_attn_metadata: CommonAttentionMetadata,

282 + sampling_metadata: SamplingMetadata,

283 + mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,

284 + num_rejected_tokens_gpu: torch.Tensor | None = None,

285 + slot_mappings: dict[str, torch.Tensor]

286 + | list[dict[str, torch.Tensor]]

287 + | None = None,

288 + ) -> torch.Tensor:

289 + self._last_draft_probs = None

290 + batch_size = common_attn_metadata.batch_size()

291 +

292 + num_tokens, token_indices_to_sample, common_attn_metadata = (

293 + self.set_inputs_first_pass(

294 + target_token_ids=target_token_ids,

295 + next_token_ids=next_token_ids,

296 + target_positions=target_positions,

297 + target_hidden_states=target_hidden_states,

298 + token_indices_to_sample=token_indices_to_sample,

299 + cad=common_attn_metadata,

300 + num_rejected_tokens_gpu=num_rejected_tokens_gpu,

301 + )

302 + )

303 +

304 + per_group_attn_metadata, per_layer_attn_metadata = (

305 + self.build_per_group_and_layer_attn_metadata(common_attn_metadata)

306 + )

307 +

308 + cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (

309 + self._determine_batch_execution_and_padding(num_tokens)

310 + )

311 +

312 + model_kwargs, slot_mapping_size = self.build_model_inputs_first_pass(

313 + num_tokens, num_input_tokens, mm_embed_inputs

314 + )

315 + model_kwargs["spec_step_idx"] = 0

316 +

317 + with set_forward_context(

318 + per_layer_attn_metadata,

319 + self.vllm_config,

320 + num_tokens=num_input_tokens,

321 + num_tokens_across_dp=num_tokens_across_dp,

322 + cudagraph_runtime_mode=cudagraph_runtime_mode,

323 + slot_mapping=self._get_slot_mapping(

324 + slot_mapping_size, common_attn_metadata.slot_mapping

325 + ),

326 + ):

327 + ret_hidden_states = self.model(**model_kwargs)

328 + if not self.model_returns_tuple():

329 + last_hidden_states = ret_hidden_states

330 + hidden_states = last_hidden_states

331 + else:

332 + last_hidden_states, hidden_states = ret_hidden_states

333 +

334 + sample_hidden_states = last_hidden_states[token_indices_to_sample]

335 +

336 + if self.num_speculative_tokens == 1 or self.parallel_drafting:

337 + draft_token_ids, draft_probs = self._sample_draft_tokens_for_step(

338 + sample_hidden_states, sampling_metadata, spec_step_idx=0

339 + )

340 + if draft_probs is not None:

341 + self._last_draft_probs = draft_probs.view(

342 + -1, self.num_speculative_tokens, draft_probs.shape[-1]

343 + ).contiguous()

344 + return draft_token_ids.view(-1, self.num_speculative_tokens)

345 +

346 + if self.uses_mrope:

347 + positions = self.mrope_positions[:, token_indices_to_sample]

348 + else:

349 + positions = self.positions[token_indices_to_sample]

350 + hidden_states = hidden_states[token_indices_to_sample]

351 +

352 + if self.constant_draft_positions:

353 + self.positions[:batch_size] = positions

354 +

355 + draft_token_ids, draft_probs = self._sample_draft_tokens_for_step(

356 + sample_hidden_states, sampling_metadata, spec_step_idx=0

357 + )

358 + draft_probs_list = None if draft_probs is None else [draft_probs]

359 +

360 + if self.allowed_attn_types is not None:

361 + for group_md in per_group_attn_metadata:

362 + if not isinstance(group_md, self.allowed_attn_types):

363 + raise ValueError(

364 + f"Unsupported attention metadata type for speculative "

365 + "decoding with num_speculative_tokens > 1: "

366 + f"{type(group_md)}. Supported types are: "

367 + f"{self.allowed_attn_types}"

368 + )

369 +

370 + draft_token_ids_list = [draft_token_ids]

371 +

372 + cudagraph_runtime_mode, input_batch_size, batch_size_across_dp = (

373 + self._determine_batch_execution_and_padding(batch_size)

374 + )

375 +

376 + common_attn_metadata.num_actual_tokens = batch_size

377 + common_attn_metadata.max_query_len = 1

378 + common_attn_metadata.query_start_loc = self.arange[: batch_size + 1]

379 + common_attn_metadata.query_start_loc_cpu = torch.from_numpy(

380 + self.token_arange_np[: batch_size + 1]

381 + ).clone()

382 +

383 + if self.num_speculative_tokens > 1 and num_rejected_tokens_gpu is not None:

384 + common_attn_metadata.seq_lens -= num_rejected_tokens_gpu

385 + common_attn_metadata._seq_lens_cpu = None

386 + common_attn_metadata._num_computed_tokens_cpu = None

387 +

388 + block_size = self.block_size

389 + assert block_size > 0, "block_size has not been initialized."

390 + for token_index in range(self.num_speculative_tokens - 1):

391 + spec_step_idx = token_index + 1

392 + input_ids = draft_token_ids_list[-1].int()

393 +

394 + if not self.constant_draft_positions:

395 + positions = self._update_positions_dependent_metadata(

396 + positions,

397 + common_attn_metadata,

398 + batch_size,

399 + input_batch_size,

400 + block_size,

401 + )

402 +

403 + if not self.constant_draft_positions or token_index == 0:

404 + _, per_layer_attn_metadata = (

405 + self.build_per_group_and_layer_attn_metadata(

406 + common_attn_metadata, draft_index=spec_step_idx

407 + )

408 + )

409 +

410 + self.input_ids[:batch_size] = input_ids

411 + self.hidden_states[:batch_size] = hidden_states

412 + if self.supports_mm_inputs:

413 + self.inputs_embeds[:batch_size] = self.model.embed_input_ids(input_ids)

414 +

415 + input_ids = None

416 + inputs_embeds = self.inputs_embeds[:input_batch_size]

417 + else:

418 + input_ids = self.input_ids[:input_batch_size]

419 + inputs_embeds = None

420 +

421 + model_kwargs = {

422 + "input_ids": input_ids,

423 + "positions": self._get_positions(input_batch_size),

424 + "inputs_embeds": inputs_embeds,

425 + "spec_step_idx": spec_step_idx,

426 + }

427 + if self.pass_hidden_states_to_model:

428 + model_kwargs["hidden_states"] = self.hidden_states[:input_batch_size]

429 +

430 + with set_forward_context(

431 + per_layer_attn_metadata,

432 + self.vllm_config,

433 + num_tokens=input_batch_size,

434 + num_tokens_across_dp=batch_size_across_dp,

435 + cudagraph_runtime_mode=cudagraph_runtime_mode,

436 + slot_mapping=self._get_slot_mapping(input_batch_size),

437 + ):

438 + ret_hidden_states = self.model(**model_kwargs)

439 + if not self.model_returns_tuple():

440 + last_hidden_states = ret_hidden_states

441 + hidden_states = ret_hidden_states

442 + else:

443 + last_hidden_states, hidden_states = ret_hidden_states

444 +

445 + hidden_states = hidden_states[:batch_size]

446 + draft_token_ids, draft_probs = self._sample_draft_tokens_for_step(

447 + last_hidden_states[:batch_size],

448 + sampling_metadata,

449 + spec_step_idx=spec_step_idx,

450 + )

451 + if draft_probs is not None:

452 + assert draft_probs_list is not None

453 + draft_probs_list.append(draft_probs)

454 + draft_token_ids_list.append(draft_token_ids)

455 +

456 + draft_token_ids = torch.stack(draft_token_ids_list, dim=1)

457 + if draft_probs_list is not None:

458 + self._last_draft_probs = torch.stack(draft_probs_list, dim=1).contiguous()

459 + return draft_token_ids

vllm/v1/worker/gpu_model_runner.py

+9 -1

@@ -186,6 +186,7 @@ from vllm.v1.spec_decode.ngram_proposer_gpu import (

186 update_ngram_gpu_tensors_incremental,

187 update_scheduler_for_invalid_drafts,

188 )

189 +from vllm.v1.spec_decode.step3p5 import Step3p5MTPProposer

190 from vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer

191 from vllm.v1.spec_decode.utils import update_num_computed_tokens_for_batch_change

192 from vllm.v1.structured_output.utils import apply_grammar_bitmask

@@ -547,6 +548,7 @@ class GPUModelRunner(

548 | MedusaProposer

549 | ExtractHiddenStatesProposer

550 | Gemma4Proposer

551 + | Step3p5MTPProposer

552 )

553 if self.speculative_config.method == "custom_class":

554 self.drafter = create_custom_proposer( # type: ignore[assignment]

@@ -581,6 +583,8 @@ class GPUModelRunner(

583 )

584 elif self.speculative_config.use_gemma4_mtp():

585 self.drafter = Gemma4Proposer(self.vllm_config, self.device, self)

586 + elif self.speculative_config.use_step3p5_mtp():

587 + self.drafter = Step3p5MTPProposer(self.vllm_config, self.device, self)

588 elif self.speculative_config.use_dflash():

589 self.drafter = DFlashProposer(self.vllm_config, self.device, self)

590 self.use_aux_hidden_state_outputs = True

@@ -2428,7 +2432,11 @@ class GPUModelRunner(

2432 else:

2433 spec_decode_common_attn_metadata = cm

2434 # Capture per-group block tables for multi-group proposers.

2431 - if self.speculative_config and isinstance(self.drafter, Gemma4Proposer):

2435 + if self.speculative_config and isinstance(self.drafter, Step3p5MTPProposer):

2436 + self.drafter.set_per_group_attn_metadata(

2437 + kv_cache_gid, cm.block_table_tensor, cm.slot_mapping

2438 + )

2439 + elif self.speculative_config and isinstance(self.drafter, Gemma4Proposer):

2440 self.drafter.set_per_group_block_table(

2441 kv_cache_gid, cm.block_table_tensor

2442 )