diff --git a/atom/config.py b/atom/config.py index 04f169bd23..0f42bb2d9f 100644 --- a/atom/config.py +++ b/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 diff --git a/atom/model_engine/arg_utils.py b/atom/model_engine/arg_utils.py index 2ec68f4cb5..f67dffdda2 100644 --- a/atom/model_engine/arg_utils.py +++ b/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, diff --git a/atom/model_engine/prefill_delayer.py b/atom/model_engine/prefill_delayer.py index 36a55c8088..72b0d57c38 100644 --- a/atom/model_engine/prefill_delayer.py +++ b/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,6 +133,7 @@ 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. @@ -139,10 +141,16 @@ def should_allow_prefill( 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,7 +165,7 @@ 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, @@ -165,11 +173,11 @@ def should_allow_prefill( ) 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 diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index 295a76f3e9..6c7da953cc 100644 --- a/atom/model_engine/scheduler.py +++ b/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 diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index ea846ebed3..c663a89796 100644 --- a/atom/model_ops/attention_mla.py +++ b/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 diff --git a/atom/model_ops/moe.py b/atom/model_ops/moe.py index bd943f5498..4070ef1f4d 100644 --- a/atom/model_ops/moe.py +++ b/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.) diff --git a/atom/model_ops/topK.py b/atom/model_ops/topK.py index 2c1a3599e5..d85f482987 100644 --- a/atom/model_ops/topK.py +++ b/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: diff --git a/atom/utils/tbo/ubatch_wrapper.py b/atom/utils/tbo/ubatch_wrapper.py index b21a7eca36..e9842d1522 100644 --- a/atom/utils/tbo/ubatch_wrapper.py +++ b/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( diff --git a/tests/conftest.py b/tests/conftest.py index f875e6fd9b..999d7ec2e7 100644 --- a/tests/conftest.py +++ b/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(), diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 659d900813..a324af0c55 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -673,3 +673,49 @@ def test_normal_decode_window_unchanged(self): ) assert list(batch.scheduled_tokens) == toks[-(mtp_k + 1) :] + + +# ── Prefill dense-batch gate ─────────────────────────────────────────────── + + +class TestPrefillBatchThreshold: + def test_threshold_defaults_to_batch_budget(self): + cfg = MockConfig(max_num_batched_tokens=32) + sched = Scheduler(cfg) + # Threshold is derived from the batch-token budget, not a separate knob. + assert sched.prefill_batch_token_threshold == 32 + + def test_token_threshold_can_be_lower_than_batch_budget(self): + sched = Scheduler( + MockConfig( + max_num_batched_tokens=32, + prefill_batch_token_threshold=20, + ) + ) + assert sched.prefill_batch_token_threshold == 20 + + def test_waiting_prefill_tokens_reaches_threshold(self, seq_factory): + # chunked prefill on so a 40-token prompt is clamped to the 32 budget + # (not rejected as oversized) and counts toward the threshold. + cfg = MockConfig(max_num_batched_tokens=32, enable_chunked_prefill=True) + sched = Scheduler(cfg) + # 40 waiting tokens are clamped to the 32-token batch budget. + sched.waiting.append(seq_factory(list(range(40)))) + assert sched._waiting_prefill_tokens() >= 32 + + def test_without_delayer_underfull_prefill_is_admitted(self, seq_factory): + cfg = MockConfig( + max_num_batched_tokens=32, + enable_chunked_prefill=False, + data_parallel_size=8, + ) + sched = Scheduler(cfg) + r = seq_factory([1, 2, 3]) + r.status = SequenceStatus.RUNNING + r.type = SequenceType.DECODE + sched.block_manager.allocate(r, 0) + sched.running.append(r) + sched.waiting.append(seq_factory(list(range(8)))) + batch, _ = sched.schedule() + assert batch is not None + assert batch.total_seqs_num_prefill == 1