[Guards] Draft identity and draft PLE guard report the drafter's own state

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

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 "

vllm/v1/worker/e2_identity.py

@@ -9,8 +9,10 @@ GPU worker rank logs one line per model role it holds::
     ple_weight_dtype=<dtype> breakable=<ON|OFF> compile_mode=<mode>
     cudagraph_mode=<mode> role=<target|draft>
 
-(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:<m>,decode:<m>`` 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

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.