diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py
index 13e99cbeec..2443a17ed1 100644
--- a/atom/model_engine/model_runner.py
+++ b/atom/model_engine/model_runner.py
@@ -98,7 +98,6 @@
     UBatchWrapper,
     local_tbo_precompute,
     maybe_create_ubatch_slices,
-    sync_dp_metadata,
 )
 
 logger = logging.getLogger("atom")
@@ -413,6 +412,14 @@ def get_token_locations(self, batch: ScheduledBatch) -> TokenLocations:
         num_prev = len(prev_req_ids)
         num_cur = len(cur_req_ids)
 
+        # A fabricated batch carries nothing over -- and the DP-sync dummy
+        # reuses one id, so without this the next dummy matches it and reads a
+        # request that never existed out of `prev_token_ids`/`draft_token_ids`.
+        # Real ids are non-negative and never matched it either way.
+        if self.prev_batch.is_dummy_run:
+            none = np.empty(0, dtype=np.intp)
+            return TokenLocations(none, none, np.arange(num_cur, dtype=np.intp))
+
         prev_id_to_idx = dict(zip(prev_req_ids, range(num_prev)))
 
         deferred_curr = np.empty(num_cur, dtype=np.intp)
@@ -447,6 +454,7 @@ def get_token_locations(self, batch: ScheduledBatch) -> TokenLocations:
     def prepare_input_ids(
         self,
         batch: ScheduledBatch,
+        max_seqlen_q: int,
     ) -> torch.Tensor:
         """Prepare the input IDs for the current batch.
 
@@ -513,12 +521,13 @@ def prepare_input_ids(
         num_deferred_seqs = len(deferred_curr_indices)
         num_new_seqs = len(new_curr_indices)
 
-        # `num_spec_query_tokens` is the single source of truth for the uniform
+        # `max_seqlen_q` is the single source of truth for the uniform
         # case (= mtp_k+1 for plain MTP, or the DSpark q-bucket when shrunk this
         # step); `dynamic_spec_query_tokens_per_req` overrides it per request
-        # when DSpark runs ragged. See `ScheduledBatch.num_spec_query_tokens`.
+        # when DSpark runs ragged. See `ForwardMode.max_seqlen_q`, which the
+        # step settles and the batch no longer carries.
         _per_req = getattr(batch, "dynamic_spec_query_tokens_per_req", None)
-        tokens_per_seq = int(batch.num_spec_query_tokens) if self.use_spec else 1
+        tokens_per_seq = max_seqlen_q if self.use_spec else 1
 
         # Receive and map bonus_list to current batch order
         self.num_rejected = batch.num_rejected
@@ -676,6 +685,12 @@ def __init__(self, rank: int, config: Config):
         self._setup_device_and_distributed(rank, config)
 
         self.capture_sizes = [0]  # for eager fallback
+        # The same ladder as an ASCENDING int32 array, which is what
+        # `ForwardMode.decide` searches. Separate from the list because the list
+        # is re-sorted in both directions during capture; rebound (never
+        # mutated) once capture narrows it, so a search cannot read a transient
+        # order.
+        self.capture_sizes_np = np.asarray(self.capture_sizes, dtype=np.int32)
         # PIECEWISE cudagraph state, populated by capture_cudagraph. Empty when
         # capture never ran (enforce_eager), so the ragged-bucket paths no-op.
         self._piecewise_captured_tokens: set[int] = set()
@@ -1160,7 +1175,14 @@ def debug(self, *args: Any):
             logger.info(*args)
 
     def dummy_execution(self):
-        """Execute dummy decode batch for DP synchronization."""
+        """Execute dummy decode batch for DP synchronization.
+
+        Two mechanisms lean on the fabricated id being -1: this pass flushes
+        the previous real step's tokens (`prepare_sampled_ids` reports the
+        batch before), and its own land on the key the deferred flag then
+        overwrites. What -1 must NOT do is look like a request --
+        `get_token_locations`.
+        """
         has_drafter = hasattr(self, "drafter")
         mtp_k = self.drafter.mtp_k if has_drafter else 0
         mtp_factor = mtp_k + 1
@@ -2088,7 +2110,7 @@ def _maybe_create_tbo_slices(
         """Create TBO ubatch slices when the collective DP decision is True.
 
         With the packed-reduce path the eligibility (local + cross-DP AND)
-        is decided in ``_preprocess``; here we just realise the split.
+        is decided in ``ForwardMode.decide``; here we just realise the split.
         """
         if not tbo_collective_active:
             return None
@@ -2112,161 +2134,50 @@ def _maybe_create_tbo_slices(
             )
         return ubatch_slices
 
-    def _preprocess(
-        self,
-        batch: ScheduledBatch,
-        num_scheduled_tokens: np.ndarray | None = None,
-        dspark_shape: tuple[int, int, int] | None = None,
-    ):
-        """Per-step DP sync: token padding, prefill fan-out, TBO decision.
-
-        Thin wrapper over :func:`atom.utils.tbo.sync_dp_metadata` (the
-        actual collective) and :func:`atom.utils.tbo.local_tbo_precompute`
-        (the rank-local TBO eligibility / per-ubatch token split).
-
-        ``dspark_shape`` (local q, decode_bs, total_tokens) folds DSpark's graph-shape
-        DP-MAX into this same all_gather so the two per-step collectives become
-        one; the reduced max is returned as the 7th tuple element for the caller
-        to apply via ``_apply_dspark_shape_max``.
-
-        Returns:
-            (running_tokens, num_tokens_across_dp, dp_uniform_decode,
-             max_tokens, tbo_collective_active, ub_max_tokens_across_dp,
-             dspark_shape_max)
+    def _local_tbo_eligibility(
+        self, batch: ScheduledBatch
+    ) -> tuple[bool, bool, int, int]:
+        """This rank's TBO answer, and the PCP flags that ride with it.
+
+        `(meets_min_tokens, can_split, ub0_tokens, ub1_tokens)`. The first two
+        are reduced across DP inside the step's one packed all_gather -- OR and
+        AND respectively -- so this is the local half of that question, not a
+        decision. Sets the PCP+TBO routing flags as a side effect because they
+        come out of the same sizing.
         """
-        scheduled_tokens = batch.total_tokens_num
         is_prefill = batch.total_tokens_num_prefill > 0
-        tbo_on = self.config.enable_tbo
-        dp_size = self.config.parallel_config.data_parallel_size
+        if not self.config.enable_tbo:
+            self._pcp_tbo_balanced_active = False
+            self._pcp_bal_groups = None
+            return (False, False, 0, 0)
 
-        # Rank-local TBO precompute (needed for both dp==1 fast path and
-        # the cross-DP packed gather below). `meets_min_tokens` = this rank's
-        # prefill reached the min-token bar (e.g. 8k), OR-reduced across DP;
-        # `can_split` = structurally splittable, AND-reduced across DP.
-        local_meets_min_tokens, local_can_split, local_ub0, local_ub1 = (
-            False,
-            False,
-            0,
-            0,
+        local_tbo = local_tbo_precompute(
+            self.config, batch, is_prefill, np.asarray(batch.num_scheduled_tokens)
         )
-        if tbo_on:
-            if num_scheduled_tokens is None:
-                num_scheduled_tokens = np.asarray(batch.num_scheduled_tokens)
-            local_meets_min_tokens, local_can_split, local_ub0, local_ub1 = (
-                local_tbo_precompute(
-                    self.config, batch, is_prefill, num_scheduled_tokens
-                )
-            )
-
-        # PCP+TBO prefill: split requests into two GROUPS at a request boundary
-        # (never split a sequence's tokens), so each ubatch = "non-TBO PCP on a
-        # request subset". Requires num_reqs >= 2 (request-boundary split needs
-        # two non-empty groups); bs=1 falls back to non-TBO.
-        pcp_size = self.config.prefill_context_parallel_size
-        # True for eligible PCP+TBO request-boundary split prefill; read by
-        # build_ubatch / run_model / prepare_prefill to route the per-group path.
-        self._pcp_tbo_balanced_active = False
-        # Per-group descriptors; reset each step, set in prepare_inputs
-        # request-boundary-split branch. Guards run_model/build_ubatch against
-        # stale values.
-        self._pcp_bal_groups = None
-        if tbo_on and is_prefill and pcp_size > 1 and not batch.is_dummy_run:
-            num_prefill_reqs = batch.total_seqs_num_prefill
-            n_prefill = batch.total_tokens_num_prefill
-            # Rough local sizing for TBO eligibility. PCP is always dp=1, so the
-            # dp_size<=1 fast path below returns local_eligible verbatim as
-            # tbo_collective_active; local_ub0/ub1 are only used by the dp>1
-            # sync path (never hit under PCP).
-            local_tokens = n_prefill // pcp_size
-            local_eligible = num_prefill_reqs >= 2 and local_tokens >= 2
-            local_ub0 = local_tokens // 2
-            local_ub1 = local_tokens - local_ub0
-            self._pcp_tbo_balanced_active = local_eligible
 
         # PCP+TBO prefill: split requests into two GROUPS at a request boundary
         # (never split a sequence's tokens), so each ubatch = "non-TBO PCP on a
         # request subset". Requires num_reqs >= 2 (request-boundary split needs
         # two non-empty groups); bs=1 falls back to non-TBO.
         pcp_size = self.config.prefill_context_parallel_size
-        # True for eligible PCP+TBO request-boundary split prefill; read by
-        # build_ubatch / run_model / prepare_prefill to route the per-group path.
+        # Read by build_ubatch / run_model / prepare_prefill to route the
+        # per-group path; reset every step so a stale value cannot route one.
         self._pcp_tbo_balanced_active = False
-        # Per-group descriptors; reset each step, set in prepare_inputs
-        # request-boundary-split branch. Guards run_model/build_ubatch against
-        # stale values.
         self._pcp_bal_groups = None
-        if tbo_on and is_prefill and pcp_size > 1 and not batch.is_dummy_run:
-            num_prefill_reqs = batch.total_seqs_num_prefill
-            n_prefill = batch.total_tokens_num_prefill
-            # Rough local sizing for TBO eligibility. PCP is always dp=1, so the
-            # dp_size<=1 fast path below returns local_eligible verbatim as
-            # tbo_collective_active; local_ub0/ub1 are only used by the dp>1
-            # sync path (never hit under PCP).
-            local_tokens = n_prefill // pcp_size
-            local_eligible = num_prefill_reqs >= 2 and local_tokens >= 2
-            local_ub0 = local_tokens // 2
-            local_ub1 = local_tokens - local_ub0
-            self._pcp_tbo_balanced_active = local_eligible
-
-        if dp_size <= 1:
-            # Single-rank: TBO decision is purely local; no collective needed.
-            # Both bits must hold (reached min-tokens AND able to split).
-            # dp_uniform_decode=True mirrors the DP-disabled case in the
-            # multi-rank branch (`not enable_dp_attention` => True) and the
-            # Context default — otherwise single-GPU/TP-only decode would
-            # be forced into eager and lose the CUDAGraph decode path.
-            self._dspark_decode_replay = True
-            self._eplb_any_rank_has_prefill = None
-            # One rank pads to nobody, so running == scheduled here.
-            return (
-                scheduled_tokens,
-                None,
-                True,
-                scheduled_tokens,
-                local_meets_min_tokens and local_can_split,
-                None,
-                dspark_shape,
-            )
-
-        sync = sync_dp_metadata(
-            dp_group=get_dp_group().cpu_group,
-            dp_size=dp_size,
-            scheduled_tokens=scheduled_tokens,
-            is_prefill=is_prefill,
-            tbo_on=tbo_on,
-            local_meets_min_tokens=local_meets_min_tokens,
-            local_can_split=local_can_split,
-            local_ub_tokens=(local_ub0, local_ub1),
-            dspark_shape=dspark_shape,
-            local_is_dummy=bool(getattr(batch, "is_dummy_run", False)),
-        )
-
-        # Stash the DP-wide prefill OR for the EPLB prefill gate. Reused for free
-        # by on_forward_pass_end when the DP group == the migration (EP) group.
-        self._eplb_any_rank_has_prefill = sync.any_rank_has_prefill
-        max_tokens = int(sync.num_tokens_across_dp.max())
-        dp_uniform_decode = (not sync.any_rank_has_prefill) or (
-            not self.config.enable_dp_attention
-        )
-        # Only here does the count stop being this rank's own: under uniform
-        # decode every rank pads to the same max for a fixed-size all_gather.
-        # The variable-length path keeps its own, so the two stay equal.
-        running_tokens = max_tokens if dp_uniform_decode else scheduled_tokens
-
-        self._dspark_decode_replay = dp_uniform_decode
-
-        return (
-            running_tokens,
-            sync.num_tokens_across_dp,
-            dp_uniform_decode,
-            max_tokens,
-            sync.tbo_collective_active,
-            sync.ub_max_tokens_across_dp,
-            sync.dspark_shape_max,
-        )
-
-    def _dspark_apply_q_bucket(self, batch: ScheduledBatch) -> None:
-        """Shrink this decode step's verify length to one CUDA-graph bucket q.
+        if is_prefill and pcp_size > 1 and not batch.is_dummy_run:
+            # PCP is always dp=1, so the single-rank path returns this verbatim
+            # as `tbo_collective_active`; ub0/ub1 are only read under dp > 1.
+            local_tokens = batch.total_tokens_num_prefill // pcp_size
+            eligible = batch.total_seqs_num_prefill >= 2 and local_tokens >= 2
+            ub0 = local_tokens // 2
+            local_tbo = (eligible, eligible, ub0, local_tokens - ub0)
+            self._pcp_tbo_balanced_active = eligible
+        return local_tbo
+
+    def _dspark_apply_q_bucket(self, batch: ScheduledBatch) -> int | None:
+        """Shrink this decode step's verify length to one CUDA-graph bucket q,
+        and return it -- or None when nothing shrank and the step keeps the
+        configured `num_spec_step + 1`.
 
         q = quantize_up(max ell_i + 1) over the batch (ell_i = last step's
         per-req schedule). All seqs then forward q tokens (anchor + q-1 drafts)
@@ -2277,16 +2188,14 @@ def _dspark_apply_q_bucket(self, batch: ScheduledBatch) -> None:
         tokens truncated to q-1); KV stays reserved at mtp_k+1. No-op unless
         DSpark confidence scheduling is on and this is a pure-decode batch.
         """
-        # Idempotency guard: prepare_model calls this before prepare_input_ids,
-        # and prepare_inputs (also reachable standalone for dummy/warmup) calls
-        # it again — only the first application must shrink the batch.
-        if getattr(batch, "_dspark_q_applied", False):
-            return
+        # No idempotency guard: `prepare_model` is the one caller and runs
+        # this once per batch. It used to be two, and the guard returned
+        # early on the second -- which now means returning None, i.e. the
+        # FULL length, undoing the shrink it was written to protect.
         if not (hasattr(self, "drafter") and self.drafter.uses_confidence_schedule):
-            return
+            return None
         if batch.total_tokens_num_prefill > 0:
-            return  # mixed/prefill step: keep full length
-        batch._dspark_q_applied = True
+            return None  # mixed/prefill step: keep full length
         scheduled_bs = batch.total_seqs_num_decode
         if scheduled_bs <= 0:
             return
@@ -2315,8 +2224,7 @@ def _dspark_apply_q_bucket(self, batch: ScheduledBatch) -> None:
         # dspark.ragged_graph_sizes set. Anchor lower bound (q>=num_bonus+1)
         # is applied PER REQUEST so each seg can hold its own anchor.
         if self.config.dspark.ragged:
-            self._dspark_apply_ragged(batch, scheduled_bs, full_q, by_req)
-            return
+            return self._dspark_apply_ragged(batch, scheduled_bs, full_q, by_req)
         # ====================================================================
 
         # ---- Q-BUCKET path (older batch-uniform padding scheme) ------------
@@ -2375,13 +2283,11 @@ def _dspark_apply_q_bucket(self, batch: ScheduledBatch) -> None:
         batch.num_scheduled_tokens = nst
         batch.total_tokens_num_decode = int(nst[:scheduled_bs].sum())
         batch.total_tokens_num = prefill_tok + batch.total_tokens_num_decode
-        # Publish the chosen q as the single source of truth (see ScheduledBatch
-        # .num_spec_query_tokens). All downstream length consumers read this.
-        batch.num_spec_query_tokens = q
         # Truncate each request's draft block to q-1 (regular matrix: all seqs q-1).
         spec = batch.scheduled_spec_decode_tokens
         if spec is not None and getattr(spec, "size", 0) > 0:
             batch.scheduled_spec_decode_tokens = np.ascontiguousarray(spec[:, : q - 1])
+        return q
 
     def _dspark_apply_ragged(self, batch, scheduled_bs, full_q, by_req):
         """DSpark per-request RAGGED verify (paper §5.2 avoid-padding).
@@ -2440,11 +2346,11 @@ def _dspark_apply_ragged(self, batch, scheduled_bs, full_q, by_req):
         # batch is rewritten. Shrinking the flat token layout is only safe if the
         # replay can follow it down; when it cannot, the rebuild is what makes
         # the step unsafe, so the decision has to come first.
-        #   * num_spec_query_tokens (scalar) q_eff : the PER-SEQ length bound
+        #   * max_seqlen_q (scalar) q_eff : the PER-SEQ length bound
         #     (>= max(new_len), quantized up to a captured bucket). Per-seq
         #     structures (compressor grid, rectangular indexer) size by it, so no
         #     seq can overflow them. It is NOT the total compute size -- that is
-        #     the flat num_tokens bucket (dynamic_num_tokens_pad), sized to the
+        #     the flat num_tokens bucket (running_tokens), sized to the
         #     real sum, so a long-tail seq no longer inflates the batch row count.
         buckets = resolve_q_buckets(self.config.dspark.ragged_graph_sizes, full_q)
         if self.enforce_eager:
@@ -2455,16 +2361,23 @@ def _dspark_apply_ragged(self, batch, scheduled_bs, full_q, by_req):
             # Graph: q_eff = smallest bucket >= the real MAX per-seq len, so no
             # seq ever exceeds q_eff. Per-seq structures (compressor grid,
             # rectangular indexer) size by q_eff and can't overflow -- no separate
-            # full_q cap needed. The TOTAL compute size is chosen apart from this
-            # by the flat num_tokens bucket (dynamic_num_tokens_pad, sized to the
-            # real sum), so q_eff no longer needs to track the sum/avg.
+            # full_q cap needed. The TOTAL compute size is `running_tokens`,
+            # settled apart from this by `ForwardMode`, so q_eff no longer needs
+            # to track the sum/avg.
             q_eff = (
                 quantize_to_bucket(int(new_len.max()), buckets)
                 if scheduled_bs > 0
                 else full_q
             )
-            if not self._ragged_flat_bucket_exists(int(new_len.sum()), q_eff):
-                return  # replay cannot follow the shrink -> stay rectangular
+            # Under FULL cudagraphs nothing flat is recorded, so a packed run has
+            # no width to land on and shrinking buys nothing -- the replay is
+            # rectangular either way. WHICH width it lands on is not asked here:
+            # `ForwardMode.decide` settles that on the DP-agreed total, and a
+            # second search on this rank's own count is how the two came to
+            # disagree. A miss there forwards the run eagerly, which is what the
+            # branch above already does under `enforce_eager`.
+            if not self._piecewise_cg_active():
+                return  # stay rectangular
 
         # Rebuild scheduled_tokens (flat) to the ragged per-seq layout: keep the
         # first new_len[i] of each seq's old segment (token[0]=anchor, rest=draft
@@ -2491,178 +2404,64 @@ def _dspark_apply_ragged(self, batch, scheduled_bs, full_q, by_req):
         # Two sources of truth (TRUE FLAT, paper §5.2): tokens are flat-packed
         # [0:Σ] with the per-seq ragged new_len.
         #   * dynamic_spec_query_tokens_per_req : the true ragged per-seq lengths.
-        #   * num_spec_query_tokens (scalar) q_eff : computed above, before the
-        #     rebuild, together with the replay-shape feasibility check.
-        batch.num_spec_query_tokens = int(q_eff)
+        #   * the scalar q_eff, RETURNED to `prepare_model` -- computed above,
+        #     before the rebuild, together with the replay-shape feasibility
+        #     check.
         batch.dynamic_spec_query_tokens_per_req = new_len
 
         # (No flat scheduled_spec_decode_tokens is built here: the ragged
         # input_ids are assembled downstream by `fill_deferred_decode_ids` from
         # prev_token_ids (anchor) + draft_token_ids, which never consults
         # scheduled_spec_decode_tokens.)
-
-    def _dspark_local_shape(self, batch: ScheduledBatch) -> tuple[int, int, int] | None:
-        """Local (q, decode_bs, total_tokens) for the DSpark DP graph-shape sync,
-        or None when the sync does not apply (single-DP or non-DSpark).
-
-        Symmetric across ranks every step regardless of prefill/decode: on a
-        prefill step this returns (1, 0, 0) so the DP-MAX reduction still has a
-        well-defined identity contribution and every rank participates in the
-        same collective (matching the pre-merge all_reduce semantics)."""
-        if self.config.parallel_config.data_parallel_size <= 1:
-            return None
-        drafter = getattr(self, "drafter", None)
-        if drafter is None or not drafter.is_block_drafter:
-            return None
-        local_q = int(getattr(batch, "num_spec_query_tokens", 1))
-        local_bs = int(getattr(batch, "total_seqs_num_decode", 0))
-        # Also DP-max the ragged decode token total (total_tokens_num_decode).
-        # PIECEWISE 1D-ragged replays at a num_tokens bucket sized to this total;
-        # to keep the MoE all_gather row count identical across ranks (else RCCL
-        # deadlock) every rank must pick the SAME bucket, so bisect on the DP-max
-        # total, not the local one.
-        local_total_tokens = int(getattr(batch, "total_tokens_num_decode", 0))
-        return local_q, local_bs, local_total_tokens
-
-    def _apply_dspark_shape_max(
-        self, batch: ScheduledBatch, shape_max: tuple[int, int, int] | None
-    ) -> None:
-        """Adopt the DP-MAX (q, decode_bs, total_tokens) across DP ranks so every
-        rank captures/replays an identical decode graph shape (else the MoE
-        all_gather pads to divergent rows -> RCCL deadlock). Raising q/bs/total
-        only enlarges graph capacity (real tokens stay flat-packed in [0:total])
-        and leaves per-request num_scheduled_tokens untouched, so it is always
-        lossless. The DP-max is computed via ``_dspark_local_shape`` and the
-        collective is folded into ``sync_dp_metadata``'s packed all_gather."""
-        if shape_max is None:
-            return
-        batch.num_spec_query_tokens = int(shape_max[0])
-        batch.dspark_dp_bs = int(shape_max[1])
-        batch.dspark_dp_total_tokens = int(shape_max[2])
+        return int(q_eff)
 
     def prepare_inputs(
         self,
         batch: ScheduledBatch,
-        input_ids: torch.Tensor = None,
-        preprocessed: tuple | None = None,
+        input_ids: torch.Tensor,
+        forward_mode: ForwardMode,
     ):
-        # NOTE: DSpark q-bucket shrink happens in prepare_model BEFORE
-        # prepare_input_ids, so the batch is already reduced when we get here.
-        # ``preprocessed``: when prepare_model already ran the merged DP collective
-        # (to get the DSpark DP-max q before prepare_input_ids), it passes the
-        # cached _preprocess tuple here so we DON'T issue a second all_gather.
+        # Always supplied, settled in `prepare_model` (which is where the reason
+        # lives). The q-bucket shrink ran there too, so `batch` is already
+        # reduced here.
         is_prefill = batch.total_tokens_num_prefill > 0
-        bs = batch.total_seqs_num
+        scheduled_bs = batch.total_seqs_num
         scheduled_tokens = batch.total_tokens_num
         num_scheduled_tokens = np.asarray(batch.num_scheduled_tokens)
         cu_seqlens_q, _arange = self._get_cumsum_and_arange(num_scheduled_tokens)
-        if preprocessed is None:
-            preprocessed = self._preprocess(
-                batch,
-                num_scheduled_tokens=num_scheduled_tokens,
-                dspark_shape=self._dspark_local_shape(batch),
-            )
-            self._apply_dspark_shape_max(batch, preprocessed[6])
-        (
-            running_tokens,
-            num_tokens_across_dp,
-            dp_uniform_decode,
-            _max_tokens,
-            tbo_collective_active,
-            ub_max_tokens_across_dp,
-            _dspark_shape_max,
-        ) = preprocessed
-        # NOTE: self._dspark_decode_replay is set inside _preprocess (it needs
-        # the DP sync result), so it's already current here for build()/run_model.
-
-        # Precompute the flat replay token count once here (before attn build +
-        # run_model) so the attn builder's positions padding matches it.
-        batch.dynamic_num_tokens_pad = self._dynamic_num_tokens_pad(batch)
+        sync = forward_mode.sync
+        num_tokens_across_dp = None if sync is None else sync.num_tokens_across_dp
+        tbo_collective_active = forward_mode.tbo_collective_active
+        ub_max_tokens_across_dp = None if sync is None else sync.ub_max_tokens_across_dp
+        running_tokens_are_unified = forward_mode.running_tokens_are_unified
 
         if not tbo_collective_active:
             self._pcp_tbo_balanced_active = False
 
-        self.forward_vars["cu_seqlens_q"].np[1 : bs + 1] = cu_seqlens_q
-
-        # mtp_step = per-seq decode token count, used by ForwardMode.decide to
-        # recover batch size as running_tokens // mtp_step. This is exactly
-        # the batch's single-source-of-truth decode length (= mtp_k+1, or the
-        # DSpark q-bucket when shrunk); running_tokens = scheduled_bs *
-        # num_spec_query_tokens, so the division recovers bs correctly. Prefill
-        # has no drafter / uses 1.
-        decide_running_tokens = running_tokens
-        dp_bs = batch.dspark_dp_bs
-        is_ragged = (
-            getattr(batch, "dynamic_spec_query_tokens_per_req", None) is not None
-        )
-        if not is_prefill and (dp_bs is not None or is_ragged):
-            # DSpark: the real Σtokens is irregular, but we
-            # replay the rectangular (bs, q_eff) graph whose capacity is bs*q_eff
-            # (q_eff = num_spec_query_tokens, the quantized bucket). Feed
-            # ForwardMode the GRAPH-CAPACITY token count so it recovers
-            # padded_scheduled_bs = bs*q_eff // q_eff = bs and picks the matching
-            # (bs, q_eff) graph; the real ragged tokens sit in [0:Σ], the tail is
-            # -1 padding (CTAs bail).
-            q_eff = int(batch.num_spec_query_tokens)
-            eff_bs = dp_bs if dp_bs is not None else batch.total_seqs_num_decode
-            mtp_step = q_eff
-            decide_running_tokens = int(eff_bs) * q_eff
-        elif not is_prefill and hasattr(self, "drafter"):
-            mtp_step = batch.num_spec_query_tokens
-        else:
-            mtp_step = (self.drafter.mtp_k + 1) if hasattr(self, "drafter") else 1
-        forward_mode = ForwardMode.decide(
-            is_prefill=is_prefill,
-            # `total_seqs_num`, not `total_seqs_num_prefill`: equal today, and
-            # the one that stays right once a prefill batch may also carry
-            # decode rows -- they go through the same attention call.
-            scheduled_bs=(
-                batch.total_seqs_num if is_prefill else batch.total_seqs_num_decode
-            ),
-            running_tokens=decide_running_tokens,
-            dp_uniform_decode=dp_uniform_decode,
-            enforce_eager=self.enforce_eager,
-            capture_sizes=self.capture_sizes,
-            mtp_step=mtp_step,
-        )
+        self.forward_vars["cu_seqlens_q"].np[1 : scheduled_bs + 1] = cu_seqlens_q
 
         if not is_prefill:
-            scheduled_bs = forward_mode.scheduled_bs
-            # Attention reads the padded width only where its tensors carry the
-            # padding; eager keeps the scheduled rows so aiter's t == t_slot
-            # holds against the local input_ids.
-            bs = (
-                forward_mode.running_bs
-                if forward_mode.attn_tensors_are_padded
-                else scheduled_bs
-            )
             assert forward_mode.running_bs >= scheduled_bs, (
                 f"running_bs={forward_mode.running_bs} < "
                 f"scheduled_bs={scheduled_bs}; ForwardMode.decide invariant violated"
             )
-            # Only pad cu_seqlens_q out to the cudagraph capture size if we
-            # actually grew bs. Eager (bs == scheduled_bs) leaves the slice
-            # empty so no overwrite happens.
-            if bs > scheduled_bs:
-                self.forward_vars["cu_seqlens_q"].np[scheduled_bs + 1 : bs + 1] = (
-                    self.forward_vars["cu_seqlens_q"].np[scheduled_bs]
-                )
-        attn_metadata, positions = self.attn_metadata_builder.build(batch=batch, bs=bs)
-
-        # The step's padded shape, in the two units its consumers read. Both
-        # must be identical on every DP rank: MoE pads hidden_states to
-        # `running_tokens` before the cross-DP all_gather, and a captured graph
-        # holds its collective at `running_bs`.
+            # Flat-extend the cumsum over the fabricated rows. Attention runs at
+            # `running_bs` on every decode step, replayed or not: keeping the
+            # scheduled rows when eager made a third width nobody could name,
+            # and every consumer had to know which of the two it was holding.
+            self.forward_vars["cu_seqlens_q"].np[
+                scheduled_bs + 1 : forward_mode.running_bs + 1
+            ] = self.forward_vars["cu_seqlens_q"].np[scheduled_bs]
+        # The step's two units, read from where they were settled: `running_bs`
+        # sizes everything per-sequence, `running_tokens` everything per-row.
         running_bs = forward_mode.running_bs
-        if not is_prefill:
-            # Prefill keeps whatever `_preprocess` settled on -- padded to the
-            # group max under uniform decode, its own count otherwise. Decode
-            # recomputes from the batch `decide` rounded up to, which is wider
-            # than that count whenever it landed on a captured size.
-            running_tokens = running_bs * int(attn_metadata.max_seqlen_q)
-        drafter = getattr(self, "drafter", None)
-        if not is_prefill and drafter is not None and drafter.is_block_drafter:
-            running_tokens = self._dspark_ragged_running_tokens(batch, running_tokens)
+        running_tokens = forward_mode.running_tokens
+        attn_metadata, positions = self.attn_metadata_builder.build(
+            batch=batch,
+            running_bs=running_bs,
+            running_tokens=running_tokens,
+            max_seqlen_q=forward_mode.max_seqlen_q,
+        )
         context = Context(
             positions=positions,
             is_prefill=is_prefill,
@@ -2671,7 +2470,7 @@ def prepare_inputs(
             scheduled_tokens=scheduled_tokens,
             running_bs=running_bs,
             running_tokens=running_tokens,
-            dp_uniform_decode=dp_uniform_decode,
+            running_tokens_are_unified=running_tokens_are_unified,
             forward_mode=forward_mode,
         )
 
@@ -2768,28 +2567,45 @@ def prepare_sample(
         return temperatures, top_ks, top_ps, all_greedy, needs_independent_noise
 
     def prepare_model(self, batch: ScheduledBatch):
-        self._dspark_apply_q_bucket(batch)
-        # DSpark-only early DP sync: only DSpark under DP needs the DP-max q
-        # BEFORE prepare_input_ids (to size input_ids). Run the merged packed
-        # all_gather (TBO + DSpark [q, bs, total_tokens]) ONCE here and reuse it in
-        # prepare_inputs, so the step issues a single cross-DP collective.
-        dspark_shape = self._dspark_local_shape(batch)
-        preprocessed = None
-        if dspark_shape is not None:
-            preprocessed = self._preprocess(
-                batch,
-                num_scheduled_tokens=np.asarray(batch.num_scheduled_tokens),
-                dspark_shape=dspark_shape,
-            )
-            self._apply_dspark_shape_max(batch, preprocessed[6])
+        shrunk_q = self._dspark_apply_q_bucket(batch)
+        # The step's shape, settled once. Here rather than in prepare_inputs
+        # because DSpark under DP needs the reduced query length BEFORE
+        # prepare_input_ids sizes the buffer, and because one call site is the
+        # only way the step is guaranteed a single cross-DP collective.
+        dp_size = self.config.parallel_config.data_parallel_size
+        forward_mode = ForwardMode.decide(
+            batch=batch,
+            dp_size=dp_size,
+            dp_group=get_dp_group().cpu_group if dp_size > 1 else None,
+            enforce_eager=self.enforce_eager,
+            capture_sizes=self.capture_sizes_np,
+            captured_tokens=(
+                self._piecewise_sorted_tokens if self._piecewise_cg_active() else None
+            ),
+            is_block_drafter=(
+                hasattr(self, "drafter") and self.drafter.is_block_drafter
+            ),
+            tbo_on=self.config.enable_tbo,
+            local_tbo=self._local_tbo_eligibility(batch),
+            max_seqlen_q=(batch.num_spec_step + 1 if shrunk_q is None else shrunk_q),
+        )
+        # Stash the DP-wide prefill OR for the EPLB prefill gate; reused free by
+        # on_forward_pass_end when the DP group == the migration (EP) group.
+        self._eplb_any_rank_has_prefill = (
+            None
+            if forward_mode.sync is None
+            else forward_mode.sync.any_rank_has_prefill
+        )
         total_tokens_num = batch.total_tokens_num
         assert total_tokens_num > 0
 
         temperatures, top_ks, top_ps, all_greedy, needs_independent_noise = (
             self.prepare_sample(batch)
         )
-        input_ids = self.tokenID_processor.prepare_input_ids(batch)
-        self.prepare_inputs(batch, input_ids, preprocessed=preprocessed)
+        input_ids = self.tokenID_processor.prepare_input_ids(
+            batch, forward_mode.max_seqlen_q
+        )
+        self.prepare_inputs(batch, input_ids, forward_mode=forward_mode)
 
         # Stage the speculative inputs while this forward's normal staging
         # window is still open.  Both buffers are pinned and reused, so copying
@@ -3032,18 +2848,18 @@ def run_model(
         is_prefill = context.is_prefill
         positions = context.positions
 
-        # Dispatch is owned by ForwardMode.decide() (called in prepare_inputs).
+        # Dispatch is owned by ForwardMode.decide() (called in prepare_model).
         # Every run_model caller MUST go through prepare_inputs first, so
         # forward_mode is always set here.
         forward_mode = context.forward_mode
         assert forward_mode is not None, (
             "context.forward_mode is None; run_model invoked without going "
-            "through prepare_inputs. Add ForwardMode.decide() at the new "
+            "through prepare_model. Add ForwardMode.decide() at the new "
             "entry point instead of re-deriving the 4-OR dispatch here."
         )
 
         # Single canonical shape check; contract owned by ForwardMode, which
-        # internally short-circuits for prefill / cudagraph.
+        # short-circuits only on prefill. The padded step is the one it is for.
         forward_mode.assert_shape_contract(input_ids, forward_context.attn_metadata)
 
         # Profiler label. Kind (prefix) distinguishes real/dummy and
@@ -3181,58 +2997,63 @@ def run_model(
             # dummy_decode[...] — see build_run_label.
             with record_function(label):
                 running_bs = context.running_bs
-                max_q_len = forward_context.attn_metadata.max_seqlen_q
-                num_tokens = context.scheduled_bs * max_q_len  # real (output slice)
+                running_tokens = forward_mode.running_tokens
+                scheduled_tokens = context.scheduled_tokens
 
                 if self._piecewise_cg_active():
-                    num_tokens_pad, real_tokens, _captured = (
-                        self._piecewise_replay_shape(batch, running_bs, max_q_len)
-                    )
-                    # Pad tail to a legal vocab id / position (builder fills to
-                    # graph_cap >= num_tokens_pad, so a no-op safety net).
-                    if num_tokens_pad > real_tokens:
+                    # Pad tail to a legal vocab id / position, from THIS rank's
+                    # own rows out to the width the step settled on. A group-max
+                    # lower bound leaves `[scheduled, max)` holding the previous
+                    # step's ids on every rank below the max, and those reach the
+                    # draft's Markov lookup as out-of-range indices.
+                    if running_tokens > scheduled_tokens:
                         self.forward_vars["input_ids"].gpu[
-                            real_tokens:num_tokens_pad
+                            scheduled_tokens:running_tokens
                         ].zero_()
                         self.forward_vars["positions"].gpu[
-                            real_tokens:num_tokens_pad
+                            scheduled_tokens:running_tokens
                         ].zero_()
                     _pos = (
-                        self._mrope_positions_view(num_tokens_pad)
+                        self._mrope_positions_view(running_tokens)
                         if self.use_mrope
-                        else self.forward_vars["positions"].gpu[:num_tokens_pad]
+                        else self.forward_vars["positions"].gpu[:running_tokens]
                     )
                     forward_context.cudagraph_runtime_mode = (
-                        CUDAGraphMode.PIECEWISE if _captured else CUDAGraphMode.NONE
+                        CUDAGraphMode.PIECEWISE
+                        if forward_mode.piecewise_captured
+                        else CUDAGraphMode.NONE
                     )
                     forward_context.batch_descriptor = BatchDescriptor(
-                        num_tokens=num_tokens_pad
+                        num_tokens=running_tokens
                     )
                     model_output = self.model(
-                        self.forward_vars["input_ids"].gpu[:num_tokens_pad], _pos
+                        self.forward_vars["input_ids"].gpu[:running_tokens], _pos
                     )
                     forward_context.cudagraph_runtime_mode = CUDAGraphMode.NONE
                     forward_context.batch_descriptor = None
                     # model_output is always a plain Tensor; drafter aux capture
                     # (if any) already wrote its own buffers inside the forward.
-                    # Spec keeps the padded [0:Σ] layout (postprocess/draft
-                    # re-gather to bs via next_token_locs); non-spec slices to the
-                    # real num_tokens so pad rows never leak into sampled_token_ids
-                    # -> prev_token_ids -> next-step shape mismatch.
-                    hidden_states = model_output
-                    _is_spec = hasattr(self, "drafter")
-                    _slice_len = num_tokens_pad if _is_spec else num_tokens
-                    hidden_states = hidden_states[:_slice_len]
+                    # Spec keeps the padded layout (postprocess/draft re-gather to
+                    # bs via next_token_locs); non-spec cuts to the scheduled rows
+                    # so pad rows never leak into sampled_token_ids ->
+                    # prev_token_ids -> next-step shape mismatch.
+                    hidden_states = model_output[
+                        : (
+                            running_tokens
+                            if hasattr(self, "drafter")
+                            else scheduled_tokens
+                        )
+                    ]
                     logits = self.model.compute_logits(hidden_states)
                     return logits, hidden_states
 
-                graph_key = (running_bs, max_q_len)
+                graph_key = (running_bs, forward_context.attn_metadata.max_seqlen_q)
                 self.graphs[graph_key].replay()
-                hidden_states = self.forward_vars["outputs"][:num_tokens]
+                hidden_states = self.forward_vars["outputs"][:scheduled_tokens]
                 # Drafter aux buffers (if any) refresh on replay: their in-place
                 # copy ops were captured into the graph.
                 if self.logits_in_graph:
-                    logits = self.graph_logits[graph_key][:num_tokens]
+                    logits = self.graph_logits[graph_key][:scheduled_tokens]
                 else:
                     logits = self.model.compute_logits(hidden_states)
 
@@ -3259,6 +3080,15 @@ def postprocess(
         spec_decode_metadata = get_forward_context().spec_decode_metadata
         bs = batch.total_seqs_num
         if spec_decode_metadata is None:
+            # The LM head emitted one row per sequence the step FORWARDED,
+            # which prefill pads to `running_bs` for the draft pass that
+            # follows it. Cut to the scheduled batch here and nowhere else:
+            # this is the boundary where a padded forward becomes a
+            # per-request result, and everything below counts requests -- the
+            # sampler's per-row parameters and the logprob gather both. Cut
+            # after either and the pad rows divide `[running_bs, V]` by
+            # `[scheduled_bs, 1]`.
+            logits = logits[:bs]
             sampled_tokens = self.sampler(
                 logits,
                 temperatures,
@@ -3499,7 +3329,7 @@ def _is_pure_middle_chunk(batch) -> bool:
     def _dp_draft_lockstep_active(self) -> bool:
         """Are this rank's draft passes bound to what the DP peers run?
 
-        Only under DP attention -- `_refresh_dp_metadata` returns early at
+        Only under DP attention -- `_publish_draft_shape` returns early at
         `data_parallel_size <= 1`, where an output-less batch legitimately
         drafts nothing. PP is excluded via `is_deferred_out`.
         """
@@ -3602,7 +3432,7 @@ def propose_draft_token_ids(
             and batch.total_tokens_num_prefill == 0
         ):
             full_q = self.drafter.mtp_k + 1
-            q_actual = batch.num_spec_query_tokens
+            q_actual = forward_context.context.forward_mode.max_seqlen_q
             if 1 <= q_actual < full_q:
                 last_token_offset = last_token_offset - (full_q - q_actual)
 
@@ -3740,126 +3570,6 @@ def _force_aiter_unreg_capture_for_piecewise(self):
                     getter,
                 )
 
-    def _ragged_flat_bucket_exists(self, total_tokens: int, q: int) -> bool:
-        """Can a captured flat num_tokens bucket hold a ragged step of
-        ``total_tokens`` at per-seq bound ``q``?
-
-        Same predicate ``_dynamic_num_tokens_pad`` searches with -- keep the two
-        in step. Asked BEFORE ``_dspark_apply_ragged`` rewrites the batch,
-        because when no bucket matches that function returns None and every
-        caller falls back to ``bs * max_seqlen_q``: a replay over MORE tokens
-        than the ragged rebuild populated. The ``[total : bs*q]`` tail is then
-        whatever the previous step left in the input buffer, and those stale ids
-        reach the draft's Markov transition-table lookup as out-of-range
-        indices -- a device-side trap attributed to whatever kernel happened to
-        be in flight, nowhere near here.
-
-        The usual way to get here is FULL (non-PIECEWISE) cudagraphs, where no
-        flat bucket set is captured at all: then NO ragged shrink is
-        representable and the caller must stay rectangular. Shrinking buys
-        nothing in that mode anyway -- the replay is rectangular either way.
-        """
-        if not self._piecewise_cg_active():
-            return False
-        from atom.spec_decode.dspark_scheduler import flat_bucket_fits
-
-        return flat_bucket_fits(
-            total_tokens, q, getattr(self, "_piecewise_sorted_tokens", None)
-        )
-
-    def _dynamic_num_tokens_pad(self, batch) -> int | None:
-        """Flat PIECEWISE replay token count for a ragged decode step, or None so
-        callers fall back to ``bs * max_seqlen_q``.
-
-        Mirrors the ragged branch of ``_piecewise_replay_shape``: the smallest
-        captured q-divisible num_tokens bucket >= the (DP-max) real token total.
-        None for non-ragged / non-piecewise / prefill / no-match.
-
-        Dummy batches are INCLUDED: a DP-lockstep dummy rank must pad to the same
-        flat token count as the real ranks or it cannot replay alongside them."""
-        if (
-            batch is None
-            or batch.total_tokens_num_prefill > 0
-            or not self._piecewise_cg_active()
-            or not self._piecewise_sorted_tokens
-            or not hasattr(self, "drafter")
-        ):
-            return None
-        dp_total = batch.dspark_dp_total_tokens
-        real_tokens = (
-            int(dp_total)
-            if dp_total is not None
-            else int(batch.total_tokens_num_decode)
-        )
-        q = int(batch.num_spec_query_tokens)
-        for b in self._piecewise_sorted_tokens:
-            if b >= real_tokens and q > 0 and b % q == 0:
-                return int(b)
-        return None
-
-    def _piecewise_replay_shape(self, batch, running_bs, max_q_len):
-        """Pick the PIECEWISE replay token count for one decode step.
-
-        Returns ``(num_tokens_pad, real_tokens, captured)``:
-        - ``num_tokens_pad``: token count to forward (the captured bucket size).
-        - ``real_tokens``: real tokens present in ``[0:real_tokens]``.
-        - ``captured``: whether a matching cudagraph bucket exists (else eager).
-
-        DSpark (has drafter, TP-only) replays at a flat num_tokens bucket sized
-        to the REAL ragged token total (= total_tokens_num_decode) so MoE/linear
-        shrink with it (dynamic verify length); attention is eager on the flat
-        [0:total] tokens, the [total:pad] tail is masked. Non-spec (or DP) uses
-        the rectangular bucket num_tokens == bs.
-        """
-        use_ragged_bucket = (
-            batch is not None
-            and self._piecewise_sorted_tokens
-            and hasattr(self, "drafter")
-        )
-        if use_ragged_bucket:
-            dp_total_tokens = batch.dspark_dp_total_tokens
-            real_tokens = (
-                int(dp_total_tokens)
-                if dp_total_tokens is not None
-                else int(batch.total_tokens_num_decode)
-            )
-            buckets = self._piecewise_sorted_tokens
-            q = int(max_q_len)
-            # Pick the q-divisible bucket N. q is num_spec_query_tokens, already
-            # quantized to a captured bucket by _dspark_apply_ragged, so the
-            # divisibility test is satisfiable by construction.
-            _replay = bool(getattr(self, "_dspark_decode_replay", False))
-            for b in buckets:
-                if b >= real_tokens and q > 0 and b % q == 0:
-                    return b, real_tokens, _replay
-            return max(real_tokens, running_bs * max_q_len), real_tokens, False
-
-        num_tokens_pad = running_bs * max_q_len
-        captured = num_tokens_pad in self._piecewise_captured_tokens
-        return num_tokens_pad, num_tokens_pad, captured
-
-    def _dspark_ragged_running_tokens(self, batch, default_running_tokens):
-        """``running_tokens`` for a DSpark ragged PIECEWISE decode step.
-
-        A ragged step replays at a flat num_tokens bucket, not a rectangular
-        bs*q grid, so the height MoE pads to IS that bucket. Picking the same
-        one ``_piecewise_replay_shape`` will (hence the identical q-divisibility
-        filter) is what keeps the eager ranks' all_gather equal to the size
-        baked into the replaying ranks' graph.
-        """
-        dp_total_tokens = batch.dspark_dp_total_tokens
-        if (
-            not self._piecewise_cg_active()
-            or dp_total_tokens is None
-            or not self._piecewise_sorted_tokens
-        ):
-            return default_running_tokens
-        q = int(batch.num_spec_query_tokens)
-        for b in self._piecewise_sorted_tokens:
-            if b >= int(dp_total_tokens) and q > 0 and b % q == 0:
-                return int(b)
-        return default_running_tokens
-
     def _dspark_capture_q_buckets(self, full_q: int) -> list[int]:
         """DSpark query-length buckets to capture graphs for (paper Phase 2).
 
@@ -4320,7 +4030,8 @@ def pause_gc():
                     ),
                     capture_ctx.stream,
                 )
-        self.capture_sizes.sort(reverse=False)
+        self.capture_sizes.sort()
+        self.capture_sizes_np = np.asarray(self.capture_sizes, dtype=np.int32)
 
         # PIECEWISE: sorted 1D num_tokens buckets for run_model's round_up_1d(Σ)
         # dispatch (bisect_left over this to pick the tightest captured shape).
diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py
index 2073d1c4e0..42d08cfef7 100644
--- a/atom/model_engine/scheduler.py
+++ b/atom/model_engine/scheduler.py
@@ -789,21 +789,11 @@ def __init__(
 
         self.is_dummy_run = is_dummy_run
         self.num_spec_step = num_spec_step
-        self.num_spec_query_tokens = num_spec_step + 1
         # DSpark RAGGED (paper §5.2): per-request decode query lengths [bs]
         # (ell_r + 1). None unless _dspark_apply_ragged set it this step; when
         # set, consumers use it (per-seq) instead of the scalar above.
         self.dynamic_spec_query_tokens_per_req = None
 
-        # DSpark DP graph-shape sync (see model_runner._apply_dspark_shape_max):
-        # DP-max decode bs / ragged token total, so every DP rank replays the same
-        # cudagraph shape. None outside DSpark-under-DP steps.
-        self.dspark_dp_bs = None
-        self.dspark_dp_total_tokens = None
-        # Flat PIECEWISE replay token count for this decode step (set in
-        # prepare_inputs). None when N/A -> consumers use bs*max_seqlen_q.
-        self.dynamic_num_tokens_pad = None
-
         # Detailed attention aggregates (set by Scheduler.compute_detailed_aggregates
         # when profiling is active and ATOM_ENABLE_DETAILED_ANNOTATION is set).
         # None on the normal path; consumed by ModelRunner.run_model to extend
@@ -836,6 +826,10 @@ def produces_output(self) -> bool:
         Decode batches always do. A pure-prefill batch yields a token only
         when at least one seq is on its final chunk; a batch of middle chunks
         produces nothing.
+
+        A DP-sync dummy answers True and must: reporting lags a step, so the
+        dummy is what flushes the last real one. Skip its `postprocess` and
+        the request waking up after it has an anchor on neither side.
         """
         if self.total_seqs_num_decode > 0:
             return True
diff --git a/atom/model_ops/attention_mha.py b/atom/model_ops/attention_mha.py
index cf4601beac..840c960209 100644
--- a/atom/model_ops/attention_mha.py
+++ b/atom/model_ops/attention_mha.py
@@ -819,13 +819,16 @@ def prefill_attention_triton(
             # Raw K/V is fed as a block_size=1 flash-layout cache, never shuffled.
             shuffled_kv_cache = False
 
+        n_block_rows = attn_metadata.block_tables.shape[0]
         unified_attention(
             q,
             k_for_attn,
             v_for_attn,
             o,
-            cu_seqlens_q=attn_metadata.cu_seqlens_q,
-            seqused_k=attn_metadata.context_lens,
+            # Cut to the rows of the block table this indexes with them: the
+            # prefill builder pads both past the requests it scheduled.
+            cu_seqlens_q=attn_metadata.cu_seqlens_q[: n_block_rows + 1],
+            seqused_k=attn_metadata.context_lens[:n_block_rows],
             max_seqlen_q=attn_metadata.max_seqlen_q,
             max_seqlen_k=attn_metadata.max_seqlen_k,
             softmax_scale=self.scale,
diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py
index 8cd3676b69..e7b0816f20 100644
--- a/atom/model_ops/attention_mla.py
+++ b/atom/model_ops/attention_mla.py
@@ -1425,7 +1425,11 @@ def _forward_prefill_cached_chunked(
                     q=prefill_q,
                     k=k_chunk,
                     v=v_chunk,
-                    cu_seqlens_q=attn_metadata.cu_seqlens_q,
+                    # As many q-cums as k-cums -- varlen takes its batch from
+                    # the q side, and the chunk's were built unpadded.
+                    cu_seqlens_q=attn_metadata.cu_seqlens_q[
+                        : chunk_meta.cu_seqlens_k[c].shape[0]
+                    ],
                     cu_seqlens_k=chunk_meta.cu_seqlens_k[c],
                     max_seqlen_q=attn_metadata.max_seqlen_q,
                     max_seqlen_k=chunk_meta.max_seqlen_k[c],
@@ -1550,7 +1554,9 @@ def _dcp_compute_prefill_context(
                 q=prefill_q,
                 k=k,
                 v=v,
-                cu_seqlens_q=attn_metadata.cu_seqlens_q,
+                cu_seqlens_q=attn_metadata.cu_seqlens_q[
+                    : chunk_meta.cu_seqlens_k[c].shape[0]
+                ],
                 cu_seqlens_k=chunk_meta.cu_seqlens_k[c],
                 max_seqlen_q=attn_metadata.max_seqlen_q,
                 max_seqlen_k=chunk_meta.max_seqlen_k[c],
@@ -1732,7 +1738,10 @@ def _forward_prefill_mla(
             device=q.device,
         )
 
-        paged_cu_seqlens_q = attn_metadata.cu_seqlens_q
+        # The paged kernels take their batch from the q-cums; cut them to a
+        # per-seq array the prefill builder left unpadded.
+        n_seqs = attn_metadata.kv_last_page_lens.shape[0]
+        paged_cu_seqlens_q = attn_metadata.cu_seqlens_q[: n_seqs + 1]
         paged_kv_indptr = attn_metadata.kv_indptr
         paged_kv_indices = attn_metadata.kv_indices
         kv_last_page_lens = attn_metadata.kv_last_page_lens
diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py
index b80461cb3f..b6c3d89c45 100644
--- a/atom/model_ops/attentions/aiter_attention.py
+++ b/atom/model_ops/attentions/aiter_attention.py
@@ -873,8 +873,10 @@ def prepare_mtp_decode(
             )
         return workinfos
 
-    def prepare_prefill(self, batch: ScheduledBatch):
-        attn_metadata, positions = CommonAttentionBuilder.prepare_prefill(self, batch)
+    def prepare_prefill(self, batch: ScheduledBatch, running_bs: int):
+        attn_metadata, positions = CommonAttentionBuilder.prepare_prefill(
+            self, batch, running_bs
+        )
         if self._has_sparse_attention and not attn_metadata.has_cached:
             bs = batch.total_seqs_num_prefill
             self.prepare_block_tables(batch)
@@ -905,9 +907,11 @@ def prepare_prefill(self, batch: ScheduledBatch):
                 attn_metadata.context_lens[:bs],
                 bs,
             )
+            # Sliced like the block tables above -- `num_prefills` counts
+            # requests, and these two are padded past it.
             attn_metadata.sparse_attention_metadata = make_sparse_prefill_metadata(
-                cu_seqlens_q=attn_metadata.cu_seqlens_q,
-                seq_lens=attn_metadata.context_lens,
+                cu_seqlens_q=attn_metadata.cu_seqlens_q[: bs + 1],
+                seq_lens=attn_metadata.context_lens[:bs],
                 block_table=sparse_block_tables,
                 slot_mapping=attn_metadata.slot_mapping,
                 max_query_len=attn_metadata.max_seqlen_q,
@@ -1063,11 +1067,16 @@ def _attach_tbo_token_split_straddle_prefix(self, ub_attn, ub_slice):
         )
         ub_attn.max_seqlen_k = int(ctx_lens.max())
 
-    def prepare_decode(self, batch: ScheduledBatch, bs: int):
+    def prepare_decode(
+        self,
+        batch: ScheduledBatch,
+        running_bs: int,
+        running_tokens: int,
+        max_seqlen_q: int,
+    ):
         scheduled_bs = batch.total_seqs_num_decode
         self.total_blocks = 0
         dropout_p = 0.0
-        max_seqlen_q = batch.num_spec_step + 1
         min_seqlen_q = 0
 
         context_lens = np.asarray(batch.context_lens, dtype=np.int32)
@@ -1103,7 +1112,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
 
         var = self.model_runner.forward_vars
         sum_scheduled_tokens = batch.total_tokens_num_decode
-        var["slot_mapping"].np[: bs * max_seqlen_q] = -1
+        var["slot_mapping"].np[:running_tokens] = -1
         if not batch.is_dummy_run:
             var["slot_mapping"].np[:sum_scheduled_tokens] = slot_mapping[
                 :sum_scheduled_tokens
@@ -1111,7 +1120,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
 
         var["positions"].np[:sum_scheduled_tokens] = positions
         var["context_lens"].np[:scheduled_bs] = context_lens
-        var["context_lens"].np[scheduled_bs:bs] = 0
+        var["context_lens"].np[scheduled_bs:running_bs] = 0
 
         # Prepare kv_indptr and kv_indices for persistent attention
         num_blocks_per_seq = cdiv(context_lens, self.block_size)
@@ -1120,19 +1129,19 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
 
         var["kv_indptr"].np[0] = 0
         var["kv_indptr"].np[1 : scheduled_bs + 1] = kv_indptr
-        var["kv_indptr"].np[scheduled_bs + 1 : bs + 1] = sum_blocks
+        var["kv_indptr"].np[scheduled_bs + 1 : running_bs + 1] = sum_blocks
 
         vars_used = [
-            ("slot_mapping", bs * max_seqlen_q),
-            ("context_lens", bs),
-            ("cu_seqlens_q", bs + 1),
-            ("block_tables", bs),
-            ("kv_indptr", bs + 1),
+            ("slot_mapping", running_tokens),
+            ("context_lens", running_bs),
+            ("cu_seqlens_q", running_bs + 1),
+            ("block_tables", running_bs),
+            ("kv_indptr", running_bs + 1),
         ]
 
         ctx = {el: var[el].copy_to_gpu(num) for el, num in vars_used}
         if self.block_size in (256, 1024):
-            ctx_pa_ps = self.set_aiter_persistent_worker_buffers(bs)
+            ctx_pa_ps = self.set_aiter_persistent_worker_buffers(running_bs)
             ctx.update(ctx_pa_ps)
 
         ctx["kv_indices"] = var["kv_indices"].gpu
@@ -1178,10 +1187,10 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
             positions = mrope_positions
         else:
             positions = var["positions"].copy_to_gpu(sum_scheduled_tokens)
-        if self.model_runner.config.enable_tbo_decode and bs >= 2:
+        if self.model_runner.config.enable_tbo_decode and running_bs >= 2:
             self._prepare_ubatch_decode(
                 scheduled_bs,
-                bs,
+                running_bs,
                 max_seqlen_q,
                 context_lens,
             )
diff --git a/atom/model_ops/attentions/aiter_mla.py b/atom/model_ops/attentions/aiter_mla.py
index 575b836412..dcc3810ab8 100644
--- a/atom/model_ops/attentions/aiter_mla.py
+++ b/atom/model_ops/attentions/aiter_mla.py
@@ -1226,8 +1226,10 @@ def _build_dcp_indexer_prefill_meta(self, attn_metadata, bs: int, counts, var):
             src.astype(np.int32)
         ).to(dev, non_blocking=True)
 
-    def prepare_prefill(self, batch: ScheduledBatch):
-        attn_metadata, positions = CommonAttentionBuilder.prepare_prefill(self, batch)
+    def prepare_prefill(self, batch: ScheduledBatch, running_bs: int):
+        attn_metadata, positions = CommonAttentionBuilder.prepare_prefill(
+            self, batch, running_bs
+        )
         bs = batch.total_seqs_num_prefill
         sum_scheduled_tokens = batch.total_tokens_num_prefill
         var = self.model_runner.forward_vars
@@ -1359,8 +1361,9 @@ def prepare_prefill(self, batch: ScheduledBatch):
             attn_metadata.kv_indices = var["kv_indices"].gpu
             attn_metadata.kv_indptr = var["kv_indptr"].gpu[: bs + 1]
             attn_metadata.kv_indptr[0] = 0
+            # `context_lens` is padded past the requests this indptr counts.
             attn_metadata.kv_indptr[1 : bs + 1] = torch.cumsum(
-                attn_metadata.context_lens, 0
+                attn_metadata.context_lens[:bs], 0
             )
             attn_metadata.kv_last_page_lens = var["kv_last_page_lens"].gpu[:bs]
 
@@ -1734,12 +1737,15 @@ def _dcp_round_robin_slot(self, block_table, pos: int) -> int:
             + dcp_local_index(pos, W, S) % block_size
         )
 
-    def prepare_decode(self, batch: ScheduledBatch, bs: int):
+    def prepare_decode(
+        self,
+        batch: ScheduledBatch,
+        running_bs: int,
+        running_tokens: int,
+        max_seqlen_q: int,
+    ):
         scheduled_bs = batch.total_seqs_num_decode
         dropout_p = 0.0
-        max_seqlen_q = 1
-        if hasattr(self.model_runner, "drafter"):
-            max_seqlen_q = self.model_runner.drafter.mtp_k + 1
 
         var = self.model_runner.forward_vars
         context_lens = np.asarray(batch.context_lens, dtype=np.int32)
@@ -1790,12 +1796,12 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
 
         # Use scheduled_bs since in dummy run, total_seqs_num_decode is 1.
         sum_scheduled_tokens = scheduled_bs * max_seqlen_q
-        var["slot_mapping"].np[: bs * max_seqlen_q] = -1
+        var["slot_mapping"].np[:running_tokens] = -1
         if not batch.is_dummy_run:
             var["slot_mapping"].np[:sum_scheduled_tokens] = slot_mapping
         var["positions"].np[:sum_scheduled_tokens] = positions
         var["context_lens"].np[:scheduled_bs] = context_lens
-        var["context_lens"].np[scheduled_bs:bs] = 0
+        var["context_lens"].np[scheduled_bs:running_bs] = 0
 
         if self.dcp_world_size > 1:
             from atom.model_ops.dcp_ops import get_dcp_local_seq_lens
@@ -1809,7 +1815,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
             # Publish it: the sparse indexer used to re-derive this on device with
             # 8 elementwise kernels per full-index layer.
             var["dcp_local_context_lens"].np[:scheduled_bs] = local_context_lens
-            var["dcp_local_context_lens"].np[scheduled_bs:bs] = 0
+            var["dcp_local_context_lens"].np[scheduled_bs:running_bs] = 0
             num_blocks_per_seq = cdiv(local_context_lens, self.block_size)
         elif any(batch.is_first_decode_without_local_prefill):
             num_blocks_per_seq = [
@@ -1832,7 +1838,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
 
         self.prepare_block_tables(batch)
         var["kv_indptr"].np[1 : scheduled_bs + 1] = kv_indptr
-        var["kv_indptr"].np[scheduled_bs + 1 : bs + 1] = sum_blocks
+        var["kv_indptr"].np[scheduled_bs + 1 : running_bs + 1] = sum_blocks
         if self.dcp_world_size > 1:
             # Global (un-sharded) token-level kv_indptr for the round-robin CP
             # causal mask (page_size=1 -> token granularity). context_lens is the
@@ -1840,7 +1846,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
             g_kv_indptr = np.cumsum(context_lens[:scheduled_bs], dtype=np.int32)
             var["g_kv_indptr"].np[0] = 0
             var["g_kv_indptr"].np[1 : scheduled_bs + 1] = g_kv_indptr
-            var["g_kv_indptr"].np[scheduled_bs + 1 : bs + 1] = (
+            var["g_kv_indptr"].np[scheduled_bs + 1 : running_bs + 1] = (
                 g_kv_indptr[-1] if scheduled_bs > 0 else 0
             )
         if self.dcp_world_size > 1 and self.block_size != 1:
@@ -1854,14 +1860,14 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
             var["kv_last_page_lens"].np[:scheduled_bs] = (
                 batch.last_block_num_tokens if self.block_size != 1 else 1
             )
-        var["kv_last_page_lens"].np[scheduled_bs:bs] = 0
+        var["kv_last_page_lens"].np[scheduled_bs:running_bs] = 0
         vars_used = [
-            ("slot_mapping", bs * max_seqlen_q),
-            ("context_lens", bs),
-            ("cu_seqlens_q", bs + 1),
-            # ("kv_indptr", bs + 1),
-            ("kv_last_page_lens", bs),
-            ("block_tables", bs),
+            ("slot_mapping", running_tokens),
+            ("context_lens", running_bs),
+            ("cu_seqlens_q", running_bs + 1),
+            # ("kv_indptr", running_bs + 1),
+            ("kv_last_page_lens", running_bs),
+            ("block_tables", running_bs),
         ]
         metadata_deps = {
             "cu_seqlens_q",
@@ -1884,7 +1890,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
                 var["sparse_kv_indptr"].np[1 : sum_scheduled_tokens + 1] = np.cumsum(
                     sparse_per_token_lens, dtype=np.int32
                 )
-                sum_tokens = bs * max_seqlen_q
+                sum_tokens = running_tokens
                 var["sparse_kv_indptr"].np[
                     sum_scheduled_tokens + 1 : sum_tokens + 1
                 ] = var["sparse_kv_indptr"].np[sum_scheduled_tokens]
@@ -1893,15 +1899,15 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
                 metadata_deps.add("sparse_kv_indptr")
             else:
                 sparse_context_lens = np.clip(
-                    var["context_lens"].np[:bs], None, index_topk
+                    var["context_lens"].np[:running_bs], None, index_topk
                 )
-                var["sparse_kv_indptr"].np[1 : bs + 1] = np.cumsum(
+                var["sparse_kv_indptr"].np[1 : running_bs + 1] = np.cumsum(
                     sparse_context_lens, dtype=np.int32
                 )
-                var["sparse_kv_indptr"].np[scheduled_bs : bs + 1] = var[
+                var["sparse_kv_indptr"].np[scheduled_bs : running_bs + 1] = var[
                     "sparse_kv_indptr"
                 ].np[scheduled_bs]
-                vars_used.append(("sparse_kv_indptr", bs + 1))
+                vars_used.append(("sparse_kv_indptr", running_bs + 1))
                 metadata_deps.add("sparse_kv_indptr")
 
         vars_for_metadata = [(el, num) for el, num in vars_used if el in metadata_deps]
@@ -1916,7 +1922,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
         # async overlap.
         disagg = self.model_runner.config.enable_rapidserve
         ctx = {}
-        ctx["kv_indptr"] = var["kv_indptr"].copy_to_gpu(bs + 1)
+        ctx["kv_indptr"] = var["kv_indptr"].copy_to_gpu(running_bs + 1)
         if disagg:
             ctx_rest = {el: var[el].copy_to_gpu(num) for el, num in vars_remaining}
             ctx.update(ctx_rest)
@@ -1950,8 +1956,10 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
         ctx.update({el: var[el].copy_to_gpu(num) for el, num in vars_for_metadata})
 
         if is_sparse_mtp:
-            sum_tokens = bs * max_seqlen_q
-            ctx_mla_ps = self.set_mla_persistent_worker_buffers(bs, max_seqlen_q)
+            sum_tokens = running_tokens
+            ctx_mla_ps = self.set_mla_persistent_worker_buffers(
+                running_bs, max_seqlen_q
+            )
             ctx_mla_ps_sparse = self._set_mla_persistent_worker_buffers_sparse_mtp(
                 sum_tokens
             )
@@ -1961,13 +1969,15 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
             # masks on global positions g(j)=j*W+r instead of local-causal trim.
             cp_round_robin = self.dcp_world_size > 1 and max_seqlen_q > 1
             # For cprr, build metadata over the REAL requests (scheduled_bs), not
-            # the cudagraph-padded bs. The padded tail is 0-query, which the cprr
+            # the cudagraph-padded running_bs. The padded tail is 0-query, which the cprr
             # metadata kernel mishandles (num_qo_tiles=0 -> div-by-zero; and its
             # is_cp_round_robin "no trim" is overridden by the nhead=128 M-fold
             # qk_batch_ratio branch). The persistent kernel is work_indptr-driven,
-            # so scheduled_bs metadata naturally skips the padding rows. Eager has
-            # scheduled_bs == bs, so this is a no-op there.
-            meta_bs = scheduled_bs if cp_round_robin else bs
+            # so scheduled_bs metadata naturally skips the padding rows. A step
+            # that padded nothing reaches here with the two equal, and then this
+            # is a no-op -- but eager no longer implies that, since a DP-unified
+            # batch pads off the graph path too.
+            meta_bs = scheduled_bs if cp_round_robin else running_bs
             ctx_mla_ps = self.set_mla_persistent_worker_buffers(
                 meta_bs, max_seqlen_q, is_cp_round_robin=cp_round_robin
             )
@@ -1987,12 +1997,14 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
         # non-DCP / qlen=1 paths keep the plain kernel). Consumed by
         # _forward_decode -> mla_decode_fwd for the MTP (max_q_len>1) cprr mask.
         attn_metadata.g_kv_indptr = (
-            var["g_kv_indptr"].copy_to_gpu(bs + 1) if self.dcp_world_size > 1 else None
+            var["g_kv_indptr"].copy_to_gpu(running_bs + 1)
+            if self.dcp_world_size > 1
+            else None
         )
         # Layer-invariant local KV lengths for the DCP sparse indexer (see the
         # buffer's declaration). None off DCP so the consumer falls back.
         attn_metadata.dcp_local_context_lens = (
-            var["dcp_local_context_lens"].copy_to_gpu(bs)
+            var["dcp_local_context_lens"].copy_to_gpu(running_bs)
             if self.dcp_world_size > 1
             else None
         )
@@ -2002,7 +2014,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
                 setattr(attn_metadata, k, v)
 
         if is_sparse_mtp:
-            sum_tokens = bs * max_seqlen_q
+            sum_tokens = running_tokens
             attn_metadata.sparse_cu_seqlens_q = var["sparse_cu_seqlens_q"].gpu[
                 : sum_tokens + 1
             ]
@@ -2022,13 +2034,14 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
             # the kernel over-read past the written sparse-index region).
             attn_metadata.sparse_kv_last_page_lens = var[
                 "sparse_kv_last_page_lens"
-            ].gpu[:bs]
+            ].gpu[:running_bs]
 
-        # Use bs (running_bs) >= 2 instead of scheduled_bs >= 2 to avoid accuracy issue:
-        if self.model_runner.config.enable_tbo_decode and bs >= 2:
+        # running_bs, not scheduled_bs: the padded rows have to be split into the
+        # ubatches too, or accuracy drifts.
+        if self.model_runner.config.enable_tbo_decode and running_bs >= 2:
             self._prepare_ubatch_decode(
                 scheduled_bs,
-                bs,
+                running_bs,
                 max_seqlen_q,
                 context_lens,
             )
diff --git a/atom/model_ops/attentions/backends.py b/atom/model_ops/attentions/backends.py
index a466af4d97..76f6659bfa 100644
--- a/atom/model_ops/attentions/backends.py
+++ b/atom/model_ops/attentions/backends.py
@@ -94,15 +94,27 @@ def __init__(self, block_size: int) -> None:
         raise NotImplementedError
 
     @abstractmethod
-    def prepare_decode(self, batch: ScheduledBatch, bs: int):
+    def prepare_decode(
+        self,
+        batch: ScheduledBatch,
+        running_bs: int,
+        running_tokens: int,
+        max_seqlen_q: int,
+    ):
         raise NotImplementedError
 
     @abstractmethod
-    def prepare_prefill(self, batch: ScheduledBatch):
+    def prepare_prefill(self, batch: ScheduledBatch, running_bs: int):
         raise NotImplementedError
 
     @abstractmethod
-    def build(self, batch: ScheduledBatch, bs: int):
+    def build(
+        self,
+        batch: ScheduledBatch,
+        running_bs: int,
+        running_tokens: int,
+        max_seqlen_q: int,
+    ):
         raise NotImplementedError
 
     def prepare_mtp_decode(
@@ -446,8 +458,8 @@ def _build_mrope_decode_positions(
 
         return self._copy_mrope_to_gpu(total_tokens)
 
-    def prepare_prefill(self, batch: ScheduledBatch):
-        bs = batch.total_seqs_num_prefill
+    def prepare_prefill(self, batch: ScheduledBatch, running_bs: int):
+        scheduled_bs = batch.total_seqs_num_prefill
         sum_scheduled_tokens = batch.total_tokens_num_prefill
         var = self.model_runner.forward_vars
         positions = []
@@ -458,8 +470,8 @@ def prepare_prefill(self, batch: ScheduledBatch):
         slot_mapping = []
         has_cached = False
         # seqs = list(batch.seqs.values())
-        # seqs = seqs[:bs]
-        for i in range(bs):
+        # seqs = seqs[:scheduled_bs]
+        for i in range(scheduled_bs):
             seqlen = batch.context_lens[i]
             cached_seqlen = batch.num_cached_tokens[i]
             if cached_seqlen > 0:
@@ -526,28 +538,38 @@ def prepare_prefill(self, batch: ScheduledBatch):
         var["positions"].np[:sum_scheduled_tokens] = positions
         var["slot_mapping"].np[:sum_scheduled_tokens] = -1
         var["slot_mapping"].np[: len(slot_mapping)] = slot_mapping
-        var["cu_seqlens_q"].np[: bs + 1] = cu_seqlens_q
-        var["cu_seqlens_k"].np[: bs + 1] = cu_seqlens_k
-        var["context_lens"].np[:bs] = batch.context_lens[:bs]
+        var["cu_seqlens_q"].np[: scheduled_bs + 1] = cu_seqlens_q
+        var["cu_seqlens_k"].np[: scheduled_bs + 1] = cu_seqlens_k
+        var["context_lens"].np[:scheduled_bs] = batch.context_lens[:scheduled_bs]
+        # Pad the per-sequence tail out to `running_bs`, the width a draft pass
+        # that follows this step runs at and the one `prepare_decode` fills to.
+        # A flat cumsum gives each fabricated row zero query length and a zero
+        # context makes it read nothing; left unwritten they hold the previous
+        # step's, and the drafter attends to a since-freed request's blocks.
+        var["cu_seqlens_q"].np[scheduled_bs + 1 : running_bs + 1] = cu_seqlens_q[-1]
+        var["cu_seqlens_k"].np[scheduled_bs + 1 : running_bs + 1] = cu_seqlens_k[-1]
+        var["context_lens"].np[scheduled_bs:running_bs] = 0
         min_seqlen_q = 0
         dropout_p = 0.0
         vars_used = [
-            ("cu_seqlens_q", bs + 1),
-            ("cu_seqlens_k", bs + 1),
+            ("cu_seqlens_q", running_bs + 1),
+            ("cu_seqlens_k", running_bs + 1),
             ("slot_mapping", sum_scheduled_tokens),
-            ("context_lens", bs),
+            ("context_lens", running_bs),
         ]
         if has_cached:
-            vars_used.append(("block_tables", bs))
-            vars_used.append(("seq_starts", bs))
+            vars_used.append(("block_tables", scheduled_bs))
+            vars_used.append(("seq_starts", scheduled_bs))
 
         ctx = {el: var[el].copy_to_gpu(num) for el, num in vars_used}
         num_cached_tokens = None
         if has_cached:
             num_cached_tokens = torch.tensor(
-                batch.num_cached_tokens[:bs], dtype=torch.int32, pin_memory=True
+                batch.num_cached_tokens[:scheduled_bs],
+                dtype=torch.int32,
+                pin_memory=True,
             ).cuda(non_blocking=True)
-            total_tokens = sum(batch.context_lens[:bs])
+            total_tokens = sum(batch.context_lens[:scheduled_bs])
         total_kv = total_tokens if has_cached else sum_scheduled_tokens
         attn_metadata = AttentionMetaData(
             # Cast to python int — numpy.int32 leaks in via batch.context_lens
@@ -601,7 +623,29 @@ def _attach_tbo_prefill_cpu_lens(
             attn_metadata, "cu_seqlens_k", var["cu_seqlens_k"].np[: bs + 1].copy()
         )
 
-    def build(self, batch: ScheduledBatch, bs: int):
+    def build(
+        self,
+        batch: ScheduledBatch,
+        running_bs: int,
+        running_tokens: int,
+        max_seqlen_q: int,
+    ):
+        """Build this step's metadata at the shape `prepare_inputs` settled.
+
+        The step's two units and nothing else: `running_bs` sizes everything
+        per-sequence, `running_tokens` everything per-row. Both arrive as
+        arguments because neither is derivable from `batch` -- the scheduler
+        built it before the DP sync and before any graph was picked -- and
+        neither is derivable from the other, since a ragged step replays a flat
+        token bucket rather than a `running_bs * q` grid.
+
+        Prefill takes only the first: its rows are the prompt's, so there is no
+        `running_tokens` to hand it, but its per-sequence tail still pads to
+        `running_bs` for the draft pass that follows. That padding is visible
+        past attention -- `cu_seqlens_q` doubles as the selector for the
+        sampler's rows, so the LM head emits `running_bs` of them. `postprocess`
+        cuts back to the scheduled batch, and that is the one place it does.
+        """
         # Run state maintenance on the compute stream before the forward.
         state_ops = batch.state_maintenance_ops
         if state_ops.relocations:
@@ -612,9 +656,9 @@ def build(self, batch: ScheduledBatch, bs: int):
             )
         is_prefill = batch.total_tokens_num_prefill > 0
         if is_prefill:
-            return self.prepare_prefill(batch)
+            return self.prepare_prefill(batch, running_bs)
         else:
-            return self.prepare_decode(batch, bs)
+            return self.prepare_decode(batch, running_bs, running_tokens, max_seqlen_q)
 
 
 class AttentionImpl(nn.Module):
diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py
index 20ac25987f..d9bce39acb 100644
--- a/atom/model_ops/attentions/deepseek_v4_attn.py
+++ b/atom/model_ops/attentions/deepseek_v4_attn.py
@@ -2454,7 +2454,13 @@ def prepare_mtp_decode(
         #   - v4 indexer meta (Indexer — only present when ratio == 4)
         return {}
 
-    def prepare_decode(self, batch: ScheduledBatch, bs: int):
+    def prepare_decode(
+        self,
+        batch: ScheduledBatch,
+        running_bs: int,
+        running_tokens: int,
+        max_seqlen_q: int,
+    ):
         """V4-style decode prep: populates positions, cu_seqlens_q,
         block_tables, and state_slot_out.
 
@@ -2465,11 +2471,10 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
         var = self.model_runner.forward_vars
         scheduled_bs = batch.total_seqs_num_decode
         context_lens_np = np.asarray(batch.context_lens, dtype=np.int32)
-        # Per-seq decode forward length: single source of truth on the batch
-        # (= num_spec_step+1 for plain MTP, or the DSpark q-bucket when shrunk).
-        # positions/attn use this so the (bs, q) graph is selected. See
-        # ScheduledBatch.num_spec_query_tokens.
-        max_seqlen_q = getattr(batch, "num_spec_query_tokens", batch.num_spec_step + 1)
+        # Per-seq decode forward length, settled by the step rather than carried
+        # on the batch (= num_spec_step+1 for plain MTP, or the DSpark q-bucket
+        # when shrunk). positions/attn use this so the (running_bs, q) graph is
+        # selected. See `ForwardMode.max_seqlen_q`.
         # MTP: roll back ctx by `num_rejected` so this fwd's positions overwrite
         # last fwd's rejected-draft slots (matches aiter_mla.py:701 /
         # aiter_attention.py:542). `batch.context_lens` = `seq.num_tokens`
@@ -2509,15 +2514,10 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
             ) + np.repeat(context_lens_np - full_q, max_seqlen_q)
         sum_scheduled_tokens = batch.total_tokens_num_decode
 
-        # Pad positions up to the graph's replayed token count. Prefer the real
-        # flat bucket run_model actually replays (dynamic_num_tokens_pad); else
-        # bs*max_seqlen_q (non-ragged / non-piecewise).
-        dynamic_num_tokens_pad = getattr(batch, "dynamic_num_tokens_pad", None)
-        graph_cap_tokens = (
-            int(dynamic_num_tokens_pad)
-            if dynamic_num_tokens_pad is not None
-            else int(bs) * int(max_seqlen_q)
-        )
+        # Pad positions to the rows this forward runs. That is
+        # `running_tokens` by definition -- a ragged step's flat bucket and a
+        # rectangular step's `running_bs * q` are the same field.
+        graph_cap_tokens = int(running_tokens)
         if graph_cap_tokens > sum_scheduled_tokens:
             _pad_positions = np.zeros(graph_cap_tokens, dtype=positions_np.dtype)
             _pad_positions[:sum_scheduled_tokens] = positions_np
@@ -2543,13 +2543,14 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
         si_buf.np[:scheduled_bs] = self._state_slot_in_np(
             batch, scheduled_bs, state_slot_np
         )
-        # Published at the PADDED `bs`, like the ubatch path below already does,
-        # so a consumer that runs the padded batch -- a speculative drafter -- can
-        # slice to it. Every reader either masks the pad tail out
-        # (`batch_id_per_token = -1`) or discards those rows, so 0 is as good a
-        # filler here as it is there.
-        ss_buf.np[scheduled_bs:bs] = 0
-        si_buf.np[scheduled_bs:bs] = 0
+        # Published at the PADDED `running_bs`, like the ubatch path below, so
+        # a consumer that runs the padded batch -- a speculative drafter -- can
+        # slice to it; `cu_seqlens_q` is padded to the same batch, so a consumer
+        # inferring bs from these slots stays consistent. Every reader either
+        # masks the pad tail out (`batch_id_per_token = -1`) or discards those
+        # rows, so 0 is a legal filler.
+        ss_buf.np[scheduled_bs:running_bs] = 0
+        si_buf.np[scheduled_bs:running_bs] = 0
 
         # ---- fire H2D on prep_stream ----
         # NB: this runs inside attn_metadata_builder.build(), BEFORE
@@ -2559,11 +2560,11 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
         prep_stream.wait_stream(current_stream)
         with torch.cuda.stream(prep_stream):
             positions = var["positions"].copy_to_gpu(sum_scheduled_tokens_padded)
-            cu_seqlens_q_gpu = var["cu_seqlens_q"].copy_to_gpu(bs + 1)
+            cu_seqlens_q_gpu = var["cu_seqlens_q"].copy_to_gpu(running_bs + 1)
             context_lens_gpu = var["context_lens"].copy_to_gpu(scheduled_bs)
             block_tables_gpu = var["block_tables"].copy_to_gpu(scheduled_bs)
-            state_slot_gpu = ss_buf.copy_to_gpu(bs)
-            state_slot_in_gpu = si_buf.copy_to_gpu(bs)
+            state_slot_gpu = ss_buf.copy_to_gpu(running_bs)
+            state_slot_in_gpu = si_buf.copy_to_gpu(running_bs)
 
         # ---- CPU numpy work, overlapped with prep_stream H2D ----
         # RAGGED: per-seq extend lengths (else uniform max_seqlen_q). compress
@@ -2576,7 +2577,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
         compress_plans = self._build_compress_plans(
             extend_lens_np,
             context_lens_np,
-            running_bs=bs,
+            running_bs=running_bs,
             max_q_len=max_seqlen_q,
         )
 
@@ -2602,7 +2603,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
         attn_metadata.state_slot_out_cpu = state_slot_np
         attn_metadata.compress_plans = compress_plans
         # DSpark RAGGED: pass per-seq verify lengths + full_q to the (rectangular-
-        # only) decode indexer so it can pad Q back to [bs, full_q].
+        # only) decode indexer so it can pad Q back to [running_bs, full_q].
         _drafter = getattr(self.model_runner, "drafter", None)
         _dspark_ragged_graph = (
             self.model_runner.config.dspark.ragged
@@ -2613,11 +2614,11 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
             # Pinned staging, not `torch.as_tensor(np, device=cuda)`: that is a
             # pageable H2D and syncs here, which was the ragged decode bubble.
             attn_metadata.dspark_ragged_lens_gpu = self._stage_dspark_ragged_lens(
-                extend_lens_np, self._dspark_ragged_lens_pad_to(bs)
+                extend_lens_np, self._dspark_ragged_lens_pad_to(running_bs)
             )
             attn_metadata.dspark_full_q = int(full_q)
 
-        running_bs = int(bs)
+        running_bs = int(running_bs)
         self._attach_v4_per_fwd_meta(
             attn_metadata,
             extend_lens_np,  # = np.full(scheduled_bs, max_seqlen_q) for decode
@@ -2643,7 +2644,7 @@ def prepare_decode(self, batch: ScheduledBatch, bs: int):
         ):
             self._prepare_ubatch_decode(
                 scheduled_bs=scheduled_bs,
-                bs=bs,
+                running_bs=running_bs,
                 max_seqlen_q=max_seqlen_q,
                 context_lens_np=context_lens_np,
                 state_slot_np=state_slot_np,
@@ -2660,7 +2661,7 @@ def _prepare_ubatch_decode(
         self,
         *,
         scheduled_bs: int,
-        bs: int,
+        running_bs: int,
         max_seqlen_q: int,
         context_lens_np: np.ndarray,
         state_slot_np: np.ndarray,
@@ -2695,7 +2696,8 @@ def _prepare_ubatch_decode(
 
             ctx = get_forward_context()
             padded_list = [
-                UBatchWrapper._decode_ub_running_bs(ctx, i, N, bs) for i in range(N)
+                UBatchWrapper._decode_ub_running_bs(ctx, i, N, running_bs)
+                for i in range(N)
             ]
             # Real-request ranges partition scheduled_bs; each ubatch owns up to
             # its padded capacity, the tail ubatch takes the remainder. Pad rows
@@ -2716,7 +2718,10 @@ def _prepare_ubatch_decode(
         np.cumsum(extend_lens_np[:scheduled_bs], out=token_offsets[1:])
         for ub_idx, (req_start, req_end) in enumerate(ub_ranges):
             p = f"ub{ub_idx}_"
-            running_bs = padded_list[ub_idx]
+            # THIS ubatch's width, not the step's -- the two are different
+            # numbers and both are called a running_bs, so they get distinct
+            # names (cf. `UBatchWrapper._decode_ub_running_bs`, which supplies it).
+            ub_running_bs = padded_list[ub_idx]
             # Real requests that fall into this ubatch's [req_start, req_end),
             # clamped to scheduled_bs (cudagraph pad rows beyond scheduled_bs
             # carry sentinels, exercised only during capture's synthetic batch).
@@ -2729,53 +2734,58 @@ def _prepare_ubatch_decode(
             # ---- per-seq slices into ub buffers ----
             ub_ctx_np = context_lens_np[req_start : req_start + ub_real_reqs]
             var[f"{p}context_lens"].np[:ub_real_reqs] = ub_ctx_np
-            var[f"{p}context_lens"].np[ub_real_reqs:running_bs] = 0
+            var[f"{p}context_lens"].np[ub_real_reqs:ub_running_bs] = 0
 
             ub_state_np = state_slot_np[req_start : req_start + ub_real_reqs]
             if len(ub_state_np) < ub_real_reqs:
                 ub_state_np = np.zeros(ub_real_reqs, dtype=np.int32)
             var[f"{p}v4_meta_state_slot_out"].np[:ub_real_reqs] = ub_state_np
-            var[f"{p}v4_meta_state_slot_out"].np[ub_real_reqs:running_bs] = 0
-            state_slot_np_ub = var[f"{p}v4_meta_state_slot_out"].np[:running_bs].copy()
+            var[f"{p}v4_meta_state_slot_out"].np[ub_real_reqs:ub_running_bs] = 0
+            state_slot_np_ub = (
+                var[f"{p}v4_meta_state_slot_out"].np[:ub_running_bs].copy()
+            )
             ub_state_in_np = state_slot_in_np[req_start : req_start + ub_real_reqs]
             if len(ub_state_in_np) < ub_real_reqs:
                 ub_state_in_np = np.zeros(ub_real_reqs, dtype=np.int32)
             var[f"{p}v4_meta_state_slot_in"].np[:ub_real_reqs] = ub_state_in_np
-            var[f"{p}v4_meta_state_slot_in"].np[ub_real_reqs:running_bs] = 0
+            var[f"{p}v4_meta_state_slot_in"].np[ub_real_reqs:ub_running_bs] = 0
 
             var[f"{p}block_tables"].np[:ub_real_reqs] = var["block_tables"].np[
                 req_start : req_start + ub_real_reqs
             ]
-            var[f"{p}block_tables"].np[ub_real_reqs:running_bs] = 0
+            var[f"{p}block_tables"].np[ub_real_reqs:ub_running_bs] = 0
 
             # positions: copy the ubatch's token slice (values match the global
             # positions slice the UBatchWrapper Context will expose).
+            ub_running_tokens = ub_running_bs * max_seqlen_q
             ub_positions_np = positions_np[tok_start : tok_start + ub_real_tokens]
             var[f"{p}positions"].np[:ub_real_tokens] = ub_positions_np
-            var[f"{p}positions"].np[ub_real_tokens : running_bs * max_seqlen_q] = 0
+            var[f"{p}positions"].np[ub_real_tokens:ub_running_tokens] = 0
 
             # cu_seqlens_q: exact per-request lengths, padded tail flat.
             cu = np.zeros(ub_real_reqs + 1, dtype=np.int32)
             np.cumsum(ub_extend_lens_np, out=cu[1:])
             var[f"{p}cu_seqlens_q"].np[: ub_real_reqs + 1] = cu
             var[f"{p}cu_seqlens_q"].np[
-                ub_real_reqs + 1 : running_bs + 1
+                ub_real_reqs + 1 : ub_running_bs + 1
             ] = ub_real_tokens
 
             # ---- H2D ----
             ub_sum_tokens = max(ub_real_tokens, 1)
-            positions_gpu = var[f"{p}positions"].copy_to_gpu(running_bs * max_seqlen_q)
-            cu_seqlens_q_gpu = var[f"{p}cu_seqlens_q"].copy_to_gpu(running_bs + 1)
-            context_lens_gpu = var[f"{p}context_lens"].copy_to_gpu(running_bs)
-            block_tables_gpu = var[f"{p}block_tables"].copy_to_gpu(running_bs)
-            state_slot_gpu = var[f"{p}v4_meta_state_slot_out"].copy_to_gpu(running_bs)
+            positions_gpu = var[f"{p}positions"].copy_to_gpu(ub_running_tokens)
+            cu_seqlens_q_gpu = var[f"{p}cu_seqlens_q"].copy_to_gpu(ub_running_bs + 1)
+            context_lens_gpu = var[f"{p}context_lens"].copy_to_gpu(ub_running_bs)
+            block_tables_gpu = var[f"{p}block_tables"].copy_to_gpu(ub_running_bs)
+            state_slot_gpu = var[f"{p}v4_meta_state_slot_out"].copy_to_gpu(
+                ub_running_bs
+            )
 
             # ---- compress plans (per ubatch buffer set) ----
             ctx_for_plan = context_lens_np[req_start : req_start + ub_real_reqs]
             compress_plans = self._build_compress_plans(
                 ub_extend_lens_np,
                 ctx_for_plan,
-                running_bs=running_bs,
+                running_bs=ub_running_bs,
                 max_q_len=max_seqlen_q,
                 buf_prefix_ubatch=p,
             )
@@ -2796,12 +2806,12 @@ def _prepare_ubatch_decode(
             )
             attn_metadata.state_slot_out = state_slot_gpu
             attn_metadata.state_slot_in = var[f"{p}v4_meta_state_slot_in"].copy_to_gpu(
-                running_bs
+                ub_running_bs
             )
             attn_metadata.state_slot_out_cpu = state_slot_np_ub
             attn_metadata.compress_plans = compress_plans
             if dspark_ragged:
-                ragged_lens_buf = np.zeros(running_bs, dtype=np.int32)
+                ragged_lens_buf = np.zeros(ub_running_bs, dtype=np.int32)
                 ragged_lens_buf[:ub_real_reqs] = ub_extend_lens_np
                 attn_metadata.dspark_ragged_lens_gpu = torch.as_tensor(
                     ragged_lens_buf, device=positions_gpu.device
@@ -2814,7 +2824,7 @@ def _prepare_ubatch_decode(
                 state_slot_np_ub,
                 ub_real_reqs,
                 ub_real_tokens,
-                running_bs=running_bs,
+                running_bs=ub_running_bs,
                 max_q_len=max_seqlen_q,
                 buf_prefix_ubatch=p,
             )
@@ -2837,7 +2847,7 @@ def build_ubatch_metadata(
         )
         return self._ubatch_decode_meta[ubatch_idx]
 
-    def prepare_prefill(self, batch: ScheduledBatch):
+    def prepare_prefill(self, batch: ScheduledBatch, running_bs: int):
         """V4 prefill prep: extends parent to always populate block_tables
         and state_slot_out.
 
@@ -2848,7 +2858,7 @@ def prepare_prefill(self, batch: ScheduledBatch):
         Also publishes CPU mirrors (`v4_*_cpu`) consumed by the V4 forward
         path to avoid `.item()` / `.tolist()` syncs (PR-A Phase 2).
         """
-        base_md, positions = super().prepare_prefill(batch)
+        base_md, positions = super().prepare_prefill(batch, running_bs)
         # Promote to V4 typed metadata so V4-specific attribute assignments
         # below are well-typed. Safe because AttentionMetaData_DSV4 only adds
         # fields with defaults; the parent dataclass is non-slotted.
@@ -2863,11 +2873,11 @@ def prepare_prefill(self, batch: ScheduledBatch):
                 batch, scheduled_bs
             )
         state_slot_gpu, state_slot_np = self._populate_state_slot_mappings(
-            batch, scheduled_bs, return_cpu=True
+            batch, scheduled_bs, running_bs, return_cpu=True
         )
         attn_metadata.state_slot_out = state_slot_gpu
         attn_metadata.state_slot_in = self._populate_state_slot_in(
-            batch, scheduled_bs, state_slot_np
+            batch, scheduled_bs, running_bs, state_slot_np
         )
         # PR-A Phase 2 CPU mirrors (generic, not V4-specific). The parent
         # populated forward_vars CPU buffers; read them back as numpy slices.
@@ -4114,9 +4124,13 @@ def _populate_block_tables(
         return self.model_runner.forward_vars["block_tables"].copy_to_gpu(scheduled_bs)
 
     def _populate_state_slot_mappings(
-        self, batch: ScheduledBatch, scheduled_bs: int, return_cpu: bool = False
+        self,
+        batch: ScheduledBatch,
+        scheduled_bs: int,
+        running_bs: int,
+        return_cpu: bool = False,
     ):
-        """Build `[scheduled_bs]` int32 tensor of per-request state-cache slots.
+        """Build the `[running_bs]` int32 tensor of per-request state slots.
 
         The state class declares `entries_per_req=1`, so a seq holds exactly one
         pool slot and `state_slots_committed` is the whole story. This
@@ -4124,9 +4138,21 @@ def _populate_state_slot_mappings(
         (the per-request state pool, distinct from the per-token paged-KV
         `slot_mapping`).
 
-        When `return_cpu=True`, returns `(gpu_tensor, cpu_numpy)`. The CPU
-        copy is consumed by the V4 forward path to avoid `.tolist()` syncs
-        (PR-A Phase 2).
+        When `return_cpu=True`, returns `(gpu_tensor, cpu_numpy)`; the CPU copy
+        stays at `scheduled_bs`, being the real rows the forward path reads to
+        avoid `.tolist()` syncs.
+
+        Published at `running_bs`, the width a padded draft pass reads and the
+        one `prepare_decode` uses -- the block slices this to that batch, and a
+        Python slice past the end truncates rather than raising, so a narrower
+        view would hand the kernel fewer slot rows than query rows. The
+        fabricated tail is slot 0, the inert filler `prepare_decode` uses, and
+        `mask_pad_tail` keeps those rows from scattering into it.
+
+        The width moves together with `cu_seqlens_q`: `swa_write` takes
+        `bs = state_slot_per_seq.shape[0]` and asserts
+        `cu_seqlens_q.shape[0] >= bs + 1`, so `prepare_prefill` pads that tail
+        to the same `running_bs`.
         """
         pool_np = np.asarray(batch.state_slots_committed[:scheduled_bs], dtype=np.int32)
         # Warmup / dummy_run batches don't allocate per_req_cache slots
@@ -4137,7 +4163,7 @@ def _populate_state_slot_mappings(
         if len(pool_np) < scheduled_bs:
             pool_np = np.zeros(scheduled_bs, dtype=np.int32)
         slots_np = self._physical_slots(pool_np)
-        gpu = self._stage("v4_meta_state_slot_out", slots_np)
+        gpu = self._stage("v4_meta_state_slot_out", slots_np, pad_to=running_bs)
         if return_cpu:
             return gpu, slots_np
         return gpu
@@ -4167,7 +4193,11 @@ def _physical_slots(self, pool_slots: np.ndarray) -> np.ndarray:
         return np.int32(self.pool_geometry.slot_positions - 1) - pool_slots
 
     def _populate_state_slot_in(
-        self, batch: ScheduledBatch, scheduled_bs: int, out_np: np.ndarray
+        self,
+        batch: ScheduledBatch,
+        scheduled_bs: int,
+        running_bs: int,
+        out_np: np.ndarray,
     ) -> torch.Tensor:
         """Read-side slot per seq: the fork source where set, else the write slot.
 
@@ -4184,6 +4214,7 @@ def _populate_state_slot_in(
         return self._stage(
             "v4_meta_state_slot_in",
             self._state_slot_in_np(batch, scheduled_bs, out_np),
+            pad_to=running_bs,
         )
 
     def build_for_cudagraph_capture(
@@ -4373,7 +4404,7 @@ def build_for_cudagraph_capture(
         if self.model_runner.config.enable_tbo_decode and bs > 2:
             self._prepare_ubatch_decode(
                 scheduled_bs=bs,
-                bs=bs,
+                running_bs=bs,
                 max_seqlen_q=max_q_len,
                 context_lens_np=context_lens_np,
                 state_slot_np=state_slot_np,
@@ -4700,11 +4731,15 @@ def _dest_row_buffers(self, buf_prefix_ubatch: str = "") -> dict[int, torch.Tens
             for ratio, name in _DEST_ROW_BUFFERS.items()
         }
 
-    def _stage(self, name: str, arr) -> torch.Tensor:
+    def _stage(self, name: str, arr, pad_to: int | None = None) -> torch.Tensor:
         """Write numpy `arr` into `forward_vars[name]` (CpuGpuBuffer) and
         return its GPU view sliced to len(arr). Asserts the buffer is large
         enough and that `arr.dtype` matches the buffer dtype (callers must
         cast to the buffer dtype before staging).
+
+        `pad_to` zero-fills the tail out to a wider view -- the padded batch a
+        drafter runs. Zero is a real slot, so the caller owes those rows a
+        reason they are never read.
         """
         buf = self.model_runner.forward_vars[name]
         n = arr.shape[0] if arr.ndim > 0 else 1
@@ -4721,8 +4756,14 @@ def _stage(self, name: str, arr) -> torch.Tensor:
             f"but got arr with dtype {arr.dtype}. Cast arr to the correct "
             f"dtype before calling _stage."
         )
+        width = n if pad_to is None else pad_to
+        assert width <= cap, (
+            f"V4 buffer {name!r} too small: need {width} padded, have {cap}. "
+            f"Increase the corresponding bound in _alloc_v4_metadata_buffers."
+        )
+        buf.np[n:width] = 0
         buf.np[:n] = arr
-        return buf.copy_to_gpu(n)
+        return buf.copy_to_gpu(width)
 
     @staticmethod
     def _make_gather_slot(
diff --git a/atom/model_ops/attentions/gdn_attn.py b/atom/model_ops/attentions/gdn_attn.py
index b88f7f5b65..4ae0c463b1 100644
--- a/atom/model_ops/attentions/gdn_attn.py
+++ b/atom/model_ops/attentions/gdn_attn.py
@@ -1020,7 +1020,12 @@ def prepare_gdn_metadata(
         if prepare_block_tables:
             self.prepare_block_tables(batch)
 
+        # Prefill only: `prepare_prefill` pads this past the requests it
+        # scheduled, for a draft pass that follows. Decode's padding is older
+        # than that and everything below is written for it.
         query_start_loc = attn_metadata.cu_seqlens_q
+        if is_prefill:
+            query_start_loc = query_start_loc[: num_prefills + 1]
         nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
         if not self.use_spec_decode or is_prefill:
             self.prepare_state_indices(batch, with_spec=False)
@@ -1414,8 +1419,9 @@ def build_kv_cache_tensor(self, layer_id: int, module):
     def prepare_prefill(  # type: ignore[override]
         self,
         batch: ScheduledBatch,
+        running_bs: int,
     ) -> GDNAttentionMetadata:
-        attn_metadata, positions = super().prepare_prefill(batch)
+        attn_metadata, positions = super().prepare_prefill(batch, running_bs)
         if batch.block_tables == []:
             attn_metadata.gdn_metadata = None
             return attn_metadata, positions
@@ -1438,16 +1444,20 @@ def prepare_prefill(  # type: ignore[override]
     def prepare_decode(  # type: ignore[override]
         self,
         batch: ScheduledBatch,
-        bs: int,
+        running_bs: int,
+        running_tokens: int,
+        max_seqlen_q: int,
     ) -> GDNAttentionMetadata:
-        attn_metadata, positions = super().prepare_decode(batch, bs)
+        attn_metadata, positions = super().prepare_decode(
+            batch, running_bs, running_tokens, max_seqlen_q
+        )
         self.model_runner.forward_vars["cu_seqlens_q"].cpu[
-            bs:
+            running_bs:
         ] = batch.total_tokens_num_decode
         # we fill the attn_metadata cu_seqlens_q here since aiter attn won't calc it for decode
         attn_metadata.cu_seqlens_q = self.model_runner.forward_vars[
             "cu_seqlens_q"
-        ].copy_to_gpu(bs + 1)
+        ].copy_to_gpu(running_bs + 1)
 
         self._attach_gdn_decode_metadata(batch, attn_metadata)
         return attn_metadata, positions
diff --git a/atom/model_ops/attentions/kimi_mla_gdn_attn.py b/atom/model_ops/attentions/kimi_mla_gdn_attn.py
index 083c271b2f..a43f6b7d1e 100644
--- a/atom/model_ops/attentions/kimi_mla_gdn_attn.py
+++ b/atom/model_ops/attentions/kimi_mla_gdn_attn.py
@@ -273,8 +273,8 @@ def build_kv_cache_tensor(self, layer_id: int, module):
 
         return None
 
-    def prepare_prefill(self, batch: ScheduledBatch):
-        attn_metadata, positions = super().prepare_prefill(batch)
+    def prepare_prefill(self, batch: ScheduledBatch, running_bs: int):
+        attn_metadata, positions = super().prepare_prefill(batch, running_bs)
         if batch.block_tables == []:
             attn_metadata.gdn_metadata = None
             return attn_metadata, positions
@@ -286,8 +286,16 @@ def prepare_prefill(self, batch: ScheduledBatch):
         )
         return attn_metadata, positions
 
-    def prepare_decode(self, batch: ScheduledBatch, bs: int):
-        attn_metadata, positions = super().prepare_decode(batch, bs)
+    def prepare_decode(
+        self,
+        batch: ScheduledBatch,
+        running_bs: int,
+        running_tokens: int,
+        max_seqlen_q: int,
+    ):
+        attn_metadata, positions = super().prepare_decode(
+            batch, running_bs, running_tokens, max_seqlen_q
+        )
         self._attach_gdn_decode_metadata(
             batch,
             attn_metadata,
diff --git a/atom/model_ops/attentions/triton_mha.py b/atom/model_ops/attentions/triton_mha.py
index e29f949f91..228159fc4f 100644
--- a/atom/model_ops/attentions/triton_mha.py
+++ b/atom/model_ops/attentions/triton_mha.py
@@ -41,8 +41,8 @@ class TritonMHAMetadataBuilder(AiterAttentionMetadataBuilder):
     for both prefill and decode.
     """
 
-    def prepare_prefill(self, batch: ScheduledBatch):
-        attn_metadata, positions = super().prepare_prefill(batch)
+    def prepare_prefill(self, batch: ScheduledBatch, running_bs: int):
+        attn_metadata, positions = super().prepare_prefill(batch, running_bs)
 
         # When there are no cached tokens, the base builder leaves
         # `block_tables=None` because AiterBackend's prefill consumes raw q/k/v
@@ -70,7 +70,8 @@ def prepare_prefill(self, batch: ScheduledBatch):
                 # raw K/V as a kv_cache. row i = [cu_seqlens_k[i], ...,
                 # cu_seqlens_k[i]+max-1].
                 cu_k = attn_metadata.cu_seqlens_k
-                num_seqs = cu_k.shape[0] - 1
+                # `cu_k` is padded past them; this table is one row per request.
+                num_seqs = batch.total_seqs_num_prefill
                 offsets = cu_k[:num_seqs]
                 attn_metadata.block_tables = offsets.unsqueeze(1) + torch.arange(
                     attn_metadata.max_seqlen_k, dtype=torch.int32, device=cu_k.device
diff --git a/atom/model_ops/attentions/triton_mla.py b/atom/model_ops/attentions/triton_mla.py
index d656947b06..f30872cf3b 100644
--- a/atom/model_ops/attentions/triton_mla.py
+++ b/atom/model_ops/attentions/triton_mla.py
@@ -93,8 +93,16 @@ def _cut_decode_buffers(
             "triton_lse": var["triton_lse"][:rows],
         }
 
-    def prepare_decode(self, batch: ScheduledBatch, bs: int):
-        attn_metadata, positions = super().prepare_decode(batch, bs)
+    def prepare_decode(
+        self,
+        batch: ScheduledBatch,
+        running_bs: int,
+        running_tokens: int,
+        max_seqlen_q: int,
+    ):
+        attn_metadata, positions = super().prepare_decode(
+            batch, running_bs, running_tokens, max_seqlen_q
+        )
         for name, buf in self._cut_decode_buffers(
             batch.total_seqs_num_decode,
             attn_metadata.max_seqlen_k,
@@ -139,8 +147,8 @@ def prepare_mtp_decode(
             var["kv_indptr"].gpu[: running_bs + 1],
         )
 
-    def prepare_prefill(self, batch: ScheduledBatch):
-        attn_metadata, positions = super().prepare_prefill(batch)
+    def prepare_prefill(self, batch: ScheduledBatch, running_bs: int):
+        attn_metadata, positions = super().prepare_prefill(batch, running_bs)
 
         if envs.ATOM_USE_TRITON_MLA_SHUFFLE_KV and attn_metadata.has_cached:
             # The shuffled cached-prefix gather (gather_kv_b_proj with
diff --git a/atom/model_ops/dcp_ops.py b/atom/model_ops/dcp_ops.py
index babc8254bc..628dceb1a0 100644
--- a/atom/model_ops/dcp_ops.py
+++ b/atom/model_ops/dcp_ops.py
@@ -1046,22 +1046,26 @@ def dcp_decode_candidate_exchange(
     from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits
 
     dcp_world_size = get_dcp_world_size()
-    batch_size, next_n = padded_q_fp8_decode_tokens.shape[:2]
-    num_rows = batch_size * next_n
-    num_padded_tokens = num_rows
+    next_n = padded_q_fp8_decode_tokens.shape[1]
     assert attn_metadata.max_seqlen_q == 1, (
         "DCP + DeepSeek-V3.2 sparse indexer (DSA) currently supports "
         "qlen=1 decode only (MTP verify not yet supported)."
     )
     local_ctx = dcp_local_context_lens(
-        attn_metadata, dcp_rank, dcp_world_size, cp_kv_cache_interleave_size, num_rows
+        attn_metadata,
+        dcp_rank,
+        dcp_world_size,
+        cp_kv_cache_interleave_size,
+        num_decode_tokens,
     )
     l_max = (max_model_len + dcp_world_size - 1) // dcp_world_size
-    local_logits = torch.empty([num_rows, l_max], dtype=torch.float32, device="cuda")
+    local_logits = torch.empty(
+        [num_decode_tokens, l_max], dtype=torch.float32, device="cuda"
+    )
     deepgemm_fp8_paged_mqa_logits(
         padded_q_fp8_decode_tokens,
         kv_cache,
-        weights[:num_padded_tokens],
+        weights[:num_decode_tokens],
         local_logits,
         local_ctx,
         attn_metadata.block_tables,
@@ -1075,14 +1079,14 @@ def dcp_decode_candidate_exchange(
     # (-inf, -1) padding, which the merge drops.
     k_loc = topk_tokens
     local_idx = torch.empty(
-        num_rows, k_loc, dtype=torch.int32, device=local_logits.device
+        num_decode_tokens, k_loc, dtype=torch.int32, device=local_logits.device
     )
     top_k_per_row_decode(
         local_logits,
         next_n,
         local_ctx,
         local_idx,
-        num_rows,
+        num_decode_tokens,
         local_logits.stride(0),
         local_logits.stride(1),
         k_loc,
@@ -1090,7 +1094,7 @@ def dcp_decode_candidate_exchange(
     )
     # [2, rows, k_loc]: plane 0 = score, plane 1 = int32 gid bits.
     send = torch.empty(
-        2, num_rows, k_loc, dtype=torch.float32, device=local_logits.device
+        2, num_decode_tokens, k_loc, dtype=torch.float32, device=local_logits.device
     )
     dcp_pack_topk_candidates(
         local_logits,
@@ -1107,7 +1111,7 @@ def dcp_decode_candidate_exchange(
     recv = (
         get_dcp_group()
         .all_gather(send.view(torch.int32), dim=0)
-        .view(dcp_world_size, 2, num_rows, k_loc)
+        .view(dcp_world_size, 2, num_decode_tokens, k_loc)
     )
     dcp_merge_candidates(recv, topk_indices[:num_decode_tokens, :topk_tokens])
 
diff --git a/atom/model_ops/embed_head.py b/atom/model_ops/embed_head.py
index 83a07da475..5abfffba88 100644
--- a/atom/model_ops/embed_head.py
+++ b/atom/model_ops/embed_head.py
@@ -2,7 +2,6 @@
 # Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
 
 import torch
-import torch.nn as nn
 import torch.nn.functional as F
 import triton
 import triton.language as tl
@@ -10,6 +9,7 @@
 from aiter.dist.parallel_state import get_dp_group, get_tp_group
 from aiter.jit.utils.torch_guard import torch_compile_guard
 from aiter.tuned_gemm import tgemm
+from torch import nn
 
 from atom.model_ops.lm_head_argmax import lm_head_argmax_pack
 from atom.model_ops.utils import atom_parameter
@@ -282,7 +282,12 @@ def compute_argmax_token(self, x: torch.Tensor) -> torch.Tensor:
         # Pack (val, idx) as fp32 — idx < 2^24 is exact — and all-gather only the
         # per-rank reductions ([N, 2]) instead of the full logits.
         packed = lm_head_argmax_pack(logits, self.vocab_start_idx)
-        gathered = get_tp_group().all_gather(packed, dim=0).view(self.tp_size, -1, 2)
+        # Custom, like the logits path above: `graph_capture()` arms only that
+        # one, and a draft pass records this. The RCCL path is what made the
+        # head un-capturable on HIP at TP > 1.
+        use_custom = envs.ATOM_USE_CUSTOM_ALL_GATHER
+        gathered = get_tp_group().all_gather(packed, dim=0, use_custom=use_custom)
+        gathered = gathered.view(self.tp_size, -1, 2)
         winner = gathered[:, :, 0].argmax(dim=0)  # [N] winning rank (ties -> lowest)
         token = gathered[:, :, 1].gather(0, winner.unsqueeze(0)).squeeze(0)  # [N] fp32
         return token.to(torch.long)
@@ -303,6 +308,11 @@ def _can_use_dp_sharded_head(self, context) -> bool:
         Every DP rank must reach the SAME verdict from globally-synced state, or
         the fixed-size collective in `_dp_sharded_logits` deadlocks. The strategy
         (all-gather vs all2all) is read from ATOM_DP_LM_HEAD_MODE there.
+
+        `is_prefill` below is per-rank, which is safe only because
+        `running_tokens_are_unified` is the DP-reduced form of the same
+        question: a prefilling peer drives it False on every rank, so no rank
+        can answer True here while another answers False.
         """
         dp_group = get_dp_group()
         # Static: enabled, pure DP, vocab evenly shardable over a >1 DP group.
@@ -313,14 +323,15 @@ def _can_use_dp_sharded_head(self, context) -> bool:
             or self.num_embeddings % dp_group.world_size != 0
         ):
             return False
-        # Per-step: all ranks in an equal-length decode. Exclude draft (eagle
-        # keeps dp_uniform_decode=True despite ragged rows) and prefill (x is
-        # sliced to 1 row/seq, mismatching the token-count pad target).
+        # Per-step: all ranks at one height. `is_draft` is its own question,
+        # not shorthand for the flag below -- a drafter states uniformity per
+        # pass, so a draft arrives with either answer. Prefill is excluded
+        # because x is sliced to 1 row/seq, mismatching the pad target.
         if (
             context is None
             or getattr(context, "is_draft", False)
             or getattr(context, "is_prefill", False)
-            or not getattr(context, "dp_uniform_decode", False)
+            or not getattr(context, "running_tokens_are_unified", False)
         ):
             return False
         return get_forward_context().dp_metadata is not None
diff --git a/atom/model_ops/fused_moe/modular_kernel.py b/atom/model_ops/fused_moe/modular_kernel.py
index cd15b9278f..288f065c6f 100644
--- a/atom/model_ops/fused_moe/modular_kernel.py
+++ b/atom/model_ops/fused_moe/modular_kernel.py
@@ -308,8 +308,10 @@ def _maybe_trim_dispatch_output(
         dp_size = get_dp_group().world_size
         # running_tokens keeps the trimmed shape consistent capture-to-replay.
         total_valid_tokens = context.running_tokens * dp_size
-        all_ranks_decode = getattr(context, "dp_uniform_decode", not context.is_prefill)
-        if total_valid_tokens < dispatch_a1.shape[0] and all_ranks_decode:
+        tokens_unified = getattr(
+            context, "running_tokens_are_unified", not context.is_prefill
+        )
+        if total_valid_tokens < dispatch_a1.shape[0] and tokens_unified:
             dispatch_a1 = dispatch_a1[:total_valid_tokens]
             dispatch_ids = dispatch_ids[:total_valid_tokens]
             dispatch_weights = dispatch_weights[:total_valid_tokens]
diff --git a/atom/model_ops/fused_moe/mori_v2_prepare_finalize.py b/atom/model_ops/fused_moe/mori_v2_prepare_finalize.py
index fea6c4c0cb..4817525091 100644
--- a/atom/model_ops/fused_moe/mori_v2_prepare_finalize.py
+++ b/atom/model_ops/fused_moe/mori_v2_prepare_finalize.py
@@ -528,8 +528,10 @@ def _decode_recv_bound(self, topk_ids: torch.Tensor, arena_rows: int) -> int | N
         context = get_forward_context().context
         if context is None:
             return None
-        all_ranks_decode = getattr(context, "dp_uniform_decode", not context.is_prefill)
-        if not all_ranks_decode:
+        tokens_unified = getattr(
+            context, "running_tokens_are_unified", not context.is_prefill
+        )
+        if not tokens_unified:
             return None
         bound = context.running_tokens * topk_ids.shape[1] * get_dp_group().world_size
         return bound if bound < arena_rows else None
diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py
index 3455194f8c..77804a5ff7 100644
--- a/atom/model_ops/moe.py
+++ b/atom/model_ops/moe.py
@@ -4295,7 +4295,7 @@ def forward_impl_graph(
         if use_dp_gather_scatter:
             ctx = get_forward_context()
             dp_group = get_dp_group()
-            dp_eager_mode = not ctx.context.dp_uniform_decode
+            dp_eager_mode = not ctx.context.running_tokens_are_unified
 
             from atom.utils.tbo.ubatching import tbo_active
 
diff --git a/atom/models/deepseek_v4.py b/atom/models/deepseek_v4.py
index 927d166b6e..a83ebbc6d9 100644
--- a/atom/models/deepseek_v4.py
+++ b/atom/models/deepseek_v4.py
@@ -3271,12 +3271,12 @@ def _sparse_attention(
         fc = get_forward_context()
         # Row count off the Q, NOT off `positions`. The Q descends from the
         # hidden state through `wqkv_a`/`wq_b`, so it carries the same token
-        # count the residual stream does. `positions` does not have to: under
-        # `--enable-dp-attention`, a step where any rank is prefilling takes the
-        # variable-length path (`model_runner.py`, `dp_uniform_decode`) and the
-        # two can differ. Sizing the attention output by `positions` then hands
-        # the mHC residual a mismatched `m` -- an aiter shape assert deep inside
-        # a compiled piece, far from here.
+        # count the residual stream does. `positions` does not have to: a step
+        # where any DP rank is prefilling takes the variable-length path (see
+        # `running_tokens_are_unified` in `forward_context.py`) and the two can
+        # differ. Sizing the attention output by `positions` then hands the mHC
+        # residual a mismatched `m` -- an aiter shape assert deep inside a
+        # compiled piece, far from here.
         q_rows = qkn.q_sa if qkn.q_sa is not None else qkn.q_packed
         assert q_rows is not None, (
             "_sparse_attention got no Q: `_qk_norm_rope` did not run upstream. "
@@ -3879,7 +3879,7 @@ def _gather_ids_for_dp(ids: torch.Tensor, ctx) -> torch.Tensor:
 
         ids_2d = ids.unsqueeze(-1)
         dp_eager_mode = (
-            not ctx.context.dp_uniform_decode
+            not ctx.context.running_tokens_are_unified
         ) and ctx.dp_metadata is not None
         if dp_eager_mode:
             from atom.model_ops.moe import all_gatherv
diff --git a/atom/plugin/sglang/attention_backend/attention_gdn.py b/atom/plugin/sglang/attention_backend/attention_gdn.py
index d65242c9af..3226006df9 100644
--- a/atom/plugin/sglang/attention_backend/attention_gdn.py
+++ b/atom/plugin/sglang/attention_backend/attention_gdn.py
@@ -454,7 +454,7 @@ def _build_context(forward_batch: Any) -> tuple[Context, int]:
         effective_forward_mode = (
             global_forward_mode if global_forward_mode is not None else mode
         )
-        dp_uniform_decode = not enable_dp_attention or bool(
+        running_tokens_are_unified = not enable_dp_attention or bool(
             effective_forward_mode.is_decode_or_idle()
         )
         return (
@@ -468,7 +468,7 @@ def _build_context(forward_batch: Any) -> tuple[Context, int]:
                 # per-request multiplier is needed (nor available here).
                 scheduled_tokens=num_tokens,
                 running_tokens=num_tokens,
-                dp_uniform_decode=dp_uniform_decode,
+                running_tokens_are_unified=running_tokens_are_unified,
             ),
             num_tokens,
         )
diff --git a/atom/plugin/sglang/runtime/forward_context.py b/atom/plugin/sglang/runtime/forward_context.py
index 7fe1142abb..84895b0a39 100644
--- a/atom/plugin/sglang/runtime/forward_context.py
+++ b/atom/plugin/sglang/runtime/forward_context.py
@@ -152,7 +152,7 @@ def _resolve_num_tokens_across_dp(
     return num_tokens_across_dp
 
 
-def _resolve_dp_uniform_decode(
+def _resolve_running_tokens_are_unified(
     atom_config: Any,
     forward_batch: ForwardBatch,
 ) -> bool:
@@ -796,7 +796,9 @@ def _set_atom_forward_context(
     # ubatch `max_seqlen_q` times too wide. Prefill's value is unread there.
     running_bs = batch_size if is_prefill else max(1, running_tokens // max_q)
 
-    dp_uniform_decode = _resolve_dp_uniform_decode(atom_config, forward_batch)
+    running_tokens_are_unified = _resolve_running_tokens_are_unified(
+        atom_config, forward_batch
+    )
     context = Context(
         positions=positions,
         is_prefill=is_prefill,
@@ -805,7 +807,7 @@ def _set_atom_forward_context(
         scheduled_tokens=num_tokens,
         running_bs=running_bs,
         running_tokens=running_tokens,
-        dp_uniform_decode=dp_uniform_decode,
+        running_tokens_are_unified=running_tokens_are_unified,
     )
     set_forward_context(
         attn_metadata=attn_metadata,
diff --git a/atom/plugin/sglang/tbo/ubatch_wrapper.py b/atom/plugin/sglang/tbo/ubatch_wrapper.py
index 65369f4400..56d7bfcfe6 100644
--- a/atom/plugin/sglang/tbo/ubatch_wrapper.py
+++ b/atom/plugin/sglang/tbo/ubatch_wrapper.py
@@ -56,7 +56,7 @@ def _make_sglang_ubatch_context(
             running_bs=running_bs,
             running_tokens=running_tokens,
             is_draft=ctx.context.is_draft,
-            dp_uniform_decode=ctx.context.dp_uniform_decode,
+            running_tokens_are_unified=ctx.context.running_tokens_are_unified,
             forward_mode=ctx.context.forward_mode,
             input_ids=input_ids,
             ubatch_token_offset=ub_slice.token_slice.start,
diff --git a/atom/plugin/vllm/attention/layer_sparse_mla.py b/atom/plugin/vllm/attention/layer_sparse_mla.py
index ebb2676281..049f6b5360 100644
--- a/atom/plugin/vllm/attention/layer_sparse_mla.py
+++ b/atom/plugin/vllm/attention/layer_sparse_mla.py
@@ -457,14 +457,13 @@ def sparse_attn_indexer_plugin_mode(
         batch_size = padded_q_fp8_decode_tokens.shape[0]
         next_n = padded_q_fp8_decode_tokens.shape[1]
         assert batch_size == decode_metadata.seq_lens.shape[0]
-        num_padded_tokens = batch_size * next_n
         logits = torch.empty(
-            [batch_size * next_n, max_model_len], dtype=torch.float32, device="cuda"
+            [num_decode_tokens, max_model_len], dtype=torch.float32, device="cuda"
         )
         deepgemm_fp8_paged_mqa_logits(
             padded_q_fp8_decode_tokens,
             kv_cache,
-            weights[:num_padded_tokens],
+            weights[:num_decode_tokens],
             logits,
             decode_metadata.seq_lens,
             decode_metadata.block_table,
@@ -475,7 +474,6 @@ def sparse_attn_indexer_plugin_mode(
             WavePerEU=2,
         )
 
-        num_rows = logits.shape[0]
         assert topk_tokens == 2048, "top_k_per_row assumes size 2048"
         topk_indices_decode = topk_indices[:num_decode_tokens, :topk_tokens]
         top_k_per_row_decode(
@@ -483,7 +481,7 @@ def sparse_attn_indexer_plugin_mode(
             next_n,
             decode_metadata.seq_lens,
             topk_indices_decode,
-            num_rows,
+            num_decode_tokens,
             logits.stride(0),
             logits.stride(1),
             stable=stable_topk,
diff --git a/atom/plugin/vllm/mori_patch.py b/atom/plugin/vllm/mori_patch.py
index 5c2fb6b891..7bdd98cd96 100644
--- a/atom/plugin/vllm/mori_patch.py
+++ b/atom/plugin/vllm/mori_patch.py
@@ -77,8 +77,10 @@ def _trim_qwen35_uniform_full_graph(
     dp_size = get_dp_group().world_size
     topk = topk_ids.shape[1]
     total_valid_tokens = context.graph_bs * topk * dp_size
-    all_ranks_decode = getattr(context, "dp_uniform_decode", not context.is_prefill)
-    if total_valid_tokens < dispatch_a1.shape[0] and all_ranks_decode:
+    tokens_unified = getattr(
+        context, "running_tokens_are_unified", not context.is_prefill
+    )
+    if total_valid_tokens < dispatch_a1.shape[0] and tokens_unified:
         dispatch_a1 = dispatch_a1[:total_valid_tokens]
         dispatch_ids = dispatch_ids[:total_valid_tokens]
         dispatch_weights = dispatch_weights[:total_valid_tokens]
diff --git a/atom/rollout/model_runner_ext.py b/atom/rollout/model_runner_ext.py
index 5bafb82279..026234346a 100644
--- a/atom/rollout/model_runner_ext.py
+++ b/atom/rollout/model_runner_ext.py
@@ -95,7 +95,7 @@ def _setup_device_and_distributed(self, rank: int, config):
 
         # DP is handled at the EngineCore level (DPEngineCoreProc), not
         # within ModelRunner. Override so downstream code (get_dp_padding,
-        # _preprocess/sync_dp_metadata) sees dp_size=1 and skips cross-DP
+        # ForwardMode.decide/sync_dp_metadata) sees dp_size=1 and skips cross-DP
         # collectives that would fail on the isolated process group.
         config.parallel_config.data_parallel_size = 1
 
diff --git a/atom/spec_decode/draft_graph.py b/atom/spec_decode/draft_graph.py
index 2b6ea02fc1..719912e35f 100644
--- a/atom/spec_decode/draft_graph.py
+++ b/atom/spec_decode/draft_graph.py
@@ -19,6 +19,19 @@
     from atom.config import Config
 
 
+def _owned(value: Any) -> Any:
+    """A copy the graph pool cannot re-issue.
+
+    Whole return, not the one output known to outlive its step: which that is
+    belongs to the flavor, not to this seam.
+    """
+    if isinstance(value, torch.Tensor):
+        return value.clone()
+    if isinstance(value, tuple):
+        return tuple(_owned(v) for v in value)
+    return value
+
+
 @dataclass(frozen=True)
 class StagedInput:
     """One staged input's type: everything but the leading batch axis."""
@@ -39,10 +52,9 @@ class DraftGraph:
 
     * **warmup** -- run it once per ``capture_sizes`` entry at startup, paying its
       per-shape JIT there rather than mid-serve. Needs nothing.
-    * **pad** -- run at the batch the target just ran, which the startup sweep
+    * **pad** -- run at the batch the step settled on, which the startup sweep
       warmed, rather than at whatever this step's batch happens to be. Needs
-      ``inputs`` (fixed addresses to pad into) and ``pads`` (an assertion that
-      the fabricated rows reach nothing).
+      ``inputs``: fixed addresses for the fabricated rows to land in.
     * **capture** -- record it. Needs pad.
 
     The pass is ``epilogue ∘ forward``. Warmup always runs both, so nothing
@@ -54,16 +66,16 @@ class DraftGraph:
     ``warmup_inputs`` makes that startup batch plausible, since warming on zeros
     compiles a shape no real batch draws.
 
-    Everything here counts SEQUENCES, in the target's own words: ``bs`` is the
-    scheduled batch, ``running_bs`` is what it actually runs at, and
-    ``capture_sizes`` is the ladder the startup sweep warms.
+    Everything here counts SEQUENCES, in the target's own words:
+    ``scheduled_bs`` is what this rank was handed, ``running_bs`` is the
+    DP-agreed batch a recording holds, and ``capture_sizes`` is the ladder
+    the startup sweep warms.
     """
 
     forward: Callable[..., Any]  # (running_bs, **staged) -> Any
     epilogue: Callable[..., Any] | None = None  # (fwd_out, running_bs, **staged)
     capture_epilogue: bool = False
     inputs: "Mapping[str, StagedInput]" = field(default_factory=dict)
-    pads: bool = False
     warmup_inputs: Callable[..., None] | None = None
 
     _buffers: dict[str, torch.Tensor] = field(init=False, default_factory=dict)
@@ -87,9 +99,9 @@ def bind(self, config: "Config", device) -> "DraftGraph":
         their addresses, and one born on the capture step would come out of the
         graph's own private pool.
         """
-        assert not self.pads or self.inputs, (
-            f"{self.name} says it may pad but stages nothing, so there is no "
-            f"buffer for the fabricated rows to land in"
+        assert self.inputs, (
+            f"{self.name} stages nothing, so a fabricated row has no buffer "
+            f"to land in -- declare no pass rather than an unpaddable one"
         )
         self._buffers = {
             role: torch.zeros(
@@ -103,46 +115,12 @@ def bind(self, config: "Config", device) -> "DraftGraph":
 
     @property
     def will_capture(self) -> bool:
-        """Whether warmup also captures. One gate for every flavor, since the
-        graph is made by the shared ``warmup``."""
-        return self.pads and envs.ATOM_DRAFT_CUDAGRAPH
-
-    def target_running_bs(self, bs: int, context) -> int:
-        """The batch the target just ran at, or ``bs`` when it pinned none.
-
-        Never a batch the drafter picks. ``context.running_bs`` IS what the
-        target ran: ``model_runner`` builds the attention metadata with it, so
-        every per-sequence buffer already reaches that far -- ``slot_mapping``
-        -1, ``context_lens`` 0, ring slots published -- and a pad row fabricates
-        nothing the backend has not already accounted for, on any backend and
-        without the drafter knowing each one's convention. It is a warmed shape
-        by construction too: on the cudagraph path ``ForwardMode.decide`` picks
-        it out of ``capture_sizes``, the list the startup sweep warms. And it is
-        DP-unified, which is what lets the draft's own collectives line up.
-
-        Off the cudagraph path nothing is padded -- ``running_bs`` is then the
-        scheduled batch, and that is the only safe answer anyway, since nobody
-        sized the target's metadata past the real batch there.
-
-        An empty batch is never widened: a pad row is a copy of the last real
-        one, and there is no last real one to copy. Declining the padding rather
-        than the pass is deliberate -- under DP every rank still has to reach the
-        draft's collectives.
-
-        ``is_dummy_run`` is deliberately NOT read, for the reason spelled out in
-        :meth:`_replays`: it is per-rank, so any term keyed on it splits the DP
-        group into two shapes.
+        """Whether warmup also captures. The name `envs.ATOM_DRAFT_CUDAGRAPH`
+        points at, and the only thing that decides it."""
+        return envs.ATOM_DRAFT_CUDAGRAPH
 
-        """
-        mode = context.forward_mode
-        if not bs or not self.pads:
-            return bs
-        if mode is None or not mode.use_cudagraph:
-            return bs
-        return max(bs, int(context.running_bs))
-
-    def label(self, real_bs: int, running_bs: int, context) -> str:
-        """``bs=<real>/<pad>`` plus a trailing ``graph`` when this step replays.
+    def label(self, scheduled_bs: int, running_bs: int) -> str:
+        """``bs=<scheduled>/<running>``, plus ``graph`` when this step replays.
 
         One convention, one implementation: "did the draft get into a graph" is
         then a single grep across every flavor, rather than each flavor writing
@@ -150,8 +128,8 @@ def label(self, real_bs: int, running_bs: int, context) -> str:
         predicate `run` does, so the mark cannot claim a replay that did not
         happen.
         """
-        return f"bs={real_bs}/{running_bs}" + (
-            " graph" if self._replays(running_bs, context) else ""
+        return f"bs={scheduled_bs}/{running_bs}" + (
+            " graph" if self.is_captured(running_bs) else ""
         )
 
     def stage(
@@ -159,12 +137,7 @@ def stage(
     ) -> dict[str, torch.Tensor]:
         """Copy every input into its fixed buffer, widening the batch to ``running_bs``.
 
-        The only way in, because ``pads`` is the part of the contract nothing
-        else can check: a pass that lies about it still runs and still returns
-        the right shape -- its fabricated rows just land in another sequence's
-        KV.
-
-        How many rows are real comes from the sources, not from the caller: a
+        How many rows are real comes from the sources, never from the caller: a
         count that cannot disagree with the tensor it describes.
         """
         counts = {t.shape[0] for t in srcs.values()}
@@ -172,11 +145,6 @@ def stage(
             f"{self.name} was handed batches of {sorted(counts)}; they "
             f"describe one step, so one of them is not this step's"
         )
-        bs = counts.pop() if counts else running_bs
-        assert running_bs == bs or self.pads, (
-            f"{self.name} was asked for {running_bs} rows against {bs} real ones, "
-            f"but never said its fabricated rows are inert"
-        )
         return {
             role: self._stage_one(role, running_bs, src) for role, src in srcs.items()
         }
@@ -206,54 +174,58 @@ def _stage_one(self, role: str, running_bs: int, src: torch.Tensor | None):
             f"{tuple(buf.shape[1:])} per sequence. MRoPE positions ([3, N]) are "
             f"the case that reaches here"
         )
-        bs = src.shape[0]
-        out[:bs].copy_(src)
-        if running_bs > bs:
-            out[bs:].copy_(src[-1])  # copy_ broadcasts; no expand needed
+        scheduled_bs = src.shape[0]
+        out[:scheduled_bs].copy_(src)
+        if running_bs > scheduled_bs:
+            out[scheduled_bs:].copy_(src[-1])  # broadcasts; no expand needed
         return out
 
     def is_captured(self, running_bs: int) -> bool:
-        """Whether a recording exists at this shape. Not whether it may be used."""
-        return running_bs in self._cuda_graphs
-
-    def _replays(self, running_bs: int, context) -> bool:
-        """Whether THIS step stands in for a recording.
-
-        A recording holds the collective at the width it was made for, so every
-        term here must be one the whole DP group agrees on -- which is why
-        `is_dummy_run` may not enter. A DP-sync dummy is per-rank, so gating on
-        it makes the rank with work replay while the rest issue eagerly, and
-        all of them wait forever (measured: V4-Flash-DSpark tp8 + DPA hung 8/8
-        on the first real decode). Replaying is safe on a dummy anyway --
-        `warmup` asserts the recording was made on a REAL context.
+        """Whether a recording exists at this batch -- and so whether this step
+        replays, the two having become one question.
+
+        Nothing about the step enters. A recording holds its collective at one
+        width, so the decision has to be one the whole DP group agrees on, and
+        ``running_bs`` was reduced in ``sync_dp_metadata`` before it got here.
+        That is what a second term used to buy: ``is_dummy_run`` and
+        ``is_prefill`` are both per-rank, and either splits the group into two
+        widths (measured -- V4-Flash-DSpark tp8 + DPA hung 8/8 on the first
+        real decode). Neither is reachable from an agreed batch, so neither is
+        available to get wrong. Safe on a dummy and on a prefill step for one
+        reason: ``Drafter.warmup_draft_graphs`` records on a real decode-shaped
+        context, and this pass is that shape whatever the target just did.
         """
-        mode = context.forward_mode
-        return mode is not None and mode.use_cudagraph and self.is_captured(running_bs)
+        return running_bs in self._cuda_graphs
 
     @property
     def _to_capture(self) -> Callable[..., Any]:
         """What a graph holds -- and so what a replay stands in for."""
         return self._forward_and_epilogue if self.capture_epilogue else self.forward
 
-    def run(self, running_bs: int, context, **staged) -> Any:
-        """Replay this batch size's graph, then whatever it did not cover.
+    def run(self, running_bs: int, **staged) -> Any:
+        """Replay this batch's graph, then whatever it did not cover.
+
+        The batch is the whole argument. `running_bs` is the reduction's own
+        answer rounded on a ladder every rank shares, so every rank answers
+        alike and nothing about the step needs to reach here -- see
+        :meth:`is_captured`.
 
-        Takes the step, not just its shape: a recording stands in for the branch
-        the flavor took while it was made, and a dummy took a different one --
-        ``warmup`` asserts it ran on a real context precisely so the recording
-        holds the real branch. Refusing the padding is not enough, because a
-        dummy's own batch can be a captured size on its own.
+        What leaves here is a VALUE. A replay hands back the tensors the
+        capture allocated, and the pool re-issues those addresses to sizes
+        captured later -- a reference does not hold them the way the eager
+        allocator's does. Only ``capture_epilogue`` can put a survivor in
+        there; the eager path allocates normally and needs no copy.
         """
         captured = (
-            self._cuda_graphs.get(running_bs)
-            if self._replays(running_bs, context)
-            else None
+            self._cuda_graphs.get(running_bs) if self.is_captured(running_bs) else None
         )
         if captured is None:
             out = self._to_capture(running_bs, **staged)
         else:
             graph, out = captured
             graph.replay()
+            if self.capture_epilogue:
+                out = _owned(out)
         return (
             out
             if self.capture_epilogue
diff --git a/atom/spec_decode/drafter.py b/atom/spec_decode/drafter.py
index fb0cd857cc..d08bd3c2a0 100644
--- a/atom/spec_decode/drafter.py
+++ b/atom/spec_decode/drafter.py
@@ -230,6 +230,11 @@ def warmup_draft_graphs(self, build_context, stream) -> None:
             for bs in capture_sizes:
                 attn_metadata, context = build_context(bs=bs)
                 context.is_draft = True
+                # Every warmed pass is decode-shaped and DP-uniform: the graphs
+                # belong to the DSpark block and Eagle's mid-steps, and step 0
+                # has none. Stated, not inherited, per the rule above.
+                context.is_prefill = False
+                context.running_tokens_are_unified = True
                 # The synthetic batch is full, so the pass's scheduled and
                 # running counts are the same number.
                 local_tokens = bs * self.draft_tokens_per_seq
@@ -551,7 +556,12 @@ def load_model(self, target_model: nn.Module) -> None:
         )
 
     def _publish_draft_shape(
-        self, forward_context, scheduled_tokens: int, running_tokens: int
+        self,
+        forward_context,
+        scheduled_tokens: int,
+        running_tokens: int,
+        *,
+        running_tokens_are_unified: bool,
     ) -> None:
         """Re-point the forward context at the pass about to run.
 
@@ -559,19 +569,35 @@ def _publish_draft_shape(
         tokens -- not the draft's own width. Anything sizing a collective off
         the context (MoE's `pad_for_all_gather` above all) would then use the
         wrong height, so the draft states its shape here rather than letting
-        each consumer special-case `is_draft`. Both units: DP's variable-length
-        gather takes `running_tokens`, the pad needs to know how much is real.
+        each consumer special-case `is_draft`.
+
+        Both units: `scheduled_tokens` is how many rows carry a real request,
+        `running_tokens` how many the pass runs -- they differ exactly when the
+        batch was widened.
+
+        Both flags below are the caller's to answer, because this seam cannot
+        see either. `running_tokens_are_unified` claims every OTHER rank runs
+        this height too -- true where it is `context.running_bs` (`decide`'s
+        reduction) times a config width, false where it came off this rank's
+        own batch or token stream. `is_prefill` is left alone for the same
+        reason: a pass on its own `[bs, T]` shape must clear it, one carrying
+        the target's stream must not.
         """
         context = forward_context.context
         context.scheduled_tokens = scheduled_tokens
         context.running_tokens = running_tokens
         parallel_config = self.config.parallel_config
+        # A group of one is uniform whatever it runs; only the table needs peers.
         if parallel_config.data_parallel_size <= 1:
             return
-        context.dp_uniform_decode = False
+        context.running_tokens_are_unified = running_tokens_are_unified
+        # The answer travels, not the table it implies -- `make` owns both ways
+        # of reaching one, and skipping its all_reduce is why this is worth
+        # stating at all.
         forward_context.dp_metadata = DPMetadata.make(
             parallel_config,
             running_tokens,
+            unified=running_tokens_are_unified,
         )
 
     def prepare_inputs(
diff --git a/atom/spec_decode/dspark_proposer.py b/atom/spec_decode/dspark_proposer.py
index 021dfe4813..3326bebd7e 100644
--- a/atom/spec_decode/dspark_proposer.py
+++ b/atom/spec_decode/dspark_proposer.py
@@ -79,7 +79,6 @@ def _declare_draft_graphs(self):
                 "anchor_ids": StagedInput(dtype=torch.int32),
                 "anchor_positions": StagedInput(dtype=torch.int64),
             },
-            pads=self._block_may_pad,
             warmup_inputs=self._block_warmup_inputs,
         )
         return (self.block,)
@@ -98,16 +97,6 @@ def _block_warmup_inputs(self, running_bs, *, anchor_positions, **_):
         )
         anchor_positions.fill_(int(self.model.window_size))
 
-    @property
-    def _block_may_pad(self) -> bool:
-        """Whether the block may take fabricated rows.
-
-        Off under EPLB: pad rows carry fabricated tokens through the draft's full
-        MoE, and ``select_experts_with_record`` counts every row it routes into
-        the expert-load histogram that drives online expert migration.
-        """
-        return not self.config.eplb_enable
-
     def _block_head(self, out, running_bs, *, anchor_ids, **_):
         """The block's epilogue: LM head, then the sequential Markov sampler.
 
@@ -117,6 +106,10 @@ def _block_head(self, out, running_bs, *, anchor_ids, **_):
         it does all_gather the vocab shard; ``capture_epilogue`` says whether
         that one collective is captured with the rest.
 
+        Its ids are what `DraftGraph.run` copies out of the pool: deferred
+        output feeds them back as the next step's `input_ids`, so they are the
+        one thing here that outlives the pass.
+
         It must be WARMED regardless: the head has its own per-shape flydsl
         builder, and leaving it out of the warm is exactly how
         `hipModuleLoadData` went 0 -> 4 on the reproducer once.
@@ -483,9 +476,11 @@ def compute_draft_kv(
         if aux_hidden_states is None:
             return
         forward_context = get_forward_context()
-        bs = forward_context.context.scheduled_bs
+        scheduled_bs = forward_context.context.scheduled_bs
         main_hidden_all = torch.cat(aux_hidden_states, dim=-1)
-        with record_function(f"draft_kv[bs={bs} tok={main_hidden_all.shape[0]}]"):
+        with record_function(
+            f"draft_kv[bs={scheduled_bs} tok={main_hidden_all.shape[0]}]"
+        ):
             self.model.write_context_kv(main_hidden_all, positions)
 
     def propose(
@@ -517,7 +512,7 @@ def propose(
         context = forward_context.context
         attn_metadata = forward_context.attn_metadata
         context.is_draft = True
-        bs = context.scheduled_bs
+        scheduled_bs = context.scheduled_bs
 
         # Drafter-owned aux: our own forward-hook capture buffers, row-aligned to
         # the target hidden states.
@@ -542,7 +537,7 @@ def propose(
             return self._propose_with_draft(
                 forward_context,
                 attn_metadata,
-                bs,
+                scheduled_bs,
                 anchor_ids,
                 anchor_positions,
             )
@@ -563,22 +558,23 @@ def propose(
         num_draft = self.draft_tokens_per_seq
         # forward_spec sizes the block off anchor_ids. context.scheduled_bs counts
         # only one half of a mixed prefill+decode step, so it is not that B.
-        real_bs = anchor_ids.shape[0]
+        scheduled_bs = anchor_ids.shape[0]
         # Agreed first, and on EVERY step. The batch a pass runs at has to be
         # one number for the whole DP group, or half of it replays a recorded
         # collective while the rest issue a differently sized one. Not
         # conditioned on prefill-vs-decode either: which of the two a rank is
         # doing is its own business, so a rank that skipped this would leave
         # the others waiting in the exchange.
-        running_bs = self.block.target_running_bs(real_bs, context)
-        # The target already replayed a padded graph, but none of that padding
-        # reaches here: its pad rows end at the graph boundary. `anchor_ids`
-        # comes from the sampler, which runs after the graph and only over real
-        # rows. So the block pads its own inputs. The target's `state_slot_out`
-        # needs no help, and by an identity rather than an agreement:
-        # `prepare_decode` publishes it at the `bs` it was built with, which is
-        # `context.running_bs` -- the same number `target_running_bs` just
-        # returned.
+        running_bs = context.running_bs
+        # The target's own padding does not reach here: its pad rows end at the
+        # graph boundary, and `anchor_ids` comes from the sampler, which runs
+        # after the graph over real rows only. So the block pads its own inputs.
+        # `state_slot_out` needs no help though: `prepare_decode` and
+        # `prepare_prefill` publish it at this same number, so the block's
+        # `[:B]` covers every row it runs. It stopped covering them the moment
+        # the draft ran at a batch the target had not published to -- a Python
+        # slice past the end truncates rather than raising, so the kernels got
+        # a slot table shorter than their own grid and read past it.
         staged = self.block.stage(
             running_bs,
             {"anchor_ids": anchor_ids, "anchor_positions": anchor_positions},
@@ -589,21 +585,29 @@ def propose(
         # and not inside the block: `run` may REPLAY, and then nothing in the
         # block's Python runs at all.
         self.model.model.index_buffers(num_draft, window, self.device).mask_pad_tail(
-            self.runner.attn_metadata_builder.row_ids, real_bs, running_bs
+            self.runner.attn_metadata_builder.row_ids, scheduled_bs, running_bs
         )
+        # The block pass is `[bs, T]` however the target ran, and
+        # `pad_for_all_gather` reads `is_prefill` to pick which count below to
+        # check its rows against -- left True it checks the scheduled one and
+        # asserts as soon as the batch is widened.
+        context.is_prefill = False
         self._publish_draft_shape(
             forward_context,
-            scheduled_tokens=real_bs * num_draft,
+            scheduled_tokens=scheduled_bs * num_draft,
             running_tokens=running_bs * num_draft,
+            # `running_bs` is `decide`'s reduction and `num_draft` config, so
+            # every rank reaches this height on its own.
+            running_tokens_are_unified=True,
         )
-        label = self.block.label(real_bs, running_bs, context)
+        label = self.block.label(scheduled_bs, running_bs)
         with record_function(f"propose_dspark[{label} T={num_draft}]"):
-            draft_token_ids, confidence = self.block.run(running_bs, context, **staged)
-        draft_token_ids = draft_token_ids[:real_bs, : self.mtp_k]
+            draft_token_ids, confidence = self.block.run(running_bs, **staged)
+        draft_token_ids = draft_token_ids[:scheduled_bs, : self.mtp_k]
         if confidence is not None:
             # compute_ell takes its batch size from confidence.shape and zips the
             # resulting ell against batch.req_ids by position.
-            confidence = confidence[:real_bs]
+            confidence = confidence[:scheduled_bs]
         # Confidence-scheduled verification. The hardware-aware prefix scheduler
         # consumes the confidence head to pick a per-request verify length
         # ell_r. We compute ell here and stash it; the actual variable-length
@@ -611,7 +615,7 @@ def propose(
         # request's scheduled spec tokens to ell_r, which frees batch capacity
         # instead of the no-op in-block masking of Level A.
         if self.verify_scheduler is not None and confidence is not None:
-            with record_function(f"dspark_sched[bs={real_bs}]"):
+            with record_function(f"dspark_sched[bs={scheduled_bs}]"):
                 self.verify_scheduler.set_last_ell(
                     self.verify_scheduler.compute_ell(confidence[:, : self.mtp_k])
                 )
@@ -625,9 +629,9 @@ def _propose_with_draft(
         self,
         forward_context,
         attn_metadata,
-        bs: int,
-        anchor_ids: torch.Tensor,  # [bs]
-        anchor_positions: torch.Tensor,  # [bs]
+        scheduled_bs: int,
+        anchor_ids: torch.Tensor,  # [scheduled_bs]
+        anchor_positions: torch.Tensor,  # [scheduled_bs]
     ) -> torch.Tensor:
         """Kimi-K3 DSpark: one non-causal block pass over the paged latent cache.
 
@@ -653,18 +657,20 @@ def _propose_with_draft(
         # never leaks back.
         attn_metadata.causal = False
 
-        block_positions = self._blk_positions[:bs]  # [bs, T] view, stable
+        block_positions = self._blk_positions[
+            :scheduled_bs
+        ]  # [scheduled_bs, T] view, stable
         torch.add(
-            anchor_positions.view(bs, 1),
+            anchor_positions.view(scheduled_bs, 1),
             self._blk_offsets.view(1, T),
             out=block_positions,
         )
 
         if not is_dummy:
-            block_tables = self.runner.forward_vars["block_tables"].gpu[:bs]
-            slots = self._blk_slots[: bs * T].view(bs, T)  # stable
+            block_tables = self.runner.forward_vars["block_tables"].gpu[:scheduled_bs]
+            slots = self._blk_slots[: scheduled_bs * T].view(scheduled_bs, T)  # stable
             # Each request's KV spans [0, anchor] (context) ++ the T block rows.
-            ctx_lens = self._blk_ctx_lens[:bs]  # stable
+            ctx_lens = self._blk_ctx_lens[:scheduled_bs]  # stable
             global_lens = anchor_positions + (1 + T)
             if self.dcp_world_size > 1:
                 # DCP shards KV token-wise round-robin, so one block table entry
@@ -709,13 +715,13 @@ def _propose_with_draft(
 
             attn_metadata.slot_mapping = slots.view(-1)
             attn_metadata.block_tables = block_tables
-            attn_metadata.cu_seqlens_q = self._blk_cu_seqlens_q[: bs + 1]
+            attn_metadata.cu_seqlens_q = self._blk_cu_seqlens_q[: scheduled_bs + 1]
             attn_metadata.context_lens = ctx_lens
             attn_metadata.max_seqlen_q = T
             # Upper bound rather than a .max() host sync: every context_len is
             # at most the target pass's longest sequence plus the block.
             attn_metadata.max_seqlen_k = int(attn_metadata.max_seqlen_k) + T
-            kv_indptr = self._blk_kv_indptr[: bs + 1]
+            kv_indptr = self._blk_kv_indptr[: scheduled_bs + 1]
             # kv_indptr[0] stays 0 (zero-init, never written). cumsum promotes
             # integers to int64, so land it through copy_ rather than out=.
             kv_indptr[1:].copy_(torch.cumsum(ctx_lens, dim=0))
@@ -740,7 +746,7 @@ def _propose_with_draft(
             )
             attn_metadata.kv_indptr = kv_indptr
             attn_metadata.kv_indices = self._blk_kv_indices
-            attn_metadata.kv_last_page_lens = self._blk_last_page_lens[:bs]
+            attn_metadata.kv_last_page_lens = self._blk_last_page_lens[:scheduled_bs]
 
             # Build persistent (ps=1) MLA-decode metadata
             from atom.model_ops.attentions.aiter_mla import get_mla_metadata_v1
@@ -777,27 +783,23 @@ def _propose_with_draft(
             attn_metadata.reduce_final_map = ps["reduce_final_map"]
             attn_metadata.reduce_partial_map = ps["reduce_partial_map"]
 
-        # The separate-draft path pads nothing, so both heights are `bs * T`.
+        # `[bs, T]` against a paged KV cache however the target just ran. Left
+        # True, every MLA layer takes its prefill branch, which reads
+        # cu_seqlens_k / chunk_meta / _gather_cached_kv_b_proj -- none of which
+        # the retarget above touches, because none of them describe this batch.
+        # Not restored: the Context is rebuilt per forward, like `is_draft`.
+        forward_context.context.is_prefill = False
+        # The separate-draft path pads nothing, so both heights are
+        # `scheduled_bs * T` -- this rank's own request count, which the group
+        # can only discover. Nothing rules this path out under DP either:
+        # `--enable-dp-attention` folds TP into the DP size.
         self._publish_draft_shape(
-            forward_context, scheduled_tokens=bs * T, running_tokens=bs * T
+            forward_context,
+            scheduled_tokens=scheduled_bs * T,
+            running_tokens=scheduled_bs * T,
+            running_tokens_are_unified=False,
         )
 
-        # The block pass is ALWAYS decode-shaped -- T queries per request against
-        # a paged KV cache -- even on a step where the target just prefilled. But
-        # `is_prefill` is a property of the step, not of the model being run, so
-        # on a prefill step it is still True here and every MLA layer would take
-        # its prefill branch (attention_mla.py:1238, and again at :1443), which
-        # reads cu_seqlens_k / chunk_meta / _gather_cached_kv_b_proj -- none of
-        # which the retarget above touches, because none of them describe this
-        # batch. Force the decode shape.
-        #
-        # EagleProposer does the same (eagle_proposer.py:358) but only from its
-        # SECOND draft step: its first step deliberately reuses the target's own
-        # layout. DSpark has exactly one block pass and it is never that shape,
-        # so this is unconditional. Not restored afterwards -- the Context is
-        # rebuilt per forward, the same reason `is_draft` above is not restored.
-        forward_context.context.is_prefill = False
-
         dtype_q = self._blk_dtype_q
         if dtype_q is None:
             dtype_q, final = self._resolve_dtype_q(forward_context)
@@ -806,7 +808,7 @@ def _propose_with_draft(
         forward_context.attn_metadata.dtype_q = dtype_q
 
         # ---- 3. Block pass + Markov sampling ---------------------------------
-        with record_function(f"propose_dspark[bs={bs} T={T}]"):
+        with record_function(f"propose_dspark[bs={scheduled_bs} T={T}]"):
             draft_token_ids, confidence = self.model.forward_spec(
                 anchor_ids,
                 block_positions.view(-1),
diff --git a/atom/spec_decode/dspark_scheduler.py b/atom/spec_decode/dspark_scheduler.py
index 2cb6d2d3b2..d1d82e20cd 100644
--- a/atom/spec_decode/dspark_scheduler.py
+++ b/atom/spec_decode/dspark_scheduler.py
@@ -21,7 +21,7 @@
 
 from __future__ import annotations
 
-from typing import Sequence
+from collections.abc import Sequence
 
 import torch
 
@@ -328,32 +328,6 @@ def resolve_q_buckets(spec: str, max_q: int) -> list[int]:
     return sorted(out)
 
 
-def flat_bucket_fits(total_tokens: int, q: int, buckets) -> bool:
-    """Can a captured flat num_tokens bucket hold ``total_tokens`` at per-seq
-    bound ``q``?
-
-    The predicate ``ModelRunner._dynamic_num_tokens_pad`` searches with, split
-    out so the ragged shrink can ask it BEFORE rewriting the batch. When no
-    bucket matches, that lookup returns None and every caller falls back to
-    ``bs * max_seqlen_q`` -- a replay over MORE tokens than the ragged rebuild
-    populated, whose tail still holds the previous step's ids. Those reach the
-    draft's Markov transition-table lookup as out-of-range indices: a
-    device-side trap attributed to whatever kernel was in flight, nowhere near
-    the shrink that caused it.
-
-    Args:
-        total_tokens: real flat token total the step would carry.
-        q: per-seq length bound (``num_spec_query_tokens``). Non-positive never
-            matches -- it would make the divisibility test meaningless.
-        buckets: captured flat num_tokens buckets. Empty/None means no flat
-            bucket set exists (FULL rather than PIECEWISE cudagraphs, or
-            pre-warmup), so no ragged shrink is representable.
-    """
-    if q <= 0 or not buckets:
-        return False
-    return any(b >= total_tokens and b % q == 0 for b in buckets)
-
-
 def quantize_to_bucket(q: int, buckets: list[int]) -> int:
     """Round a desired query length UP to the nearest available bucket.
 
diff --git a/atom/spec_decode/eagle_proposer.py b/atom/spec_decode/eagle_proposer.py
index 9771418018..9d4cd99a40 100644
--- a/atom/spec_decode/eagle_proposer.py
+++ b/atom/spec_decode/eagle_proposer.py
@@ -117,8 +117,9 @@ def _declare_draft_graphs(self):
         }
         self.step = DraftGraph(
             forward=self._step_forward,
+            epilogue=self._step_head,
+            capture_epilogue=True,
             inputs=inputs,
-            pads=True,
             warmup_inputs=self._step_warmup_inputs,
         )
         return (self.step,)
@@ -138,6 +139,16 @@ def _step_forward(self, running_bs, *, input_ids, positions, hidden_states):
             input_ids=input_ids, positions=positions, hidden_states=hidden_states
         )
 
+    def _step_head(self, out, running_bs, **_):
+        """The mid-step's draft ids, recorded with the backbone that made them.
+
+        Both come back because the next mid-step reads the hidden states and
+        they exist only inside the recording. Every row, never
+        `[:scheduled_bs]`: a capture bakes the length it was made at. Slicing
+        is the caller's, on the way out.
+        """
+        return out, self.model.compute_draft_ids(out)
+
     def _build_draft_model(self, model_class) -> nn.Module:
         draft_model_hf_config = self.speculative_config.draft_model_hf_config
         if self.speculative_config.method == "eagle3":
@@ -270,17 +281,17 @@ def compute_draft_kv(
             return
         forward_context = get_forward_context()
         context = forward_context.context
-        bs = context.scheduled_bs
-        anchors = next_token_ids[:bs]
+        scheduled_bs = context.scheduled_bs
+        anchors = next_token_ids[:scheduled_bs]
         if any(t < 0 for t in anchors):
             return
 
         # Anchor row per sequence = `cu_seqlens_q[1:] - 1`, the rule
         # `propose_draft_token_ids` uses on a pure prefill step.
-        last_token_indices = self.prepare_inputs(bs, 1)
+        last_token_indices = self.prepare_inputs(scheduled_bs, 1)
         anchor_ids = forward_context.context.draft_anchor_overrides
         assert anchor_ids is not None
-        anchor_ids = anchor_ids[:bs]
+        anchor_ids = anchor_ids[:scheduled_bs]
 
         # `positions` is the padded forward buffer; the target's own output row
         # count is this batch's real token count, and all three inputs below
@@ -344,7 +355,12 @@ def _step_warmup_inputs(self, running_bs, **staged):
         )
 
     def _enter_decode_metadata(
-        self, bs, running_bs, positions, last_token_indices, num_reject_tokens
+        self,
+        scheduled_bs,
+        running_bs,
+        positions,
+        last_token_indices,
+        num_reject_tokens,
     ):
         """Rewrite the target's attn_metadata to one row per sequence.
 
@@ -401,11 +417,11 @@ def _enter_decode_metadata(
         cu_seqlens_q[: running_bs + 1] = builder.row_ids[: running_bs + 1]
         if target_uses_mla and has_flat_kv:
             # MLA: block_size=1, kv_indptr tracks tokens
-            # Per REAL request: `num_reject_tokens` is bs-long and a pad row
+            # Per REAL request: `num_reject_tokens` is scheduled_bs-long, a pad row
             # rejected nothing. Their `kv_indptr` keeps what the target left,
             # which is one of its own valid ranges, so their reads stay in
             # bounds; their WRITES are what has to be neutralized, below.
-            kv_indptr[1 : bs + 1] -= torch.cumsum(num_reject_tokens, dim=0)
+            kv_indptr[1 : scheduled_bs + 1] -= torch.cumsum(num_reject_tokens, dim=0)
         if positions.ndim == 1:
             positions = torch.index_select(positions, 0, last_token_indices)
         else:
@@ -435,11 +451,12 @@ def propose(
         forward_context = get_forward_context()
         context = forward_context.context
         attn_metadata = forward_context.attn_metadata
-        bs = context.scheduled_bs
-        # Steps 1+ run at the row count the target just replayed, so every
-        # mid-step lands on a shape the startup sweep already warmed. No pass
-        # declared (MRoPE) means the loop runs the model directly, at `bs`.
-        running_bs = self.step.target_running_bs(bs, context) if self.step else bs
+        scheduled_bs = context.scheduled_bs
+        # Mid-steps run at the recorded shape the startup sweep warmed -- but
+        # only when one exists at this batch. Widening is for reaching a
+        # recording; without one it just hands the variable-length MoE gather
+        # padded rows it is contracted not to receive.
+        running_bs = context.running_bs if self.step is not None else scheduled_bs
         context.is_draft = True
 
         assert self.runner is not None
@@ -461,7 +478,10 @@ def propose(
             hidden_states = target_hidden_states
 
         draft_token_ids = torch.empty(
-            bs, self.mtp_k, dtype=next_token_ids.dtype, device=next_token_ids.device
+            scheduled_bs,
+            self.mtp_k,
+            dtype=next_token_ids.dtype,
+            device=next_token_ids.device,
         )
         if envs.ATOM_DEBUG_FORCE_SKIP_DRAFT_MODEL:
             draft_token_ids.fill_(-1)
@@ -476,7 +496,7 @@ def propose(
         # cache is indexed by runner.block_size blocks.
         if draft_uses_mha:
             attn_metadata.slot_mapping = var["slot_mapping"].gpu[: len(input_ids)]
-            attn_metadata.block_tables = var["block_tables"].gpu[:bs]
+            attn_metadata.block_tables = var["block_tables"].gpu[:scheduled_bs]
         elif attn_metadata.slot_mapping is not None:
             # Make MLA draft slot_mapping == q rows. DeepSeek-V4 uses
             # block_tables + context_lens (slot_mapping is None) — nothing to
@@ -492,34 +512,31 @@ def propose(
 
         # Mid-steps run padded and may replay; step 0 does neither, so it is
         # labelled as the plain batch it is.
-        step0_label = f"bs={bs}"
+        step0_label = f"bs={scheduled_bs}"
         mid_label = (
-            self.step.label(bs, running_bs, context) if self.step else step0_label
+            self.step.label(scheduled_bs, running_bs) if self.step else step0_label
         )
         for i in range(self.mtp_k):
-            # `tok` is this step's real row count, which is NOT `bs`: step 0 runs
-            # over the target's whole token stream (a prefill chunk, or
-            # bs*(mtp_k+1) on decode) and only steps 1+ run one row per sequence.
-            # Labelling both as `bs` collapsed the draft's two shape axes into
-            # one. Taken pre-PCP-split on purpose -- the sharded count is a
-            # rank-local detail, and this is the count `_refresh_dp_metadata`
-            # reports.
+            # `tok` is this step's real row count, NOT `scheduled_bs`: step 0
+            # runs the target's whole token stream (a prefill chunk, or
+            # scheduled_bs*(mtp_k+1) on decode) and only steps 1+ run one row
+            # per sequence -- labelling both as `scheduled_bs` collapsed the two
+            # shape axes into one. Taken pre-PCP-split on purpose: the sharded
+            # count is a rank-local detail.
             with record_function(
                 f"propose_eagle[{i}/{self.mtp_k} tok={input_ids.shape[0]} "
                 f"{mid_label if i else step0_label}]"
             ):
-                # Re-sync DP token
-                # The count the forward RUNS, which at i>=1 is the padded one
-                # -- DP sizes its all_gatherv from this, and aiter asserts the
-                # tensor matches. DSpark reports `running_bs * num_draft` for the
-                # same reason.
-                # Step 0 carries the target's whole token stream and pads
-                # nothing; steps 1+ run the padded row count, of which `bs`
-                # rows are real.
                 self._publish_draft_shape(
                     forward_context,
-                    scheduled_tokens=bs if i else input_ids.shape[0],
+                    # Steps 1+ run one row per sequence, padded to `running_bs`
+                    # with `scheduled_bs` of them real -- DP sizes its gather
+                    # from the count the forward RUNS. Step 0 carries the
+                    # target's stream and pads nothing.
+                    scheduled_tokens=scheduled_bs if i else input_ids.shape[0],
                     running_tokens=running_bs if i else input_ids.shape[0],
+                    # Only a widened mid-step runs a height the group agreed on.
+                    running_tokens_are_unified=i > 0 and self.step is not None,
                 )
                 # ---- Prefill Context Parallel (draft i==0 prefill) --------
                 # The draft's first pass is a prefill that reuses the target's
@@ -557,10 +574,10 @@ def propose(
                     self.model.model.set_skip_topk(False)
                 if i and self.step is not None:
                     # Steps 1+ are the declared pass: one row per sequence, at
-                    # a batch the startup sweep already warmed.
-                    ret_hidden_states = self.step.run(
+                    # a batch the startup sweep already warmed. The head rides
+                    # in the recording, so this hands back both of its outputs.
+                    ret_hidden_states, graphed_ids = self.step.run(
                         running_bs,
-                        context,
                         **self._stage_step_inputs(
                             running_bs, d_input_ids, d_positions, d_hidden
                         ),
@@ -571,6 +588,7 @@ def propose(
                         positions=d_positions,
                         hidden_states=d_hidden,
                     )
+                    graphed_ids = None
                 if pcp_draft_prefill:
                     ret_hidden_states = pcp_allgather_rerange(
                         ret_hidden_states, pcp_ws
@@ -585,8 +603,10 @@ def propose(
                 sample_hidden_states = (
                     torch.index_select(ret_hidden_states, 0, last_token_indices)
                     if i == 0
-                    else ret_hidden_states[:bs]
+                    else ret_hidden_states[:scheduled_bs]
                 )
+                # Only step 0 and a flavor with no declared pass land here; a
+                # recorded mid-step produced its ids inside the graph.
                 # Every draft model EagleProposer can build implements this --
                 # the DSpark archs in support_draft_model_arch_dict do not, but
                 # they are DSparkProposer's and never reach this loop. All of
@@ -595,7 +615,11 @@ def propose(
                 # compute_logits().argmax(-1) here because is_draft suppresses
                 # the LM head's prefill last-token slice. How the ids are
                 # produced stays the model's business, not this loop's.
-                new_draft_ids = self.model.compute_draft_ids(sample_hidden_states)
+                new_draft_ids = (
+                    graphed_ids[:scheduled_bs]
+                    if graphed_ids is not None
+                    else self.model.compute_draft_ids(sample_hidden_states)
+                )
                 draft_token_ids[:, i] = new_draft_ids
 
                 if i < self.mtp_k - 1:
@@ -606,7 +630,7 @@ def propose(
                     )
                     if i == 0:
                         i0_max_seqlen_q, positions = self._enter_decode_metadata(
-                            bs,
+                            scheduled_bs,
                             running_bs,
                             positions,
                             last_token_indices,
@@ -642,11 +666,11 @@ def propose(
                         attn_metadata.context_lens[:running_bs] += 1
                         positions += 1
                         mtp_decode_kwargs = {}
-                    # `bs` stays the REAL sequence count; the builder reads
-                    # the padded row count off `positions`, which is the
+                    # `scheduled_bs` stays the REAL sequence count; the builder
+                    # reads the padded row count off `positions`, which is the
                     # pass's own buffer and therefore already `running_bs` long.
                     workinfos = self.runner.attn_metadata_builder.prepare_mtp_decode(
-                        bs,
+                        scheduled_bs,
                         (
                             attn_metadata.max_seqlen_q
                             if not do_attn_metadata_update
@@ -693,8 +717,8 @@ def propose(
                     # Outside the branch above: a backend that RETURNS a
                     # slot_mapping computed it over `running_bs` rows too, so leaving
                     # this inside was how those pad rows kept a live slot.
-                    if running_bs > bs:
-                        attn_metadata.slot_mapping[bs:] = -1
+                    if running_bs > scheduled_bs:
+                        attn_metadata.slot_mapping[scheduled_bs:] = -1
 
                     input_ids = new_draft_ids
                     hidden_states = sample_hidden_states
diff --git a/atom/utils/forward_context.py b/atom/utils/forward_context.py
index 74347c483a..199c864d07 100644
--- a/atom/utils/forward_context.py
+++ b/atom/utils/forward_context.py
@@ -83,21 +83,36 @@ def make(
         # attn_metadata: Any,
         num_tokens: int,
         num_tokens_across_dp: torch.Tensor | None = None,
+        *,
+        unified: bool = False,
     ) -> "DPMetadata":
+        """Per-rank token counts for this pass, and the sizes derived from them.
+
+        Three ways to a table, and the caller picks by what it knows. Hand one
+        over and it is taken as given. Say `unified` and every rank runs
+        `num_tokens`, so it is that number repeated. Say neither and it is
+        discovered, at the cost of one CPU all_reduce.
 
+        `unified` cannot be checked here -- the assert below reads this rank's
+        own entry, which a repeated fill satisfies by construction. A wrong
+        claim surfaces later, as a fixed-size collective posted at mismatched
+        heights.
+        """
         assert parallel_config.data_parallel_size > 1
         dp_size = parallel_config.data_parallel_size
         dp_rank = parallel_config.data_parallel_rank
         batchsize = num_tokens
 
-        # If num_tokens_across_dp is None, it will be computed by all_reduce
-        # Otherwise, num_tokens_across_dp[dp_rank] should be equal to batchsize
-        assert (
-            num_tokens_across_dp is None or num_tokens_across_dp[dp_rank] == batchsize
+        # A supplied table already says what every rank runs: its own entry is
+        # this pass's height, and saying it again can only disagree.
+        assert num_tokens_across_dp is None or (
+            not unified and num_tokens_across_dp[dp_rank] == batchsize
         )
         if num_tokens_across_dp is None:
-            num_tokens_across_dp = DPMetadata.num_tokens_across_dp(
-                batchsize, dp_size, dp_rank
+            num_tokens_across_dp = (
+                torch.full((dp_size,), batchsize, dtype=torch.int32, device="cpu")
+                if unified
+                else DPMetadata.num_tokens_across_dp(batchsize, dp_size, dp_rank)
             )
         max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp)
         cu_tokens_across_dp_cpu = torch.cumsum(num_tokens_across_dp, dim=0)
@@ -163,130 +178,292 @@ class SpecDecodeMetadata:
 
 @dataclass(frozen=True)
 class ForwardMode:
-    """Per-step dispatch, and the two batch sizes ATOM names.
-
-    ``scheduled_bs`` is what this rank was handed; ``running_bs`` is what the
-    GPU is made to run. The second is DP-group-UNIFIED and that is its whole
-    point: MoE pads to it and a captured graph holds its collective at that
-    width, so a rank arriving with a different one hangs the group.
-
-    An op reads whichever its own tensors carry. They differ in one corner --
-    uniform decode running eager -- where MoE pads while attention keeps its
-    unpadded rows for aiter's ``t == t_slot``.
+    """One step's shape, settled once, in the only two units ATOM has.
+
+    ``running_bs`` sizes everything per-sequence (attention above all);
+    ``running_tokens`` everything per-row (MoE above all). ``scheduled_*`` is
+    this rank's own, and what a result slices back to.
+
+    They are agreed across DP to different extents, and the difference is the
+    point. ``running_bs`` is agreed on EVERY step -- it is the reduction on a
+    ladder every rank shares -- because a captured graph holds its collective
+    at that batch and a rank arriving with another one hangs the group.
+    ``running_tokens`` is agreed exactly when ``running_tokens_are_unified``:
+    off it, MoE takes the variable-length gather, whose contract is that the
+    rows it was handed are the rows THIS rank scheduled, so a group height
+    there would claim rows nobody allocated. An over-wide per-sequence array
+    costs a few sentinel rows; an over-wide height is a contract break.
+
+    Nothing downstream recomputes a width. ``running_bs * q`` is NOT
+    ``running_tokens`` on a ragged step, and every site that rederived one from
+    the other is a site where the two came to disagree.
+
+    ``decide`` owns the whole chain -- the collective, adopting its reductions
+    onto the batch, the ladder, both widths -- because the steps are not
+    separable: the reduced query length rewrites the batch, and everything else
+    is read off it afterwards.
     """
 
     use_cudagraph: bool
+    is_prefill: bool
     scheduled_bs: int
+    scheduled_tokens: int
     running_bs: int
-    is_prefill: bool
+    running_tokens: int
+    running_tokens_are_unified: bool
+    # Query rows per sequence this step. `num_spec_step + 1` unless a DSpark
+    # shrink cut it, then DP-MAX'd -- so it is the group's, like the widths.
+    max_seqlen_q: int
+    # A PIECEWISE recording exists at `running_tokens`, so this step may replay.
+    piecewise_captured: bool
+    # TBO's cross-DP AND, or the local answer on one rank -- there being no peer
+    # to veto it. Settled here so the gate has one source at every dp_size.
+    tbo_collective_active: bool
+    # What else the one per-step collective returned: per-rank token counts, the
+    # DSpark shape max. Not shape. None on a single rank.
+    sync: Any | None = None
 
     @classmethod
     def decide(
         cls,
         *,
-        is_prefill: bool,
-        scheduled_bs: int,
-        running_tokens: int,
-        dp_uniform_decode: bool,
+        batch,
+        dp_size: int,
+        dp_group,
         enforce_eager: bool,
         capture_sizes: list[int],
-        mtp_step: int = 1,
+        captured_tokens: list[int] | None,
+        is_block_drafter: bool,
+        tbo_on: bool,
+        local_tbo: tuple[bool, bool, int, int],
+        max_seqlen_q: int,
     ) -> "ForwardMode":
-        """Pick dispatch and the step's ``running_bs``. Any new force-eager
-        condition belongs here, not in caller-side checks."""
-        if is_prefill:
-            # Prefill is always eager, so nothing is padded and the two are the
-            # same number. Per-token sums ride cu_seqlens_q, not this.
-            return cls(
-                use_cudagraph=False,
-                scheduled_bs=scheduled_bs,
-                running_bs=scheduled_bs,
-                is_prefill=True,
-            )
-
-        # The sequences the step will run, recovered from the token count:
-        # DP-unified in uniform mode (`running_tokens` is the group max, set in
-        # ModelRunner._preprocess), local-equivalent otherwise.
-        unified_bs = (running_tokens + mtp_step - 1) // mtp_step
-
-        if not dp_uniform_decode:
-            # MoE goes through all_gatherv (variable-length, per-rank sizes),
-            # so pad_for_all_gather is never reached and there is nothing to
-            # unify. Running at the scheduled batch pads nothing.
-            return cls(
-                use_cudagraph=False,
-                scheduled_bs=scheduled_bs,
-                running_bs=scheduled_bs,
-                is_prefill=False,
-            )
+        """Run the step's DP collective and settle its shape from the result.
 
-        # From here on uniform decode: MoE WILL pad_for_all_gather, so
-        # running_bs MUST be the same number on every rank.
-        if enforce_eager or unified_bs > capture_sizes[-1]:
-            # Eager under uniform decode, either forced or above the largest
-            # captured size. This is the one corner where the two differ:
-            # attention keeps its own scheduled batch while MoE pads to the
-            # unified one.
-            return cls(
-                use_cudagraph=False,
+        Any new force-eager condition belongs here, not in a caller-side check.
+        """
+        # Lazy: `atom.utils.tbo`'s package init reaches back into this module.
+        from atom.utils.tbo.ubatching import sync_dp_metadata
+
+        is_prefill = batch.total_tokens_num_prefill > 0
+        scheduled_tokens = batch.total_tokens_num
+        # The whole batch, never a per-kind count: `schedule` returns its
+        # prefill batch before it can add a decode row, so the two are equal
+        # today, and this is the one that survives a mixed batch.
+        scheduled_bs = batch.total_seqs_num
+        meets_min, can_split, ub0, ub1 = local_tbo
+
+        sync = None
+        if dp_size > 1:
+            sync = sync_dp_metadata(
+                dp_group=dp_group,
+                dp_size=dp_size,
+                scheduled_tokens=scheduled_tokens,
                 scheduled_bs=scheduled_bs,
-                running_bs=unified_bs,
-                is_prefill=False,
+                is_prefill=is_prefill,
+                tbo_on=tbo_on,
+                local_meets_min_tokens=meets_min,
+                local_can_split=can_split,
+                local_ub_tokens=(ub0, ub1),
+                # Only a block drafter needs the group to agree on it; every
+                # other flavor keeps its own and the two extra wire rows are
+                # not sent.
+                max_seqlen_q=max_seqlen_q if is_block_drafter else None,
             )
-
-        # CUDAGraph: the smallest captured size that fits. Captured sizes are
-        # a static list every rank shares, so this is unified by construction.
+            # The group's query length, taken BEFORE the rows are read off it:
+            # a rank still on its local q settles on a different `running_tokens`
+            # than its peers.
+            if sync.max_seqlen_q_across_dp is not None:
+                max_seqlen_q = sync.max_seqlen_q_across_dp
+            unified_bs = sync.max_bs_across_dp
+            tbo_collective_active = sync.tbo_collective_active
+            # One meaning and only one: every rank is decoding, so the group max
+            # pads them level and `scheduled_tokens` reduces to a decode count on
+            # every rank. A prefilling peer makes the step ragged across the
+            # group, and then each rank runs its own count through the
+            # variable-length gather.
+            unified = not sync.any_rank_has_prefill
+        else:
+            # One rank pads to nobody, so it has one height by construction, and
+            # its TBO answer stands unvetoed.
+            unified_bs, unified = scheduled_bs, True
+            tbo_collective_active = tbo_on and meets_min and can_split
+
+        # ---- the batch: one rule, every step ---------------------------
+        # `unified_bs` is the reduction's own answer and the ladder is a rule
+        # every rank shares, so this is DP-agreed by construction -- not by a
+        # case analysis a future branch could get wrong. That is the whole
+        # property: a graph key, a draft pass and an attention plan all read
+        # this one field and cannot end up holding two widths.
+        #
+        # One lookup answers both of the step's questions. `capture_sizes` is
+        # ascending int32 (`ModelRunner.capture_sizes_np`, sorted where it is
+        # built and never mutated), so the insertion point IS the smallest size
+        # that holds the batch, and running off the end IS "no recording can".
+        rung = int(np.searchsorted(capture_sizes, unified_bs))
+        on_ladder = rung < len(capture_sizes)
+        running_bs = int(capture_sizes[rung]) if on_ladder else unified_bs
+
+        # ---- dispatch: a different question, asked separately ----------
+        # Whether the TARGET replays. It no longer decides how wide anything
+        # is; when the batch answered four ways depending on this, `running_bs`
+        # was per-rank on two of them and every consumer had to know which.
+        use_cudagraph = not is_prefill and unified and not enforce_eager and on_ladder
+
+        running_tokens, piecewise_captured = cls._running_tokens(
+            is_prefill=is_prefill,
+            unified=unified,
+            running_bs=running_bs,
+            scheduled_tokens=scheduled_tokens,
+            num_tokens_across_dp=None if sync is None else sync.num_tokens_across_dp,
+            captured_tokens=captured_tokens,
+            is_block_drafter=is_block_drafter,
+            max_seqlen_q=max_seqlen_q,
+        )
         return cls(
-            use_cudagraph=True,
+            use_cudagraph=use_cudagraph,
+            is_prefill=is_prefill,
             scheduled_bs=scheduled_bs,
-            running_bs=next((x for x in capture_sizes if x >= unified_bs), unified_bs),
-            is_prefill=False,
+            scheduled_tokens=scheduled_tokens,
+            running_bs=running_bs,
+            running_tokens=running_tokens,
+            running_tokens_are_unified=unified,
+            max_seqlen_q=max_seqlen_q,
+            piecewise_captured=piecewise_captured,
+            tbo_collective_active=tbo_collective_active,
+            sync=sync,
         )
 
-    @property
-    def attn_tensors_are_padded(self) -> bool:
-        """True iff per-token attention tensors carry padding this step
-        (cudagraph today; update if other padded layouts are added)."""
-        return self.use_cudagraph
+    @staticmethod
+    def _running_tokens(
+        *,
+        is_prefill: bool,
+        unified: bool,
+        running_bs: int,
+        scheduled_tokens: int,
+        num_tokens_across_dp,
+        captured_tokens,
+        is_block_drafter: bool,
+        max_seqlen_q: int,
+    ) -> tuple[int, bool]:
+        """``(running_tokens, piecewise_captured)``.
+
+        `pad_for_all_gather` asserts the first equals the rows the tensor really
+        has, so it is not "the biggest count anyone reported": it is what THIS
+        forward will produce.
+        """
+        if not unified:
+            # Nothing padded the tensor -- the varlen gather sends each rank its
+            # own count -- so a height from the group would claim rows nobody
+            # allocated.
+            return scheduled_tokens, False
+        if is_prefill:
+            # Reachable on one rank only: a group of one is unified whatever it
+            # runs, while under DP any prefill clears `unified` above. Prompts
+            # are ragged, so no product recovers a height and this rank's own
+            # count is the whole answer.
+            return scheduled_tokens, False
+
+        # The batch was rounded onto the ladder AFTER the sync, so this runs rows
+        # no rank reported; group-agreed all the same, both factors having come
+        # out of the reduction.
+        q = max_seqlen_q
+        running_tokens = running_bs * q
+        if not captured_tokens:
+            return running_tokens, False
+        if not is_block_drafter:
+            # Rectangular layout: one row per (seq, query), which every per-token
+            # buffer spans and `slot_mapping[i * q + j]` addresses through. Anything
+            # under the product would leave the tail of that rectangle unwritten.
+            return running_tokens, running_tokens in captured_tokens
+
+        # Marker-driven layout: the rows are one flat run, so the width is the
+        # smallest recorded one that holds it, not a product -- q-divisible
+        # because the per-seq structures tile the run at that stride. Every rank
+        # is decoding here, so `scheduled_tokens` IS a decode count on each of
+        # them and row 0's reduction is the group's run. `None` is one rank.
+        run = (
+            scheduled_tokens
+            if num_tokens_across_dp is None
+            else int(num_tokens_across_dp.max())
+        )
+        recorded = next((n for n in captured_tokens if n >= run and n % q == 0), None)
+        if recorded is None:
+            # Nothing recorded holds the run, so forward it eagerly rather than
+            # claim a graph nobody made at this width.
+            return max(run, running_tokens), False
+        return recorded, True
 
     def assert_shape_contract(
         self,
         input_ids: "torch.Tensor",
         attn_metadata: "AttentionMetaData",
     ) -> None:
-        """Validate ``input_ids`` / ``slot_mapping`` against this ForwardMode.
-
-        Skips prefill (variable-length) and cudagraph (graph-internal shape);
-        callers invoke unconditionally to keep dispatch decisions centralised.
+        """Tie what the builder published to what this step settled.
+
+        Covers the PADDED step, which is the only one that can be wrong: a
+        width is either both units' or neither's, and every way they have come
+        apart -- a builder on a stale batch, a per-seq array widened while the
+        per-token ones were not -- shows up as one of the equalities below. The
+        version that skipped whenever ``running_bs != scheduled_bs`` exempted
+        exactly that case, so a mismatch reached the model and surfaced as lost
+        accuracy rather than as a failure here.
+
+        Shapes only, never values: reading a device tensor's contents on this
+        path is a D2H sync per step.
+
+        Prefill is still out. Its rows are the prompt's, so no product of a
+        batch and a query length describes them, and its per-seq tail is padded
+        for a following draft pass rather than for anything this step runs.
         """
-        if self.is_prefill or self.attn_tensors_are_padded:
+        if self.is_prefill or input_ids is None or attn_metadata is None:
             return
-        if (
-            input_ids is None
-            or attn_metadata is None
-            or attn_metadata.slot_mapping is None
-        ):
-            return
-        # Past the early-return nothing is padded, so the rows attention runs
-        # are exactly the scheduled ones.
-        max_q = attn_metadata.max_seqlen_q
-        expected = self.scheduled_bs * max_q
-        actual_in = input_ids.shape[0]
-        actual_slot = attn_metadata.slot_mapping.shape[0]
-        assert actual_in == expected, (
-            f"eager input_ids length {actual_in} != scheduled_bs*max_q="
-            f"{expected} ({self})"
+
+        def _rows(name):
+            t = getattr(attn_metadata, name, None)
+            return None if t is None else int(t.shape[0])
+
+        # `input_ids` is the argument, this rank's own rows -- the cudagraph
+        # branch re-slices the buffer to `running_tokens` itself.
+        assert input_ids.shape[0] == self.scheduled_tokens, (
+            f"input_ids length {input_ids.shape[0]} != scheduled_tokens="
+            f"{self.scheduled_tokens} ({self})"
+        )
+        assert self.scheduled_tokens <= self.running_tokens, (
+            f"running_tokens={self.running_tokens} is below the rows this rank "
+            f"scheduled ({self.scheduled_tokens}) ({self})"
         )
-        assert actual_slot == expected, (
-            f"eager slot_mapping length {actual_slot} != scheduled_bs*max_q="
-            f"{expected} ({self}); attn_metadata_builder used a stale bs"
+        # Per-token: sized by `running_tokens`, the height everything per-row runs.
+        slot_rows = _rows("slot_mapping")
+        assert slot_rows in (None, self.running_tokens), (
+            f"slot_mapping length {slot_rows} != running_tokens="
+            f"{self.running_tokens} ({self}); the builder used another width"
         )
+        # Per-sequence: sized by `running_bs`. `cu_seqlens_q` carries the extra
+        # boundary entry; the state slots are V4's and absent elsewhere.
+        cu_rows = _rows("cu_seqlens_q")
+        assert cu_rows in (None, self.running_bs + 1), (
+            f"cu_seqlens_q length {cu_rows} != running_bs+1="
+            f"{self.running_bs + 1} ({self})"
+        )
+        slot_out_rows = _rows("state_slot_out")
+        assert slot_out_rows in (None, self.running_bs), (
+            f"state_slot_out length {slot_out_rows} != running_bs="
+            f"{self.running_bs} ({self})"
+        )
+        # No rectangle assertion: `_running_tokens` deliberately returns a
+        # height ABOVE `running_bs * max_seqlen_q` when nothing recorded holds
+        # the run.
 
 
 def running_tokens_from_bs(bs: int, *, is_prefill: bool, attn_metadata) -> int:
-    """``running_tokens`` for a bridge that only knows a request count.
+    """``running_tokens`` for a BRIDGE that only knows a request count.
+
+    Native ATOM never calls this -- it reads `ForwardMode.running_tokens`. A
+    plugin bridge has no ForwardMode, so it reconstructs the rectangle here and
+    nowhere else.
 
-    Prefill returns that count verbatim, which is NOT a height: prompt lengths
+    Prefill returns the count verbatim, which is NOT a height: prompt lengths
     are ragged, so no multiplier recovers one, and a number below the real
     height just leaves the batch unpadded -- correct there, since a prefill
     all_gather is variable-length.
@@ -310,19 +487,20 @@ class Context:
     # The step's DP-unified padded shape. `running_bs` counts SEQUENCES (graph
     # identity, the draft's pad width), `running_tokens` the hidden_states rows
     # MoE pads to. Both stored, because the ratio is not always max_seqlen_q --
-    # a DSpark ragged step runs a flat bucket no rectangular bs*q recovers.
+    # a DSpark ragged step runs a packed width no rectangular bs*q recovers.
     running_bs: int = 0
     running_tokens: int = 0
     is_draft: bool = False
-    # True iff all DP ranks are running pure decode this step (DP-disabled
-    # case is treated as True). Mirrors vLLM's `uniform_decode` flag and
-    # gates DP-specific variable-length all_gather/scatter paths.
-    dp_uniform_decode: bool = True
-    # Single source of truth for cudagraph vs eager dispatch + the step's
-    # scheduled_bs / running_bs.
-    # Set by prepare_inputs via ForwardMode.decide(). None only on legacy
-    # paths that haven't been routed through it (run_model falls back to
-    # the original four-OR derivation in that case for back-compat).
+    # Every rank of the group is decoding, so `running_tokens` above is the
+    # group's number and not this rank's own. That is what picks MoE's
+    # fixed-size all_gather over the variable-length one, and so what makes a
+    # captured graph usable at all. One meaning, deliberately: it once also
+    # stood for "dp-attention is off", and readers needing the strong claim --
+    # `_can_use_dp_sharded_head` above all -- silently got the weak one.
+    running_tokens_are_unified: bool = True
+    # The step's whole shape decision. Set by `prepare_model` via
+    # `ForwardMode.decide`; None only on a capture context, which declares its
+    # shape rather than reading one.
     forward_mode: ForwardMode | None = None
     # Optional flat token ids for the current forward. Read by callbacks
     # invoked inside Dynamo-opaque custom ops (e.g. V4 MoE hash routing)
@@ -351,7 +529,7 @@ def __init__(
         running_bs: int = 0,
         running_tokens: int = 0,
         is_draft: bool = False,
-        dp_uniform_decode: bool = True,
+        running_tokens_are_unified: bool = True,
         forward_mode: ForwardMode | None = None,
         input_ids: torch.Tensor | None = None,
         ubatch_token_offset: int = 0,
@@ -366,7 +544,7 @@ def __init__(
         self.running_bs = running_bs
         self.running_tokens = running_tokens
         self.is_draft = is_draft
-        self.dp_uniform_decode = dp_uniform_decode
+        self.running_tokens_are_unified = running_tokens_are_unified
         self.forward_mode = forward_mode
         self.input_ids = input_ids
         self.ubatch_token_offset = ubatch_token_offset
@@ -552,8 +730,8 @@ class ForwardContext:
 
     ubatch_slices: list[Any] | None = None
 
-    # Cross-DP MAX of per-ubatch token counts, precomputed in
-    # ``ModelRunner._preprocess`` and propagated here so
+    # Cross-DP MAX of per-ubatch token counts, reduced in the step's one packed
+    # all_gather (``ForwardMode.decide``) and propagated here so
     # ``UBatchWrapper._compute_ub_graph_bs`` no longer needs its own
     # per-ubatch all_reduce. Shape: tuple of length N == len(ubatch_slices).
     # None when DP is off or when TBO is not active this step.
diff --git a/atom/utils/tbo/ubatch_wrapper.py b/atom/utils/tbo/ubatch_wrapper.py
index ca4cdca92d..915ac827fb 100644
--- a/atom/utils/tbo/ubatch_wrapper.py
+++ b/atom/utils/tbo/ubatch_wrapper.py
@@ -459,7 +459,7 @@ def _decode_ub_running_bs(
         pads each ubatch to this size, so a per-rank-local split (which differs
         when ranks carry different decode batch sizes, e.g. during drain)
         desyncs the collective and faults. Derive it from the DP-unified
-        ``ub_max_tokens_across_dp`` (MAX-reduced in ModelRunner._preprocess),
+        ``ub_max_tokens_across_dp`` (MAX-reduced in ``ForwardMode.decide``),
         converting the per-ubatch token max back to a request count via
         ``max_seqlen_q``. Falls back to the local split only when DP is off or
         the precomputed value is unavailable.
@@ -487,7 +487,7 @@ def _compute_ub_running_tokens(
         ``dp_size``: pad_for_all_gather gathers across ranks itself.
 
         For prefill (eager only): the cross-DP per-ubatch token MAX that
-            ``ModelRunner._preprocess`` already packed into the single DP
+            ``ForwardMode.decide`` already packed into the single DP
             all_reduce (``ctx.ub_max_tokens_across_dp``). Falls back to
             local sizes when DP is off / value not precomputed.
         For decode: the per-ubatch padded request count times ``max_seqlen_q``.
diff --git a/atom/utils/tbo/ubatching.py b/atom/utils/tbo/ubatching.py
index 69fe5911d5..5fe4dbbbb9 100644
--- a/atom/utils/tbo/ubatching.py
+++ b/atom/utils/tbo/ubatching.py
@@ -24,7 +24,7 @@ def tbo_enabled() -> bool:
 
 # =====================================================================
 # DP-side TBO helpers (formerly ModelRunner._local_tbo_precompute and
-# the inline sync block inside ModelRunner._preprocess).
+# the inline sync block inside ForwardMode.decide).
 # =====================================================================
 
 
@@ -175,6 +175,11 @@ class DPSyncResult:
 
     # [dp_size] int32 CPU tensor — each rank's input token count.
     num_tokens_across_dp: torch.Tensor
+    # DP-MAX of the scheduled SEQUENCE count -- the companion to
+    # `num_tokens_across_dp` in the other unit. Reduced on EVERY step, not just
+    # uniform decode: a captured graph holds its collective at one batch, so a
+    # per-rank count is precisely what splits the group into two widths.
+    max_bs_across_dp: int
     # True iff ANY rank has at least one prefill seq this step.
     any_rank_has_prefill: bool
     # True iff TBO on AND OR(meets_min_tokens) AND AND(can_split) AND uniform.
@@ -184,17 +189,11 @@ class DPSyncResult:
     tbo_collective_active: bool
     # (ub0_max, ub1_max) across DP — only set when tbo_collective_active.
     ub_max_tokens_across_dp: tuple[int, int] | None
-    # DP-MAX of the DSpark decode shape (q, decode_bs, total_tokens) — only set when
-    # ``dspark_shape`` was passed in (DSpark active + DP). Folded into this
-    # packed all_gather so DSpark's graph-shape sync no longer needs its own
-    # separate all_reduce (halves the per-step DP round-trips).
-    dspark_shape_max: tuple[int, int, int] | None = None
-    # True iff EVERY DP rank is running a dummy forward this step (AND-reduce of
-    # per-rank is_dummy). Used to gate DSpark cudagraph dummy-replay: an all-dummy
-    # step has no real rank replaying to rendezvous with, so all ranks must stay
-    # eager (vLLM never runs an all-idle forward). Only meaningful when
-    # ``dspark_shape`` was passed (DSpark + DP); False otherwise.
-    any_dummy: bool = False
+    # DP-MAX of the per-seq decode length — only set when ``max_seqlen_q`` was
+    # passed in (a block drafter under DP). Folded into this packed all_gather so
+    # the graph-shape sync no longer needs its own separate all_reduce (halves
+    # the per-step DP round-trips).
+    max_seqlen_q_across_dp: int | None = None
 
 
 def sync_dp_metadata(
@@ -202,64 +201,66 @@ def sync_dp_metadata(
     dp_group,
     dp_size: int,
     scheduled_tokens: int,
+    scheduled_bs: int,
     is_prefill: bool,
     tbo_on: bool,
     local_meets_min_tokens: bool = False,
     local_can_split: bool = False,
     local_ub_tokens: tuple[int, int] = (0, 0),
-    dspark_shape: tuple[int, int, int] | None = None,
-    local_is_dummy: bool = False,
+    max_seqlen_q: int | None = None,
 ) -> DPSyncResult:
     """Single packed DP all_gather over all per-rank scalars a decode/prefill
     step needs synchronized: DP token padding, the prefill fan-out, the
     cross-DP TBO gate, and (when active) the DSpark graph-shape MAX.
 
-    DSpark folds its [q, bs, total_tokens] shape sync in here too (3 extra
-    fields), so the step issues one collective instead of two.
+    A block drafter folds its per-seq length sync in here too (one extra field),
+    so the step issues one collective instead of two.
 
     Pre-Plan-B this required up to 3 separate all_reduces per step
     (``get_dp_padding`` / this sync / a third inside
     ``UBatchWrapper``). Now one all_gather of ``n_fields`` int32 values
-    per rank suffices. When TBO is off only the first 2 fields are
-    exchanged (saves 60 % payload + skips :func:`local_tbo_precompute`
+    per rank suffices. When TBO is off only the first 3 fields are
+    exchanged (saves 57 % payload + skips :func:`local_tbo_precompute`
     at the call site).
 
     Layout (``sync`` is ``[n_fields, dp_size]``):
 
-      row 0 : scheduled_tokens        -> num_tokens_across_dp
-      row 1 : is_prefill (0/1)         -> any_rank_has_prefill (OR)
-      row 2 : meets_min_tokens (0/1)   -> OR  -> any rank reached the min-token bar [TBO only]
-      row 3 : can_split (0/1)          -> AND -> every rank can split              [TBO only]
-      row 4 : ub0_tokens               -> ub_max_tokens_across_dp[0]              [TBO only]
-      row 5 : ub1_tokens               -> ub_max_tokens_across_dp[1]              [TBO only]
-      row k+0 : dspark q               -> dspark_shape_max[0] (MAX)     [DSpark only]
-      row k+1 : dspark decode_bs       -> dspark_shape_max[1] (MAX)     [DSpark only]
-      row k+2 : dspark total_tokens    -> dspark_shape_max[2] (MAX)     [DSpark only]
+      row 0 : scheduled_tokens         -> num_tokens_across_dp
+      row 1 : scheduled_bs             -> max_bs_across_dp (MAX)
+      row 2 : is_prefill (0/1)         -> any_rank_has_prefill (OR)
+      row 3 : meets_min_tokens (0/1)   -> OR  -> any rank reached the min-token bar [TBO only]
+      row 4 : can_split (0/1)          -> AND -> every rank can split              [TBO only]
+      row 5 : ub0_tokens               -> ub_max_tokens_across_dp[0]              [TBO only]
+      row 6 : ub1_tokens               -> ub_max_tokens_across_dp[1]              [TBO only]
+      row k+0 : max_seqlen_q           -> max_seqlen_q_across_dp (MAX)  [DSpark only]
+
+    Rows 0 and 1 are the same batch in ATOM's two units; both ride this one
+    collective because agreeing on only one of them still leaves two shapes.
 
     Gate: ``active = OR(meets_min_tokens) AND AND(can_split) AND uniform``.
 
-    ``dspark_shape`` folds DSpark's per-step graph-shape sync (the DP-MAX of
-    (q, decode_bs, total_tokens) from ``_dspark_local_shape``) into this same
-    packed all_gather. It is global-config gated (DSpark on + DP), so every
-    rank appends the same 3 fields and the payload stays symmetric.
+    ``max_seqlen_q`` folds a block drafter's per-step graph-shape sync into this
+    same packed all_gather. It is global-config gated (DSpark on + DP), so every
+    rank appends the same field and the payload stays symmetric.
+
+    Neither the batch nor the token total needs a row of its own here: rows 1 and
+    0 already carry them, and a step that reaches the reduction with every rank
+    decoding has `scheduled_tokens` == its decode total on every rank.
     """
-    tbo_fields = 6 if tbo_on else 2
-    dspark_on = dspark_shape is not None
-    # +1 extra DSpark field: per-rank is_dummy (OR-reduced -> any_dummy).
-    n_fields = tbo_fields + (4 if dspark_on else 0)
+    tbo_fields = 7 if tbo_on else 3
+    dspark_on = max_seqlen_q is not None
+    n_fields = tbo_fields + (1 if dspark_on else 0)
     local = torch.zeros(n_fields, dtype=torch.int32, device="cpu")
     local[0] = scheduled_tokens
-    local[1] = 1 if is_prefill else 0
+    local[1] = scheduled_bs
+    local[2] = 1 if is_prefill else 0
     if tbo_on:
-        local[2] = 1 if local_meets_min_tokens else 0
-        local[3] = 1 if local_can_split else 0
-        local[4] = local_ub_tokens[0]
-        local[5] = local_ub_tokens[1]
+        local[3] = 1 if local_meets_min_tokens else 0
+        local[4] = 1 if local_can_split else 0
+        local[5] = local_ub_tokens[0]
+        local[6] = local_ub_tokens[1]
     if dspark_on:
-        local[tbo_fields + 0] = dspark_shape[0]
-        local[tbo_fields + 1] = dspark_shape[1]
-        local[tbo_fields + 2] = dspark_shape[2]
-        local[tbo_fields + 3] = 1 if local_is_dummy else 0
+        local[tbo_fields + 0] = max_seqlen_q
 
     gathered = [
         torch.empty(n_fields, dtype=torch.int32, device="cpu") for _ in range(dp_size)
@@ -268,7 +269,8 @@ def sync_dp_metadata(
     sync = torch.stack(gathered, dim=1)  # [n_fields, dp_size]
 
     num_tokens_across_dp = sync[0]
-    any_rank_has_prefill = bool(sync[1].any())
+    max_bs_across_dp = int(sync[1].max())
+    any_rank_has_prefill = bool(sync[2].any())
     tbo_collective_active = False
     ub_max_tokens_across_dp: tuple[int, int] | None = None
     if tbo_on:
@@ -277,40 +279,33 @@ def sync_dp_metadata(
         # that rank would run 1 ubatch while peers run 2 → per-ubatch collective
         # size mismatch → RCCL hang. Under-filled-but-splittable ranks are then
         # force-split (see maybe_create_ubatch_slices force=True) to stay aligned.
-        tbo_collective_active = bool(sync[2].any()) and bool(sync[3].all())
+        tbo_collective_active = bool(sync[3].any()) and bool(sync[4].all())
         # Mixed-mode guard: ALWAYS require a uniform batch mode (all prefill or
         # all decode) across DP. A prefill rank running 2 ubatches alongside a
         # decode rank running 2 ubatches still issues different collectives per
         # ubatch → hang. (Previously gated behind require_uniform_mode; with the
         # OR-reduce this must be unconditional.)
         if tbo_collective_active:
-            prefill_rank_count = int(sync[1].sum())
+            prefill_rank_count = int(sync[2].sum())
             uniform_mode = prefill_rank_count == 0 or prefill_rank_count == dp_size
             tbo_collective_active = uniform_mode
         if tbo_collective_active:
             ub_max_tokens_across_dp = (
-                int(sync[4].max()),
                 int(sync[5].max()),
+                int(sync[6].max()),
             )
 
-    dspark_shape_max: tuple[int, int, int] | None = None
-    any_dummy = False
+    max_seqlen_q_across_dp: int | None = None
     if dspark_on:
-        dspark_shape_max = (
-            int(sync[tbo_fields + 0].max()),
-            int(sync[tbo_fields + 1].max()),
-            int(sync[tbo_fields + 2].max()),
-        )
-        # OR-reduce: True if ANY rank is a dummy this step.
-        any_dummy = bool(sync[tbo_fields + 3].any())
+        max_seqlen_q_across_dp = int(sync[tbo_fields + 0].max())
 
     return DPSyncResult(
         num_tokens_across_dp=num_tokens_across_dp,
+        max_bs_across_dp=max_bs_across_dp,
         any_rank_has_prefill=any_rank_has_prefill,
         tbo_collective_active=tbo_collective_active,
         ub_max_tokens_across_dp=ub_max_tokens_across_dp,
-        dspark_shape_max=dspark_shape_max,
-        any_dummy=any_dummy,
+        max_seqlen_q_across_dp=max_seqlen_q_across_dp,
     )
 
 
diff --git a/docs/environment_variables.md b/docs/environment_variables.md
index 99f04e40e4..be971c8a3c 100644
--- a/docs/environment_variables.md
+++ b/docs/environment_variables.md
@@ -142,7 +142,7 @@ reachable shape one set rather than two lists that drift. The switch below decid
 
 | Variable | Type | Default | Description |
 |----------|------|---------|-------------|
-| **ATOM_DRAFT_CUDAGRAPH** | bool | 1 (true) | Capture each declared draft pass into a per-`capture_sizes` CUDAGraph as it is warmed, so a draft pass replays instead of relaunching every kernel. `0` keeps the warmup (and therefore the JIT saving) but drafts eagerly. Only passes that can pad are captured — a graph is one shape — so this is inert wherever padding is declined (EPLB, the separate-draft Kimi-K3 path). A DP-sync dummy DOES replay, in lockstep with the ranks holding work — `is_dummy_run` is per-rank, so gating on it splits one DP group across two collectives. Measured on V4-Flash-DSpark tp1: GSM8K 0.9527 / acceptance 65.25% captured against 0.9497 / 65.21% eager, i.e. indistinguishable; on tp4 with the LM head inside the capture, draft kernel launches went 30 → 0 per pass and draft wall time 915.8 → 118.9 µs. Read per pass at warmup time, so set it before the server starts. Grep a trace for a trailing ` graph` in a `propose_*` label to confirm which passes replayed. |
+| **ATOM_DRAFT_CUDAGRAPH** | bool | 1 (true) | Capture each declared draft pass into a per-`capture_sizes` CUDAGraph as it is warmed, so a draft pass replays instead of relaunching every kernel. `0` keeps the warmup (and therefore the JIT saving) but drafts eagerly. Only passes that declare a graph are captured — the separate-draft Kimi-K3 path declares none, so this is inert there. EPLB no longer declines the padding: the target pads on every cudagraph decode step and its rows reach the same expert-load recorder, so declining on the draft protected nothing. A DP-sync dummy DOES replay, in lockstep with the ranks holding work — `is_dummy_run` is per-rank, so gating on it splits one DP group across two collectives. Measured on V4-Flash-DSpark tp1: GSM8K 0.9527 / acceptance 65.25% captured against 0.9497 / 65.21% eager, i.e. indistinguishable; on tp4 with the LM head inside the capture, draft kernel launches went 30 → 0 per pass and draft wall time 915.8 → 118.9 µs. Read per pass at warmup time, so set it before the server starts. Grep a trace for a trailing ` graph` in a `propose_*` label to confirm which passes replayed. |
 
 ### DSpark drafting
 
diff --git a/tests/plugin/test_sglang_gdn_forward_context.py b/tests/plugin/test_sglang_gdn_forward_context.py
index d78afdf576..87284e10fe 100644
--- a/tests/plugin/test_sglang_gdn_forward_context.py
+++ b/tests/plugin/test_sglang_gdn_forward_context.py
@@ -78,7 +78,7 @@ def test_build_context_falls_back_to_local_decode_mode(monkeypatch):
 
     context, num_tokens = SGLangGDNForwardContext._build_context(forward_batch)
 
-    assert context.dp_uniform_decode
+    assert context.running_tokens_are_unified
     assert num_tokens == 2
 
 
diff --git a/tests/test_dp_metadata.py b/tests/test_dp_metadata.py
new file mode 100644
index 0000000000..3f4d95927c
--- /dev/null
+++ b/tests/test_dp_metadata.py
@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: MIT
+"""How `DPMetadata.make` arrives at the per-rank token counts.
+
+Three ways in, and the choice is a correctness question, not a performance one:
+the table sizes every DP collective in the forward, so a table that does not
+describe what the ranks will actually run shows up as a fixed-size collective
+posted at mismatched heights -- a hang on eight ranks, nothing here.
+
+No aiter and no DP group: the all_reduce lives behind
+`DPMetadata.num_tokens_across_dp`, which these tests replace. That is also what
+makes "did it ask the group" observable at all.
+"""
+
+import types
+
+import pytest
+import torch
+
+from atom.utils.forward_context import DPMetadata
+
+
+def _cfg(dp_size=4, dp_rank=1):
+    return types.SimpleNamespace(data_parallel_size=dp_size, data_parallel_rank=dp_rank)
+
+
+@pytest.fixture
+def no_group(monkeypatch):
+    """Make the ask itself fail, rather than count calls.
+
+    A version that asks would pass against a stub that answers.
+    """
+
+    def _refuse(*a, **k):
+        raise AssertionError("asked the group for a height it was told")
+
+    monkeypatch.setattr(DPMetadata, "num_tokens_across_dp", staticmethod(_refuse))
+
+
+def test_a_unified_height_is_written_out_rather_than_reduced(no_group):
+    """Every rank runs `num_tokens`, so every rank can write the whole table."""
+    md = DPMetadata.make(_cfg(dp_size=4), 336, unified=True)
+
+    assert md.get_sizes_across_dp() == [336, 336, 336, 336]
+    assert md.max_tokens_across_dp == 336
+    assert md.cu_tokens_across_dp_cpu.tolist() == [336, 672, 1008, 1344]
+
+
+def test_a_height_not_declared_unified_is_still_discovered(monkeypatch):
+    """The default is unchanged: no table and no claim means one all_reduce.
+
+    Its answer is honoured whole -- the reduction, not this rank's own count,
+    is what the derived sizes come from.
+    """
+    asked = {}
+
+    def _reduce(num_tokens, dp_size, dp_rank):
+        asked["args"] = (num_tokens, dp_size, dp_rank)
+        return torch.tensor([10, 336, 7, 1], dtype=torch.int32)
+
+    monkeypatch.setattr(DPMetadata, "num_tokens_across_dp", staticmethod(_reduce))
+
+    md = DPMetadata.make(_cfg(dp_size=4, dp_rank=1), 336)
+
+    assert asked["args"] == (336, 4, 1)
+    assert md.get_sizes_across_dp() == [10, 336, 7, 1]
+    assert md.max_tokens_across_dp == 336
+
+
+def test_a_supplied_table_and_a_unified_claim_cannot_both_be_given(no_group):
+    """Handing over the table already says what every rank runs.
+
+    Saying it again can only agree redundantly or disagree silently, and the
+    silent case is the one that costs eight ranks.
+    """
+    table = torch.tensor([336, 336, 336, 336], dtype=torch.int32)
+
+    with pytest.raises(AssertionError):
+        DPMetadata.make(_cfg(dp_size=4), 336, table, unified=True)
+
+    # ...and the table alone is still taken as given.
+    md = DPMetadata.make(_cfg(dp_size=4), 336, table)
+    assert md.get_sizes_across_dp() == [336, 336, 336, 336]
+
+
+def test_a_unified_claim_is_not_checked_against_the_group(no_group):
+    """It cannot be, and the pass that makes it has to know that.
+
+    `make` sees one rank, and the entry it asserts is the one a repeated fill
+    satisfies by construction -- so a caller declaring uniformity off a
+    rank-local number gets a well-formed table describing a group that does not
+    exist. Pinned because the guarantee lives entirely in the callers.
+    """
+    md = DPMetadata.make(_cfg(dp_size=4, dp_rank=3), 7, unified=True)
+
+    assert md.get_sizes_across_dp() == [7, 7, 7, 7]
diff --git a/tests/test_dp_sync_layout.py b/tests/test_dp_sync_layout.py
new file mode 100644
index 0000000000..c701d293ad
--- /dev/null
+++ b/tests/test_dp_sync_layout.py
@@ -0,0 +1,164 @@
+# SPDX-License-Identifier: MIT
+"""The packed DP all_gather's row layout.
+
+Every per-step scalar a DP group must agree on rides one collective as a row of
+one int32 tensor, so the rows are a wire format: pack and unpack are two pieces
+of index arithmetic that have to stay in step. Nothing else checks them -- a
+shifted row reads a plausible number out of the wrong field and shows up as a
+hang on eight ranks, not as a failure here.
+
+The collective itself is faked, so this runs on CPU with no DP group.
+"""
+
+import torch
+
+from atom.utils.tbo.ubatching import sync_dp_metadata
+
+
+class _FakeGroup:
+    """Stands in for the ranks. Each entry is one rank's field vector."""
+
+    def __init__(self, peers):
+        self.peers = peers
+
+
+def _sync(monkeypatch, *, peers, **kw):
+    """Run `sync_dp_metadata` against `peers` other ranks' field vectors.
+
+    `peers` is a list of callables taking this rank's packed tensor and
+    returning that rank's -- built relative to ours so a test says only how a
+    peer DIFFERS.
+    """
+
+    def fake_all_gather(out_list, local, group=None):
+        out_list[0].copy_(local)
+        for i, make in enumerate(peers, start=1):
+            out_list[i].copy_(make(local))
+
+    monkeypatch.setattr(torch.distributed, "all_gather", fake_all_gather)
+    return sync_dp_metadata(
+        dp_group=_FakeGroup(peers),
+        dp_size=1 + len(peers),
+        **kw,
+    )
+
+
+def _base(**kw):
+    base = {
+        "scheduled_tokens": 10,
+        "scheduled_bs": 2,
+        "is_prefill": False,
+        "tbo_on": False,
+    }
+    return {**base, **kw}
+
+
+def _with(idx, value):
+    """A peer identical to us except in row `idx`."""
+
+    def make(local):
+        peer = local.clone()
+        peer[idx] = value
+        return peer
+
+    return make
+
+
+def test_the_batch_is_reduced_by_max_and_not_by_something_that_looks_like_it(
+    monkeypatch,
+):
+    """A rank holding 2 sequences alongside one holding 7 must leave with 7.
+
+    Chosen so no other reduction lands on the same answer: sum is 9, min is 2,
+    ours is 2. Only MAX gives 7, so a row read out of the wrong field or
+    reduced the wrong way cannot pass.
+    """
+    r = _sync(monkeypatch, peers=[_with(1, 7)], **_base(scheduled_bs=2))
+    assert r.max_bs_across_dp == 7
+
+
+def test_the_batch_and_the_token_count_do_not_read_each_other(monkeypatch):
+    """They are the same step in two units and sit in adjacent rows.
+
+    Given distinct values on both, a one-row shift would swap them and every
+    downstream shape would be wrong by a factor of the per-seq token count --
+    which on a decode step is a plausible-looking number.
+    """
+    r = _sync(
+        monkeypatch,
+        peers=[_with(0, 99)],
+        **_base(scheduled_tokens=10, scheduled_bs=3),
+    )
+    assert int(r.num_tokens_across_dp.max()) == 99
+    assert r.max_bs_across_dp == 3
+
+
+def test_the_prefill_flag_still_reads_as_a_flag_after_the_batch_row(monkeypatch):
+    """`is_prefill` moved down one row when the batch was folded in.
+
+    It OR-reduces, so a row shift onto the batch would make it True whenever
+    any rank had a non-empty batch -- which is every step, and it silently
+    disables the uniform-decode path for the whole group.
+    """
+    assert not _sync(
+        monkeypatch, peers=[], **_base(is_prefill=False)
+    ).any_rank_has_prefill
+    assert _sync(monkeypatch, peers=[], **_base(is_prefill=True)).any_rank_has_prefill
+    # A peer prefilling turns it on for everyone; our own batch does not.
+    r = _sync(
+        monkeypatch, peers=[_with(2, 1)], **_base(is_prefill=False, scheduled_bs=5)
+    )
+    assert r.any_rank_has_prefill and r.max_bs_across_dp == 5
+
+
+def test_the_tbo_rows_survive_the_batch_being_folded_in(monkeypatch):
+    """TBO's four fields shifted down with everything else.
+
+    `can_split` AND-reduces and `meets_min_tokens` OR-reduces, so reading each
+    other's row inverts the gate. Pinned with the two disagreeing.
+    """
+    r = _sync(
+        monkeypatch,
+        peers=[],
+        **_base(
+            tbo_on=True,
+            local_meets_min_tokens=True,
+            local_can_split=True,
+            local_ub_tokens=(4, 6),
+        ),
+    )
+    assert r.tbo_collective_active
+    assert r.ub_max_tokens_across_dp == (4, 6)
+
+    # One rank that cannot split vetoes the whole group.
+    vetoed = _sync(
+        monkeypatch,
+        peers=[_with(4, 0)],
+        **_base(
+            tbo_on=True,
+            local_meets_min_tokens=True,
+            local_can_split=True,
+            local_ub_tokens=(4, 6),
+        ),
+    )
+    assert not vetoed.tbo_collective_active
+
+
+def test_the_dspark_block_starts_after_the_tbo_block_in_both_widths(monkeypatch):
+    """DSpark's rows are placed at an offset computed from `tbo_on`.
+
+    Both widths are exercised because the offset is the thing that moved, and
+    with TBO off it is the only arithmetic between the fixed head and DSpark.
+    The length is read off a PEER carrying a value no other row holds, so a
+    shifted offset reads one of the head's fields and reports it.
+    """
+    for tbo_on in (False, True):
+        head = 7 if tbo_on else 3
+        r = _sync(
+            monkeypatch,
+            peers=[_with(head, 9)],
+            **_base(tbo_on=tbo_on, max_seqlen_q=3),
+        )
+        assert r.max_seqlen_q_across_dp == 9, f"tbo_on={tbo_on}"
+        alone = _sync(monkeypatch, peers=[], **_base(tbo_on=tbo_on, max_seqlen_q=3))
+        assert alone.max_seqlen_q_across_dp == 3, f"tbo_on={tbo_on}"
diff --git a/tests/test_draft_graph.py b/tests/test_draft_graph.py
index 384d46c41e..a3082c19bb 100644
--- a/tests/test_draft_graph.py
+++ b/tests/test_draft_graph.py
@@ -8,7 +8,6 @@
 """
 
 import dataclasses
-import types
 
 import pytest
 import torch
@@ -19,7 +18,6 @@
 def _graph(**kw):
     """A pass over one int32 input, bound to a stub config."""
     kw.setdefault("inputs", {"row": StagedInput(dtype=torch.int32)})
-    kw.setdefault("pads", True)
     g = DraftGraph(
         forward=kw.pop("forward", lambda running_bs, **rows: running_bs), **kw
     )
@@ -27,76 +25,46 @@ def _graph(**kw):
     return g.bind(config, torch.device("cpu"))
 
 
-def _ctx(*, running_bs, use_cudagraph=True, dummy=False, running_tokens=None):
-    """A forward context carrying only what `target_running_bs` reads."""
-    return types.SimpleNamespace(
-        is_dummy_run=dummy,
-        running_bs=running_bs,
-        running_tokens=running_bs if running_tokens is None else running_tokens,
-        forward_mode=types.SimpleNamespace(
-            use_cudagraph=use_cudagraph, running_bs=running_bs
-        ),
-    )
-
-
-@pytest.mark.parametrize("bs,ran_at,want", [(44, 48, 48), (50, 64, 64), (64, 64, 64)])
-def test_the_pad_batch_is_the_one_the_target_ran(bs, ran_at, want):
-    """Never a batch the drafter picks. The target sized its per-sequence
-    metadata with `running_bs`, so that is exactly how far a pad row may
-    reach; anything wider reads a row nobody wrote this step."""
-    assert _graph().target_running_bs(bs, _ctx(running_bs=ran_at)) == want
-
-
-def test_an_empty_batch_is_never_widened():
-    """A pad row is a copy of the last real row, and there is none.
+def test_an_unwidened_empty_batch_stages_nothing():
+    """`stage` tail-repeats `src[-1]`, which a zero-row source does not have.
 
-    `stage` would reach `src[-1]` on a zero-row source and raise IndexError,
-    which is a much worse way to learn that a rank was kept alive purely to
-    reach the draft's collectives.
+    No caller can ask for one: `prepare_model` asserts the batch carries tokens,
+    so every drafter sees at least one row. Pinned as the boundary anyway --
+    entering at the real count must stay a no-op, not an IndexError.
     """
     g = _graph()
-    assert g.target_running_bs(0, _ctx(running_bs=48)) == 0
     assert g.stage(0, {"row": torch.zeros(0, dtype=torch.int32)})["row"].shape[0] == 0
 
 
-def test_the_pad_batch_counts_sequences_and_never_rows():
-    """The context carries the step's padded shape in both units, and only one
-    of them is a batch.
+def test_owning_a_recording_says_nothing_about_a_batch_it_is_not_at():
+    """The proposers widen to `running_bs` only when `is_captured` says that
+    batch replays, so this answer decides whether a pad row is ever fabricated.
 
-    `running_tokens` is what MoE pads hidden_states to; on a DSpark ragged step
-    it is a flat token bucket that no bs*q recovers, so reading it here would
-    stage hundreds of fabricated sequences. Pinned with the two far apart --
-    equal values would prove nothing about which one is read.
+    It has to be about the batch it was ASKED about. Answering out of inventory
+    -- "I own 48, so yes" -- would widen a 256-batch step to rows no recording
+    covers, and the pass would then hand the variable-length MoE gather padded
+    rows on a step that replays nothing.
     """
-    ctx = _ctx(running_bs=48, running_tokens=16336)
-    assert _graph().target_running_bs(44, ctx) == 48
-
-
-def test_nothing_is_padded_where_the_target_pinned_no_batch():
     g = _graph()
-    # Planted first: owning a graph at 48 is NOT a licence to pad to it. Only a
-    # replayed target sizes its per-sequence metadata past the real batch, and
-    # a pad row that reaches further takes its ring slot -- which DSpark's block
-    # SCATTERS draft KV through -- from a row nobody wrote this step.
     g._cuda_graphs[48] = ("graph", "out")
-    ctx = _ctx(running_bs=256, use_cudagraph=False)
-    assert g.target_running_bs(44, ctx) == 44, "eager: nobody sized metadata past bs"
+    assert g.is_captured(48)
+    assert not g.is_captured(256)
+    assert not g.is_captured(44)
 
 
-def test_a_pass_that_cannot_pad_stays_at_the_real_batch():
-    """`pads` gates the padding, not just the capture -- EPLB turns it off
-    because pad rows would route through the draft's MoE and land in the
-    expert-load histogram."""
-    g = _graph(pads=False)
-    assert g.target_running_bs(44, _ctx(running_bs=48)) == 44
-    assert not g.will_capture
+def test_nothing_recorded_means_no_batch_replays():
+    """`ATOM_DRAFT_CUDAGRAPH=0` leaves the inventory empty, and then every
+    batch must answer no -- that is what keeps the eager pass unpadded."""
+    g = _graph()
+    assert not any(g.is_captured(n) for n in (0, 1, 44, 48, 256))
 
 
-def test_declaring_pads_without_inputs_is_refused():
-    """`pads` promises the fabricated rows land somewhere inert; with nothing
-    staged there is no buffer for them to land in at all."""
+def test_a_pass_that_stages_nothing_is_refused():
+    """A fabricated row has to land somewhere. Every declared pass may pad now,
+    so staging nothing is not an unpaddable pass -- it is no pass at all, and
+    the flavor must decline the declaration instead."""
     with pytest.raises(AssertionError, match="stages nothing"):
-        _graph(inputs={}, pads=True)
+        _graph(inputs={})
 
 
 def test_staged_tail_repeats_a_real_row_rather_than_zero_filling():
@@ -137,22 +105,6 @@ def test_staging_a_source_whose_leading_axis_is_not_the_batch_fails_loudly():
         g._stage_one("row", 8, torch.zeros(3, 8, dtype=torch.int32))
 
 
-def test_padding_a_pass_that_never_called_its_pad_rows_inert_is_refused():
-    """`pads` is the only part of the contract nothing else can check.
-
-    A pass that lies about it still runs and still returns the right shape; the
-    fabricated rows just land in another sequence's KV.
-    """
-    src = {"row": torch.zeros(3, dtype=torch.int32)}
-    assert _graph(pads=True).stage(4, src)["row"].shape[0] == 4
-
-    g = _graph(pads=False)
-    with pytest.raises(AssertionError, match="fabricated rows are inert"):
-        g.stage(4, src)
-    # Unpadded entry stays legal -- the contract is about fabricated rows only.
-    assert g.stage(3, src)["row"].shape[0] == 3
-
-
 def test_sources_that_disagree_on_the_batch_are_refused():
     """They describe one step, so one of them is not this step's."""
     g = _graph(
@@ -188,17 +140,17 @@ def test_warmup_runs_the_epilogue_too_not_just_the_capturable_forward(monkeypatc
     assert ran == ["forward", "epilogue"]
 
 
-def test_warmup_happens_before_the_pad_and_capture_gates(monkeypatch):
-    """A pass that pads nothing and captures nothing is still WARMED.
+def test_warmup_happens_before_the_capture_gate(monkeypatch):
+    """A pass that captures nothing is still WARMED.
 
     That ordering is deliberate -- the JIT is worth paying at startup either
-    way -- but it means declining to pad does not keep a flavor's forward off
-    the startup sweep. A flavor whose model cannot answer it must decline the
-    whole pass.
+    way -- but it means turning capture off does not keep a flavor's forward
+    off the startup sweep. A flavor whose model cannot answer it must decline
+    the whole pass.
     """
     monkeypatch.setenv("ATOM_DRAFT_CUDAGRAPH", "0")
     ran = []
-    g = _graph(pads=False, forward=lambda running_bs, **rows: ran.append("forward"))
+    g = _graph(forward=lambda running_bs, **rows: ran.append("forward"))
     assert not g.will_capture
     g.warmup(8)
     assert ran == ["forward"]
@@ -226,17 +178,70 @@ def replay(self):
     g = _graph(forward=lambda running_bs, **rows: forwards.append(running_bs))
     src = {"row": torch.zeros(4, dtype=torch.int32)}
 
-    real = _ctx(running_bs=4)
     assert not g.is_captured(4)
-    g.run(4, real, **g.stage(4, src))
+    g.run(4, **g.stage(4, src))
     assert (forwards, replays) == ([4], []), "no recording: must run the forward"
 
     g._cuda_graphs[4] = (_Graph(), "recorded")
     assert g.is_captured(4)
-    assert g.run(4, real, **g.stage(4, src)) == "recorded"
+    assert g.run(4, **g.stage(4, src)) == "recorded"
     assert (forwards, replays) == ([4], [1]), "recorded: must replay, not re-run"
 
 
+def test_a_replay_hands_back_a_value_not_the_recordings_own_storage():
+    """What `run` returns has to survive the next replay of any other size.
+
+    The recording's tensors are allocated inside the capture, so they live in
+    the pool every captured size shares; the pool packs a later size's capture
+    into memory the earlier ones released, and a replay rewrites every address
+    it recorded whatever holds it now. Holding a Python reference does not
+    stop that -- which is exactly why this cannot be left to the caller.
+
+    Written as "the recording changing underneath must not change what was
+    handed out", because that is the failure: DSpark's draft ids are read a
+    step later as the next forward's `input_ids`, and another size's replay
+    had turned them into that size's activations. The target's embedding
+    gathered on those floats and faulted all eight ranks.
+    """
+
+    class _Graph:
+        def replay(self):
+            pass
+
+    ids = torch.full((4, 3), 9, dtype=torch.int32)
+    conf = torch.zeros(4, 3)
+    g = _graph(epilogue=lambda out, running_bs, **rows: out, capture_epilogue=True)
+    g._cuda_graphs[4] = (_Graph(), (ids, conf))
+
+    got_ids, got_conf = g.run(
+        4, **g.stage(4, {"row": torch.zeros(4, dtype=torch.int32)})
+    )
+    assert got_ids.tolist() == ids.tolist(), "the replay's values, unchanged"
+
+    # A later replay at another size lands on the recording's storage.
+    ids.fill_(1234)
+    conf.fill_(5.0)
+    assert got_ids.tolist() != ids.tolist(), "handed out a window into the recording"
+    assert got_conf.tolist() != conf.tolist(), "every output, not just the ids"
+
+
+def test_the_eager_path_hands_back_what_it_allocated():
+    """...and only a replay needs the copy.
+
+    Without a recording the pass allocates normally, where holding the
+    reference is what keeps the storage. Copying there would be pure cost, so
+    the identity is asserted rather than left to inspection.
+    """
+    made = torch.zeros(2, 3)
+    g = _graph(
+        forward=lambda running_bs, **rows: made,
+        epilogue=lambda out, running_bs, **rows: out,
+        capture_epilogue=True,
+    )
+    assert not g.is_captured(2)
+    assert g.run(2, **g.stage(2, {"row": torch.zeros(2, dtype=torch.int32)})) is made
+
+
 def test_a_dummy_replays_in_lockstep_with_the_real_ranks():
     """`is_dummy_run` is per-rank, so no decision that feeds a collective may
     read it.
@@ -263,30 +268,32 @@ def replay(self):
     g._cuda_graphs[1] = (_Graph(), "recorded")
     src = {"row": torch.zeros(1, dtype=torch.int32)}
 
-    for dummy in (False, True):
-        ctx = _ctx(running_bs=1, dummy=dummy)
-        assert g.run(1, ctx, **g.stage(1, src)) == "recorded"
-        assert " graph" in g.label(1, 1, ctx)
+    for _dummy in (False, True):
+        assert g.run(1, **g.stage(1, src)) == "recorded"
+        assert " graph" in g.label(1, 1)
     assert (forwards, replays) == ([], [1, 1]), "both ranks replay, neither runs"
 
 
-@pytest.mark.parametrize("use_cudagraph", [True, False])
-def test_the_replay_decision_never_differs_between_a_dummy_and_a_real_step(
-    use_cudagraph,
-):
-    """The property the case above pins, stated over the whole decision.
+def test_the_replay_decision_is_a_function_of_the_batch_and_nothing_else():
+    """Every term the decision reads has to be one the whole DP group agrees on.
 
-    Every term `_replays` reads has to be one the whole DP group agrees on;
-    `use_cudagraph` and the batch are (``ForwardMode.decide`` derives both from
-    DP-unified counts), and nothing else may enter. Asserting the equality
-    rather than the two values is what makes a future third term fail here.
+    `running_bs` is, having been reduced in `sync_dp_metadata`. The step's own
+    kind is not: `is_dummy_run` is per-rank, and so is `is_prefill` (one rank
+    prefills while another decodes). Either would split the group into two
+    collective widths -- the first one measured, V4-Flash-DSpark tp8 + DPA
+    hanging 8/8 on the first real decode.
+
+    Pinned as an interface property rather than a value: the decision takes the
+    batch and nothing else, so a future term cannot be added without changing
+    this signature.
     """
-    g = _graph()
-    g._cuda_graphs[8] = ("graph", "out")
-    real = _ctx(running_bs=8, use_cudagraph=use_cudagraph)
-    dummy = _ctx(running_bs=8, use_cudagraph=use_cudagraph, dummy=True)
-    assert g.target_running_bs(8, real) == g.target_running_bs(8, dummy)
-    assert g.label(8, 8, real) == g.label(8, 8, dummy)
+    import inspect
+
+    params = list(inspect.signature(DraftGraph.is_captured).parameters)
+    assert params == ["self", "running_bs"], (
+        f"is_captured grew a term beyond the agreed batch: {params}. Anything "
+        "keyed on the step's kind is per-rank and hangs the DP group."
+    )
 
 
 def test_the_capture_gate_reaches_the_env_that_names_it(monkeypatch):
diff --git a/tests/test_dspark.py b/tests/test_dspark.py
index 9fcfc2f982..f4e699029e 100644
--- a/tests/test_dspark.py
+++ b/tests/test_dspark.py
@@ -9,10 +9,19 @@
 
 import pytest
 
-pytest.importorskip("aiter", reason="the compiled draft imports aiter at module load")
+# Named on the module, not on "aiter": a test collected after one that stubs
+# `sys.modules["aiter"]` (tests/test_pd_pp.py does) sees the name import fine
+# and dies on a symbol, which is a collection error and takes the suite with
+# it. `exc_type` is what turns that ImportError into a skip.
+pytest.importorskip(
+    "atom.spec_decode.drafter",
+    reason="the compiled draft imports aiter at module load",
+    exc_type=ImportError,
+)
 
 import torch
 
+import atom.spec_decode.drafter as mod_drafter
 from atom.model_ops.v4_kernels.dspark_fp8_indices import DSparkIndexBuffers
 from atom.models.deepseek_v4_dspark import (
     DSparkConfidenceHead,
@@ -970,7 +979,7 @@ def head_and_sample(normed, hc_hidden, anchor_ids):
 _GRAPH_BS = [1, 2, 4, 8, 16, 32, 48, 64, 128, 256]
 
 
-def _stub_forward_context(*, scheduled_bs, target_bs, use_cudagraph=True):
+def _stub_forward_context(*, scheduled_bs, target_bs):
     context = types.SimpleNamespace(
         scheduled_bs=scheduled_bs,
         running_bs=target_bs,
@@ -980,10 +989,8 @@ def _stub_forward_context(*, scheduled_bs, target_bs, use_cudagraph=True):
         running_tokens=target_bs * 337,
         is_dummy_run=False,
         is_draft=False,
+        is_prefill=False,
         positions=None,
-        forward_mode=types.SimpleNamespace(
-            use_cudagraph=use_cudagraph, running_bs=target_bs
-        ),
     )
     # `prepare_decode` publishes the ring slots at the PADDED batch, so the stub
     # does too -- the block slices to that length and nothing stages it.
@@ -994,19 +1001,30 @@ def _stub_forward_context(*, scheduled_bs, target_bs, use_cudagraph=True):
     return types.SimpleNamespace(context=context, attn_metadata=attn_metadata)
 
 
-def _proposer_with_graph_bs(monkeypatch, *, eplb=False, mtp_k=5, window=128):
-    """A DSparkProposer carrying only what the block pass reads."""
+def _proposer_with_graph_bs(
+    monkeypatch, *, eplb=False, mtp_k=5, window=128, captured=True
+):
+    """A DSparkProposer carrying only what the block pass reads.
+
+    ``captured`` is the fixture keeping its own name's promise: the startup
+    sweep records every capture size, and `propose` widens only for a batch
+    that has a recording. Answered through `is_captured` rather than by filling
+    `_cuda_graphs`, so `run` still reaches the backbone -- the only place these
+    tests can watch the width. `captured=False` is the no-recording case, where
+    widening would hand the variable-length MoE gather rows nothing replays.
+    """
+    from atom.spec_decode.draft_graph import DraftGraph
     from atom.spec_decode.dspark_proposer import DSparkProposer
 
-    monkeypatch.setattr(DSparkProposer, "_with_draft", False, raising=False)
     monkeypatch.setattr(
-        DSparkProposer, "aux_for", lambda self, h: [torch.zeros(1)], raising=False
+        DraftGraph,
+        "is_captured",
+        lambda self, running_bs: captured and running_bs in _GRAPH_BS,
     )
+
+    monkeypatch.setattr(DSparkProposer, "_with_draft", False, raising=False)
     monkeypatch.setattr(
-        DSparkProposer,
-        "_refresh_dp_metadata",
-        lambda self, fc, n: None,
-        raising=False,
+        DSparkProposer, "aux_for", lambda self, h: [torch.zeros(1)], raising=False
     )
     monkeypatch.setattr(DSparkProposer, "verify_scheduler", None, raising=False)
 
@@ -1031,7 +1049,7 @@ def _proposer_with_graph_bs(monkeypatch, *, eplb=False, mtp_k=5, window=128):
     return p
 
 
-def _run_propose(p, fc, real_bs, monkeypatch, seen):
+def _run_propose(p, fc, scheduled_bs, monkeypatch, seen):
     import atom.spec_decode.dspark_proposer as mod
 
     def _backbone(ids, pos, num_draft):
@@ -1067,18 +1085,18 @@ def head_and_sample(normed, hc_hidden, anchor_ids):
     monkeypatch.setattr(mod, "get_forward_context", lambda: fc)
     return p.propose(
         target_token_ids=None,
-        target_positions=torch.arange(real_bs, dtype=torch.int64) * 7 + 3,
-        target_hidden_states=torch.zeros(real_bs, 4),
+        target_positions=torch.arange(scheduled_bs, dtype=torch.int64) * 7 + 3,
+        target_hidden_states=torch.zeros(scheduled_bs, 4),
         num_reject_tokens=None,
-        next_token_ids=torch.full((real_bs,), 5, dtype=torch.int32),
-        last_token_indices=torch.arange(real_bs, dtype=torch.int64),
+        next_token_ids=torch.full((scheduled_bs,), 5, dtype=torch.int32),
+        last_token_indices=torch.arange(scheduled_bs, dtype=torch.int64),
     )
 
 
 @pytest.mark.parametrize(
-    "real_bs,expect_B", [(44, 48), (50, 64), (1, 1), (64, 64), (35, 48)]
+    "scheduled_bs,expect_B", [(44, 48), (50, 64), (1, 1), (64, 64), (35, 48)]
 )
-def test_propose_drafts_at_the_captured_graph_bs(monkeypatch, real_bs, expect_B):
+def test_propose_drafts_at_the_captured_graph_bs(monkeypatch, scheduled_bs, expect_B):
     """The block runs at the target's running_bs, not the live batch size.
 
     Without this the drafter hands aiter a fresh width on every distinct decode
@@ -1087,42 +1105,257 @@ def test_propose_drafts_at_the_captured_graph_bs(monkeypatch, real_bs, expect_B)
     """
     seen = {}
     p = _proposer_with_graph_bs(monkeypatch)
-    fc = _stub_forward_context(scheduled_bs=real_bs, target_bs=expect_B)
-    out = _run_propose(p, fc, real_bs, monkeypatch, seen)
+    fc = _stub_forward_context(scheduled_bs=scheduled_bs, target_bs=expect_B)
+    out = _run_propose(p, fc, scheduled_bs, monkeypatch, seen)
 
     assert seen["B"] == expect_B
     assert seen["positions_B"] == expect_B
     # The block does not touch the target's ring slots: they arrive already at
     # the padded length, so there is nothing to install and nothing to restore.
     assert seen["slots"].shape[0] >= expect_B
-    assert seen["slots"][:real_bs].tolist() == list(range(100, 100 + real_bs))
-    assert out.shape[0] == real_bs
+    assert seen["slots"][:scheduled_bs].tolist() == list(range(100, 100 + scheduled_bs))
+    assert out.shape[0] == scheduled_bs
 
     # ...but the rows it fabricated must scatter no draft KV. Their ring slot is
     # the 0 `prepare_decode` fills that tail with, and 0 is a real position, so
     # an unmasked pad row writes into another request's window.
     t = p.mtp_k
     ids = seen["batch_ids"]
-    assert ids[: real_bs * t].tolist() == [i // t for i in range(real_bs * t)]
-    assert ids[real_bs * t : expect_B * t].tolist() == [-1] * (expect_B - real_bs) * t
+    assert ids[: scheduled_bs * t].tolist() == [i // t for i in range(scheduled_bs * t)]
+    assert (
+        ids[scheduled_bs * t : expect_B * t].tolist()
+        == [-1] * (expect_B - scheduled_bs) * t
+    )
 
 
-@pytest.mark.parametrize(
-    "cudagraph,eplb,why",
-    [
-        (False, False, "eager: the target pinned no wider batch to follow"),
-        (True, True, "pad rows would poison the expert-load histogram"),
-    ],
-)
-def test_propose_leaves_the_batch_alone_when_nothing_pins_a_wider_one(
-    monkeypatch, cudagraph, eplb, why
-):
+def test_propose_states_both_units_of_its_own_pass(monkeypatch):
+    """The draft's shape, and it is the pass's own, not the target's.
+
+    `scheduled_tokens` is how many rows carry a real request, `running_tokens`
+    how many the pass runs, and they DIFFER exactly when the batch was widened.
+    Pinned 44 against 48 so a version that reports one number twice cannot pass.
+    """
+    seen = {}
+    p = _proposer_with_graph_bs(monkeypatch, captured=False)
+    fc = _stub_forward_context(scheduled_bs=44, target_bs=48)
+    out = _run_propose(p, fc, 44, monkeypatch, seen)
+
+    t = p.mtp_k
+    assert seen["B"] == 48, "the pass runs the agreed batch, recording or not"
+    assert fc.context.scheduled_tokens == 44 * t
+    assert fc.context.running_tokens == 48 * t
+    assert out.shape[0] == 44, "sliced back to the real requests on the way out"
+    # ...and the fabricated rows still scatter no draft KV.
+    ids = seen["batch_ids"]
+    assert ids[44 * t : 48 * t].tolist() == [-1] * (4 * t)
+
+
+def test_propose_sizes_the_draft_gather_without_asking_the_group(monkeypatch):
+    """No collective here. The step already ran its one.
+
+    Every rank's draft height is the same number by construction --
+    `running_bs` came out of that reduction and the draft width is config -- so
+    each can write the whole table itself. `DPMetadata.make` all_reduces it
+    when the table is not supplied, which cost one CPU collective per draft
+    pass: `mtp_k` of them per step for the serial drafters.
+
+    Watched by making the ask itself fail, not by counting calls: a version
+    that asks would otherwise pass here on a stub that answers.
+    """
+    seen = {}
+    p = _proposer_with_graph_bs(monkeypatch)
+    p.config.parallel_config = types.SimpleNamespace(
+        data_parallel_size=2, data_parallel_rank=0
+    )
+
+    def _must_not_ask(*a, **k):
+        raise AssertionError("draft asked the group for a height it already knows")
+
+    monkeypatch.setattr(
+        mod_drafter.DPMetadata, "num_tokens_across_dp", staticmethod(_must_not_ask)
+    )
+    fc = _stub_forward_context(scheduled_bs=44, target_bs=48)
+    fc.dp_metadata = None
+    _run_propose(p, fc, 44, monkeypatch, seen)
+
+    t = p.mtp_k
+    assert fc.dp_metadata.get_sizes_across_dp() == [48 * t, 48 * t]
+
+
+def test_propose_puts_the_pass_on_the_path_its_own_shape_is_on(monkeypatch):
+    """The block pass is uniform across DP and is not a prefill.
+
+    Its height is `running_bs * draft width`, and both factors are the step's
+    own reduction or config -- so every rank runs the same number of rows,
+    whatever `decide` said about the TARGET's counts. That is what lets the
+    table below be declared rather than asked for.
+
+    `is_prefill` is a SEPARATE write by this same caller, not a consequence of
+    publishing a shape -- the padded gather reads it to choose which of the two
+    heights to check the rows against. Asserted here rather than one seam down
+    because the drafter reusing the target's prefill layout needs the opposite.
+    """
+    seen = {}
+    p = _proposer_with_graph_bs(monkeypatch)
+    p.config.parallel_config = types.SimpleNamespace(
+        data_parallel_size=2, data_parallel_rank=0
+    )
+    monkeypatch.setattr(
+        mod_drafter.DPMetadata, "make", staticmethod(lambda *a, **k: "dp_meta")
+    )
+    fc = _stub_forward_context(scheduled_bs=44, target_bs=48)
+    fc.context.is_prefill = True
+    fc.dp_metadata = None
+    _run_propose(p, fc, 44, monkeypatch, seen)
+
+    assert fc.context.running_tokens_are_unified is True
+    assert fc.context.is_prefill is False
+    assert fc.dp_metadata == "dp_meta"
+
+
+class _BareFlavor(mod_drafter.Drafter):
+    """The abstract methods, stubbed. `_publish_draft_shape` reads neither."""
+
+    _resolve_mtp_k = propose = staticmethod(lambda *a, **k: None)
+
+
+def _bare_drafter(dp_size):
+    """A Drafter carrying only what `_publish_draft_shape` reads.
+
+    Not a DSparkProposer: the tests below are about the seam every flavor goes
+    through, and the one that gets it wrong is the one with no fixture.
+    """
+    drafter = object.__new__(_BareFlavor)
+    drafter.config = types.SimpleNamespace(
+        parallel_config=types.SimpleNamespace(
+            data_parallel_size=dp_size, data_parallel_rank=0
+        )
+    )
+    return drafter
+
+
+def test_a_height_the_group_did_not_agree_on_is_asked_for_not_declared(monkeypatch):
+    """A draft pass sized off this rank alone still costs the collective.
+
+    `running_bs * width` is the same number everywhere; a pass carrying this
+    rank's own token stream is not, and declaring it uniform states an equality
+    no reduction ever checked.
+
+    What travels is the answer, not the table it implies -- how `make` then
+    builds one is tested against the real thing in tests/test_dp_metadata.py.
+    The two counts are pinned apart so forwarding the wrong one fails here.
+    """
     seen = {}
-    p = _proposer_with_graph_bs(monkeypatch, eplb=eplb)
-    fc = _stub_forward_context(scheduled_bs=44, target_bs=48, use_cudagraph=cudagraph)
+
+    def _make(cfg, batchsize, num_tokens_across_dp=None, *, unified=False):
+        seen["batchsize"] = batchsize
+        seen["table"] = num_tokens_across_dp
+        seen["unified"] = unified
+        return "dp_meta"
+
+    monkeypatch.setattr(mod_drafter.DPMetadata, "make", staticmethod(_make))
+    fc = _stub_forward_context(scheduled_bs=44, target_bs=48)
+
+    _bare_drafter(2)._publish_draft_shape(
+        fc,
+        scheduled_tokens=308,
+        running_tokens=336,
+        running_tokens_are_unified=False,
+    )
+    assert fc.context.running_tokens_are_unified is False
+    assert seen["unified"] is False, "a ragged height has to be discovered"
+
+    _bare_drafter(2)._publish_draft_shape(
+        fc,
+        scheduled_tokens=308,
+        running_tokens=336,
+        running_tokens_are_unified=True,
+    )
+    assert fc.context.running_tokens_are_unified is True
+    assert seen["batchsize"] == 336, "the group is told what the pass RUNS"
+    assert seen["unified"] is True, "an agreed height is worth no ask"
+    assert seen["table"] is None, "the answer travels, not a table built here"
+
+
+def test_publishing_a_shape_does_not_decide_whether_it_is_a_prefill():
+    """Two different questions, and one drafter answers them differently.
+
+    A draft step that reuses the target's own prefill layout -- Eagle's first
+    -- runs the target's ragged token stream through metadata that still
+    describes a prefill, so clearing the flag here would send every MLA layer,
+    the indexer and PCP down the decode branch on a pass that is not one.
+    """
+    fc = _stub_forward_context(scheduled_bs=44, target_bs=48)
+    fc.context.is_prefill = True
+
+    # dp_size 1: no metadata to build, so nothing but the flag is under test.
+    _bare_drafter(1)._publish_draft_shape(
+        fc,
+        scheduled_tokens=337,
+        running_tokens=337,
+        running_tokens_are_unified=False,
+    )
+
+    assert fc.context.is_prefill is True
+
+
+def test_propose_pads_under_eplb_exactly_as_it_does_without_it(monkeypatch):
+    """EPLB is not a reason to decline the padding, though it used to be.
+
+    The stated reason was the expert-load histogram: a pad row routes through
+    the draft's MoE and `select_experts_with_record` counts it. But the TARGET
+    pads on every cudagraph decode step and its rows reach the same recorder
+    (`eplb_map_and_record_fused` masks nothing), in far greater volume -- so
+    declining here protected nothing and only cost the draft its warmed shape.
+
+    Asserting the two are EQUAL, not that either is 48: a future term that
+    reads the config back into the width fails here whichever way it leans.
+    """
+    seen_on, seen_off = {}, {}
+    _run_propose(
+        _proposer_with_graph_bs(monkeypatch, eplb=True),
+        _stub_forward_context(scheduled_bs=44, target_bs=48),
+        44,
+        monkeypatch,
+        seen_on,
+    )
+    _run_propose(
+        _proposer_with_graph_bs(monkeypatch, eplb=False),
+        _stub_forward_context(scheduled_bs=44, target_bs=48),
+        44,
+        monkeypatch,
+        seen_off,
+    )
+    assert seen_on["B"] == seen_off["B"] == 48
+
+
+def test_propose_replays_at_the_agreed_batch_even_when_it_had_to_pad(monkeypatch):
+    """The seam, end to end through `propose`.
+
+    A step whose real batch is 44 runs the block at the agreed 48 and stands in
+    for the recording there. Watched through the recording being returned and
+    the backbone never running -- asserting the label alone would pass on a
+    `run` that replays nothing.
+
+    """
+    seen = {}
+    p = _proposer_with_graph_bs(monkeypatch)
+    recorded = (
+        torch.full((48, p.mtp_k), 9, dtype=torch.int32),
+        torch.zeros(48, p.mtp_k),
+    )
+    replays = []
+    p.block._cuda_graphs[48] = (
+        types.SimpleNamespace(replay=lambda: replays.append(1)),
+        recorded,
+    )
+    fc = _stub_forward_context(scheduled_bs=44, target_bs=48)
     out = _run_propose(p, fc, 44, monkeypatch, seen)
-    assert seen["B"] == 44, why
-    assert out.shape[0] == 44
+    assert replays == [1], "must replay the recording at the agreed batch"
+    assert "B" not in seen, "the backbone must not run when the recording stands in"
+    # Sliced back to the real batch: the 4 fabricated rows are the recording's
+    # to compute and nobody's to return.
+    assert out.tolist() == recorded[0][:44, : p.mtp_k].tolist()
 
 
 def test_propose_pads_a_dp_sync_dummy_exactly_like_the_rank_with_work(monkeypatch):
@@ -1348,12 +1581,12 @@ def test_the_separate_draft_model_path_declares_no_draft_graph(monkeypatch):
     the warmup and epilogue reach for -- and `warmup` runs both BEFORE it
     consults the pad/capture gates. So an unpaddable-but-declared pass still
     takes the startup sweep through `_block_warmup_inputs` and dies with an
-    AttributeError, which is what leaving this at `pads = False` used to do.
+    AttributeError, so declining to pad was never enough to keep it out.
     """
     from atom.spec_decode.dspark_proposer import DSparkProposer
 
     p = _proposer_with_graph_bs(monkeypatch)
-    assert p.draft_graphs and p.block.pads
+    assert p.draft_graphs and p.block is not None
 
     monkeypatch.setattr(DSparkProposer, "_with_draft", True, raising=False)
     p._build_draft_graphs()
diff --git a/tests/test_dspark_scheduler.py b/tests/test_dspark_scheduler.py
index f3ad8ce85f..b4cb476e6c 100644
--- a/tests/test_dspark_scheduler.py
+++ b/tests/test_dspark_scheduler.py
@@ -14,7 +14,6 @@
     build_sps_table,
     calibrate_confidence,
     expected_throughput,
-    flat_bucket_fits,
     ragged_verify_len,
     schedule_prefix_lengths,
     survival_probabilities,
@@ -773,52 +772,3 @@ def test_both_bounds_hold_across_the_grid():
                     continue
                 assert max_nb + 1 <= li <= FULL_Q, (ell, max_nb, sched, li)
                 assert li <= sched, (ell, max_nb, sched, li)
-
-
-# ---------------------------------------------------------------------------
-# The ragged shrink is only safe if the REPLAY can follow it down.
-
-
-def test_full_cudagraphs_cannot_represent_a_ragged_shrink():
-    """No captured flat bucket set -> no representable shrink.
-
-    This is the configuration that trapped: under FULL (non-PIECEWISE)
-    cudagraphs nothing flat is captured, `_dynamic_num_tokens_pad` returns None,
-    every caller falls back to `bs * max_seqlen_q`, and the replay runs over
-    more tokens than the rebuild populated. The tail holds the previous step's
-    ids, which the draft's Markov lookup then indexes out of range.
-    """
-    assert flat_bucket_fits(3, 6, []) is False
-    assert flat_bucket_fits(3, 6, None) is False
-
-
-def test_bucket_must_cover_the_total_and_divide_by_q():
-    buckets = [6, 12, 24]
-    # 3 real tokens at q=6 -> bucket 6 holds them and 6 % 6 == 0.
-    assert flat_bucket_fits(3, 6, buckets) is True
-    assert flat_bucket_fits(6, 6, buckets) is True
-    # Bigger than every captured bucket -> nothing can hold it.
-    assert flat_bucket_fits(25, 6, buckets) is False
-    # Covers the total but is not q-divisible -> the per-seq rows would not tile.
-    assert flat_bucket_fits(3, 6, [10]) is False
-    # Degenerate q never matches (guards the b % q modulo).
-    assert flat_bucket_fits(3, 0, buckets) is False
-
-
-def test_predicate_agrees_with_dynamic_num_tokens_pad():
-    """Must stay in step with the lookup it mirrors: yes exactly when that
-    lookup finds a bucket, no exactly when it returns the None that triggers the
-    unsafe `bs * max_seqlen_q` fallback."""
-    buckets = [6, 12, 24]
-
-    def pad_lookup(total, q):
-        for b in buckets:
-            if b >= total and q > 0 and b % q == 0:
-                return b
-        return None
-
-    for total in (0, 1, 3, 6, 7, 12, 25):
-        for q in (0, 1, 3, 6):
-            assert flat_bucket_fits(total, q, buckets) is (
-                pad_lookup(total, q) is not None
-            ), (total, q)
diff --git a/tests/test_forward_mode.py b/tests/test_forward_mode.py
new file mode 100644
index 0000000000..43fa9d82ac
--- /dev/null
+++ b/tests/test_forward_mode.py
@@ -0,0 +1,433 @@
+# SPDX-License-Identifier: MIT
+"""`ForwardMode.decide`: the one place a step's shape is settled.
+
+Single-rank steps take the early return and issue no collective, so most of
+this runs with no DP group at all. The cases that need one fake the packed
+all_gather -- see `tests/test_dp_sync_layout.py` for the wire format itself.
+
+A shape that came out per-rank is how a DP group ends up with two collective
+widths, and that failure shows up as a hang rather than as a wrong number.
+"""
+
+import types
+
+import numpy as np
+import pytest
+import torch
+
+from atom.utils.forward_context import ForwardMode
+
+# What `ModelRunner.capture_sizes_np` is: ascending int32, sorted where it is
+# built. `decide` searches it, so a test that passed a plain list would be
+# exercising a type production never hands it.
+LADDER = np.asarray([1, 2, 4, 8, 16, 32, 48, 64], dtype=np.int32)
+
+
+def _batch(*, seqs, tokens=None, prefill_tokens=0, q=1):
+    """The fields `decide` reads. Decode unless `prefill_tokens` is given.
+
+    A decode-only batch has `total_tokens_num == total_tokens_num_decode`, as a
+    real one does -- the two diverging is what a mixed batch means, and no test
+    here builds one.
+    """
+    total = seqs * q if tokens is None else tokens
+    return types.SimpleNamespace(
+        total_seqs_num=seqs,
+        total_tokens_num=total,
+        total_tokens_num_prefill=prefill_tokens,
+        total_tokens_num_decode=0 if prefill_tokens else total,
+        # Not a real ScheduledBatch field any more -- kept so `_decide` can pass
+        # this stub's per-seq length the way `prepare_model` passes the shrink's.
+        max_seqlen_q=q,
+        is_dummy_run=False,
+    )
+
+
+def _decide(batch=None, **kw):
+    base = {
+        "batch": batch if batch is not None else _batch(seqs=8, q=1),
+        "dp_size": 1,
+        "dp_group": None,
+        "enforce_eager": False,
+        "capture_sizes": LADDER,
+        "captured_tokens": None,
+        "is_block_drafter": False,
+        "tbo_on": False,
+        "local_tbo": (False, False, 0, 0),
+    }
+    merged = {**base, **kw}
+    # The batch no longer carries it -- `decide` is handed the step's per-seq
+    # length, so the stub's `q` travels the same way production's does.
+    merged.setdefault("max_seqlen_q", merged["batch"].max_seqlen_q)
+    return ForwardMode.decide(**merged)
+
+
+def _fake_sync(monkeypatch, *, peer_tokens, peer_bs, peer_prefill=False):
+    """One peer rank differing from ours, without a real DP group."""
+    import atom.utils.tbo.ubatching as ub
+
+    def fake_all_gather(out_list, local, group=None):
+        out_list[0].copy_(local)
+        peer = local.clone()
+        peer[0], peer[1] = peer_tokens, peer_bs
+        peer[2] = 1 if peer_prefill else 0
+        out_list[1].copy_(peer)
+
+    monkeypatch.setattr(torch.distributed, "all_gather", fake_all_gather)
+    return ub
+
+
+# --------------------------------------------------------------------------- #
+# running_bs -- what everything per-sequence runs at.                          #
+# --------------------------------------------------------------------------- #
+
+
+def test_the_batch_is_the_reduction_on_the_ladder_on_every_kind_of_step():
+    """One rule, and the kind of step is not part of it.
+
+    Prefill, eager decode and a replayed decode all leave with the same batch:
+    the group's, rounded on a ladder every rank shares. That is what lets a
+    graph key, a draft pass and an attention plan read one field without
+    knowing which kind of step produced it -- when this branched on dispatch,
+    two of the four arms handed back a per-rank count and every consumer had
+    to know which it was holding.
+    """
+    prefill = _decide(_batch(seqs=3, tokens=900, prefill_tokens=900))
+    assert prefill.is_prefill and not prefill.use_cudagraph
+    assert prefill.running_bs == 4, "rounded up, exactly like a decode step"
+
+    decode = _decide(_batch(seqs=3, q=6))
+    assert decode.use_cudagraph and decode.running_bs == 4
+    assert prefill.running_bs == decode.running_bs
+
+
+def test_a_batch_the_ladder_cannot_hold_keeps_the_reduction_and_replays_nothing():
+    """Running off the end of the ladder is the same fact twice: there is no
+    recording that fits, so the batch stays what the reduction said and the
+    step cannot replay. One lookup has to answer both, or a step dispatches
+    PIECEWISE at a width nothing recorded."""
+    mode = _decide(_batch(seqs=100, q=1))
+    assert mode.running_bs == 100 and not mode.use_cudagraph
+
+
+def test_a_decode_step_that_fits_the_ladder_replays():
+    mode = _decide(_batch(seqs=50, q=6))
+    assert mode.use_cudagraph and mode.running_bs == 64
+
+
+def test_forcing_eager_changes_the_dispatch_and_not_the_batch():
+    """`enforce_eager` answers "does the target replay", which is a different
+    question from "how wide is the batch". Under it the runner never runs
+    `capture_cudagraph`, so the ladder stays at its `[0]` placeholder and the
+    reduction stands unrounded on its own -- by having nothing to round to,
+    not by a branch."""
+    warmed = _decide(_batch(seqs=5, q=1), enforce_eager=True)
+    assert not warmed.use_cudagraph and warmed.running_bs == 8
+
+    unwarmed = _decide(
+        _batch(seqs=5, q=1),
+        enforce_eager=True,
+        capture_sizes=np.asarray([0], dtype=np.int32),
+    )
+    assert not unwarmed.use_cudagraph and unwarmed.running_bs == 5
+
+
+def test_the_batch_ignores_what_this_rank_alone_was_handed(monkeypatch):
+    """The DP property, stated where it is decided.
+
+    Two ranks arrive at one step with different local batches. They must leave
+    with the same `running_bs`, because that is the width a recording holds its
+    collective at. Asserting the equality rather than the value is what makes a
+    future term that reads the local count fail here.
+    """
+    _fake_sync(monkeypatch, peer_tokens=40, peer_bs=40)
+    small = _decide(_batch(seqs=1, q=1), dp_size=2, dp_group=object())
+    _fake_sync(monkeypatch, peer_tokens=1, peer_bs=1)
+    large = _decide(_batch(seqs=40, q=1), dp_size=2, dp_group=object())
+    assert small.running_bs == large.running_bs == 48
+    assert small.use_cudagraph == large.use_cudagraph
+
+
+# --------------------------------------------------------------------------- #
+# The rest of what one step settles.                                           #
+# --------------------------------------------------------------------------- #
+
+
+def test_tbo_survives_having_no_peer_to_ask():
+    """One rank still gets an answer, and it is the local one.
+
+    The gate used to be read off the DP sync result, which a single rank never
+    produces -- so `--enable-tbo` was a dead switch on every TP-only deployment,
+    silently, and PCP+TBO (which is always dp==1) could never fire at all.
+    """
+    on = _decide(tbo_on=True, local_tbo=(True, True, 4, 4))
+    assert on.tbo_collective_active
+
+    # Both bits are required: reaching the min-token bar is not enough if the
+    # batch cannot be split, and vice versa.
+    for local in ((True, False, 4, 4), (False, True, 4, 4)):
+        assert not _decide(tbo_on=True, local_tbo=local).tbo_collective_active
+    assert not _decide(tbo_on=False, local_tbo=(True, True, 4, 4)).tbo_collective_active
+
+
+def test_a_peer_that_cannot_split_vetoes_tbo_for_the_group(monkeypatch):
+    """With peers the answer is theirs, not ours -- the reduction decides."""
+    _fake_sync(monkeypatch, peer_tokens=8, peer_bs=8)
+    mode = _decide(
+        _batch(seqs=8, q=1),
+        dp_size=2,
+        dp_group=object(),
+        tbo_on=True,
+        local_tbo=(True, True, 4, 4),
+    )
+    assert mode.tbo_collective_active == mode.sync.tbo_collective_active
+
+
+# --------------------------------------------------------------------------- #
+# running_tokens -- what everything per-row runs at. NOT the batch.            #
+# --------------------------------------------------------------------------- #
+
+
+def test_the_decode_height_exceeds_what_any_rank_scheduled():
+    """Why the reduced token count cannot serve on its own.
+
+    The batch is rounded onto the ladder AFTER the sync, so the graph runs rows
+    nobody reported: 50 seqs of 6 is 300 tokens, but the tensor has 64*6 = 384
+    rows. Padding to 300 leaves MoE 84 rows short.
+    """
+    mode = _decide(_batch(seqs=50, q=6))
+    assert mode.running_bs == 64 and mode.running_tokens == 384
+
+
+def test_a_prefill_step_runs_its_own_rows(monkeypatch):
+    """Prompts are ragged: no `bs * q` recovers a height, and no group max
+    stands in for one either.
+
+    `running_bs` here is 8, which would give a plausible-looking 8*6=48 if the
+    decode rule leaked in, and the peer reports 700, which would win if the
+    group max did. The answer is this rank's own 100 -- a prefilling rank is
+    exactly the case where the ranks are NOT level, so each runs what it has.
+    """
+    _fake_sync(monkeypatch, peer_tokens=700, peer_bs=2, peer_prefill=True)
+    mode = _decide(
+        _batch(seqs=8, tokens=100, prefill_tokens=100, q=6),
+        dp_size=2,
+        dp_group=object(),
+    )
+    assert mode.is_prefill and not mode.running_tokens_are_unified
+    assert mode.running_tokens == 100
+
+
+def test_the_variable_length_path_runs_this_rank_alone_in_both_units(monkeypatch):
+    """A prefilling peer puts every rank on the variable-length gather, and
+    that gather's contract is that nothing was padded -- so NEITHER unit may
+    come from the group. Unifying the batch here while leaving the token count
+    local is what cost 0.9530 -> 0.9447: every published width stayed
+    self-consistent, just consistent with a batch this rank never scheduled.
+
+    Reachable only with a mixed step, which no single-DP run can show."""
+    _fake_sync(monkeypatch, peer_tokens=9999, peer_bs=48, peer_prefill=True)
+    mode = _decide(_batch(seqs=3, q=6), dp_size=2, dp_group=object())
+    assert not mode.running_tokens_are_unified
+    assert mode.running_tokens == 18 and mode.scheduled_tokens == 18
+    # ...while the BATCH is still the group's. The two units answer to
+    # different consumers: an over-wide per-sequence array costs a few sentinel
+    # rows, an over-wide height is what the gather asserts against.
+    assert mode.running_bs == 48 and mode.scheduled_bs == 3
+
+
+def test_a_prefilling_peer_makes_the_whole_group_ragged(monkeypatch):
+    """The one meaning `running_tokens_are_unified` now has.
+
+    It used to also be True whenever dp-attention was off -- one flag standing
+    for two questions, which let a decoding rank answer "we are all level" while
+    a peer prefilled. Everything downstream that reads it (MoE's fixed-size
+    gather, the DP-sharded LM head) needs the strong claim, so the flag makes
+    only the strong claim and a mixed step drops every rank to eager.
+    """
+    _fake_sync(monkeypatch, peer_tokens=5, peer_bs=5, peer_prefill=True)
+    mixed = _decide(_batch(seqs=8, q=1), dp_size=2, dp_group=object())
+    _fake_sync(monkeypatch, peer_tokens=5, peer_bs=5, peer_prefill=False)
+    level = _decide(_batch(seqs=8, q=1), dp_size=2, dp_group=object())
+
+    assert not mixed.running_tokens_are_unified and not mixed.use_cudagraph
+    assert level.running_tokens_are_unified and level.use_cudagraph
+
+
+def test_a_flat_step_takes_the_bucket_and_not_a_product():
+    """A block drafter's rows are a packed run, so they are a bucket rather than
+    `running_bs * q`. Chosen here and read back everywhere -- two searches is how
+    the height and the replay came to disagree."""
+    mode = _decide(
+        _batch(seqs=10, q=6), captured_tokens=[48, 96, 192], is_block_drafter=True
+    )
+    assert mode.running_bs == 16, "the batch still rounds onto the seq ladder"
+    assert mode.running_tokens == 96, "smallest q-divisible bucket >= 60"
+    assert mode.piecewise_captured
+
+
+def test_a_rectangular_layout_never_takes_a_bucket_below_its_own_rectangle():
+    """Only the flat layout may shrink below `running_bs * q`.
+
+    Every per-token buffer of a rectangular backend spans one row per
+    (seq, query) and `slot_mapping[i * q + j]` addresses through it, so a bucket
+    under that leaves the tail of the rectangle unwritten. Same batch as above,
+    same buckets -- only the layout differs.
+    """
+    mode = _decide(
+        _batch(seqs=10, q=6), captured_tokens=[48, 96, 192], is_block_drafter=False
+    )
+    assert mode.running_tokens == 16 * 6
+
+
+def test_no_fitting_bucket_declines_the_graph_rather_than_claiming_one():
+    """The miss must not be reported as captured.
+
+    `running_bs * q` is the fallback precisely because no bucket covered the
+    step, so it is not in the captured set either -- and dispatching PIECEWISE
+    at a width nothing recorded makes the wrapper capture one mid-serve.
+    """
+    mode = _decide(
+        _batch(seqs=10, q=6), captured_tokens=[12, 24], is_block_drafter=True
+    )
+    assert mode.running_tokens == 16 * 6
+    assert not mode.piecewise_captured
+
+
+def test_the_rectangle_replays_only_when_it_was_captured():
+    fits = _decide(
+        _batch(seqs=8, q=2), captured_tokens=[16, 32], is_block_drafter=False
+    )
+    assert fits.running_tokens == 16 and fits.piecewise_captured
+    misses = _decide(_batch(seqs=8, q=2), captured_tokens=[24], is_block_drafter=False)
+    assert misses.running_tokens == 16 and not misses.piecewise_captured
+
+
+# --------------------------------------------------------------------------- #
+# Which recorded width a packed run lands on. Asked through `decide`, there    #
+# being no second entry point to ask it through.                               #
+# --------------------------------------------------------------------------- #
+
+
+def _packed(*, scheduled_tokens, q, captured):
+    """A block-drafter step whose packed run is `scheduled_tokens` long, over
+    `q`-wide sequences. `seqs=1` keeps the rectangle out of the answer's way."""
+    return _decide(
+        _batch(seqs=1, q=q, tokens=scheduled_tokens),
+        captured_tokens=captured,
+        is_block_drafter=True,
+    )
+
+
+def test_nothing_recorded_means_no_packed_width_to_land_on():
+    """Under FULL (non-PIECEWISE) cudagraphs nothing flat is recorded, so the
+    run has nowhere to land and the step stays rectangular."""
+    for captured in ([], None):
+        mode = _packed(scheduled_tokens=3, q=6, captured=captured)
+        assert mode.running_tokens == 1 * 6 and not mode.piecewise_captured
+
+
+def test_the_width_covers_the_run_and_tiles_it():
+    captured = [6, 12, 24]
+    assert _packed(scheduled_tokens=3, q=6, captured=captured).running_tokens == 6
+    assert _packed(scheduled_tokens=6, q=6, captured=captured).running_tokens == 6
+    assert (
+        _packed(scheduled_tokens=7, q=6, captured=captured).running_tokens == 12
+    ), "smallest that holds it"
+
+    # Longer than anything recorded -> forwarded eagerly at its own length.
+    over = _packed(scheduled_tokens=25, q=6, captured=captured)
+    assert over.running_tokens == 25 and not over.piecewise_captured
+
+    # Covers the run but does not tile the per-seq rows at stride q.
+    untiled = _packed(scheduled_tokens=3, q=6, captured=[10])
+    assert untiled.running_tokens == 1 * 6 and not untiled.piecewise_captured
+
+
+# --------------------------------------------------------------------------- #
+# assert_shape_contract -- the padded step, which is the one that can be wrong. #
+# --------------------------------------------------------------------------- #
+
+
+def _md(*, mode, slot_mapping_rows=None, cu_rows=None, slot_out_rows=None):
+    """What a builder publishes, defaulting to what `mode` says it should.
+
+    Defaulting rather than requiring each width keeps a case to the ONE it is
+    about; a test that had to restate all three would pass while contradicting
+    itself, which is how the pair this replaces came to be built on a
+    combination no builder emits.
+    """
+
+    def _z(n):
+        return torch.zeros(n, dtype=torch.int32)
+
+    return types.SimpleNamespace(
+        slot_mapping=_z(
+            mode.running_tokens if slot_mapping_rows is None else slot_mapping_rows
+        ),
+        cu_seqlens_q=_z(mode.running_bs + 1 if cu_rows is None else cu_rows),
+        state_slot_out=_z(mode.running_bs if slot_out_rows is None else slot_out_rows),
+        max_seqlen_q=mode.max_seqlen_q,
+    )
+
+
+def _ids(mode):
+    return torch.zeros(mode.scheduled_tokens, dtype=torch.int32)
+
+
+def test_a_padded_step_whose_widths_all_agree_passes():
+    """50 sequences of 6 run at 64: per-token structures span the rectangle,
+    per-sequence ones the batch. The check has to admit this -- it is what
+    every cudagraph decode step looks like."""
+    mode = _decide(_batch(seqs=50, q=6))
+    assert mode.running_bs == 64 and mode.running_tokens == 384
+    mode.assert_shape_contract(_ids(mode), _md(mode=mode))
+
+
+def test_per_token_structures_left_at_the_scheduled_rows_are_caught():
+    """The defect this check exists for: the batch was widened and the token
+    dimension was not. Publishing `slot_mapping` at `scheduled_tokens` on a
+    step running `running_tokens` rows leaves the tail of the rectangle
+    addressing memory the builder never wrote."""
+    mode = _decide(_batch(seqs=50, q=6))
+    with pytest.raises(AssertionError, match="slot_mapping"):
+        mode.assert_shape_contract(
+            _ids(mode), _md(mode=mode, slot_mapping_rows=mode.scheduled_tokens)
+        )
+
+
+def test_per_sequence_structures_left_at_the_scheduled_batch_are_caught():
+    """The mirror image, and the one that stays silent at runtime: a
+    `cu_seqlens_q` published at `scheduled_bs` while the step runs
+    `running_bs` hands attention a boundary list short of its own batch."""
+    mode = _decide(_batch(seqs=50, q=6))
+    with pytest.raises(AssertionError, match="cu_seqlens_q"):
+        mode.assert_shape_contract(
+            _ids(mode), _md(mode=mode, cu_rows=mode.scheduled_bs + 1)
+        )
+    with pytest.raises(AssertionError, match="state_slot_out"):
+        mode.assert_shape_contract(
+            _ids(mode), _md(mode=mode, slot_out_rows=mode.scheduled_bs)
+        )
+
+
+def test_a_ragged_group_is_checked_on_widths_but_not_on_the_rectangle(monkeypatch):
+    """A prefilling peer drops every rank to its own rows, so `scheduled_tokens`
+    need not tile at `max_seqlen_q` and no rectangle bounds it. The width
+    equalities still hold, and still have to be checked -- this is the step
+    where the batch and the token count are reduced differently."""
+    _fake_sync(monkeypatch, peer_tokens=9999, peer_bs=48, peer_prefill=True)
+    mode = _decide(_batch(seqs=3, q=6), dp_size=2, dp_group=object())
+    assert not mode.running_tokens_are_unified
+    mode.assert_shape_contract(_ids(mode), _md(mode=mode))
+    with pytest.raises(AssertionError, match="cu_seqlens_q"):
+        mode.assert_shape_contract(_ids(mode), _md(mode=mode, cu_rows=2))
+
+
+def test_a_backend_that_publishes_none_is_not_invented_a_width_for():
+    """V4 carries no `slot_mapping` on decode; absent is not zero-length."""
+    mode = _decide(_batch(seqs=50, q=6))
+    md = _md(mode=mode)
+    md.slot_mapping = None
+    md.state_slot_out = None
+    mode.assert_shape_contract(_ids(mode), md)
diff --git a/tests/test_mori_dispatch_trim_bound.py b/tests/test_mori_dispatch_trim_bound.py
index e65886a7d3..8bce729ef9 100644
--- a/tests/test_mori_dispatch_trim_bound.py
+++ b/tests/test_mori_dispatch_trim_bound.py
@@ -26,7 +26,7 @@
 
 def _trim(monkeypatch, *, running_tokens, dp_size, topk, recv_rows):
     context = SimpleNamespace(
-        running_tokens=running_tokens, is_prefill=False, dp_uniform_decode=True
+        running_tokens=running_tokens, is_prefill=False, running_tokens_are_unified=True
     )
     monkeypatch.setattr(
         mk, "get_forward_context", lambda: SimpleNamespace(context=context)
diff --git a/tests/test_mtp_deferred_status_queue.py b/tests/test_mtp_deferred_status_queue.py
index 818c574d39..abc8234079 100644
--- a/tests/test_mtp_deferred_status_queue.py
+++ b/tests/test_mtp_deferred_status_queue.py
@@ -4,7 +4,13 @@
 import numpy as np
 import pytest
 
-pytest.importorskip("aiter", reason="requires AITER to import model_runner")
+# On the module, not on "aiter" -- see tests/test_dspark.py for why a name-only
+# guard passes on the non-GPU runner and then aborts collection.
+pytest.importorskip(
+    "atom.model_engine.model_runner",
+    reason="requires AITER to import model_runner",
+    exc_type=ImportError,
+)
 
 from atom.model_engine.model_runner import tokenIDProcessor
 from atom.model_engine.scheduler import ScheduledBatch
@@ -22,6 +28,72 @@ def _prefill_batch(is_final_chunk: list[bool]) -> ScheduledBatch:
     return batch
 
 
+def _batch_with_ids(req_ids: list[int], *, is_dummy_run: bool) -> ScheduledBatch:
+    """The two fields `get_token_locations` reads, and nothing else."""
+    batch = object.__new__(ScheduledBatch)
+    batch.req_ids = req_ids
+    batch.is_dummy_run = is_dummy_run
+    return batch
+
+
+def _locations(prev: ScheduledBatch, cur: ScheduledBatch):
+    processor = object.__new__(tokenIDProcessor)
+    processor.prev_batch = prev
+    return tokenIDProcessor.get_token_locations(processor, cur)
+
+
+def test_a_dummy_carries_nothing_over_to_the_next_dummy():
+    """Back-to-back DP-sync dummies must not look like one carried-over request.
+
+    `dummy_execution` fabricates its sequence with a fixed id, so a second
+    dummy matches the first by id and is taken for a request that ran last
+    step. It then reads that "request's" anchor and drafts out of
+    `prev_token_ids` / `draft_token_ids`, which by then hold whatever the draft
+    pass's graph-pool storage was overwritten with -- out-of-vocab ids that
+    fault the target's embedding gather. A rank idling behind someone else's
+    long chunked prefill runs dummies back to back, so this is the ordinary
+    case at high concurrency, not a corner.
+    """
+    dummy_id = -1  # what dummy_execution builds; the collision is with ITSELF
+    first = _batch_with_ids([dummy_id], is_dummy_run=True)
+    second = _batch_with_ids([dummy_id], is_dummy_run=True)
+
+    locs = _locations(first, second)
+
+    assert locs.deferred_curr.tolist() == []
+    assert locs.deferred_prev.tolist() == []
+    assert locs.new_curr.tolist() == [0]
+
+
+def test_a_real_batch_after_a_dummy_takes_the_host_path():
+    """...and so does a real request that wakes up after one.
+
+    Not a consolation prize for the line above: the dummy's own `postprocess`
+    reports the PREVIOUS batch's tokens, so by the time this request is
+    scheduled its anchor is already on the host. Reading `prev_token_ids`
+    instead would be reading the dummy's forward.
+    """
+    dummy = _batch_with_ids([-1], is_dummy_run=True)
+    real = _batch_with_ids([7, 9], is_dummy_run=False)
+
+    locs = _locations(dummy, real)
+
+    assert locs.deferred_curr.tolist() == []
+    assert locs.new_curr.tolist() == [0, 1]
+
+
+def test_a_real_batch_still_carries_over_from_a_real_batch():
+    """The dummy rule must not cost an ordinary decode step its deferred rows."""
+    prev = _batch_with_ids([7, 9], is_dummy_run=False)
+    cur = _batch_with_ids([9, 11], is_dummy_run=False)
+
+    locs = _locations(prev, cur)
+
+    assert locs.deferred_curr.tolist() == [0]  # req 9 is position 0 now...
+    assert locs.deferred_prev.tolist() == [1]  # ...and was position 1 before
+    assert locs.new_curr.tolist() == [1]
+
+
 def _processor() -> tokenIDProcessor:
     processor = object.__new__(tokenIDProcessor)
     processor.input_ids = SimpleNamespace(
@@ -45,8 +117,8 @@ def test_middle_prefills_preserve_status_until_mixed_final_batch():
 
     # Pure middle chunks skip postprocess, so neither the deferred-token queue
     # nor its matching MTP-status queue may advance.
-    tokenIDProcessor.prepare_input_ids(processor, _prefill_batch([False]))
-    tokenIDProcessor.prepare_input_ids(processor, _prefill_batch([False, False]))
+    tokenIDProcessor.prepare_input_ids(processor, _prefill_batch([False]), 1)
+    tokenIDProcessor.prepare_input_ids(processor, _prefill_batch([False, False]), 1)
 
     processor.recv_mtp_status_async.assert_not_called()
     np.testing.assert_array_equal(processor.prev_rejected_num, [7])
@@ -55,7 +127,7 @@ def test_middle_prefills_preserve_status_until_mixed_final_batch():
     # If any request reaches its final chunk, the batch runs postprocess. Its
     # status dequeue must therefore happen exactly once, even though another
     # request in the same batch is still a middle chunk.
-    tokenIDProcessor.prepare_input_ids(processor, _prefill_batch([False, True]))
+    tokenIDProcessor.prepare_input_ids(processor, _prefill_batch([False, True]), 1)
 
     processor.recv_mtp_status_async.assert_called_once_with()
     np.testing.assert_array_equal(processor.prev_rejected_num, [2])
diff --git a/tests/test_postprocess_width.py b/tests/test_postprocess_width.py
new file mode 100644
index 0000000000..5c14d5b07c
--- /dev/null
+++ b/tests/test_postprocess_width.py
@@ -0,0 +1,181 @@
+# SPDX-License-Identifier: MIT
+"""What `postprocess` hands the sampler, in rows.
+
+The LM head emits one row per sequence the step FORWARDED. On a prefill that is
+`running_bs` -- the ladder-rounded, DP-agreed batch -- while the sampler's
+per-request parameters are `scheduled_bs` wide, because they describe requests.
+Every step between the two has to happen on the same side of the cut.
+
+Skipped where `model_engine.model_runner` cannot import (it needs aiter at
+module load); nothing here touches a GPU.
+"""
+
+import types
+
+import numpy as np
+import pytest
+import torch
+
+# Not `importorskip("aiter")`: the non-GPU runner has `aiter` as a namespace
+# package, so that import succeeds and `model_runner` still fails on a symbol.
+# Naming the module we actually need catches both, and `exc_type` is what makes
+# a non-ModuleNotFoundError ImportError a skip rather than a collection error --
+# one of which takes the whole suite down with it. See tests/import_guard.py.
+mod = pytest.importorskip(
+    "atom.model_engine.model_runner",
+    reason="model_runner imports aiter at module load",
+    exc_type=ImportError,
+)
+
+
+def _batch(seqs):
+    return types.SimpleNamespace(
+        total_seqs_num=seqs,
+        total_tokens_num=seqs,
+        req_ids=list(range(seqs)),
+        return_logprobs=[False] * seqs,
+    )
+
+
+def _runner(seen):
+    """A ModelRunner carrying only what the no-spec branch of postprocess reads."""
+    runner = object.__new__(mod.ModelRunner)
+
+    def _sampler(logits, temperatures, top_ks, top_ps, all_greedy, **kw):
+        seen["logits_rows"] = logits.shape[0]
+        seen["temperature_rows"] = temperatures.shape[0]
+        return torch.zeros(logits.shape[0], dtype=torch.int)
+
+    runner.sampler = _sampler
+    runner.forward_done_event = types.SimpleNamespace(record=lambda: None)
+    runner.tokenID_processor = types.SimpleNamespace(
+        is_deferred_out=False,
+        prev_batch=None,
+        default_num_rejected_tokens=torch.zeros(64, dtype=torch.int32),
+        prepare_sampled_ids=lambda *a, **k: ({}, {}),
+    )
+    return runner
+
+
+def test_the_sampler_gets_as_many_rows_as_it_has_parameters(monkeypatch):
+    """A prefill forwards more sequences than it scheduled, and only one of the
+    two numbers describes a request.
+
+    `decide` rounds every step onto the capture ladder, so a 3-sequence prefill
+    runs 4 rows, and `prepare_prefill` pads `cu_seqlens_q` to match -- which is
+    what the LM head slices by. `prepare_sample` sizes `temperatures` at the
+    scheduled count and never pads. Sampling before the cut divides
+    `[4, V]` by `[3, 1]`, which raises on the top-k/top-p path and reads past
+    the parameter buffer on the temperature one.
+    """
+    seen = {}
+    monkeypatch.setattr(
+        mod,
+        "get_forward_context",
+        lambda: types.SimpleNamespace(spec_decode_metadata=None),
+    )
+    monkeypatch.setattr(
+        mod, "get_tp_group", lambda: types.SimpleNamespace(world_size=1)
+    )
+
+    scheduled_bs, running_bs, vocab = 3, 4, 8
+    out = _runner(seen).postprocess(
+        batch=_batch(scheduled_bs),
+        logits=torch.zeros(running_bs, vocab),
+        temperatures=torch.ones(scheduled_bs),
+        top_ks=None,
+        top_ps=None,
+        all_greedy=False,
+        hidden_states=None,
+    )
+
+    assert seen["logits_rows"] == seen["temperature_rows"] == scheduled_bs
+    assert out.num_rejected.tolist() == [0] * scheduled_bs
+
+
+def test_a_step_that_padded_nothing_is_unaffected(monkeypatch):
+    """The cut is a slice, so the common case has to stay a no-op."""
+    seen = {}
+    monkeypatch.setattr(
+        mod,
+        "get_forward_context",
+        lambda: types.SimpleNamespace(spec_decode_metadata=None),
+    )
+    monkeypatch.setattr(
+        mod, "get_tp_group", lambda: types.SimpleNamespace(world_size=1)
+    )
+
+    bs = 5
+    _runner(seen).postprocess(
+        batch=_batch(bs),
+        logits=torch.zeros(bs, 8),
+        temperatures=torch.ones(bs),
+        top_ks=None,
+        top_ps=None,
+        all_greedy=False,
+        hidden_states=None,
+    )
+
+    assert seen["logits_rows"] == bs
+
+
+def test_the_logprob_gather_reads_the_cut_logits(monkeypatch):
+    """`log_probs.gather` indexes with the sampled ids, which are per-request.
+
+    Left on the padded logits it is `[running_bs, V]` gathered by `[scheduled_bs,
+    1]` -- a second break from the same uncut tensor, and the one a client asking
+    for logprobs hits.
+    """
+    seen = {}
+    monkeypatch.setattr(
+        mod,
+        "get_forward_context",
+        lambda: types.SimpleNamespace(spec_decode_metadata=None),
+    )
+    monkeypatch.setattr(
+        mod, "get_tp_group", lambda: types.SimpleNamespace(world_size=1)
+    )
+
+    scheduled_bs, running_bs = 3, 4
+    batch = _batch(scheduled_bs)
+    batch.return_logprobs = [True] * scheduled_bs
+    captured = {}
+    runner = _runner(seen)
+    runner.tokenID_processor.prepare_sampled_ids = lambda b, ids, ev, lp: (
+        captured.update(logprobs=lp) or ({}, {})
+    )
+
+    runner.postprocess(
+        batch=batch,
+        logits=torch.zeros(running_bs, 8),
+        temperatures=torch.ones(scheduled_bs),
+        top_ks=None,
+        top_ps=None,
+        all_greedy=False,
+        hidden_states=None,
+    )
+
+    assert captured["logprobs"].shape == (scheduled_bs,)
+
+
+def test_prefill_really_does_run_wider_than_it_scheduled():
+    """The premise, from the decider itself -- not a number invented here."""
+    mode = mod.ForwardMode.decide(
+        batch=types.SimpleNamespace(
+            total_tokens_num_prefill=9,
+            total_tokens_num=9,
+            total_seqs_num=3,
+            is_dummy_run=False,
+        ),
+        dp_size=1,
+        dp_group=None,
+        enforce_eager=False,
+        capture_sizes=np.array([1, 2, 4, 8], dtype=np.int32),
+        captured_tokens=None,
+        is_block_drafter=False,
+        tbo_on=False,
+        local_tbo=(False, False, 0, 0),
+        max_seqlen_q=1,
+    )
+
+    assert (mode.scheduled_bs, mode.running_bs) == (3, 4)