From eb61cb28233ba7959155761c72d4f3298eec97ad Mon Sep 17 00:00:00 2001 From: "Wang, FangYuan" <39615225+At1a8@users.noreply.github.com> Date: Mon, 17 Aug 2026 13:40:29 +0800 Subject: [PATCH] [AMD] Support prefill context parallel two batch overlap for DeepSeek V4 (#33480) --- .../srt/batch_overlap/operations_strategy.py | 28 +- .../srt/batch_overlap/two_batch_overlap.py | 7 + python/sglang/srt/distributed/bootstrap.py | 6 + .../sglang/srt/distributed/parallel_state.py | 66 +++++ .../srt/layers/attention/dsv4/compressor.py | 36 +++ python/sglang/srt/layers/dp_attention.py | 9 + python/sglang/srt/layers/utils/cp_utils.py | 50 ++++ python/sglang/srt/models/deepseek_v4.py | 250 ++++++++++++++++-- python/sglang/srt/server_args.py | 6 + .../amd/test_deepseek_v4_pro_fp4_cp_tbo.py | 155 +++++++++++ 10 files changed, 594 insertions(+), 19 deletions(-) create mode 100644 test/registered/amd/test_deepseek_v4_pro_fp4_cp_tbo.py diff --git a/python/sglang/srt/batch_overlap/operations_strategy.py b/python/sglang/srt/batch_overlap/operations_strategy.py index af2171e3d..281e9c232 100644 --- a/python/sglang/srt/batch_overlap/operations_strategy.py +++ b/python/sglang/srt/batch_overlap/operations_strategy.py @@ -34,6 +34,7 @@ class OperationsStrategy: def init_new_tbo( layers: torch.nn.ModuleList, forward_mode: ForwardMode, + use_cp: bool = False, ) -> "OperationsStrategy": layer_name = layers[0].__class__.__name__ if layer_name == "DeepseekV2DecoderLayer": @@ -67,7 +68,7 @@ class OperationsStrategy: return OperationsStrategy.concat( [ _compute_moe_deepseek_v4_layer_operations_strategy_tbo( - layer, forward_mode + layer, forward_mode, use_cp=use_cp ) for layer in layers ] @@ -170,9 +171,10 @@ def _compute_moe_deepseek_blog_decode(layer): def _compute_moe_deepseek_v4_layer_operations_strategy_tbo( layer: torch.nn.Module, forward_mode: ForwardMode, + use_cp: bool = False, ) -> OperationsStrategy: if forward_mode == ForwardMode.EXTEND: - return _compute_moe_deepseek_v4_prefill(layer) + return _compute_moe_deepseek_v4_prefill(layer, use_cp=use_cp) else: # Decode TBO for DSV4 is not implemented yet (ATOM data: decode TBO # regresses; needs cuda-graph capture work). Prefill-only for now. @@ -181,10 +183,28 @@ def _compute_moe_deepseek_v4_layer_operations_strategy_tbo( ) -def _compute_moe_deepseek_v4_prefill(layer): +def _compute_moe_deepseek_v4_prefill(layer, use_cp: bool = False): from sglang.srt.layers.moe import get_moe_a2a_backend - if get_moe_a2a_backend().is_none(): + if use_cp: + assert get_moe_a2a_backend().is_none(), ( + "DSA prefill CP + TBO is only wired for the non-EP TP-MoE path " + "(moe_a2a_backend == none)." + ) + ops = [ + layer.op_mhc_prepare_attn, + layer.self_attn.op_attn, + layer.op_mhc_post_attn_pre_mlp, + layer.op_cp_gather_a, + operations.YieldOperation(), + layer.op_cp_gather_b, + layer.op_cp_moe, + layer.op_cp_combine_a, + operations.YieldOperation(), + layer.op_cp_combine_b, + layer.op_mhc_postprocess, + ] + elif get_moe_a2a_backend().is_none(): # Non-EP DP TP-MoE: overlap the DP all_gatherv (gather) + reduce_scatterv # (combine) with the other ubatch's attn+MoE compute (ATOM's DSV4 path). ops = [ diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index ed13cdb2b..f26955de4 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -3,6 +3,7 @@ from __future__ import annotations import copy import dataclasses import logging +import math from dataclasses import replace from typing import TYPE_CHECKING, Dict, List, Optional, Sequence @@ -678,6 +679,12 @@ class TboForwardBatchPreparer: _tbo_padded_len = ( (end_token_index - start_token_index - 1) // attention_tp_size + 1 ) * attention_tp_size + if _is_hip: + from sglang.srt.layers.cp.padding import get_cp_padding_align_size + + align = math.lcm(attention_tp_size, get_cp_padding_align_size()) + n_tokens = end_token_index - start_token_index + _tbo_padded_len = ((n_tokens + align - 1) // align) * align output_dict["tbo_padded_len"] = _tbo_padded_len for key in [ diff --git a/python/sglang/srt/distributed/bootstrap.py b/python/sglang/srt/distributed/bootstrap.py index 173b5ee61..9927355ef 100644 --- a/python/sglang/srt/distributed/bootstrap.py +++ b/python/sglang/srt/distributed/bootstrap.py @@ -32,6 +32,7 @@ from sglang.srt.server_args import ServerArgs from sglang.srt.utils import ( cpu_has_amx_support, get_available_gpu_memory, + is_hip, is_host_cpu_arm64, is_npu, monkey_patch_p2p_access_check, @@ -235,6 +236,11 @@ def _init_parallel_groups( moe_data_model_parallel_size=moe_dp_size, decode_context_parallel_size=dcp_size, duplicate_tp_group=server_args.enable_pdmux, + duplicate_attn_cp_group=( + is_hip() + and server_args.enable_two_batch_overlap + and server_args.enable_dsa_prefill_context_parallel + ), enable_symm_mem=server_args.enable_symm_mem, recovered_rank=is_ep_joiner, rank_offset=rank_offset, diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 903fdff79..b383c5404 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -1928,6 +1928,7 @@ def init_model_parallel_group( _TP: Optional[GroupCoordinator] = None _ATTN_TP: Optional[GroupCoordinator] = None _ATTN_CP: Optional[GroupCoordinator] = None +_ATTN_CP_OVERLAP: Optional[GroupCoordinator] = None _DCP: Optional[GroupCoordinator] = None # duplicate GroupCoordinator for prefill in PD-Multiplexing @@ -1965,6 +1966,55 @@ def get_attn_cp_group() -> GroupCoordinator: return _ATTN_CP +def get_attn_cp_overlap_group() -> GroupCoordinator: + return _ATTN_CP_OVERLAP if _ATTN_CP_OVERLAP is not None else get_attn_cp_group() + + +def _init_attn_cp_overlap_group( + *, + world_size: int, + attn_cp_size: int, + attn_tp_size: int, + backend: Optional[str], + recovered_rank: bool, + rank_offset: int, + max_world_size: Optional[int], +) -> None: + """Second communicator over the attention CP ranks; RCCL deadlocks when one + communicator is driven from two streams at once.""" + global _ATTN_CP_OVERLAP + assert ( + _ATTN_CP_OVERLAP is None + ), "attention context parallel overlap group is already initialized" + if attn_cp_size <= 1: + return + + span = attn_tp_size * attn_cp_size + group_ranks = [ + list(range(base + i, base + i + span, attn_tp_size)) + for base in range(0, world_size, span) + for i in range(attn_tp_size) + ] + rank = torch.distributed.get_rank() + mine = next(ranks for ranks in group_ranks if rank in ranks) + assert mine == get_attn_cp_group().ranks, ( + f"attn_cp_overlap partition {mine} does not match attn_cp " + f"{get_attn_cp_group().ranks}; the two communicators must span the " + "same ranks or the overlapped collectives will not pair up" + ) + + _ATTN_CP_OVERLAP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=False, + group_name="attn_cp_overlap", + recovered_rank=recovered_rank, + rank_offset=rank_offset, + max_world_size=max_world_size, + ) + + def get_dcp_group_no_assert() -> Optional[GroupCoordinator]: return _DCP @@ -2292,6 +2342,7 @@ def initialize_model_parallel( decode_context_parallel_size: int = 1, backend: Optional[str] = None, duplicate_tp_group: bool = False, + duplicate_attn_cp_group: bool = False, enable_symm_mem: bool = False, recovered_rank: bool = False, rank_offset: int = 0, @@ -2486,6 +2537,17 @@ def initialize_model_parallel( max_world_size=max_world_size, ) + if duplicate_attn_cp_group and is_hip(): + _init_attn_cp_overlap_group( + world_size=world_size, + attn_cp_size=attn_cp_size, + attn_tp_size=attn_tp_size, + backend=backend, + recovered_rank=recovered_rank, + rank_offset=rank_offset, + max_world_size=max_world_size, + ) + from sglang.srt.layers.sampler import SYNC_TOKEN_IDS_ACROSS_TP global _ATTN_TP @@ -2887,12 +2949,16 @@ def destroy_model_parallel(): _MOE_TP = None global _ATTN_CP + global _ATTN_CP_OVERLAP global _MOE_DP # Destroy _MOE_DP before _ATTN_CP since it may alias _ATTN_CP. # Only destroy if not aliasing another group. if _MOE_DP and _MOE_DP is not _ATTN_CP and _MOE_DP is not _TP: _MOE_DP.destroy() _MOE_DP = None + if _ATTN_CP_OVERLAP: + _ATTN_CP_OVERLAP.destroy() + _ATTN_CP_OVERLAP = None if _ATTN_CP: _ATTN_CP.destroy() _ATTN_CP = None diff --git a/python/sglang/srt/layers/attention/dsv4/compressor.py b/python/sglang/srt/layers/attention/dsv4/compressor.py index ac041d599..fa45ea50b 100644 --- a/python/sglang/srt/layers/attention/dsv4/compressor.py +++ b/python/sglang/srt/layers/attention/dsv4/compressor.py @@ -26,12 +26,17 @@ from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp from sglang.srt.layers.cp.utils import cp_materialize_global_token_order from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ReplicatedLinear +from sglang.srt.layers.utils.cp_utils import ( + cp_all_gather_rerange_finish, + cp_all_gather_rerange_launch, +) from sglang.srt.mem_cache.deepseek_v4_compress_state import ( CompressStatePool, ) from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.model_executor.forward_context import get_attn_backend from sglang.srt.models.deepseek_v2 import _is_hip +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import add_prefix, is_npu, set_weight_attrs _is_npu = is_npu() @@ -420,7 +425,38 @@ class Compressor(BaseFusedOp): assert isinstance(ret, CompressStatePool) return ret + def _pending_key(self): + return ("kv_score", self.layer_id, self.is_in_indexer) + + def prelaunch_kv_score(self, x: torch.Tensor, forward_batch: ForwardBatch): + """Compute kv_score and start its CP all-gather, without waiting. + + kv_score only needs `x`, which the attention already has at entry, so the + gather can be issued before the q/kv projections and collected later in + compute_kv_score -- that projection work is what hides it. Caller must + guarantee a matching compute_kv_score in the same op (see + DeepseekV4Attention._forward_prepare). + """ + if not _is_hip: + return + comm_stream = getattr(forward_batch, "_cp_prefetch_comm_stream", None) + if comm_stream is None or not dsa_use_prefill_cp(forward_batch): + return + kv_score = linear_bf16_fp32(x, self.wkv_gate.weight) + # Keyed by forward_batch: each TBO ubatch carries its own, so the two + # ubatches cannot collect each other's gather. + pending = forward_batch.__dict__.setdefault("_cp_pending_gathers", {}) + pending[self._pending_key()] = cp_all_gather_rerange_launch( + kv_score, get_parallel().attn_cp_size, comm_stream, self._pending_key() + ) + def compute_kv_score(self, x: torch.Tensor, forward_batch: ForwardBatch): + if _is_hip: + pending = getattr(forward_batch, "_cp_pending_gathers", None) + handle = pending.pop(self._pending_key(), None) if pending else None + if handle is not None: + return cp_all_gather_rerange_finish(handle) + kv_score = linear_bf16_fp32(x, self.wkv_gate.weight) # CUDA path: delegate to backend diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index a89a839d0..355ed42e5 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -13,6 +13,7 @@ import triton.language as tl from sglang.srt.distributed import ( GroupCoordinator, get_attn_cp_group, + get_attn_cp_overlap_group, get_attn_tensor_model_parallel_rank, get_attn_tensor_model_parallel_world_size, get_attn_tp_group, @@ -969,6 +970,14 @@ def attn_cp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor): return get_attn_cp_group().all_gather_into_tensor(output, input) +def attn_cp_overlap_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor): + return get_attn_cp_overlap_group().all_gather_into_tensor(output, input) + + +def attn_cp_overlap_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor): + return get_attn_cp_overlap_group().reduce_scatter_tensor(output, input) + + def get_moe_cp_group() -> GroupCoordinator: """Returns the MOE_DP group, which includes CP partners when attn_cp_size > moe_dp_size.""" return _get_moe_dp_group() diff --git a/python/sglang/srt/layers/utils/cp_utils.py b/python/sglang/srt/layers/utils/cp_utils.py index 08abfcd3e..2a692eeb4 100644 --- a/python/sglang/srt/layers/utils/cp_utils.py +++ b/python/sglang/srt/layers/utils/cp_utils.py @@ -9,7 +9,9 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) from sglang.srt.layers.dp_attention import ( + _tbo_event, attn_cp_all_gather_into_tensor, + attn_cp_overlap_all_gather_into_tensor, is_allocation_symmetric, ) from sglang.srt.layers.moe import get_moe_a2a_backend @@ -286,6 +288,54 @@ def cp_all_gather_reorganized_into_tensor_kv_cache( return outputs +def cp_all_gather_rerange_launch(input_tensor, cp_size, comm_stream, event_key): + """Start a round-robin CP all-gather on `comm_stream`; do NOT wait for it. + + Pair with cp_all_gather_rerange_finish(). Splitting launch from wait is the + only way an attention-side CP gather can overlap anything: the collectives + inside op_attn are consumed a few statements later, so issuing and waiting + at the same point just moves the queue (measured in perf_sweep_report ยง4.6). + + The handle keeps both buffers alive until finish(); without that reference + the allocator can hand the input block back to the compute stream before the + comm-stream kernel has read it. + """ + from sglang.srt.distributed.parallel_state import ( + get_attn_cp_group, + get_attn_cp_overlap_group, + ) + + group = get_attn_cp_overlap_group() + assert group is not get_attn_cp_group(), ( + "the comm-stream path needs the duplicate attn_cp_overlap communicator; " + "driving one communicator from two streams deadlocks RCCL" + ) + + input_tensor = input_tensor.contiguous() + with use_symmetric_memory(group, disabled=not is_allocation_symmetric()): + output_tensor = input_tensor.new_empty( + (input_tensor.shape[0] * cp_size, *input_tensor.shape[1:]), + ) + comm_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(comm_stream): + attn_cp_overlap_all_gather_into_tensor(output_tensor, input_tensor) + event = _tbo_event(event_key) + event.record(comm_stream) + return (output_tensor, input_tensor, event, cp_size) + + +def cp_all_gather_rerange_finish(handle): + """Wait for a launched gather on the current stream, then rerange.""" + output_tensor, _keepalive, event, cp_size = handle + torch.cuda.current_stream().wait_event(event) + out_shape = output_tensor.shape + return ( + output_tensor.view(cp_size, -1, *out_shape[1:]) + .transpose(0, 1) + .reshape(out_shape) + ) + + def cp_all_gather_rerange_output(input_tensor, cp_size, forward_batch, stream): """ # for in-seq-split diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index b5e131e6a..a557671b3 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -66,6 +66,8 @@ from sglang.srt.layers.cp.utils import ( ) from sglang.srt.layers.dp_attention import ( _tbo_event, + attn_cp_overlap_all_gather_into_tensor, + attn_cp_overlap_reduce_scatter_tensor, attn_tp_all_gather, attn_tp_all_reduce, dp_gather_partial, @@ -96,6 +98,8 @@ from sglang.srt.layers.quantization.fp8_utils import ( from sglang.srt.layers.rotary_embedding import get_rope_wrapper from sglang.srt.layers.utils import PPMissingLayer, get_layer_id from sglang.srt.layers.utils.cp_utils import ( + cp_all_gather_rerange_finish, + cp_all_gather_rerange_launch, cp_all_gather_rerange_output, cp_round_robin_input_ids, cp_split_and_rebuild_data, @@ -1076,6 +1080,15 @@ class MQALayer(MqaAttentionBase): x_quant=None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: x_linear = x_quant if x_quant is not None else x + # kv_score depends only on x, so its CP all-gather can start before the + # projections and be collected inside forward_core_compressor below -- + # the projections are what hides it. No-op unless the CP+TBO path armed + # _cp_prefetch_comm_stream. + if _is_hip and self.compressor is not None: + self.compressor.prelaunch_kv_score(x, forward_batch) + if self.indexer is not None: + self.indexer.compressor.prelaunch_kv_score(x, forward_batch) + if self.fuse_wqa_wkv: qkv_a, _ = self.wqkv_a(x_linear) q_lora = qkv_a[..., : self.q_lora_rank] @@ -1085,6 +1098,7 @@ class MQALayer(MqaAttentionBase): use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch) kv: Optional[torch.Tensor] + kv_handle = None from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, @@ -1207,11 +1221,22 @@ class MQALayer(MqaAttentionBase): # unified_kv + DSA CP: the 2-source prefill path needs the # FULL current-chunk KV (extend source + ring write), so # all-gather the per-rank bf16 KV across the CP group. - kv = cp_materialize_global_token_order( - kv.contiguous(), - forward_batch, - torch.cuda.current_stream(), + comm_stream = getattr( + forward_batch, "_cp_prefetch_comm_stream", None ) + if comm_stream is not None: + # kv is not read again until this function returns, so the + # indexer + compressor below can run while it gathers. + kv_handle = cp_all_gather_rerange_launch( + kv, self.cp_size, comm_stream, ("kv", self.layer_id) + ) + kv = None + else: + kv = cp_materialize_global_token_order( + kv.contiguous(), + forward_batch, + torch.cuda.current_stream(), + ) elif use_cp: # NSA CP: keep bf16 kv around for the cross-rank all-gather, then # write to the FlashMLA cache after gather. @@ -1249,6 +1274,9 @@ class MQALayer(MqaAttentionBase): self.compressor, ) + if _is_hip and kv_handle is not None: + kv = cp_all_gather_rerange_finish(kv_handle) + return q, kv def forward( @@ -2230,6 +2258,67 @@ class DeepseekV4DecoderLayer(nn.Module): hidden = hidden + shared_local[:n] state.hidden_states_mlp_output = hidden + def _cp_tbo_launch(self, state, x, key, out_rows, collective): + assert _is_hip, "CP+TBO MoE overlap is HIP-only" + x = x.contiguous() + sub = state.tbo_subbatch_index + out = get_tbo_persistent_buffer( + (key, sub), out_rows, x.shape[1], x.dtype, x.device + ) + comm = get_dp_tbo_comm_stream() + comm.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(comm): + collective(out, x) + event = _tbo_event((key, sub)) + event.record(comm) + return out, event, x + + def op_cp_gather_a(self, state): + local = state.pop("hidden_states_mlp_input") + out, event, keepalive = self._cp_tbo_launch( + state, + local, + "cpgh", + local.shape[0] * get_parallel().attn_cp_size, + attn_cp_overlap_all_gather_into_tensor, + ) + state.global_hidden = out + state.cp_gather_event = event + state.cp_gather_keepalive = keepalive + + def op_cp_gather_b(self, state): + torch.cuda.current_stream().wait_event(state.pop("cp_gather_event")) + state.pop("cp_gather_keepalive") + + def op_cp_moe(self, state): + fb = state.forward_batch + global_ids = fb._cp_moe_input_ids + with get_forward().scoped(mlp_reduce_scatter=True): + state.global_expert_out = self.mlp( + state.pop("global_hidden"), + fb, + input_ids=global_ids, + input_ids_global=global_ids, + ) + + def op_cp_combine_a(self, state): + global_out = state.pop("global_expert_out") + out, event, keepalive = self._cp_tbo_launch( + state, + global_out, + "cplo", + global_out.shape[0] // get_parallel().attn_cp_size, + attn_cp_overlap_reduce_scatter_tensor, + ) + state.local_out = out + state.cp_combine_event = event + state.cp_combine_keepalive = keepalive + + def op_cp_combine_b(self, state): + torch.cuda.current_stream().wait_event(state.pop("cp_combine_event")) + state.pop("cp_combine_keepalive") + state.hidden_states_mlp_output = state.pop("local_out") + class DeepseekV4Model(nn.Module): fall_back_to_pt_during_load = False @@ -2335,17 +2424,44 @@ class DeepseekV4Model(nn.Module): hc_eps=self.hc_eps, ) + def _cp_children_splittable(self, forward_batch: ForwardBatch) -> bool: + children = forward_batch.tbo_children + if not children: + return False + cp_size = get_parallel().attn_cp_size + for child in children: + if child.batch_size <= 0 or child.extend_seq_lens_cpu is None: + return False + if sum(child.extend_seq_lens_cpu) < cp_size: + return False + return True + def _can_run_tbo(self, forward_batch: ForwardBatch) -> bool: """DSV4 prefill-only two-batch-overlap gate. TBO batch prep (tbo_split_seq_index / tbo_children) is populated model-agnostically when --enable-two-batch-overlap is set and the DP-attention preparer allows it (mori `normal` mode permits prefill - TBO). We additionally restrict to: prefill (EXTEND), single PP, and the - non-CP path, which is the only case the DSV4 op strategy implements. + TBO). We additionally restrict to: prefill (EXTEND), single PP, and a + path the DSV4 op strategy implements -- the non-CP path everywhere, plus + the round-robin DSA prefill CP path on HIP. """ from sglang.srt.layers.moe import is_tbo_enabled + if dsa_use_prefill_cp(forward_batch): + path_ok = ( + _is_hip + and not is_cp_v2_active(forward_batch) + and is_dsa_prefill_cp_round_robin_split() + and get_moe_a2a_backend().is_none() + and self._cp_children_splittable(forward_batch) + ) + else: + path_ok = ( + not _is_hip + or not get_moe_a2a_backend().is_none() + or get_parallel().attn_dp_size > 1 + ) return ( is_tbo_enabled() and forward_batch.can_run_tbo @@ -2354,7 +2470,7 @@ class DeepseekV4Model(nn.Module): # MTP target-verify also reports is_extend(); only real prefill # should enter the prefill TBO strategy. and forward_batch.global_forward_mode.is_extend_without_speculative() - and not dsa_use_prefill_cp(forward_batch) + and path_ok and self.pp_group.world_size == 1 ) @@ -2371,6 +2487,13 @@ class DeepseekV4Model(nn.Module): _model_forward_tbo_merge_outputs, ) + if _is_hip and dsa_use_prefill_cp(forward_batch): + return self._forward_layers_tbo_cp( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + layers = [self.layers[i] for i in range(self.start_layer, self.end_layer)] operations_strategy = OperationsStrategy.init_new_tbo( layers, forward_batch.global_forward_mode @@ -2447,6 +2570,97 @@ class DeepseekV4Model(nn.Module): ) return hidden_states + def _setup_child_cp_metadata(self, child: ForwardBatch, child_backend) -> None: + cp_rank = get_parallel().attn_cp_rank + cp_size = get_parallel().attn_cp_size + child.attn_cp_metadata = prepare_context_parallel_metadata( + len(child.input_ids), + cp_rank, + cp_size, + child.seq_lens_cpu.tolist(), + extend_seqs_len=child.extend_seq_lens_cpu, + ) + if is_dsa_prefill_cp_round_robin_split(): + metadata = child_backend.forward_metadata + core_meta = metadata.core_attn_metadata + core_meta.apply_cp_reindex() + core_meta.init_flashmla_related(is_prefill=True) + if metadata.indexer_metadata is not None: + metadata.indexer_metadata = child_backend.init_forward_metadata_indexer( + core_meta + ) + + def _forward_layers_tbo_cp( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + assert _is_hip, "CP+TBO prefill path is HIP-only" + + from sglang.srt.batch_overlap.operations import execute_overlapped_operations + from sglang.srt.batch_overlap.operations_strategy import OperationsStrategy + from sglang.srt.batch_overlap.two_batch_overlap import ( + _model_forward_filter_inputs, + _model_forward_tbo_merge_outputs, + ) + + original_len = hidden_states.shape[0] + cp_size = get_parallel().attn_cp_size + layers = [self.layers[i] for i in range(self.start_layer, self.end_layer)] + operations_strategy = OperationsStrategy.init_new_tbo( + layers, forward_batch.global_forward_mode, use_cp=True + ) + + attn_backend = get_attn_backend() + children = forward_batch.tbo_children + # Attention-side CP gathers run two-phase (launch early on the comm + # stream / collect right before their consumer). Only the MoE + # collectives are splittable across a YieldOperation, so without this the + # ~2.5 attention-side collectives per layer would stay on the compute + # stream and defeat most of TBO's overlap. + prefetch_comm_stream = get_dp_tbo_comm_stream() + + inputs_arr = [] + for idx, child in enumerate(children): + child_inputs = _model_forward_filter_inputs( + hidden_states=hidden_states, + residual=None, + positions=positions, + output_forward_batch=child, + tbo_subbatch_index=idx, + ) + self._setup_child_cp_metadata(child, attn_backend.children[idx]) + if self.pp_group.is_first_rank: + child_inputs["hidden_states"] = cp_split_and_rebuild_data( + child, child_inputs["hidden_states"] + ) + child_inputs["positions"] = cp_split_and_rebuild_position( + child, child_inputs["positions"] + ) + child._cp_moe_input_ids = cp_round_robin_input_ids(child.input_ids) + child._cp_prefetch_comm_stream = prefetch_comm_stream + inputs_arr.append(child_inputs) + + outputs_arr = execute_overlapped_operations( + inputs_arr=inputs_arr, + operations_arr=[operations_strategy.operations] * 2, + delta_stages=[0, operations_strategy.tbo_delta_stages], + ) + + if self.pp_group.is_last_rank: + for idx, child in enumerate(children): + outputs_arr[idx]["hidden_states"] = cp_all_gather_rerange_output( + outputs_arr[idx]["hidden_states"], + cp_size, + child, + torch.cuda.current_stream(), + ) + hidden_states, _ = _model_forward_tbo_merge_outputs( + outputs_arr[0], outputs_arr[1], original_len + ) + return hidden_states + def forward( self, input_ids: torch.Tensor, @@ -2488,7 +2702,13 @@ class DeepseekV4Model(nn.Module): else: input_ids_global = input_ids - if use_prefill_cp: + capture_dspark = self.dspark_layers_to_capture is not None + dspark_aux_hidden_states: List[torch.Tensor] = [] + # DSpark aux capture needs the per-layer eager loop (TBO's overlapped + # execution cannot expose per-layer completed hidden states), so skip + # TBO when capturing -- a perf-only downgrade, not a correctness one. + run_tbo = self._can_run_tbo(forward_batch) and not capture_dspark + if use_prefill_cp and not run_tbo: if cp_v2_active: input_ids = cp_round_robin_input_ids_v2(input_ids, forward_batch) else: @@ -2504,12 +2724,7 @@ class DeepseekV4Model(nn.Module): for _attr in ("freqs_cis_c4", "freqs_cis_c128"): if hasattr(forward_batch, _attr): delattr(forward_batch, _attr) - capture_dspark = self.dspark_layers_to_capture is not None - dspark_aux_hidden_states: List[torch.Tensor] = [] - # DSpark aux capture needs the per-layer eager loop (TBO's overlapped - # execution cannot expose per-layer completed hidden states), so skip - # TBO when capturing -- a perf-only downgrade, not a correctness one. - if self._can_run_tbo(forward_batch) and not capture_dspark: + if run_tbo: # Two-batch-overlap prefill (EP / mori). Cross-layer mHC fusion is # disabled here (each layer self-contained), so no trailing hc_post. hidden_states = self._forward_layers_tbo( @@ -2554,7 +2769,12 @@ class DeepseekV4Model(nn.Module): ) # CP all-gather only on the last PP rank; PP IPC carries CP-split tensors. - if self.pp_group.is_last_rank and use_prefill_cp and not cp_v2_active: + if ( + self.pp_group.is_last_rank + and use_prefill_cp + and not cp_v2_active + and not run_tbo + ): stream = torch.cuda.current_stream() hidden_states = cp_all_gather_rerange_output( hidden_states, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index f48d478d5..589db0b1b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -9049,10 +9049,16 @@ class ServerArgs: # DP TP-MoE path (overlapping the DP all_gatherv / reduce_scatterv with # the other ubatch's compute), which requires DP attention. Enabling it # there needs no extra opt-in env flag. + cp_tbo = ( + is_hip() + and self.enable_dsa_prefill_context_parallel + and self.dsa_prefill_cp_mode == "round-robin-split" + ) if ( self.enable_two_batch_overlap and self.moe_a2a_backend == "none" and not self.enable_dp_attention + and not cp_tbo ): raise ValueError( "When enabling two batch overlap without an EP a2a backend " diff --git a/test/registered/amd/test_deepseek_v4_pro_fp4_cp_tbo.py b/test/registered/amd/test_deepseek_v4_pro_fp4_cp_tbo.py new file mode 100644 index 000000000..24655ced6 --- /dev/null +++ b/test/registered/amd/test_deepseek_v4_pro_fp4_cp_tbo.py @@ -0,0 +1,155 @@ +"""MI35x DeepSeek-V4-Pro FP4 prefill context-parallel (CP) + two-batch-overlap (TBO) +accuracy test (8-GPU). + +Same launch conventions as test_deepseek_v4_pro_fp4_cp.py (prefill CP over the +unified_kv backend via ``--enable-prefill-cp --cp-strategy interleave``), plus +``--enable-two-batch-overlap``. This exercises the CP TBO op strategy +(``op_cp_gather`` / ``op_cp_moe`` / ``op_cp_combine`` driven by +``DeepseekV4Model._forward_layers_tbo_cp``), which splits each prefill batch into +two token-range ubatches, round-robin splits each one across the CP group +independently, and overlaps one ubatch's CP MoE all-gather / reduce-scatter with +the other ubatch's attention + expert compute. + +The overlap runs on a duplicate CP communicator (``attn_cp_overlap``) so it can +execute concurrently with the attention-internal CP all-gathers, which stay on +the compute stream and the primary CP communicator. This test guards that the +combination stays numerically equivalent to CP-only (>0.92 on GSM8K, same bar as +test_deepseek_v4_pro_fp4_cp.py) and that neither the per-ubatch CP metadata setup +nor the concurrent CP communicators deadlock or corrupt state. + +Registry: nightly-amd-8-gpu-mi35x-deepseek-v4-pro suite +""" + +import os +import unittest +from types import SimpleNamespace + +from sglang.srt.utils import kill_process_tree, set_ulimit +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + CustomTestCase, + is_in_ci, + popen_launch_server, + write_github_step_summary, +) + +register_amd_ci( + est_time=5400, suite="nightly-amd-8-gpu-mi35x-deepseek-v4-pro", nightly=True +) + +DEEPSEEK_V4_PRO_FP4_MODEL_PATH = os.environ.get( + "DEEPSEEK_V4_PRO_MODEL_PATH_FP4", "deepseek-ai/DeepSeek-V4-Pro" +) +# Pro is 1.6T; weight load + warmup is much longer than Flash 285B. +SERVER_LAUNCH_TIMEOUT = 5400 + +# Matches test_deepseek_v4_pro_fp4_cp.py; prefill CP requires unified_kv_triton. +COMMON_ENV_VARS = { + "SGLANG_DEFAULT_THINKING": "1", + "SGLANG_DSV4_REASONING_EFFORT": "max", + "SGLANG_USE_ROCM700A": "0", + "SGLANG_DP_USE_GATHERV": "1", + "SGLANG_HACK_FLASHMLA_BACKEND": "unified_kv_triton", + "AITER_BF16_FP8_MOE_BOUND": "0", + # ROCm HSA-resource stability for TBO at high concurrency. + "GPU_MAX_HW_QUEUES": "5", +} + +# FP4 variant (matches test_deepseek_v4_pro_fp4.py; V4-Pro also auto-detects it). +FP4_ENV_VARS = { + "SGLANG_DSV4_FP4_EXPERTS": "true", +} + + +class TestDeepseekV4ProFp4CPInterleaveTbo(CustomTestCase): + """DeepSeek-V4-Pro FP4 unified_kv prefill CP (round-robin-split) + TBO, tp=8.""" + + @classmethod + def setUpClass(cls): + cls.model = DEEPSEEK_V4_PRO_FP4_MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + + # GSM8K below drives 1319 concurrent requests, and the launched server + # inherits the fd limit of this process, so raise it before popen. + set_ulimit(65536) + + env = os.environ.copy() + env.update(COMMON_ENV_VARS) + env.update(FP4_ENV_VARS) + + other_args = [ + "--trust-remote-code", + "--tp", + "8", + "--dp", + "1", + "--enable-prefill-cp", + "--cp-strategy", + "interleave", + "--enable-two-batch-overlap", + "--disable-radix-cache", + "--attention-backend", + "dsv4", + "--max-running-requests", + "256", + "--page-size", + "256", + "--mem-fraction-static", + "0.90", + "--swa-full-tokens-ratio", + "0.1", + # TBO halves the per-ubatch MoE rows, so it only pays off once the + # chunk is large; 32768 also keeps num_q_tokens under the compress + # prefill plan's uint16 token limit for a single full chunk. + "--chunked-prefill-size", + "32768", + "--disable-shared-experts-fusion", + "--tool-call-parser", + "deepseekv4", + "--reasoning-parser", + "deepseek-v4", + ] + + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=other_args, + env=env, + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + def test_a_gsm8k( + self, + ): # Append an "a" to make this test run first (alphabetically) to warm up the server + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + api="completion", + max_tokens=512, + num_examples=1319, + num_threads=1319, + num_shots=5, + ) + metrics = run_eval(args) + print(f"{metrics=}") + + if is_in_ci(): + write_github_step_summary( + f"### test_a_gsm8k (deepseek-v4-pro-fp4-cp-interleave-tbo)\n" + f'{metrics["score"]=:.3f}\n' + ) + # CP + TBO must stay numerically equivalent to CP-only (same bar as + # test_deepseek_v4_pro_fp4_cp.py). + self.assertGreater(metrics["score"], 0.92) + + +if __name__ == "__main__": + unittest.main()