minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[4/N][CP] Support interleave strategy for cp v2 (#30482)

Browse Source 
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
Zhangheng
2026-07-30 01:45:32 -07:00
committed by GitHub
co-authored by Xinyuan Tong Baizhou Zhang
parent c192145830
commit f46d5f25b4
15 changed files with 907 additions and 134 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+23 -8
View File
@@ -111,6 +111,8 @@ from sglang.srt.distributed import (
from sglang.srt.distributed.parallel_state import get_pp_group
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.communicator import ScatterMode
from sglang.srt.layers.cp.base import get_cp_strategy
from sglang.srt.layers.cp.utils import is_cp_v2_active
from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.rotary_embedding import get_rope_wrapper
@@ -630,13 +632,20 @@ class Indexer(MultiPlatformOp):
self.alt_stream.wait_stream(current_stream)
query = self._maybe_rotate(query)
# Gather the full key on alt_stream so the CP all-gather overlaps
# with the query rotate above on the current stream.
with torch.cuda.stream(self.alt_stream):
key = cp_all_gather_rerange_output(
key.contiguous(),
self.cp_size,
forward_batch,
torch.cuda.current_stream(),
)
if is_cp_v2_active(forward_batch):
key = get_cp_strategy().materialize_full_indexer_k_cache(
key, forward_batch
)
else:
key = cp_all_gather_rerange_output(
key.contiguous(),
self.cp_size,
forward_batch,
torch.cuda.current_stream(),
)
current_stream.wait_stream(self.alt_stream)
return query, key, weights_raw
else:
@@ -644,7 +653,9 @@ class Indexer(MultiPlatformOp):
key = self._maybe_rotate(key)
# allgather+rerrange
if forward_batch.attn_cp_metadata is not None and self.dsa_enable_prefill_cp:
if is_cp_v2_active(forward_batch):
key = get_cp_strategy().materialize_full_indexer_k_cache(key, forward_batch)
elif forward_batch.attn_cp_metadata is not None and self.dsa_enable_prefill_cp:
key = cp_all_gather_rerange_output(
key.contiguous(),
self.cp_size,
@@ -1870,7 +1881,11 @@ class Indexer(MultiPlatformOp):
weights = self._apply_q_scale_and_softmax_scale(weights, q_scale)
else:
query, key, weights_raw = self._get_q_k_bf16(
q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
q_lora,
x,
positions,
enable_dual_stream,
forward_batch=forward_batch,
)
if enable_dual_stream:
python/sglang/srt/layers/attention/dsa/utils.py
+21 -12
View File
@@ -82,7 +82,17 @@ def should_use_dsa_fused_topk(
def is_dsa_enable_prefill_cp():
return get_server_args().enable_dsa_prefill_context_parallel
if not envs.SGLANG_ENABLE_CP_V2.get():
return get_parallel().enable_dsa_prefill_context_parallel
# Derive from the runtime CP topology + model arch rather than the legacy
# flag under CP-v2: DSA prefill CP is active when the CP group is on for a
# DeepSeek Sparse Attention model.
if get_parallel().attn_cp_size <= 1:
return False
from sglang.srt.configs.model_config import is_deepseek_dsa
return is_deepseek_dsa(get_server_args().get_model_config().hf_config)
def is_dsa_prefill_cp_in_seq_split():
@@ -206,6 +216,15 @@ def pad_dsa_cache_seqlens(forward_batch: "ForwardBatch", dsa_cache_seqlens):
def can_dsa_cp_split(seq_len: int, cp_size: int, use_dsa: bool, forward_batch):
if (
cp_size <= 1
or not use_dsa
or not forward_batch.forward_mode.is_context_parallel_extend()
or not is_dsa_enable_prefill_cp()
or sum(forward_batch.extend_seq_lens_cpu) < cp_size
):
return False
if is_dsa_prefill_cp_round_robin_split():
cur_cp_seq_len = seq_len // cp_size
assert (
@@ -216,17 +235,7 @@ def can_dsa_cp_split(seq_len: int, cp_size: int, use_dsa: bool, forward_batch):
# Note: (self.cp_size * 2) To achieve load balancing for seq computation,
# the seq data needs to be divided and recombined at twice the size of cp_size.
cur_cp_seq_len = seq_len // (cp_size * 2)
if (
cur_cp_seq_len != 0
and cp_size > 1
and use_dsa
and forward_batch.forward_mode.is_context_parallel_extend()
and is_dsa_enable_prefill_cp()
and sum(forward_batch.extend_seq_lens_cpu) >= cp_size
):
return True
else:
return False
return cur_cp_seq_len != 0
from sglang.kernels.ops.attention.dsa.cp_split import (
python/sglang/srt/layers/attention/dsa_backend.py
+45 -9
View File
@@ -62,6 +62,8 @@ from sglang.srt.layers.attention.trtllm_mla_backend import (
grow_multi_ctas_kv_counter_buffer_if_needed,
make_persistent_multi_ctas_kv_counter_buffer,
)
from sglang.srt.layers.cp.base import get_cp_strategy
from sglang.srt.layers.cp.utils import is_cp_v2_active
from sglang.srt.layers.utils.cp_utils import (
cp_all_gather_rerange_output,
cp_split_and_rebuild_position,
@@ -111,6 +113,23 @@ def _all_gather_dsa_trtllm_fp8_kv(
return kv.split((kv_lora_rank, qk_rope_head_dim), dim=-1)
def materialize_full_kv_cp(
attn_mla,
forward_batch: ForwardBatch,
latent_cache: torch.Tensor,
k_nope: torch.Tensor,
k_pe: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
if is_cp_v2_active(forward_batch):
return get_cp_strategy().materialize_full_mla_kv(
forward_batch,
attn_mla.attn_mqa,
k_nope,
k_pe,
)
return attn_mla.rebuild_cp_kv_cache(latent_cache, forward_batch, k_nope, k_pe)
_is_hip = is_hip()
if _is_hip:
@@ -987,12 +1006,19 @@ class DeepseekSparseAttnBackend(
)
if can_dsa_prefill_cp_round_robin_split(forward_batch):
seqlens_expanded = dsa_cp_round_robin_split_data(seqlens_expanded)
extend_seq_lens_cpu, extend_seq_lens, bs_idx_cpu, bs_idx = (
dsa_cp_round_robin_split_q_seqs(
extend_seq_lens_cpu, extend_seq_lens
if is_cp_v2_active(forward_batch):
strategy = get_cp_strategy()
seqlens_expanded = strategy.shard_local_tokens(seqlens_expanded)
extend_seq_lens_cpu, extend_seq_lens, bs_idx_cpu, bs_idx = (
strategy.shard_per_request(extend_seq_lens_cpu, extend_seq_lens)
)
else:
seqlens_expanded = dsa_cp_round_robin_split_data(seqlens_expanded)
extend_seq_lens_cpu, extend_seq_lens, bs_idx_cpu, bs_idx = (
dsa_cp_round_robin_split_q_seqs(
extend_seq_lens_cpu, extend_seq_lens
)
)
)
indexer_seq_lens_cpu = indexer_seq_lens_cpu[bs_idx_cpu]
indexer_seq_lens = indexer_seq_lens[bs_idx]
cache_seqlens_int32 = cache_seqlens_int32[bs_idx]
@@ -1198,9 +1224,14 @@ class DeepseekSparseAttnBackend(
token_to_batch_idx = torch.cat(token_to_batch_idx, dim=0)
if bs_idx is not None:
assert can_dsa_prefill_cp_round_robin_split(forward_batch)
ks = dsa_cp_round_robin_split_data(ks)
ke = dsa_cp_round_robin_split_data(ke)
token_to_batch_idx = dsa_cp_round_robin_split_data(token_to_batch_idx)
split_per_token = (
get_cp_strategy().shard_local_tokens
if is_cp_v2_active(forward_batch)
else dsa_cp_round_robin_split_data
)
ks = split_per_token(ks)
ke = split_per_token(ke)
token_to_batch_idx = split_per_token(token_to_batch_idx)
return (ks, ke), token_to_batch_idx
def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
@@ -3187,7 +3218,12 @@ class DeepseekSparseAttnBackend(
self.qk_rope_head_dim,
)
if save_kv_cache and dsa_use_prefill_cp(forward_batch):
k, k_rope = _all_gather_dsa_trtllm_fp8_kv(forward_batch, k, k_rope)
if is_cp_v2_active(forward_batch):
k, k_rope = get_cp_strategy().all_gather_dsa_trtllm_fp8_kv(
forward_batch, k, k_rope
)
else:
k, k_rope = _all_gather_dsa_trtllm_fp8_kv(forward_batch, k, k_rope)
merge_query = False
# Save KV cache if requested
python/sglang/srt/layers/cp/base.py
+38 -6
View File
@@ -67,17 +67,20 @@ class CPAttentionBackendKind(IntEnum):
"""Attention backend calling convention used by CP strategy dispatch."""
FLASH_ATTENTION = 0
TRTLLM_MHA = 1
DSA = 1
TRTLLM_MHA = 2
@classmethod
def from_string(cls, value: str) -> CPAttentionBackendKind:
if value in ("fa3", "fa4", "flashinfer"):
return cls.FLASH_ATTENTION
if value in ("dsa"):
return cls.DSA
if value == "trtllm_mha":
return cls.TRTLLM_MHA
raise ValueError(
f"Unsupported attention_backend={value!r} for CP strategy; expected one "
"of {'fa3', 'fa4', 'flashinfer', 'trtllm_mha'}"
"of {'fa3', 'fa4', 'flashinfer', 'dsa', 'trtllm_mha'}"
)
@@ -153,6 +156,25 @@ class ContextParallelStrategy(ABC):
f"{self.name} strategy does not support per-request sharding"
)
def shard_local_tokens(self, input_: Any) -> Any:
raise NotImplementedError(
f"{self.name} strategy does not support local-token sharding"
)
def materialize_full_indexer_k_cache(
self, key: Any, forward_batch: ForwardBatch
) -> Any:
raise NotImplementedError(
f"{self.name} strategy does not support DSA indexer key gather"
)
def all_gather_dsa_trtllm_fp8_kv(
self, forward_batch: ForwardBatch, k: Any, k_rope: Any
) -> Any:
raise NotImplementedError(
f"{self.name} strategy does not support DSA trtllm FP8 KV gather"
)
def split_before_forward(
self,
forward_batch: ForwardBatch,
@@ -185,13 +207,23 @@ class ContextParallelStrategy(ABC):
def materialize_full_kv(
self,
forward_batch: ForwardBatch,
layer: Any,
k: Any,
v: Any,
layer: Any = None,
k: Any = None,
v: Any = None,
swa_loc: Optional[Any] = None,
) -> None:
) -> Any:
"""Write full-layout K/V to the backend cache if needed."""
@abstractmethod
def materialize_full_mla_kv(
self,
forward_batch: ForwardBatch,
layer: Any,
k_nope: Any,
k_rope: Any,
) -> Any:
"""Materialize full-layout MLA K/V for the strategy."""
def reindex_attn_metadata(self, core_attn_metadata: Any) -> None:
"""Optional attention metadata rewrite for strategies that need it."""
return None
python/sglang/srt/layers/cp/interleave.py
+186 -21
View File
@@ -31,17 +31,30 @@ from __future__ import annotations
from dataclasses import dataclass
from typing import Any, List, Optional
import torch
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.layers.cp.base import (
BaseContextParallelMetadata,
ContextParallelStrategy,
ContextParallelStrategyKind,
CPAttentionBackendKind,
)
from sglang.srt.layers.cp.padding import pad_local_rows
from sglang.srt.layers.dp_attention import (
attn_cp_all_gather_into_tensor,
is_allocation_symmetric,
)
from sglang.srt.runtime_context import get_parallel
@dataclass
class InterleaveContextParallelMetadata(BaseContextParallelMetadata):
"""Interleave has no per-forward zigzag permutation payload."""
per_rank_actual_token: Optional[List[int]] = None
max_rank_len: Optional[List[int]] = None
per_rank_logical_token: Optional[List[int]] = None
class InterleaveCPStrategy(ContextParallelStrategy):
@@ -49,10 +62,11 @@ class InterleaveCPStrategy(ContextParallelStrategy):
kind = ContextParallelStrategyKind.INTERLEAVE
def can_apply(self, num_tokens: int, forward_batch) -> bool:
if self.cp_size <= 1 or num_tokens < self.cp_size:
if not forward_batch.forward_mode.is_context_parallel_extend():
return False
forward_mode = getattr(forward_batch, "forward_mode", None)
return forward_mode is None or forward_mode.is_context_parallel_extend()
cp_size = self.cp_size
seq_len = sum(forward_batch.extend_seq_lens_cpu)
return seq_len > 0 and seq_len >= cp_size and cp_size > 1
def build_metadata(
self,
@@ -60,32 +74,150 @@ class InterleaveCPStrategy(ContextParallelStrategy):
seqs_len: Optional[List[int]],
extend_seqs_len: Optional[List[int]] = None,
) -> InterleaveContextParallelMetadata:
if extend_seqs_len is None:
extend_seqs_len = seqs_len or [num_tokens]
extend_seqs_len = [int(x) for x in extend_seqs_len]
pad_len = int(num_tokens) - sum(extend_seqs_len)
if pad_len > 0:
extend_seqs_len[-1] += pad_len
total_seq_lens = sum(extend_seqs_len)
base_len, extra = divmod(total_seq_lens, self.cp_size)
per_rank_actual_token = [
base_len + (rank < extra) for rank in range(self.cp_size)
]
return InterleaveContextParallelMetadata(
total_seq_lens=sum(extend_seqs_len or seqs_len or [num_tokens]),
bs=len(extend_seqs_len or seqs_len or [num_tokens]),
per_rank_actual_token=per_rank_actual_token,
max_rank_len=[max(per_rank_actual_token)] * self.cp_size,
total_seq_lens=total_seq_lens,
bs=len(extend_seqs_len),
)
def shard_hidden_states(self, x: Any, forward_batch) -> Any:
raise NotImplementedError(
"Interleave hidden-state sharding will land in a follow-up PR"
)
metadata = forward_batch.attn_cp_metadata
local_x = self._interleave_shard(x[: metadata.total_seq_lens])
return pad_local_rows(local_x, metadata, dim=0)
def shard_position_ids(self, positions: Any, forward_batch) -> Any:
raise NotImplementedError(
"Interleave position-id sharding will land in a follow-up PR"
metadata = forward_batch.attn_cp_metadata
local_positions = self._interleave_shard(positions[: metadata.total_seq_lens])
return pad_local_rows(local_positions, metadata, dim=0)
def _interleave_shard(self, input_: Any) -> Any:
cp_size = self.cp_size
cp_rank = self.cp_rank
if isinstance(input_, (tuple, list)):
indices = range(cp_rank, len(input_), cp_size)
return input_[indices]
tokens = len(input_)
if tokens % cp_size != 0:
cur_len = tokens // cp_size + (tokens % cp_size > cp_rank)
if cur_len == 0:
return input_.new_empty(0, *input_.shape[1:])
indices = torch.arange(cp_rank, tokens, cp_size, device=input_.device)
return input_[indices]
return input_.view(-1, cp_size, *input_.shape[1:])[:, cp_rank].contiguous()
def shard_local_tokens(self, input_: Any) -> Any:
return self._interleave_shard(input_)
def shard_per_request(
self,
extend_seqs_cpu: List[int],
extend_seqs: Any,
):
"""Build device outputs in the shared kernel to keep the split graph-safe."""
from sglang.kernels.ops.attention.dsa.cp_split import (
dsa_cp_round_robin_split_q_seqs_kernel,
)
cp_size = self.cp_size
cp_rank = self.cp_rank
extra_seq = 0
q_lens_cpu: List[int] = []
for cur_len in extend_seqs_cpu:
cur_len += extra_seq
cur_seq = cur_len // cp_size + int(cur_len % cp_size > cp_rank)
q_lens_cpu.append(cur_seq)
extra_seq = cur_len - cur_seq * cp_size
bs_idx_cpu = [i for i, q_len in enumerate(q_lens_cpu) if q_len > 0]
q_lens_cpu = [q_len for q_len in q_lens_cpu if q_len > 0]
q_lens = torch.empty(
(len(bs_idx_cpu),), device=extend_seqs.device, dtype=extend_seqs.dtype
)
bs_idx = torch.empty(
(len(bs_idx_cpu),), device=extend_seqs.device, dtype=torch.int32
)
dsa_cp_round_robin_split_q_seqs_kernel[(1,)](
extend_seqs, q_lens, bs_idx, len(extend_seqs), cp_size, cp_rank
)
return q_lens_cpu, q_lens, bs_idx_cpu, bs_idx
def gather_hidden_states(
self, x: Any, forward_batch, stream: Optional[Any] = None
) -> Any:
raise NotImplementedError(
"Interleave hidden-state gather will land in a follow-up PR"
)
return self._gather_interleaved_tensor(x, forward_batch)
def gather_kv_cache(
self, x: Any, forward_batch, stream: Optional[Any] = None
) -> Any:
raise NotImplementedError("Interleave KV gather will land in a follow-up PR")
return self._gather_interleaved_tensor(x, forward_batch)
def _gather_interleaved_tensor(self, x: Any, forward_batch) -> Any:
metadata = getattr(forward_batch, "attn_cp_metadata", None)
if metadata is None:
raise RuntimeError("Interleave CP gather requires attn_cp_metadata.")
total_tokens = int(metadata.total_seq_lens)
if total_tokens < 0:
raise RuntimeError(
f"Invalid interleave CP total_seq_lens={total_tokens}; expected >= 0."
)
logical_rank_lens = (
metadata.per_rank_logical_token or metadata.per_rank_actual_token
)
local_logical_len = logical_rank_lens[self.cp_rank]
if x.shape[0] < local_logical_len:
raise RuntimeError(
"Interleave CP gather received an unexpected local token count: "
f"rank={self.cp_rank}, got={x.shape[0]}, "
f"expected_at_least={local_logical_len}, "
f"total={total_tokens}, cp_size={self.cp_size}."
)
physical_rank_len = max(metadata.per_rank_actual_token)
if physical_rank_len == 0:
return x.new_empty((0, *x.shape[1:]))
padded_x = x.new_zeros((physical_rank_len, *x.shape[1:]))
padded_x[:local_logical_len] = x[:local_logical_len]
with use_symmetric_memory(
get_parallel().attn_cp_group, disabled=not is_allocation_symmetric()
):
gathered = x.new_empty((self.cp_size * physical_rank_len, *x.shape[1:]))
attn_cp_all_gather_into_tensor(gathered, padded_x.contiguous())
flat_indices = torch.arange(total_tokens, device=x.device)
gather_indices = (
flat_indices % self.cp_size
) * physical_rank_len + flat_indices // self.cp_size
return gathered.index_select(0, gather_indices)
def get_supported_attention_backend(self):
return [CPAttentionBackendKind.DSA]
def materialize_full_indexer_k_cache(self, key: Any, forward_batch) -> Any:
return self.gather_kv_cache(
key.contiguous(), forward_batch, torch.cuda.current_stream()
)
def run_attention(
self,
@@ -94,14 +226,47 @@ class InterleaveCPStrategy(ContextParallelStrategy):
device: Any,
attn_fn,
attention_backend: CPAttentionBackendKind = CPAttentionBackendKind.FLASH_ATTENTION,
**kwargs,
) -> Any:
raise NotImplementedError(
"Interleave attention dispatch will land in a follow-up PR"
)
# No-op: run_attention is the FlashAttention/zigzag dispatch hook.
# Interleave serves the DSA backend, which runs attention itself.
return None
def all_gather_dsa_trtllm_fp8_kv(self, forward_batch, k: Any, k_rope: Any) -> Any:
kv_lora_rank = k.shape[-1]
qk_rope_head_dim = k_rope.shape[-1]
kv_dtype = k.dtype
# Pack → gather in raw bytes to avoid dtype issues with FP8
kv = torch.cat((k, k_rope), dim=-1).view(torch.uint8)
kv = self.gather_kv_cache(
kv.contiguous(), forward_batch, torch.cuda.current_stream()
).view(kv_dtype)
return kv.split((kv_lora_rank, qk_rope_head_dim), dim=-1)
def materialize_full_kv(
self, forward_batch, layer: Any, k: Any, v: Any, swa_loc: Optional[Any] = None
) -> None:
self,
forward_batch,
layer: Any = None,
k: Any = None,
v: Any = None,
swa_loc: Optional[Any] = None,
) -> Any:
raise NotImplementedError(
"Interleave KV materialization will land in a follow-up PR"
f"{self.name} strategy does not support dense K/V materialization"
)
def materialize_full_mla_kv(
self,
forward_batch,
layer: Any,
k_nope: Any,
k_rope: Any,
) -> Any:
kv_lora_rank = k_nope.shape[-1]
latent_cache = torch.cat([k_nope, k_rope], dim=-1).squeeze(1)
full_latent = self.gather_kv_cache(
latent_cache.contiguous(), forward_batch, torch.cuda.current_stream()
)
k_nope = full_latent[..., :kv_lora_rank].unsqueeze(1)
k_rope = full_latent[..., kv_lora_rank:].unsqueeze(1)
return k_nope, k_rope
python/sglang/srt/layers/cp/utils.py
+59 -7
View File
@@ -14,6 +14,7 @@
"""Public import facade and runtime helpers for context parallel strategies."""
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, Optional, Tuple
from sglang.srt.layers.cp.base import (
@@ -40,6 +41,8 @@ if TYPE_CHECKING:
CP_V2_DEFAULT_MODEL_CLASSES = frozenset(
{
"DeepseekV32ForCausalLM",
"GlmMoeDsaForCausalLM",
"GptOssForCausalLM",
"MiMoV2FlashForCausalLM",
"MiMoV2ForCausalLM",
@@ -176,9 +179,7 @@ def prepare_cp_forward(forward_batch) -> None:
from sglang.srt.layers.dp_attention import set_local_dp_buffer_len
set_local_dp_buffer_len(
forward_batch.attn_cp_metadata.per_rank_actual_token[
get_parallel().attn_cp_rank
]
sum(forward_batch.attn_cp_metadata.per_rank_actual_token)
)
if getattr(forward_batch, "out_cache_loc", None) is not None:
@@ -191,15 +192,31 @@ def cp_split_before_forward(
forward_batch,
) -> Tuple[Optional[Any], Optional[Any]]:
"""Shard embeddings and positions for CP-v2 model-runner forwarding."""
assert is_cp_v2_active(forward_batch)
assert complete_hidden_states is not None
assert getattr(forward_batch, "attn_cp_metadata", None) is not None
return (
cp_shard_hidden_states(complete_hidden_states, forward_batch),
cp_shard_position_ids(complete_position_ids, forward_batch),
)
def cp_shard_hidden_states(complete_hidden_states: Any, forward_batch):
assert is_cp_v2_active(forward_batch)
strategy = get_cp_strategy()
assert strategy is not None
assert complete_hidden_states is not None
assert getattr(forward_batch, "attn_cp_metadata", None) is not None
return (
strategy.shard_hidden_states(complete_hidden_states, forward_batch),
strategy.shard_position_ids(complete_position_ids, forward_batch),
)
return strategy.shard_hidden_states(complete_hidden_states, forward_batch)
def cp_shard_position_ids(complete_position_ids: Any, forward_batch):
assert is_cp_v2_active(forward_batch)
strategy = get_cp_strategy()
assert strategy is not None
assert complete_position_ids is not None
assert getattr(forward_batch, "attn_cp_metadata", None) is not None
return strategy.shard_position_ids(complete_position_ids, forward_batch)
def cp_gather_after_forward(x: Any, forward_batch, stream: Optional[Any] = None):
@@ -221,6 +238,38 @@ def cp_gather_after_forward(x: Any, forward_batch, stream: Optional[Any] = None)
return strategy.gather_hidden_states(x, forward_batch, stream)
@contextmanager
def cp_shard_model_inputs(
complete_hidden_states: Any,
complete_position_ids: Any,
forward_batch,
):
"""Restore the shared batch so logits processing keeps full-batch metadata."""
assert is_cp_v2_active(forward_batch)
sharded_hidden_states = cp_shard_hidden_states(
complete_hidden_states, forward_batch
)
sharded_positions = cp_shard_position_ids(complete_position_ids, forward_batch)
spec_info = getattr(forward_batch, "spec_info", None)
spec_hidden_states = getattr(spec_info, "hidden_states", None)
spec_hidden_states_backup = None
if (
spec_hidden_states is not None
and spec_hidden_states.shape[0] == complete_hidden_states.shape[0]
):
spec_hidden_states_backup = spec_hidden_states
spec_info.hidden_states = cp_shard_hidden_states(
spec_hidden_states, forward_batch
)
try:
yield sharded_hidden_states, sharded_positions
finally:
if spec_hidden_states_backup is not None:
spec_info.hidden_states = spec_hidden_states_backup
def _to_int_list(values) -> Optional[list[int]]:
if values is None:
return None
@@ -244,6 +293,9 @@ __all__ = [
"get_cp_strategy",
"is_cp_v2_active",
"cp_gather_after_forward",
"cp_shard_hidden_states",
"cp_shard_model_inputs",
"cp_shard_position_ids",
"cp_split_before_forward",
"prepare_cp_forward",
"is_glm_dsa_cache_layer_split_enabled",
python/sglang/srt/layers/cp/zigzag.py
+7 -2
View File
@@ -363,8 +363,13 @@ class ZigzagCPStrategy(ContextParallelStrategy):
return result
def materialize_full_kv(
self, forward_batch, layer: Any, k: Any, v: Any, swa_loc: Optional[Any] = None
) -> None:
self,
forward_batch,
layer: Any = None,
k: Any = None,
v: Any = None,
swa_loc: Optional[Any] = None,
) -> Any:
cache_loc = (
forward_batch.out_cache_loc
if not layer.is_cross_attention
python/sglang/srt/model_executor/runner/eager_runner.py
+17 -14
View File
@@ -26,7 +26,7 @@ from sglang.srt.dllm.config import DllmConfig
from sglang.srt.environ import envs
from sglang.srt.layers.cp.utils import (
cp_gather_after_forward,
cp_split_before_forward,
cp_shard_model_inputs,
is_cp_v2_active,
prepare_cp_forward,
)
@@ -259,7 +259,11 @@ class EagerRunner(BaseRunner):
if not self.enable_pdmux:
forward_batch = self.load_batch(forward_batch, pp_proxy_tensors)
if forward_batch.needs_forward_metadata_init():
cp_v2_active = is_cp_v2_active(forward_batch)
if cp_v2_active:
prepare_cp_forward(forward_batch)
if forward_batch.needs_forward_metadata_init() or cp_v2_active:
if model_runner.dcp_size > 1 and hasattr(
model_runner.model, "prepare_context_parallel_metadata_for_dcp"
):
@@ -285,7 +289,6 @@ class EagerRunner(BaseRunner):
model_runner.model.prepare_forward_batch(forward_batch)
model_runner.attn_backend.init_forward_metadata(forward_batch)
cp_v2_active = is_cp_v2_active(forward_batch)
if not cp_v2_active:
forward_batch.attn_cp_metadata = None
@@ -341,21 +344,21 @@ class EagerRunner(BaseRunner):
"""
model = self.model_runner.model
prepare_cp_forward(forward_batch)
input_embeds = kwargs.get("input_embeds")
if input_embeds is None:
input_embeds = model.get_input_embeddings()(forward_batch.input_ids)
input_embeds, positions = cp_split_before_forward(
with cp_shard_model_inputs(
input_embeds, forward_batch.positions, forward_batch
)
hidden_states = model.model(
forward_batch.input_ids,
positions,
forward_batch,
input_embeds=input_embeds,
pp_proxy_tensors=kwargs.get("pp_proxy_tensors"),
)
) as (sharded_input_embeds, sharded_positions):
model_kwargs = {"input_embeds": sharded_input_embeds}
if (pp_proxy_tensors := kwargs.get("pp_proxy_tensors")) is not None:
model_kwargs["pp_proxy_tensors"] = pp_proxy_tensors
hidden_states = model.model(
forward_batch.input_ids,
sharded_positions,
forward_batch,
**model_kwargs,
)
capture_aux_hidden_states = getattr(model, "capture_aux_hidden_states", False)
aux_hidden_states = None
if capture_aux_hidden_states:
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py
+15 -6
View File
@@ -774,15 +774,24 @@ class DeepseekMLAForwardMixin:
dsa_prefill_cp=dsa_prefill_cp,
fuse_rope_for_trtllm_mla=fuse_rope_for_trtllm_mla,
)
if (
(dsa_prefill_cp or mla_prefill_cp)
and not defer_kv_gather_until_after_rope
and not is_cp_v2_active(forward_batch)
):
if dsa_prefill_cp and not defer_kv_gather_until_after_rope:
from sglang.srt.layers.attention.dsa_backend import materialize_full_kv_cp
k_nope, k_pe = materialize_full_kv_cp(
self,
forward_batch,
latent_cache,
k_nope,
k_pe,
)
elif mla_prefill_cp and not is_cp_v2_active(forward_batch):
# CP-v1 gathers the latent here; CP-v2 gathers it in the attention
# backend via the strategy (materialize_full_mla_kv).
k_nope, k_pe = self.rebuild_cp_kv_cache(
latent_cache, forward_batch, k_nope, k_pe
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.
python/sglang/srt/models/deepseek_nextn.py
+31 -23
View File
@@ -35,7 +35,7 @@ from sglang.srt.layers.attention.dsa.utils import (
is_dsa_enable_prefill_cp,
is_dsa_prefill_cp_round_robin_split,
)
from sglang.srt.layers.cp.utils import is_cp_v2_active
from sglang.srt.layers.cp.utils import cp_gather_after_forward, is_cp_v2_active
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.logits_processor import LogitsProcessor
@@ -252,11 +252,12 @@ class DeepseekModelNextN(nn.Module):
else:
hidden_states = self.eh_proj(eh_input)
# CP-v2 shards/gathers at the eager-runner boundary instead.
# CP-v2 shards/gathers hidden states at the eager-runner boundary.
cp_v2_active = is_cp_v2_active(forward_batch)
use_cp_v1 = (
dsa_use_prefill_cp(forward_batch, self.dsa_enable_prefill_cp)
or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp)
) and not is_cp_v2_active(forward_batch)
) and not cp_v2_active
if use_cp_v1:
hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
positions = cp_split_and_rebuild_position(forward_batch, positions)
@@ -304,6 +305,12 @@ class DeepseekModelNextN(nn.Module):
forward_batch,
torch.cuda.current_stream(),
)
elif (
cp_v2_active
and should_update_dsa_topk_indices
and topk_indices is not None
):
topk_indices = cp_gather_after_forward(topk_indices, forward_batch)
if should_update_dsa_topk_indices and topk_indices is not None:
if forward_batch.reuse_dsa_topk_indices:
forward_batch.spec_info.dsa_topk_indices = topk_indices
@@ -389,26 +396,27 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
forward_batch: ForwardBatch,
) -> torch.Tensor:
# TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2
if self.dsa_enable_prefill_cp:
if can_dsa_cp_split(
len(input_ids), self.cp_size, self.use_dsa, forward_batch
):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len(input_ids),
self.cp_rank,
self.cp_size,
forward_batch.seq_lens_cpu.tolist(),
extend_seqs_len=forward_batch.extend_seq_lens_cpu,
)
elif self.mla_enable_prefill_cp:
if can_cp_split(len(input_ids), self.cp_size, forward_batch):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len(input_ids),
self.cp_rank,
self.cp_size,
forward_batch.seq_lens_cpu.tolist(),
extend_seqs_len=forward_batch.extend_seq_lens_cpu,
)
if not is_cp_v2_active(forward_batch):
if self.dsa_enable_prefill_cp:
if can_dsa_cp_split(
len(input_ids), self.cp_size, self.use_dsa, forward_batch
):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len(input_ids),
self.cp_rank,
self.cp_size,
forward_batch.seq_lens_cpu.tolist(),
extend_seqs_len=forward_batch.extend_seq_lens_cpu,
)
elif self.mla_enable_prefill_cp:
if can_cp_split(len(input_ids), self.cp_size, forward_batch):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len(input_ids),
self.cp_rank,
self.cp_size,
forward_batch.seq_lens_cpu.tolist(),
extend_seqs_len=forward_batch.extend_seq_lens_cpu,
)
hidden_states = self.model(input_ids, positions, forward_batch)
return self.logits_processor(
input_ids, hidden_states, self.lm_head, forward_batch
python/sglang/srt/models/deepseek_v2.py
+22 -20
View File
@@ -3045,6 +3045,7 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
input_embeds: torch.Tensor = None,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> torch.Tensor:
# Multi-modal: input_ids may be None (use input_embeds).
# Non-first PP ranks: both are None (activations via pp_proxy_tensors).
if input_ids is not None:
@@ -3053,26 +3054,27 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
len_input_ids = input_embeds.shape[0]
else:
len_input_ids = pp_proxy_tensors["hidden_states"].shape[0]
if self.dsa_enable_prefill_cp:
if can_dsa_cp_split(
len_input_ids, self.cp_size, self.use_dsa, forward_batch
):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len_input_ids,
self.cp_rank,
self.cp_size,
forward_batch.seq_lens_cpu.tolist(),
extend_seqs_len=forward_batch.extend_seq_lens_cpu,
)
elif self.mla_enable_prefill_cp:
if can_cp_split(len_input_ids, self.cp_size, forward_batch):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len_input_ids,
self.cp_rank,
self.cp_size,
forward_batch.seq_lens_cpu.tolist(),
extend_seqs_len=forward_batch.extend_seq_lens_cpu,
)
if not is_cp_v2_active(forward_batch):
if self.dsa_enable_prefill_cp:
if can_dsa_cp_split(
len_input_ids, self.cp_size, self.use_dsa, forward_batch
):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len_input_ids,
self.cp_rank,
self.cp_size,
forward_batch.seq_lens_cpu.tolist(),
extend_seqs_len=forward_batch.extend_seq_lens_cpu,
)
elif self.mla_enable_prefill_cp:
if can_cp_split(len_input_ids, self.cp_size, forward_batch):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len_input_ids,
self.cp_rank,
self.cp_size,
forward_batch.seq_lens_cpu.tolist(),
extend_seqs_len=forward_batch.extend_seq_lens_cpu,
)
with get_attn_tp_context().maybe_input_scattered(forward_batch):
hidden_states = self.model(
python/sglang/srt/server_args.py
+17 -6
View File
@@ -6189,15 +6189,26 @@ class ServerArgs:
def _handle_context_parallelism(self):
if parse_connector_type(self.model_path) != ConnectorType.INSTANCE:
from sglang.srt.configs.model_config import is_deepseek_dsa
from sglang.srt.layers.cp.utils import CP_V2_DEFAULT_MODEL_CLASSES
model_config = self.get_model_config()
model_arch = model_config.hf_config.architectures[0]
if (
model_arch in CP_V2_DEFAULT_MODEL_CLASSES
and not envs.SGLANG_ENABLE_CP_V2.is_set()
):
envs.SGLANG_ENABLE_CP_V2.set(True)
hf_config = model_config.hf_config
model_arch = hf_config.architectures[0]
if model_arch in CP_V2_DEFAULT_MODEL_CLASSES:
if getattr(hf_config, "index_share_for_mtp_iteration", False):
# GLM 5.2 (DSA index-share MTP): CP-v2 is not ready for it
# yet, so default the env to off and keep the legacy CP path.
if not envs.SGLANG_ENABLE_CP_V2.is_set():
envs.SGLANG_ENABLE_CP_V2.set(False)
else:
is_dsa_default_model = is_deepseek_dsa(hf_config)
# DSA CP-v2 currently supports only the interleave strategy.
enable_default_cp_v2 = not is_dsa_default_model or (
self.enable_prefill_cp and self.cp_strategy == "interleave"
)
if enable_default_cp_v2 and not envs.SGLANG_ENABLE_CP_V2.is_set():
envs.SGLANG_ENABLE_CP_V2.set(True)
if (
self.enable_prefill_cp