@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4aab2b0ebe

[Model Runner V2] Use MRV2 for all models by default (#53183)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

Nick Hill committed Aug 27, 2026 at 11:59 UTC 4aab2b0ebed20343efe543c633f71b3c1336d5b8

5 files changed +122 -117

tests/model_executor/layers/test_fused_shared_expert.py

@@ -487,6 +487,7 @@ def test_models_fse_init(

487 quantization_config=None,

488 runner_type="generate",

489 is_moe=True,

490 + logits_processors=None,

491 )

492 vllm_config.parallel_config.enable_expert_parallel = False

493 if model_type == "deepseek_v4":

@@ -545,6 +546,8 @@ def test_models_fse_init(

546 vllm_config.speculative_config = SimpleNamespace(

547 draft_model_config=SimpleNamespace(hf_config=config),

548 method="mtp",

549 + parallel_drafting=False,

550 + enable_adaptive_verification=False,

551 )

552 mtp = DeepSeekV4MTP(vllm_config=vllm_config)

553 assert model.is_fused_shared_expert_enabled is (fse_enabled and not exclude)

tests/test_config.py

+54 -31

@@ -132,25 +132,30 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected):

132 def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch):

133 """ROCm keeps DeepSeek V3.2 and V4 on their compiled MRV1 paths."""

134 from vllm.config.vllm import (

135 + ROCM_DEFAULT_MRV1_ARCHITECTURES,

136 default_breakable_cudagraph_architectures,

136 - default_v2_model_runner_architectures,

137 )

138 from vllm.platforms import current_platform

139

140 monkeypatch.setattr(current_platform, "is_rocm", lambda: True)

141 # The lookup is lru_cached against a fixed platform.

142 - default_v2_model_runner_architectures.cache_clear()

142 default_breakable_cudagraph_architectures.cache_clear()

143 try:

145 - v2_architectures = default_v2_model_runner_architectures()

146 - breakable_architectures = default_breakable_cudagraph_architectures()

144 + assert "DeepseekV32ForCausalLM" in ROCM_DEFAULT_MRV1_ARCHITECTURES

145 + assert "DeepseekV4ForCausalLM" in ROCM_DEFAULT_MRV1_ARCHITECTURES

146

148 - assert "DeepseekV32ForCausalLM" not in v2_architectures

149 - assert "DeepseekV4ForCausalLM" not in v2_architectures

147 + breakable_architectures = default_breakable_cudagraph_architectures()

148 assert "DeepseekV32ForCausalLM" not in breakable_architectures

149 assert "DeepseekV32MTPModel" not in breakable_architectures

150 +

151 + # The carve-out takes effect via the runner-selection property

152 + # (warning_once args must be hashable for its lru_cache).

153 + monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False)

154 + config = SimpleNamespace(

155 + model_config=SimpleNamespace(architectures=["DeepseekV32ForCausalLM"])

156 + )

157 + assert VllmConfig.use_v2_model_runner.fget(config) is False

158 finally:

153 - default_v2_model_runner_architectures.cache_clear()

159 default_breakable_cudagraph_architectures.cache_clear()

160

161

@@ -169,17 +174,13 @@ def test_dsa_models_default_to_mrv2_and_breakable_cudagraph(

174 from vllm.compilation.breakable_cudagraph import (

175 is_breakable_cudagraph_enabled,

176 )

172 - from vllm.config.vllm import (

173 - default_breakable_cudagraph_architectures,

174 - default_v2_model_runner_architectures,

175 - )

177 + from vllm.config.vllm import default_breakable_cudagraph_architectures

178 from vllm.platforms import current_platform

179

180 monkeypatch.delenv("VLLM_USE_BREAKABLE_CUDAGRAPH", raising=False)

181 monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False)

182 monkeypatch.setattr(vllm_config_module, "HAS_TRITON", True)

183 monkeypatch.setattr(current_platform, "is_rocm", lambda: False)

182 - default_v2_model_runner_architectures.cache_clear()

184 default_breakable_cudagraph_architectures.cache_clear()

185

186 model_config = SimpleNamespace(

@@ -200,10 +201,6 @@ def test_dsa_models_default_to_mrv2_and_breakable_cudagraph(

201 ),

202 )

203 config._dflash_needs_multi_kv_group = lambda: False

203 - config._is_dflash2_draft = lambda: False

204 - config._is_default_v2_model_runner_model = lambda: (

205 - VllmConfig._is_default_v2_model_runner_model(config)

206 - )

204 config._get_v2_model_runner_unsupported_features = lambda: []

205 config._uses_breakable_cudagraph_by_default = lambda: (

206 VllmConfig._uses_breakable_cudagraph_by_default(config)

@@ -217,7 +214,6 @@ def test_dsa_models_default_to_mrv2_and_breakable_cudagraph(

214 assert config.compilation_config.cudagraph_mode.has_piecewise_cudagraphs()

215 finally:

216 os.environ.pop("VLLM_USE_BREAKABLE_CUDAGRAPH", None)

220 - default_v2_model_runner_architectures.cache_clear()

217 default_breakable_cudagraph_architectures.cache_clear()

218

219

@@ -387,7 +383,7 @@ def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling():

383 [

384 (

385 SimpleNamespace(

390 - model="Qwen/Qwen3-1.7B-Base",

386 + model="Qwen/Qwen3-32B",

387 architectures=["Qwen3ForCausalLM"],

388 runner_type="generate",

389 is_moe=False,

@@ -397,8 +393,8 @@ def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling():

393 ),

394 (

395 SimpleNamespace(

400 - model="Qwen/Qwen3-32B",

401 - architectures=["Qwen3ForCausalLM"],

396 + model="Qwen/Qwen2-7B-Instruct",

397 + architectures=["Qwen2ForCausalLM"],

398 runner_type="generate",

399 is_moe=False,

400 is_quantized=False,

@@ -465,6 +461,16 @@ def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling():

461 ),

462 True,

463 ),

464 + (

465 + SimpleNamespace(

466 + model="deepseek-ai/DeepSeek-V3",

467 + architectures=["DeepseekV3ForCausalLM"],

468 + runner_type="generate",

469 + is_moe=True,

470 + is_quantized=False,

471 + ),

472 + True,

473 + ),

474 (

475 SimpleNamespace(

476 model="deepseek-ai/DeepSeek-V4-Flash",

@@ -533,7 +539,7 @@ def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling():

539 is_moe=True,

540 is_quantized=False,

541 ),

536 - False,

542 + True,

543 ),

544 (

545 SimpleNamespace(

@@ -554,7 +560,7 @@ def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling():

560 is_quantized=False,

561 is_hybrid=True,

562 ),

557 - False,

563 + True,

564 ),

565 (

566 SimpleNamespace(

@@ -565,7 +571,7 @@ def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling():

571 is_quantized=False,

572 is_attention_free=True,

573 ),

568 - False,

574 + True,

575 ),

576 (

577 SimpleNamespace(

@@ -602,20 +608,37 @@ def test_resolve_cudagraph_mode_skips_mamba_block_check_while_profiling():

608 ),

609 ],

610 )

605 -def test_is_default_v2_model_runner_model(model_config, expected, monkeypatch):

606 - from vllm.config.vllm import default_v2_model_runner_architectures

611 +def test_models_default_to_v2_model_runner(model_config, expected, monkeypatch):

612 from vllm.platforms import current_platform

613

614 # The expectations below are the platform-independent defaults; ROCm's

610 - # DeepSeek V4 carve-out is covered by test_rocm_defaults_deepseek_v4_to_mrv1.

615 + # DeepSeek carve-out is covered by test_rocm_keeps_compiled_deepseek_defaults.

616 + monkeypatch.delenv("VLLM_USE_V2_MODEL_RUNNER", raising=False)

617 + monkeypatch.setattr(vllm_config_module, "HAS_TRITON", True)

618 monkeypatch.setattr(current_platform, "is_rocm", lambda: False)

612 - default_v2_model_runner_architectures.cache_clear()

619 config = SimpleNamespace(model_config=model_config)

620 + config._get_v2_model_runner_unsupported_features = lambda: []

621

615 - try:

616 - assert VllmConfig._is_default_v2_model_runner_model(config) is expected

617 - finally:

618 - default_v2_model_runner_architectures.cache_clear()

622 + assert VllmConfig.use_v2_model_runner.fget(config) is expected

623 +

624 +

625 +def test_v1_model_runner_rejects_v2_only_features():

626 + config = SimpleNamespace(

627 + parallel_config=SimpleNamespace(

628 + prefill_context_parallel_size=2,

629 + enable_batch_sharded_sampling=False,

630 + ),

631 + speculative_config=None,

632 + model_config=None,

633 + )

634 + config._dflash_needs_multi_kv_group = lambda: False

635 + config._is_dflash2_draft = lambda: False

636 + config._get_v1_model_runner_unsupported_features = lambda: (

637 + VllmConfig._get_v1_model_runner_unsupported_features(config)

638 + )

639 +

640 + with pytest.raises(ValueError, match="prefill context parallel"):

641 + VllmConfig._validate_v1_model_runner(config)

642

643

644 @pytest.mark.skip_global_cleanup

vllm/compilation/decorators.py

@@ -591,6 +591,15 @@ def _support_torch_compile(

591 **kwargs,

592 )

593

594 + # Tensors that enter the graph through the forward context rather than

595 + # model args (e.g. the MoE padding mask) need their batch dim marked

596 + # dynamic too, otherwise AOT-compiled artifacts bake a static guard

597 + # for the compile-time batch size.

598 + if is_forward_context_available():

599 + is_padding = get_forward_context().is_padding

600 + if is_padding is not None:

601 + torch._dynamo.mark_dynamic(is_padding, 0)

602 +

603 original_code_object = self.original_code_object()

604 logger.debug("Start compiling function %s", original_code_object)

605

vllm/config/vllm.py

+52 -85

@@ -66,19 +66,10 @@ else:

67 logger = init_logger(__name__)

69 -DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(

70 - {

71 - "DeepseekV2ForCausalLM",

72 - "DeepseekV32ForCausalLM",

73 - "DeepseekV4ForCausalLM",

74 - "GlmMoeDsaForCausalLM",

75 - "GraniteMoeForCausalLM",

76 - "InklingForCausalLM",

77 - "InklingForConditionalGeneration",

78 - "KimiK3ForConditionalGeneration",

79 - "LongcatFlashNgramForCausalLM",

80 - "Qwen2MoeForCausalLM",

81 - }

69 +# TODO(rocm): These models are either unsupported by MRV2 or slower with

70 +# MRV2 on AMD GPUs.

71 +ROCM_DEFAULT_MRV1_ARCHITECTURES = frozenset(

72 + {"DeepseekV32ForCausalLM", "DeepseekV4ForCausalLM"}

73 )

75 DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset(

@@ -99,21 +90,6 @@ DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset(

90 )

102 -@lru_cache

103 -def default_v2_model_runner_architectures() -> frozenset[str]:

104 - """Architectures defaulting to the V2 model runner on this platform."""

105 - from vllm.platforms import current_platform

106 -

107 - if current_platform.is_rocm():

108 - # TODO(rocm): These models are either unsupported by MRV2 or slower with

109 - # MRV2 on AMD GPUs.

110 - return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES - {

111 - "DeepseekV32ForCausalLM",

112 - "DeepseekV4ForCausalLM",

113 - }

114 - return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES

115 -

116 -

93 @lru_cache

94 def default_breakable_cudagraph_architectures() -> frozenset[str]:

95 """Architectures defaulting to breakable CUDA graphs on this platform."""

@@ -646,36 +622,17 @@ class VllmConfig:

622 if use_v2_model_runner is not None:

623 return use_v2_model_runner

624

649 - # PCP runtime support is implemented only by the V2 model runner.

650 - if self.parallel_config.prefill_context_parallel_size > 1:

651 - return True

652 -

653 - # DSpark is implemented only by the V2 GPU model runner, and DeepSeek-V4

654 - # is not otherwise a default-V2 architecture, so force V2 for it. If V2

655 - # is unsupported for the rest of the config, _validate_v2_model_runner

656 - # raises rather than silently falling back to V1 (which can't run dspark).

657 - if (

658 - self.speculative_config is not None

659 - and self.speculative_config.method == "dspark"

660 - ):

661 - return True

662 -

663 - # Mixed sliding/full DFlash drafts need multiple KV groups (V2 only);

664 - # force V2 as for dspark, since a hybrid target otherwise defaults to V1.

665 - if self._dflash_needs_multi_kv_group():

666 - return True

667 -

668 - # The DFlash2 candidate selector exists only in the V2 speculator. On V1

669 - # the same checkpoint drafts through DFlashProposer, which never calls

670 - # it, so the draft degrades to DFlash1 silently. Force V2 as for dspark.

671 - if self._is_dflash2_draft():

672 - return True

673 -

674 - if self.model_config is not None and self.model_config.is_diffusion:

675 - return True

625 + from vllm.platforms import current_platform

626

677 - if not self._is_default_v2_model_runner_model():

678 - return False

627 + model_config = self.model_config

628 + if model_config is not None and current_platform.is_rocm():

629 + architectures = getattr(model_config, "architectures", ())

630 + if any(arch in ROCM_DEFAULT_MRV1_ARCHITECTURES for arch in architectures):

631 + logger.warning_once(

632 + "Defaulting to V1 model runner on ROCm for model architectures: %s",

633 + ", ".join(architectures),

634 + )

635 + return False

636

637 if not HAS_TRITON:

638 logger.warning_once(

@@ -717,26 +674,6 @@ class VllmConfig:

674 num_sliding = sum(lt == "sliding_attention" for lt in layer_types)

675 return 0 < num_sliding < len(layer_types)

676

720 - def _is_default_v2_model_runner_model(self) -> bool:

721 - model_config = self.model_config

722 - if model_config is None:

723 - return False

724 -

725 - architectures = getattr(model_config, "architectures", [])

726 - default_architectures = default_v2_model_runner_architectures()

727 - is_default_v2_architecture = any(

728 - arch in default_architectures for arch in architectures

729 - )

730 -

731 - if getattr(model_config, "is_hybrid", False) and (

732 - not is_default_v2_architecture

733 - ):

734 - return False

735 -

736 - if getattr(model_config, "is_attention_free", False):

737 - return False

738 - return is_default_v2_architecture or not model_config.is_moe

739 -

677 def _uses_breakable_cudagraph_by_default(self) -> bool:

678 model_config = self.model_config

679 if model_config is None:

@@ -1666,11 +1603,8 @@ class VllmConfig:

1603

1604 if self.use_v2_model_runner:

1605 self._validate_v2_model_runner()

1669 - elif self.parallel_config.prefill_context_parallel_size > 1:

1670 - raise ValueError(

1671 - "Prefill context parallelism requires Model Runner V2. "

1672 - "Remove VLLM_USE_V2_MODEL_RUNNER=0."

1673 - )

1606 + else:

1607 + self._validate_v1_model_runner()

1608

1609 self._validate_batch_sharded_sampling()

1610

@@ -2561,6 +2495,35 @@ class VllmConfig:

2495

2496 return unsupported

2497

2498 + def _get_v1_model_runner_unsupported_features(self) -> list[str]:

2499 + unsupported: list[str] = []

2500 +

2501 + # PCP runtime support is implemented only by the V2 model runner.

2502 + if self.parallel_config.prefill_context_parallel_size > 1:

2503 + unsupported.append("prefill context parallel")

2504 +

2505 + # DSpark is implemented only by the V2 GPU model runner.

2506 + if self.speculative_config and self.speculative_config.method == "dspark":

2507 + unsupported.append("dspark speculative decoding")

2508 +

2509 + # Mixed sliding/full DFlash drafts need multiple KV groups (V2 only).

2510 + if self._dflash_needs_multi_kv_group():

2511 + unsupported.append("mixed sliding/full dflash drafts")

2512 +

2513 + # The DFlash2 candidate selector exists only in the V2 speculator. On

2514 + # V1 the same checkpoint drafts through DFlashProposer, which never

2515 + # calls it, so the draft would degrade to DFlash1 silently.

2516 + if self._is_dflash2_draft():

2517 + unsupported.append("dflash2 drafts")

2518 +

2519 + if self.model_config is not None and self.model_config.is_diffusion:

2520 + unsupported.append("diffusion models")

2521 +

2522 + if self.parallel_config.enable_batch_sharded_sampling:

2523 + unsupported.append("batch-sharded sampling")

2524 +

2525 + return unsupported

2526 +

2527 def _validate_batch_sharded_sampling(self) -> None:

2528 """Validate `enable_batch_sharded_sampling` against the rest of the config."""

2529 if not self.parallel_config.enable_batch_sharded_sampling:

@@ -2571,9 +2534,6 @@ class VllmConfig:

2534 blockers: list[str] = []

2535 tp_size = self.parallel_config.tensor_parallel_size

2536

2574 - if not self.use_v2_model_runner:

2575 - blockers.append("it is only implemented by Model Runner V2")

2576 -

2537 if tp_size <= 1:

2538 blockers.append("tensor_parallel_size is 1, so there is nothing to shard")

2539 elif self.scheduler_config.max_num_seqs < tp_size:

@@ -2623,6 +2583,13 @@ class VllmConfig:

2583 f"Model Runner V2 does not yet support: {', '.join(unsupported)}"

2584 )

2585

2586 + def _validate_v1_model_runner(self) -> None:

2587 + unsupported = self._get_v1_model_runner_unsupported_features()

2588 + if unsupported:

2589 + raise ValueError(

2590 + f"Model Runner V1 does not support: {', '.join(unsupported)}"

2591 + )

2592 +

2593 def validate_block_size(self) -> None:

2594 """Validate block_size against DCP and mamba constraints.

2595

vllm/v1/attention/backends/utils.py

+4 -1

@@ -17,6 +17,7 @@ from typing_extensions import runtime_checkable

18 from vllm.config import CacheConfig, VllmConfig, get_layers_from_vllm_config

19 from vllm.config.cache import _layout_from_name

20 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

21 from vllm.utils.math_utils import cdiv

22 from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d, np_to_pinned_tensor

23 from vllm.v1.kv_cache_interface import KVCacheLayout, KVCacheSpec, MambaSpec

@@ -475,7 +476,9 @@ def make_local_attention_virtual_batches(

476 block_size: int = 0,

477 ) -> tuple[CommonAttentionMetadata, Callable[[torch.Tensor], torch.Tensor]]:

478 query_start_loc_np = common_attn_metadata.query_start_loc_cpu.numpy()

478 - seq_lens_np = common_attn_metadata.seq_lens_cpu.numpy()

479 + with gpu_sync_allowed():

480 + # TODO see https://github.com/vllm-project/vllm/pull/31852

481 + seq_lens_np = common_attn_metadata.seq_lens_cpu.numpy()

482 block_table = common_attn_metadata.block_table_tensor

483 device = common_attn_metadata.query_start_loc.device

484