From 35e25f53567c3290f3d96aebc21019b87ced710f Mon Sep 17 00:00:00 2001 From: Thanhhao <31717833+thanhhao98@users.noreply.github.com> Date: Fri, 24 Jul 2026 14:21:53 +0700 Subject: [PATCH] [Feature] DCP: A2A + FlashInfer-MNNVL comm backends and q-replicate (Helix) (#21637) Co-authored-by: Hao Phan Co-authored-by: Claude Opus 4.8 Co-authored-by: kpham-sgl Co-authored-by: Cursor Co-authored-by: Baizhou Zhang --- .../kernels/ops/attention/dcp_kernels.py | 197 +++++++- .../device_communicators/pynccl.py | 54 +++ .../sglang/srt/distributed/parallel_state.py | 7 +- python/sglang/srt/layers/activation.py | 49 +- python/sglang/srt/layers/dcp/__init__.py | 4 + python/sglang/srt/layers/dcp/comm.py | 248 +++++++++- .../sglang/srt/model_executor/model_runner.py | 44 ++ .../srt/model_executor/runner/base_runner.py | 14 + .../attention_forward_methods/forward_mla.py | 71 ++- python/sglang/srt/models/deepseek_v2.py | 5 + python/sglang/srt/server_args.py | 47 ++ .../kernels/test_dcp_lse_combine.py | 452 ++++++++++++++++++ .../unit/server_args/test_dcp_config.py | 102 ++++ 13 files changed, 1276 insertions(+), 18 deletions(-) create mode 100644 test/registered/kernels/test_dcp_lse_combine.py create mode 100644 test/registered/unit/server_args/test_dcp_config.py diff --git a/python/sglang/kernels/ops/attention/dcp_kernels.py b/python/sglang/kernels/ops/attention/dcp_kernels.py index 81dd49391..395ec28de 100644 --- a/python/sglang/kernels/ops/attention/dcp_kernels.py +++ b/python/sglang/kernels/ops/attention/dcp_kernels.py @@ -20,7 +20,7 @@ Consolidated from the two merged DCP implementations: - _correct_attn_cp_out_kernel / correct_attn_out / CPTritonContext (PR #14194) """ -from typing import Optional +from typing import Optional, Tuple import torch import triton @@ -331,3 +331,198 @@ def correct_attn_out( ctx.call_kernel(_correct_attn_cp_out_kernel, grid, *regular_args, **const_args) return new_output, lse + + +# A2A DCP reduce: LSE-weighted combine of N partial attention outputs +# (used by the a2a / fi_a2a communication backends, see comm.py). + + +def _lse_pack_dim(output_dtype: torch.dtype) -> int: + """Number of output-dtype elements needed to store one fp32 LSE value.""" + return torch.finfo(torch.float32).bits // torch.finfo(output_dtype).bits + + +@triton.jit +def _dcp_lse_combine_kernel( + recv_output_ptr, + recv_lse_ptr, + out_ptr, + out_lse_ptr, + recv_output_stride_N, + recv_output_stride_B, + recv_output_stride_H, + recv_output_stride_D, + recv_lse_stride_N, + recv_lse_stride_B, + recv_lse_stride_H, + out_stride_B, + out_stride_H, + out_stride_D, + N: tl.constexpr, + HEAD_DIM: tl.constexpr, + IS_BASE_E: tl.constexpr, + RETURN_LSE: tl.constexpr, +): + """Combine N partial attention outputs weighted by their LSE values. + + Grid: (B, H_local). + Each program handles one (batch, head) position across all N shards. + + Two-pass approach: + Pass 1: find max LSE and weight sum across shards + Pass 2: accumulate weighted outputs + """ + batch_idx = tl.program_id(0).to(tl.int64) + head_idx = tl.program_id(1).to(tl.int64) + d_offsets = tl.arange(0, HEAD_DIM) + + lse_base = batch_idx * recv_lse_stride_B + head_idx * recv_lse_stride_H + + # Pass 1: find max LSE across N shards + lse_max = tl.load(recv_lse_ptr + lse_base).to(tl.float32) + lse_max = tl.where( + (lse_max != lse_max) | (lse_max == float("inf")), -float("inf"), lse_max + ) + for i in tl.static_range(1, N): + lse_i = tl.load(recv_lse_ptr + lse_base + i * recv_lse_stride_N).to(tl.float32) + lse_i = tl.where( + (lse_i != lse_i) | (lse_i == float("inf")), -float("inf"), lse_i + ) + lse_max = tl.where(lse_i > lse_max, lse_i, lse_max) + + lse_max = tl.where(lse_max == -float("inf"), 0.0, lse_max) + + # Pass 2: accumulate weighted outputs + weight_sum = tl.zeros([], dtype=tl.float32) + acc = tl.zeros([HEAD_DIM], dtype=tl.float32) + + for i in tl.static_range(N): + lse_i = tl.load(recv_lse_ptr + lse_base + i * recv_lse_stride_N).to(tl.float32) + lse_i = tl.where( + (lse_i != lse_i) | (lse_i == float("inf")), -float("inf"), lse_i + ) + centered = lse_i - lse_max + if IS_BASE_E: + w = tl.exp(centered) + else: + w = tl.exp2(centered) + weight_sum += w + + o_offsets = ( + i * recv_output_stride_N + + batch_idx * recv_output_stride_B + + head_idx * recv_output_stride_H + + d_offsets * recv_output_stride_D + ) + partial_out = tl.load(recv_output_ptr + o_offsets).to(tl.float32) + acc += partial_out * w + + acc = acc / weight_sum + + out_offsets = ( + batch_idx * out_stride_B + head_idx * out_stride_H + d_offsets * out_stride_D + ) + tl.store(out_ptr + out_offsets, acc.to(out_ptr.dtype.element_ty)) + + if RETURN_LSE: + if IS_BASE_E: + global_lse = tl.log(weight_sum) + lse_max + else: + global_lse = tl.log2(weight_sum) + lse_max + out_lse_offset = batch_idx * recv_lse_stride_B + head_idx * recv_lse_stride_H + tl.store(out_lse_ptr + out_lse_offset, global_lse) + + +def dcp_lse_combine_triton( + recv_output: torch.Tensor, + recv_lse: torch.Tensor, + is_lse_base_on_e: bool = True, + return_lse: bool = False, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Launch the Triton LSE-combine kernel. + + Args: + recv_output: [N, B, H_local, D] partial outputs from each DCP rank. + recv_lse: [N, B, H_local] log-sum-exp from each DCP rank. + is_lse_base_on_e: True if LSE uses base-e (FlashAttention), + False if base-2 (FlashInfer). + return_lse: If True, also return the combined global LSE. + + Returns: + (combined_output [B, H_local, D], combined_lse [B, H_local] or None) + """ + N, B, H_local, D = recv_output.shape + out = torch.empty( + (B, H_local, D), device=recv_output.device, dtype=recv_output.dtype + ) + out_lse = ( + torch.empty((B, H_local), device=recv_lse.device, dtype=recv_lse.dtype) + if return_lse + else recv_lse.new_empty(0) + ) + + grid = (B, H_local) + _dcp_lse_combine_kernel[grid]( + recv_output, + recv_lse, + out, + out_lse, + recv_output.stride(0), + recv_output.stride(1), + recv_output.stride(2), + recv_output.stride(3), + recv_lse.stride(0), + recv_lse.stride(1), + recv_lse.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + N=N, + HEAD_DIM=D, + IS_BASE_E=is_lse_base_on_e, + RETURN_LSE=return_lse, + ) + return out, (out_lse if return_lse else None) + + +def _lse_weighted_combine_cpu( + partial_outputs: torch.Tensor, + partial_lses: torch.Tensor, + is_lse_base_on_e: bool = True, +) -> torch.Tensor: + """CPU reference: combine N partial attention outputs using LSE weights. + + Args: + partial_outputs: [N, B, H_local, D] + partial_lses: [N, B, H_local] + is_lse_base_on_e: base-e (True) or base-2 (False) + + Returns: + [B, H_local, D] combined output + """ + N, B, H_local, D = partial_outputs.shape + partial_outputs = partial_outputs.float() + partial_lses = partial_lses.float() + + # Sanitize + partial_lses = torch.where( + torch.isnan(partial_lses) | torch.isinf(partial_lses), + torch.full_like(partial_lses, float("-inf")), + partial_lses, + ) + + # max LSE for numerical stability + lse_max, _ = partial_lses.max(dim=0) + lse_max = torch.where(lse_max == float("-inf"), torch.zeros_like(lse_max), lse_max) + + centered = partial_lses - lse_max.unsqueeze(0) + if is_lse_base_on_e: + weights = torch.exp(centered) + else: + weights = torch.pow(2.0, centered) + + weight_sum = weights.sum(dim=0, keepdim=True) + weights = weights / weight_sum + + combined = (partial_outputs * weights.unsqueeze(-1)).sum(dim=0) + return combined diff --git a/python/sglang/srt/distributed/device_communicators/pynccl.py b/python/sglang/srt/distributed/device_communicators/pynccl.py index 53eafe6d5..cb74d4d35 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl.py +++ b/python/sglang/srt/distributed/device_communicators/pynccl.py @@ -303,6 +303,60 @@ class PyNcclCommunicator: cudaStream_t(stream.cuda_stream), ) + def all_to_all_single( + self, + output_tensor: torch.Tensor, + input_tensor: torch.Tensor, + ): + """All-to-All over the flattened leading dim: each rank sends the i-th + equal-sized chunk to rank i and receives rank i's chunk into output + position i. Uses ncclGroupStart/End to fuse the sends/recvs into a + single NCCL operation, which is CUDA-graph-capturable (used by the DCP + a2a communication backend).""" + if self.disabled: + return + assert input_tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {input_tensor.device}" + ) + assert output_tensor.device == self.device, ( + f"this nccl communicator is created to work on {self.device}, " + f"but the output tensor is on {output_tensor.device}" + ) + stream = self._resolve_stream() + # Equal-split all-to-all: fail loudly instead of silently truncating the tail. + assert input_tensor.numel() == output_tensor.numel(), ( + f"all_to_all_single: input numel ({input_tensor.numel()}) != output " + f"numel ({output_tensor.numel()})" + ) + assert input_tensor.numel() % self.world_size == 0, ( + f"all_to_all_single: input numel ({input_tensor.numel()}) not " + f"divisible by world_size ({self.world_size})" + ) + chunk_size = input_tensor.numel() // self.world_size + dtype = ncclDataTypeEnum.from_torch(input_tensor.dtype) + self.nccl.ncclGroupStart() + for i in range(self.world_size): + send_buf = input_tensor.narrow(0, i * chunk_size, chunk_size) + self.nccl.ncclSend( + buffer_type(send_buf.data_ptr()), + chunk_size, + dtype, + i, + self.comm, + cudaStream_t(stream.cuda_stream), + ) + recv_buf = output_tensor.narrow(0, i * chunk_size, chunk_size) + self.nccl.ncclRecv( + buffer_type(recv_buf.data_ptr()), + chunk_size, + dtype, + i, + self.comm, + cudaStream_t(stream.cuda_stream), + ) + self.nccl.ncclGroupEnd() + def broadcast(self, tensor: torch.Tensor, src: int): if self.disabled: return diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index b5db39b72..88cdb7485 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -1072,7 +1072,12 @@ class GroupCoordinator: return True def _all_to_all_single(self, output: torch.Tensor, input: torch.Tensor) -> None: - torch.distributed.all_to_all_single(output, input, group=self.device_group) + # pynccl path keeps the a2a exchange CUDA-graph-capturable (DCP a2a backend). + pynccl_comm = self.pynccl_comm + if pynccl_comm is not None and not pynccl_comm.disabled: + pynccl_comm.all_to_all_single(output, input) + else: + torch.distributed.all_to_all_single(output, input, group=self.device_group) def all_to_all_single(self, output: torch.Tensor, input: torch.Tensor): if self.world_size == 1: diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py index 9bd538397..0e97fe5fb 100644 --- a/python/sglang/srt/layers/activation.py +++ b/python/sglang/srt/layers/activation.py @@ -57,12 +57,53 @@ _is_xpu = is_xpu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip if _is_cuda: + from sgl_kernel import gelu_and_mul as _sgl_gelu_and_mul + from sgl_kernel import gelu_tanh_and_mul as _sgl_gelu_tanh_and_mul + from sgl_kernel import silu_and_mul as _sgl_silu_and_mul + from sglang.kernels.ops.activation.activation import ( - gelu_and_mul, - gelu_tanh_and_mul, - relu2, - silu_and_mul, + gelu_and_mul as _jit_gelu_and_mul, ) + from sglang.kernels.ops.activation.activation import ( + gelu_tanh_and_mul as _jit_gelu_tanh_and_mul, + ) + from sglang.kernels.ops.activation.activation import ( + relu2, + ) + from sglang.kernels.ops.activation.activation import ( + silu_and_mul as _jit_silu_and_mul, + ) + + # The jit act-and-mul kernel requires the per-rank hidden size to be a + # multiple of the vector width (kMaxVecBytes/dtype: 32B on SM100+, else 16B -- + # RuntimeCheck "hidden size must be divisible by vector size" in + # kernels/jit/csrc/elementwise/activation.cuh). Route unaligned shapes to the + # sgl_kernel implementation (e.g. DeepSeek-V2-Lite dense 10944/8 = 1368 at + # tp8, 1368 % 16 != 0). + _jit_act_max_vec_bytes: Optional[int] = None + + def _jit_act_supported(out: torch.Tensor) -> bool: + global _jit_act_max_vec_bytes + if _jit_act_max_vec_bytes is None: + major, _ = torch.cuda.get_device_capability() + _jit_act_max_vec_bytes = 32 if major >= 10 else 16 + return out.shape[-1] % (_jit_act_max_vec_bytes // out.dtype.itemsize) == 0 + + def _act_and_mul(jit_fn, sgl_fn, input: torch.Tensor, out=None) -> torch.Tensor: + if out is None: + out = input.new_empty(*input.shape[:-1], input.shape[-1] // 2) + (jit_fn if _jit_act_supported(out) else sgl_fn)(input, out) + return out + + def silu_and_mul(input: torch.Tensor, out=None) -> torch.Tensor: + return _act_and_mul(_jit_silu_and_mul, _sgl_silu_and_mul, input, out) + + def gelu_and_mul(input: torch.Tensor, out=None) -> torch.Tensor: + return _act_and_mul(_jit_gelu_and_mul, _sgl_gelu_and_mul, input, out) + + def gelu_tanh_and_mul(input: torch.Tensor, out=None) -> torch.Tensor: + return _act_and_mul(_jit_gelu_tanh_and_mul, _sgl_gelu_tanh_and_mul, input, out) + elif _is_xpu: from sgl_kernel import gelu_and_mul, gelu_tanh_and_mul, silu_and_mul elif _is_hip: diff --git a/python/sglang/srt/layers/dcp/__init__.py b/python/sglang/srt/layers/dcp/__init__.py index eb6aeea65..54ae19195 100644 --- a/python/sglang/srt/layers/dcp/__init__.py +++ b/python/sglang/srt/layers/dcp/__init__.py @@ -40,9 +40,11 @@ from sglang.srt.layers.dcp.comm import ( all_gather_q_for_mla_decode, cp_lse_ag_out_rs_mha, cp_lse_ag_out_rs_mla, + dcp_a2a_lse_reduce, dcp_enabled, get_attention_dcp_rank, get_attention_dcp_world_size, + init_fi_a2a_workspace, ) from sglang.srt.layers.dcp.layout import ( filter_dcp_local_kv_indices, @@ -61,6 +63,8 @@ from sglang.srt.layers.dcp.metadata import DecodeContextParallelMetadata __all__ = [ "DecodeContextParallelMetadata", + "dcp_a2a_lse_reduce", + "init_fi_a2a_workspace", "all_gather_kv_cache_for_dcp", "all_gather_kv_cache_for_mha_chunk_extend", "all_gather_kv_cache_for_mha_extend", diff --git a/python/sglang/srt/layers/dcp/comm.py b/python/sglang/srt/layers/dcp/comm.py index 577950834..518801709 100644 --- a/python/sglang/srt/layers/dcp/comm.py +++ b/python/sglang/srt/layers/dcp/comm.py @@ -25,7 +25,12 @@ from typing import Optional import torch -from sglang.kernels.ops.attention.dcp_kernels import CPTritonContext, correct_attn_out +from sglang.kernels.ops.attention.dcp_kernels import ( + CPTritonContext, + _lse_pack_dim, + correct_attn_out, + dcp_lse_combine_triton, +) from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) @@ -141,7 +146,14 @@ def _all_gather_dcp_kv_cache(kv_a: torch.Tensor): gathered_kv_a = kv_a.new_empty( (kv_a.shape[0] * dcp_world_size, *kv_a.shape[1:]), ) - parallel.dcp_group.all_gather_into_tensor(gathered_kv_a, kv_a) + # pynccl has no fp8 dtype; all-gather is a byte copy, so transport an fp8 KV + # cache as raw bytes via a uint8 view (works with --kv-cache-dtype fp8_*). + if kv_a.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + parallel.dcp_group.all_gather_into_tensor( + gathered_kv_a.view(torch.uint8), kv_a.contiguous().view(torch.uint8) + ) + else: + parallel.dcp_group.all_gather_into_tensor(gathered_kv_a, kv_a) gathered_kv_a = ( gathered_kv_a.reshape((dcp_world_size,) + kv_a.shape) .transpose(0, 1) @@ -193,6 +205,12 @@ def all_gather_kv_cache_for_mha_extend( [kv_a.shape[-1], k_pe.shape[-1]], dim=-1 ) prefix_kv_a = prefix_kv_a.squeeze(1) + # torch.cat can't promote fp8 (gathered prefix) + bf16 (current extend), so + # align dtypes first (dequant the fp8 prefix; exact for the scale=1.0 default). + if prefix_kv_a.dtype != kv_a.dtype: + prefix_kv_a = prefix_kv_a.to(kv_a.dtype) + if prefix_k_pe.dtype != k_pe.dtype: + prefix_k_pe = prefix_k_pe.to(k_pe.dtype) # re-organize kv with query orders prefix_lens_cu = torch.zeros( len(seq_lens) + 1, @@ -342,3 +360,229 @@ def all_gather_kv_cache_for_dcp( gatherd_kv_cache = torch.cat(kv_cache_tuple, dim=0) return gatherd_kv_cache + + +# --------------------------------------------------------------------------- +# A2A communication backend for DCP decode (alternative to AG+RS above): exchange +# per-head partial outputs + LSEs across DCP ranks, then combine locally with the +# Triton LSE kernel. fi_a2a delegates the exchange to FlashInfer MNNVL (#2951). +# --------------------------------------------------------------------------- + +# Per-process singleton: MNNVL workspace + this rank's cp position. Populated +# once, pre-CUDA-graph-capture, by init_fi_a2a_workspace(). +_FI_A2A_STATE: Optional[dict] = None + + +def init_fi_a2a_workspace(cp_group: "GroupCoordinator") -> None: + # Call once per process BEFORE CUDA-graph capture: the FlashInfer init syncs + # the stream and barriers cross-rank, neither of which is capturable. + global _FI_A2A_STATE + if _FI_A2A_STATE is not None: + return + if cp_group.world_size == 1: + return + + import torch.distributed as dist + + try: + from flashinfer.comm.dcp_alltoall import ( + decode_cp_a2a_allocate_mnnvl_workspace, + decode_cp_a2a_init_workspace, + ) + from flashinfer.comm.mapping import Mapping + from flashinfer.comm.mnnvl import MnnvlConfig, is_mnnvl_fabric_supported + except ImportError as e: + raise ImportError( + "--dcp-comm-backend fi_a2a requires FlashInfer with the DCP " + "all-to-all kernel (flashinfer #2951); could not import " + "flashinfer.comm.dcp_alltoall." + ) from e + + # Reuse the MoE adapter: its Split() returns a CommBackend (what FlashInfer's + # Mapping expects); the flashinfer_comm_fusion copy has drifted to return a + # raw ProcessGroup, so don't swap without re-checking the Split() contract. + from sglang.srt.layers.moe.token_dispatcher.flashinfer_utils import ( + TorchDistributedCommBackend, + ) + + if not is_mnnvl_fabric_supported(torch.cuda.current_device()): + raise RuntimeError( + "--dcp-comm-backend fi_a2a requires MNNVL fabric memory (e.g. " + "GB200 NVL72); is_mnnvl_fabric_supported() returned False. Use " + "--dcp-comm-backend a2a or ag_rs on clusters without MNNVL." + ) + + cp_size = cp_group.world_size + cp_rank = cp_group.rank_in_group + mapping = Mapping( + world_size=cp_size, + rank=cp_rank, + gpus_per_node=torch.cuda.device_count(), + cp_size=cp_size, + tp_size=1, + pp_size=1, + ) + workspace = decode_cp_a2a_allocate_mnnvl_workspace( + mapping, + mnnvl_config=MnnvlConfig( + comm_backend=TorchDistributedCommBackend(cp_group.device_group) + ), + ) + decode_cp_a2a_init_workspace(workspace, cp_rank, cp_size) + # REQUIRED barrier before the first alltoall: every rank must finish init, + # else a rank writes a peer's FIFO before it is ready -> deadlock. + dist.barrier(group=cp_group.device_group) + _FI_A2A_STATE = { + "workspace": workspace, + "cp_rank": cp_rank, + } + + +def dcp_a2a_lse_reduce( + cp_attn_out: torch.Tensor, + cp_attn_lse: torch.Tensor, + cp_group: "GroupCoordinator", + is_lse_base_on_e: bool = True, + cuda_graph_buffers: Optional[dict] = None, + comm_backend: str = "a2a", +) -> torch.Tensor: + """A2A DCP reduce: all-to-all exchange of head partials, then local Triton + combine. Output + fp32 LSE are packed into ONE all_to_all (LSE reinterpreted + as output-dtype columns along D) -> 1 NCCL call/layer instead of 2. + is_lse_base_on_e: True=base-e (FlashAttention), False=base-2 (FlashInfer-MLA). + """ + if cp_group.world_size == 1: + return cp_attn_out + + if comm_backend == "fi_a2a": + return _dcp_fi_a2a_lse_reduce( + cp_attn_out, cp_attn_lse, cp_group, is_lse_base_on_e + ) + + N = cp_group.world_size + B, H, D = cp_attn_out.shape + assert H % N == 0, f"num_heads ({H}) must be divisible by dcp_size ({N})" + H_per_rank = H // N + out_dtype = cp_attn_out.dtype + lpd = _lse_pack_dim(out_dtype) # 2 for bf16/fp16 + + # Reshape [B, H, D] -> [N, B, H/N, D] — split heads across ranks + reshaped_out = cp_attn_out.view(B, N, H_per_rank, D).permute(1, 0, 2, 3) + reshaped_lse = cp_attn_lse.view(B, N, H_per_rank).permute(1, 0, 2) + + if cuda_graph_buffers is not None: + # CUDA graph path with pre-allocated fused buffers. + send_combined = cuda_graph_buffers["send_combined"] + recv_combined = cuda_graph_buffers["recv_combined"] + send_lse_stg = cuda_graph_buffers["send_lse"] + recv_lse_stg = cuda_graph_buffers["recv_lse"] + + send_combined[:, :B, :, :D].copy_(reshaped_out) + send_lse_stg[:, :B, :].copy_(reshaped_lse) + send_combined[:, :, :, D:].copy_( + send_lse_stg.view(out_dtype).view(N, -1, H_per_rank, lpd) + ) + + cp_group.all_to_all_single( + recv_combined.reshape(-1).view(torch.uint8), + send_combined.reshape(-1).view(torch.uint8), + ) + recv_output = recv_combined[:, :B, :, :D] + recv_lse_stg.view(out_dtype).view(N, -1, H_per_rank, lpd).copy_( + recv_combined[:, :, :, D:] + ) + recv_lse = recv_lse_stg[:, :B, :] + else: + send_lse_contig = reshaped_lse.contiguous() # [N, B, H_per_rank] fp32 + send_combined = torch.empty( + N, + B, + H_per_rank, + D + lpd, + dtype=out_dtype, + device=cp_attn_out.device, + ) + recv_combined = torch.empty_like(send_combined) + + send_combined[:, :, :, :D].copy_(reshaped_out) + send_combined[:, :, :, D:].copy_( + send_lse_contig.view(out_dtype).view(N, B, H_per_rank, lpd) + ) + + # Transport as raw bytes (uint8): the output may be fp8 (fp8 KV cache), + # which pynccl's dtype enum can't send; byte a2a is exact for equal chunks. + cp_group.all_to_all_single( + recv_combined.reshape(-1).view(torch.uint8), + send_combined.reshape(-1).view(torch.uint8), + ) + + recv_output = recv_combined[:, :, :, :D] + recv_lse_stg = torch.empty( + N, + B, + H_per_rank, + dtype=torch.float32, + device=cp_attn_out.device, + ) + recv_lse_stg.view(out_dtype).view(N, B, H_per_rank, lpd).copy_( + recv_combined[:, :, :, D:] + ) + recv_lse = recv_lse_stg + + combined, _ = dcp_lse_combine_triton( + recv_output, recv_lse, is_lse_base_on_e=is_lse_base_on_e + ) + return combined + + +def _dcp_fi_a2a_lse_reduce( + cp_attn_out: torch.Tensor, + cp_attn_lse: torch.Tensor, + cp_group: "GroupCoordinator", + is_lse_base_on_e: bool = True, +) -> torch.Tensor: + """fi_a2a: delegate only the cross-rank exchange to FlashInfer's MNNVL kernel, + then reuse the local Triton LSE combine. FlashInfer takes output + LSE as + separate tensors: partial_o [B, H_per_rank, cp_size, D] (peer axis 2nd-to-last), + softmax_stats [B, H_per_rank, cp_size, 2] fp32 (S padded 1->2). + """ + from flashinfer.comm.dcp_alltoall import decode_cp_a2a_alltoall + + state = _FI_A2A_STATE + assert state is not None, ( + "fi_a2a workspace not initialized — call init_fi_a2a_workspace(dcp_group) " + "at model-runner init (before CUDA graph capture)." + ) + + N = cp_group.world_size + B, H, D = cp_attn_out.shape + assert H % N == 0, f"num_heads ({H}) must be divisible by dcp_size ({N})" + H_per_rank = H // N + + # FlashInfer sends partial_o[..., peer, :] to `peer`; head h -> peer h//H_per_rank, + # so the peer axis is the outer head split: [B,N,H_pr,D] -> [B,H_pr,N,D]. + partial_o = cp_attn_out.view(B, N, H_per_rank, D).permute(0, 2, 1, 3).contiguous() + # softmax_stats: fp32 [B, H_per_rank, N, S=2] (FI requires S>=2 & even); + # carry the LSE in lane 0, lane 1 is ignored by the combine. + lse_view = cp_attn_lse.view(B, N, H_per_rank).permute(0, 2, 1) # [B,H_pr,N] + softmax_stats = torch.zeros( + B, H_per_rank, N, 2, dtype=torch.float32, device=cp_attn_out.device + ) + softmax_stats[..., 0] = lse_view + + o_out, stats_out = decode_cp_a2a_alltoall( + partial_o, + softmax_stats, + state["workspace"], + state["cp_rank"], + N, + ) + + # o_out[b,hpr,src] = rank src's partial for local head hpr -> combine layout. + recv_output = o_out.permute(2, 0, 1, 3).contiguous() # [N, B, H_per_rank, D] + recv_lse = stats_out[..., 0].permute(2, 0, 1).contiguous() # [N, B, H_per_rank] + + combined, _ = dcp_lse_combine_triton( + recv_output, recv_lse, is_lse_base_on_e=is_lse_base_on_e + ) + return combined diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index fdf2e69bd..207946c64 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -162,6 +162,7 @@ from sglang.srt.model_executor.runner import ( from sglang.srt.platforms import current_platform from sglang.srt.runtime_context import ( get_global_dwdp_manager, + get_parallel, get_server_args, set_global_dwdp_manager, ) @@ -856,6 +857,49 @@ class ModelRunner: self.prefill_attention_backend_str = backends.prefill_attention_backend_str self.decode_attention_backend_str = backends.decode_attention_backend_str + if self.server_args.dcp_size > 1 and self.server_args.dcp_replicate_q_proj: + self._prepare_replicated_q_proj() + + def _prepare_replicated_q_proj(self) -> None: + # --dcp-replicate-q-proj: gather each rank's attn_tp head-shard of + # q_b_proj / w_kc into full-head buffers once here (pre-capture) so the + # MLA decode path can skip the per-layer Q all-gather. bf16/fp16 only. + from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod + from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA + + dcp_group = get_parallel().dcp_group + if dcp_group.world_size <= 1: + return + n_prepared = 0 + for m in self.model.modules(): + if not isinstance(m, DeepseekV2AttentionMLA): + continue + if m.w_kc is None: + continue + qp = m.q_b_proj if m.has_q_b_proj else m.q_proj + # q-replicate only supports the unquantized bf16/fp16 absorb path; + # quantized q-proj (packed weights) and non-16-bit w_kc keep the + # per-layer Q all-gather. + if ( + m.w_kc.dtype not in (torch.bfloat16, torch.float16) + or not isinstance(qp.quant_method, UnquantizedLinearMethod) + or qp.weight.dtype not in (torch.bfloat16, torch.float16) + ): + logger.warning( + "dcp_replicate_q_proj: skipping quantized q-proj/w_kc " + "(bf16/fp16 only); this layer keeps the Q all-gather." + ) + continue + m.w_kc_qrep = dcp_group.all_gather(m.w_kc.contiguous(), dim=0) + m.q_b_proj_qrep_weight = dcp_group.all_gather( + qp.weight.data.contiguous(), dim=0 + ) + n_prepared += 1 + logger.info( + "dcp_replicate_q_proj: prepared full-head Q weights for %d MLA layers", + n_prepared, + ) + def init_cuda_graphs(self, capture_decode_cuda_graph: bool = True): capture = capture_cuda_graphs( model_runner=self, capture_decode_cuda_graph=capture_decode_cuda_graph diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index ae9244ca2..952731036 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -218,6 +218,7 @@ class BaseRunner(ABC): return self._pre_initialize_flashinfer_allreduce_workspace() + self._pre_initialize_fi_a2a_workspace() if should_run_flashinfer_autotune(self.model_runner): buffers, batch_size = self._autotune_buffers() @@ -256,6 +257,19 @@ class BaseRunner(ABC): dtype=mr.dtype, ) + def _pre_initialize_fi_a2a_workspace(self): + """Allocate the FlashInfer MNNVL all-to-all workspace for the fi_a2a DCP + comm backend; must run before CG capture (it syncs the stream + barriers + cross-rank, uncapturable) and raises early on non-MNNVL platforms. + """ + mr = self.model_runner + if mr.server_args.dcp_size <= 1 or mr.server_args.dcp_comm_backend != "fi_a2a": + return + + from sglang.srt.layers.dcp import init_fi_a2a_workspace + + init_fi_a2a_workspace(get_parallel().dcp_group) + def _flashinfer_autotune(self, *, buffers, batch_size): """Run flashinfer autotune. diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 644003173..823fbe2eb 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -24,6 +24,7 @@ from sglang.srt.layers.dcp import ( all_gather_kv_cache_for_mla_extend, all_gather_q_for_mla_decode, cp_lse_ag_out_rs_mla, + dcp_a2a_lse_reduce, ) from sglang.srt.layers.quantization.fp8_utils import ( materialize_bpreshuffle_fp8_scale_tuple, @@ -249,6 +250,19 @@ class DeepseekMLAForwardMixin: self.q_lora_rank is not None and self._can_fuse_bmm_into_attention(forward_batch) ) + # --dcp-replicate-q-proj: project full-head Q locally from pre-gathered + # weights and skip the per-layer Q all-gather (bf16 decode absorb only). + q_replicate_active = ( + get_server_args().dcp_replicate_q_proj + and get_parallel().dcp_enabled + and forward_batch.forward_mode.is_decode() + and not self.use_deep_gemm_bmm + and self.w_kc_qrep is not None + and self.q_b_proj_qrep_weight is not None + ) + if q_replicate_active: + # force standard absorb so the full-head w_kc bmm runs + fuse_bmm_attention = False q_lora = None topk_indices = None q_nope = None @@ -350,6 +364,7 @@ class DeepseekMLAForwardMixin: and get_is_capture_mode() and forward_batch.forward_mode.is_decode_or_idle() and q_lora is not None + and not q_replicate_active ): current_stream = torch.cuda.current_stream() self.alt_stream.wait_stream(current_stream) @@ -373,7 +388,15 @@ class DeepseekMLAForwardMixin: current_stream.wait_stream(self.alt_stream) else: k_nope = k_nope.unsqueeze(1) - q = self.q_b_proj_forward(q) + if q_replicate_active: + # full-head Q from the gathered weight (skips Q all-gather) + q = torch.nn.functional.linear(q, self.q_b_proj_qrep_weight).view( + -1, + self.num_local_heads * get_parallel().attn_dcp_size, + self.qk_head_dim, + ) + else: + q = self.q_b_proj_forward(q) # Hoist these above the DSA indexer split op so the indexer # and the composite bmm+attention split op are adjacent in FX. @@ -395,9 +418,18 @@ class DeepseekMLAForwardMixin: self.layer_id, prev_topk_indices ) else: - q = self.q_proj(hidden_states)[0].view( - -1, self.num_local_heads, self.qk_head_dim - ) + if q_replicate_active: + q = torch.nn.functional.linear( + hidden_states, self.q_b_proj_qrep_weight + ).view( + -1, + self.num_local_heads * get_parallel().attn_dcp_size, + self.qk_head_dim, + ) + else: + q = self.q_proj(hidden_states)[0].view( + -1, self.num_local_heads, self.qk_head_dim + ) latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0] k_nope = latent_cache[..., : self.kv_lora_rank] k_nope = self.kv_a_layernorm(k_nope).unsqueeze(1) @@ -406,7 +438,14 @@ class DeepseekMLAForwardMixin: q_nope, q_pe, k_pe = self._split_q_nope_pe(q, latent_cache) _kvb_q = None - if fusion_plan is not None: + if q_replicate_active: + # full-head absorb with the pre-gathered w_kc (q_nope already full-head) + q_nope_out = ( + torch.bmm(q_nope.transpose(0, 1), self.w_kc_qrep) + .transpose(0, 1) + .contiguous() + ) + elif fusion_plan is not None: # The composite split op fills q_nope_out_buf and attention reads # this transposed alias directly. q_nope_out = fusion_plan.q_nope_out_view @@ -556,7 +595,7 @@ class DeepseekMLAForwardMixin: # all_gather q_pe, q_nope_out,take tp8 as an example， q_pe [B, H, ROPE_DIM], q_nope_out [B, H, NOPE_DIM] gathered to [B, H * dcp_world_size, ROPE_DIM] [B, H * dcp_world_size, NOPE_DIM] for decode batch, and all gather k_pe, k_nope for extend batch. if get_parallel().dcp_enabled: - if forward_batch.forward_mode.is_decode(): + if forward_batch.forward_mode.is_decode() and not q_replicate_active: # if forward_batch.forward_mode is decode, gather q q_nope_out, q_pe = all_gather_q_for_mla_decode( q_nope_out=q_nope_out, @@ -792,10 +831,22 @@ class DeepseekMLAForwardMixin: self.num_local_heads * get_parallel().attn_dcp_size, self.kv_lora_rank, ) - attn_output = cp_lse_ag_out_rs_mla( - attn_output, lse, get_parallel().dcp_group - ) - attn_output = attn_output.transpose(0, 1) + dcp_comm_backend = get_server_args().dcp_comm_backend + if dcp_comm_backend in ("a2a", "fi_a2a"): + # A2A exchange of head partials + LSE, then local Triton combine. + # MLA decode LSE is base-2 (FlashInfer-MLA/FlashMLA) -> base_on_e=False. + attn_output = dcp_a2a_lse_reduce( + attn_output.contiguous(), + lse.contiguous(), + get_parallel().dcp_group, + is_lse_base_on_e=False, + comm_backend=dcp_comm_backend, + ) + else: + attn_output = cp_lse_ag_out_rs_mla( + attn_output, lse, get_parallel().dcp_group + ) + attn_output = attn_output.transpose(0, 1) attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank) _kvb_v = None diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index be6577806..b914dccae 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1761,6 +1761,11 @@ class DeepseekV2AttentionMLA( self.w_vc = None self.w_scale = 1.0 + # Full-head Q/absorb weights for --dcp-replicate-q-proj, gathered once + # pre-CUDA-graph-capture by the model runner; None unless replicate is on. + self.w_kc_qrep = None + self.q_b_proj_qrep_weight = None + self.w_scale_k = None self.w_scale_v = None self.use_deep_gemm_bmm = False diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 2386c26bc..1fbc56bb9 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1038,6 +1038,29 @@ class ServerArgs: ), NS("parallel"), ] = 1 + dcp_comm_backend: A[ + str, + Arg( + help="Communication backend for the decode context-parallel (DCP) " + "attention reduction: 'ag_rs' (AllGather + ReduceScatter), 'a2a' " + "(fused NCCL All-to-All exchange of output+LSE + local Triton LSE " + "combine), or 'fi_a2a' (FlashInfer MNNVL All-to-All kernel; requires " + "SM90+ and MNNVL fabric memory, e.g. GB200 NVL72).", + choices=["ag_rs", "a2a", "fi_a2a"], + ), + NS("parallel"), + ] = "ag_rs" + dcp_replicate_q_proj: A[ + bool, + Arg( + help="For MLA decode context parallelism with the a2a/fi_a2a " + "backend: replicate the Q projection so each DCP rank computes the " + "full-head query locally (redundant projection compute), eliminating " + "the per-layer head-dim all-gather of Q. Trades a small amount of " + "extra GEMM for one fewer collective per layer.", + ), + NS("parallel"), + ] = False enable_prefill_cp: A[ bool, "Enable context parallelism for the prefill phase. Select the layout with --cp-strategy.", @@ -3587,6 +3610,30 @@ class ServerArgs: "--decode-context-parallel-size) must be >= 1, but got " f"dcp_size={self.dcp_size}." ) + if self.dcp_comm_backend in ("a2a", "fi_a2a") and self.dcp_size <= 1: + raise ValueError( + f"--dcp-comm-backend {self.dcp_comm_backend} only affects the " + "decode context-parallel attention reduction and therefore " + "requires --dcp-size / --decode-context-parallel-size > 1, but " + f"got dcp_size={self.dcp_size}." + ) + if self.dcp_comm_backend == "fi_a2a" and not is_cuda(): + raise ValueError( + "--dcp-comm-backend fi_a2a delegates the exchange to FlashInfer's " + "MNNVL All-to-All kernel, which requires an NVIDIA CUDA platform " + "with SM90+ and MNNVL fabric memory (e.g. GB200 NVL72). The " + "authoritative fabric probe runs at model-runner init; use 'a2a' " + "or 'ag_rs' on clusters without MNNVL." + ) + if self.dcp_replicate_q_proj: + if self.dcp_size <= 1: + raise ValueError("--dcp-replicate-q-proj requires --dcp-size > 1.") + if self.dcp_comm_backend not in ("a2a", "fi_a2a"): + raise ValueError( + "--dcp-replicate-q-proj only applies to the a2a/fi_a2a DCP " + "communication backend (it removes the head-dim Q all-gather); " + f"got --dcp-comm-backend={self.dcp_comm_backend}." + ) if not self.dcp_size > 1: return if is_hip(): diff --git a/test/registered/kernels/test_dcp_lse_combine.py b/test/registered/kernels/test_dcp_lse_combine.py new file mode 100644 index 000000000..d50f46184 --- /dev/null +++ b/test/registered/kernels/test_dcp_lse_combine.py @@ -0,0 +1,452 @@ +"""Tests for DCP LSE combine kernels. + +Covers: +1. Triton LSE combine kernel correctness vs CPU reference (base-e and base-2) +2. Various DCP world sizes (N=1,2,4,8) +3. Edge cases: single shard, dominant LSE, equal LSE, NaN/inf +4. return_lse mode +5. dcp_a2a_lse_reduce with pre-allocated CUDA graph buffers +""" + +import unittest +from unittest.mock import MagicMock + +import torch + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-large") + + +class TestLSECombineTritonVsCPU(CustomTestCase): + """Test Triton LSE combine kernel against CPU reference.""" + + @classmethod + def setUpClass(cls): + if not torch.cuda.is_available(): + raise unittest.SkipTest("CUDA required for Triton kernel tests") + cls.device = "cuda" + + def _run_combine_test( + self, N, B, H_local, D, is_base_e, dtype=torch.bfloat16, atol=1e-2 + ): + from sglang.kernels.ops.attention.dcp_kernels import ( + _lse_weighted_combine_cpu, + dcp_lse_combine_triton, + ) + + torch.manual_seed(42) + + partial_outputs = torch.randn(N, B, H_local, D, device=self.device, dtype=dtype) + if is_base_e: + partial_lses = torch.randn( + N, B, H_local, device=self.device, dtype=torch.float32 + ) + else: + partial_lses = ( + torch.randn(N, B, H_local, device=self.device, dtype=torch.float32) + * 5.0 + ) + + cpu_result = _lse_weighted_combine_cpu( + partial_outputs.cpu(), + partial_lses.cpu(), + is_lse_base_on_e=is_base_e, + ) + + triton_result, _ = dcp_lse_combine_triton( + partial_outputs, + partial_lses, + is_lse_base_on_e=is_base_e, + return_lse=False, + ) + + torch.testing.assert_close( + triton_result.float().cpu(), + cpu_result.float(), + atol=atol, + rtol=1e-2, + ) + + def test_n2_base_e(self): + self._run_combine_test(N=2, B=4, H_local=8, D=64, is_base_e=True) + + def test_n2_base_2(self): + self._run_combine_test(N=2, B=4, H_local=8, D=64, is_base_e=False) + + def test_n4_base_e(self): + self._run_combine_test(N=4, B=8, H_local=16, D=128, is_base_e=True) + + def test_n4_base_2(self): + self._run_combine_test(N=4, B=8, H_local=16, D=128, is_base_e=False) + + def test_n8_base_e(self): + self._run_combine_test(N=8, B=4, H_local=8, D=128, is_base_e=True) + + def test_n8_base_2(self): + self._run_combine_test(N=8, B=4, H_local=8, D=512, is_base_e=False) + + def test_n2_large_batch(self): + self._run_combine_test(N=2, B=64, H_local=16, D=128, is_base_e=False) + + def test_n4_large_head_dim(self): + self._run_combine_test(N=4, B=8, H_local=8, D=512, is_base_e=True) + + +class TestLSECombineSingleShard(CustomTestCase): + """N=1 should return input unchanged.""" + + @classmethod + def setUpClass(cls): + if not torch.cuda.is_available(): + raise unittest.SkipTest("CUDA required") + cls.device = "cuda" + + def test_single_shard(self): + from sglang.kernels.ops.attention.dcp_kernels import dcp_lse_combine_triton + + N, B, H_local, D = 1, 4, 8, 64 + partial_outputs = torch.randn( + N, B, H_local, D, device=self.device, dtype=torch.bfloat16 + ) + partial_lses = torch.randn( + N, B, H_local, device=self.device, dtype=torch.float32 + ) + + triton_result, _ = dcp_lse_combine_triton( + partial_outputs, partial_lses, is_lse_base_on_e=True + ) + + torch.testing.assert_close( + triton_result.float().cpu(), + partial_outputs.squeeze(0).float().cpu(), + atol=1e-3, + rtol=1e-3, + ) + + +class TestLSECombineReturnLSE(CustomTestCase): + """Verify return_lse=True produces valid global LSE.""" + + @classmethod + def setUpClass(cls): + if not torch.cuda.is_available(): + raise unittest.SkipTest("CUDA required") + cls.device = "cuda" + + def test_return_lse(self): + from sglang.kernels.ops.attention.dcp_kernels import dcp_lse_combine_triton + + N, B, H_local, D = 2, 4, 8, 64 + partial_outputs = torch.randn( + N, B, H_local, D, device=self.device, dtype=torch.bfloat16 + ) + partial_lses = torch.randn( + N, B, H_local, device=self.device, dtype=torch.float32 + ) + + triton_result, triton_lse = dcp_lse_combine_triton( + partial_outputs, partial_lses, is_lse_base_on_e=True, return_lse=True + ) + + self.assertIsNotNone(triton_lse) + self.assertEqual(triton_lse.shape, (B, H_local)) + self.assertFalse(torch.isnan(triton_lse).any()) + + +class TestLSECombineEdgeCases(CustomTestCase): + """Test edge cases for LSE combine.""" + + @classmethod + def setUpClass(cls): + if not torch.cuda.is_available(): + raise unittest.SkipTest("CUDA required") + cls.device = "cuda" + + def test_one_shard_dominant(self): + """One shard has much larger LSE -- output should be close to that shard.""" + from sglang.kernels.ops.attention.dcp_kernels import ( + _lse_weighted_combine_cpu, + dcp_lse_combine_triton, + ) + + N, B, H_local, D = 2, 1, 1, 64 + partial_outputs = torch.randn( + N, B, H_local, D, device=self.device, dtype=torch.bfloat16 + ) + partial_lses = torch.tensor( + [[[100.0]], [[-100.0]]], device=self.device, dtype=torch.float32 + ) + + triton_result, _ = dcp_lse_combine_triton( + partial_outputs, partial_lses, is_lse_base_on_e=True + ) + cpu_result = _lse_weighted_combine_cpu( + partial_outputs.cpu(), partial_lses.cpu(), is_lse_base_on_e=True + ) + + torch.testing.assert_close( + triton_result.float().cpu(), cpu_result.float(), atol=1e-2, rtol=1e-2 + ) + torch.testing.assert_close( + triton_result.float().cpu(), + partial_outputs[0].float().cpu(), + atol=1e-2, + rtol=1e-2, + ) + + def test_equal_lse(self): + """Equal LSE across shards -- output should be mean of outputs.""" + from sglang.kernels.ops.attention.dcp_kernels import dcp_lse_combine_triton + + N, B, H_local, D = 2, 1, 1, 64 + partial_outputs = torch.randn( + N, B, H_local, D, device=self.device, dtype=torch.bfloat16 + ) + partial_lses = torch.tensor( + [[[5.0]], [[5.0]]], device=self.device, dtype=torch.float32 + ) + + triton_result, _ = dcp_lse_combine_triton( + partial_outputs, partial_lses, is_lse_base_on_e=True + ) + expected = partial_outputs.float().mean(dim=0) + + torch.testing.assert_close( + triton_result.float().cpu(), expected.cpu(), atol=1e-2, rtol=1e-2 + ) + + +class TestCPUReference(CustomTestCase): + """Test the CPU reference implementation independently.""" + + def test_basic_combine(self): + from sglang.kernels.ops.attention.dcp_kernels import _lse_weighted_combine_cpu + + N, B, H, D = 2, 2, 4, 8 + outputs = torch.randn(N, B, H, D) + lses = torch.randn(N, B, H) + + result = _lse_weighted_combine_cpu(outputs, lses, is_lse_base_on_e=True) + self.assertEqual(result.shape, (B, H, D)) + self.assertFalse(torch.isnan(result).any()) + + def test_base2_vs_base_e(self): + from sglang.kernels.ops.attention.dcp_kernels import _lse_weighted_combine_cpu + + N, B, H, D = 2, 2, 4, 8 + outputs = torch.randn(N, B, H, D) + lses = torch.randn(N, B, H) * 3.0 + + result_e = _lse_weighted_combine_cpu(outputs, lses, is_lse_base_on_e=True) + result_2 = _lse_weighted_combine_cpu(outputs, lses, is_lse_base_on_e=False) + + self.assertFalse(torch.allclose(result_e, result_2, atol=1e-3)) + + def test_nan_lse_handled(self): + from sglang.kernels.ops.attention.dcp_kernels import _lse_weighted_combine_cpu + + N, B, H, D = 2, 1, 1, 8 + outputs = torch.randn(N, B, H, D) + lses = torch.tensor([[[5.0]], [[float("nan")]]]) + + result = _lse_weighted_combine_cpu(outputs, lses, is_lse_base_on_e=True) + self.assertFalse(torch.isnan(result).any()) + + def test_inf_lse_handled(self): + from sglang.kernels.ops.attention.dcp_kernels import _lse_weighted_combine_cpu + + N, B, H, D = 2, 1, 1, 8 + outputs = torch.randn(N, B, H, D) + lses = torch.tensor([[[5.0]], [[float("inf")]]]) + + result = _lse_weighted_combine_cpu(outputs, lses, is_lse_base_on_e=True) + self.assertFalse(torch.isnan(result).any()) + + +class TestDCPA2AReduceWithCUDAGraphBuffers(CustomTestCase): + """Test dcp_a2a_lse_reduce with pre-allocated CUDA graph buffers.""" + + @classmethod + def setUpClass(cls): + if not torch.cuda.is_available(): + raise unittest.SkipTest("CUDA required") + cls.device = "cuda" + + def _make_mock_group(self, world_size): + group = MagicMock() + group.world_size = world_size + + def identity_a2a(output, input_): + output.copy_(input_) + + group.all_to_all_single = MagicMock(side_effect=identity_a2a) + return group + + def _make_cuda_graph_buffers(self, N, max_bs, H_per_rank, D, lpd=2): + """Create fused CUDA graph buffers matching dcp_a2a_lse_reduce API.""" + return { + "send_combined": torch.empty( + N, max_bs, H_per_rank, D + lpd, dtype=torch.bfloat16, device=self.device + ), + "recv_combined": torch.empty( + N, max_bs, H_per_rank, D + lpd, dtype=torch.bfloat16, device=self.device + ), + "send_lse": torch.empty( + N, max_bs, H_per_rank, dtype=torch.float32, device=self.device + ), + "recv_lse": torch.empty( + N, max_bs, H_per_rank, dtype=torch.float32, device=self.device + ), + } + + def test_cuda_graph_buffers_same_as_dynamic(self): + from sglang.srt.layers.dcp import dcp_a2a_lse_reduce + + torch.manual_seed(123) + N, B, H_per_rank, D = 2, 4, 8, 128 + H = H_per_rank * N + max_bs = 16 + + group = self._make_mock_group(N) + + attn_out = torch.randn(B, H, D, device=self.device, dtype=torch.bfloat16) + attn_lse = torch.randn(B, H, device=self.device, dtype=torch.float32) + + result_dynamic = dcp_a2a_lse_reduce( + attn_out.clone(), attn_lse.clone(), group, is_lse_base_on_e=True + ) + + cuda_graph_buffers = self._make_cuda_graph_buffers(N, max_bs, H_per_rank, D) + + result_graph = dcp_a2a_lse_reduce( + attn_out.clone(), + attn_lse.clone(), + group, + is_lse_base_on_e=True, + cuda_graph_buffers=cuda_graph_buffers, + ) + + torch.testing.assert_close( + result_graph.float().cpu(), + result_dynamic.float().cpu(), + atol=1e-5, + rtol=1e-5, + ) + + def test_cuda_graph_buffers_n4(self): + from sglang.srt.layers.dcp import dcp_a2a_lse_reduce + + torch.manual_seed(456) + N, B, H_per_rank, D = 4, 2, 4, 64 + H = H_per_rank * N + max_bs = 8 + + group = self._make_mock_group(N) + + attn_out = torch.randn(B, H, D, device=self.device, dtype=torch.bfloat16) + attn_lse = torch.randn(B, H, device=self.device, dtype=torch.float32) + + result_dynamic = dcp_a2a_lse_reduce( + attn_out.clone(), attn_lse.clone(), group, is_lse_base_on_e=True + ) + + cuda_graph_buffers = self._make_cuda_graph_buffers(N, max_bs, H_per_rank, D) + + result_graph = dcp_a2a_lse_reduce( + attn_out.clone(), + attn_lse.clone(), + group, + is_lse_base_on_e=True, + cuda_graph_buffers=cuda_graph_buffers, + ) + + torch.testing.assert_close( + result_graph.float().cpu(), + result_dynamic.float().cpu(), + atol=1e-5, + rtol=1e-5, + ) + + def test_cuda_graph_buffers_partial_batch(self): + """Buffer max_bs > actual B -- should correctly slice.""" + from sglang.srt.layers.dcp import dcp_a2a_lse_reduce + + torch.manual_seed(789) + N, B, H_per_rank, D = 2, 3, 8, 128 + H = H_per_rank * N + max_bs = 32 + + group = self._make_mock_group(N) + + attn_out = torch.randn(B, H, D, device=self.device, dtype=torch.bfloat16) + attn_lse = torch.randn(B, H, device=self.device, dtype=torch.float32) + + cuda_graph_buffers = self._make_cuda_graph_buffers(N, max_bs, H_per_rank, D) + + result = dcp_a2a_lse_reduce( + attn_out, + attn_lse, + group, + is_lse_base_on_e=True, + cuda_graph_buffers=cuda_graph_buffers, + ) + + self.assertEqual(result.shape, (B, H_per_rank, D)) + self.assertFalse(torch.isnan(result).any()) + + def test_a2a_reduce_allocates_when_no_buffers(self): + """Without cuda_graph_buffers, dcp_a2a_lse_reduce still works (eager mode).""" + from sglang.srt.layers.dcp import dcp_a2a_lse_reduce + + N, B, H_per_rank, D = 2, 4, 8, 64 + H = H_per_rank * N + + group = self._make_mock_group(N) + + attn_out = torch.randn(B, H, D, device=self.device, dtype=torch.bfloat16) + attn_lse = torch.randn(B, H, device=self.device, dtype=torch.float32) + + result = dcp_a2a_lse_reduce( + attn_out, + attn_lse, + group, + is_lse_base_on_e=True, + cuda_graph_buffers=None, + ) + + self.assertEqual(result.shape, (B, H_per_rank, D)) + self.assertFalse(torch.isnan(result).any()) + + def test_buffers_have_fixed_data_ptrs(self): + """Pre-allocated buffer data_ptr must not change -- required for graph replay.""" + from sglang.srt.layers.dcp import dcp_a2a_lse_reduce + + N, B, H_per_rank, D = 2, 4, 8, 64 + H = H_per_rank * N + max_bs = 16 + + group = self._make_mock_group(N) + + buffers = self._make_cuda_graph_buffers(N, max_bs, H_per_rank, D) + send_ptr = buffers["send_combined"].data_ptr() + recv_ptr = buffers["recv_combined"].data_ptr() + + attn_out = torch.randn(B, H, D, device=self.device, dtype=torch.bfloat16) + attn_lse = torch.randn(B, H, device=self.device, dtype=torch.float32) + + dcp_a2a_lse_reduce( + attn_out, + attn_lse, + group, + is_lse_base_on_e=True, + cuda_graph_buffers=buffers, + ) + + self.assertEqual(buffers["send_combined"].data_ptr(), send_ptr) + self.assertEqual(buffers["recv_combined"].data_ptr(), recv_ptr) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/server_args/test_dcp_config.py b/test/registered/unit/server_args/test_dcp_config.py new file mode 100644 index 000000000..dc3a8ebf7 --- /dev/null +++ b/test/registered/unit/server_args/test_dcp_config.py @@ -0,0 +1,102 @@ +"""Unit tests for DCP (Decode Context Parallelism) server args configuration. + +Covers the ``--dcp-comm-backend`` field ({ag_rs, a2a, fi_a2a}) and its +validation in ``ServerArgs._handle_dcp_validation``: + - a2a / fi_a2a require --dcp-size > 1 + - fi_a2a requires a CUDA platform (the authoritative MNNVL fabric probe runs + later, at model-runner init) + - dcp>1 requires CUDA or HIP (base behavior from the merged DCP PR) + +Tests construct with safe defaults (dcp_size=1) then mutate the fields and call +``_handle_dcp_validation`` directly, so construction never trips the platform +gate; is_cuda / is_hip are patched per-test to pin the platform deterministically +(these are CPU-CI tests, where the real is_cuda() is False). +""" + +import dataclasses +import unittest +from unittest.mock import patch + +from sglang.srt.server_args import ServerArgs +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +_mock_device = patch("sglang.srt.server_args.get_device", return_value="cuda") +_mock_device.start() + + +class TestDCPFieldDefaults(CustomTestCase): + """Verify DCP-related dataclass fields exist with correct defaults.""" + + def test_dcp_size_field_exists(self): + fields = {f.name for f in dataclasses.fields(ServerArgs)} + self.assertIn("dcp_size", fields) + + def test_dcp_comm_backend_field_exists(self): + fields = {f.name for f in dataclasses.fields(ServerArgs)} + self.assertIn("dcp_comm_backend", fields) + + def test_dcp_size_default(self): + self.assertEqual(ServerArgs.dcp_size, 1) + + def test_dcp_comm_backend_default(self): + self.assertEqual(ServerArgs.dcp_comm_backend, "ag_rs") + + +class TestDCPCommBackendValidation(CustomTestCase): + """Verify ``_handle_dcp_validation`` accepts/rejects the right combos.""" + + @staticmethod + def _make_args(dcp_size, dcp_comm_backend): + # Construct with safe defaults (dcp_size=1) so __post_init__ never trips + # the dcp>1 platform gate, then set the fields under test. + args = ServerArgs(model_path="dummy") + args.dcp_size = dcp_size + args.dcp_comm_backend = dcp_comm_backend + return args + + def test_a2a_requires_dcp_size_gt_1(self): + args = self._make_args(dcp_size=1, dcp_comm_backend="a2a") + with self.assertRaises(ValueError): + args._handle_dcp_validation() + + def test_fi_a2a_requires_dcp_size_gt_1(self): + args = self._make_args(dcp_size=1, dcp_comm_backend="fi_a2a") + with self.assertRaises(ValueError): + args._handle_dcp_validation() + + @patch("sglang.srt.server_args.is_hip", return_value=False) + @patch("sglang.srt.server_args.is_cuda", return_value=True) + def test_a2a_with_dcp_size_2_on_cuda_passes(self, *_): + args = self._make_args(dcp_size=2, dcp_comm_backend="a2a") + args._handle_dcp_validation() # no raise + self.assertEqual(args.dcp_comm_backend, "a2a") + + @patch("sglang.srt.server_args.is_hip", return_value=False) + @patch("sglang.srt.server_args.is_cuda", return_value=True) + def test_fi_a2a_with_dcp_size_2_on_cuda_passes_server_args(self, *_): + # server_args accepts fi_a2a on CUDA; the MNNVL fabric probe is deferred + # to model-runner init (init_fi_a2a_workspace). + args = self._make_args(dcp_size=2, dcp_comm_backend="fi_a2a") + args._handle_dcp_validation() # no raise + self.assertEqual(args.dcp_comm_backend, "fi_a2a") + + @patch("sglang.srt.server_args.is_hip", return_value=False) + @patch("sglang.srt.server_args.is_cuda", return_value=False) + def test_fi_a2a_on_non_cuda_raises(self, *_): + args = self._make_args(dcp_size=2, dcp_comm_backend="fi_a2a") + with self.assertRaises(ValueError): + args._handle_dcp_validation() + + @patch("sglang.srt.server_args.is_hip", return_value=False) + @patch("sglang.srt.server_args.is_cuda", return_value=True) + def test_ag_rs_with_dcp_size_8_on_cuda_passes(self, *_): + args = self._make_args(dcp_size=8, dcp_comm_backend="ag_rs") + args._handle_dcp_validation() # no raise + self.assertEqual(args.dcp_size, 8) + + +if __name__ == "__main__": + unittest.main()