From 3575cdb1264646afd15f53620d148cf6fca6f3cd Mon Sep 17 00:00:00 2001 From: MOVZX Date: Wed, 23 Sep 2026 11:11:01 -0400 Subject: [PATCH] [Guards] Draft identity and draft PLE guard report the drafter's own state Co-Authored-By: Hinata Signed-off-by: MOVZX --- vllm/v1/e2_guards.py | 37 ++++++++- vllm/v1/worker/e2_identity.py | 150 ++++++++++++++++++++++++++++++---- vllm/v1/worker/gpu_worker.py | 11 ++- 3 files changed, 177 insertions(+), 21 deletions(-) diff --git a/vllm/v1/e2_guards.py b/vllm/v1/e2_guards.py index 7624ae1d96..947dadab0d 100644 --- a/vllm/v1/e2_guards.py +++ b/vllm/v1/e2_guards.py @@ -484,6 +484,38 @@ def resolve_ple_mode(vllm_config) -> str: return PLE_MODE_UVA if cpu_offload else PLE_MODE_DEVICE +def _role_ple_layer_ids(vllm_config, role: str) -> list[int]: + """The 1-based PLE layer ids the role's own decoder layers carry. + + ``role="target"``: ``model_config.hf_text_config.ple_layer_ids``. + ``role="draft"``: the drafter is constructed with the target's VllmConfig + (the trap an earlier commit fixed for the architectures), so the target's + model_config names the TARGET's PLE layers. The drafter's decoder layers + are ``num_hidden_layers + i`` for ``i < mtp_num_hidden_layers`` + (``Qwen4ExpMultiTokenPredictor``), and a layer builds a PLE block only if + ``layer_idx + 1`` is in ``ple_layer_ids`` (``Qwen4ExpDecoderLayer``). The + draft line reports exactly that intersection, read from + ``speculative_config.draft_model_config`` -- empty for this checkpoint, + whose ids must name backbone layers (``Qwen4ExpTextConfig``). + """ + model_config = getattr(vllm_config, "model_config", None) + if role == "draft": + spec = getattr(vllm_config, "speculative_config", None) + draft_config = getattr(spec, "draft_model_config", None) + if draft_config is not None: + model_config = draft_config + text_config = getattr(model_config, "hf_text_config", None) + ids = list(getattr(text_config, "ple_layer_ids", None) or []) + if role != "draft": + return ids + num_hidden_layers = getattr(text_config, "num_hidden_layers", None) + if num_hidden_layers is None: + return ids + num_mtp = int(getattr(text_config, "mtp_num_hidden_layers", 1)) + own = range(int(num_hidden_layers) + 1, int(num_hidden_layers) + num_mtp + 1) + return [layer_id for layer_id in ids if layer_id in own] + + def assert_ple_mode(vllm_config, *, role: str = "target") -> str: """Startup guard: resolve, assert and log the effective PLE mode. @@ -492,11 +524,10 @@ def assert_ple_mode(vllm_config, *, role: str = "target") -> str: target's EngramConfig and the same process environment, so a mixed configuration fails there too). NOT disabled by VLLM_E2_ALLOW_BREAKABLE: that opt-out is about graph identity, not about owning the PLE table twice. + The logged ``ple_layer_ids`` are the role's own (:func:`_role_ple_layer_ids`). """ mode = resolve_ple_mode(vllm_config) - model_config = getattr(vllm_config, "model_config", None) - text_config = getattr(model_config, "hf_text_config", None) - ple_layer_ids = list(getattr(text_config, "ple_layer_ids", None) or []) + ple_layer_ids = _role_ple_layer_ids(vllm_config, role) engram_config = getattr(vllm_config, "engram_config", None) logger.info( "e2-guard[%s]: ple_mode=%s pull_transport=%d engram_cpu_offload=%s " diff --git a/vllm/v1/worker/e2_identity.py b/vllm/v1/worker/e2_identity.py index e457b09d74..e2a9527dc7 100644 --- a/vllm/v1/worker/e2_identity.py +++ b/vllm/v1/worker/e2_identity.py @@ -9,8 +9,10 @@ GPU worker rank logs one line per model role it holds:: ple_weight_dtype= breakable= compile_mode= cudagraph_mode= role= -(one line in the log). ``role=draft`` is logged only where a draft model is -loaded (the last PP stage). Every value is read from the live objects: +(one line in the log). ``role=target`` is logged after model loading; +``role=draft`` is logged only where a draft model is loaded (the last PP +stage), once the drafter's graph modes are resolved (after the KV cache is +initialised -- see below). Every value is read from the live objects: * ``ple_backend``: the class of each PLE layer's ``ple_embedding`` module, and the transport that serves its rows on this rank, read from the placeholder's @@ -25,8 +27,28 @@ loaded (the last PP stage). Every value is read from the live objects: (``get_offload_output_dtype``). * ``breakable``: ``is_breakable_cudagraph_enabled()``, the switch the capture path itself reads. -* ``compile_mode`` / ``cudagraph_mode``: the role's ``CompilationConfig`` - (config level, i.e. before KV-cache-time capture-mode resolution). +* ``compile_mode`` / ``cudagraph_mode`` for ``role=target``: the engine's + ``CompilationConfig`` (config level, i.e. before KV-cache-time capture-mode + resolution). +* ``compile_mode`` / ``cudagraph_mode`` for ``role=draft``: the DRAFTER's own + resolved state, not the config. The drafter is built with the target's + VllmConfig and there is one ``CompilationConfig`` in the engine, so the + config fields name the target's profile (the trap an earlier commit fixed for + the draft guard's architectures). What the drafter actually runs is: + + - ``compile_mode``: the mode of the draft model's live + ``@support_torch_compile`` wrapper(s) (``do_not_compile`` false -- the + property the capture path tests, ``gpu/cudagraph_utils.py`` + ``has_compiled_submodule``); ``NONE`` when no wrapper is live. + - ``cudagraph_mode``: the mode each speculator cudagraph manager holds after + ``init_cudagraph_manager`` (V2 runner): ``prefill:,decode:`` for the + autoregressive / MTP speculator, the bare mode for a single-manager + speculator; for a V1 drafter, its dispatcher's mode. These are the values + ``e2-guard[runtime]`` checks before capture. ``unresolved`` if a manager + does not exist yet. + + Hence the draft line is logged from ``Worker.initialize_from_config`` after + ``initialize_kv_cache``, not at load time. A role with no PLE layer reports ``none`` for the three PLE fields; several distinct values on one rank are joined with ``,``. @@ -102,46 +124,142 @@ def _rank_tag() -> str: def _draft(model_runner: Any) -> tuple[nn.Module | None, Any]: - """The loaded draft model and its config (V2 speculator or V1 drafter).""" + """The loaded draft model and its holder (V2 speculator or V1 drafter).""" for name in ("speculator", "drafter"): holder = getattr(model_runner, name, None) model = getattr(holder, "model", None) if isinstance(model, nn.Module): - return model, getattr(holder, "vllm_config", None) + return model, holder return None, None +#: Written where a draft graph mode has no resolved source yet. +UNRESOLVED = "unresolved" +#: Speculator cudagraph-manager attributes (the set ``e2-guard[runtime]`` +#: reads), with the phase label each carries on the draft line. +_DRAFT_MANAGER_PHASES = ( + ("prefill_cudagraph_manager", "prefill"), + ("decode_cudagraph_manager", "decode"), + ("cudagraph_manager", None), +) + + +def _draft_compile_mode(model: nn.Module) -> str: + """Compile mode of the draft model's live ``@support_torch_compile`` wrappers. + + ``NONE`` when no wrapper is live: the drafter then runs eager, whatever + the shared ``CompilationConfig.mode`` says. + """ + from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper + + modes = { + _enum_name(getattr(getattr(m, "compilation_config", None), "mode", None)) + for m in model.modules() + if isinstance(m, TorchCompileWithNoGuardsWrapper) + and not getattr(m, "do_not_compile", True) + } + return _join(modes) if modes else "NONE" + + +def _draft_cudagraph_mode(holder: Any) -> str: + """The drafter's resolved cudagraph mode(s), read from its managers.""" + parts = [] + for attr, phase in _DRAFT_MANAGER_PHASES: + if not hasattr(holder, attr): + continue + manager = getattr(holder, attr) + mode = ( + UNRESOLVED + if manager is None + else _enum_name(getattr(manager, "cudagraph_mode", None)) + ) + parts.append(mode if phase is None else f"{phase}:{mode}") + if parts: + return ",".join(parts) + # V1 drafter: one dispatcher, keyed in _check_and_update_cudagraph_mode. + dispatcher = getattr(holder, "cudagraph_dispatcher", None) + if dispatcher is not None and getattr(dispatcher, "keys_initialized", False): + return _enum_name(getattr(dispatcher, "cudagraph_mode", None)) + return UNRESOLVED + + def identity_line( - tag: str, model: nn.Module | None, vllm_config: Any, role: str + tag: str, + model: nn.Module | None, + vllm_config: Any, + role: str, + *, + compile_mode: str | None = None, + cudagraph_mode: str | None = None, ) -> str: + """One identity line. ``compile_mode`` / ``cudagraph_mode`` override the + config-level values (the draft role passes its resolved state).""" from vllm.compilation.breakable_cudagraph import is_breakable_cudagraph_enabled embeddings = _ple_embeddings(model) backends = {f"{type(e).__name__}@{_transport(e)}" for e in embeddings} compilation = getattr(vllm_config, "compilation_config", None) + if compile_mode is None: + compile_mode = _enum_name(getattr(compilation, "mode", None)) + if cudagraph_mode is None: + cudagraph_mode = _enum_name(getattr(compilation, "cudagraph_mode", None)) return ( f"e2-identity[{tag}]: " f"ple_backend={_join(backends)} " f"ple_method={_join({_method(e) for e in embeddings})} " f"ple_weight_dtype={_join({_weight_dtype(e) for e in embeddings})} " f"breakable={'ON' if is_breakable_cudagraph_enabled() else 'OFF'} " - f"compile_mode={_enum_name(getattr(compilation, 'mode', None))} " - f"cudagraph_mode={_enum_name(getattr(compilation, 'cudagraph_mode', None))} " + f"compile_mode={compile_mode} " + f"cudagraph_mode={cudagraph_mode} " f"role={role}" ) +def draft_identity_line(tag: str, draft_model: nn.Module, holder: Any) -> str: + """The ``role=draft`` line: PLE fields from the draft modules, graph fields + from the drafter's resolved state (never the shared CompilationConfig).""" + return identity_line( + tag, + draft_model, + getattr(holder, "vllm_config", None), + "draft", + compile_mode=_draft_compile_mode(draft_model), + cudagraph_mode=_draft_cudagraph_mode(holder), + ) + + def log_launch_identity(model_runner: Any, vllm_config: Any) -> None: - """Log the identity line(s) for this rank. Never raises.""" + """Log this rank's ``role=target`` line after model load. Never raises. + + The ``role=draft`` line comes later, from :func:`log_draft_identity`, once + the drafter's graph modes exist. + """ tag = "PP?_TP?" try: tag = _rank_tag() - roles = [("target", model_runner.get_model(), vllm_config)] - draft_model, draft_config = _draft(model_runner) - if draft_model is not None: - roles.append(("draft", draft_model, draft_config or vllm_config)) - for role, model, config in roles: - logger.info("%s", identity_line(tag, model, config, role)) + logger.info( + "%s", identity_line(tag, model_runner.get_model(), vllm_config, "target") + ) + except Exception as error: # log-only: identity must never break a boot + logger.warning( + "e2-identity[%s]: unavailable (%s: %s)", tag, type(error).__name__, error + ) + + +def log_draft_identity(model_runner: Any) -> None: + """Log this rank's ``role=draft`` line, if it holds a drafter. Never raises. + + Call after ``initialize_kv_cache``: that is where the V2 runner resolves + the target's mode and ``init_cudagraph_manager`` derives the drafter's + (and where the V1 runner keys the drafter's dispatcher). + """ + tag = "PP?_TP?" + try: + draft_model, holder = _draft(model_runner) + if draft_model is None: + return + tag = _rank_tag() + logger.info("%s", draft_identity_line(tag, draft_model, holder)) except Exception as error: # log-only: identity must never break a boot logger.warning( "e2-identity[%s]: unavailable (%s: %s)", tag, type(error).__name__, error diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 44732ced1f..c0ca5aff61 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -740,8 +740,9 @@ class Worker(WorkerBase): model_runner = cast("GPUModelRunnerV2", self.model_runner) model_runner._setup_ple_offload(self.parallel_config._ple_offload_ipc_path) - # Log-only: this rank's PLE / capture identity for target and drafter, - # read after the offload placeholders are armed (v1/worker/e2_identity.py). + # Log-only: this rank's PLE / capture identity for the target, read + # after the offload placeholders are armed (v1/worker/e2_identity.py). + # The drafter's line follows initialize_kv_cache (initialize_from_config). from vllm.v1.worker.e2_identity import log_launch_identity log_launch_identity(self.model_runner, self.vllm_config) @@ -1034,6 +1035,12 @@ class Worker(WorkerBase): kv_cache_allocation_context=mem_pool_context, ) + # Log-only: the drafter's identity line, now that initialize_kv_cache + # has resolved its cudagraph managers (v1/worker/e2_identity.py). + from vllm.v1.worker.e2_identity import log_draft_identity + + log_draft_identity(self.model_runner) + # Build KV-zero metadata outside the CuMem pool so the bookkeeping # GPU tensors (seg_addrs, block-id buffers) use the standard PyTorch # allocator and are not discarded during sleep/wake cycles.