fix(plugin): gate dual-stream MoE by runtime graph mode

atom/model_ops/module_dispatch_ops.py

@@ -25,9 +25,10 @@
 
 import torch
 
-from atom.config import get_current_atom_config
+from atom.config import CUDAGraphMode, get_current_atom_config
 from atom.utils import envs
 from atom.utils.custom_register import direct_register_custom_op
+from atom.utils.forward_context import get_current_cudagraph_runtime_mode
 
 # ---------------------------------------------------------------------------
 # Dual-stream MoE dispatch (V2 / V3.2 / V4)
@@ -54,13 +55,12 @@ def maybe_dual_stream_forward(
     # Under TBO the two micro-batches already overlap on separate threads
     from atom.utils.tbo.ubatching import tbo_active
 
-    # PIECEWISE cudagraph only: dual_stream_moe_forward forks work onto
-    # `alt_stream` and does a caching-allocator alloc there; under PIECEWISE
-    # per-piece capture close the dual stream
-    compilation_config = get_current_atom_config().compilation_config
-    cudagraph_mode = getattr(compilation_config, "cudagraph_mode", None)
+    # Graph ownership belongs to the active frontend.  Only a concrete
+    # PIECEWISE runtime decision is unsafe here: per-piece capture closes over
+    # the main stream while this forward forks work onto `alt_stream`.  Eager
+    # NONE and whole-model FULL capture both support the fork/join topology.
     is_piecewise_cudagraph = (
-        cudagraph_mode is not None and cudagraph_mode.requires_piecewise_compilation()
+        get_current_cudagraph_runtime_mode() == CUDAGraphMode.PIECEWISE
     )
 
     if (

atom/utils/forward_context.py

@@ -11,7 +11,7 @@
 import numpy as np
 import torch
 
-from atom.config import Config, KVCacheTensor, ParallelConfig
+from atom.config import Config, CUDAGraphMode, KVCacheTensor, ParallelConfig
 
 
 class AttnState(Enum):
@@ -600,6 +600,57 @@ def get_forward_context() -> ForwardContext:
     return _forward_context
 
 
+def _normalize_cudagraph_runtime_mode(mode: Any) -> CUDAGraphMode | None:
+    """Normalize a frontend runtime mode to ATOM's concrete enum.
+
+    Frontends own their graph dispatch and therefore use distinct enum
+    classes.  Match by name rather than value so their enum layouts can evolve
+    independently.  Composite configuration modes are deliberately rejected:
+    a forward context must describe the concrete NONE/PIECEWISE/FULL decision
+    for the current batch.
+    """
+    name = mode if isinstance(mode, str) else getattr(mode, "name", None)
+    if name not in {"NONE", "PIECEWISE", "FULL"}:
+        return None
+    return CUDAGraphMode[name]
+
+
+def get_current_cudagraph_runtime_mode() -> CUDAGraphMode:
+    """Return the concrete graph mode for the active model forward.
+
+    In vLLM plugin mode graph capture/replay is owned by vLLM, so its forward
+    context is authoritative.  Native ATOM records the same decision on its
+    own ForwardContext.  An unavailable/unknown context is treated as NONE:
+    eager dual-stream execution is valid, and some vLLM runners expose NONE
+    while a whole-model FULL graph is being captured.  Replay does not execute
+    this Python dispatcher.
+    """
+    from atom.plugin import is_vllm
+
+    if is_vllm():
+        try:
+            from vllm.forward_context import (
+                get_forward_context as get_vllm_forward_context,
+            )
+            from vllm.forward_context import (
+                is_forward_context_available,
+            )
+
+            if is_forward_context_available():
+                mode = _normalize_cudagraph_runtime_mode(
+                    get_vllm_forward_context().cudagraph_runtime_mode
+                )
+                if mode is not None:
+                    return mode
+        except (ImportError, AttributeError, AssertionError):
+            pass
+
+    mode = _normalize_cudagraph_runtime_mode(
+        getattr(get_forward_context(), "cudagraph_runtime_mode", None)
+    )
+    return mode if mode is not None else CUDAGraphMode.NONE
+
+
 def set_forward_context(
     attn_metadata: AttentionMetaData,
     atom_config: Config,