Settle a step's shape once, and fix the three defects that exposed

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

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,

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

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,15 +1112,15 @@ 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
             ]
 
         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,
             )