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..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 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,