[Bugfix] Guard DeepSeek V4 MRV1 piecewise CUDA graphs (#51768)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com>

tests/test_config.py

@@ -66,6 +66,43 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected):
     assert envs.VLLM_USE_V2_MODEL_RUNNER is expected
 
 
+@pytest.mark.parametrize(
+    "cudagraph_mode",
+    [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL_AND_PIECEWISE],
+)
+def test_deepseek_v4_rejects_mrv1_piecewise_cudagraph(cudagraph_mode):
+    config = SimpleNamespace(
+        use_v2_model_runner=False,
+        model_config=SimpleNamespace(architectures=["DeepseekV4ForCausalLM"]),
+        compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode),
+    )
+
+    with pytest.raises(ValueError, match="DeepSeek V4 does not support PIECEWISE"):
+        VllmConfig._validate_mrv1_piecewise_cudagraph(config)
+
+
+@pytest.mark.parametrize(
+    ("use_v2_model_runner", "architecture", "cudagraph_mode"),
+    [
+        (True, "DeepseekV4ForCausalLM", CUDAGraphMode.PIECEWISE),
+        (False, "DeepseekV4ForCausalLM", CUDAGraphMode.NONE),
+        (False, "DeepseekV4ForCausalLM", CUDAGraphMode.FULL),
+        (False, "DeepseekV4ForCausalLM", CUDAGraphMode.FULL_DECODE_ONLY),
+        (False, "LlamaForCausalLM", CUDAGraphMode.PIECEWISE),
+    ],
+)
+def test_mrv1_piecewise_cudagraph_allowed(
+    use_v2_model_runner, architecture, cudagraph_mode
+):
+    config = SimpleNamespace(
+        use_v2_model_runner=use_v2_model_runner,
+        model_config=SimpleNamespace(architectures=[architecture]),
+        compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode),
+    )
+
+    VllmConfig._validate_mrv1_piecewise_cudagraph(config)
+
+
 @pytest.mark.parametrize(
     ("use_v2_model_runner", "expected_capture_sizes"),
     [
@@ -179,6 +216,16 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1(
             ),
             True,
         ),
+        (
+            SimpleNamespace(
+                model="deepseek-ai/DeepSeek-V4-Flash",
+                architectures=["DeepseekV4ForCausalLM"],
+                runner_type="generate",
+                is_moe=True,
+                is_quantized=True,
+            ),
+            True,
+        ),
         (
             SimpleNamespace(
                 model="Qwen/Qwen1.5-MoE-A2.7B",

vllm/config/vllm.py

@@ -66,9 +66,14 @@ else:
 
 logger = init_logger(__name__)
 
+MRV1_UNSUPPORTED_PIECEWISE_CUDAGRAPH_ARCHITECTURES = frozenset(
+    {"DeepseekV4ForCausalLM"}
+)
+
 DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(
     {
         "DeepseekV2ForCausalLM",
+        "DeepseekV4ForCausalLM",
         "GraniteMoeForCausalLM",
         "InklingForCausalLM",
         "InklingForConditionalGeneration",
@@ -702,6 +707,25 @@ class VllmConfig:
             return False
         return is_default_v2_architecture or not model_config.is_moe
 
+    def _validate_mrv1_piecewise_cudagraph(self) -> None:
+        if self.use_v2_model_runner:
+            return
+        model_config = self.model_config
+        if model_config is None:
+            return
+        if not self.compilation_config.cudagraph_mode.has_piecewise_cudagraphs():
+            return
+        architectures = getattr(model_config, "architectures", [])
+        if any(
+            arch in MRV1_UNSUPPORTED_PIECEWISE_CUDAGRAPH_ARCHITECTURES
+            for arch in architectures
+        ):
+            raise ValueError(
+                "DeepSeek V4 does not support PIECEWISE CUDA graphs with "
+                "Model Runner V1. Use Model Runner V2 or disable PIECEWISE "
+                "CUDA graphs."
+            )
+
     @property
     def needs_dp_coordinator(self) -> bool:
         """
@@ -1589,6 +1613,8 @@ class VllmConfig:
                         "pipeline parallelism",
                     )
 
+        self._validate_mrv1_piecewise_cudagraph()
+
         # final check of cudagraph mode after all possible updates
         if current_platform.is_cuda_alike():
             if (