From e017f7492f2b38b119cf7b5bb3a0d78259220652 Mon Sep 17 00:00:00 2001 From: xytpai Date: Tue, 4 Aug 2026 06:04:33 -0500 Subject: [PATCH 1/2] fix(plugin): gate dual-stream MoE by runtime graph mode Use the frontend's per-forward graph decision so FULL decode can retain overlap while PIECEWISE capture remains safely single-stream. Co-authored-by: Cursor --- atom/model_ops/module_dispatch_ops.py | 14 +++---- atom/utils/forward_context.py | 53 ++++++++++++++++++++++++++- 2 files changed, 59 insertions(+), 8 deletions(-) diff --git a/atom/model_ops/module_dispatch_ops.py b/atom/model_ops/module_dispatch_ops.py index a1141bb6ae..29dd9ba37f 100644 --- a/atom/model_ops/module_dispatch_ops.py +++ b/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 ( diff --git a/atom/utils/forward_context.py b/atom/utils/forward_context.py index b5126f6398..8361a063bf 100644 --- a/atom/utils/forward_context.py +++ b/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 CUDAGraphMode, Config, 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, From 670a25d33b2672b01bbdbca850c033127d9f1e9b Mon Sep 17 00:00:00 2001 From: xytpai Date: Tue, 4 Aug 2026 06:10:55 -0500 Subject: [PATCH 2/2] style: sort forward context imports Match the repository's Ruff import ordering rules. Co-authored-by: Cursor --- atom/utils/forward_context.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/atom/utils/forward_context.py b/atom/utils/forward_context.py index 8361a063bf..18c9596ff9 100644 --- a/atom/utils/forward_context.py +++ b/atom/utils/forward_context.py @@ -11,7 +11,7 @@ import numpy as np import torch -from atom.config import CUDAGraphMode, Config, KVCacheTensor, ParallelConfig +from atom.config import Config, CUDAGraphMode, KVCacheTensor, ParallelConfig class AttnState(Enum):