From 2e68aba2a7b7f35b76a6bab931e863a7871ef92f Mon Sep 17 00:00:00 2001 From: MOVZX Date: Wed, 23 Sep 2026 04:45:12 -0400 Subject: [PATCH] [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 Signed-off-by: MOVZX --- vllm/config/compilation.py | 7 +- vllm/config/vllm.py | 9 +- vllm/models/qwen4_exp/nvidia/model.py | 40 +- vllm/models/qwen4_exp/nvidia/mtp.py | 20 +- .../qwen4_exp/nvidia/ngram_embedding.py | 215 +++++++++- vllm/models/qwen4_exp/nvidia/ple_layer.py | 46 ++- vllm/v1/e2_guards.py | 387 ++++++++++++------ vllm/v1/worker/gpu/model_runner.py | 7 + 8 files changed, 586 insertions(+), 145 deletions(-) diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 3a492843fe..4e87161ed9 100644 --- a/vllm/config/compilation.py +++ b/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", diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 52868756b0..ff7319dc15 100644 --- a/vllm/config/vllm.py +++ b/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). } ) diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py index 255aee6ae5..d6e6bc11ca 100644 --- a/vllm/models/qwen4_exp/nvidia/model.py +++ b/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: diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index d4f29257a9..4109867067 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/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: diff --git a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py b/vllm/models/qwen4_exp/nvidia/ngram_embedding.py index 15b8f96f80..0187a05590 100644 --- a/vllm/models/qwen4_exp/nvidia/ngram_embedding.py +++ b/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", ] diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py index b48c1d9bd2..35fdaf794b 100644 --- a/vllm/models/qwen4_exp/nvidia/ple_layer.py +++ b/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", diff --git a/vllm/v1/e2_guards.py b/vllm/v1/e2_guards.py index 12a0d4864e..1917885117 100644 --- a/vllm/v1/e2_guards.py +++ b/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 ''} " - 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 "", + _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', '')}" + ) + 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"] = "" + 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}"] = ( + "" + if manager is None + else _name(getattr(manager, "cudagraph_mode", None)) + ) + if not found: + observed["draft"] = f"" + + 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, "") 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 ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index fbb630243f..187c6925b8 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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()