@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6fe67cbbf3

[Spec][V2] Support MTP speculative decoding under pipeline parallelism (#46994)

Signed-off-by: Chris Eastwood <chris.eastwood@pwn4g3.dev> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Chris Eastwood committed Sep 11, 2026 at 13:55 UTC 6fe67cbbf3e43da89bebf6ab0eeaca4ba6c75663

12 files changed +167 -49

tests/v1/attention/test_sparse_mla_backends.py

+45

@@ -24,6 +24,7 @@ from vllm.config import set_current_vllm_config

24 from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha

25 from vllm.model_executor.layers.attention.sparse_mla_attention import (

26 GLOBAL_TOPK_MASK_MAX_BYTES,

27 + SparseMLACommonImpl,

28 _masked_mha_workspace_fits,

29 _topk_mask_shape,

30 )

@@ -44,6 +45,9 @@ import vllm.v1.attention.backends.mla.flashinfer_mla_sparse as flashinfer_sparse

45 from vllm.model_executor.layers.attention.mla_attention import (

46 _canonicalize_sparse_mla_kv_cache_dtype,

47 )

48 +from vllm.model_executor.layers.attention.sparse_mla_attention import (

49 + SharedTopkIndicesBuffer,

50 +)

51 from vllm.utils.math_utils import cdiv

52 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (

53 FlashInferMLASparseImpl,

@@ -1864,3 +1868,44 @@ def test_fp8_mixed_batch_dcp_neutralizes_empty_rows(monkeypatch):

1868 assert out.is_contiguous()

1869 assert not out.isnan().any()

1870 assert not lse.isnan().any()

1871 +

1872 +

1873 +def test_sparse_impl_observes_repointed_indexer_buffer():

1874 + """The MTP proposer repoints the draft's indexer at the target model's buffer

1875 + after the backend impl is built, so the impl must resolve the buffer per read.

1876 + Snapshotting it in __init__ leaves the layer reading indices nothing writes."""

1877 + impl = object.__new__(FlashInferMLASparseImpl)

1878 + own = torch.zeros(4, 8, dtype=torch.int32)

1879 + target = torch.ones(4, 8, dtype=torch.int32)

1880 + indexer = SimpleNamespace(topk_indices_buffer=own)

1881 + impl.init_topk_indices_buffer(indexer, None)

1882 +

1883 + assert impl.topk_indices_buffer is own

1884 +

1885 + indexer.topk_indices_buffer = target

1886 +

1887 + assert impl.topk_indices_buffer is target

1888 +

1889 +

1890 +def test_explicit_topk_buffer_supersedes_indexer():

1891 + """Backbone skip-topk layers have no indexer, and the proposer also assigns the

1892 + shared buffer directly onto draft submodules."""

1893 + impl = object.__new__(FlashInferMLASparseImpl)

1894 + indexer = SimpleNamespace(

1895 + topk_indices_buffer=torch.zeros(2, 2, dtype=torch.int32),

1896 + )

1897 + impl.init_topk_indices_buffer(indexer, None)

1898 + explicit = torch.ones(2, 2, dtype=torch.int32)

1899 +

1900 + impl.topk_indices_buffer = explicit

1901 +

1902 + assert impl.topk_indices_buffer is explicit

1903 +

1904 +

1905 +def test_sparse_mla_common_impl_resolves_buffer_lazily():

1906 + """Guards the whole SparseMLACommonImpl family at once: a backend that

1907 + snapshots the buffer instead silently loses MTP buffer sharing."""

1908 + assert issubclass(SparseMLACommonImpl, SharedTopkIndicesBuffer)

1909 + assert isinstance(SparseMLACommonImpl.topk_indices_buffer, property), (

1910 + "topk_indices_buffer must stay a lazily-resolved property"

1911 + )

tests/v1/e2e/spec_decode/test_mtp_parallel_load.py

+1 -8

@@ -56,14 +56,7 @@ INLINE_CONFIGS = [

56 enable_expert_parallel=True,

57 enable_eplb=True,

58 ),

59 - InlineConfig(

60 - id="pp2",

61 - pp_size=2,

62 - skip_reason=(

63 - "DeepSeek MTP pipeline-parallel support is in flight upstream; "

64 - "see https://github.com/vllm-project/vllm/pull/44698"

65 - ),

66 - ),

59 + InlineConfig(id="pp2", pp_size=2),

60 ]

vllm/model_executor/layers/attention/sparse_mla_attention.py

+35 -9

@@ -501,7 +501,40 @@ def _build_topk_mask(

501 return out[:batch_size, :max_q_len]

502

503

504 -class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):

504 +class SharedTopkIndicesBuffer:

505 + """Resolves the shared top-k index buffer for sparse MLA implementations.

506 +

507 + The indexer owns the buffer, but `LLMBaseProposer.load_model` repoints the

508 + draft's indexers at the target's buffer after the impls are constructed, so

509 + it must be resolved per read rather than snapshotted. Backbone skip-topk

510 + layers have no indexer and pass the buffer explicitly.

511 + """

512 +

513 + _indexer: object | None = None

514 + _topk_indices_buffer: torch.Tensor | None = None

515 +

516 + def init_topk_indices_buffer(

517 + self,

518 + indexer: object | None,

519 + topk_indices_buffer: torch.Tensor | None,

520 + ) -> None:

521 + self._indexer = indexer

522 + self._topk_indices_buffer = topk_indices_buffer

523 +

524 + @property

525 + def topk_indices_buffer(self) -> torch.Tensor | None:

526 + if self._indexer is not None:

527 + return self._indexer.topk_indices_buffer # type: ignore[attr-defined]

528 + return self._topk_indices_buffer

529 +

530 + @topk_indices_buffer.setter

531 + def topk_indices_buffer(self, buffer: torch.Tensor | None) -> None:

532 + # An explicit assignment supersedes the indexer.

533 + self._indexer = None

534 + self._topk_indices_buffer = buffer

535 +

536 +

537 +class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic[T]):

538 """Sparse MLA base with dense and masked-MHA prefill paths."""

539

540 is_sparse = True

@@ -543,14 +576,7 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):

576 kv_b_proj,

577 )

578

546 - # The indexer carries the shared buffer for normal layers and tests;

547 - # the explicitly-passed buffer covers backbone skip layers, whose

548 - # indexer is not constructed (see deepseek_v2.py).

549 - self.topk_indices_buffer: torch.Tensor | None = (

550 - indexer.topk_indices_buffer # type: ignore[attr-defined]

551 - if indexer is not None

552 - else topk_indices_buffer

553 - )

579 + self.init_topk_indices_buffer(indexer, topk_indices_buffer)

580 self._use_flashinfer_concat_mla_k = (

581 has_flashinfer()

582 and which("ninja") is not None

vllm/model_executor/models/deepseek_mtp.py

+16 -2

@@ -39,9 +39,11 @@ from .deepseek_v2 import (

39 DeepseekV2MoE,

40 _try_load_fp8_indexer_wk,

41 )

42 +from .interfaces import SupportsPP

43 from .utils import (

44 get_pp_missing_layer_names,

45 get_spec_layer_idx_from_weight_name,

46 + make_empty_intermediate_tensors_factory,

47 maybe_prefix,

48 )

@@ -230,7 +232,7 @@ class DeepSeekMultiTokenPredictor(nn.Module):

232

233

234 @support_torch_compile

233 -class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):

235 +class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts, SupportsPP):

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

237 super().__init__()

238 self.config = vllm_config.model_config.hf_config

@@ -240,6 +242,9 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):

242 )

243 # Set MoE hyperparameters

244 self.set_moe_parameters()

245 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

246 + ["hidden_states", "residual"], self.config.hidden_size

247 + )

248

249 def set_moe_parameters(self):

250 self.num_moe_layers = self.config.num_nextn_predict_layers

@@ -266,7 +271,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):

271 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

272 return self.model.embed_input_ids(input_ids)

273

269 - def forward(

274 + def forward( # type: ignore[override]

275 self,

276 input_ids: torch.Tensor | None,

277 positions: torch.Tensor,

@@ -328,6 +333,15 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):

333 continue

334 spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)

335 if spec_layer is None:

336 + # A tied top-level embed_tokens has no spec layer to rewrite

337 + # from; the draft needs its own copy under PP.

338 + param = params_dict.get(name) if "embed_tokens" in name else None

339 + if param is not None:

340 + weight_loader = getattr(

341 + param, "weight_loader", default_weight_loader

342 + )

343 + weight_loader(param, loaded_weight)

344 + loaded_params.add(name)

345 continue

346 is_fusion_moe_shared_experts_layer = (

347 self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)

vllm/model_executor/models/qwen3_5_mtp.py

+9 -3

@@ -38,6 +38,7 @@ from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeTextConfig

38 from .interfaces import (

39 MultiModalEmbeddings,

40 SupportsMultiModal,

41 + SupportsPP,

42 _require_is_multimodal,

43 )

44 from .utils import (

@@ -152,7 +153,9 @@ class Qwen3_5MultiTokenPredictor(nn.Module):

153 inputs_embeds: torch.Tensor | None = None,

154 spec_step_idx: int = 0,

155 ) -> torch.Tensor:

155 - if get_pp_group().is_first_rank:

156 + # Branch on the inputs, not the rank: the drafter is built entirely on

157 + # the last PP stage, where `is_first_rank` is False.

158 + if intermediate_tensors is None:

159 if inputs_embeds is None:

160 inputs_embeds = self.embed_input_ids(input_ids)

161 assert hidden_states.shape[-1] == inputs_embeds.shape[-1]

@@ -162,7 +165,6 @@ class Qwen3_5MultiTokenPredictor(nn.Module):

165 hidden_states = self.fc(hidden_states)

166 residual = None

167 else:

165 - assert intermediate_tensors is not None

168 hidden_states = intermediate_tensors["hidden_states"]

169 residual = intermediate_tensors["residual"]

170

@@ -212,7 +214,7 @@ class Qwen3_5MultiTokenPredictor(nn.Module):

214 "hidden_states": 0,

215 }

216 )

215 -class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal):

217 +class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal, SupportsPP):

218 packed_modules_mapping = {

219 "qkv_proj": [

220 "q_proj",

@@ -254,6 +256,10 @@ class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal):

256

257 self.logits_processor = LogitsProcessor(config.vocab_size)

258

259 + self.make_empty_intermediate_tensors = (

260 + self.model.make_empty_intermediate_tensors

261 + )

262 +

263 def embed_input_ids(

264 self,

265 input_ids: torch.Tensor,

vllm/models/deepseek_v32/amd/mtp.py

+16 -2

@@ -32,8 +32,10 @@ from vllm.model_executor.models.deepseek_v2 import (

32 _try_load_fp8_indexer_wk,

33 get_spec_layer_idx_from_weight_name,

34 )

35 +from vllm.model_executor.models.interfaces import SupportsPP

36 from vllm.model_executor.models.utils import (

37 get_pp_missing_layer_names,

38 + make_empty_intermediate_tensors_factory,

39 maybe_prefix,

40 )

41 from vllm.models.deepseek_v32.common.kernels import fused_eh_norm

@@ -160,7 +162,7 @@ class DeepseekV32MultiTokenPredictor(nn.Module):

162 return self.logits_processor(mtp_layer.shared_head.head, hidden_states)

163

164

163 -class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):

165 +class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts, SupportsPP):

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

167 super().__init__()

168 self.config = vllm_config.model_config.hf_config

@@ -169,6 +171,9 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):

171 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

172 )

173 self.set_moe_parameters()

174 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

175 + ["hidden_states", "residual"], self.config.hidden_size

176 + )

177 self.is_fused_shared_expert_enabled = is_model_fused_shared_expert_compatible(

178 self.model.layers.values(),

179 DeepseekV2MoE,

@@ -192,7 +197,7 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):

197 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

198 return self.model.embed_input_ids(input_ids)

199

195 - def forward(

200 + def forward( # type: ignore[override]

201 self,

202 input_ids: torch.Tensor | None,

203 positions: torch.Tensor,

@@ -268,6 +273,15 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):

273 continue

274 spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)

275 if spec_layer is None:

276 + # A tied top-level embed_tokens has no spec layer to rewrite

277 + # from; the draft needs its own copy under PP.

278 + param = params_dict.get(name) if "embed_tokens" in name else None

279 + if param is not None:

280 + weight_loader = getattr(

281 + param, "weight_loader", default_weight_loader

282 + )

283 + weight_loader(param, loaded_weight)

284 + loaded_params.add(name)

285 continue

286 is_fusion_moe_shared_experts_layer = (

287 self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name)

vllm/models/deepseek_v32/nvidia/model.py

+3 -2

@@ -182,7 +182,8 @@ class DeepseekV32Model(torch.nn.Module):

182 # DSA is always sparse (has index_topk); allocate the shared top-k

183 # buffer the indexer writes and the sparse MLA backend reads.

184 self.is_v32 = True

185 - topk_indices_buffer = torch.empty(

185 + # On the model, not a local: the MTP proposer shares it with the draft.

186 + self.topk_indices_buffer = torch.empty(

187 vllm_config.scheduler_config.max_num_batched_tokens,

188 config.index_topk,

189 dtype=torch.int32,

@@ -207,7 +208,7 @@ class DeepseekV32Model(torch.nn.Module):

208 lambda prefix: DeepseekV32DecoderLayer(

209 vllm_config=vllm_config,

210 prefix=prefix,

210 - topk_indices_buffer=topk_indices_buffer,

211 + topk_indices_buffer=self.topk_indices_buffer,

212 ),

213 prefix=f"{prefix}.layers",

214 )

vllm/models/deepseek_v32/nvidia/mtp.py

+16 -2

@@ -33,8 +33,10 @@ from vllm.model_executor.models.deepseek_v2 import (

33 _try_load_fp8_indexer_wk,

34 get_spec_layer_idx_from_weight_name,

35 )

36 +from vllm.model_executor.models.interfaces import SupportsPP

37 from vllm.model_executor.models.utils import (

38 get_pp_missing_layer_names,

39 + make_empty_intermediate_tensors_factory,

40 maybe_prefix,

41 )

42 from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm

@@ -260,7 +262,7 @@ class DeepseekV32MultiTokenPredictor(nn.Module):

262 )

263

264

263 -class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):

265 +class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts, SupportsPP):

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

267 super().__init__()

268 self.config = vllm_config.model_config.hf_config

@@ -271,6 +273,9 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):

273 if self.config.model_type == "glm_moe_dsa":

274 enable_glm52_low_latency_gemm(self, vllm_config.model_config.dtype)

275 self.set_moe_parameters()

276 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(

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

278 + )

279

280 def set_moe_parameters(self):

281 self.num_moe_layers = self.config.num_nextn_predict_layers

@@ -289,7 +294,7 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):

294 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

295 return self.model.embed_input_ids(input_ids)

296

292 - def forward(

297 + def forward( # type: ignore[override]

298 self,

299 input_ids: torch.Tensor | None,

300 positions: torch.Tensor,

@@ -368,6 +373,15 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):

373 continue

374 spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)

375 if spec_layer is None:

376 + # A tied top-level embed_tokens has no spec layer to rewrite

377 + # from; the draft needs its own copy under PP.

378 + param = params_dict.get(name) if "embed_tokens" in name else None

379 + if param is not None:

380 + weight_loader = getattr(

381 + param, "weight_loader", default_weight_loader

382 + )

383 + weight_loader(param, loaded_weight)

384 + loaded_params.add(name)

385 continue

386 name = self._rewrite_spec_layer_name(spec_layer, name)

387

vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py

+7 -6

@@ -6,6 +6,9 @@ from typing import TYPE_CHECKING, cast

7 import torch

9 +from vllm.model_executor.layers.attention.sparse_mla_attention import (

10 + SharedTopkIndicesBuffer,

11 +)

12 from vllm.v1.attention.backend import (

13 AttentionLayer,

14 AttentionType,

@@ -29,7 +32,9 @@ def _kv_scale_format_for_model(model_type: str | None) -> str:

32 return "pow2_fp32"

32 -class FlashInferMLASparseSM120Impl(MLAAttentionImpl[FlashInferMLASparseMetadata]):

35 +class FlashInferMLASparseSM120Impl(

36 + MLAAttentionImpl[FlashInferMLASparseMetadata], SharedTopkIndicesBuffer

37 +):

38 """SM120 FlashInfer sparse-MLA implementation."""

40 is_sparse = True

@@ -86,11 +91,7 @@ class FlashInferMLASparseSM120Impl(MLAAttentionImpl[FlashInferMLASparseMetadata]

92 # Skip-topk layers are built with indexer=None and get the shared

93 # buffer via mla_args instead (cf. FLASHMLA_SPARSE).

89 - self.topk_indices_buffer: torch.Tensor | None = (

90 - indexer.topk_indices_buffer

91 - if indexer is not None

92 - else mla_args.get("topk_indices_buffer")

93 - )

94 + self.init_topk_indices_buffer(indexer, mla_args.get("topk_indices_buffer"))

95 from vllm.utils.flashinfer import has_flashinfer_sparse_mla_sm120

97 if not has_flashinfer_sparse_mla_sm120():

vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py

+7 -7

@@ -15,6 +15,9 @@ from vllm.logger import init_logger

15 from vllm.model_executor.layers.attention.mla_attention import (

16 get_mla_dims,

17 )

18 +from vllm.model_executor.layers.attention.sparse_mla_attention import (

19 + SharedTopkIndicesBuffer,

20 +)

21 from vllm.platforms import current_platform

22 from vllm.triton_utils import tl, triton

23 from vllm.utils.gpu_sync_debug import gpu_sync_allowed

@@ -750,7 +753,9 @@ def reference_mla_sparse_prefill(

753 return (result, lse)

754

755

753 -class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):

756 +class ROCMAiterMLASparseImpl(

757 + MLAAttentionImpl[ROCMAiterMLASparseMetadata], SharedTopkIndicesBuffer

758 +):

759 is_sparse = True

760 supports_dense_mha_prefill = False

761 supports_dcp = False

@@ -795,12 +800,7 @@ class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):

800 self.kv_cache_dtype = kv_cache_dtype

801 self.kv_lora_rank: int = mla_args["kv_lora_rank"]

802 self.softmax_scale = scale

798 - # The indexer carries the shared buffer for normal layers and tests;

799 - # the explicitly-passed buffer covers backbone skip layers, whose

800 - # indexer is not constructed (see deepseek_v2.py).

801 - self.topk_indices_buffer: torch.Tensor | None = (

802 - indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer

803 - )

803 + self.init_topk_indices_buffer(indexer, topk_indices_buffer)

804

805 vllm_config = get_current_vllm_config()

806 max_tokens = vllm_config.scheduler_config.max_num_batched_tokens

vllm/v1/attention/backends/mla/xpu_mla_sparse.py

+5 -7

@@ -13,6 +13,9 @@ from vllm.logger import init_logger

13 from vllm.model_executor.layers.attention.mla_attention import (

14 get_mla_dims,

15 )

16 +from vllm.model_executor.layers.attention.sparse_mla_attention import (

17 + SharedTopkIndicesBuffer,

18 +)

19 from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor

20 from vllm.v1.attention.backend import (

21 AttentionBackend,

@@ -183,7 +186,7 @@ class XPUMLASparseMetadataBuilder(AttentionMetadataBuilder[XPUMLASparseMetadata]

186 return metadata

187

188

186 -class XPUMLASparseImpl(MLAAttentionImpl[XPUMLASparseMetadata]):

189 +class XPUMLASparseImpl(MLAAttentionImpl[XPUMLASparseMetadata], SharedTopkIndicesBuffer):

190 is_sparse = True

191

192 def __init__(

@@ -210,12 +213,7 @@ class XPUMLASparseImpl(MLAAttentionImpl[XPUMLASparseMetadata]):

213 self.kv_cache_dtype = kv_cache_dtype

214 self.kv_lora_rank: int = mla_args["kv_lora_rank"]

215 self.softmax_scale = scale

213 - # The indexer carries the shared buffer for normal layers and tests;

214 - # the explicitly-passed buffer covers backbone skip layers, whose

215 - # indexer is not constructed (see deepseek_v2.py).

216 - self.topk_indices_buffer: torch.Tensor | None = (

217 - indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer

218 - )

216 + self.init_topk_indices_buffer(indexer, topk_indices_buffer)

217

218 def _forward_bf16_kv(

219 self,

vllm/v1/worker/gpu/model_runner.py

+7 -1

@@ -2030,7 +2030,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):

2030 )

2031

2032 mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None

2033 - if self.speculator is not None and self.speculator.supports_mm_inputs:

2033 + # The encoder runner exists only on the first PP rank, so later ranks

2034 + # have no cached embeddings to gather.

2035 + if (

2036 + self.speculator is not None

2037 + and self.speculator.supports_mm_inputs

2038 + and self.model_state.supports_mm_inputs

2039 + ):

2040 # Get cached multimodal embeddings for draft forward.

2041 # NOTE: This is done here because postprocess updates

2042 # num_computed_prefill_tokens.