[Qwen4Exp] Restore torch.compile on the NVIDIA path; exact breakable pin; runtime graph-mode assertion

- torch.compile restored on the NVIDIA Qwen4Exp target and MTP so the
  production profile (VLLM_COMPILE + FULL_AND_PIECEWISE, breakable off) runs.
- The production pin requires VLLM_USE_BREAKABLE_CUDAGRAPH to be exactly "0".
- assert_resolved_cudagraph_modes() runs at the top of capture_model.
- CPU tests for the combined profile, the exact pin and the resolved modes.

Co-Authored-By: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>

vllm/config/compilation.py

@@ -767,7 +767,12 @@ class CompilationConfig:
         "vllm::mamba_mixer2",
         "vllm::mamba_mixer",
         "vllm::short_conv",
-        # Qwen4Exp's AMD backend still uses these splitting ops.
+        # Qwen4Exp: the AMD backend and (fork, compile-profile guard) the compile-restored
+        # NVIDIA backend. #55272 dropped the NVIDIA-only ngram-ids split; the
+        # prefetch pair is the UVA PLE side-stream fork/join (#54371).
+        "vllm::qwen4_exp_compute_ple_ngram_ids",
+        "vllm::qwen4_exp_ple_start_prefetch",
+        "vllm::qwen4_exp_ple_finalize_prefetch",
         "vllm::qwen4_exp_ple_short_conv",
         "vllm::qwen4_exp_qsa_with_output",
         "vllm::linear_attention",

vllm/config/vllm.py

@@ -106,9 +106,12 @@ DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset(
         "KimiLinearForCausalLM",
         "MiniMaxM3SparseForCausalLM",
         "MiniMaxM3SparseForConditionalGeneration",
-        "Qwen4ExpForCausalLM",
-        "Qwen4ExpForConditionalGeneration",
-        "Qwen4ExpMTP",
+        # Fork (compile-profile guard): Qwen4ExpForCausalLM,
+        # Qwen4ExpForConditionalGeneration and Qwen4ExpMTP are deliberately
+        # NOT here. #55272 added them when it removed torch.compile from the
+        # NVIDIA path; this tree restores compile, so the resolved profile for
+        # them is VLLM_COMPILE + FULL_AND_PIECEWISE with breakable capture off
+        # (vllm/v1/e2_guards.py asserts it at model construction).
     }
 )
 

vllm/models/qwen4_exp/nvidia/model.py

@@ -10,6 +10,7 @@ import torch
 from torch import nn
 from transformers import Qwen4ExpConfig, Qwen4ExpTextConfig
 
+from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
 from vllm.distributed import get_pp_group
 from vllm.logger import init_logger
@@ -78,7 +79,11 @@ from vllm.multimodal.inputs import MultiModalFeatureSpec
 from vllm.sequence import IntermediateTensors
 from vllm.tokenizers.registry import cached_tokenizer_from_config
 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
-from vllm.v1.e2_guards import assert_graph_profile_is_executable, assert_ple_mode
+from vllm.v1.e2_guards import (
+    assert_compiled_submodule,
+    assert_graph_profile_is_executable,
+    assert_ple_mode,
+)
 from vllm.v1.kv_cache_interface import MambaSpec
 
 from ..common.mtp import make_mtp_hidden_buffer
@@ -472,6 +477,25 @@ class Qwen4ExpMixtureOfExperts(MixtureOfExperts):
             moe.experts.update_expert_map()
 
 
+# Compile-profile guard: torch.compile restored on the NVIDIA
+# path, reverting #55272's removal. Every request-dependent region is a custom
+# op listed in CompilationConfig._attention_ops (qsa_with_output,
+# ple_short_conv, compute_ple_ngram_ids, ple_{start,finalize}_prefetch), so
+# VLLM_COMPILE + FULL_AND_PIECEWISE is executable again with breakable capture
+# off. query_start_loc / ngram_context are fixed-capacity since #55272
+# (model_state.py); they stay marked dynamic as on v2.0.1 and only ever reach
+# the opaque PLE ops.
+@support_torch_compile(
+    dynamic_arg_dims={
+        "input_ids": 0,
+        "positions": -1,
+        "intermediate_tensors": 0,
+        "inputs_embeds": 0,
+        "query_start_loc": 0,
+        "ngram_context": 0,
+        "deepstack_input_embeds": 0,
+    }
+)
 class Qwen4ExpModel(nn.Module):
     hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER
 
@@ -760,12 +784,11 @@ class Qwen4ExpForCausalLM(
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         super().__init__()
-        # V0.30.0 auto-enables breakable cudagraphs
-        # for this architecture when VLLM_USE_BREAKABLE_CUDAGRAPH is unset, and
-        # VLLM_COMPILE + a PIECEWISE mode dies at capture because #55272 left no
-        # compiled submodule. Fail here rather than let a reference-host run qualify an
-        # identity nobody chose. Qwen4ExpForConditionalGeneration constructs
-        # this class; Qwen4ExpMTP does not, and guards itself.
+        # Compile-profile guard: the only qualified profile on this tree is VLLM_COMPILE +
+        # FULL_AND_PIECEWISE with breakable capture off (the v2.0.1 identity).
+        # Fail at construction, before profiling or capture, on anything else.
+        # Qwen4ExpForConditionalGeneration constructs this class; Qwen4ExpMTP
+        # does not, and guards itself.
         assert_graph_profile_is_executable(vllm_config, role="target")
         assert_ple_mode(vllm_config, role="target")
         config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
@@ -783,6 +806,9 @@ class Qwen4ExpForCausalLM(
         self.model = Qwen4ExpModel(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
         )
+        # has_compiled_submodule() at gpu/cudagraph_utils.py:605 is what capture
+        # actually tests; check the same property here.
+        assert_compiled_submodule(self.model, vllm_config, role="target")
         # PP ownership: logits are computed on the
         # last stage only; earlier stages carried an unused 0.59 GiB head.
         if get_pp_group().is_last_rank:

vllm/models/qwen4_exp/nvidia/mtp.py

@@ -20,6 +20,7 @@ import torch
 from torch import nn
 from transformers import Qwen4ExpTextConfig
 
+from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig, replace, set_current_vllm_config
 from vllm.distributed import get_pp_group
 from vllm.model_executor.layers.fused_moe.utils import (
@@ -51,7 +52,11 @@ from vllm.model_executor.models.utils import (
     maybe_prefix,
 )
 from vllm.sequence import IntermediateTensors
-from vllm.v1.e2_guards import assert_graph_profile_is_executable, assert_ple_mode
+from vllm.v1.e2_guards import (
+    assert_compiled_submodule,
+    assert_graph_profile_is_executable,
+    assert_ple_mode,
+)
 
 from .hyperconnection import GatedResidual, HyperConnectionConfig
 from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
@@ -168,6 +173,18 @@ def _make_draft_vllm_config(
     return draft_vllm_config
 
 
+# Compile restored on the drafter too (reverts #55272 here), as
+# on v2.0.1: both classes wrapped; the outer one compiles, the inner is
+# inlined while Dynamo traces.
+@support_torch_compile(
+    dynamic_arg_dims={
+        "input_ids": 0,
+        "positions": -1,
+        "intermediate_tensors": 0,
+        "inputs_embeds": 0,
+        "hidden_states": 0,
+    }
+)
 class Qwen4ExpMultiTokenPredictor(nn.Module):
     hf_to_vllm_mapper = Qwen3_5Model.hf_to_vllm_mapper | _EXTRA_WEIGHTS_MAPPER
 
@@ -411,6 +428,7 @@ class Qwen4ExpMTP(
             vllm_config=vllm_config,
             prefix=maybe_prefix(prefix, "mtp"),
         )
+        assert_compiled_submodule(self.model, vllm_config, role="draft")
 
         if get_pp_group().is_last_rank:
             if config.tie_word_embeddings:

vllm/models/qwen4_exp/nvidia/ngram_embedding.py

@@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig,
 )
 from vllm.model_executor.models.utils import AutoWeightsLoader
+from vllm.utils.torch_utils import direct_register_custom_op
 
 from ..common.ngram_embedding import (
     Qwen4ExpPLEDeviceEmbedding,
@@ -177,6 +178,15 @@ class Qwen4ExpPLEFileGatherEmbedding(Qwen4ExpPLEEmbedding):
             self._host_rows[:num_tokens], non_blocking=True
         )
 
+    @property
+    def prefetch_buffer(self) -> torch.Tensor:
+        """Side-stream lookup destination (declared mutated by the start op)."""
+        return self._prefetch_buffer
+
+    def new_prefetch_output(self, num_tokens: int) -> torch.Tensor:
+        """Graph-owned output storage for the finalize op."""
+        return self._prefetch_buffer.new_empty((num_tokens, self._output_dim))
+
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
         """Return the rows staged for this step."""
         return self._staging[: hidden_states.shape[0]].flatten(-2)
@@ -296,10 +306,16 @@ class Qwen4ExpNGramEmbedding(nn.Module):
         *,
         data_parallel_rank: int,
         prefix: str,
+        layer_name: str | None = None,
         quant_config: QuantizationConfig | None = None,
         params_dtype: torch.dtype | None = None,
     ) -> None:
         super().__init__()
+        # Compile-profile guard: the owning Qwen4ExpPLELayer's static_forward_context key.
+        # When set, request-dependent work runs through the custom ops at the
+        # bottom of this file so VLLM_COMPILE never traces it (#55272 removed
+        # them together with @support_torch_compile). None = direct calls.
+        self.layer_name = layer_name
         self.embedding_dim = embedding_dim
         self.ngram_size = int(config.ngram_size)
         self.heads_per_ngram = int(config.heads_per_ngram)
@@ -531,12 +547,43 @@ class Qwen4ExpNGramEmbedding(nn.Module):
         ngram_context: torch.Tensor,
     ) -> torch.Tensor:
         embedding = self.ngram_embedding
-        if embedding.supports_prefetch or isinstance(
-            embedding, Qwen4ExpPLEFileGatherEmbedding
-        ):
+        if isinstance(embedding, Qwen4ExpPLEFileGatherEmbedding):
+            # Fork: this base's third backend stages rows in prepare_inputs,
+            # so it has no graph-owned lookup op.
             return embedding(hidden_states)
-        ngram_ids = self.compute_ngram_ids(input_ids, query_start_loc, ngram_context)
-        return self.ngram_embedding(ngram_ids).flatten(-2)
+        if self.layer_name is None:
+            if embedding.supports_prefetch:
+                return embedding(hidden_states)
+            ngram_ids = self.compute_ngram_ids(
+                input_ids, query_start_loc, ngram_context
+            )
+            return self.ngram_embedding(ngram_ids).flatten(-2)
+
+        if embedding.supports_prefetch:
+            output = embedding.new_prefetch_output(hidden_states.shape[0])
+            torch.ops.vllm.qwen4_exp_ple_finalize_prefetch(output, self.layer_name)
+            return output
+        input_ids = input_ids.reshape(-1)
+        # Request-layout-dependent ID generation stays outside PIECEWISE graphs
+        # (a splitting op), as on v2.0.1; the op writes graph-owned storage.
+        ngram_ids = input_ids.new_empty(
+            (input_ids.numel(), self.ngram_heads), dtype=torch.long
+        )
+        torch.ops.vllm.qwen4_exp_compute_ple_ngram_ids(
+            input_ids,
+            query_start_loc,
+            ngram_context,
+            ngram_ids,
+            self.layer_name,
+        )
+        # The table lookup is opaque so Inductor never sees (or autotunes
+        # against) the full PLE table -- #55272's ~50 GB activation peak.
+        output = ngram_ids.new_empty(
+            (ngram_ids.shape[0], self.ngram_heads * self.head_dim),
+            dtype=embedding.weight.dtype,
+        )
+        torch.ops.vllm.qwen4_exp_ple_device_lookup(ngram_ids, output, self.layer_name)
+        return output
 
     def start_prefetch(
         self,
@@ -549,6 +596,32 @@ class Qwen4ExpNGramEmbedding(nn.Module):
         embedding = self.ngram_embedding
         if not embedding.supports_prefetch:
             return
+        if self.layer_name is not None:
+            # Forks onto the prefetch side stream; the join is in the finalize
+            # op. Both are splitting ops, so neither end of the fork is ever
+            # inside a piecewise capture.
+            torch.ops.vllm.qwen4_exp_ple_start_prefetch(
+                hidden_states,
+                input_ids,
+                query_start_loc,
+                ngram_context,
+                embedding.prefetch_buffer,
+                self.layer_name,
+            )
+            return
+        self._start_prefetch_eager(
+            hidden_states, input_ids, query_start_loc, ngram_context
+        )
+
+    def _start_prefetch_eager(
+        self,
+        hidden_states: torch.Tensor,
+        input_ids: torch.Tensor,
+        query_start_loc: torch.Tensor,
+        ngram_context: torch.Tensor,
+    ) -> None:
+        """Compute IDs and launch the pinned lookup (runs outside compile)."""
+        embedding = self.ngram_embedding
         ngram_ids = self.compute_ngram_ids(
             input_ids,
             query_start_loc,
@@ -646,6 +719,138 @@ def stage_checkpoint_rows(
         module.ngram_embedding.stage_rows(num_tokens)
 
 
+# --- compile support ------------------------------------------------
+# The four ops below are the NVIDIA PLE's boundary with torch.compile. Each
+# resolves its Qwen4ExpPLELayer through the forward context by name, as
+# v2.0.1's qwen4_exp_compute_ple_ngram_ids did before #55272 removed it.
+# Splitting ops (vllm/config/compilation.py ``_attention_ops``):
+#   qwen4_exp_compute_ple_ngram_ids, qwen4_exp_ple_start_prefetch,
+#   qwen4_exp_ple_finalize_prefetch.
+# qwen4_exp_ple_device_lookup is opaque but NOT a split: a gather plus the ETP
+# all-reduce is capturable; it only has to be hidden from Inductor.
+
+
+def _ple_embedding_of(layer_name: str) -> Qwen4ExpNGramEmbedding:
+    layer = get_forward_context().no_compile_layers[layer_name]
+    embedding = getattr(layer, "ple_embedding", None)
+    if not isinstance(embedding, Qwen4ExpNGramEmbedding):
+        raise TypeError(f"{layer_name} does not own a Qwen4Exp PLE embedding")
+    if hasattr(type(embedding), "dequantize_pulled"):
+        # Fork (pull transport + compile profile): the pull transport's GPU placeholder owns no
+        # table and dispatches through PleOffloadLayer.forward; it must never
+        # reach these device/UVA ops (ple_layer.py builds it without a
+        # layer_name). Same marker ple_layer.py uses to select dequantisation.
+        raise TypeError(
+            f"{layer_name}: PLE pull transport module routed through an "
+            "upstream device/UVA PLE op"
+        )
+    return embedding
+
+
+def qwen4_exp_compute_ple_ngram_ids(
+    input_ids: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    ngram_context: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: str,
+) -> None:
+    """Compute request-dependent PLE n-gram IDs outside piecewise graphs."""
+    _ple_embedding_of(layer_name).compute_ngram_ids(
+        input_ids, query_start_loc, ngram_context, output
+    )
+
+
+def qwen4_exp_ple_device_lookup(
+    ngram_ids: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: str,
+) -> None:
+    """Look up device-resident PLE rows; keeps the table out of Inductor."""
+    embedding = _ple_embedding_of(layer_name)
+    output.copy_(embedding.ngram_embedding(ngram_ids).flatten(-2))
+
+
+def qwen4_exp_ple_start_prefetch(
+    hidden_states: torch.Tensor,
+    input_ids: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    ngram_context: torch.Tensor,
+    prefetch_buffer: torch.Tensor,
+    layer_name: str,
+) -> None:
+    """Start the pinned-host (UVA) lookup on the side stream."""
+    embedding = _ple_embedding_of(layer_name)
+    owned = embedding.ngram_embedding.prefetch_buffer
+    if prefetch_buffer.data_ptr() != owned.data_ptr():
+        raise RuntimeError(f"{layer_name}: PLE prefetch buffer identity changed")
+    embedding._start_prefetch_eager(
+        hidden_states, input_ids, query_start_loc, ngram_context
+    )
+
+
+def qwen4_exp_ple_finalize_prefetch(output: torch.Tensor, layer_name: str) -> None:
+    """Join the side stream and write the pinned lookup into ``output``."""
+    pinned = _ple_embedding_of(layer_name).ngram_embedding
+    pinned._finalize_prefetch(pinned.prefetch_buffer, output)
+
+
+def _compute_ple_ngram_ids_fake(
+    input_ids: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    ngram_context: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: str,
+) -> None:
+    return None
+
+
+def _ple_device_lookup_fake(
+    ngram_ids: torch.Tensor, output: torch.Tensor, layer_name: str
+) -> None:
+    return None
+
+
+def _ple_start_prefetch_fake(
+    hidden_states: torch.Tensor,
+    input_ids: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    ngram_context: torch.Tensor,
+    prefetch_buffer: torch.Tensor,
+    layer_name: str,
+) -> None:
+    return None
+
+
+def _ple_finalize_prefetch_fake(output: torch.Tensor, layer_name: str) -> None:
+    return None
+
+
+direct_register_custom_op(
+    op_name="qwen4_exp_compute_ple_ngram_ids",
+    op_func=qwen4_exp_compute_ple_ngram_ids,
+    mutates_args=["output"],
+    fake_impl=_compute_ple_ngram_ids_fake,
+)
+direct_register_custom_op(
+    op_name="qwen4_exp_ple_device_lookup",
+    op_func=qwen4_exp_ple_device_lookup,
+    mutates_args=["output"],
+    fake_impl=_ple_device_lookup_fake,
+)
+direct_register_custom_op(
+    op_name="qwen4_exp_ple_start_prefetch",
+    op_func=qwen4_exp_ple_start_prefetch,
+    mutates_args=["prefetch_buffer"],
+    fake_impl=_ple_start_prefetch_fake,
+)
+direct_register_custom_op(
+    op_name="qwen4_exp_ple_finalize_prefetch",
+    op_func=qwen4_exp_ple_finalize_prefetch,
+    mutates_args=["output"],
+    fake_impl=_ple_finalize_prefetch_fake,
+)
+
+
 __all__ = [
     "Qwen4ExpNGramEmbedding",
 ]

vllm/models/qwen4_exp/nvidia/ple_layer.py

@@ -19,6 +19,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
     MambaStateShapeCalculator,
     is_conv_state_dim_first,
 )
+from vllm.utils.torch_utils import direct_register_custom_op
 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
 from vllm.v1.attention.backends.short_conv_attn import (
     PleShortConvAttentionBackend,
@@ -95,6 +96,10 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
         if envs.VLLM_E2_PLE_PULL_TRANSPORT:
             # Fork: the E2 pull transport owns the table in a CPU-only
             # offload process; upstream's device/UVA backends are never built.
+            # Pull transport with the compile-profile guard: no layer_name -- the pull module never routes
+            # through the n-gram custom ops. Its GPU placeholder dispatches
+            # through PleOffloadLayer.forward (the ple_offload_wait_copy op) and
+            # owns no table; the offload process calls forward_impl directly.
             from .ple_pull import Qwen4ExpNGramEmbeddingPull
 
             self.ple_embedding = Qwen4ExpNGramEmbeddingPull.build(
@@ -115,6 +120,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
                 vllm_config.scheduler_config.max_num_batched_tokens,
                 data_parallel_rank=vllm_config.parallel_config.data_parallel_rank,
                 prefix=f"{prefix}.ple_embedding",
+                layer_name=prefix,
                 quant_config=quant_config,
                 params_dtype=model_config.dtype,
             )
@@ -161,11 +167,12 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
         output_dtype: torch.dtype,
     ) -> torch.Tensor:
         """Dequantize PLE lookup output."""
-        return self.ple_embedding.ngram_embedding.dequantize(
+
         dequantize_pulled = getattr(self.ple_embedding, "dequantize_pulled", None)
         if dequantize_pulled is not None:
             # Fork pull transport: no local table; FP8 scale kept here.
             return dequantize_pulled(embeddings, output_dtype)
+        return self.ple_embedding.ngram_embedding.dequantize(
             embeddings,
             output_dtype,
         )
@@ -450,11 +457,44 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
             self.norm_key.eps,
         )
         # State routing depends on runtime request metadata and remains outside
-        # the graph; short convolution accumulates into gated_output.
-        self._short_conv(conv_input, gated_output, hidden_states)
+        # the graph; short convolution accumulates into gated_output. Compile-profile guard:
+        # a splitting op again (as on v2.0.1), so VLLM_COMPILE never traces it.
+        torch.ops.vllm.qwen4_exp_ple_short_conv(
+            conv_input, gated_output, hidden_states, self.prefix
+        )
         return gated_output
 
 
+def qwen4_exp_ple_short_conv(
+    inputs: torch.Tensor,
+    residual_output: torch.Tensor,
+    outer_residual: torch.Tensor,
+    layer_name: str,
+) -> None:
+    """PLE short convolution + state routing, outside piecewise graphs."""
+    layer = get_forward_context().no_compile_layers[layer_name]
+    if not isinstance(layer, Qwen4ExpPLELayer):
+        raise TypeError(f"{layer_name} is not a Qwen4Exp PLE layer")
+    layer._short_conv(inputs, residual_output, outer_residual)
+
+
+def _qwen4_exp_ple_short_conv_fake(
+    inputs: torch.Tensor,
+    residual_output: torch.Tensor,
+    outer_residual: torch.Tensor,
+    layer_name: str,
+) -> None:
+    return None
+
+
+direct_register_custom_op(
+    op_name="qwen4_exp_ple_short_conv",
+    op_func=qwen4_exp_ple_short_conv,
+    mutates_args=["residual_output"],
+    fake_impl=_qwen4_exp_ple_short_conv_fake,
+)
+
+
 __all__ = [
     "Qwen4ExpPLEGroupedNorm",
     "Qwen4ExpPLELayer",

vllm/v1/e2_guards.py

@@ -3,33 +3,30 @@
 """Fork startup guards that fail closed on a graph profile we never qualified.
 
 These run at model construction, on every rank, for the target model and for
-the MTP drafter, before profiling or capture. They exist because a *documented*
-launch flag is not a mechanism: the first reference-host run of a rebase will otherwise
-qualify a compile/capture identity nobody chose
-(see below).
+the MTP drafter, before profiling or capture.
 
-What this module does NOT do: choose a graph profile. That is the maintainer's
-call and it is open (the graph-profile choice). What it does is refuse the
-two states that are silently wrong on vLLM v0.30.0 --
+Compile-profile guard: this tree restores torch.compile on the
+NVIDIA Qwen4Exp path (reverting vLLM #55272) and removes Qwen4Exp from
+``DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES``. The ONE qualified profile is the
+v2.0.1 production identity:
 
-  * ``VLLM_USE_BREAKABLE_CUDAGRAPH`` left unset, which lets
-    ``VllmConfig._maybe_enable_breakable_cudagraph`` turn it on for this fork's
-    architectures and force ``CompilationMode.NONE``; and
-  * a PIECEWISE-bearing cudagraph mode combined with ``VLLM_COMPILE`` and
-    breakable-off, which upstream accepts at config time and then rejects at
-    capture, because #55272 removed every ``@support_torch_compile`` wrapper
-    from the NVIDIA Qwen4Exp model and MTP.
+  * ``VLLM_USE_BREAKABLE_CUDAGRAPH`` explicitly pinned to exactly ``0``, and
+    resolved OFF;
+  * ``CompilationMode.VLLM_COMPILE``;
+  * ``CUDAGraphMode.FULL_AND_PIECEWISE``;
+  * every Qwen4Exp request-dependent custom op in ``splitting_ops``
+    (:data:`REQUIRED_SPLITTING_OPS`);
+  * a live ``@support_torch_compile`` submodule
+    (:func:`assert_compiled_submodule`), which is the property
+    ``vllm/v1/worker/gpu/cudagraph_utils.py:605`` actually tests at capture.
 
--- and to make the resolved profile visible and assertable.
+Anything else -- eager, NONE + FULL, FULL_DECODE_ONLY, breakable PIECEWISE --
+fails loud. ``VLLM_E2_ALLOW_BREAKABLE=1`` is the single, logged opt-out for
+deliberate non-production boots (debug, probes); a production launch must not set it.
 
-Timing limit, stated honestly: this runs at model construction, i.e. AFTER
-``VllmConfig.__post_init__`` has already called
-``_maybe_enable_breakable_cudagraph`` (``config/vllm.py:1561``) and
-``_normalize_piecewise_cudagraph_mode`` (``:1673``, ``:1844``). It is early
-enough to stop profiling and capture; it cannot un-mutate ``os.environ`` or
-recover the originally requested cudagraph mode. Requiring the pin to be
-*explicitly present* is the part of "validate before auto-enable" that is
-checkable from here: absence is exactly the state the auto-enable feeds on.
+Timing limit: this runs AFTER ``VllmConfig.__post_init__`` resolved the
+profile. It stops the boot before profiling and capture; it cannot recover
+the originally requested profile if the config layer rewrote it.
 """
 
 import os
@@ -42,14 +39,32 @@ BREAKABLE_ENV = "VLLM_USE_BREAKABLE_CUDAGRAPH"
 #: Opt-out for every guard in this module. Recorded in the launch identity and
 #: PROHIBITED in production profiles; a production launch must not set it.
 ALLOW_BREAKABLE_ENV = "VLLM_E2_ALLOW_BREAKABLE"
-#: Optional expected resolved cudagraph mode, by name (e.g. "FULL",
-#: "FULL_DECODE_ONLY"). When set, a mismatch is an error. This is how a launch
-#: profile pins the graph identity once the maintainer has chosen one.
+#: Optional expected resolved cudagraph mode, by name. When set, a mismatch is
+#: an error (in addition to the fixed requirement below).
 EXPECTED_MODE_ENV = "VLLM_E2_EXPECTED_CUDAGRAPH_MODE"
 
-#: The architectures vLLM v0.30.0 auto-enables breakable cudagraphs for, out of
-#: ``vllm/config/vllm.py:DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES``, that are
-#: ours. Recorded so the guard's message can name the upstream mechanism.
+#: The qualified profile (the compile-profile guard).
+REQUIRED_COMPILATION_MODE_NAME = "VLLM_COMPILE"
+REQUIRED_CUDAGRAPH_MODE_NAME = "FULL_AND_PIECEWISE"
+#: The only raw VLLM_USE_BREAKABLE_CUDAGRAPH value accepted as the production
+#: pin. Resolution semantics stay upstream's (_upstream_truth).
+PRODUCTION_BREAKABLE_PIN = "0"
+
+#: Custom ops that must stay graph breaks on the NVIDIA Qwen4Exp path. Each
+#: reads per-step request metadata (or forks/joins a side stream) and would be
+#: baked into a piecewise graph otherwise.
+REQUIRED_SPLITTING_OPS = (
+    "vllm::qwen4_exp_qsa_with_output",
+    "vllm::qwen4_exp_ple_short_conv",
+    "vllm::qwen4_exp_compute_ple_ngram_ids",
+    "vllm::qwen4_exp_ple_start_prefetch",
+    "vllm::qwen4_exp_ple_finalize_prefetch",
+    "vllm::qwen_gdn_attention_core",
+    "vllm::qwen_gdn_attention_core_fused_norm_packed",
+)
+
+#: Upstream v0.30.0 auto-enabled breakable cudagraphs for these; this tree
+#: does not (vllm/config/vllm.py). Kept to name the mechanism in messages.
 OUR_AUTO_BREAKABLE_ARCHITECTURES = frozenset(
     {
         "Qwen4ExpForCausalLM",
@@ -62,10 +77,7 @@ OUR_AUTO_BREAKABLE_ARCHITECTURES = frozenset(
 def _upstream_truth(raw: str | None) -> bool:
     """Upstream's exact semantics for the breakable switch.
 
-    ``envs.py:937`` is ``os.environ.get(..., "0") == "1"``. Nothing else is on:
-    ``"2"``, ``"true"``, ``"FALSE"`` and ``""`` are all off. Using one parser
-    that matches it is the point -- a stricter local parser would reject boots
-    upstream treats as off, which is a false rejection, not extra safety.
+    ``envs.py:937`` is ``os.environ.get(..., "0") == "1"``. Nothing else is on.
     """
     return raw == "1"
 
@@ -73,17 +85,8 @@ def _upstream_truth(raw: str | None) -> bool:
 def breakable_resolved() -> bool:
     """Whether breakable cudagraphs are on, as the engine will see it.
 
-    Two readings of ONE source, not three independent sources:
-    ``vllm.compilation.breakable_cudagraph.is_breakable_cudagraph_enabled()``
-    is literally ``bool(envs.VLLM_USE_BREAKABLE_CUDAGRAPH)``
-    (``breakable_cudagraph.py:53``), so it adds nothing over the envs read.
-    The envs read is kept beside the raw one because
-    ``_maybe_enable_breakable_cudagraph`` mutates ``os.environ`` at config time
-    while the envs accessor is cached, so the two can legitimately disagree;
-    either saying "on" means on.
-
-    No exception is swallowed. A guard that cannot read its own input must fail
-    the boot, not report False.
+    The raw environment and the (cached) envs accessor are two readings of one
+    source; either saying "on" means on. No exception is swallowed.
     """
     if _upstream_truth(os.environ.get(BREAKABLE_ENV)):
         return True
@@ -92,8 +95,12 @@ def breakable_resolved() -> bool:
     return bool(envs.VLLM_USE_BREAKABLE_CUDAGRAPH)
 
 
-def _profile(vllm_config) -> tuple[list[str], object, object]:
-    """Resolved (architectures, compilation mode, cudagraph mode). Fails closed."""
+def _name(value) -> str:
+    return getattr(value, "name", str(value))
+
+
+def _profile(vllm_config) -> tuple[list[str], object, object, list[str] | None]:
+    """Resolved (architectures, mode, cudagraph mode, splitting ops)."""
     model_config = getattr(vllm_config, "model_config", None)
     if model_config is None:
         raise RuntimeError("e2-guard: vllm_config has no model_config to check")
@@ -105,112 +112,242 @@ def _profile(vllm_config) -> tuple[list[str], object, object]:
         architectures,
         compilation_config.mode,
         compilation_config.cudagraph_mode,
+        getattr(compilation_config, "splitting_ops", None),
     )
 
 
+def _opted_out(role: str, detail: str) -> bool:
+    if os.environ.get(ALLOW_BREAKABLE_ENV, "") != "1":
+        return False
+    logger.warning(
+        "e2-guard[%s]: %s=1 -- guards disabled by request (%s). This boot is "
+        "NOT a promotion-grade boot; a production launch must not set this flag.",
+        role,
+        ALLOW_BREAKABLE_ENV,
+        detail,
+    )
+    return True
+
+
 def assert_graph_profile_is_executable(vllm_config, *, role: str = "target") -> None:
-    """Refuse the two silently-wrong startup states on v0.30.0.
+    """Require the qualified profile: VLLM_COMPILE + FULL_AND_PIECEWISE.
 
     ``role`` is ``"target"`` or ``"draft"``; both are checked, because
     ``Qwen4ExpMTP`` builds ``Qwen4ExpMultiTokenPredictor`` rather than
-    ``Qwen4ExpForCausalLM`` (``mtp.py:397``, ``:411``) from a config rebuilt at
-    ``mtp.py:159``, so the drafter's profile can differ from the target's.
-
-    Intentional eager and intentional ``CompilationMode.NONE`` are PERMITTED.
-    On v0.30.0's NVIDIA Qwen4Exp path ``NONE`` is a normal resolved state, not
-    evidence of the auto-enable: #55272 removed the compile wrappers, so
-    ``_normalize_piecewise_cudagraph_mode`` legitimately lands there.
+    ``Qwen4ExpForCausalLM``, from a config rebuilt in ``mtp.py``.
     """
-    from vllm.config import CompilationMode
-    architectures, mode, cudagraph_mode = _profile(vllm_config)
-
-    if os.environ.get(ALLOW_BREAKABLE_ENV, "") == "1":
-        logger.warning(
-            "e2-guard[%s]: %s=1 -- guards disabled by request. architectures=%s "
-            "breakable=%s compilation_mode=%s cudagraph_mode=%s. This boot is "
-            "NOT a promotion-grade boot; a production launch must not set this flag.",
-            role,
-            ALLOW_BREAKABLE_ENV,
-            architectures,
-            os.environ.get(BREAKABLE_ENV),
-            mode,
-            cudagraph_mode,
-        )
+    architectures, mode, cudagraph_mode, splitting_ops = _profile(vllm_config)
+    detail = (
+        f"architectures={architectures} breakable={os.environ.get(BREAKABLE_ENV)} "
+        f"compilation_mode={_name(mode)} cudagraph_mode={_name(cudagraph_mode)}"
+    )
+    if _opted_out(role, detail):
         return
 
-    # 1. The pin must be explicit. Absence is the state _maybe_enable_breakable_
-    #    cudagraph feeds on (config/vllm.py:758: "not in os.environ").
-    if BREAKABLE_ENV not in os.environ:
+    # 1. The pin must be explicit, so a future upstream merge that re-adds us
+    #    to DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES cannot flip it silently.
+    raw_breakable = os.environ.get(BREAKABLE_ENV)
+    if raw_breakable is None:
         raise RuntimeError(
-            f"e2-guard[{role}]: {BREAKABLE_ENV} is not set. vLLM v0.30.0 lists "
-            f"{sorted(OUR_AUTO_BREAKABLE_ARCHITECTURES)} in "
-            "DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES and turns the switch on "
-            "by itself when it is unset (vllm/config/vllm.py:756-774), forcing "
-            "CompilationMode.NONE with it. Pin it in the launcher: "
-            f"export {BREAKABLE_ENV}=0 (or =1 deliberately), or set "
-            f"{ALLOW_BREAKABLE_ENV}=1 to take an unpinned boot on purpose."
+            f"e2-guard[{role}]: {BREAKABLE_ENV} is not set. Pin it in the "
+            f"launcher: export {BREAKABLE_ENV}=0, or set {ALLOW_BREAKABLE_ENV}=1 "
+            "to take an unpinned boot on purpose."
         )
 
-    breakable = breakable_resolved()
-
-    # 2. Breakable on is an identity this fork has never qualified.
-    if breakable:
+    # 2. Breakable capture is not the qualified identity.
+    if raw_breakable == "1" or breakable_resolved():
         raise RuntimeError(
-            f"e2-guard[{role}]: {BREAKABLE_ENV} resolved ON for "
-            f"{architectures or '<unknown architecture>'} "
-            f"(compilation_mode={mode}, cudagraph_mode={cudagraph_mode}). That "
-            "is not the identity the v2.0.1 results were measured on. Export "
-            f"{BREAKABLE_ENV}=0, or set {ALLOW_BREAKABLE_ENV}=1 to take it "
-            "deliberately."
+            f"e2-guard[{role}]: {BREAKABLE_ENV} resolved ON ({detail}). The "
+            "qualified profile is breakable-off VLLM_COMPILE + "
+            f"FULL_AND_PIECEWISE. Export {BREAKABLE_ENV}=0, or set "
+            f"{ALLOW_BREAKABLE_ENV}=1 to take it deliberately."
         )
 
-    # 3. PIECEWISE with breakable off needs a torch-compiled submodule, and
-    #    #55272 removed every @support_torch_compile wrapper from
-    #    vllm/models/qwen4_exp/nvidia/*. Upstream only rewrites the mode when
-    #    compilation mode is NOT VLLM_COMPILE
-    #    (config/vllm.py:1010-1030), so VLLM_COMPILE + PIECEWISE + breakable-off
-    #    survives config time and then dies at
-    #    gpu/cudagraph_utils.py:605 in the middle of capture. Fail here instead,
-    #    with the candidates named.
+    # 2b. The production pin is the exact string "0".
+    #     Upstream's accessor reads anything but "1" as off, so "2", "true",
+    #     "" or "00" would resolve off -- but a mistyped launcher value must
+    #     not be recorded as a qualified breakable-off profile. The logged
+    #     debug opt-out (VLLM_E2_ALLOW_BREAKABLE=1, handled above) is separate.
+    if raw_breakable != PRODUCTION_BREAKABLE_PIN:
+        raise RuntimeError(
+            f"e2-guard[{role}]: {BREAKABLE_ENV}={raw_breakable!r} is not the "
+            f"production pin; it must be exactly {PRODUCTION_BREAKABLE_PIN!r}. "
+            "Upstream would read it as off, but the launch profile does not "
+            f"say what was meant. Export {BREAKABLE_ENV}=0, or set "
+            f"{ALLOW_BREAKABLE_ENV}=1 for a deliberate non-production boot."
+        )
+
+    # 3. The profile itself: mode AND cudagraph mode, by name.
     if (
-        cudagraph_mode.requires_piecewise_compilation()
-        and mode is CompilationMode.VLLM_COMPILE
+        _name(mode) != REQUIRED_COMPILATION_MODE_NAME
+        or _name(cudagraph_mode) != REQUIRED_CUDAGRAPH_MODE_NAME
     ):
         raise RuntimeError(
-            f"e2-guard[{role}]: cudagraph_mode={cudagraph_mode} needs piecewise "
-            "capture, compilation mode is VLLM_COMPILE and breakable cudagraphs "
-            "are off -- but vLLM PR #55272 removed every @support_torch_compile "
-            "wrapper from vllm/models/qwen4_exp/nvidia/, so "
-            "has_compiled_submodule() is False and capture will raise at "
-            "vllm/v1/worker/gpu/cudagraph_utils.py:605. Executable candidates: "
-            "(a) breakable off + -O0/CompilationMode.NONE, which normalises "
-            "FULL_AND_PIECEWISE to FULL_DECODE_ONLY (changed prefill "
-            "behaviour -- must be qualified); (b) breakable off + an explicit "
-            "cudagraph_mode=FULL or NONE; (c) breakable PIECEWISE, i.e. "
-            f"{BREAKABLE_ENV}=1, which is a different identity. See "
-            "the graph-profile choice."
+            f"e2-guard[{role}]: resolved compilation_mode={_name(mode)} "
+            f"cudagraph_mode={_name(cudagraph_mode)}; the qualified profile is "
+            f"compilation_mode={REQUIRED_COMPILATION_MODE_NAME} "
+            f"cudagraph_mode={REQUIRED_CUDAGRAPH_MODE_NAME} (compile-profile guard). Check "
+            "-O / --compilation-config / --enforce-eager / TORCH_COMPILE_DISABLE, "
+            f"or set {ALLOW_BREAKABLE_ENV}=1 for a deliberate non-production boot."
         )
 
-    # 4. Optional exact pin of the resolved mode, for a launch profile.
+    # 4. The graph breaks that make PIECEWISE correct on this model.
+    if splitting_ops is None:
+        raise RuntimeError(
+            f"e2-guard[{role}]: compilation_config.splitting_ops is unresolved; "
+            "cannot confirm the Qwen4Exp graph breaks."
+        )
+    missing = [op for op in REQUIRED_SPLITTING_OPS if op not in splitting_ops]
+    if missing:
+        raise RuntimeError(
+            f"e2-guard[{role}]: splitting_ops is missing {missing}. Each of "
+            "these reads per-step request metadata (or forks/joins a side "
+            "stream) and would be frozen into a piecewise CUDA graph."
+        )
+
+    # 5. Optional exact pin of the resolved mode, for a launch profile.
     expected = os.environ.get(EXPECTED_MODE_ENV, "").strip()
-    if expected:
-        actual = getattr(cudagraph_mode, "name", str(cudagraph_mode))
-        if actual != expected:
-            raise RuntimeError(
-                f"e2-guard[{role}]: resolved cudagraph_mode={actual} but "
-                f"{EXPECTED_MODE_ENV}={expected}. The launch profile and the "
-                "engine disagree about the graph identity."
-            )
+    if expected and _name(cudagraph_mode) != expected:
+        raise RuntimeError(
+            f"e2-guard[{role}]: resolved cudagraph_mode={_name(cudagraph_mode)} "
+            f"but {EXPECTED_MODE_ENV}={expected}. The launch profile and the "
+            "engine disagree about the graph identity."
+        )
 
     logger.info(
         "e2-guard[%s]: architectures=%s breakable=OFF compilation_mode=%s "
-        "cudagraph_mode=%s expected=%s. Recorded, not certified -- this guard "
-        "checks startup states, not that the profile was qualified.",
+        "cudagraph_mode=%s splitting_ops=+%d required. Recorded, not "
+        "certified -- this checks startup state, not qualification.",
         role,
         architectures,
-        mode,
-        cudagraph_mode,
-        expected or "<unset>",
+        _name(mode),
+        _name(cudagraph_mode),
+        len(REQUIRED_SPLITTING_OPS),
+    )
+
+
+def assert_compiled_submodule(module, vllm_config, *, role: str = "target") -> None:
+    """Require a live ``@support_torch_compile`` wrapper on ``module``.
+
+    This is the property ``has_compiled_submodule()`` checks before PIECEWISE
+    capture (``gpu/cudagraph_utils.py:132``, ``:605``). Checking it at
+    construction turns a mid-capture failure into a startup one.
+    """
+    from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
+
+    compiled = isinstance(module, TorchCompileWithNoGuardsWrapper) and not getattr(
+        module, "do_not_compile", True
+    )
+    if compiled:
+        return
+    detail = (
+        f"{type(module).__name__} is_wrapper="
+        f"{isinstance(module, TorchCompileWithNoGuardsWrapper)} "
+        f"do_not_compile={getattr(module, 'do_not_compile', '<absent>')}"
+    )
+    if _opted_out(role, detail):
+        return
+    raise RuntimeError(
+        f"e2-guard[{role}]: {detail}: no live torch.compile wrapper, so "
+        "PIECEWISE capture would raise at gpu/cudagraph_utils.py:605. The "
+        "@support_torch_compile decorator (compile-profile guard) is missing or disabled."
+    )
+
+
+# --------------------------------------------------------------------------
+# Resolved runtime graph modes.
+# --------------------------------------------------------------------------
+
+#: Target architectures whose V2 capture is checked. The construction guard
+#: above lives in the NVIDIA Qwen4Exp model; this tree targets sm_86 only.
+QWEN4EXP_TARGET_ARCHITECTURES = frozenset(
+    {"Qwen4ExpForCausalLM", "Qwen4ExpForConditionalGeneration"}
+)
+
+#: What each V2 cudagraph manager may hold immediately before capture.
+#: The construction guard checks ``compilation_config.cudagraph_mode``; the
+#: V2 runner later calls ``resolve_cudagraph_mode_and_sizes`` (and, with
+#: adaptive verification, ``resolve_adaptive_cudagraph_mode``), which may
+#: downgrade to FULL_DECODE_ONLY / PIECEWISE / NONE or drop the piecewise half
+#: (attention CG support, spec-decode uniform batches, no compiled submodule).
+#: No such downgrade is allowed for the target. The drafter's managers are
+#: derived from the target's resolved mode by the speculator itself:
+#:   * autoregressive / MTP (``AutoRegressiveSpeculator.init_cudagraph_manager``):
+#:     prefill = the target's mode unchanged; decode = FULL_DECODE_ONLY
+#:     ("PIECEWISE cudagraphs are not supported for draft decodes");
+#:   * multi-module MTP (``MultiModuleMTPSpeculator``): one manager,
+#:     FULL_AND_PIECEWISE -> FULL_DECODE_ONLY by design.
+#: Those decode-specific modes are the only drafter deviations accepted.
+EXPECTED_RESOLVED_MODES = {
+    "target": REQUIRED_CUDAGRAPH_MODE_NAME,
+    "draft.prefill_cudagraph_manager": REQUIRED_CUDAGRAPH_MODE_NAME,
+    "draft.decode_cudagraph_manager": "FULL_DECODE_ONLY",
+    "draft.cudagraph_manager": "FULL_DECODE_ONLY",
+}
+_DRAFT_MANAGER_ATTRS = (
+    "prefill_cudagraph_manager",
+    "decode_cudagraph_manager",
+    "cudagraph_manager",
+)
+
+
+def assert_resolved_cudagraph_modes(
+    vllm_config, target_manager, speculator=None
+) -> None:
+    """Fail loud if the V2 runner is about to capture a downgraded mode.
+
+    Called from ``GPUModelRunner.capture_model`` (``vllm/v1/worker/gpu``)
+    before any capture, including the ``profile_only`` capture and the
+    "nothing to capture" early return, so a resolver downgrade to NONE cannot
+    skip it. Applies only to a Qwen4Exp target; any other model is untouched.
+    ``VLLM_E2_ALLOW_BREAKABLE=1`` skips it, logged, like the other graph
+    guards.
+    """
+    model_config = getattr(vllm_config, "model_config", None)
+    architectures = list(getattr(model_config, "architectures", None) or [])
+    if not QWEN4EXP_TARGET_ARCHITECTURES & set(architectures):
+        return
+
+    observed: dict[str, str] = {}
+    if target_manager is None:
+        observed["target"] = "<no manager>"
+    else:
+        observed["target"] = _name(getattr(target_manager, "cudagraph_mode", None))
+    if speculator is not None:
+        found = False
+        for attr in _DRAFT_MANAGER_ATTRS:
+            if not hasattr(speculator, attr):
+                continue
+            found = True
+            manager = getattr(speculator, attr)
+            observed[f"draft.{attr}"] = (
+                "<no manager>"
+                if manager is None
+                else _name(getattr(manager, "cudagraph_mode", None))
+            )
+        if not found:
+            observed["draft"] = f"<unrecognised speculator {type(speculator).__name__}>"
+
+    detail = " ".join(f"{key}={value}" for key, value in observed.items())
+    if _opted_out("runtime", detail):
+        return
+    wrong = {
+        key: value
+        for key, value in observed.items()
+        if EXPECTED_RESOLVED_MODES.get(key) != value
+    }
+    if wrong:
+        expected = {key: EXPECTED_RESOLVED_MODES.get(key, "<none>") for key in wrong}
+        raise RuntimeError(
+            f"e2-guard[runtime]: resolved cudagraph modes {wrong} differ from "
+            f"the qualified profile {expected} ({detail}). The V2 runner's "
+            "resolver downgraded the graph identity after the construction "
+            "guard passed (attention CG support, spec-decode uniform batch, "
+            "or no compiled submodule); see the resolver's warning above. Set "
+            f"{ALLOW_BREAKABLE_ENV}=1 only for a deliberate non-production boot."
+        )
+    logger.info(
+        "e2-guard[runtime]: resolved cudagraph modes before capture: %s.", detail
     )
 
 

vllm/v1/worker/gpu/model_runner.py

@@ -64,6 +64,7 @@ from vllm.utils.gc_utils import freeze_gc_for_cudagraph_capture
 from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE, async_tensor_h2d
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
+from vllm.v1.e2_guards import assert_resolved_cudagraph_modes
 from vllm.v1.kv_cache_interface import (
     KVCacheConfig,
     MambaSpec,
@@ -1106,6 +1107,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
     @torch.inference_mode()
     def capture_model(self, *, profile_only: bool = False) -> int:
         assert self.cudagraph_manager is not None
+        # Fork (0.30 rebase): the construction guard saw the configured
+        # mode; this asserts what the resolver actually handed the target and
+        # drafter managers, before any capture and before the early return.
+        assert_resolved_cudagraph_modes(
+            self.vllm_config, self.cudagraph_manager, self.speculator
+        )
         capture_encoder = (
             self.model_state.supports_mm_inputs
             and self.model_state.encoder_runner.has_cudagraph()