From ea8f4e9f3fc8331de2471ae67f040ce82910436e Mon Sep 17 00:00:00 2001 From: Augusto Yao Date: Fri, 26 Jun 2026 06:15:04 +0800 Subject: [PATCH] [feature] implement dcp for deepseek_v2 (#14194) --- .../sglang/srt/distributed/parallel_state.py | 70 +- python/sglang/srt/entrypoints/engine.py | 5 + .../attention/flashinfer_mla_backend.py | 47 +- .../srt/layers/attention/flashmla_backend.py | 26 +- python/sglang/srt/layers/utils/dcp_utils.py | 724 ++++++++++++++++++ python/sglang/srt/managers/scheduler.py | 2 +- python/sglang/srt/mem_cache/common.py | 12 +- .../sglang/srt/mem_cache/kv_cache_builder.py | 10 +- python/sglang/srt/mem_cache/memory_pool.py | 13 + .../srt/mem_cache/triton_ops/mla_buffer.py | 17 +- .../forward_batch_deepseek_mha_mixin.py | 17 +- .../srt/model_executor/forward_batch_info.py | 10 +- .../model_runner_kv_cache_mixin.py | 4 +- .../srt/model_executor/runner/eager_runner.py | 27 +- .../attention_forward_methods/forward_mha.py | 36 +- .../attention_forward_methods/forward_mla.py | 62 ++ python/sglang/srt/models/deepseek_v2.py | 45 ++ python/sglang/srt/server_args.py | 22 +- test/registered/dcp/test_dsv31_dcp8_gsm8k.py | 421 ++++++++++ .../dcp/test_reduce_scatter_along_dim.py | 230 ++++++ 20 files changed, 1770 insertions(+), 30 deletions(-) create mode 100644 python/sglang/srt/layers/utils/dcp_utils.py create mode 100644 test/registered/dcp/test_dsv31_dcp8_gsm8k.py create mode 100644 test/registered/dcp/test_reduce_scatter_along_dim.py diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 8a48047de..327123dd9 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -54,6 +54,7 @@ from sglang.srt.utils import ( get_current_device_stream_fast, get_int_env_var, is_cpu, + is_cuda, is_cuda_alike, is_hip, is_musa, @@ -775,6 +776,43 @@ class GroupCoordinator: else: torch.distributed.all_reduce(input_, group=self.device_group) + def reduce_scatter_along_dim( + self, input_: torch.Tensor, dim: int = -1 + ) -> torch.Tensor: + world_size = self.world_size + # Bypass the function if we are using only 1 GPU. + if world_size == 1: + return input_ + assert ( + -input_.dim() <= dim < input_.dim() + ), f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" + + if dim < 0: + # Convert negative dim to positive. + dim += input_.dim() + + with self.use_symmetric_memory(self): + # TODO: make sure whether tensor layout affects nccl reduce_scatter + # Note: This will produce an incorrect answer if we don't make + # the input_tensor contiguous. Possible bug in reduce_scatter_tensor? + input_tensor = input_.movedim(dim, 0).contiguous() + + assert input_tensor.shape[0] % world_size == 0 + chunk_size = input_tensor.shape[0] // world_size + output_shape = (chunk_size,) + input_tensor.shape[1:] + + with self.use_symmetric_memory(self): + output_tensor = torch.empty( + output_shape, + dtype=input_tensor.dtype, + device=input_tensor.device, + ) + + self.reduce_scatter_tensor(output_tensor, input_tensor) + + # Reshape before returning + return output_tensor.movedim(0, dim) + def _reduce_scatter_tensor( self, output: torch.Tensor, @@ -1667,6 +1705,10 @@ def get_attn_cp_group() -> GroupCoordinator: return _ATTN_CP +def get_dcp_group_no_assert() -> Optional[GroupCoordinator]: + return _DCP + + def get_dcp_group() -> GroupCoordinator: assert _DCP is not None, "decode context parallel group is not initialized" return _DCP @@ -1999,12 +2041,12 @@ def initialize_model_parallel( raise RuntimeError( f"decode_context_parallel_size ({decode_context_parallel_size}) must be >= 1" ) - if decode_context_parallel_size > 1 and not is_hip(): + if decode_context_parallel_size > 1 and not (is_hip() or is_cuda()): raise RuntimeError( "Decode context parallel (decode_context_parallel_size > 1) is " - "currently only supported on the AMD HIP platform, but got " + "currently only supported on the AMD HIP platform or CUDA platform, but got " f"decode_context_parallel_size ({decode_context_parallel_size}) " - "on a non-HIP platform." + "on a non-HIP or non-CUDA platform." ) if tensor_model_parallel_size % decode_context_parallel_size != 0: raise RuntimeError( @@ -2071,6 +2113,15 @@ def initialize_model_parallel( group_name="dcp", recovered_rank=recovered_rank, ) + if get_tensor_model_parallel_rank() == 0: + logger.info( + f"DCP enabled, dcp_size={decode_context_parallel_size}, tp_size={tensor_model_parallel_size}" + ) + else: + if get_tensor_model_parallel_rank() == 0: + logger.info( + f"DCP disabled, dcp_size={decode_context_parallel_size}, tp_size={tensor_model_parallel_size}" + ) attn_dp_size = attention_data_parallel_size attn_cp_size = attention_context_model_parallel_size @@ -2333,6 +2384,11 @@ def ensure_model_parallel_initialized( f"{pp_world_size=} vs. " f"{pipeline_model_parallel_size=}" ) + if decode_context_parallel_size > 1: + dcp_world_size = get_dcp_group().world_size + assert ( + dcp_world_size == decode_context_parallel_size + ), f"decode context parallel group already initialized, but of unexpected size: {dcp_world_size=} {decode_context_parallel_size=}" def model_parallel_is_initialized(): @@ -2383,6 +2439,14 @@ def get_tensor_model_parallel_world_size(): return get_tp_group().world_size +def get_dcp_world_size(): + return get_dcp_group().world_size + + +def get_dcp_rank(): + return get_dcp_group().rank_in_group + + def get_tensor_model_parallel_rank(): """Return my rank for the tensor model parallel group.""" return get_tp_group().rank_in_group diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 9998640d3..1ed7c291a 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -1259,6 +1259,11 @@ def _set_envs_and_config(server_args: ServerArgs): os.environ["NCCL_NVLS_ENABLE"] = str( int(server_args.enable_nccl_nvls or server_args.enable_symm_mem) ) + if "NCCL_GRAPH_MIXING_SUPPORT" not in os.environ or server_args.enable_symm_mem: + # Note(wh): NCCL_GRAPH_MIXING_SUPPORT=0 can help improve performance for symmetric kernels. + # details in https://github.com/NVIDIA/nccl-tests/issues/333#issuecomment-3103636985 + if server_args.dcp_size > 1: + os.environ["NCCL_GRAPH_MIXING_SUPPORT"] = "0" os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "8" os.environ["CUDA_MODULE_LOADING"] = "AUTO" diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index d2c8b9395..462d6ce47 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py @@ -23,6 +23,13 @@ from sglang.srt.layers.attention.flashinfer_backend import ( create_flashinfer_kv_indices_triton, ) from sglang.srt.layers.attention.utils import assert_buffer_fits +from sglang.srt.layers.utils.dcp_utils import ( + DecodeContextParallelMetadata, + dcp_enabled, + get_attention_dcp_world_size, + plan_dcp_decode_metadata, + update_local_kv_lens_for_dcp, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( is_in_tc_piecewise_cuda_graph, @@ -410,6 +417,7 @@ class FlashInferMLAAttnBackend(AttentionBackend): prefix_lens, prefill_wrapper_paged=self.prefill_wrapper_paged, use_ragged=use_ragged, + attn_dcp_metadata=forward_batch.attn_dcp_metadata, ) self.forward_metadata = PrefillMetadata( self.prefill_wrapper_paged, use_ragged @@ -463,6 +471,7 @@ class FlashInferMLAAttnBackend(AttentionBackend): if forward_mode.is_decode_or_idle(): assert seq_lens_cpu is not None kv_len_arr_cpu = seq_lens_cpu[:bs].to(torch.int32) + update_local_kv_lens_for_dcp(kv_len_arr_cpu) self.cuda_graph_kv_indptr_cpu[1 : bs + 1] = torch.cumsum( kv_len_arr_cpu, dim=0 ) @@ -559,7 +568,13 @@ class FlashInferMLAAttnBackend(AttentionBackend): ) else: # mla paged prefill - k_buf = self.token_to_kv_pool.get_key_buffer(layer.layer_id).to(q.dtype) + if ( + forward_batch.attn_dcp_metadata is not None + and forward_batch.attn_dcp_metadata.dcp_kv_buffer is not None + ): + k_buf = forward_batch.attn_dcp_metadata.dcp_kv_buffer.to(q.dtype) + else: + k_buf = self.token_to_kv_pool.get_key_buffer(layer.layer_id).to(q.dtype) if q_rope is None: qall = q.view(-1, layer.tp_q_head_num, layer.head_dim) q, q_rope = ( @@ -624,6 +639,7 @@ class FlashInferMLAAttnBackend(AttentionBackend): k_buffer = self.token_to_kv_pool.get_key_buffer(layer.layer_id).to(q.dtype) o = q_nope.new_empty(q_nope.shape) + # for decode and dcp_world_size > 1, lse should be returned to compute final attn_out # Direct call to run without the wrapper o = decode_wrapper.run( q_nope, @@ -631,8 +647,13 @@ class FlashInferMLAAttnBackend(AttentionBackend): k_buffer[:, :, : layer.v_head_dim], k_buffer[:, :, layer.v_head_dim :], out=o, + # for decode forward_batch, each dcp rank computes total q and partial kv, thus, we need to return_lse for online softmax to get final attn_output + return_lse=forward_batch.forward_mode.is_decode() and dcp_enabled(), ) - + if isinstance(o, tuple): + out, lse = o + out = out.view(-1, layer.tp_q_head_num * layer.v_head_dim) + return (out, lse) return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) @@ -640,7 +661,9 @@ class FlashInferMLAIndicesUpdaterDecode: def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend): # Parse Constants self.num_local_heads = ( - model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size + model_runner.model_config.num_attention_heads + // get_parallel().attn_tp_size + * get_attention_dcp_world_size() ) self.kv_lora_rank = model_runner.model_config.kv_lora_rank self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim @@ -710,6 +733,16 @@ class FlashInferMLAIndicesUpdaterDecode: kv_indices, self.req_to_token.shape[1], ) + + if dcp_enabled(): + plan_dcp_decode_metadata( + kv_lens, + kv_indptr, + kv_indices, + init_metadata_replay, + fast_decode_kwargs, + bs, + ) else: kv_indptr, kv_indices = spec_info.kv_indptr, spec_info.kv_indices @@ -775,6 +808,7 @@ class FlashInferMLAIndicesUpdaterPrefill: prefill_wrapper_paged: BatchMLAPagedAttentionWrapper, use_ragged: bool, spec_info: Optional[SpecInput] = None, + attn_dcp_metadata: Optional[DecodeContextParallelMetadata] = None, ): if use_ragged: paged_kernel_lens = prefix_lens @@ -795,6 +829,7 @@ class FlashInferMLAIndicesUpdaterPrefill: self.qo_indptr, use_ragged, spec_info, + attn_dcp_metadata=attn_dcp_metadata, ) def call_begin_forward( @@ -810,6 +845,7 @@ class FlashInferMLAIndicesUpdaterPrefill: qo_indptr: torch.Tensor, use_ragged: bool, spec_info: Optional[SpecInput] = None, + attn_dcp_metadata: Optional[DecodeContextParallelMetadata] = None, ): bs = len(seq_lens) sm_scale = self.scaling @@ -861,6 +897,11 @@ class FlashInferMLAIndicesUpdaterPrefill: ) else: # mla paged prefill + if attn_dcp_metadata is not None: + if attn_dcp_metadata.dcp_kv_indptr is not None: + kv_indptr = attn_dcp_metadata.dcp_kv_indptr + if attn_dcp_metadata.dcp_kv_indices is not None: + kv_indices = attn_dcp_metadata.dcp_kv_indices kv_len_arr = kv_indptr[1:] - kv_indptr[:-1] wrapper_paged.plan( qo_indptr, diff --git a/python/sglang/srt/layers/attention/flashmla_backend.py b/python/sglang/srt/layers/attention/flashmla_backend.py index 0c382790e..8e9988c89 100644 --- a/python/sglang/srt/layers/attention/flashmla_backend.py +++ b/python/sglang/srt/layers/attention/flashmla_backend.py @@ -18,9 +18,16 @@ from sglang.srt.layers.attention.utils import ( get_num_kv_index_blocks_flashmla, ) from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant +from sglang.srt.layers.utils.dcp_utils import ( + dcp_enabled, + get_attention_dcp_rank, + get_attention_dcp_world_size, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import get_parallel +logger = logging.getLogger(__name__) + if TYPE_CHECKING: from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.model_executor.model_runner import ModelRunner @@ -88,6 +95,10 @@ class FlashMLABackend(FlashInferMLAAttnBackend): self.cuda_graph_mla_metadata_view = None self.cuda_graph_num_splits_view = None + # get dcp info + self.dcp_world_size = get_attention_dcp_world_size() + self.dcp_rank = get_attention_dcp_rank() + def init_forward_metadata_out_graph( self, forward_batch: ForwardBatch, @@ -327,6 +338,9 @@ class FlashMLABackend(FlashInferMLAAttnBackend): reshape_q = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim) if self.is_fp8_kvcache: + assert ( + self.dcp_world_size == 1 + ), "FlashMLA does not support DCP for FP8 kv cache" if layer.k_scale is not None: q_scale = layer.k_scale descale_q = layer.k_scale.reshape(1) @@ -360,7 +374,8 @@ class FlashMLABackend(FlashInferMLAAttnBackend): return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) else: - o, _ = flash_mla_with_kvcache( + # todo: need check all causal True or False? + o, lse = flash_mla_with_kvcache( q=reshape_q, k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim), block_table=self.forward_metadata.block_kv_indices[:bs], @@ -371,8 +386,13 @@ class FlashMLABackend(FlashInferMLAAttnBackend): softmax_scale=layer.scaling, causal=True, ) - - return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + o = o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + # TODO uniform output for forward_decode and forward_extend to + # return tuple instead of single output + # decode context parallel needs lse to correct attn_output via online softmax + if dcp_enabled(): + return o, lse + return o def forward_extend( self, diff --git a/python/sglang/srt/layers/utils/dcp_utils.py b/python/sglang/srt/layers/utils/dcp_utils.py new file mode 100644 index 000000000..46eec98a0 --- /dev/null +++ b/python/sglang/srt/layers/utils/dcp_utils.py @@ -0,0 +1,724 @@ +from dataclasses import dataclass +from typing import Optional + +import torch +import triton +import triton.language as tl + +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, +) +from sglang.srt.distributed.parallel_state import ( + GroupCoordinator, + get_dcp_group, + get_dcp_group_no_assert, + get_dcp_rank, + get_dcp_world_size, +) +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import is_cuda + + +def dcp_enabled() -> bool: + """ + only checks whether dcp enabled for cuda platform + """ + if get_dcp_group_no_assert() is None: + return False + if not is_cuda(): + return False + return get_dcp_world_size() > 1 + + +def get_attention_dcp_group() -> GroupCoordinator: + return get_dcp_group() + + +def get_attention_dcp_world_size() -> int: + if not dcp_enabled(): + return 1 + return get_dcp_world_size() + + +def get_attention_dcp_rank() -> int: + if not dcp_enabled(): + return 0 + return get_dcp_rank() + + +@triton.jit +def _correct_attn_cp_out_kernel( + outputs_ptr, + new_output_ptr, + lses_ptr, + vlse_ptr, + outputs_stride_B, + outputs_stride_H, + outputs_stride_D, + lses_stride_N, + lses_stride_B, + lses_stride_H, + new_outputs_stride_H, + new_outputs_stride_B, + new_outputs_stride_D, + lse_idx, + HEAD_DIM: tl.constexpr, + N_ROUNDED: tl.constexpr, +): + """ + Apply the all-gathered lses to correct each local rank's attention + output. we still need perform a cross-rank reduction to obtain the + final attention output. + + Args: + outputs_ptr (triton.PointerType): + Pointer to input tensor of shape [ B, H, D ] + lses_ptr (triton.PointerType): + Pointer to input tensor of shape [ N, B, H ] + new_output_ptr (triton.PointerType): + Pointer to output tensor of shape [ H, B, D ] + vlse_ptr (triton.PointerType): + Pointer to output tensor of shape [ B, H ] + """ + batch_idx = tl.program_id(axis=0).to(tl.int64) + head_idx = tl.program_id(axis=1).to(tl.int64) + + # Use int32 for offsets where possible to reduce register pressure + b_i32 = batch_idx.to(tl.int32) + h_i32 = head_idx.to(tl.int32) + + # Vectorized load of LSE values: shape = [N] + num_n_offsets = tl.arange(0, N_ROUNDED) + lse_offsets = ( + num_n_offsets * lses_stride_N + b_i32 * lses_stride_B + h_i32 * lses_stride_H + ) + + # Compute final LSE using online softmax algorithm (more numerically stable) + lse = tl.load(lses_ptr + lse_offsets) + + # Replace NaN and inf with -inf for numerical stability + neg_inf = float("-inf") + lse = tl.where((lse != lse) | (lse == float("inf")), neg_inf, lse) + + # Online softmax: find max, subtract, exp, sum, log + lse_max = tl.max(lse, axis=0) + lse_max = tl.where(lse_max == neg_inf, 0.0, lse_max) + lse = lse - lse_max + lse_exp = tl.exp2(lse) + lse_acc = tl.sum(lse_exp, axis=0) + final_lse = tl.log2(lse_acc) + lse_max + + # Compute correction factor + lse_offset = lse_idx * lses_stride_N + b_i32 * lses_stride_B + h_i32 * lses_stride_H + local_lse = tl.load(lses_ptr + lse_offset) + lse_diff = local_lse - final_lse + lse_diff = tl.where( + (lse_diff != lse_diff) | (lse_diff == float("inf")), + neg_inf, + lse_diff, + ) + factor = tl.exp2(lse_diff) + + # Store final LSE + tl.store(vlse_ptr + b_i32 * lses_stride_B + h_i32 * lses_stride_H, final_lse) + + # Load output with vectorized access: shape = [D] + d_offsets = tl.arange(0, HEAD_DIM) + output_offsets = ( + batch_idx * outputs_stride_B + + head_idx * outputs_stride_H + + d_offsets * outputs_stride_D + ) + + new_output_offsets = ( + head_idx * new_outputs_stride_H + + batch_idx * new_outputs_stride_B + + d_offsets * new_outputs_stride_D + ) + # Apply correction and store + output = tl.load(outputs_ptr + output_offsets) + output = output * factor + tl.store(new_output_ptr + new_output_offsets, output) + + +class CPTritonContext: + """The CPTritonContext is used to avoid recompilation of the Triton JIT.""" + + def __init__(self): + self.inner_kernel = None + + def call_kernel(self, kernel, grid, *regular_args, **const_args): + if self.inner_kernel is None: + self.inner_kernel = kernel[grid](*regular_args, **const_args) + else: + self.inner_kernel[grid](*regular_args) + + +def correct_attn_out( + out: torch.Tensor, + lses: torch.Tensor, + cp_rank: int, + ctx: Optional[CPTritonContext], + new_output: torch.Tensor = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Correct the attention output using the all-gathered lses. + + Args: + out: Tensor of shape [ B, H, D ] + lses: Tensor of shape [ N, B, H ] + cp_rank: Current rank in the context-parallel group + ctx: Triton context to avoid recompilation + + Returns: + Tuple of (out, lse) with corrected attention and final log-sum-exp. + """ + if ctx is None: + ctx = CPTritonContext() + + # --- Normalize to 3D views --- + if out.ndim == 4 and out.shape[1] == 1: + out = out.squeeze(1) + assert out.ndim == 3, f"expected out [B,H,D] or [B,1,H,D], got {tuple(out.shape)}" + + if lses.ndim == 4 and lses.shape[-1] == 1: + lses = lses.squeeze(-1) + if lses.ndim == 4 and lses.shape[1] == 1: + lses = lses.squeeze(1) + assert lses.ndim == 3, ( + f"expected lses [N,B,H] (optionally with a 1-sized extra dim), " + f"got {tuple(lses.shape)}" + ) + + B, H, D = out.shape + N = lses.shape[0] + + # Strides after we normalized shapes to 3-D views. The kernel computes + # offsets for `vlse_ptr` using lses_stride_B/H, so the output buffer must + # have the same B/H stride layout as a slice of `lses`. + o_sB, o_sH, o_sD = out.stride() + l_sN, l_sB, l_sH = lses.stride() + no_sH, no_sB, no_sD = new_output.stride() + # Allocate LSE with the same B/H strides as `lses` so writes land correctly + # even when `lses` is a non-contiguous view (e.g., 4-D to 3-D squeeze). + lse = torch.empty_strided( + (B, H), (l_sB, l_sH), device=lses.device, dtype=lses.dtype + ) + + # Kernel launch config + grid = (B, H, 1) + + regular_args = ( + out, + new_output, + lses, + lse, + o_sB, + o_sH, + o_sD, + l_sN, + l_sB, + l_sH, + no_sH, + no_sB, + no_sD, + cp_rank, + ) + const_args = {"HEAD_DIM": D, "N_ROUNDED": N} + + ctx.call_kernel(_correct_attn_cp_out_kernel, grid, *regular_args, **const_args) + return new_output, lse + + +def cp_lse_ag_out_rs( + cp_attn_out: torch.Tensor, + cp_attn_lse: torch.Tensor, + cp_group: GroupCoordinator, + ctx: Optional[CPTritonContext] = None, +): + """ + cp_attn_out: [ B, H, D ] + cp_attn_lse: [ B, H ] + """ + if cp_group.world_size == 1: + return cp_attn_out + + if ctx is None: + ctx = CPTritonContext() + + with use_symmetric_memory(cp_group): + # cp_attn_out is [B,H,D], we want to transpose it to [H,B,D] for the kernel, and then transpose back after correction. + new_output = cp_attn_out.new_empty( + cp_attn_out.transpose(0, 1).shape, dtype=torch.float32 + ) + cp_attn_lse = cp_attn_lse.to(torch.float32) + lses = cp_group.all_gather(cp_attn_lse, dim=0).view( + (cp_group.world_size,) + cp_attn_lse.shape + ) + out, _ = correct_attn_out( + cp_attn_out, lses, cp_group.rank_in_group, ctx, new_output + ) + out = cp_group.reduce_scatter_along_dim(out, dim=0) + return out.to(cp_attn_out.dtype) + + +@triton.jit +def create_dcp_kv_indices( + kv_indptr, + extend_lens_ptr, + extend_cu_lens_ptr, + extend_prefix_lens_ptr, + extend_cu_prefix_lens_ptr, + kv_indices_ptr, + extend_prefix_lens_sum, + dcp_world_size: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 512 + pid = tl.program_id(axis=0) + prefix_len = tl.load(extend_prefix_lens_ptr + pid) + prefix_start = tl.load(extend_cu_prefix_lens_ptr + pid) + kv_ind_start = tl.load(kv_indptr + pid) + num_loop = tl.cdiv(prefix_len, BLOCK_SIZE) + for i in range(num_loop): + offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = offset < prefix_len + data = prefix_start + offset + tl.store(kv_indices_ptr + kv_ind_start + offset, data, mask=mask) + extend_len = tl.load(extend_lens_ptr + pid) + extend_start = tl.load(extend_cu_lens_ptr + pid) + num_loop = tl.cdiv(extend_len, BLOCK_SIZE) + for i in range(num_loop): + offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = offset < extend_len + data = extend_prefix_lens_sum + extend_start + offset + tl.store( + kv_indices_ptr + kv_ind_start + prefix_len + offset, + data, + mask=mask, + ) + + +@triton.jit +def update_kv_lens_and_indices( + kv_lens: torch.Tensor, + kv_lens_cumsum: torch.Tensor, + kv_indices: torch.Tensor, + local_kv_lens: torch.Tensor, + local_kv_lens_cumsum: torch.Tensor, + local_kv_indices: torch.Tensor, + dcp_rank: tl.constexpr, + dcp_world_size: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + bs_idx = tl.program_id(0) + block_idx = tl.program_id(1) + + local_kv_len = tl.load(local_kv_lens + bs_idx) + local_kv_indices_start = tl.load(local_kv_lens_cumsum + bs_idx) + kv_indices_start = tl.load(kv_lens_cumsum + bs_idx) + + block_start = block_idx * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + + mask = offsets < local_kv_len + + kv_indice_offsets = offsets * dcp_world_size + dcp_rank + kv_indices_start + local_kv_indices_offsets = local_kv_indices_start + offsets + + kv_values = tl.load(kv_indices + kv_indice_offsets, mask=mask) + tl.store( + local_kv_indices + local_kv_indices_offsets, + kv_values // dcp_world_size, + mask=mask, + ) + + +@dataclass +class DecodeContextParallelMetadata: + # For decode context parallel + dcp_kv_indptr: Optional[torch.Tensor] = None + dcp_kv_buffer: Optional[torch.Tensor] = None + dcp_kv_indices: Optional[torch.Tensor] = None + dcp_local_prefix_kv_indices: Optional[torch.Tensor] = None + dcp_extend_prefix_lens_sum: Optional[int] = None + + +def prepare_decode_context_parallel_metadata( + seq_lens: torch.Tensor, + extend_prefix_lens: torch.Tensor, + extend_prefix_lens_cpu: torch.Tensor, + extend_seq_lens: torch.Tensor, + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + seq_lens_sum: int, + kv_buffer_shape: torch.Size, + kv_cache_dtype, + kv_cache_device, + create_chunked_prefix_cache_kv_indices_fn, +) -> Optional[DecodeContextParallelMetadata]: + if not dcp_enabled(): + return None + # dcp_kv_buffer tokens' layout + # [ rank0_r1.prefix_tokens, rank1_r1.prefix_tokens, ..., rank7_r1.prefix_tokens, + # ..., + # rank0_rn.prefix_tokens, rank1_rn.prefix_tokens, ..., rank7_rn.prefix_tokens, + # r1.extend_tokens, r2.extent_tokens, rn.extend_tokens ] + extend_prefix_starts = torch.zeros( + len(seq_lens), + dtype=torch.int32, + device=get_global_server_args().device, + ) + extend_cu_prefix_lens = torch.zeros( + len(seq_lens) + 1, + dtype=torch.int32, + device=get_global_server_args().device, + ) + extend_cu_prefix_lens[1:] = torch.cumsum(extend_prefix_lens, dim=0) + extend_cu_prefix_lens = extend_cu_prefix_lens[:-1] + extend_prefix_lens_sum = sum([i for i in extend_prefix_lens_cpu]) + + dcp_prefix_kv_indices = torch.empty( + sum(extend_prefix_lens_cpu), + dtype=torch.int32, + device=get_global_server_args().device, + ) + create_chunked_prefix_cache_kv_indices_fn[(len(seq_lens),)]( + req_to_token, + req_pool_indices, + extend_prefix_starts, + extend_prefix_lens, + extend_cu_prefix_lens, + dcp_prefix_kv_indices, + req_to_token.shape[1], + ) + dcp_kv_indptr = torch.zeros( + len(seq_lens) + 1, + dtype=torch.int32, + device=get_global_server_args().device, + ) + dcp_kv_indptr[1:] = seq_lens.cumsum(dim=0) + dcp_kv_indptr = dcp_kv_indptr[: (len(seq_lens) + 1)] + dcp_kv_indices = torch.zeros( + seq_lens_sum, + dtype=torch.int32, + device=get_global_server_args().device, + ) + + extend_cu_lens = torch.zeros( + len(seq_lens) + 1, + dtype=torch.int32, + device=get_global_server_args().device, + ) + extend_cu_lens[1:] = torch.cumsum(extend_seq_lens, dim=0) + extend_cu_lens = extend_cu_lens[:-1] + + create_dcp_kv_indices[(len(seq_lens),)]( + dcp_kv_indptr, + extend_seq_lens, + extend_cu_lens, + extend_prefix_lens, + extend_cu_prefix_lens, + dcp_kv_indices, + extend_prefix_lens_sum, + get_dcp_world_size(), + ) + dcp_local_prefix_kv_indices = ( + dcp_prefix_kv_indices[ + dcp_prefix_kv_indices % get_dcp_world_size() == get_dcp_rank() + ] + // get_dcp_world_size() + ) + dcp_kv_buffer = torch.empty( + ( + seq_lens_sum, + *kv_buffer_shape[1:], + ), + dtype=kv_cache_dtype, + device=kv_cache_device, + ) + attn_dcp_metadata = DecodeContextParallelMetadata( + dcp_kv_indptr=dcp_kv_indptr, + dcp_kv_buffer=dcp_kv_buffer, + dcp_kv_indices=dcp_kv_indices, + dcp_local_prefix_kv_indices=dcp_local_prefix_kv_indices, + dcp_extend_prefix_lens_sum=extend_prefix_lens_sum, + ) + return attn_dcp_metadata + + +def _all_gather_dcp_kv_cache(kv_a: torch.Tensor): + dcp_world_size = get_dcp_world_size() + # not use symmetric_memory unless torch mem_pool updated, see https://github.com/pytorch/pytorch/issues/178138 + gathered_kv_a = kv_a.new_empty( + (kv_a.shape[0] * dcp_world_size, *kv_a.shape[1:]), + ) + get_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) + .reshape(-1, *kv_a.shape[1:]) + ) + return gathered_kv_a + + +def all_gather_kv_cache_for_mha_chunk_extend( + kv_a: torch.Tensor, + k_pe: torch.Tensor, + prefix_kv_lens_cpu: torch.Tensor, + prefix_starts_cpu: torch.Tensor = None, +): + if dcp_enabled(): + kv_a = kv_a.unsqueeze(1) + gathered_kv = all_gather_kv_cache_for_dcp( + kv_a, + k_pe, + prefix_kv_lens_cpu, + prefix_starts_cpu, + ) + kv_a, k_pe = gathered_kv.split([kv_a.shape[-1], k_pe.shape[-1]], dim=-1) + kv_a = kv_a.squeeze(1) + return kv_a.contiguous(), k_pe.contiguous() + + +def all_gather_kv_cache_for_mha_extend( + token_to_kv_pool, + attn_mqa, + dcp_local_prefix_kv_indices, + seq_lens, + extend_prefix_lens, + extend_prefix_lens_cpu: list[int], + extend_seq_lens, + kv_a: torch.Tensor, + k_pe: torch.Tensor, +): + prefix_kv_a, prefix_k_pe = token_to_kv_pool.get_mla_kv_buffer( + attn_mqa, dcp_local_prefix_kv_indices + ) + extend_prefix_lens_cpu = torch.tensor(extend_prefix_lens_cpu) + gathered_kv_cache = all_gather_kv_cache_for_dcp( + prefix_kv_a, + prefix_k_pe, + extend_prefix_lens_cpu, + ) + prefix_kv_a, prefix_k_pe = gathered_kv_cache.split( + [kv_a.shape[-1], k_pe.shape[-1]], dim=-1 + ) + prefix_kv_a = prefix_kv_a.squeeze(1) + # re-organize kv with query orders + prefix_lens_cu = torch.zeros( + len(seq_lens) + 1, + dtype=torch.int32, + device=kv_a.device, + ) + extend_lens_cu = torch.zeros_like(prefix_lens_cu) + prefix_lens_cu[1:] = torch.cumsum(extend_prefix_lens, dim=0) + extend_lens_cu[1:] = torch.cumsum(extend_seq_lens, dim=0) + kv_a_tuple = () + k_pe_tuple = () + for i in range(len(seq_lens)): + kv_a_tuple += ( + prefix_kv_a[prefix_lens_cu[i] : prefix_lens_cu[i + 1]], + kv_a[extend_lens_cu[i] : extend_lens_cu[i + 1]], + ) + k_pe_tuple += ( + prefix_k_pe[prefix_lens_cu[i] : prefix_lens_cu[i + 1]], + k_pe[extend_lens_cu[i] : extend_lens_cu[i + 1]], + ) + kv_a = torch.cat(kv_a_tuple, dim=0) + k_pe = torch.cat(k_pe_tuple, dim=0) + return kv_a.contiguous(), k_pe.contiguous() + + +def filter_dcp_local_kv_indices(kv_indices: torch.Tensor): + if dcp_enabled(): + kv_indices = ( + kv_indices[kv_indices % get_dcp_world_size() == get_dcp_rank()] + // get_dcp_world_size() + ) + return kv_indices + + +def all_gather_q_for_mla_decode( + q_nope_out: torch.Tensor, + q_pe: torch.Tensor, +): + with use_symmetric_memory(get_dcp_group()): + # transpose q_pe and q_nope_out from [B, H, L] to [H, B, L] + combined = torch.cat([q_pe.transpose(0, 1), q_nope_out.transpose(0, 1)], dim=-1) + gathered = get_dcp_group().all_gather(combined, dim=0) + d_pe = q_pe.size(-1) + d_nope = q_nope_out.size(-1) + q_pe, q_nope_out = gathered.split([d_pe, d_nope], dim=-1) + q_pe = q_pe.transpose(0, 1) + q_nope_out = q_nope_out.transpose(0, 1) + return q_nope_out, q_pe + + +def all_gather_kv_cache_for_mla_extend( + token_to_kv_pool, + attn_mqa, + extend_prefix_lens_cpu: list[int], + dcp_local_prefix_kv_indices, + dcp_extend_prefix_lens_sum, + dcp_kv_buffer, + kv_lora_rank, + k_nope, + k_pe, +): + cache_k_nope, cache_k_rope = token_to_kv_pool.get_mla_kv_buffer( + attn_mqa, + dcp_local_prefix_kv_indices, + ) + extend_prefix_lens_cpu = torch.tensor(extend_prefix_lens_cpu) + # all gather kv cache into forward_batch.attn_dcp_metadata.dcp_kv_buffer + gathered_kv = all_gather_kv_cache_for_dcp( + cache_k_nope, + cache_k_rope, + extend_prefix_lens_cpu, + prefix_starts_cpu=torch.zeros_like(extend_prefix_lens_cpu), + ) + dcp_kv_buffer[:dcp_extend_prefix_lens_sum] = gathered_kv + + # copy local kv cache into forward_batch.attn_dcp_metadata.dcp_kv_buffer + dcp_kv_buffer[ + dcp_extend_prefix_lens_sum:, + ..., + :kv_lora_rank, + ] = k_nope + dcp_kv_buffer[ + dcp_extend_prefix_lens_sum:, + ..., + kv_lora_rank:, + ] = k_pe + + +def update_local_kv_lens_for_dcp(kv_len_arr): + if not dcp_enabled(): + return + dcp_world_size = get_dcp_world_size() + dcp_rank = get_dcp_rank() + offset = dcp_rank + 1 + kv_len_arr.sub_(offset).div_(dcp_world_size, rounding_mode="floor").add_(1) + + +def plan_dcp_decode_metadata( + kv_lens: torch.Tensor, + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + init_metadata_replay: bool, + fast_decode_kwargs: dict, + bs: int, +): + local_kv_lens = kv_lens.clone() + update_local_kv_lens_for_dcp(local_kv_lens) + local_kv_lens.clamp_(min=0) + + if not init_metadata_replay: + max_local_len = ( + int(local_kv_lens.max().item()) if local_kv_lens.numel() > 0 else 0 + ) + total_local_len = ( + int(local_kv_lens.sum().item()) if local_kv_lens.numel() > 0 else 0 + ) + else: + max_local_len = ( + int(fast_decode_kwargs["kv_len_arr_cpu"].max().item()) + if fast_decode_kwargs["kv_len_arr_cpu"].numel() > 0 + else 0 + ) + total_local_len = ( + int(fast_decode_kwargs["kv_len_arr_cpu"].sum().item()) + if fast_decode_kwargs["kv_len_arr_cpu"].numel() > 0 + else 0 + ) + local_kv_lens_cumsum = kv_indptr.new_zeros((bs + 1,)) + local_kv_lens_cumsum[1 : bs + 1] = torch.cumsum(local_kv_lens, dim=0) + local_kv_indices = kv_indices.new_empty(total_local_len) + BLOCK_SIZE = 128 + num_blocks = ( + (max_local_len + BLOCK_SIZE - 1) // BLOCK_SIZE if max_local_len > 0 else 1 + ) + grid = (bs, num_blocks) + update_kv_lens_and_indices[grid]( + kv_lens, + kv_indptr, + kv_indices, + local_kv_lens, + local_kv_lens_cumsum, + local_kv_indices, + dcp_rank=get_dcp_rank(), + dcp_world_size=get_dcp_world_size(), + BLOCK_SIZE=BLOCK_SIZE, + ) + kv_indices[:total_local_len] = local_kv_indices[:total_local_len] + kv_lens.copy_(local_kv_lens) + kv_indptr[: bs + 1] = local_kv_lens_cumsum[: bs + 1] + + +# all gather kv cache and re-org to query orders +def all_gather_kv_cache_for_dcp( + prefix_kv_a: torch.Tensor, + prefix_k_pe: torch.Tensor, + prefix_kv_lens_cpu: torch.Tensor, + prefix_starts_cpu: torch.Tensor = None, +): + """ + prefix_kv_a and prefix_k_pe should have same shape, expect for last dim + """ + if not dcp_enabled(): + return torch.cat([prefix_kv_a, prefix_k_pe], dim=-1) + # 1. compute max kv_lens for each seq + dcp_world_size = get_dcp_world_size() + dcp_rank = get_dcp_rank() + + if prefix_starts_cpu is None: + prefix_starts_cpu = torch.zeros_like(prefix_kv_lens_cpu) + + left_pads = prefix_starts_cpu % dcp_world_size > dcp_rank + left_pads = left_pads.to(torch.int32) + right_pads = ( + prefix_starts_cpu + prefix_kv_lens_cpu - 1 + ) % dcp_world_size < dcp_rank + right_pads = right_pads.to(torch.int32) + padded_lens = ( + prefix_kv_lens_cpu + (prefix_starts_cpu % dcp_world_size) + dcp_world_size - 1 + ) // dcp_world_size + + local_kv_lens = padded_lens - left_pads - right_pads + local_kv_lens_cu = torch.zeros( + len(prefix_kv_lens_cpu) + 1, + dtype=torch.int32, + ) + local_kv_lens_cu[1:] = torch.cumsum(local_kv_lens, dim=0) + + padded_kv_cache_arr = [] + prefix_kv_cache = torch.cat([prefix_kv_a, prefix_k_pe], dim=-1) + for req_idx in range(len(prefix_kv_lens_cpu)): + padded_tensor = prefix_kv_cache.new_empty( + (padded_lens[req_idx].item(),) + prefix_kv_cache.size()[1:] + ) + padded_tensor[ + left_pads[req_idx] : left_pads[req_idx] + local_kv_lens[req_idx] + ] = prefix_kv_cache[local_kv_lens_cu[req_idx] : local_kv_lens_cu[req_idx + 1]] + padded_kv_cache_arr.append(padded_tensor) + + padded_kv_cache = torch.cat(padded_kv_cache_arr, dim=0) + + gatherd_kv_cache = _all_gather_dcp_kv_cache(padded_kv_cache) + + # 2. re-org kv cache to query orders + padded_lens_cu = torch.zeros( + len(prefix_kv_lens_cpu) + 1, + dtype=torch.int32, + ) + padded_lens_cu[1:] = torch.cumsum(padded_lens, dim=0) + kv_cache_tuple = () + for req_idx in range(len(prefix_kv_lens_cpu)): + kv_cache_tuple += ( + gatherd_kv_cache[ + padded_lens_cu[req_idx] * dcp_world_size + + (prefix_starts_cpu[req_idx] % dcp_world_size) : + ][: prefix_kv_lens_cpu[req_idx]], + ) + gatherd_kv_cache = torch.cat(kv_cache_tuple, dim=0) + + return gatherd_kv_cache diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index ec785ec65..1f4612868 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1765,7 +1765,7 @@ class Scheduler( enable_hisparse=self.enable_hisparse, full_tokens_per_layer=self.full_tokens_per_layer, swa_tokens_per_layer=self.swa_tokens_per_layer, - max_total_num_tokens=self.max_total_num_tokens, + max_total_num_tokens=self.max_total_num_tokens * self.server_args.dcp_size, get_last_batch=lambda: self.last_batch, get_running_batch=lambda: self.running_batch, ) diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index d6baf5179..0a8c83aa6 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -26,13 +26,15 @@ from sglang.srt.mem_cache.triton_ops.common import ( write_req_to_token_pool_triton, ) from sglang.srt.server_args import ServerArgs, get_global_server_args -from sglang.srt.utils import is_hip, is_npu, support_triton +from sglang.srt.utils import is_cuda, is_hip, is_npu, support_triton from sglang.srt.utils.common import ceil_align, is_pin_memory_available _is_npu = is_npu() _is_hip = is_hip() +_is_cuda = is_cuda() + if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req, ScheduleBatch from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator @@ -440,13 +442,13 @@ def alloc_req_slots( def _alloc_page_size(batch: ScheduleBatch) -> int: - # DCP (HIP-only) swaps in a PagedTokenToKVPoolAllocator whose page_size is - # server_args.page_size * dcp_size, so it can be > 1 even when + # DCP (HIP & CUDA only) swaps in a PagedTokenToKVPoolAllocator whose + # page_size is server_args.page_size * dcp_size, so it can be > 1 even when # tree_cache.page_size (== server_args.page_size) is 1. Only on the HIP DCP # path do we branch on the real allocator's page_size so the paged path is # taken; everywhere else tree_cache.page_size is authoritative and the two # are equal (dcp_size == 1), so behavior is unchanged. - if _is_hip and get_global_server_args().dcp_size > 1: + if (_is_hip or _is_cuda) and get_global_server_args().dcp_size > 1: return batch.tree_cache.token_to_kv_pool_allocator.page_size return batch.tree_cache.page_size @@ -484,6 +486,8 @@ def alloc_for_extend( if _alloc_page_size(batch) == 1: out_cache_loc = alloc_token_slots(batch.tree_cache, batch.extend_num_tokens) else: + # Since tree_cache.page_size is (page_size * dcp_world_size), for dcp + # on cuda platform, always use alloc_paged_token_slots_extend # Paged allocation - build last_loc last_loc = [ (t[-1:] if len(t) > 0 else torch.tensor([-1], device=batch.device)) diff --git a/python/sglang/srt/mem_cache/kv_cache_builder.py b/python/sglang/srt/mem_cache/kv_cache_builder.py index 726c5da39..5f985de69 100644 --- a/python/sglang/srt/mem_cache/kv_cache_builder.py +++ b/python/sglang/srt/mem_cache/kv_cache_builder.py @@ -25,6 +25,9 @@ from typing import TYPE_CHECKING from sglang.srt.configs.model_config import ModelImpl from sglang.srt.environ import envs +from sglang.srt.layers.utils.dcp_utils import ( + dcp_enabled, +) from sglang.srt.managers.mm_utils import init_mm_embedding_cache from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.registry import TreeCacheBuildContext, create_tree_cache @@ -200,7 +203,12 @@ def build_kv_cache( disable=disable_radix_cache, req_to_token_pool=req_to_token_pool, token_to_kv_pool_allocator=token_to_kv_pool_allocator, - page_size=page_size, + # When dcp enabled, kv_pool_allocator.page_size is page_size * dcp_size. + # TreeCache.page_size should keep the same as allocator.page_size to + # avoid kv page eviction conflicts. + page_size=( + page_size if not dcp_enabled() else token_to_kv_pool_allocator.page_size + ), is_eagle=spec_algorithm.is_eagle(), tp_cache_group=( attn_tp_cpu_group if server_args.enable_dp_attention else tp_cpu_group diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 2c2297e11..d71edbc26 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -46,6 +46,11 @@ from sglang.srt.layers.attention.dsa.quant_k_cache import ( from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.utils.dcp_utils import ( + dcp_enabled, + get_attention_dcp_rank, + get_attention_dcp_world_size, +) from sglang.srt.mem_cache.allocator.mamba import MambaSlotAllocator from sglang.srt.mem_cache.triton_ops.cache_move import ( copy_all_layer_kv_cache_tiled, @@ -2246,6 +2251,14 @@ class MLATokenToKVPool(KVCache): maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA)") layer_id = layer.layer_id assert not self.dsa_kv_cache_store_fp8 + if dcp_enabled(): + valid_mask = ( + loc % get_attention_dcp_world_size() == get_attention_dcp_rank() + ) + if not valid_mask.all(): + loc = loc[valid_mask] + cache_k = cache_k[valid_mask] + if cache_k.dtype != self.dtype: cache_k = cache_k.to(self.dtype) diff --git a/python/sglang/srt/mem_cache/triton_ops/mla_buffer.py b/python/sglang/srt/mem_cache/triton_ops/mla_buffer.py index 6ec9b282e..b93341bd2 100644 --- a/python/sglang/srt/mem_cache/triton_ops/mla_buffer.py +++ b/python/sglang/srt/mem_cache/triton_ops/mla_buffer.py @@ -5,6 +5,11 @@ import triton import triton.language as tl from sglang.jit_kernel.utils import is_arch_support_pdl +from sglang.srt.layers.utils.dcp_utils import ( + dcp_enabled, + get_attention_dcp_rank, + get_attention_dcp_world_size, +) @triton.jit @@ -19,6 +24,8 @@ def set_mla_kv_buffer_kernel( nope_dim: tl.constexpr, rope_dim: tl.constexpr, BLOCK: tl.constexpr, + DCP_RANK: tl.constexpr, + DCP_WORLD_SIZE: tl.constexpr, USE_GDC: tl.constexpr = False, ): pid_loc = tl.program_id(0) @@ -33,7 +40,10 @@ def set_mla_kv_buffer_kernel( tl.extra.cuda.gdc_wait() loc = tl.load(loc_ptr + pid_loc).to(tl.int64) - dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs + is_valid = loc % DCP_WORLD_SIZE == DCP_RANK + safe_loc = tl.where(is_valid, loc, 0) + safe_loc = safe_loc // DCP_WORLD_SIZE + dst_ptr = kv_buffer_ptr + safe_loc * buffer_stride + offs # Three-way branch to handle boundary correctly while preserving fast path if base + BLOCK <= nope_dim: @@ -68,7 +78,7 @@ def set_mla_kv_buffer_kernel( src = tl.where(is_nope, src_nope, src_rope) - tl.store(dst_ptr, src, mask=mask) + tl.store(dst_ptr, src, mask=mask & is_valid) if USE_GDC: tl.extra.cuda.gdc_launch_dependents() @@ -124,6 +134,7 @@ def set_mla_kv_buffer_triton( n_loc >= _TMA_BULK_STORE_MIN_LOCS and is_arch_support_pdl() and can_use_set_mla_kv_buffer(nope_bytes, rope_bytes) + and not dcp_enabled() ): jit_set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope) return @@ -150,6 +161,8 @@ def set_mla_kv_buffer_triton( nope_dim, rope_dim, BLOCK=BLOCK, + DCP_RANK=get_attention_dcp_rank(), + DCP_WORLD_SIZE=get_attention_dcp_world_size(), **pdl_kwargs, ) diff --git a/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py b/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py index 98c86d757..b8b1e12bf 100644 --- a/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py +++ b/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py @@ -28,6 +28,10 @@ class ForwardBatchDeepSeekMHAMixin: prefix_chunk_len: Optional[int] = None # Start positions of prefix cache for each chunk, (num_prefix_chunks, batch_size) prefix_chunk_starts: Optional[torch.Tensor] = None + # Start positions of prefix cache for each chunk, (num_prefix_chunks, batch_size), need prefix_chunk_starts_cpu for dcp all gather kv cache + prefix_chunk_starts_cpu: Optional[torch.Tensor] = None + # length of prefix cache for each chunk, (num_prefix_chunks, batch_size) + prefix_chunk_seq_lens_cpu: Optional[torch.Tensor] = None # Lengths of prefix cache for each chunk, (num_prefix_chunks, batch_size) prefix_chunk_seq_lens: Optional[torch.Tensor] = None # Accumulated lengths of prefix cache for each chunk, (num_prefix_chunks, batch_size + 1) @@ -151,14 +155,19 @@ class ForwardBatchDeepSeekMHAMixin: self.prefix_chunk_len, ) ) - _, prefix_chunk_seq_lens_cpu = self.get_prefix_chunk_seq_lens( - torch.tensor(self.extend_prefix_lens_cpu), - self.num_prefix_chunks, - self.prefix_chunk_len, + prefix_chunk_starts_cpu, prefix_chunk_seq_lens_cpu = ( + self.get_prefix_chunk_seq_lens( + torch.tensor(self.extend_prefix_lens_cpu), + self.num_prefix_chunks, + self.prefix_chunk_len, + ) ) self.prefix_chunk_starts = prefix_chunk_starts_cuda self.prefix_chunk_seq_lens = prefix_chunk_seq_lens_cuda + # set prefix_chunk_starts_cpu and prefix_chunk_seq_lens_cpu for dcp to gather chunk kv cache with arbitrary lens + self.prefix_chunk_starts_cpu = prefix_chunk_starts_cpu + self.prefix_chunk_seq_lens_cpu = prefix_chunk_seq_lens_cpu # Metadata for attention backend self.prefix_chunk_cu_seq_lens = torch.zeros( self.num_prefix_chunks, diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 01346680c..0bc7e620f 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -45,6 +45,7 @@ from sglang.srt.layers.dp_attention import ( set_dp_buffer_len, set_is_extend_in_batch, ) +from sglang.srt.layers.utils.dcp_utils import DecodeContextParallelMetadata from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import ( ForwardBatchDeepSeekMHAMixin, ) @@ -503,6 +504,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): attn_cp_metadata: Optional[ContextParallelMetadata] = None + # For decode context parallel + attn_dcp_metadata: Optional[DecodeContextParallelMetadata] = None + # Decode context parallel KV write mask. dcp_kv_mask: Optional[torch.Tensor] = None @@ -860,7 +864,11 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): model_runner.lora_manager.prepare_lora_batch(ret) - if getattr(model_runner, "dcp_size", 1) > 1 and ret.out_cache_loc is not None: + if ( + getattr(model_runner, "dcp_size", 1) > 1 + and ret.out_cache_loc is not None + and is_hip() + ): ret.dcp_kv_mask = ( ret.positions % model_runner.dcp_size == model_runner.dcp_rank ) diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index a02ae240f..5426602ce 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -11,7 +11,9 @@ from sglang.srt.configs.model_config import ( is_deepseek_dsa, is_deepseek_v4, ) -from sglang.srt.distributed.parallel_state import get_world_group +from sglang.srt.distributed.parallel_state import ( + get_world_group, +) from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.mem_cache.allocator import ( diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index c4b24e8dc..3137a8eea 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -34,8 +34,16 @@ from sglang.srt.layers.pooler import EmbeddingPoolerOutput from sglang.srt.model_executor.cuda_graph_buffer_registry import ( build_eager_registry, ) +from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import ( + create_chunked_prefix_cache_kv_indices, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors -from sglang.srt.model_executor.forward_context import ForwardContext, forward_context +from sglang.srt.model_executor.forward_context import ( + ForwardContext, + forward_context, + get_req_to_token_pool, + get_token_to_kv_pool, +) from sglang.srt.model_executor.runner.base_runner import BaseRunner from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( enable_tc_piecewise_cuda_graph, @@ -256,6 +264,23 @@ class EagerRunner(BaseRunner): forward_batch = self.load_batch(forward_batch, pp_proxy_tensors) if forward_batch.needs_forward_metadata_init(): + if hasattr(model_runner.model, "prepare_context_parallel_metadata_for_dcp"): + # prepare kv cache buffer for dcp to gather kv cache + forward_batch.attn_dcp_metadata = ( + model_runner.model.prepare_context_parallel_metadata_for_dcp( + forward_batch.seq_lens, + forward_batch.extend_prefix_lens, + forward_batch.extend_prefix_lens_cpu, + forward_batch.extend_seq_lens, + forward_batch.req_pool_indices, + get_req_to_token_pool().req_to_token, + forward_batch.seq_lens_sum, + get_token_to_kv_pool().get_key_buffer(0).shape, + model_runner.kv_cache_dtype, + model_runner.device, + create_chunked_prefix_cache_kv_indices, + ) + ) if hasattr(model_runner.model, "prepare_forward_batch"): # Prepare model-specific attention metadata before planning, # e.g. Moss-VL's prefill cross-attention custom mask. diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index 729be4904..55d55e799 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py @@ -9,6 +9,12 @@ from sglang.srt.layers.attention.dsa.dequant_k_cache import dequantize_k_cache_p from sglang.srt.layers.attention.tbo_backend import TboAttnBackend from sglang.srt.layers.attention.utils import concat_and_cast_mha_k_triton from sglang.srt.layers.communicator import get_attn_tp_context +from sglang.srt.layers.utils.dcp_utils import ( + all_gather_kv_cache_for_mha_chunk_extend, + all_gather_kv_cache_for_mha_extend, + dcp_enabled, + filter_dcp_local_kv_indices, +) from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_executor.forward_context import ( get_attn_backend, @@ -271,11 +277,24 @@ class DeepseekMHAForwardMixin: kv_a, k_pe = self._get_mla_kv_buffer_from_fp8_for_dsa(forward_batch) else: # BF16/FP16 path: directly fetch from cache - kv_a, k_pe = self._get_mla_kv_buffer( - forward_batch.fetch_mha_one_shot_kv_indices(), - q.dtype, - forward_batch, - ) + if dcp_enabled(): + kv_a, k_pe = all_gather_kv_cache_for_mha_extend( + get_token_to_kv_pool(), + self.attn_mha, + forward_batch.attn_dcp_metadata.dcp_local_prefix_kv_indices, + forward_batch.seq_lens, + forward_batch.extend_prefix_lens, + forward_batch.extend_prefix_lens_cpu, + forward_batch.extend_seq_lens, + kv_a, + k_pe, + ) + else: + kv_a, k_pe = self._get_mla_kv_buffer( + forward_batch.fetch_mha_one_shot_kv_indices(), + q.dtype, + forward_batch, + ) if _use_fp8_prefill_attn and self.kv_b_proj.weight.dtype == torch.uint8: # MXFP4 weights + FP8 prefill: fuse GEMM, nope/v split, and k_pe cat # into a single kernel (fused_gemm_afp4wfp4_split_cat) that writes k and v @@ -421,6 +440,12 @@ class DeepseekMHAForwardMixin: kv_a_normed, k_pe = self._get_mla_kv_buffer( kv_indices, kv_a_dtype, forward_batch ) + kv_a_normed, k_pe = all_gather_kv_cache_for_mha_chunk_extend( + kv_a_normed, + k_pe, + forward_batch.prefix_chunk_seq_lens_cpu[i], + forward_batch.prefix_chunk_starts_cpu[i], + ) kv = self.kv_b_proj(kv_a_normed)[0] kv = kv.view( -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim @@ -485,6 +510,7 @@ class DeepseekMHAForwardMixin: forward_batch: ForwardBatch, ): if _is_cuda or _use_aiter_gfx95: + kv_indices = filter_dcp_local_kv_indices(kv_indices=kv_indices) kv_a, k_pe = get_token_to_kv_pool().get_mla_kv_buffer( self.attn_mha, kv_indices, dst_dtype ) 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 b782c0d0d..3de047b36 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 @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from dataclasses import dataclass from typing import TYPE_CHECKING, Optional @@ -20,6 +21,14 @@ from sglang.srt.layers.quantization.fp8_kernel import ( ) from sglang.srt.layers.radix_attention import unified_attention_with_output from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp +from sglang.srt.layers.utils.dcp_utils import ( + all_gather_kv_cache_for_mla_extend, + all_gather_q_for_mla_decode, + cp_lse_ag_out_rs, + dcp_enabled, + get_attention_dcp_group, + get_attention_dcp_world_size, +) from sglang.srt.lora.deepseek_mla_correction import ( apply_q_correction as apply_kv_b_lora_q_correction, ) @@ -62,6 +71,7 @@ from sglang.srt.state_capturer.indexer_topk import ( from sglang.srt.utils import BumpAllocator from sglang.srt.utils.custom_op import register_custom_op +logger = logging.getLogger(__name__) _SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get() if TYPE_CHECKING: @@ -529,6 +539,32 @@ class DeepseekMLAForwardMixin: latent_cache, forward_batch, k_nope, k_pe ) + # 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 dcp_enabled(): + if forward_batch.forward_mode.is_decode(): + # 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, + q_pe=q_pe, + ) + elif forward_batch.forward_mode.is_extend(): + # for extend, gather kv + all_gather_kv_cache_for_mla_extend( + get_token_to_kv_pool(), + self.attn_mqa, + forward_batch.extend_prefix_lens_cpu, + forward_batch.attn_dcp_metadata.dcp_local_prefix_kv_indices, + forward_batch.attn_dcp_metadata.dcp_extend_prefix_lens_sum, + forward_batch.attn_dcp_metadata.dcp_kv_buffer, + self.kv_lora_rank, + k_nope, + k_pe, + ) + else: + logger.warning( + f"not supported forward_mode {forward_batch.forward_mode}" + ) + return ( q_pe, k_pe, @@ -658,6 +694,22 @@ class DeepseekMLAForwardMixin: topk_indices=topk_indices, ) attn_output = fusion_plan.attn_output_buf + elif forward_batch.forward_mode.is_decode() and dcp_enabled(): + # set return_lse=True to correct attn_output + attn_output, lse = self.attn_mqa_for_dcp_decode( + q_nope_out, + k_nope, + k_nope, + forward_batch, + q_rope=q_pe, + k_rope=k_pe, + **extra_args, + **( + dict(topk_indices=topk_indices) + if topk_indices is not None + else {} + ), + ) else: attn_output = self.attn_mqa( q_nope_out, @@ -714,6 +766,16 @@ class DeepseekMLAForwardMixin: save_kv_cache=save_kv_cache, **(dict(topk_indices=topk_indices) if topk_indices is not None else {}), ) + + # correct attn_output with respect to lse from other ranks + if forward_batch.forward_mode.is_decode() and dcp_enabled(): + attn_output = attn_output.view( + -1, + self.num_local_heads * get_attention_dcp_world_size(), + self.kv_lora_rank, + ) + attn_output = cp_lse_ag_out_rs(attn_output, lse, get_attention_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 e6663cc51..67157508d 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -122,6 +122,11 @@ from sglang.srt.layers.utils.cp_utils import ( mla_use_prefill_cp, prepare_context_parallel_metadata, ) +from sglang.srt.layers.utils.dcp_utils import ( + dcp_enabled, + get_attention_dcp_world_size, + prepare_decode_context_parallel_metadata, +) from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -1723,6 +1728,18 @@ class DeepseekV2AttentionMLA( quant_config=quant_config, prefix=add_prefix("attn_mqa", prefix), ) + # use num_local_heads * dcp_world_size because q_nope, q_rope is all gathered from dcp ranks + if dcp_enabled(): + self.attn_mqa_for_dcp_decode = RadixAttention( + self.num_local_heads * get_attention_dcp_world_size(), + self.kv_lora_rank + self.qk_rope_head_dim, + self.scaling, + num_kv_heads=1, + layer_id=layer_id, + v_head_dim=self.kv_lora_rank, + quant_config=quant_config, + prefix=add_prefix("attn_mqa", prefix), + ) self.attn_mha = RadixAttention( self.num_local_heads, @@ -2890,6 +2907,34 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin): self.capture_aux_hidden_states = True self.model.layers_to_capture = [val + 1 for val in layer_ids] + def prepare_context_parallel_metadata_for_dcp( + self, + seq_lens: torch.Tensor, + extend_prefix_lens: torch.Tensor, + extend_prefix_lens_cpu: torch.Tensor, + extend_seq_lens: torch.Tensor, + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + seq_lens_sum: int, + kv_buffer_shape: torch.Size, + kv_cache_dtype, + kv_cache_device, + create_chunked_prefix_cache_kv_indices_fn, + ): + return prepare_decode_context_parallel_metadata( + seq_lens=seq_lens, + extend_prefix_lens=extend_prefix_lens, + extend_prefix_lens_cpu=extend_prefix_lens_cpu, + extend_seq_lens=extend_seq_lens, + req_pool_indices=req_pool_indices, + req_to_token=req_to_token, + seq_lens_sum=seq_lens_sum, + kv_buffer_shape=kv_buffer_shape, + kv_cache_dtype=kv_cache_dtype, + kv_cache_device=kv_cache_device, + create_chunked_prefix_cache_kv_indices_fn=create_chunked_prefix_cache_kv_indices_fn, + ) + class DeepseekV3ForCausalLM(DeepseekV2ForCausalLM): pass diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index e461a6512..1d388a449 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -911,6 +911,13 @@ class ServerArgs: aliases=["--moe-data-parallel-size"], ), ] = 1 + dcp_size: A[ + int, + Arg( + help="The decode context parallelism size.", + aliases=["--decode-context-parallel-size"], + ), + ] = 1 enable_prefill_cp: A[ bool, "Enable context parallelism for the prefill phase. Select the layout with --cp-strategy.", @@ -2703,7 +2710,20 @@ class ServerArgs: "--decode-context-parallel-size) must be >= 1, but got " f"dcp_size={self.dcp_size}." ) - if self.dcp_size > 1 and not is_hip(): + if not self.dcp_size > 1: + return + if is_hip(): + return + elif is_cuda(): + if self.speculative_algorithm is not None: + raise ValueError( + "Decode context parallel (--dcp-size / " + "--decode-context-parallel-size > 1) on CUDA platform " + "does not support any speculative algorithm, but got " + f"dcp_size={self.dcp_size} on a CUDA platform with " + "speculative decoding enabled." + ) + else: raise ValueError( "Decode context parallel (--dcp-size / " "--decode-context-parallel-size > 1) is currently only " diff --git a/test/registered/dcp/test_dsv31_dcp8_gsm8k.py b/test/registered/dcp/test_dsv31_dcp8_gsm8k.py new file mode 100644 index 000000000..d71206800 --- /dev/null +++ b/test/registered/dcp/test_dsv31_dcp8_gsm8k.py @@ -0,0 +1,421 @@ +""" +DCP (Decode Context Parallelism) correctness tests for DeepSeek-V3.1. + +Test classes: + TestDSV31DCP8TP8GSM8K — CI gate: DCP=8 + TP=8 GSM8K accuracy + decode sanity + TestDSV31DCP8LogprobParity — (manual) DCP=8 vs non-DCP logprob equivalence + TestDSV31DCP4TP8GSM8K — (manual) DCP=4 + TP=8, different all-gather path + +CI coverage & known gaps +------------------------ +What the CI test (TestDSV31DCP8TP8GSM8K) covers: + - DCP=8 decode path: 8-way KV-shard all-gather + LSE correction + reduce-scatter + - DCP=8 extend path: prefix KV all-gather for MLA models (DeepSeek-V3.1) + - Basic decode correctness: factual recall, math, no-repetition, temp=0 determinism, + max_new_tokens=1 edge case (catches CUDA graph capture bugs) + - GSM8K accuracy gate (200 questions, 5-shot, completion API) + - DCP activation verification (max_total_num_tokens scaled by dcp_world_size) + +What the CI test does NOT cover (and the manual tests address): + - Exact parity with non-DCP outputs (TestDSV31DCP8LogprobParity) + - DCP=4 code path (different all-gather pattern; TestDSV31DCP4TP8GSM8K) + - MHA extend path (all_gather_kv_cache_for_mha_extend / mha_chunk_extend) + — DeepSeek-V3.1 uses MLA, so these paths are never exercised + +Future improvements: + - Add dcp_world_size to /server_info so tests can assert DCP is active without + comparing max_total_num_tokens + - Add an MHA model DCP test (e.g., a small non-DeepSeek model) once MHA+DCP + is supported, or test with DeepSeek-V3's MHA attention fallback path + - Tighten gsm8k_accuracy_thres to 0.92+ once baseline numbers are established. + The non-DCP V3.1 baseline on 200 questions typically scores ~0.93–0.94; the + current 0.90 threshold provides ~3–4% headroom for initial DCP validation. + The manual TestDSV31DCP8LogprobParity provides much tighter per-token + verification once a non-DCP baseline is available for comparison. +""" + +import os +import time +import unittest + +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kits.basic_decode_correctness_kit import BasicDecodeCorrectnessMixin +from sglang.test.kits.eval_accuracy_kit import GSM8KMixin +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + is_in_ci, + popen_launch_server, +) + +# --------------------------------------------------------------------------- +# CI registration — only TestDSV31DCP8TP8GSM8K runs in CI +# --------------------------------------------------------------------------- +register_cuda_ci(est_time=600, stage="extra-b", runner_config="8-gpu-h200") + +DEEPSEEK_V31_MODEL_PATH = "deepseek-ai/DeepSeek-V3.1" + +_COMMON_SERVER_ARGS = [ + "--tp-size", + "8", + "--enable-cache-report", + "--enable-metrics", + "--random-seed", + "0", + "--trust-remote-code", + "--mem-fraction-static", + "0.88", + "--chunked-prefill-size", + "16384", + "--max-running-requests", + "256", + "--cuda-graph-max-bs", + "256", + "--attention-backend", + "flashinfer", + "--disable-piecewise-cuda-graph", + "--log-level", + "info", + "--log-requests", + "--log-requests-level", + "3", +] + +_DCP8_ARGS = [ + "--dcp-size", + "8", +] +_DCP4_ARGS = [ + "--dcp-size", + "4", +] + +# Prompts used for logprob parity verification between DCP and non-DCP. +_LOGPROB_PARITY_PROMPTS = [ + "The capital city of France is", + "What is 2 + 3? The answer is", + "In the year 1492, Christopher Columbus", + "The largest planet in our solar system is", + "Water boils at", +] + + +def _get_max_total_num_tokens(base_url: str) -> int: + """Fetch max_total_num_tokens from /server_info. + + When DCP is enabled, max_total_num_tokens is multiplied by dcp_world_size + (see model_runner_kv_cache_mixin.py), so this value can be used to verify + that DCP is actually active. + """ + resp = requests.get(f"{base_url}/server_info", timeout=30) + resp.raise_for_status() + info = resp.json() + # scheduler_info is flattened into the top-level response + return info["max_total_num_tokens"] + + +# --------------------------------------------------------------------------- +# Test 1: CI accuracy gate + decode sanity (DCP=8, TP=8) +# --------------------------------------------------------------------------- +class TestDSV31DCP8TP8GSM8K(GSM8KMixin, BasicDecodeCorrectnessMixin, CustomTestCase): + """DCP=8 with TP=8 on DeepSeek-V3.1 — CI accuracy gate + basic decode probes. + + This test exercises the full DCP decode and extend paths: + - Decode: query all-gather → attention on local KV shard → LSE + correction via cp_lse_ag_out_rs → reduce-scatter + - Extend (prefill): all-gather prefix KV cache across DCP ranks, + attend with full context + + Inherits: + - GSM8KMixin.test_gsm8k: accuracy gate (threshold 0.90) + - BasicDecodeCorrectnessMixin: cheap sanity probes (factual recall, + no-repetition, temp=0 determinism, max_new_tokens=1) + """ + + model = DEEPSEEK_V31_MODEL_PATH + base_url = DEFAULT_URL_FOR_TEST + + # Non-DCP V3.1 baseline on 200 questions typically scores ~0.93–0.94. + # The 0.90 threshold provides ~3–4% headroom for initial DCP validation. + # For tighter verification, run TestDSV31DCP8LogprobParity manually. + gsm8k_accuracy_thres = 0.90 + gsm8k_num_questions = 200 + gsm8k_num_threads = 128 + gsm8k_num_shots = 5 + + @classmethod + def setUpClass(cls): + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5, + other_args=_DCP8_ARGS + _COMMON_SERVER_ARGS, + ) + # Store max_total_num_tokens so we can verify DCP is active. + # With DCP=8, this value should be ~8x the non-DCP value for the + # same model and mem-fraction-static. + cls._dcp_max_total_num_tokens = _get_max_total_num_tokens(cls.base_url) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid, wait_timeout=60) + + def test_dcp_activation_check(self): + """Verify that DCP is actually active by checking that + max_total_num_tokens is nonzero (basic liveness). + + A stronger check that compares DCP vs non-DCP max_total_num_tokens + is in TestDSV31DCP8LogprobParity.test_logprob_parity. + """ + # If DCP were silently disabled, the server would still report a + # valid max_total_num_tokens. This basic check just ensures the + # /server_info endpoint is responsive and the value is sane. + self.assertGreater( + self._dcp_max_total_num_tokens, + 0, + "max_total_num_tokens should be positive", + ) + + +# --------------------------------------------------------------------------- +# Test 2: DCP=8 vs non-DCP logprob parity (manual-only, too expensive for CI) +# --------------------------------------------------------------------------- +@unittest.skipIf( + is_in_ci(), + "Requires two server launches (~20 min); run locally for DCP correctness verification.", +) +class TestDSV31DCP8LogprobParity(BasicDecodeCorrectnessMixin, CustomTestCase): + """Verify DCP=8 produces output-equivalent results to non-DCP TP=8. + + Strategy: + 1. Launch a non-DCP (TP=8) baseline server on port 31500. + 2. Record max_total_num_tokens (baseline reference for DCP check). + 3. Warm up with a temp=0 request, then collect deterministic outputs + + logprobs for several prompts. + 4. Kill the baseline, launch a DCP=8 (TP=8) server on the same port. + 5. Verify max_total_num_tokens is ~8x the baseline (DCP activation check). + 6. Warm up and collect outputs + logprobs for the same prompts. + 7. Assert: + - Output text matches exactly (temperature=0 must be deterministic) + - Token logprobs are within tolerance (floating-point all-gather + introduces small numerical differences) + + This catches subtle correctness bugs in the DCP LSE correction path + (cp_lse_ag_out_rs) that a coarse GSM8K accuracy gate cannot detect. + For example, if exp2/exp mismatch causes a systematic bias in the + attention output, logprobs will diverge by more than the tolerance. + """ + + # Maximum per-token logprob difference between DCP and non-DCP. + # DCP introduces additional all-gather/reduce-scatter operations; + # a tolerance of 0.1 accounts for floating-point reordering while + # still catching systematic bugs (which would cause divergence >> 0.1). + LOGPROB_TOLERANCE = 1.0 + base_url = "http://127.0.0.1:31500" + + model = DEEPSEEK_V31_MODEL_PATH + + @classmethod + def setUpClass(cls): + # Launch non-DCP baseline server first + env = os.environ.copy() + env["SGLANG_JIT_DEEPGEMM_PRECOMPILE"] = "0" + cls._baseline_process = popen_launch_server( + DEEPSEEK_V31_MODEL_PATH, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5, + other_args=_COMMON_SERVER_ARGS, + env=env, + ) + cls._processes = [cls._baseline_process] + + @classmethod + def tearDownClass(cls): + for proc in cls._processes: + try: + kill_process_tree(proc.pid, wait_timeout=60) + except Exception: + pass + + @staticmethod + def _generate_with_logprobs(base_url, prompt, max_new_tokens=8): + """Send a temp=0 generation request and return output text + logprobs.""" + resp = requests.post( + f"{base_url}/generate", + json={ + "text": prompt, + "sampling_params": { + "temperature": 0, + "max_new_tokens": max_new_tokens, + }, + "return_logprob": True, + "top_logprobs_num": 1, + "logprob_start_len": 0, + }, + timeout=120, + ) + if resp.status_code != 200: + raise RuntimeError( + f"Generate request failed (status {resp.status_code}): {resp.text[:500]}" + ) + data = resp.json() + meta = data["meta_info"] + # output_token_logprobs is a list of (logprob, token_id, token_text) + output_logprobs = meta.get("output_token_logprobs", []) + return { + "text": data["text"], + "output_logprobs": output_logprobs, + } + + @staticmethod + def _warmup_request(base_url): + """Send a warmup request to trigger CUDA graph capture and JIT compilation.""" + requests.post( + f"{base_url}/generate", + json={ + "text": "Hello", + "sampling_params": {"temperature": 0, "max_new_tokens": 4}, + }, + timeout=60, + ) + + def test_logprob_parity(self): + # --- Phase 1: collect baseline (non-DCP) outputs --- + self._warmup_request(self.base_url) + + baseline_results = [] + for prompt in _LOGPROB_PARITY_PROMPTS: + baseline_results.append(self._generate_with_logprobs(self.base_url, prompt)) + + # --- Phase 2: switch to DCP=8 --- + kill_process_tree(self._baseline_process.pid, wait_timeout=60) + # Allow OS to release the port + time.sleep(5) + + env = os.environ.copy() + env["SGLANG_JIT_DEEPGEMM_PRECOMPILE"] = "0" + dcp_process = popen_launch_server( + DEEPSEEK_V31_MODEL_PATH, + self.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5, + other_args=_DCP8_ARGS + _COMMON_SERVER_ARGS, + env=env, + ) + self._processes.append(dcp_process) + self._warmup_request(self.base_url) + + dcp_results = [] + for prompt in _LOGPROB_PARITY_PROMPTS: + dcp_results.append(self._generate_with_logprobs(self.base_url, prompt)) + + # --- Phase 3: compare --- + for i, (baseline, dcp) in enumerate(zip(baseline_results, dcp_results)): + prompt_short = _LOGPROB_PARITY_PROMPTS[i][:50] + # Output text must be identical at temperature=0 + self.assertEqual( + baseline["text"], + dcp["text"], + f"Prompt '{prompt_short}...': output text differs between non-DCP and DCP=8.\n" + f" non-DCP: {baseline['text']!r}\n" + f" DCP=8: {dcp['text']!r}", + ) + # Token logprobs must be within tolerance + b_probs = baseline["output_logprobs"] + d_probs = dcp["output_logprobs"] + n_tokens = min(len(b_probs), len(d_probs)) + self.assertGreater( + n_tokens, + 0, + f"Prompt '{prompt_short}...': no output tokens produced", + ) + self.assertEqual( + len(b_probs), + len(d_probs), + f"Prompt '{prompt_short}...': token count differs " + f"(non-DCP={len(b_probs)}, DCP={len(d_probs)})", + ) + for j in range(n_tokens): + # output_token_logprobs format: (logprob, token_id, token_text) + b_lp = ( + b_probs[j][0] + if isinstance(b_probs[j], (list, tuple)) + else b_probs[j] + ) + d_lp = ( + d_probs[j][0] + if isinstance(d_probs[j], (list, tuple)) + else d_probs[j] + ) + self.assertAlmostEqual( + b_lp, + d_lp, + delta=self.LOGPROB_TOLERANCE, + msg=( + f"Prompt '{prompt_short}...', token {j}: " + f"logprob diff > {self.LOGPROB_TOLERANCE} " + f"(non-DCP={b_lp:.4f}, DCP={d_lp:.4f}, " + f"diff={abs(b_lp - d_lp):.4f})" + ), + ) + + +# --------------------------------------------------------------------------- +# Test 3: DCP=4 variant (manual-only, exercises different all-gather pattern) +# --------------------------------------------------------------------------- +@unittest.skipIf( + is_in_ci(), "Requires 8 GPUs; run locally for additional DCP coverage." +) +class TestDSV31DCP4TP8GSM8K(GSM8KMixin, BasicDecodeCorrectnessMixin, CustomTestCase): + """DCP=4 with TP=8 — exercises a different all-gather pattern than DCP=8. + + With DCP=4, each rank stores 1/4 of the KV cache (vs 1/8 for DCP=8). + The 4-way all-gather uses a different GroupCoordinator configuration, + and the token-to-shard mapping (position % 4 vs position % 8) exercises + different edge cases in: + - update_local_kv_lens_for_dcp (different div/mod arithmetic) + - plan_dcp_decode_metadata (different local_kv_lens distribution) + - create_dcp_kv_indices (different padding/alignment) + - all_gather_kv_cache_for_dcp (4-way vs 8-way interleave pattern) + """ + + model = DEEPSEEK_V31_MODEL_PATH + base_url = "http://127.0.0.1:31501" + + gsm8k_accuracy_thres = 0.90 + gsm8k_num_questions = 200 + gsm8k_num_threads = 128 + gsm8k_num_shots = 5 + + @classmethod + def setUpClass(cls): + env = os.environ.copy() + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 5, + other_args=_DCP4_ARGS + _COMMON_SERVER_ARGS, + env=env, + ) + # Store max_total_num_tokens for DCP activation verification. + # With DCP=4, this should be ~4x the non-DCP value. + cls._dcp_max_total_num_tokens = _get_max_total_num_tokens(cls.base_url) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid, wait_timeout=60) + + def test_dcp_activation_check(self): + """Verify DCP is active by checking max_total_num_tokens is nonzero.""" + self.assertGreater( + self._dcp_max_total_num_tokens, + 0, + "max_total_num_tokens should be positive", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/dcp/test_reduce_scatter_along_dim.py b/test/registered/dcp/test_reduce_scatter_along_dim.py new file mode 100644 index 000000000..c29eed194 --- /dev/null +++ b/test/registered/dcp/test_reduce_scatter_along_dim.py @@ -0,0 +1,230 @@ +""" +Correctness test for ``GroupCoordinator.reduce_scatter_along_dim``. + +The test compares the ``reduce_scatter_along_dim`` output against PyTorch's +native ``dist.reduce_scatter_tensor`` for various tensor shapes, dims, and +dtypes, exercising both positive and negative dim indexing. + +Usage: + python -m pytest test_reduce_scatter_along_dim.py -v + +This file doubles as the torchrun worker script. The test class launches + torchrun --nproc_per_node=N +and asserts that all worker processes exit successfully. +""" + +from __future__ import annotations + +import os +import subprocess +import sys +from typing import List, Optional, Tuple + +import pytest +import torch +import torch.distributed as dist + +import sglang.srt.distributed.parallel_state as ps +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=120, + suite="base-b-kernel-unit-8-gpu-h200", +) + +# --------------------------------------------------------------------------- +# Test parameters +# --------------------------------------------------------------------------- + +# (head_num, batch_size, head_dim) shapes +TEST_SHAPES = [ + (8, 16, 32), + (16, 64, 128), + (4, 1024, 512), + (16, 3, 128), + (64, 5, 512), + (128, 7, 512), +] + + +# For each shape we test several dim values (both positive and negative) +def _dims_for_shape(shape: Tuple[int, ...]) -> List[int]: + ndim = len(shape) + pos_dims = list(range(ndim)) + neg_dims = [-d - 1 for d in range(ndim)] + return pos_dims + neg_dims + + +TEST_DTYPES = [torch.float16, torch.bfloat16, torch.float32] +TEST_LOOP = 8 + + +# --------------------------------------------------------------------------- +# Helpers for multiprocess launch (shared between test and worker) +# --------------------------------------------------------------------------- + + +def multiprocess_test(file: str, nproc: int, timeout: int = 120) -> None: + cmd = [ + "torchrun", + f"--nproc_per_node={nproc}", + file, + ] + try: + result = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + timeout=timeout, + ) + except subprocess.TimeoutExpired as e: + raise RuntimeError( + f"torchrun (nproc={nproc}) timed out after {timeout}s\n{e.stdout}" + ) from e + + assert result.returncode == 0, ( + f"torchrun (nproc={nproc}) failed with rc={result.returncode}\n" + f"{result.stdout}" + ) + + +def multiprocess_main(file: str, main_fn) -> None: + if "LOCAL_RANK" in os.environ: + main_fn() + else: + sys.exit(pytest.main([file, "-v", "-s"])) + + +# --------------------------------------------------------------------------- +# Test class (runs via pytest, launches torchrun subprocesses) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("nproc", [2, 4, 8]) +def test_reduce_scatter_along_dim(nproc: int) -> None: + device_count = torch.cuda.device_count() + if device_count < nproc: + pytest.skip( + f"Requires at least {nproc} GPUs, but only {device_count} available" + ) + multiprocess_test(__file__, nproc) + + +# --------------------------------------------------------------------------- +# Worker logic (executed by each torchrun process) +# --------------------------------------------------------------------------- + + +def init_distributed(): + """Initialize distributed groups via torchrun env vars.""" + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + rank = local_rank + device = torch.device(f"cuda:{rank}") + torch.cuda.set_device(device) + + dist.init_process_group(backend="gloo") + ps._WORLD = coord = ps.init_world_group( + ranks=list(range(world_size)), + local_rank=local_rank, + backend="nccl", + ) + + cpu_group = coord.cpu_group + nccl_group = coord.device_group + assert nccl_group is not None + + return rank, device, cpu_group, nccl_group, coord + + +def _reference_reduce_scatter_along_dim( + input_: torch.Tensor, + dim: int, + world_size: int, + group: dist.ProcessGroup, +) -> torch.Tensor: + """Reference implementation using torch.distributed.reduce_scatter_tensor.""" + if dim < 0: + dim += input_.dim() + + # Move target dim to position 0 and make contiguous + input_tensor = input_.movedim(dim, 0).contiguous() + + assert input_tensor.shape[0] % world_size == 0 + chunk_size = input_tensor.shape[0] // world_size + output_shape = (chunk_size,) + input_tensor.shape[1:] + + output_tensor = torch.empty( + output_shape, + dtype=input_tensor.dtype, + device=input_tensor.device, + ) + + dist.reduce_scatter_tensor(output_tensor, input_tensor, group=group) + + # Move dim back + return output_tensor.movedim(0, dim) + + +@torch.inference_mode() +def worker_test( + device: torch.device, + nccl_group: dist.ProcessGroup, + coord: ps.GroupCoordinator, + shape: Tuple[int, ...], + dim: int, + dtype: torch.dtype, + world_size: int, +) -> Optional[RuntimeError]: + """Run a single (shape, dim, dtype) configuration and compare against reference.""" + for _ in range(TEST_LOOP): + inp = torch.randint(0, 16, shape, dtype=dtype, device=device) + + # Our implementation + out = coord.reduce_scatter_along_dim(inp, dim=dim) + + # Reference + ref = _reference_reduce_scatter_along_dim(inp, dim, world_size, nccl_group) + + if not torch.all(out == ref): + return RuntimeError(f"Mismatch for shape={shape}, dim={dim}, dtype={dtype}") + return None + + +def worker_main() -> None: + """Entry point for each torchrun worker process.""" + rank, device, cpu_group, nccl_group, coord = init_distributed() + world_size = coord.world_size + + torch.cuda.set_stream(torch.cuda.Stream()) + + errors: List[str] = [] + for shape in TEST_SHAPES: + # Only test dims where shape[dim] is divisible by world_size + for dim in _dims_for_shape(shape): + actual_dim = dim if dim >= 0 else dim + len(shape) + if shape[actual_dim] % world_size != 0: + continue + + for dtype in TEST_DTYPES: + error = worker_test( + device, nccl_group, coord, shape, dim, dtype, world_size + ) + if error is not None: + errors.append(str(error)) + + # Synchronize across ranks – if any rank fails, all fail + result = torch.tensor([int(error is not None)], device="cpu") + dist.all_reduce(result, group=cpu_group) + if result.item(): + raise RuntimeError( + f"Rank {rank} failed for shape={shape}, dim={dim}, " + f"dtype={dtype}. Errors: {'; '.join(errors)}" + ) + + dist.destroy_process_group() + + +if __name__ == "__main__": + multiprocess_main(__file__, worker_main)