Enable TBO Support & Fix Accuracy Regressions for Kimi K2.5

atom/config.py

@@ -1032,6 +1032,7 @@ class Config:
     long_prefill_token_threshold: int = 0
     attn_prefill_chunk_size: int = 16384
     scheduler_delay_factor: float = 0.0
+    prefill_batch_token_threshold: int = 0
     max_num_seqs: int = 512
     max_model_len: int | None = None
     gpu_memory_utilization: float = 0.9

atom/model_engine/arg_utils.py

@@ -43,6 +43,7 @@ class EngineArgs:
     attn_prefill_chunk_size: int = 16384
     enable_chunked_prefill: bool = True
     scheduler_delay_factor: float = 0.0
+    prefill_batch_token_threshold: int = 0
     max_num_seqs: int = 512
     gpu_memory_utilization: float = 0.9
     cudagraph_capture_sizes: str = "[1,2,4,8,16,32,48,64,128,256]"
@@ -212,6 +213,15 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
             default=16384,
             help="Maximum number of tokens to batch together in async engine",
         )
+        parser.add_argument(
+            "--prefill-batch-token-threshold",
+            type=int,
+            default=0,
+            help=(
+                "Hold new prefills until this many eligible tokens are waiting. "
+                "0 uses max-num-batched-tokens."
+            ),
+        )
         parser.add_argument(
             "--long-prefill-token-threshold",
             type=int,

atom/model_engine/prefill_delayer.py

@@ -7,7 +7,7 @@
 
 Mechanism (per scheduler tick):
   1. Each DP rank reports its local state via cpu all_gather:
-       (local_prefillable, watermark_force_allow)
+       (local_prefillable, local_prefill_sufficient, watermark_force_allow)
   2. Compute `prefillable_status` ∈ {all, none, mixed}:
        - "all"   → every rank has a new prefill ready  → allow (8-way aligned)
        - "none"  → no rank has any prefill             → allow (vacuous)
@@ -97,10 +97,11 @@ def __init__(
         # Encoding:
         #   slot 0 = local_prefillable          (MAX → "any rank prefillable")
         #   slot 1 = local_force                (MAX → "any rank forces allow")
-        #   slot 2 = NOT local_prefillable      (MAX → "any rank lacks prefill")
+        #   slot 2 = NOT local_prefill_sufficient
+        #            (MAX → "any rank lacks the configured dense batch")
         # Then prefillable_status:
-        #   any_prefillable AND any_not_prefillable → "mixed"
-        #   any_prefillable AND NOT any_not_prefillable → "all"
+        #   any_prefillable AND any_not_sufficient → "mixed"
+        #   any_prefillable AND NOT any_not_sufficient → "all"
         #   NOT any_prefillable → "none"
         # Single all_reduce, 3 int64s on cpu — negligible overhead.
         self._reduce_buf = torch.zeros(3, dtype=torch.int64, device="cpu")
@@ -132,17 +133,24 @@ def should_allow_prefill(
         self,
         local_prefillable: bool,
         token_usage: float,
+        local_prefill_sufficient: Optional[bool] = None,
     ) -> bool:
         """
         Returns True iff this rank is allowed to admit new prefills this tick.
 
         Args:
             local_prefillable: this rank has at least one new prefill ready
                 (i.e. self.waiting non-empty and admission would succeed).
+            local_prefill_sufficient: this rank also meets the configured
+                token/request density target. Defaults to local_prefillable
+                for compatibility with existing callers.
             token_usage: fraction of KV cache blocks currently in use
                 (used_blocks / total_blocks ∈ [0, 1]). Used by the
                 low-watermark safety valve.
         """
+        if local_prefill_sufficient is None:
+            local_prefill_sufficient = local_prefillable
+
         # Local "force allow" if KV cache is underutilized — don't delay
         # when GPU is starving. Only meaningful if this rank actually has
         # a prefill to push through (otherwise force_allow is a no-op).
@@ -157,19 +165,19 @@ def should_allow_prefill(
         # Cross-DP MAX-reduce: 3 booleans encoded as int64.
         self._reduce_buf[0] = 1 if local_prefillable else 0
         self._reduce_buf[1] = 1 if force else 0
-        self._reduce_buf[2] = 0 if local_prefillable else 1
+        self._reduce_buf[2] = 0 if local_prefill_sufficient else 1
         torch.distributed.all_reduce(
             self._reduce_buf,
             op=torch.distributed.ReduceOp.MAX,
             group=self.cpu_group,
         )
         any_prefillable = int(self._reduce_buf[0].item()) > 0
         force_max = int(self._reduce_buf[1].item())
-        any_not_prefillable = int(self._reduce_buf[2].item()) > 0
+        any_not_sufficient = int(self._reduce_buf[2].item()) > 0
 
         # Derive 3-way status: all / none / mixed.
         prefillable_max = 1 if any_prefillable else 0
-        prefillable_min = 0 if any_not_prefillable else 1
+        prefillable_min = 0 if any_not_sufficient else 1
 
         # Watermark short-circuit: ANY rank below the watermark forces all
         # ranks to allow this tick. Without this the delayer can stall a
@@ -211,7 +219,7 @@ def should_allow_prefill(
                     f"[PrefillDelayer] DELAY: count={self._delayed_count} "
                     f"elapsed={elapsed_ms:.1f}ms "
                     f"any_prefillable={any_prefillable} "
-                    f"any_not_prefillable={any_not_prefillable}"
+                    f"any_not_sufficient={any_not_sufficient}"
                 )
             self._maybe_log()
             return False

atom/model_engine/scheduler.py

@@ -447,6 +447,11 @@ def __init__(self, config: Config):
         # Latency of the last prompt step
         self.last_prompt_latency = 0.0
         self.delay_factor = config.scheduler_delay_factor
+        self.prefill_batch_token_threshold = (
+            config.prefill_batch_token_threshold
+            if config.prefill_batch_token_threshold > 0
+            else self.max_num_batched_tokens
+        )
 
         # Speculative decoding
         self.use_spec = config.speculative_config is not None
@@ -559,6 +564,33 @@ def _can_admit_head_prefill(self) -> bool:
             return True
         return False
 
+    def _waiting_prefill_tokens(self) -> int:
+        """Sum of admissible new-prefill tokens sitting in the waiting queue,
+        capped at max_num_batched_tokens (we only care whether a *dense* batch
+        can be filled). Skips unschedulable / remote-KV-waiting entries and
+        clamps each seq to the chunked-prefill / budget limit, mirroring the
+        Phase-2 admission math so the count reflects what would actually pack
+        into one prefill step."""
+        total = 0
+        cap = self.max_num_batched_tokens
+        for seq in self.waiting:
+            if self._unschedulable_reason(seq) is not None:
+                continue
+            if seq.status == SequenceStatus.WAITING_FOR_REMOTE_KVS:
+                continue
+            n = seq.num_tokens - seq.num_cached_tokens
+            if (
+                self.enable_chunked_prefill
+                and 0 < self.long_prefill_token_threshold < n
+            ):
+                n = self.long_prefill_token_threshold
+            if n > cap:
+                n = cap
+            total += n
+            if total >= cap:
+                return cap
+        return total
+
     def _kv_usage(self) -> float:
         """Fraction of KV-cache blocks currently in use ∈ [0, 1].
 
@@ -721,24 +753,41 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]:
 
         # should_allow_prefill() runs a cross-DP all_reduce and MUST be called
         # every tick on every rank for lockstep — hence before the early-return.
+        _delayer_allows_prefill = True
         if self.prefill_delayer is not None:
-            delayer_allows = self.prefill_delayer.should_allow_prefill(
-                local_prefillable=self._can_admit_head_prefill(),
+            # A rank counts as "prefillable" for cross-DP alignment only if it
+            # can admit a prefill AND has a full batch's worth of waiting tokens.
+            # This makes all ranks align on firing dense prefills together
+            # instead of straggling partials.
+            _local_prefillable = self._can_admit_head_prefill()
+            _local_prefill_sufficient = (
+                _local_prefillable
+                and self._waiting_prefill_tokens()
+                >= self.prefill_batch_token_threshold
+            )
+            _delayer_allows_prefill = self.prefill_delayer.should_allow_prefill(
+                local_prefillable=_local_prefillable,
+                local_prefill_sufficient=_local_prefill_sufficient,
                 token_usage=self._kv_usage(),
             )
-        else:
-            delayer_allows = True
 
         if not self.running and not self.waiting:
             return None
 
+        # PrefillDelayer's result is a cross-rank decision. Do not apply a
+        # second rank-local density gate here: ranks could otherwise disagree
+        # after negotiation and recreate the mixed prefill/decode step the
+        # delayer is meant to prevent. Without a delayer, preserve immediate
+        # prefill admission rather than holding decode behind a local gate.
+        _new_prefill_allowed = _delayer_allows_prefill
+
         # ---- Phase 1: resume partial prefills from running ----
         # Gated by `delayer_allows` so cross-DP alignment still holds when one
         # rank is mid-chunked-prefill: a delayer veto skips both Phase 1 and
         # Phase 2 in lockstep. Inside that, skip the running-queue scan entirely
         # when no seq is mid-prefill — the common steady-state decode case —
         # using the counter maintained by postprocess / preempt / finished-removal.
-        if delayer_allows and self._partial_prefill_count > 0:
+        if _delayer_allows_prefill and self._partial_prefill_count > 0:
             for seq in self.running:
                 if num_seqs_prefill >= self.max_num_seqs:
                     break
@@ -759,7 +808,7 @@ def schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]]:
 
         # ---- Phase 2: new requests from waiting ----
         while (
-            delayer_allows
+            _new_prefill_allowed
             and (self.delay_factor <= 0 or self._passed_delay(time.time()))
             and self.waiting
             and num_seqs_prefill < self.max_num_seqs

atom/model_ops/attention_mla.py

@@ -1051,7 +1051,15 @@ def _forward_decode(
                     paged_kv_indices = self.sparse_kv_indices_buffer
 
             dp_size = get_dp_group().world_size
-            use_persistent_mode = not (dp_size > 1)
+            # for DPA, AITER has no non-persistent kernel: fp8 Q + fp8 KV with a
+            # GQA ratio of 64.
+            gqa_ratio = self.padded_num_heads // self.num_kv_heads
+            requires_persistent_mode = (
+                q.dtype == dtypes.fp8
+                and kv_buffer.dtype == dtypes.fp8
+                and gqa_ratio == 64
+            )
+            use_persistent_mode = dp_size == 1 or requires_persistent_mode
             if envs.ATOM_MLA_PAGE_SIZE > 1:
                 use_persistent_mode = False
 

atom/model_ops/moe.py

@@ -87,6 +87,9 @@ def from_model_config(a_quant_dtype: str | None) -> "MoEActivationQuant":
         return MoEActivationQuant.BF16
 
 
+_TBO_KEEPALIVE: dict[tuple[str, int], tuple[torch.Tensor, ...]] = {}
+
+
 class FusedMoeWeightScaleSupported(Enum):
     """Supported quantization strategies for MoE weight scales."""
 
@@ -2477,6 +2480,17 @@ def __init__(
                 ),
                 dim=0,
             )
+        # In the DP-attn fallback path (dp>1, no MORI all2all), MoE runs
+        # after all_gather_with_padding, so the token dim can be dp_size times
+        # the per-rank max.
+        moe_max_num_tokens = atom_config.max_num_batched_tokens
+        if (
+            self.moe_parallel_config.dp_size > 1
+            and not self.moe_parallel_config.use_all2all_kernels
+            and atom_config.enable_dp_attention
+        ):
+            moe_max_num_tokens *= self.moe_parallel_config.dp_size
+
         if fuse_shared_experts and self.num_fused_shared_experts > 0:
             init_aiter_topK_meta_data(
                 n_routed_experts=self.global_num_experts,
@@ -2489,7 +2503,7 @@ def __init__(
                     if is_rocm_aiter_fuse_routed_scaling_factor()
                     else 1 / self.routed_scaling_factor
                 ),
-                max_num_tokens=atom_config.max_num_batched_tokens,
+                max_num_tokens=moe_max_num_tokens,
                 is_EP=self.use_ep,
             )
         if fuse_shared_experts:
@@ -2529,7 +2543,7 @@ def __init__(
             moe_parallel_config=self.moe_parallel_config,
             in_dtype=atom_config.torch_dtype,
             a_quant_dtype=a_quant_dtype,
-            max_num_tokens=atom_config.max_num_batched_tokens,
+            max_num_tokens=moe_max_num_tokens,
             has_bias=self.has_bias,
             # is_act_and_mul=True,
             is_lora_enabled=False,
@@ -3448,6 +3462,26 @@ def forward(self, hidden_states: torch.Tensor, router_logits: torch.Tensor):
             hidden_states, router_logits, self.layer_name
         )
 
+    def _tbo_keepalive_slot(self) -> int:
+        try:
+            from atom.utils.tbo.ubatching import tbo_current_ubatch_id
+
+            return tbo_current_ubatch_id()
+        except Exception:
+            return 0
+
+    def _hold_tbo_keepalive(self, role: str, *tensors: torch.Tensor) -> None:
+        tensors = tuple(tensor for tensor in tensors if tensor is not None)
+        if tensors:
+            # Keep one previous tensor set per ubatch/role alive globally.
+            # The next same-role hold, often in the next MoE layer, happens
+            # after this ubatch has waited on the prior comm work, so
+            # overwriting here is the delayed safe release point.
+            key = (role, self._tbo_keepalive_slot())
+            if key in _TBO_KEEPALIVE:
+                del _TBO_KEEPALIVE[key]
+            _TBO_KEEPALIVE[key] = tensors
+
     def forward_impl_graph(
         self, hidden_states: torch.Tensor, router_logits: torch.Tensor
     ):
@@ -3478,6 +3512,7 @@ def forward_impl_graph(
                 )
 
                 tbo_yield_and_switch_from_compute_to_comm()
+                self._hold_tbo_keepalive("ag_source", hidden_states, router_logits)
 
             (
                 hidden_states,
@@ -3490,6 +3525,7 @@ def forward_impl_graph(
 
             if _tbo:
                 tbo_switch_to_compute_sync()
+                self._hold_tbo_keepalive("ag_output", hidden_states, router_logits)
 
         # Matrix multiply.
         final_hidden_states = self.quant_method.apply(
@@ -3515,6 +3551,7 @@ def forward_impl_graph(
         if use_dp_gather_scatter:
             if _tbo:
                 tbo_yield_and_switch_from_compute_to_comm()
+                self._hold_tbo_keepalive("rs_source", final_hidden_states)
             if dp_eager_mode:
                 final_hidden_states = reduce_scatterv(
                     final_hidden_states, sizes, dp_group
@@ -3525,6 +3562,7 @@ def forward_impl_graph(
                 )
             if _tbo:
                 tbo_switch_to_compute_sync()
+                self._hold_tbo_keepalive("rs_output", final_hidden_states)
 
         if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1):
             # Default set to False. (May have to add shared expert outputs.)

atom/model_ops/topK.py

@@ -25,7 +25,13 @@ def is_rocm_aiter_fusion_shared_expert_enabled_for_quant_config(
     # layout (set by the vLLM plugin under DP+EP); disable it there.
     if dp_size > 1 and config.moe_ep_flatten_tp_across_dp:
         return False
-    if dp_size > 1 and _has_module("mori") and config.enable_dp_attention:
+    use_mori_all2all = (
+        dp_size > 1
+        and _has_module("mori")
+        and config.enable_dp_attention
+        and config.enable_expert_parallel
+    )
+    if use_mori_all2all:
         return False
 
     if quant_config is not None and shared_expert_prefix is not None:

atom/utils/tbo/ubatch_wrapper.py

@@ -542,6 +542,7 @@ def _make_ubatch_context(
             batch_size=ub_num_reqs,
             graph_bs=graph_bs,
             is_draft=ctx.context.is_draft,
+            dp_uniform_decode=ctx.context.dp_uniform_decode,
         )
 
         return ForwardContext(

tests/conftest.py

@@ -52,7 +52,8 @@ def __init__(self, **kwargs):
 class _StubParallelConfig:
     """Placeholder for ParallelConfig."""
 
-    pass
+    def __init__(self, data_parallel_size: int = 1):
+        self.data_parallel_size = data_parallel_size
 
 
 _atom_config.Config = _StubConfig
@@ -145,7 +146,14 @@ def __init__(self, **overrides):
             eos_token_id=2,
             stop_token_ids=[],
             scheduler_delay_factor=0.0,
+            prefill_batch_token_threshold=0,
             speculative_config=None,
+            # DP size gates the dense-batch prefill hold (see Scheduler). Default
+            # 1 (gate off) so unrelated tests keep legacy behavior; gate tests
+            # pass data_parallel_size>1.
+            parallel_config=_StubParallelConfig(
+                data_parallel_size=overrides.pop("data_parallel_size", 1)
+            ),
             # Scheduler.__init__ reads config.hf_config.architectures for V4
             # SWA-warmup detection; a non-V4 stub keeps that path inert.
             hf_config=_MockHFConfig(),