refactor: one graph mechanism per region, settled DP step shapes, cheaper V4.1 decode prep

.claude/commands/debug-guide.md

@@ -91,7 +91,7 @@
 ## ForwardContext
 
 - **"Forward context is not set":** `set_forward_context()` was not called before model forward. Lifecycle bug — check `ModelRunner` forward path
-- **[COMMON] DP metadata wrong:** `DPMetadata.make()` does CPU `all_reduce` via Gloo. If one rank has 0 tokens, mismatch causes NCCL hang
+- **[COMMON] DP metadata wrong:** `DPMetadata.gather_running_tokens()` does CPU `all_reduce` via Gloo. If one rank has 0 tokens, mismatch causes NCCL hang
 - **Stale context after exception:** If forward throws, `reset_forward_context()` may not be called. Wrap forward in try/finally
 
 ## Speculative Decoding / MTP
@@ -146,7 +146,7 @@
 | Block manager | `model_engine/block_manager.py` | KV cache block allocation |
 | Forward context | `utils/forward_context.py` | Module-level global, attn/DP metadata |
 | Compilation | `utils/backends.py` | VllmBackend, CompilerManager |
-| CUDA graph | `utils/cuda_graph.py` | CUDAGraphWrapper |
+| CUDA graph | `utils/region_graph.py` | TokenGraph, ReqGraph, eager_region |
 | MTP proposer | `spec_decode/eagle.py` | Safe to instrument |
 | MTP model (DeepSeek) | `models/deepseek_mtp.py` | **Compiled — do NOT modify** |
 | MTP model (Qwen3.5) | `models/qwen3_5_mtp.py` | **Compiled — do NOT modify** |

.github/benchmark/models.json

@@ -71,7 +71,7 @@
           "label": "DPA DSpark",
           "suffix": "-dpa-dspark",
           "path": "deepseek-ai/DeepSeek-V4-Pro-DSpark",
-          "extra_args": "--method dspark --num-speculative-tokens 7 --cudagraph-mode AF_PIECEWISE --enable-dp-attention --dspark-config '{\"confidence_schedule\": true, \"ragged\": true, \"ragged_graph_sizes\": \"8\"}'",
+          "extra_args": "--method dspark --num-speculative-tokens 7 --cudagraph-mode TR_HYBRID --enable-dp-attention --dspark-config '{\"confidence_schedule\": true, \"ragged\": true, \"ragged_graph_sizes\": \"8\"}'",
           "bench_args": "--use-chat-template",
           "conc_min": 64,
           "conc_max": 512

.github/benchmark/models_accuracy.json

@@ -51,14 +51,14 @@
   {
     "model_name": "DeepSeek-V4-Pro DSpark",
     "model_path": "deepseek-ai/DeepSeek-V4-Pro-DSpark",
-    "extraArgs": "--kv_cache_dtype fp8 --index_cache_dtype fp4 -tp 8 --method dspark --num-speculative-tokens 7 --cudagraph-mode PIECEWISE --enable-dp-attention --dspark-config '{\"confidence_schedule\": true, \"ragged\": true, \"ragged_graph_sizes\": \"8\"}'",
+    "extraArgs": "--kv_cache_dtype fp8 --index_cache_dtype fp4 -tp 8 --method dspark --num-speculative-tokens 7 --cudagraph-mode TR_HYBRID --enable-dp-attention --dspark-config '{\"confidence_schedule\": true, \"ragged\": true, \"ragged_graph_sizes\": \"8\"}'",
     "env_vars": "AITER_BF16_FP8_MOE_BOUND=0\nATOM_MOE_GU_ITLV=1",
     "runner": "linux-atom-do-mi350x-8",
     "test_level": "pr",
     "accuracy_threshold": 0.93,
     "accuracy_baseline": 0.96,
     "accuracy_baseline_model": "deepseek-ai/DeepSeek-V4-Pro",
-    "_baseline_note": "DSpark spec-decode (7 tokens, dp-attention, PIECEWISE cudagraph) on the DeepSeek-V4-Pro-DSpark checkpoint. Spec-decode is lossless w.r.t. the target, so baseline reuses the DeepSeek-V4-Pro FP8 base (0.96); threshold 0.93 leaves ~3pp headroom for spec-decode / dp-attention run-to-run variance. mtp_accept_threshold intentionally omitted until the first CI run reports the DSpark acceptance rate — add it once measured to guard draft-head regressions."
+    "_baseline_note": "DSpark spec-decode (7 tokens, dp-attention, TR_HYBRID cudagraph) on the DeepSeek-V4-Pro-DSpark checkpoint. Spec-decode is lossless w.r.t. the target, so baseline reuses the DeepSeek-V4-Pro FP8 base (0.96); threshold 0.93 leaves ~3pp headroom for spec-decode / dp-attention run-to-run variance. mtp_accept_threshold intentionally omitted until the first CI run reports the DSpark acceptance rate — add it once measured to guard draft-head regressions."
   },
   {
     "model_name": "DeepSeek-V4.1-Flash DSpark",

CLAUDE.md

@@ -113,7 +113,7 @@ Key entry points:
 
 - **Adding a model**: see `/add-model` for full guide
 - **Model reuse**: DeepSeek V3/V3.2/GLM-5 share `deepseek_v2.py`; MTP models in `deepseek_mtp.py`, `qwen3_next_mtp.py`, and `qwen3_5_mtp.py`
-- **Compilation levels**: `--level` 0=eager, 1=torch.compile, 2=dynamo once, 3=piecewise+CUDAGraph (default)
+- **Compilation levels**: `--level` 0=eager, 1=torch.compile, 2=dynamo once, 3=piecewise compile (default); CUDA graphs are `--cudagraph-mode`'s, independent of the level
 
 ## Dependencies
 

atom/config.py

@@ -12,7 +12,6 @@
 from typing import TYPE_CHECKING, Any, ClassVar
 
 import torch
-from torch.distributed import ProcessGroup, ReduceOp
 from transformers import AutoConfig, GenerationConfig, PretrainedConfig
 
 # plugin-related utilities
@@ -90,36 +89,24 @@ class KVCacheConfig:
 
 
 class CUDAGraphMode(enum.Enum):
-    """Constants for the cudagraph mode in CompilationConfig.
-    Meanwhile, the subset enum `NONE`, `PIECEWISE` and `FULL` are also
-    treated as concrete runtime mode for cudagraph runtime dispatching.
-    """
+    """Constants for the cudagraph mode in CompilationConfig."""
 
     NONE = 0
-    PIECEWISE = 1
+    TR_HYBRID = 1
     FULL = 2
     FULL_DECODE_ONLY = (FULL, NONE)
-    FULL_AND_PIECEWISE = (FULL, PIECEWISE)
-    # AF_PIECEWISE ("attention/FFN-wise"): PIECEWISE + attention core in its own
-    # cudagraph (DSpark). Tuple so decode/mixed_mode resolve to PIECEWISE (existing
-    # == PIECEWISE checks treat it as such); extra capture gated by is_attn_ffn_piecewise().
-    AF_PIECEWISE = (PIECEWISE, PIECEWISE)
+    FULL_AND_TR_HYBRID = (FULL, TR_HYBRID)
 
     def decode_mode(self) -> "CUDAGraphMode":
         return CUDAGraphMode(self.value[0]) if self.separate_routine() else self
 
     def mixed_mode(self) -> "CUDAGraphMode":
         return CUDAGraphMode(self.value[1]) if self.separate_routine() else self
 
-    def is_attn_ffn_piecewise(self) -> bool:
-        """True only for AF_PIECEWISE — gates the extra attention-core cudagraph
-        (attention/FFN-wise capture) on top of the standard piecewise pieces."""
-        return self is CUDAGraphMode.AF_PIECEWISE
-
-    def requires_piecewise_compilation(self) -> bool:
+    def uses_tr_hybrid(self) -> bool:
         return (
-            self.decode_mode() == CUDAGraphMode.PIECEWISE
-            or self.mixed_mode() == CUDAGraphMode.PIECEWISE
+            self.decode_mode() == CUDAGraphMode.TR_HYBRID
+            or self.mixed_mode() == CUDAGraphMode.TR_HYBRID
         )
 
     def max_cudagraph_mode(self) -> "CUDAGraphMode":
@@ -179,41 +166,32 @@ class CompilationConfig:
     The mode of the cudagraph:
 
     - NONE, no cudagraph capture.
-    - PIECEWISE. (v1 default)
+    - TR_HYBRID. (default at level 3)
     - FULL.
     - FULL_DECODE_ONLY.
-    - FULL_AND_PIECEWISE.
-    - AF_PIECEWISE.
-
-    AF_PIECEWISE ("attention/FFN-wise") mode: PIECEWISE where the attention
-    core is ALSO captured into its own cudagraph with zero-copy public buffers
-    (DeepSeek-V4 DSpark), so small-batch decode is all-replay with no eager
-    attention gap between the dense pieces. Falls back to plain PIECEWISE
-    behavior for models that don't implement the attention-core capture.
+    - FULL_AND_TR_HYBRID.
 
-    PIECEWISE mode build piecewise cudagraph only, keeping the cudagraph
-    incompatiable ops (i.e. some attention ops) outside the cudagraph
-    for general flexibility.
-    This is the default mode.
+    TR_HYBRID mode mixes token graphs and request graphs: the compiled forward
+    is recorded as a `TokenGraph` per token bucket, breaking only at the split
+    ops whose launches depend on host data (`eager_region`), which run between
+    its segments -- eagerly, or as a `ReqGraph` where the request bucket fixes
+    their launches.
 
     FULL mode: Capture full cudagraph for all batches. Can be good for small
     models or workloads with small prompts; not supported by many backends.
-    Generally for performance FULL_AND_PIECEWISE is better.
 
     FULL_DECODE_ONLY mode: Capture full cudagraph for decode batches only.
     Mixed prefill-decode batches are run without cudagraphs. Can be good for
     decode instances in a P/D setup where prefill is not as important so we
     can save some memory.
 
-    FULL_AND_PIECEWISE mode: Capture full cudagraph for decode batches and
-    piecewise cudagraph for prefill and mixed prefill-decode batches.
-    This is like the most performant mode for most models.
+    FULL_AND_TR_HYBRID mode: the native runner records it as TR_HYBRID --
+    decode replays token graphs too, not a full graph (`_tr_hybrid_active`).
 
-    Currently, the cudagraph mode is only used for the v1 engine.
     Note that the cudagraph logic is generally orthogonal to the
-    compilation logic. While piecewise cudagraphs require piecewise
-    compilation (level=PIECEWISE and non-empty splitting_ops), full
-    cudagraphs are supported with and without compilation.
+    compilation logic: full cudagraphs and token graphs are recorded with and
+    without compilation; piecewise compilation (level=PIECEWISE and non-empty
+    splitting_ops) only decides where the compiled pieces split.
 
     Warning: This flag is new and subject to change in addition
     more modes may be added.
@@ -233,7 +211,8 @@ class CompilationConfig:
     are always used, it can set this to False. Otherwise, it should
     set this to True, and the compiler will copy the input to an
     internally managed buffer. Default is False.
-    Note that this flag is only effective when cudagraph_mode is PIECEWISE.
+    Note that this flag is only effective at level 3, whenever cudagraph_mode
+    is not NONE (`VllmBackend`).
     """
 
     inductor_compile_config: dict = field(default_factory=dict)
@@ -1096,17 +1075,6 @@ def stateless_init_dp_group(self):
         )
         return dp_group
 
-    @staticmethod
-    def has_unfinished_dp(dp_group: ProcessGroup, has_unfinished: bool) -> bool:
-        tensor = torch.tensor([has_unfinished], dtype=torch.int32, device="cpu")
-        # dp rank 0: has_unfinished_seqs=True
-        # dp rank 1: has_unfinished_seqs=False
-        # aggregated: has_unfinished_seqs=True
-        # so this is an OR operation, i.e. MAX in integers
-        torch.distributed.all_reduce(tensor, op=ReduceOp.MAX, group=dp_group)
-        aggregated_has_unfinished = bool(tensor.item())
-        return aggregated_has_unfinished
-
     def compute_hash(self):
         """
         Provide a hash that uniquely identifies all the configs
@@ -2313,8 +2281,7 @@ def __post_init__(self):
             )
             if qrep_off is not None:
                 logger.warning(
-                    "dcp_config.enable_query_replication disabled: %s not "
-                    "supported.",
+                    "dcp_config.enable_query_replication disabled: %s not supported.",
                     qrep_off,
                 )
                 self.dcp_config.enable_query_replication = False
@@ -2412,14 +2379,14 @@ def __post_init__(self):
             or (self.plugin_config is not None and self.plugin_config.is_vllm)
         ) and self.compilation_config.level == CompilationLevel.PIECEWISE:
             self.compilation_config.set_splitting_ops_for_v1()
-            # Default to PIECEWISE, except V4.1 whose Engram stream fork/join
+            # Default to TR_HYBRID, except V4.1 whose Engram stream fork/join
             # currently requires a FULL graph. Explicit choices are preserved.
             # splitting_ops/sizes still configure the level-3 backend.
             if self.compilation_config.cudagraph_mode is None:
                 self.compilation_config.cudagraph_mode = (
                     CUDAGraphMode.FULL
                     if self.hf_config.model_type == "deepseek_v41_text"
-                    else CUDAGraphMode.PIECEWISE
+                    else CUDAGraphMode.TR_HYBRID
                 )
             self.compilation_config.init_with_cudagraph_sizes()
 

atom/distributed/nccl_watchdog.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""Retiring the NCCL watchdog's pending works before a CUDA graph capture."""
+
+import torch
+import torch.distributed as dist
+from torch.distributed.distributed_c10d import _world
+
+
+def drain_nccl_works() -> None:
+    """Wait until every NCCL group's watchdog has retired the works it tracks.
+
+    The watchdog polls the end event of each eager collective. A capture that
+    records a collective joins its NCCL stream into the capture, and a poll
+    landing meanwhile fails with `hipErrorCapturedEvent` and invalidates the
+    capture. Finishing the device is not enough: the watchdog drops a work only
+    on its next poll. A group with nothing pending returns at once.
+    """
+    if not dist.is_initialized():
+        return
+    torch.cuda.synchronize()
+    for group, (backend, _) in list(_world.pg_map.items()):
+        # Only ProcessGroupNCCL keeps a watchdog list; gloo refuses the call.
+        if backend == "nccl":
+            group._wait_for_pending_works()

atom/distributed/pcp_utils.py

@@ -14,10 +14,10 @@
 """
 
 import logging
-from typing import NamedTuple, Optional
+from typing import NamedTuple
 
+import numpy as np
 import torch
-
 from aiter.dist.parallel_state import (
     get_pcp_group,
     get_prefill_context_model_parallel_rank,
@@ -61,7 +61,7 @@ def pcp_is_enabled() -> bool:
 
 def pcp_pad_len(
     total_tokens: int,
-    pcp_size: Optional[int] = None,
+    pcp_size: int | None = None,
     multiple: int = 1,
 ) -> int:
     """Padded token count so the global sequence is divisible by pcp_size * multiple.
@@ -86,7 +86,7 @@ def pcp_pad_len(
 
 
 def pcp_round_robin_split(
-    input_: torch.Tensor, pcp_size: Optional[int] = None, pcp_rank: Optional[int] = None
+    input_: torch.Tensor, pcp_size: int | None = None, pcp_rank: int | None = None
 ) -> torch.Tensor:
     """Take this rank's round-robin shard along dim 0.
 
@@ -109,7 +109,7 @@ def pcp_round_robin_split(
 
 
 def pcp_allgather_rerange(
-    input_: torch.Tensor, pcp_size: Optional[int] = None
+    input_: torch.Tensor, pcp_size: int | None = None
 ) -> torch.Tensor:
     """All-gather round-robin shards along dim 0 and restore original token order.
 
@@ -150,7 +150,7 @@ def pcp_allgather_rerange(
 
 
 def pcp_allgather_rankmajor(
-    input_: torch.Tensor, pcp_size: Optional[int] = None
+    input_: torch.Tensor, pcp_size: int | None = None
 ) -> torch.Tensor:
     """Gather this rank's 1/W stripe shard into the full rank-major sequence
     via a plain all_gather (dim=0). Inverse of pcp_reduce_scatter."""
@@ -162,7 +162,7 @@ def pcp_allgather_rankmajor(
 
 
 def pcp_reduce_scatter(
-    input_: torch.Tensor, pcp_size: Optional[int] = None
+    input_: torch.Tensor, pcp_size: int | None = None
 ) -> torch.Tensor:
     """Sum the pcp-half across ranks and scatter dim0 back to this rank's 1/W
     stripe via a plain reduce_scatter (dim=0). Inverse of pcp_allgather_rankmajor."""
@@ -173,9 +173,7 @@ def pcp_reduce_scatter(
     return get_pcp_group().reduce_scatter(input_.contiguous(), dim=0)
 
 
-def pcp_all_reduce(
-    input_: torch.Tensor, pcp_size: Optional[int] = None
-) -> torch.Tensor:
+def pcp_all_reduce(input_: torch.Tensor, pcp_size: int | None = None) -> torch.Tensor:
     """All-reduce (sum) over the PCP group, no token reshaping. DECODE path:
     tokens are pcp-redundant (every rank holds the same full batch), so just sum
     the pcp-half of the intermediate that combine_outputs' tp all_reduce missed.
@@ -189,7 +187,7 @@ def pcp_all_reduce(
 
 
 def pcp_round_robin_query_indices(
-    n_global_q: int, pcp_size: Optional[int] = None, pcp_rank: Optional[int] = None
+    n_global_q: int, pcp_size: int | None = None, pcp_rank: int | None = None
 ) -> torch.Tensor:
     """Global query indices owned by this rank under round-robin split.
 
@@ -205,6 +203,21 @@ def pcp_round_robin_query_indices(
     return torch.arange(pcp_rank, n_global_q, pcp_size, dtype=torch.long)
 
 
+def pcp_owned_rows_by_request(
+    cu_seqlens_q: np.ndarray, first: int, stride: int, owned: int
+) -> np.ndarray:
+    """Starts of the rows a PCP rank owns (`first`, `first + stride`, ...,
+    `owned` of them) grouped by request (`cu_seqlens_q`), then the padding
+    rows past the last request as one more. Rows below `x` it owns number
+    `ceil((x - first) / stride)`, so this is arithmetic on the host."""
+    cu = np.asarray(cu_seqlens_q, dtype=np.int64)
+    below = (np.maximum(cu - first, 0) + stride - 1) // stride
+    starts = below - below[0]
+    if starts[-1] < owned:
+        starts = np.append(starts, owned)
+    return starts.astype(np.int32)
+
+
 # pcp_pad_indptr / pcp_pad_dense share the (tensor, n_pad) signature but pad two
 # DIFFERENT metadata shapes, so they are kept separate on purpose:
 #

atom/model_engine/arg_utils.py

@@ -340,12 +340,12 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
             "--cudagraph-mode",
             type=str,
             default="FULL",
-            choices=["NONE", "PIECEWISE", "FULL", "FULL_AND_PIECEWISE", "AF_PIECEWISE"],
+            choices=["NONE", "TR_HYBRID", "FULL", "FULL_AND_TR_HYBRID"],
             help="CUDA graph runtime mode. FULL = manual whole-forward capture "
-            "(default, existing behavior). PIECEWISE = per-piece cudagraph with "
-            "attention eager (requires --level 3). AF_PIECEWISE = PIECEWISE where "
-            "the attention core is also captured into its own cudagraph with "
-            "zero-copy buffers (DeepSeek-V4 DSpark).",
+            "(default, existing behavior). TR_HYBRID = a token graph per token "
+            "bucket, breaking only where a launch depends on host data, which "
+            "runs eagerly or as a request graph; recorded at any --level. "
+            "FULL_AND_TR_HYBRID is recorded as TR_HYBRID.",
         )
         parser.add_argument(
             "--load_dummy",