Add DCP to runtime parallel context (#30478)

This commit is contained in:
Baizhou Zhang
2026-07-10 21:23:41 -07:00
committed by GitHub
parent bbcfcaeefe
commit 07165d5daa
16 changed files with 167 additions and 105 deletions
@@ -5,11 +5,7 @@ import triton
import triton.language as tl import triton.language as tl
from sglang.jit_kernel.utils import is_arch_support_pdl from sglang.jit_kernel.utils import is_arch_support_pdl
from sglang.srt.layers.dcp import ( from sglang.srt.runtime_context import get_parallel
dcp_enabled,
get_attention_dcp_rank,
get_attention_dcp_world_size,
)
@triton.jit @triton.jit
@@ -134,7 +130,7 @@ def set_mla_kv_buffer_triton(
n_loc >= _TMA_BULK_STORE_MIN_LOCS n_loc >= _TMA_BULK_STORE_MIN_LOCS
and is_arch_support_pdl() and is_arch_support_pdl()
and can_use_set_mla_kv_buffer(nope_bytes, rope_bytes) and can_use_set_mla_kv_buffer(nope_bytes, rope_bytes)
and not dcp_enabled() and not get_parallel().dcp_enabled
): ):
jit_set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope) jit_set_mla_kv_buffer(kv_buffer, loc, cache_k_nope, cache_k_rope)
return return
@@ -161,8 +157,8 @@ def set_mla_kv_buffer_triton(
nope_dim, nope_dim,
rope_dim, rope_dim,
BLOCK=BLOCK, BLOCK=BLOCK,
DCP_RANK=get_attention_dcp_rank(), DCP_RANK=get_parallel().attn_dcp_rank,
DCP_WORLD_SIZE=get_attention_dcp_world_size(), DCP_WORLD_SIZE=get_parallel().attn_dcp_size,
**pdl_kwargs, **pdl_kwargs,
) )
@@ -25,8 +25,6 @@ from sglang.srt.layers.attention.flashinfer_backend import (
from sglang.srt.layers.attention.utils import assert_buffer_fits from sglang.srt.layers.attention.utils import assert_buffer_fits
from sglang.srt.layers.dcp import ( from sglang.srt.layers.dcp import (
DecodeContextParallelMetadata, DecodeContextParallelMetadata,
dcp_enabled,
get_attention_dcp_world_size,
update_local_kv_lens_for_dcp, update_local_kv_lens_for_dcp,
) )
from sglang.srt.layers.dcp.planner import plan_dcp_decode_metadata from sglang.srt.layers.dcp.planner import plan_dcp_decode_metadata
@@ -647,7 +645,9 @@ class FlashInferMLAAttnBackend(AttentionBackend):
k_buffer[:, :, layer.v_head_dim :], k_buffer[:, :, layer.v_head_dim :],
out=o, 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 # 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(), return_lse=(
forward_batch.forward_mode.is_decode() and get_parallel().dcp_enabled
),
) )
if isinstance(o, tuple): if isinstance(o, tuple):
out, lse = o out, lse = o
@@ -662,7 +662,7 @@ class FlashInferMLAIndicesUpdaterDecode:
self.num_local_heads = ( self.num_local_heads = (
model_runner.model_config.num_attention_heads model_runner.model_config.num_attention_heads
// get_parallel().attn_tp_size // get_parallel().attn_tp_size
* get_attention_dcp_world_size() * get_parallel().attn_dcp_size
) )
self.kv_lora_rank = model_runner.model_config.kv_lora_rank self.kv_lora_rank = model_runner.model_config.kv_lora_rank
self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim self.qk_nope_head_dim = model_runner.model_config.qk_nope_head_dim
@@ -733,7 +733,7 @@ class FlashInferMLAIndicesUpdaterDecode:
self.req_to_token.shape[1], self.req_to_token.shape[1],
) )
if dcp_enabled(): if get_parallel().dcp_enabled:
plan_dcp_decode_metadata( plan_dcp_decode_metadata(
kv_lens, kv_lens,
kv_indptr, kv_indptr,
@@ -17,11 +17,6 @@ from sglang.srt.layers.attention.utils import (
create_flashmla_kv_indices_triton, create_flashmla_kv_indices_triton,
get_num_kv_index_blocks_flashmla, get_num_kv_index_blocks_flashmla,
) )
from sglang.srt.layers.dcp import (
dcp_enabled,
get_attention_dcp_rank,
get_attention_dcp_world_size,
)
from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.runtime_context import get_parallel from sglang.srt.runtime_context import get_parallel
@@ -96,8 +91,8 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
self.cuda_graph_num_splits_view = None self.cuda_graph_num_splits_view = None
# get dcp info # get dcp info
self.dcp_world_size = get_attention_dcp_world_size() self.dcp_world_size = get_parallel().attn_dcp_size
self.dcp_rank = get_attention_dcp_rank() self.dcp_rank = get_parallel().attn_dcp_rank
def init_forward_metadata_out_graph( def init_forward_metadata_out_graph(
self, self,
@@ -390,7 +385,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
# TODO uniform output for forward_decode and forward_extend to # TODO uniform output for forward_decode and forward_extend to
# return tuple instead of single output # return tuple instead of single output
# decode context parallel needs lse to correct attn_output via online softmax # decode context parallel needs lse to correct attn_output via online softmax
if dcp_enabled(): if get_parallel().dcp_enabled:
return o, lse return o, lse
return o return o
@@ -14,7 +14,6 @@ from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
from sglang.srt.distributed.parallel_state import get_dcp_group
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.dcp import ( from sglang.srt.layers.dcp import (
@@ -166,8 +165,8 @@ class TritonAttnBackend(AttentionBackend):
and self.topk == 1 and self.topk == 1
) )
self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA
self.dcp_size = getattr(model_runner, "dcp_size", 1) self.dcp_size = get_parallel().attn_dcp_size
self.dcp_rank = getattr(model_runner, "dcp_rank", 0) self.dcp_rank = get_parallel().attn_dcp_rank
self.num_head = ( self.num_head = (
model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size model_runner.model_config.num_attention_heads // get_parallel().attn_tp_size
) * self.dcp_size ) * self.dcp_size
@@ -1387,7 +1386,7 @@ class TritonAttnBackend(AttentionBackend):
"DCP Triton extend does not support sliding window" "DCP Triton extend does not support sliding window"
) )
group = get_dcp_group() group = get_parallel().dcp_group
q_local = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim).contiguous() q_local = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim).contiguous()
total_tokens, local_heads, _ = q_local.shape total_tokens, local_heads, _ = q_local.shape
@@ -1712,7 +1711,7 @@ class TritonAttnBackend(AttentionBackend):
attn_logits = self.forward_metadata.swa_attn_logits attn_logits = self.forward_metadata.swa_attn_logits
if self.dcp_size > 1: if self.dcp_size > 1:
group = get_dcp_group() group = get_parallel().dcp_group
with use_symmetric_memory(group): with use_symmetric_memory(group):
q_for_decode = q.view( q_for_decode = q.view(
-1, layer.tp_q_head_num, layer.qk_head_dim -1, layer.tp_q_head_num, layer.qk_head_dim
+5 -1
View File
@@ -23,7 +23,11 @@ Only the symbols imported by code OUTSIDE this subpackage are re-exported here.
Package-internal helpers (the @triton.jit kernels, ``CPTritonContext``, Package-internal helpers (the @triton.jit kernels, ``CPTritonContext``,
``correct_attn_out``, ``create_dcp_kv_indices``, ``update_kv_lens_and_indices``, ``correct_attn_out``, ``create_dcp_kv_indices``, ``update_kv_lens_and_indices``,
``_all_gather_dcp_kv_cache``) stay private to their submodules — import them from ``_all_gather_dcp_kv_cache``) stay private to their submodules — import them from
``sglang.srt.layers.dcp.{kernels,comm}`` if ever needed internally.""" ``sglang.srt.layers.dcp.{kernels,comm}`` if ever needed internally.
``dcp_enabled`` / ``get_attention_dcp_*`` remain compatibility exports for
out-of-tree callers; in-tree code should use ``get_parallel().dcp_enabled`` and
``get_parallel().attn_dcp_*``."""
from sglang.srt.layers.dcp.comm import ( from sglang.srt.layers.dcp.comm import (
all_gather_kv_cache_for_dcp, all_gather_kv_cache_for_dcp,
+34 -28
View File
@@ -20,6 +20,7 @@ PR #25090 vs #14194):
- cp_lse_ag_out_rs_mla: Triton (log2/exp2) correction / reduce-scatter - cp_lse_ag_out_rs_mla: Triton (log2/exp2) correction / reduce-scatter
""" """
import warnings
from typing import Optional from typing import Optional
import torch import torch
@@ -27,37 +28,39 @@ import torch
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
from sglang.srt.distributed.parallel_state import ( from sglang.srt.distributed.parallel_state import GroupCoordinator
GroupCoordinator,
get_dcp_group,
get_dcp_group_no_assert,
)
from sglang.srt.layers.dcp.kernels import CPTritonContext, correct_attn_out from sglang.srt.layers.dcp.kernels import CPTritonContext, correct_attn_out
from sglang.srt.runtime_context import get_parallel from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import is_cuda
def _warn_deprecated_dcp_accessor(name: str, replacement: str) -> None:
warnings.warn(
f"{name} is deprecated; use {replacement} instead.",
DeprecationWarning,
stacklevel=2,
)
def dcp_enabled() -> bool: def dcp_enabled() -> bool:
""" """Deprecated: use ``get_parallel().dcp_enabled``."""
only checks whether dcp enabled for cuda platform _warn_deprecated_dcp_accessor("dcp_enabled()", "get_parallel().dcp_enabled")
""" return get_parallel().dcp_enabled
if get_dcp_group_no_assert() is None:
return False
if not is_cuda():
return False
return get_parallel().dcp_size > 1
def get_attention_dcp_world_size() -> int: def get_attention_dcp_world_size() -> int:
if not dcp_enabled(): """Deprecated: use ``get_parallel().attn_dcp_size``."""
return 1 _warn_deprecated_dcp_accessor(
return get_parallel().dcp_size "get_attention_dcp_world_size()", "get_parallel().attn_dcp_size"
)
return get_parallel().attn_dcp_size
def get_attention_dcp_rank() -> int: def get_attention_dcp_rank() -> int:
if not dcp_enabled(): """Deprecated: use ``get_parallel().attn_dcp_rank``."""
return 0 _warn_deprecated_dcp_accessor(
return get_parallel().dcp_rank "get_attention_dcp_rank()", "get_parallel().attn_dcp_rank"
)
return get_parallel().attn_dcp_rank
def _ag_lse(cp_attn_lse: torch.Tensor, cp_group: GroupCoordinator) -> torch.Tensor: def _ag_lse(cp_attn_lse: torch.Tensor, cp_group: GroupCoordinator) -> torch.Tensor:
@@ -132,12 +135,13 @@ def cp_lse_ag_out_rs_mla(
def _all_gather_dcp_kv_cache(kv_a: torch.Tensor): def _all_gather_dcp_kv_cache(kv_a: torch.Tensor):
dcp_world_size = get_parallel().dcp_size parallel = get_parallel()
dcp_world_size = parallel.dcp_size
# not use symmetric_memory unless torch mem_pool updated, see https://github.com/pytorch/pytorch/issues/178138 # not use symmetric_memory unless torch mem_pool updated, see https://github.com/pytorch/pytorch/issues/178138
gathered_kv_a = kv_a.new_empty( gathered_kv_a = kv_a.new_empty(
(kv_a.shape[0] * dcp_world_size, *kv_a.shape[1:]), (kv_a.shape[0] * dcp_world_size, *kv_a.shape[1:]),
) )
get_dcp_group().all_gather_into_tensor(gathered_kv_a, kv_a) parallel.dcp_group.all_gather_into_tensor(gathered_kv_a, kv_a)
gathered_kv_a = ( gathered_kv_a = (
gathered_kv_a.reshape((dcp_world_size,) + kv_a.shape) gathered_kv_a.reshape((dcp_world_size,) + kv_a.shape)
.transpose(0, 1) .transpose(0, 1)
@@ -152,7 +156,7 @@ def all_gather_kv_cache_for_mha_chunk_extend(
prefix_kv_lens_cpu: torch.Tensor, prefix_kv_lens_cpu: torch.Tensor,
prefix_starts_cpu: torch.Tensor = None, prefix_starts_cpu: torch.Tensor = None,
): ):
if dcp_enabled(): if get_parallel().dcp_enabled:
kv_a = kv_a.unsqueeze(1) kv_a = kv_a.unsqueeze(1)
gathered_kv = all_gather_kv_cache_for_dcp( gathered_kv = all_gather_kv_cache_for_dcp(
kv_a, kv_a,
@@ -218,10 +222,11 @@ def all_gather_q_for_mla_decode(
q_nope_out: torch.Tensor, q_nope_out: torch.Tensor,
q_pe: torch.Tensor, q_pe: torch.Tensor,
): ):
with use_symmetric_memory(get_dcp_group()): group = get_parallel().dcp_group
with use_symmetric_memory(group):
# transpose q_pe and q_nope_out from [B, H, L] to [H, B, L] # 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) 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) gathered = group.all_gather(combined, dim=0)
d_pe = q_pe.size(-1) d_pe = q_pe.size(-1)
d_nope = q_nope_out.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_nope_out = gathered.split([d_pe, d_nope], dim=-1)
@@ -278,11 +283,12 @@ def all_gather_kv_cache_for_dcp(
""" """
prefix_kv_a and prefix_k_pe should have same shape, expect for last dim prefix_kv_a and prefix_k_pe should have same shape, expect for last dim
""" """
if not dcp_enabled(): parallel = get_parallel()
if not parallel.dcp_enabled:
return torch.cat([prefix_kv_a, prefix_k_pe], dim=-1) return torch.cat([prefix_kv_a, prefix_k_pe], dim=-1)
# 1. compute max kv_lens for each seq # 1. compute max kv_lens for each seq
dcp_world_size = get_parallel().dcp_size dcp_world_size = parallel.dcp_size
dcp_rank = get_parallel().dcp_rank dcp_rank = parallel.dcp_rank
if prefix_starts_cpu is None: if prefix_starts_cpu is None:
prefix_starts_cpu = torch.zeros_like(prefix_kv_lens_cpu) prefix_starts_cpu = torch.zeros_like(prefix_kv_lens_cpu)
+7 -8
View File
@@ -17,7 +17,6 @@ the owner-rule local-index filter."""
import torch import torch
from sglang.srt.layers.dcp.comm import dcp_enabled
from sglang.srt.runtime_context import get_parallel from sglang.srt.runtime_context import get_parallel
@@ -43,10 +42,11 @@ def get_dcp_lens(
def filter_dcp_local_kv_indices(kv_indices: torch.Tensor): def filter_dcp_local_kv_indices(kv_indices: torch.Tensor):
if dcp_enabled(): parallel = get_parallel()
if parallel.dcp_enabled:
kv_indices = ( kv_indices = (
kv_indices[kv_indices % get_parallel().dcp_size == get_parallel().dcp_rank] kv_indices[kv_indices % parallel.dcp_size == parallel.dcp_rank]
// get_parallel().dcp_size // parallel.dcp_size
) )
return kv_indices return kv_indices
@@ -59,8 +59,7 @@ def update_local_kv_lens_for_dcp(kv_len_arr):
in-place mutation because callers (plan_dcp_decode_metadata, the FlashInfer-MLA in-place mutation because callers (plan_dcp_decode_metadata, the FlashInfer-MLA
cuda-graph replay path) rely on it. cuda-graph replay path) rely on it.
""" """
if not dcp_enabled(): parallel = get_parallel()
if not parallel.dcp_enabled:
return return
kv_len_arr.copy_( kv_len_arr.copy_(get_dcp_lens(kv_len_arr, parallel.dcp_size, parallel.dcp_rank))
get_dcp_lens(kv_len_arr, get_parallel().dcp_size, get_parallel().dcp_rank)
)
+8 -7
View File
@@ -20,7 +20,6 @@ from typing import Optional
import torch import torch
from sglang.srt.layers.dcp.comm import dcp_enabled
from sglang.srt.layers.dcp.kernels import ( from sglang.srt.layers.dcp.kernels import (
create_dcp_kv_indices, create_dcp_kv_indices,
update_kv_lens_and_indices, update_kv_lens_and_indices,
@@ -43,7 +42,8 @@ def prepare_decode_context_parallel_metadata(
kv_cache_device, kv_cache_device,
create_chunked_prefix_cache_kv_indices_fn, create_chunked_prefix_cache_kv_indices_fn,
) -> Optional[DecodeContextParallelMetadata]: ) -> Optional[DecodeContextParallelMetadata]:
if not dcp_enabled(): parallel = get_parallel()
if not parallel.dcp_enabled:
return None return None
# dcp_kv_buffer tokens' layout # dcp_kv_buffer tokens' layout
# [ rank0_r1.prefix_tokens, rank1_r1.prefix_tokens, ..., rank7_r1.prefix_tokens, # [ rank0_r1.prefix_tokens, rank1_r1.prefix_tokens, ..., rank7_r1.prefix_tokens,
@@ -107,13 +107,13 @@ def prepare_decode_context_parallel_metadata(
extend_cu_prefix_lens, extend_cu_prefix_lens,
dcp_kv_indices, dcp_kv_indices,
extend_prefix_lens_sum, extend_prefix_lens_sum,
get_parallel().dcp_size, parallel.dcp_size,
) )
dcp_local_prefix_kv_indices = ( dcp_local_prefix_kv_indices = (
dcp_prefix_kv_indices[ dcp_prefix_kv_indices[
dcp_prefix_kv_indices % get_parallel().dcp_size == get_parallel().dcp_rank dcp_prefix_kv_indices % parallel.dcp_size == parallel.dcp_rank
] ]
// get_parallel().dcp_size // parallel.dcp_size
) )
dcp_kv_buffer = torch.empty( dcp_kv_buffer = torch.empty(
( (
@@ -141,6 +141,7 @@ def plan_dcp_decode_metadata(
fast_decode_kwargs: dict, fast_decode_kwargs: dict,
bs: int, bs: int,
): ):
parallel = get_parallel()
local_kv_lens = kv_lens.clone() local_kv_lens = kv_lens.clone()
update_local_kv_lens_for_dcp(local_kv_lens) update_local_kv_lens_for_dcp(local_kv_lens)
local_kv_lens.clamp_(min=0) local_kv_lens.clamp_(min=0)
@@ -178,8 +179,8 @@ def plan_dcp_decode_metadata(
local_kv_lens, local_kv_lens,
local_kv_lens_cumsum, local_kv_lens_cumsum,
local_kv_indices, local_kv_indices,
dcp_rank=get_parallel().dcp_rank, dcp_rank=parallel.dcp_rank,
dcp_world_size=get_parallel().dcp_size, dcp_world_size=parallel.dcp_size,
BLOCK_SIZE=BLOCK_SIZE, BLOCK_SIZE=BLOCK_SIZE,
) )
kv_indices[:total_local_len] = local_kv_indices[:total_local_len] kv_indices[:total_local_len] = local_kv_indices[:total_local_len]
@@ -25,11 +25,11 @@ from typing import TYPE_CHECKING
from sglang.srt.configs.model_config import ModelImpl from sglang.srt.configs.model_config import ModelImpl
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.dcp import dcp_enabled
from sglang.srt.managers.mm_utils import init_mm_embedding_cache 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.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.registry import TreeCacheBuildContext, create_tree_cache from sglang.srt.mem_cache.registry import TreeCacheBuildContext, create_tree_cache
from sglang.srt.model_loader.utils import get_resolved_model_impl from sglang.srt.model_loader.utils import get_resolved_model_impl
from sglang.srt.runtime_context import get_parallel
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -204,7 +204,9 @@ def build_kv_cache(
# TreeCache.page_size should keep the same as allocator.page_size to # TreeCache.page_size should keep the same as allocator.page_size to
# avoid kv page eviction conflicts. # avoid kv page eviction conflicts.
page_size=( page_size=(
page_size if not dcp_enabled() else token_to_kv_pool_allocator.page_size page_size
if not get_parallel().dcp_enabled
else token_to_kv_pool_allocator.page_size
), ),
is_eagle=spec_algorithm.is_eagle(), is_eagle=spec_algorithm.is_eagle(),
tp_cache_group=( tp_cache_group=(
+4 -9
View File
@@ -50,11 +50,6 @@ from sglang.srt.layers.attention.dsa.quant_k_cache import (
quantize_k_cache_separate, quantize_k_cache_separate,
) )
from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa from sglang.srt.layers.attention.dsa.utils import aiter_can_use_preshuffle_paged_mqa
from sglang.srt.layers.dcp import (
dcp_enabled,
get_attention_dcp_rank,
get_attention_dcp_world_size,
)
from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz 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.radix_attention import RadixAttention
from sglang.srt.mem_cache.allocator.mamba import MambaSlotAllocator from sglang.srt.mem_cache.allocator.mamba import MambaSlotAllocator
@@ -73,6 +68,7 @@ from sglang.srt.mem_cache.utils import (
set_mla_kv_scale_buffer_triton, set_mla_kv_scale_buffer_triton,
) )
from sglang.srt.platforms import current_platform from sglang.srt.platforms import current_platform
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import ( from sglang.srt.utils import (
cpu_has_amx_support, cpu_has_amx_support,
is_cpu, is_cpu,
@@ -2834,10 +2830,9 @@ class MLATokenToKVPool(KVCache):
maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA)") maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA)")
layer_id = layer.layer_id layer_id = layer.layer_id
assert not self.dsa_kv_cache_store_fp8 assert not self.dsa_kv_cache_store_fp8
if dcp_enabled(): parallel = get_parallel()
valid_mask = ( if parallel.dcp_enabled:
loc % get_attention_dcp_world_size() == get_attention_dcp_rank() valid_mask = loc % parallel.attn_dcp_size == parallel.attn_dcp_rank
)
if not valid_mask.all(): if not valid_mask.all():
loc = loc[valid_mask] loc = loc[valid_mask]
cache_k = cache_k[valid_mask] cache_k = cache_k[valid_mask]
@@ -12,7 +12,6 @@ from sglang.srt.layers.communicator import get_attn_tp_context
from sglang.srt.layers.dcp import ( from sglang.srt.layers.dcp import (
all_gather_kv_cache_for_mha_chunk_extend, all_gather_kv_cache_for_mha_chunk_extend,
all_gather_kv_cache_for_mha_extend, all_gather_kv_cache_for_mha_extend,
dcp_enabled,
filter_dcp_local_kv_indices, filter_dcp_local_kv_indices,
) )
from sglang.srt.layers.quantization.fp8_utils import ( from sglang.srt.layers.quantization.fp8_utils import (
@@ -31,7 +30,7 @@ from sglang.srt.models.deepseek_common.utils import (
_use_aiter_bpreshuffle_gfx95, _use_aiter_bpreshuffle_gfx95,
_use_aiter_gfx95, _use_aiter_gfx95,
) )
from sglang.srt.runtime_context import get_server_args from sglang.srt.runtime_context import get_parallel, get_server_args
from sglang.srt.utils import BumpAllocator, get_bool_env_var, next_power_of_2 from sglang.srt.utils import BumpAllocator, get_bool_env_var, next_power_of_2
_use_fp8_prefill_attn = ( _use_fp8_prefill_attn = (
@@ -287,7 +286,7 @@ class DeepseekMHAForwardMixin:
kv_a, k_pe = self._get_mla_kv_buffer_from_fp8_for_dsa(forward_batch) kv_a, k_pe = self._get_mla_kv_buffer_from_fp8_for_dsa(forward_batch)
else: else:
# BF16/FP16 path: directly fetch from cache # BF16/FP16 path: directly fetch from cache
if dcp_enabled(): if get_parallel().dcp_enabled:
kv_a, k_pe = all_gather_kv_cache_for_mha_extend( kv_a, k_pe = all_gather_kv_cache_for_mha_extend(
get_token_to_kv_pool(), get_token_to_kv_pool(),
self.attn_mha, self.attn_mha,
@@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Optional
import torch import torch
from sglang.srt.compilation.compilation_config import register_split_op from sglang.srt.compilation.compilation_config import register_split_op
from sglang.srt.distributed.parallel_state import get_dcp_group
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.attention.dsa.utils import ( from sglang.srt.layers.attention.dsa.utils import (
@@ -19,8 +18,6 @@ from sglang.srt.layers.dcp import (
all_gather_kv_cache_for_mla_extend, all_gather_kv_cache_for_mla_extend,
all_gather_q_for_mla_decode, all_gather_q_for_mla_decode,
cp_lse_ag_out_rs_mla, cp_lse_ag_out_rs_mla,
dcp_enabled,
get_attention_dcp_world_size,
) )
from sglang.srt.layers.quantization.fp8_kernel import ( from sglang.srt.layers.quantization.fp8_kernel import (
fp8_dtype, fp8_dtype,
@@ -67,7 +64,7 @@ from sglang.srt.models.deepseek_common.utils import (
_use_aiter_bpreshuffle_gfx95, _use_aiter_bpreshuffle_gfx95,
_use_aiter_gfx95, _use_aiter_gfx95,
) )
from sglang.srt.runtime_context import get_server_args from sglang.srt.runtime_context import get_parallel, get_server_args
from sglang.srt.state_capturer.indexer_topk import ( from sglang.srt.state_capturer.indexer_topk import (
maybe_capture_indexer_topk, maybe_capture_indexer_topk,
) )
@@ -572,7 +569,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. # 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 get_parallel().dcp_enabled:
if forward_batch.forward_mode.is_decode(): if forward_batch.forward_mode.is_decode():
# if forward_batch.forward_mode is decode, gather q # if forward_batch.forward_mode is decode, gather q
q_nope_out, q_pe = all_gather_q_for_mla_decode( q_nope_out, q_pe = all_gather_q_for_mla_decode(
@@ -726,7 +723,10 @@ class DeepseekMLAForwardMixin:
topk_indices=topk_indices, topk_indices=topk_indices,
) )
attn_output = fusion_plan.attn_output_buf attn_output = fusion_plan.attn_output_buf
elif forward_batch.forward_mode.is_decode() and dcp_enabled(): elif (
forward_batch.forward_mode.is_decode()
and get_parallel().dcp_enabled
):
# set return_lse=True to correct attn_output # set return_lse=True to correct attn_output
attn_output, lse = self.attn_mqa_for_dcp_decode( attn_output, lse = self.attn_mqa_for_dcp_decode(
q_nope_out, q_nope_out,
@@ -800,13 +800,15 @@ class DeepseekMLAForwardMixin:
) )
# correct attn_output with respect to lse from other ranks # correct attn_output with respect to lse from other ranks
if forward_batch.forward_mode.is_decode() and dcp_enabled(): if forward_batch.forward_mode.is_decode() and get_parallel().dcp_enabled:
attn_output = attn_output.view( attn_output = attn_output.view(
-1, -1,
self.num_local_heads * get_attention_dcp_world_size(), self.num_local_heads * get_parallel().attn_dcp_size,
self.kv_lora_rank, self.kv_lora_rank,
) )
attn_output = cp_lse_ag_out_rs_mla(attn_output, lse, get_dcp_group()) 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.transpose(0, 1)
attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank) attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank)
+2 -3
View File
@@ -74,7 +74,6 @@ from sglang.srt.layers.communicator_dsa_cp import (
DSACPLayerCommunicator, DSACPLayerCommunicator,
maybe_prefetch_next_full_attention_kv, maybe_prefetch_next_full_attention_kv,
) )
from sglang.srt.layers.dcp import dcp_enabled, get_attention_dcp_world_size
from sglang.srt.layers.dcp.planner import ( from sglang.srt.layers.dcp.planner import (
prepare_decode_context_parallel_metadata, prepare_decode_context_parallel_metadata,
) )
@@ -1724,9 +1723,9 @@ class DeepseekV2AttentionMLA(
prefix=add_prefix("attn_mqa", prefix), prefix=add_prefix("attn_mqa", prefix),
) )
# use num_local_heads * dcp_world_size because q_nope, q_rope is all gathered from dcp ranks # use num_local_heads * dcp_world_size because q_nope, q_rope is all gathered from dcp ranks
if dcp_enabled(): if get_parallel().dcp_enabled:
self.attn_mqa_for_dcp_decode = RadixAttention( self.attn_mqa_for_dcp_decode = RadixAttention(
self.num_local_heads * get_attention_dcp_world_size(), self.num_local_heads * get_parallel().attn_dcp_size,
self.kv_lora_rank + self.qk_rope_head_dim, self.kv_lora_rank + self.qk_rope_head_dim,
self.scaling, self.scaling,
num_kv_heads=1, num_kv_heads=1,
+27 -3
View File
@@ -13,7 +13,7 @@
# ============================================================================== # ==============================================================================
"""A single structured accessor for process-static runtime state. """A single structured accessor for process-static runtime state.
``get_parallel()`` returns a ``ParallelContext`` whose attributes — tp / pp / ``get_parallel()`` returns a ``ParallelContext`` whose attributes — tp / dcp / pp /
moe / attn size and rank, plus the process-group handles — each delegate live to moe / attn size and rank, plus the process-group handles — each delegate live to
the canonical getter in ``distributed.parallel_state`` / ``layers.dp_attention``. the canonical getter in ``distributed.parallel_state`` / ``layers.dp_attention``.
Returned values are exactly what those getters return; this is a read-through Returned values are exactly what those getters return; this is a read-through
@@ -79,10 +79,13 @@ _PARALLEL_FIELDS = frozenset(
"attn_tp_rank", "attn_tp_rank",
"attn_cp_size", "attn_cp_size",
"attn_cp_rank", "attn_cp_rank",
"attn_dp_size", "dcp_enabled",
"attn_dp_rank",
"dcp_size", "dcp_size",
"dcp_rank", "dcp_rank",
"attn_dcp_size",
"attn_dcp_rank",
"attn_dp_size",
"attn_dp_rank",
"world_group", "world_group",
"tp_group", "tp_group",
"pp_group", "pp_group",
@@ -194,6 +197,27 @@ class ParallelContext:
def dcp_rank(self) -> int: def dcp_rank(self) -> int:
return self._v("dcp_rank", _ps().get_dcp_rank) return self._v("dcp_rank", _ps().get_dcp_rank)
@property
def dcp_enabled(self) -> bool:
def getter():
if _ps().get_dcp_group_no_assert() is None:
return False
return self.dcp_size > 1
return self._v("dcp_enabled", getter)
@property
def attn_dcp_size(self) -> int:
return self._v(
"attn_dcp_size", lambda: self.dcp_size if self.dcp_enabled else 1
)
@property
def attn_dcp_rank(self) -> int:
return self._v(
"attn_dcp_rank", lambda: self.dcp_rank if self.dcp_enabled else 0
)
@property @property
def attn_dp_size(self) -> int: def attn_dp_size(self) -> int:
return self._v("attn_dp_size", _dp().get_attention_dp_size) return self._v("attn_dp_size", _dp().get_attention_dp_size)
@@ -7,10 +7,10 @@ which gives one import, one naming scheme, and the scoped ``override()``
test primitive. Direct calls to the ``parallel_state`` size/rank getters in test primitive. Direct calls to the ``parallel_state`` size/rank getters in
these directories are regressions against that sweep. these directories are regressions against that sweep.
Exemptions, pinned by path: ``layers/dp_attention.py`` is delegation Exemptions, pinned by path: ``runtime_context.py`` and
substrate (the wrapper's attn-DP dims delegate TO it), and ``layers/dcp/`` ``layers/dp_attention.py`` are delegation substrate, while
is the DCP subsystem's own plumbing, booked for a follow-up sweep. Sweeping ``layers/dcp/comm.py`` retains deprecated DCP compatibility shims for
an exempt path must remove it from the pin. out-of-tree callers. Sweeping an exempt path must remove it from the pin.
""" """
from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.ci.ci_register import register_cpu_ci
@@ -27,7 +27,7 @@ from sglang.test.test_utils import CustomTestCase
_SRT_ROOT = Path(next(iter(sglang.srt.__path__))) _SRT_ROOT = Path(next(iter(sglang.srt.__path__)))
_BANNED_CALLS = re.compile( _BANNED_CALLS = re.compile(
r"\bget_(?:" r"\b(?:dcp_enabled|get_(?:"
r"tensor_model_parallel_(?:world_size|rank)" r"tensor_model_parallel_(?:world_size|rank)"
r"|pipeline_model_parallel_(?:world_size|rank)" r"|pipeline_model_parallel_(?:world_size|rank)"
r"|moe_expert_parallel_(?:world_size|rank)" r"|moe_expert_parallel_(?:world_size|rank)"
@@ -36,8 +36,10 @@ _BANNED_CALLS = re.compile(
r"|attn_tensor_model_parallel_(?:world_size|rank)" r"|attn_tensor_model_parallel_(?:world_size|rank)"
r"|attn_context_model_parallel_(?:world_size|rank)" r"|attn_context_model_parallel_(?:world_size|rank)"
r"|dcp_(?:world_size|rank)" r"|dcp_(?:world_size|rank)"
r"|dcp_group(?:_no_assert)?"
r"|attention_dcp_(?:world_size|rank)"
r"|attention_(?:tp|cp)_(?:group|rank|size)" r"|attention_(?:tp|cp)_(?:group|rank|size)"
r")\(\)" r"))\(\)"
) )
# The whole package is swept; the exemptions are the substrate itself. # The whole package is swept; the exemptions are the substrate itself.
@@ -45,7 +47,9 @@ _SWEPT_DIRS = ("",)
_EXEMPT = ( _EXEMPT = (
"distributed/", # parallel_state: defines the canonical getters "distributed/", # parallel_state: defines the canonical getters
"runtime_context.py", # delegates DCP reads to canonical getters
"layers/dp_attention.py", # delegation substrate for the attn-DP dims "layers/dp_attention.py", # delegation substrate for the attn-DP dims
"layers/dcp/comm.py", # deprecated out-of-tree DCP compatibility shims
# The dumper's megatron plugin calls third-party getters that share the # The dumper's megatron plugin calls third-party getters that share the
# parallel_state names (self._mpu.get_tensor_model_parallel_rank()). # parallel_state names (self._mpu.get_tensor_model_parallel_rank()).
"debug_utils/dumper.py", "debug_utils/dumper.py",
@@ -32,6 +32,8 @@ SIZE_RANK_DELEGATIONS = [
("world_rank", f"{_PS}.get_world_rank"), ("world_rank", f"{_PS}.get_world_rank"),
("tp_size", f"{_PS}.get_tensor_model_parallel_world_size"), ("tp_size", f"{_PS}.get_tensor_model_parallel_world_size"),
("tp_rank", f"{_PS}.get_tensor_model_parallel_rank"), ("tp_rank", f"{_PS}.get_tensor_model_parallel_rank"),
("dcp_size", f"{_PS}.get_dcp_world_size"),
("dcp_rank", f"{_PS}.get_dcp_rank"),
("pp_size", f"{_PS}.get_pipeline_model_parallel_world_size"), ("pp_size", f"{_PS}.get_pipeline_model_parallel_world_size"),
("pp_rank", f"{_PS}.get_pipeline_model_parallel_rank"), ("pp_rank", f"{_PS}.get_pipeline_model_parallel_rank"),
("moe_ep_size", f"{_PS}.get_moe_expert_parallel_world_size"), ("moe_ep_size", f"{_PS}.get_moe_expert_parallel_world_size"),
@@ -51,6 +53,7 @@ SIZE_RANK_DELEGATIONS = [
GROUP_DELEGATIONS = [ GROUP_DELEGATIONS = [
("world_group", f"{_PS}.get_world_group"), ("world_group", f"{_PS}.get_world_group"),
("tp_group", f"{_PS}.get_tp_group"), ("tp_group", f"{_PS}.get_tp_group"),
("dcp_group", f"{_PS}.get_dcp_group"),
("pp_group", f"{_PS}.get_pp_group"), ("pp_group", f"{_PS}.get_pp_group"),
("moe_ep_group", f"{_PS}.get_moe_ep_group"), ("moe_ep_group", f"{_PS}.get_moe_ep_group"),
("moe_dp_group", f"{_PS}.get_moe_dp_group"), ("moe_dp_group", f"{_PS}.get_moe_dp_group"),
@@ -151,6 +154,40 @@ class TestParallelOverride(_IsolatedOverrides):
self.assertEqual(p._overrides, {}) self.assertEqual(p._overrides, {})
class TestParallelDCP(_IsolatedOverrides):
def test_attn_dcp_defaults_when_group_is_uninitialized(self):
with (
patch(f"{_PS}.get_dcp_group_no_assert", return_value=None),
patch(f"{_PS}.get_dcp_world_size", side_effect=AssertionError),
patch(f"{_PS}.get_dcp_rank", side_effect=AssertionError),
):
self.assertFalse(get_parallel().dcp_enabled)
self.assertEqual(get_parallel().attn_dcp_size, 1)
self.assertEqual(get_parallel().attn_dcp_rank, 0)
def test_attn_dcp_delegates_when_enabled(self):
with (
patch(f"{_PS}.get_dcp_group_no_assert", return_value=object()),
patch(f"{_PS}.get_dcp_world_size", return_value=8),
patch(f"{_PS}.get_dcp_rank", return_value=3),
):
self.assertTrue(get_parallel().dcp_enabled)
self.assertEqual(get_parallel().attn_dcp_size, 8)
self.assertEqual(get_parallel().attn_dcp_rank, 3)
def test_dcp_enablement_is_platform_agnostic(self):
with (
patch(f"{_PS}.get_dcp_group_no_assert", return_value=object()),
patch("sglang.srt.utils.is_cuda", return_value=False) as is_cuda,
patch(f"{_PS}.get_dcp_world_size", return_value=8),
patch(f"{_PS}.get_dcp_rank", return_value=3),
):
self.assertTrue(get_parallel().dcp_enabled)
self.assertEqual(get_parallel().attn_dcp_size, 8)
self.assertEqual(get_parallel().attn_dcp_rank, 3)
is_cuda.assert_not_called()
class _IsolatedServerArgs(CustomTestCase): class _IsolatedServerArgs(CustomTestCase):
"""Save/restore the published ServerArgs around each test (the slot is """Save/restore the published ServerArgs around each test (the slot is
process-global; another test file sharing the process may have published).""" process-global; another test file sharing the process may have published)."""