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

[CP V1 Deprecation 3.5/5] Deprecate HIP/NPU/MUSA prefill CP and remove legacy implementation (#38293)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-09-07 15:39:34 -07:00
committed by GitHub
parent f4b75b5c36
commit 85d39401c8
47 changed files with 213 additions and 2127 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/attention_hook.py
+1 -1
View File
@@ -161,7 +161,7 @@ def handle_attention_backend_compatibility(server_args: Any):
if (
prefill_backend == "trtllm_mha"
and not get_platform().is_sm100
and (cfg.enable_prefill_context_parallel or cfg.attn_cp_size > 1)
and cfg.attn_cp_size > 1
):
raise ValueError(
"Prefill context parallelism with the TRTLLM MHA prefill backend "
python/sglang/srt/arg_groups/cuda_graph_hook.py
-4
View File
@@ -244,10 +244,6 @@ def disable_tc_piecewise_cudagraph_if_incompatible(server_args: Any):
lambda: resolved_view(server_args).attn_cp_size > 1,
),
("CUDA graph debug mode", lambda: cfg.debug_cuda_graph),
(
"DSA prefill context parallelism",
lambda: cfg.enable_dsa_prefill_context_parallel,
),
# Capture builds a dummy extend forward with attn_dcp_metadata=None.
(
"decode context parallel (dcp_size > 1)",
python/sglang/srt/arg_groups/deepseek_v4_hook.py
-18
View File
@@ -168,24 +168,6 @@ def validate_deepseek_v4_cp(server_args: ServerArgs) -> None:
f"DeepSeekV4 only supports interleave CP strategy, got {cfg.cp_strategy}"
)
if get_platform().is_hip or get_platform().is_npu or get_platform().is_musa:
# Protected platform implementations still consume the legacy runtime
# fields. Generic backends use enable_prefill_cp/cp_strategy directly.
declare_resolution(
server_args,
"validate_deepseek_v4_cp",
enable_dsa_prefill_context_parallel=True,
)
declare_resolution(
server_args,
"validate_deepseek_v4_cp",
enable_prefill_context_parallel=False,
)
declare_resolution(
server_args,
"validate_deepseek_v4_cp",
dsa_prefill_cp_mode="round-robin-split",
)
declare_resolution(
server_args,
"validate_deepseek_v4_cp",
python/sglang/srt/arg_groups/field_order.py
-4
View File
@@ -96,10 +96,6 @@ POSITIONAL_FIELD_ORDER = (
"enable_prefill_cp",
"cp_strategy",
"enable_dsa_cache_layer_split",
"enable_dsa_prefill_context_parallel",
"dsa_prefill_cp_mode",
"enable_prefill_context_parallel",
"prefill_cp_mode",
"enable_cp_decode_attn_tp",
"enable_dp_attention",
"enable_dp_attention_local_control_broadcast",
python/sglang/srt/arg_groups/fields/parallel.py
-4
View File
@@ -161,10 +161,6 @@ class Parallel:
bool,
"Split DSA (DeepSeek Sparse Attention) GPU KV/indexer cache layers across context-parallel ranks to reduce per-rank KV memory. Currently only supported with the mooncake transfer backend (mooncake / mooncake_tcp); mori/nixl support will be added later by the community.",
] = False
enable_dsa_prefill_context_parallel: A[bool, Arg(no_cli=True)] = False
dsa_prefill_cp_mode: A[str, Arg(no_cli=True)] = "round-robin-split"
enable_prefill_context_parallel: A[bool, Arg(no_cli=True)] = False
prefill_cp_mode: A[str, Arg(no_cli=True)] = "in-seq-split"
enable_cp_decode_attn_tp: A[
bool,
"Enable attention tensor-parallel weight slicing during decode under context parallel (cp_size>1). Slices the replicated attention linears to the local CP partition, eliminating redundant decode GEMMs.",
python/sglang/srt/arg_groups/model_hook.py
+1 -3
View File
@@ -284,9 +284,7 @@ def handle_model_specific_adjustments(server_args: Any):
):
raise ValueError(
"--enable-dsa-cache-layer-split requires "
"--enable-prefill-cp and --cp-strategy interleave "
"(or legacy --enable-nsa-prefill-context-parallel with "
"--nsa-prefill-cp-mode round-robin-split)."
"--enable-prefill-cp and --cp-strategy interleave."
)
# Layer split relies on the mooncake all-CP-rank KV/indexer
# transfer path. mori/nixl support is a temporary limitation
python/sglang/srt/arg_groups/parallel_hook.py
+13 -104
View File
@@ -28,18 +28,17 @@ logger = logging.getLogger(__name__)
def handle_context_parallelism(server_args: Any):
validate_prefill_cp_platform(server_args)
cfg = resolving_view(server_args)
if parse_connector_type(cfg.model_path) != ConnectorType.INSTANCE:
model_config = model_config_of(server_args)
hf_config = model_config.hf_config
model_arch = hf_config.architectures[0]
platform = get_platform()
if (
cfg.enable_prefill_cp
and model_arch == "DeepseekV32ForCausalLM"
and cfg.cp_strategy == "zigzag"
and not (platform.is_hip or platform.is_npu or platform.is_musa)
):
raise ValueError(
"DeepSeek V3.2 prefill CP does not support --cp-strategy "
@@ -68,16 +67,6 @@ def handle_context_parallelism(server_args: Any):
"--cp-strategy must be set when --enable-prefill-cp is enabled."
)
if cfg.enable_prefill_context_parallel and cfg.enable_dsa_prefill_context_parallel:
raise ValueError(
"--enable-prefill-context-parallel and "
"--enable-nsa-prefill-context-parallel are mutually "
"exclusive. Use --enable-nsa-prefill-context-parallel for "
"DeepSeek V3.2 (NSA) models and "
"--enable-prefill-context-parallel for MLA-based models "
"(DeepSeek V3/R1, Kimi K2.5) or MHA/GQA-based models."
)
view = resolved_view(server_args)
if view.attn_cp_size > 1:
# The tp_size is the world size, not the real tensor parallel size
@@ -560,98 +549,6 @@ def handle_eplb_and_dispatch(server_args: Any):
assert resolved_view(server_args).ep_size > 1
def handle_platform_cp_compatibility(server_args: Any):
cfg = resolving_view(server_args)
platform = get_platform()
is_protected_platform = platform.is_hip or platform.is_npu or platform.is_musa
if not is_protected_platform:
if (
cfg.enable_prefill_context_parallel
or cfg.enable_dsa_prefill_context_parallel
):
raise ValueError(
"Legacy prefill context-parallel options are supported only "
"by protected HIP, Ascend NPU, or MUSA paths. Use "
"--enable-prefill-cp with --cp-strategy."
)
return
legacy_mode_to_strategy = {
"in-seq-split": "zigzag",
"round-robin-split": "interleave",
}
if cfg.enable_prefill_context_parallel or cfg.enable_dsa_prefill_context_parallel:
declare_resolution(
server_args,
"_handle_platform_cp_compatibility",
enable_prefill_cp=True,
)
if cfg.enable_prefill_context_parallel and cfg.cp_strategy is None:
declare_resolution(
server_args,
"_handle_platform_cp_compatibility",
cp_strategy=legacy_mode_to_strategy[cfg.prefill_cp_mode],
)
if cfg.enable_dsa_prefill_context_parallel and cfg.cp_strategy is None:
declare_resolution(
server_args,
"_handle_platform_cp_compatibility",
cp_strategy=legacy_mode_to_strategy[cfg.dsa_prefill_cp_mode],
)
def handle_legacy_cp_runtime_compatibility(server_args: Any):
"""Project canonical CP settings only for protected platform runtimes."""
platform = get_platform()
if not (platform.is_hip or platform.is_npu or platform.is_musa):
return
cfg = resolving_view(server_args)
if cfg.enable_prefill_context_parallel and cfg.enable_dsa_prefill_context_parallel:
return
if not cfg.enable_prefill_cp or cfg.cp_strategy is None:
return
strategy_to_legacy_mode = {
"zigzag": "in-seq-split",
"interleave": "round-robin-split",
}
mode = strategy_to_legacy_mode[cfg.cp_strategy]
use_dsa_legacy_aliases = cfg.enable_dsa_prefill_context_parallel or getattr(
resolved_view(server_args), "attention_backend", None
) in ("dsa", "dsv4")
if use_dsa_legacy_aliases:
declare_resolution(
server_args,
"_handle_legacy_cp_runtime_compatibility",
enable_dsa_prefill_context_parallel=True,
)
declare_resolution(
server_args,
"_handle_legacy_cp_runtime_compatibility",
enable_prefill_context_parallel=False,
)
else:
declare_resolution(
server_args,
"_handle_legacy_cp_runtime_compatibility",
enable_prefill_context_parallel=True,
)
declare_resolution(
server_args,
"_handle_legacy_cp_runtime_compatibility",
dsa_prefill_cp_mode=mode,
)
declare_resolution(
server_args,
"_handle_legacy_cp_runtime_compatibility",
prefill_cp_mode=mode,
)
def handle_expert_distribution_metrics(server_args: Any):
cfg = resolving_view(server_args)
if "SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC" in os.environ:
@@ -683,3 +580,15 @@ def handle_expert_distribution_metrics(server_args: Any):
"_handle_expert_distribution_metrics",
expert_distribution_recorder_buffer_size=1000,
)
def validate_prefill_cp_platform(server_args: Any):
"""Reject deprecated platform CP before resolving models or CP topology."""
cfg = resolving_view(server_args)
platform = get_platform()
if cfg.enable_prefill_cp and (
platform.is_hip or platform.is_npu or platform.is_musa
):
raise ValueError(
"Prefill CP on HIP/NPU/MUSA is deprecated; CP support will be refactored soon."
)
python/sglang/srt/arg_groups/pipeline.py
+11 -18
View File
@@ -100,10 +100,12 @@ def run_resolution_pipeline(server_args: Any) -> None:
# Reject an explicitly enabled but incompatible hardware runtime before
# model path resolution, downloads, or the dummy-model short circuit.
from sglang.srt.arg_groups.parallel_hook import validate_prefill_cp_platform
from sglang.srt.arg_groups.platform_hook import (
handle_hardware_runtime_validation,
)
validate_prefill_cp_platform(server_args)
handle_hardware_runtime_validation()
if cfg.model_path.lower() in ["none", "dummy"]:
return
@@ -146,21 +148,6 @@ def run_resolution_pipeline(server_args: Any) -> None:
handle_pd_disaggregation(server_args)
# Normalize protected-platform CP aliases before validations or
# model-specific defaults inspect enable_prefill_cp/cp_strategy.
from sglang.srt.arg_groups.parallel_hook import (
handle_context_parallelism,
handle_data_parallelism,
handle_dcp_validation,
handle_dwdp,
handle_elastic_ep,
handle_eplb_and_dispatch,
handle_expert_distribution_metrics,
handle_legacy_cp_runtime_compatibility,
handle_platform_cp_compatibility,
)
handle_platform_cp_compatibility(server_args)
from sglang.srt.arg_groups.kv_cache_hook import (
handle_cache_compatibility,
handle_kv4_compatibility,
@@ -170,6 +157,15 @@ def run_resolution_pipeline(server_args: Any) -> None:
handle_unified_memory_pool,
validate_prefill_only_disable_kv_cache_args,
)
from sglang.srt.arg_groups.parallel_hook import (
handle_context_parallelism,
handle_data_parallelism,
handle_dcp_validation,
handle_dwdp,
handle_elastic_ep,
handle_eplb_and_dispatch,
handle_expert_distribution_metrics,
)
validate_prefill_only_disable_kv_cache_args(server_args)
handle_dcp_validation(server_args)
@@ -287,9 +283,6 @@ def run_resolution_pipeline(server_args: Any) -> None:
# Normalize load balancing defaults.
handle_load_balance_method(server_args)
# Protected runtimes still consume platform CP fields after backend selection.
handle_legacy_cp_runtime_compatibility(server_args)
# Handle context parallelism.
handle_context_parallelism(server_args)
python/sglang/srt/arg_groups/validation_hook.py
-6
View File
@@ -426,16 +426,10 @@ def check_two_batch_overlap(server_args: Any):
# there needs no extra opt-in env flag.
cfg = resolving_view(server_args)
cp_tbo = (
get_platform().is_hip
and cfg.enable_dsa_prefill_context_parallel
and cfg.dsa_prefill_cp_mode == "round-robin-split"
)
if (
cfg.enable_two_batch_overlap
and cfg.moe_a2a_backend == "none"
and not cfg.enable_dp_attention
and not cp_tbo
):
raise ValueError(
"When enabling two batch overlap without an EP a2a backend "
python/sglang/srt/batch_overlap/operations_strategy.py
+4 -24
View File
@@ -34,7 +34,6 @@ class OperationsStrategy:
def init_new_tbo(
layers: torch.nn.ModuleList,
forward_mode: ForwardMode,
use_cp: bool = False,
) -> "OperationsStrategy":
layer_name = layers[0].__class__.__name__
if layer_name == "DeepseekV2DecoderLayer":
@@ -68,7 +67,7 @@ class OperationsStrategy:
return OperationsStrategy.concat(
[
_compute_moe_deepseek_v4_layer_operations_strategy_tbo(
layer, forward_mode, use_cp=use_cp
layer, forward_mode
)
for layer in layers
]
@@ -171,10 +170,9 @@ def _compute_moe_deepseek_blog_decode(layer):
def _compute_moe_deepseek_v4_layer_operations_strategy_tbo(
layer: torch.nn.Module,
forward_mode: ForwardMode,
use_cp: bool = False,
) -> OperationsStrategy:
if forward_mode == ForwardMode.EXTEND:
return _compute_moe_deepseek_v4_prefill(layer, use_cp=use_cp)
return _compute_moe_deepseek_v4_prefill(layer)
else:
# Decode TBO for DSV4 is not implemented yet (ATOM data: decode TBO
# regresses; needs cuda-graph capture work). Prefill-only for now.
@@ -183,28 +181,10 @@ def _compute_moe_deepseek_v4_layer_operations_strategy_tbo(
)
def _compute_moe_deepseek_v4_prefill(layer, use_cp: bool = False):
def _compute_moe_deepseek_v4_prefill(layer):
from sglang.srt.layers.moe import get_moe_a2a_backend
if use_cp:
assert get_moe_a2a_backend().is_none(), (
"DSA prefill CP + TBO is only wired for the non-EP TP-MoE path "
"(moe_a2a_backend == none)."
)
ops = [
layer.op_mhc_prepare_attn,
layer.self_attn.op_attn,
layer.op_mhc_post_attn_pre_mlp,
layer.op_cp_gather_a,
operations.YieldOperation(),
layer.op_cp_gather_b,
layer.op_cp_moe,
layer.op_cp_combine_a,
operations.YieldOperation(),
layer.op_cp_combine_b,
layer.op_mhc_postprocess,
]
elif get_moe_a2a_backend().is_none():
if get_moe_a2a_backend().is_none():
# Non-EP DP TP-MoE: overlap the DP all_gatherv (gather) + reduce_scatterv
# (combine) with the other ubatch's attn+MoE compute (ATOM's DSV4 path).
ops = [
python/sglang/srt/batch_overlap/two_batch_overlap.py
-7
View File
@@ -3,7 +3,6 @@ from __future__ import annotations
import copy
import dataclasses
import logging
import math
from dataclasses import replace
from typing import TYPE_CHECKING, Dict, List, Optional, Sequence
@@ -702,12 +701,6 @@ class TboForwardBatchPreparer:
_tbo_padded_len = (
(end_token_index - start_token_index - 1) // attention_tp_size + 1
) * attention_tp_size
if _is_hip:
from sglang.srt.layers.cp.padding import get_cp_padding_align_size
align = math.lcm(attention_tp_size, get_cp_padding_align_size())
n_tokens = end_token_index - start_token_index
_tbo_padded_len = ((n_tokens + align - 1) // align) * align
output_dict["tbo_padded_len"] = _tbo_padded_len
for key in [
python/sglang/srt/disaggregation/prefill.py
-7
View File
@@ -81,7 +81,6 @@ from sglang.srt.observability.scheduler_stage_metrics import (
)
from sglang.srt.runtime_context import (
get_disagg,
get_parallel,
get_schedule,
)
from sglang.srt.utils import is_npu
@@ -201,12 +200,6 @@ class PrefillBootstrapQueue:
"SGLANG_DISAGG_STAGING_BUFFER with pp_size > 1 is only "
"supported by Mooncake."
)
if get_parallel().enable_prefill_context_parallel:
# CP rewrites index_slice per rank, breaking the chunk grid.
raise RuntimeError(
"SGLANG_DISAGG_STAGING_BUFFER does not support "
"prefill context parallelism."
)
self.kv_manager = self._init_kv_manager()
def _init_kv_manager(self) -> CommonKVManager:
python/sglang/srt/distributed/bootstrap.py
-6
View File
@@ -39,7 +39,6 @@ from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import (
cpu_has_amx_support,
get_available_gpu_memory,
is_hip,
is_host_cpu_arm64,
is_npu,
monkey_patch_p2p_access_check,
@@ -293,11 +292,6 @@ def _init_parallel_groups(
moe_data_model_parallel_size=moe_dp_size,
decode_context_parallel_size=dcp_size,
duplicate_tp_group=get_disagg().enable_pdmux,
duplicate_attn_cp_group=(
is_hip()
and get_exec().overlap.enable_two_batch_overlap
and get_parallel().enable_dsa_prefill_context_parallel
),
enable_symm_mem=get_exec().comm.enable_symm_mem,
recovered_rank=is_ep_joiner,
rank_offset=rank_offset,
python/sglang/srt/distributed/parallel_state.py
-66
View File
@@ -1938,7 +1938,6 @@ def init_model_parallel_group(
_TP: Optional[GroupCoordinator] = None
_ATTN_TP: Optional[GroupCoordinator] = None
_ATTN_CP: Optional[GroupCoordinator] = None
_ATTN_CP_OVERLAP: Optional[GroupCoordinator] = None
_DCP: Optional[GroupCoordinator] = None
# duplicate GroupCoordinator for prefill in PD-Multiplexing
@@ -1976,55 +1975,6 @@ def get_attn_cp_group() -> GroupCoordinator:
return _ATTN_CP
def get_attn_cp_overlap_group() -> GroupCoordinator:
return _ATTN_CP_OVERLAP if _ATTN_CP_OVERLAP is not None else get_attn_cp_group()
def _init_attn_cp_overlap_group(
*,
world_size: int,
attn_cp_size: int,
attn_tp_size: int,
backend: Optional[str],
recovered_rank: bool,
rank_offset: int,
max_world_size: Optional[int],
) -> None:
"""Second communicator over the attention CP ranks; RCCL deadlocks when one
communicator is driven from two streams at once."""
global _ATTN_CP_OVERLAP
assert _ATTN_CP_OVERLAP is None, (
"attention context parallel overlap group is already initialized"
)
if attn_cp_size <= 1:
return
span = attn_tp_size * attn_cp_size
group_ranks = [
list(range(base + i, base + i + span, attn_tp_size))
for base in range(0, world_size, span)
for i in range(attn_tp_size)
]
rank = torch.distributed.get_rank()
mine = next(ranks for ranks in group_ranks if rank in ranks)
assert mine == get_attn_cp_group().ranks, (
f"attn_cp_overlap partition {mine} does not match attn_cp "
f"{get_attn_cp_group().ranks}; the two communicators must span the "
"same ranks or the overlapped collectives will not pair up"
)
_ATTN_CP_OVERLAP = init_model_parallel_group(
group_ranks,
get_world_group().local_rank,
backend,
use_message_queue_broadcaster=False,
group_name="attn_cp_overlap",
recovered_rank=recovered_rank,
rank_offset=rank_offset,
max_world_size=max_world_size,
)
def get_dcp_group_no_assert() -> Optional[GroupCoordinator]:
return _DCP
@@ -2355,7 +2305,6 @@ def initialize_model_parallel(
decode_context_parallel_size: int = 1,
backend: Optional[str] = None,
duplicate_tp_group: bool = False,
duplicate_attn_cp_group: bool = False,
enable_symm_mem: bool = False,
recovered_rank: bool = False,
rank_offset: int = 0,
@@ -2565,17 +2514,6 @@ def initialize_model_parallel(
max_world_size=max_world_size,
)
if duplicate_attn_cp_group and is_hip():
_init_attn_cp_overlap_group(
world_size=world_size,
attn_cp_size=attn_cp_size,
attn_tp_size=attn_tp_size,
backend=backend,
recovered_rank=recovered_rank,
rank_offset=rank_offset,
max_world_size=max_world_size,
)
from sglang.srt.layers.sampler import SYNC_TOKEN_IDS_ACROSS_TP
global _ATTN_TP
@@ -3022,16 +2960,12 @@ def destroy_model_parallel():
_MOE_TP = None
global _ATTN_CP
global _ATTN_CP_OVERLAP
global _MOE_DP
# Destroy _MOE_DP before _ATTN_CP since it may alias _ATTN_CP.
# Only destroy if not aliasing another group.
if _MOE_DP and _MOE_DP is not _ATTN_CP and _MOE_DP is not _TP:
_MOE_DP.destroy()
_MOE_DP = None
if _ATTN_CP_OVERLAP:
_ATTN_CP_OVERLAP.destroy()
_ATTN_CP_OVERLAP = None
if _ATTN_CP:
_ATTN_CP.destroy()
_ATTN_CP = None
python/sglang/srt/hardware_backend/musa/layers/utils/cp_utils.py
+3 -38
View File
@@ -16,42 +16,7 @@ def musa_cp_attn_forward_extend(
device: torch.device,
attn_fn: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, int], torch.Tensor],
) -> torch.Tensor:
"""
Split q into prev/next zigzag halves based on CP metadata, call the
backend-specific attention function twice with appropriate per-half
metadata, and concatenate the results.
"""Retained import for the MUSA backend; legacy CP execution is deprecated."""
from sglang.srt.layers.utils.cp_utils import _deprecated_platform_cp
attn_fn signature:
attn_fn(q, cu_seqlens_q, cache_seqlens, max_seqlen_q) -> result
where only these four CP-varying parameters differ between halves.
All other backend-specific args should be captured in the closure.
"""
cp_meta = forward_batch.attn_cp_metadata
q_prev, q_next = torch.chunk(q, 2, dim=0)
cu_seqlens_q_prev = torch.tensor(
[0, cp_meta.actual_seq_q_prev], device=device, dtype=torch.int32
)
if hasattr(musa_fa_backend, "_current_prefix"):
musa_fa_backend._current_prefix = "forward_extend_cp_prev"
result_prev = attn_fn(
q_prev,
cu_seqlens_q_prev,
cp_meta.kv_len_prev_tensor,
cp_meta.actual_seq_q_prev,
)
cu_seqlens_q_next = torch.tensor(
[0, cp_meta.actual_seq_q_next], device=device, dtype=torch.int32
)
if hasattr(musa_fa_backend, "_current_prefix"):
musa_fa_backend._current_prefix = "forward_extend_cp_next"
result_next = attn_fn(
q_next,
cu_seqlens_q_next,
cp_meta.kv_len_next_tensor,
cp_meta.actual_seq_q_next,
)
return torch.concat([result_prev, result_next], dim=0)
_deprecated_platform_cp()
python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+4 -24
View File
@@ -121,7 +121,6 @@ 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
from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_output
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.forward_context import (
get_attn_backend,
@@ -530,11 +529,7 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
with torch.cuda.stream(self.alt_stream):
key = self._maybe_rotate(key)
current_stream.wait_stream(self.alt_stream)
elif (
self.alt_stream is not None
and forward_batch.attn_cp_metadata is not None
and self.dsa_enable_prefill_cp
):
elif self.alt_stream is not None and is_cp_v2_active(forward_batch):
key = self._maybe_rotate(key)
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
@@ -543,17 +538,9 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
# 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):
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(),
)
key = get_cp_strategy().materialize_full_indexer_k_cache(
key, forward_batch
)
current_stream.wait_stream(self.alt_stream)
return query, key, weights_raw
else:
@@ -563,13 +550,6 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
# allgather+rerrange
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,
forward_batch,
torch.cuda.current_stream(),
)
return query, key, weights_raw
def _get_k_bf16(
python/sglang/srt/layers/attention/dsa/utils.py
+8 -135
View File
@@ -1,5 +1,5 @@
from functools import lru_cache
from typing import TYPE_CHECKING, List, Tuple, Union
from typing import TYPE_CHECKING
import torch
import triton
@@ -19,7 +19,7 @@ from sglang.srt.runtime_context import (
process_model_config,
)
from sglang.srt.utils import get_bool_env_var, is_cuda, is_hip, is_musa, is_npu
from sglang.srt.utils.common import ceil_align, ceil_div
from sglang.srt.utils.common import ceil_div
@lru_cache(maxsize=1)
@@ -116,10 +116,10 @@ def should_use_dsa_fused_topk(seed_dsa_topk_from_draft_extend: bool) -> bool:
def is_dsa_enable_prefill_cp():
if is_hip() or is_npu() or is_musa():
return get_parallel().enable_dsa_prefill_context_parallel
return False
# Generic prefill CP derives activation from the runtime topology and model
# architecture. Protected HIP/NPU paths continue to use their legacy field.
# architecture.
if get_parallel().attn_cp_size <= 1:
return False
from sglang.srt.configs.model_config import is_deepseek_dsa, is_deepseek_v4
@@ -129,10 +129,7 @@ def is_dsa_enable_prefill_cp():
def is_dsa_prefill_cp_round_robin_split():
return (
is_dsa_enable_prefill_cp()
and get_parallel().dsa_prefill_cp_mode == "round-robin-split"
)
return is_dsa_enable_prefill_cp() and get_parallel().cp_strategy == "interleave"
# Structural surface where the graph DSA split-op dispatch (DSA indexer) and the
@@ -161,53 +158,10 @@ def can_dsa_prefill_cp_round_robin_split(forward_batch: "ForwardBatch"):
)
def dsa_cp_round_robin_split_data(input_: Union[torch.Tensor, List]):
"""
# for round-robin-split, split the tokens evenly according to the rule of token_idx % cp_size.
| +-----------before split------------+|
| token0, token1, token2, token3, token4, token5, token6, token7, ...
|
| +--------------result-------------------+
| dp_atten_tp0: token0, token4, token8, token12, token16, ... |
| dp_atten_tp1: token1, token5, token9, token13, token17, ... |
| dp_atten_tp2: token2, token6, token10, token14, token18, ... |
| dp_atten_tp3: token3, token7, token11, token15, token19, ... |
| +-------------------------+
"""
cp_size = get_parallel().attn_cp_size
cp_rank = get_parallel().attn_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]
# for torch device tensor
shard = input_.view(-1, cp_size, *input_.shape[1:])[:, cp_rank]
# .contiguous() is not sufficient here. When tokens == cp_size every rank's
# shard has a single row, and a size-1 outer dimension imposes no contiguity
# constraint, so is_contiguous() is True whatever stride(0) is and
# .contiguous() becomes a no-op. The shard then keeps the cp_size-inflated
# row pitch (cp_size * row_numel instead of row_numel), which any kernel that
# takes its row pitch from stride(0) will read as an oversized tensor.
# Compare the pitch against the parent's explicitly, so the copy happens
# exactly when the shard really is strided -- and not at all for cp_size == 1.
if shard.stride(0) != input_.stride(0):
shard = shard.clone(memory_format=torch.contiguous_format)
return shard
def cal_padded_tokens(forward_batch: "ForwardBatch"):
# Consistent with the padding calculation logic in ForwardBatch.prepare_mlp_sync_batch,
# calculate the actual token length after padding when attn_tp_size > 1 or in the MAX_LEN padding mode.
from sglang.srt.layers.cp.padding import get_cp_padding_align_size
from sglang.srt.layers.cp.utils import enable_cp_v2, is_cp_v2_active
from sglang.srt.layers.cp.utils import is_cp_v2_active
# CP-v2 already pads each rank-local shard to its physical size
if is_cp_v2_active(forward_batch):
@@ -216,18 +170,9 @@ def cal_padded_tokens(forward_batch: "ForwardBatch"):
]
global_num_tokens = forward_batch.global_num_tokens_cpu.copy()
sync_group_size = len(global_num_tokens)
attn_cp_size = get_parallel().attn_cp_size
# Must mirror ForwardBatch.prepare_mlp_sync_batch, which applies cp_align_size only when
# CP-v2 is disabled. Under enable_cp_v2() the speculative forwards (TARGET_VERIFY /
# DRAFT_EXTEND_V2) reach here with is_cp_v2_active False, and q is padded to attn_tp_size only
# (not cp-aligned). Applying cp_align here over-pads the flashmla metadata past q, so
# num_splits ends up longer than q -> fwd_kvcache_mla fails "num_splits must have shape (b+1)".
# (attn_cp analog of the attn_tp fix in PR #30642 / issue #30296.)
if not enable_cp_v2():
cp_align_size = get_cp_padding_align_size()
for i in range(sync_group_size):
global_num_tokens[i] = ceil_align(global_num_tokens[i], cp_align_size)
# Non-CP forwards (including speculative forwards) use attention-TP padding
# only, matching ForwardBatch.prepare_mlp_sync_batch.
# Reuse the mode selected when the DP buffer was prepared.
dp_padding_mode = forward_batch.dp_padding_mode
if dp_padding_mode is None:
@@ -265,78 +210,6 @@ def pad_dsa_cache_seqlens(forward_batch: "ForwardBatch", dsa_cache_seqlens):
return 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 seq_len % cp_size == 0, (
f"seq_len {seq_len} is not divisible by cp_size {cp_size} when dsa_prefill_cp_mode is round-robin-split"
)
else:
# TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2
# 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)
return cur_cp_seq_len != 0
from sglang.kernels.ops.attention.dsa.cp_split import (
dsa_cp_round_robin_split_q_seqs_kernel,
)
def dsa_cp_round_robin_split_q_seqs_cpu(extend_seqs):
cp_size = get_parallel().attn_cp_size
cp_rank = get_parallel().attn_cp_rank
extra_seq = 0
q_seqs = []
for bs, cur_len in enumerate(extend_seqs):
cur_len += extra_seq
cur_seq = cur_len // cp_size + int(cur_len % cp_size > cp_rank)
q_seqs.append(cur_seq)
extra_seq = cur_len - cur_seq * cp_size
bs_idx = list([i for i, x in enumerate(q_seqs) if x > 0])
q_seqs = [q_len for q_len in q_seqs if q_len > 0]
return q_seqs, bs_idx
def dsa_cp_round_robin_split_q_seqs(
extend_seqs_cpu, extend_seqs
) -> Tuple[List, torch.Tensor, List, torch.Tensor]:
"""
round-robin-split distributes tokens across ranks based on token_idx % cp_size.
Return:
ret_q_lens_cpu(List) and ret_q_lens(torch.Tensor): the partitioned length (excluding zeros) on the current cp rank
for each sequence after distribution across cp ranks.
bs_idx_cpu(List) and bs_idx(torch.Tensor): marks which sequences are ultimately selected,
i.e., those with a partitioned length greater than zero.
"""
cp_size = get_parallel().attn_cp_size
cp_rank = get_parallel().attn_cp_rank
# len(ret_q_lens_cpu) == len(bs_idx_cpu)
ret_q_lens_cpu, bs_idx_cpu = dsa_cp_round_robin_split_q_seqs_cpu(extend_seqs_cpu)
ret_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
)
grid = (1,)
dsa_cp_round_robin_split_q_seqs_kernel[grid](
extend_seqs, ret_q_lens, bs_idx, len(extend_seqs), cp_size, cp_rank
)
return ret_q_lens_cpu, ret_q_lens, bs_idx_cpu, bs_idx
def dsa_use_prefill_cp(forward_batch, dsa_enable_prefill_cp=None):
if dsa_enable_prefill_cp is None:
dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
python/sglang/srt/layers/attention/dsa_backend.py
+12 -80
View File
@@ -76,8 +76,6 @@ from sglang.srt.layers.attention.dsa.kpool_plan import (
from sglang.srt.layers.attention.dsa.utils import (
can_dsa_prefill_cp_round_robin_split,
compute_dsa_seqlens,
dsa_cp_round_robin_split_data,
dsa_cp_round_robin_split_q_seqs,
dsa_use_prefill_cp,
is_dsa_enable_prefill_cp,
pad_dsa_cache_seqlens,
@@ -89,10 +87,6 @@ from sglang.srt.layers.attention.trtllm_mla_backend import (
)
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,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.runtime_context import get_buffer, get_exec, get_parallel, get_spec
from sglang.srt.utils import (
@@ -122,24 +116,6 @@ if TYPE_CHECKING:
from sglang.srt.speculative.spec_info import SpecInput
def _all_gather_dsa_trtllm_fp8_kv(
forward_batch: ForwardBatch,
k: torch.Tensor,
k_rope: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
kv_lora_rank = k.shape[-1]
qk_rope_head_dim = k_rope.shape[-1]
kv_dtype = k.dtype
kv = torch.cat((k, k_rope), dim=-1).view(torch.uint8)
kv = cp_all_gather_rerange_output(
kv,
get_parallel().attn_cp_size,
forward_batch,
torch.cuda.current_stream(),
).view(kv_dtype)
return kv.split((kv_lora_rank, qk_rope_head_dim), dim=-1)
def prepare_kv_for_attention(
attn_mla,
forward_batch: ForwardBatch,
@@ -165,28 +141,6 @@ def prepare_kv_for_attention(
)
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]:
"""Materialize generic CP KV, retaining the ROCm DSA fallback."""
if is_cp_v2_active(forward_batch):
strategy = get_cp_strategy()
assert strategy is not None
return strategy.materialize_full_mla_kv(
forward_batch,
attn_mla.attn_mqa,
k_nope,
k_pe,
)
assert is_hip(), "Legacy DSA KV materialization is HIP-only"
return attn_mla.rebuild_cp_kv_cache(latent_cache, forward_batch, k_nope, k_pe)
_is_hip = is_hip()
_is_xpu = is_xpu()
@@ -1052,19 +1006,11 @@ class DeepseekSparseAttnBackend(
kpool_inputs.full_seqlens_expanded = seqlens_expanded
if can_dsa_prefill_cp_round_robin_split(forward_batch):
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
)
)
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)
)
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]
@@ -1277,11 +1223,7 @@ 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)
split_per_token = (
get_cp_strategy().shard_local_tokens
if is_cp_v2_active(forward_batch)
else dsa_cp_round_robin_split_data
)
split_per_token = get_cp_strategy().shard_local_tokens
ks = split_per_token(ks)
ke = split_per_token(ke)
token_to_batch_idx = split_per_token(token_to_batch_idx)
@@ -3500,14 +3442,9 @@ class DeepseekSparseAttnBackend(
else:
rope_positions = forward_batch.positions
if dsa_use_prefill_cp(forward_batch):
if is_cp_v2_active(forward_batch):
rope_positions = get_cp_strategy().shard_position_ids(
rope_positions, forward_batch
)
else:
rope_positions = cp_split_and_rebuild_position(
forward_batch, rope_positions
)
rope_positions = get_cp_strategy().shard_position_ids(
rope_positions, forward_batch
)
q, k, k_rope = mla_quantize_and_rope_for_fp8(
q,
@@ -3521,14 +3458,9 @@ class DeepseekSparseAttnBackend(
self.qk_rope_head_dim,
)
if save_kv_cache and dsa_use_prefill_cp(forward_batch):
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
)
k, k_rope = get_cp_strategy().all_gather_dsa_trtllm_fp8_kv(
forward_batch, k, k_rope
)
merge_query = False
# Save KV cache if requested
python/sglang/srt/layers/attention/dsv4/compressor.py
+1 -35
View File
@@ -22,17 +22,13 @@ from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
from sglang.srt.layers.cp.utils import cp_materialize_global_token_order
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.utils.cp_utils import (
cp_all_gather_rerange_finish,
cp_all_gather_rerange_launch,
)
from sglang.srt.mem_cache.deepseek_v4_compress_state import (
CompressStatePool,
)
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.model_executor.forward_context import get_attn_backend
from sglang.srt.models.deepseek_v2 import _is_hip
from sglang.srt.runtime_context import get_exec, get_parallel
from sglang.srt.runtime_context import get_exec
from sglang.srt.utils import add_prefix, is_npu, set_weight_attrs
_is_npu = is_npu()
@@ -445,37 +441,7 @@ class Compressor(BaseFusedOp):
assert isinstance(ret, CompressStatePool)
return ret
def _pending_key(self):
return ("kv_score", self.layer_id, self.is_in_indexer)
def prelaunch_kv_score(self, x: torch.Tensor, forward_batch: ForwardBatch):
"""Compute kv_score and start its CP all-gather, without waiting.
kv_score only needs `x`, which the attention already has at entry, so the
gather can be issued before the q/kv projections and collected later in
compute_kv_score -- that projection work is what hides it. Caller must
guarantee a matching compute_kv_score in the same op (see
DeepseekV4Attention._forward_prepare).
"""
if not _is_hip:
return
comm_stream = getattr(forward_batch, "_cp_prefetch_comm_stream", None)
if comm_stream is None or not dsa_use_prefill_cp(forward_batch):
return
kv_score = self._compute_wkv_gate(x)
# Keyed by forward_batch: each TBO ubatch carries its own, so the two
# ubatches cannot collect each other's gather.
pending = forward_batch.__dict__.setdefault("_cp_pending_gathers", {})
pending[self._pending_key()] = cp_all_gather_rerange_launch(
kv_score, get_parallel().attn_cp_size, comm_stream, self._pending_key()
)
def compute_kv_score(self, x: torch.Tensor, forward_batch: ForwardBatch):
if _is_hip:
pending = getattr(forward_batch, "_cp_pending_gathers", None)
handle = pending.pop(self._pending_key(), None) if pending else None
if handle is not None:
return cp_all_gather_rerange_finish(handle)
kv_score = self._compute_wkv_gate(x)
python/sglang/srt/layers/attention/flashattention_backend.py
+1 -20
View File
@@ -24,7 +24,7 @@ from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
from sglang.srt.layers.cp.base import CPAttentionBackendKind, get_cp_strategy
from sglang.srt.layers.cp.utils import enable_cp_v2, is_cp_v2_active
from sglang.srt.layers.cp.utils import is_cp_v2_active
from sglang.srt.layers.radix_attention import AttentionType
from sglang.srt.mem_cache.memory_pool import KVWriteLoc
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
@@ -1001,25 +1001,6 @@ class FlashAttentionBackend(AttentionBackend):
torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)
)
# MLA/MHA CP: prepare_mlp_sync_batch pads extend tokens up to
# lcm(attn_tp_size, attn_cp_size), so cache_seqlens_cp can exceed
# seq_lens_cpu.max(). Widen page_table by the pad delta to keep
# FA3's causal reads in-bounds; widened columns index KV slot 0
# (req_to_token is zero-init) and outputs for padding queries are
# discarded downstream.
if (
not enable_cp_v2()
and self.attn_cp_size > 1
and forward_batch.global_num_tokens_cpu is not None
and forward_batch.extend_num_tokens is not None
and forward_batch.extend_seq_lens_cpu is not None
):
padded_extend = int(forward_batch.extend_num_tokens)
real_extend = int(sum(forward_batch.extend_seq_lens_cpu))
pad_delta = padded_extend - real_extend
if pad_delta > 0:
metadata.max_seq_len_k += pad_delta
metadata.page_table = self.req_to_token_pool.req_to_token[
forward_batch.req_pool_indices, : metadata.max_seq_len_k
]
python/sglang/srt/layers/cp/utils.py
+6 -22
View File
@@ -146,20 +146,12 @@ def is_cp_v2_active(forward_batch) -> bool:
def is_mla_prefill_cp_enabled() -> bool:
"""Return whether prefill CP is configured for an MLA attention backend."""
if enable_cp_v2():
return is_cp_enabled() and uses_mla_backend()
return get_parallel().enable_prefill_context_parallel and uses_mla_backend()
return enable_cp_v2() and is_cp_enabled() and uses_mla_backend()
def mla_use_prefill_cp(forward_batch) -> bool:
"""Return whether this MLA forward batch is using prefill CP."""
if enable_cp_v2():
return is_mla_prefill_cp_enabled() and is_cp_v2_active(forward_batch)
return (
getattr(forward_batch, "attn_cp_metadata", None) is not None
and is_mla_prefill_cp_enabled()
and forward_batch.forward_mode.is_context_parallel_extend()
)
return is_mla_prefill_cp_enabled() and is_cp_v2_active(forward_batch)
def prepare_cp_forward(forward_batch) -> None:
@@ -265,18 +257,10 @@ def cp_materialize_global_token_order(
x: Any, forward_batch, stream: Optional[Any] = None
):
"""Materialize a CP tensor in the global logical token order."""
if is_cp_v2_active(forward_batch):
strategy = get_cp_strategy()
assert strategy is not None
return strategy.gather_kv_cache(x, forward_batch, stream)
# HIP/NPU still materialize their protected platform layout through the
# legacy collective until those backends migrate independently.
from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_output
return cp_all_gather_rerange_output(
x, get_parallel().attn_cp_size, forward_batch, stream
)
assert is_cp_v2_active(forward_batch)
strategy = get_cp_strategy()
assert strategy is not None
return strategy.gather_kv_cache(x, forward_batch, stream)
@contextmanager
python/sglang/srt/layers/dp_attention.py
-9
View File
@@ -17,7 +17,6 @@ from sglang.srt.arg_groups.model_override_base import (
from sglang.srt.distributed import (
GroupCoordinator,
get_attn_cp_group,
get_attn_cp_overlap_group,
get_attn_tensor_model_parallel_rank,
get_attn_tensor_model_parallel_world_size,
get_attn_tp_group,
@@ -1038,14 +1037,6 @@ def attn_cp_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor):
return get_attn_cp_group().all_gather_into_tensor(output, input)
def attn_cp_overlap_all_gather_into_tensor(output: torch.Tensor, input: torch.Tensor):
return get_attn_cp_overlap_group().all_gather_into_tensor(output, input)
def attn_cp_overlap_reduce_scatter_tensor(output: torch.Tensor, input: torch.Tensor):
return get_attn_cp_overlap_group().reduce_scatter_tensor(output, input)
def get_moe_cp_group() -> GroupCoordinator:
"""Returns the MOE_DP group, which includes CP partners when attn_cp_size > moe_dp_size."""
return _get_moe_dp_group()
python/sglang/srt/layers/utils/cp_utils.py
+10 -665
View File
@@ -1,678 +1,23 @@
"""Legacy prefill CP helpers retained for HIP, NPU, and MUSA callers."""
"""Import-only shims for deprecated platform backends awaiting CP refactoring.
from dataclasses import dataclass
from itertools import accumulate
from typing import List
import torch
import torch.nn.functional as F
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.layers.dp_attention import (
_tbo_event,
attn_cp_all_gather_into_tensor,
attn_cp_overlap_all_gather_into_tensor,
is_allocation_symmetric,
)
from sglang.srt.layers.moe import get_moe_a2a_backend
from sglang.srt.mem_cache.memory_pool import KVWriteLoc
from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
from sglang.srt.runtime_context import (
get_parallel,
uses_mla_backend,
)
The legacy CP algorithms have been removed. These names keep the retained
NPU/MUSA attention backends importable for non-CP inference; calling them fails.
"""
@dataclass
class ContextParallelMetadata:
# Layout lists have length bs * cp_segment_num (= bs * 2 * cp_size).
split_list: List[int] = None
zigzag_index: List[int] = None
cp_reverse_index: List[int] = None
reverse_split_len: List[int] = None
# Per-rank-aggregate lists have length cp_size.
# max_rank_len is a list of cp_size copies of max(per_rank_actual_token),
# kept as a list for torch.split() bucket sizes.
per_rank_actual_token: List[int] = None
max_rank_len: List[int] = None
# Per-sequence FlashAttention tensors (shape [bs] or [bs+1]).
kv_len_prev_tensor: torch.Tensor = None # [bs] int32 CUDA
kv_len_next_tensor: torch.Tensor = None # [bs] int32 CUDA
actual_seq_q_prev_tensor: torch.Tensor = None # [bs] int32 CUDA
actual_seq_q_next_tensor: torch.Tensor = None # [bs] int32 CUDA
cu_seqlens_q_prev_tensor: torch.Tensor = None # [bs+1] int32 CUDA
cu_seqlens_q_next_tensor: torch.Tensor = None # [bs+1] int32 CUDA
# Scalars derived from the per-sequence lists above.
total_q_prev_tokens: int = 0
total_q_next_tokens: int = 0
max_seqlen_q_prev: int = 0
max_seqlen_q_next: int = 0
# Per-seq CPU lists (useful for NSA indexer and diagnostics).
kv_len_prev_list: List[int] = None
kv_len_next_list: List[int] = None
actual_seq_q_prev_list: List[int] = None
actual_seq_q_next_list: List[int] = None
# Aggregate sum of extend_seq_lens across the batch.
total_seq_lens: int = 0
bs: int = 1
def is_prefill_context_parallel_enabled():
return get_parallel().enable_prefill_context_parallel
def is_mla_prefill_cp_enabled() -> bool:
return get_parallel().enable_prefill_context_parallel and uses_mla_backend()
def mla_use_prefill_cp(forward_batch, mla_enable_prefill_cp=None):
if mla_enable_prefill_cp is None:
mla_enable_prefill_cp = is_mla_prefill_cp_enabled()
return (
forward_batch.attn_cp_metadata is not None
and mla_enable_prefill_cp
and forward_batch.forward_mode.is_context_parallel_extend()
)
def can_cp_split(seq_len: int, cp_size: int, forward_batch):
# Base conditions: CP must be enabled, size > 1, and this must be a
# CP-extend (prefill) step. The seq_len // (cp_size * 2) check ensures
# the load-balancing split into 2 * cp_size blocks is non-degenerate.
from sglang.srt.model_executor.forward_batch_info import ForwardMode
cur_cp_seq_len = seq_len // (cp_size * 2)
if not (
cur_cp_seq_len != 0
and cp_size > 1
# prepare_context_parallel_metadata hard-codes bs_per_cp_group = 1;
# guard explicitly to avoid silent mis-partitioning under continuous batching.
and forward_batch.forward_mode.is_context_parallel_extend()
# is_context_parallel_extend() returns True for MIXED (prefill+decode
# in one step), but the zigzag split only makes sense on pure extend.
and forward_batch.forward_mode != ForwardMode.MIXED
and is_prefill_context_parallel_enabled()
):
return False
# Per-sequence guards for bs > 1. Every sequence must be long enough for
# the 2*cp_size-way split. A sub-threshold request reaching this point
# means the scheduler failed to filter it out and a silent non-CP
# fallback would have masked the bug -- raise instead. Per-sequence
# radix-cache prefix is supported: prefix is baked into kv_len_prev/next
# via prefix_offsets[s] inside prepare_context_parallel_metadata.
extend_lens = getattr(forward_batch, "extend_seq_lens_cpu", None)
if extend_lens is None:
return True
cp_min = cp_size * 2
for L in extend_lens:
if L < cp_min:
# A sub-threshold request cannot be zigzag-split into 2*cp_size
# blocks; fall back to a normal (non-CP) prefill for this batch
# instead of failing. Happens e.g. when a radix-cache prefix hit
# leaves only a few unique extend tokens.
return False
return True
def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
from sglang.srt.layers.attention.dsa.utils import (
dsa_cp_round_robin_split_data,
is_dsa_prefill_cp_round_robin_split,
)
if is_dsa_prefill_cp_round_robin_split():
cp_size = get_parallel().attn_cp_size
assert input_.shape[0] % cp_size == 0, (
f"Expect input shape 0 can divided by cp size, but got input shape {input_.shape}, cp size {cp_size}"
)
return dsa_cp_round_robin_split_data(input_)
input_list = list(
torch.split(input_, forward_batch.attn_cp_metadata.split_list, dim=0)
)
result = torch.cat(
[input_list[i] for i in forward_batch.attn_cp_metadata.zigzag_index], dim=0
).view(-1, input_.shape[-1])
return result
def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
from sglang.srt.layers.attention.dsa.utils import (
dsa_cp_round_robin_split_data,
is_dsa_prefill_cp_round_robin_split,
)
if is_dsa_prefill_cp_round_robin_split():
cp_size = get_parallel().attn_cp_size
assert positions.shape[0] % cp_size == 0, (
f"Expect positions shape 0 can divided by cp size, but got positions shape {positions.shape}, "
f"cp size {cp_size}"
)
return dsa_cp_round_robin_split_data(positions)
position_id_list = list(
torch.split(positions, forward_batch.attn_cp_metadata.split_list, dim=-1)
)
positions = torch.cat(
[position_id_list[i] for i in forward_batch.attn_cp_metadata.zigzag_index],
dim=-1,
)
return positions
def cp_round_robin_input_ids(input_ids):
"""
input input_ids:
rank0~7: 0,1,2,3,4,5,...
output input_ids:
a2a none:
rank0~7: 0,8,16,...,1,9,17,...,2,10,18,...
not a2a none:
rank0: 0,8,16,...
rank1: 1,9,17,...
rank2: 2,10,18,...
...
"""
cp_size = get_parallel().attn_cp_size
cp_rank = get_parallel().attn_cp_rank
if get_moe_a2a_backend().is_none():
input_ids = input_ids.reshape(-1, cp_size).T.flatten()
else:
input_ids = input_ids[cp_rank::cp_size].contiguous()
return input_ids
def cp_all_gather_reorganized_into_tensor(input_tensor, cp_size, forward_batch, stream):
"""
Allgather communication for context_parallel(kv_cache, index_k, hidden_states).
This implementation mainly consists of three parts:
Step 1, padding the input shape to unify the shape for allgather communication (the shape must be the same).
Step 2, synchronized allgather communication.
Step 3, removing the padding and reassembling the data according to the actual tokens.
"""
max_len = forward_batch.attn_cp_metadata.max_rank_len[0]
pad_size = max_len - input_tensor.shape[0]
if pad_size > 0:
input_tensor = F.pad(
input_tensor, (0, 0, 0, pad_size), mode="constant", value=0
)
group = get_parallel().attn_cp_group
with use_symmetric_memory(group, disabled=not is_allocation_symmetric()):
input_tensor_full = torch.empty(
max_len * cp_size,
input_tensor.shape[1],
device=input_tensor.device,
dtype=input_tensor.dtype,
)
group.all_gather_into_tensor(input_tensor_full, input_tensor)
outputs_list_max = list(
torch.split(
input_tensor_full, forward_batch.attn_cp_metadata.max_rank_len, dim=0
)
)
outputs = torch.cat(
[
outputs_list_max[index][:per_rank_len]
for index, per_rank_len in enumerate(
forward_batch.attn_cp_metadata.per_rank_actual_token
)
],
dim=0,
)
return outputs
def cp_all_gather_reorganized_into_tensor_kv_cache(
input_tensor, cp_size, forward_batch, stream
):
"""
Allgather communication for context_parallel KV cache.
Handles multi-dimensional tensors (e.g., [seq_len, num_heads, head_dim]).
"""
max_len = forward_batch.attn_cp_metadata.max_rank_len[0]
pad_size = max_len - input_tensor.shape[0]
if pad_size > 0:
# Pad the first dimension (seq_len). F.pad expects padding in reverse dimension order.
# For n dimensional tensor, we need 2*n values: (last_dim_left, last_dim_right, ..., first_dim_left, first_dim_right)
# To pad only the first dimension: [0, 0] * (ndim - 1) + [0, pad_size]
padding = [0, 0] * (input_tensor.ndim - 1) + [0, pad_size]
input_tensor = F.pad(input_tensor, padding, mode="constant", value=0)
# Create output tensor with proper shape for all dimensions
group = get_parallel().attn_cp_group
with use_symmetric_memory(group, disabled=not is_allocation_symmetric()):
input_tensor_full = torch.empty(
max_len * cp_size,
*input_tensor.shape[1:],
device=input_tensor.device,
dtype=input_tensor.dtype,
)
group.all_gather_into_tensor(input_tensor_full, input_tensor)
outputs_list_max = list(
torch.split(
input_tensor_full, forward_batch.attn_cp_metadata.max_rank_len, dim=0
)
)
outputs = torch.cat(
[
outputs_list_max[index][:per_rank_len]
for index, per_rank_len in enumerate(
forward_batch.attn_cp_metadata.per_rank_actual_token
)
],
dim=0,
)
return outputs
def cp_all_gather_rerange_launch(input_tensor, cp_size, comm_stream, event_key):
"""Start a round-robin CP all-gather on `comm_stream`; do NOT wait for it.
Pair with cp_all_gather_rerange_finish(). Splitting launch from wait is the
only way an attention-side CP gather can overlap anything: the collectives
inside op_attn are consumed a few statements later, so issuing and waiting
at the same point just moves the queue (measured in perf_sweep_report §4.6).
The handle keeps both buffers alive until finish(); without that reference
the allocator can hand the input block back to the compute stream before the
comm-stream kernel has read it.
"""
from sglang.srt.distributed.parallel_state import (
get_attn_cp_group,
get_attn_cp_overlap_group,
)
group = get_attn_cp_overlap_group()
assert group is not get_attn_cp_group(), (
"the comm-stream path needs the duplicate attn_cp_overlap communicator; "
"driving one communicator from two streams deadlocks RCCL"
)
input_tensor = input_tensor.contiguous()
with use_symmetric_memory(group, disabled=not is_allocation_symmetric()):
output_tensor = input_tensor.new_empty(
(input_tensor.shape[0] * cp_size, *input_tensor.shape[1:]),
)
comm_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(comm_stream):
attn_cp_overlap_all_gather_into_tensor(output_tensor, input_tensor)
event = _tbo_event(event_key)
event.record(comm_stream)
return (output_tensor, input_tensor, event, cp_size)
def cp_all_gather_rerange_finish(handle):
"""Wait for a launched gather on the current stream, then rerange."""
output_tensor, _keepalive, event, cp_size = handle
torch.cuda.current_stream().wait_event(event)
out_shape = output_tensor.shape
return (
output_tensor.view(cp_size, -1, *out_shape[1:])
.transpose(0, 1)
.reshape(out_shape)
def _deprecated_platform_cp():
raise ValueError(
"Prefill CP on HIP/NPU/MUSA is deprecated; CP support will be refactored soon."
)
def cp_all_gather_rerange_output(input_tensor, cp_size, forward_batch, stream):
"""
# for in-seq-split
| +-----------before allgather------------+|
| | dp_atten_tp0: block0, block7 |
| | dp_atten_tp1: block1, block6 |
| | dp_atten_tp2: block2, block5 |
| | dp_atten_tp3: block3, block4 |
|
| +----------before rerange---------------+|
| block0 | block7 | block1 | block6 | block2 | block5 | block3 | block4 |
|
| +--------------result-------------------+
| block0 | block1 | block2 | block3 | block4 | block5 | block6 | block7 |
| +-------------------------+
# for round-robin-split
| +-----------before allgather------------+|
| dp_atten_tp0: token0, token4, token8, token12, token16, ... |
| dp_atten_tp1: token1, token5, token9, token13, token17, ... |
| dp_atten_tp2: token2, token6, token10, token14, token18, ... |
| dp_atten_tp3: token3, token7, token11, token15, token19, ... |
|
| +--------------result-------------------+
| token0, token1, token2, token3, token4, token5, token6, token7, ...
| +-------------------------+
"""
from sglang.srt.layers.attention.dsa.utils import (
is_dsa_prefill_cp_round_robin_split,
)
if is_dsa_prefill_cp_round_robin_split():
with use_symmetric_memory(
get_parallel().attn_cp_group, disabled=not is_allocation_symmetric()
):
output_tensor = input_tensor.new_empty(
(input_tensor.shape[0] * cp_size, *input_tensor.shape[1:]),
)
attn_cp_all_gather_into_tensor(
output_tensor,
input_tensor,
)
out_shape = output_tensor.shape
output_tensor = (
output_tensor.view(cp_size, -1, *out_shape[1:])
.transpose(0, 1)
.reshape(out_shape)
)
return output_tensor
# TODO: Do we need to remove the padding here?
bs_seq_len, hidden_size = input_tensor.shape
output_tensor = cp_all_gather_reorganized_into_tensor(
input_tensor,
cp_size,
forward_batch,
stream,
)
outputs_list = list(
torch.split(
output_tensor, forward_batch.attn_cp_metadata.reverse_split_len, dim=0
)
)
output_tensor = torch.cat(
[outputs_list[i] for i in forward_batch.attn_cp_metadata.cp_reverse_index],
dim=0,
)
output_tensor = output_tensor.view(-1, hidden_size)
return output_tensor
_deprecated_platform_cp()
def cp_all_gather_rerange_kv_cache(input_tensor, cp_size, forward_batch, stream):
"""
Allgather and reorganize KV cache from all ranks in context parallel group.
# for in-seq-split
| +-----------before allgather------------+|
| | dp_atten_tp0: block0, block7 |
| | dp_atten_tp1: block1, block6 |
| | dp_atten_tp2: block2, block5 |
| | dp_atten_tp3: block3, block4 |
|
| +----------before rerange---------------+|
| block0 | block7 | block1 | block6 | block2 | block5 | block3 | block4 |
|
| +--------------result-------------------+
| block0 | block1 | block2 | block3 | block4 | block5 | block6 | block7 |
| +-------------------------+
"""
output_tensor = cp_all_gather_reorganized_into_tensor_kv_cache(
input_tensor,
cp_size,
forward_batch,
stream,
)
outputs_list = list(
torch.split(
output_tensor, forward_batch.attn_cp_metadata.reverse_split_len, dim=0
)
)
output_tensor = torch.cat(
[outputs_list[i] for i in forward_batch.attn_cp_metadata.cp_reverse_index],
dim=0,
)
# No need to reshape - output_tensor already has the correct shape [seq_len, ...]
return output_tensor
_deprecated_platform_cp()
def cp_allgather_and_save_kv_cache(forward_batch, layer, k, v, cp_size, swa_loc=None):
"""
Allgather KV cache from all CP ranks and write the full result
into each rank's local memory pool.
swa_loc is the pre-translated full->SWA write target for hybrid SWA pools.
"""
cache_loc = (
forward_batch.out_cache_loc
if not layer.is_cross_attention
else forward_batch.encoder_out_cache_loc
)
k = k.contiguous()
v = v.contiguous()
key_cache_full = cp_all_gather_rerange_kv_cache(
k, cp_size, forward_batch, torch.cuda.current_stream()
)
value_cache_full = cp_all_gather_rerange_kv_cache(
v, cp_size, forward_batch, torch.cuda.current_stream()
)
get_token_to_kv_pool().set_kv_buffer(
layer,
KVWriteLoc(cache_loc, swa_loc),
key_cache_full,
value_cache_full,
layer.k_scale,
layer.v_scale,
)
def prepare_context_parallel_metadata(
kv_len,
cp_rank,
cp_size,
seqs_len,
extend_seqs_len=None,
device="cuda",
):
from sglang.srt.layers.attention.dsa.utils import (
is_dsa_prefill_cp_round_robin_split,
)
if is_dsa_prefill_cp_round_robin_split():
return ContextParallelMetadata()
"""prepare_input_dp_with_cp_dsa-zigzag index
Example (DP_ATTENT_TP == CP_SIZE == 4, single sequence):
block0 | block1 | block2 | block3 | block4 | block5 | block6 | block7
rank 0: block0, block7
rank 1: block1, block6
rank 2: block2, block5
rank 3: block3, block4
For bs > 1, each sequence is split into cp_segment_num = 2 * cp_size
blocks independently; per-rank layout becomes:
[s0.block_r, s1.block_r, ..., s_{bs-1}.block_r,
s0.block_{2*cp_size-1-r}, ..., s_{bs-1}.block_{2*cp_size-1-r}]
i.e. all prev blocks first, then all next blocks -- so torch.split at
total_q_prev_tokens cleanly separates them.
"""
assert extend_seqs_len is not None
extend_seqs_len = [int(x) for x in extend_seqs_len]
# Update the extend_seqs_len to the padded length.
pad_len = int(kv_len) - sum(extend_seqs_len)
if pad_len > 0:
extend_seqs_len[-1] += pad_len
if seqs_len is not None and len(seqs_len) == len(extend_seqs_len):
seqs_len = list(seqs_len)
seqs_len[-1] += pad_len
bs = len(extend_seqs_len)
cp_segment_num = cp_size * 2
# Prefix offset (radix cache hit length) per sequence. For non-NSA
# (FlashAttention) the prefix is baked into kv_len_prev/next via
# prefix_offsets[s] below, so cache_seqlens correctly covers the cached
# prefix. NSA leaves bare cumulatives so its indexer can re-add the
# offset itself.
if seqs_len is not None and len(seqs_len) == bs:
prefix_offsets = [
max(int(seqs_len[s]) - extend_seqs_len[s], 0) for s in range(bs)
]
else:
prefix_offsets = [0] * bs
# Per-sequence block sizes: first (L % cp_segment_num) blocks get +1.
per_seq_block_sizes: List[List[int]] = []
split_list: List[int] = []
for s in range(bs):
L = extend_seqs_len[s]
base = L // cp_segment_num
rem = L % cp_segment_num
blk = [base + 1 if i < rem else base for i in range(cp_segment_num)]
per_seq_block_sizes.append(blk)
split_list.extend(blk)
# Per-rank aggregate: this rank owns block r and block (2*cp_size-1-r)
# of every sequence.
per_rank_actual_token = [0] * cp_size
for r in range(cp_size):
total = 0
for s in range(bs):
total += (
per_seq_block_sizes[s][r]
+ per_seq_block_sizes[s][cp_segment_num - 1 - r]
)
per_rank_actual_token[r] = total
max_single_rank = max(per_rank_actual_token) if per_rank_actual_token else 0
# Kept as cp_size copies so downstream torch.split(x, max_rank_len) still
# works directly. All entries intentionally identical.
max_rank_len = [max_single_rank] * cp_size
# Zigzag index selecting which of split_list's bs * cp_segment_num pieces
# this rank owns, in the order [all_prevs, all_nexts].
zigzag_index = list(
range(cp_rank, cp_rank + bs * cp_segment_num, cp_segment_num)
) + list(
range(
cp_segment_num - cp_rank - 1,
bs * cp_segment_num,
cp_segment_num,
)
)
# Reverse index: given the post-allgather concatenation
# [rank0_prevs_all_seqs, rank0_nexts_all_seqs,
# rank1_prevs_all_seqs, rank1_nexts_all_seqs, ...]
# produce a permutation that restores [s0_b0..s0_bN, s1_b0..s1_bN, ...].
cp_reverse_index: List[int] = []
for batch_id in range(bs):
cp_reverse_index.extend(
list(range(batch_id, cp_segment_num * bs, 2 * bs))
+ list(
range(
(cp_segment_num - 1) * bs + batch_id,
0,
-2 * bs,
)
)
)
# Split sizes matching the post-allgather concatenation order above.
reverse_split_len: List[int] = []
for r in range(cp_size):
for s in range(bs):
reverse_split_len.append(per_seq_block_sizes[s][r])
for s in range(bs):
reverse_split_len.append(per_seq_block_sizes[s][cp_segment_num - 1 - r])
# Per-sequence cumulatives used for FA cache_seqlens.
# kv_len_prev[s] = sum of seq s's blocks [0..cp_rank] (inclusive).
# kv_len_next[s] = sum of seq s's blocks [0..cp_segment_num-cp_rank-1] (inclusive).
from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
nsa_mode = is_dsa_enable_prefill_cp()
kv_len_prev_list: List[int] = []
kv_len_next_list: List[int] = []
actual_seq_q_prev_list: List[int] = []
actual_seq_q_next_list: List[int] = []
for s in range(bs):
blk = per_seq_block_sizes[s]
cum_prev = sum(blk[: cp_rank + 1])
cum_next = sum(blk[: cp_segment_num - cp_rank])
# NSA indexer re-adds prefix offset itself; leave bare cumulative.
# For non-NSA (FlashAttention), bake prefix into cache_seqlens.
if nsa_mode:
kv_len_prev_list.append(cum_prev)
kv_len_next_list.append(cum_next)
else:
kv_len_prev_list.append(prefix_offsets[s] + cum_prev)
kv_len_next_list.append(prefix_offsets[s] + cum_next)
actual_seq_q_prev_list.append(blk[cp_rank])
actual_seq_q_next_list.append(blk[cp_segment_num - cp_rank - 1])
# FlashAttention CUDA tensors (device parameterized for unit tests).
kv_len_prev_tensor = torch.tensor(
kv_len_prev_list, device=device, dtype=torch.int32
)
kv_len_next_tensor = torch.tensor(
kv_len_next_list, device=device, dtype=torch.int32
)
actual_seq_q_prev_tensor = torch.tensor(
actual_seq_q_prev_list, device=device, dtype=torch.int32
)
actual_seq_q_next_tensor = torch.tensor(
actual_seq_q_next_list, device=device, dtype=torch.int32
)
cu_prev = [0] + list(accumulate(actual_seq_q_prev_list))
cu_next = [0] + list(accumulate(actual_seq_q_next_list))
cu_seqlens_q_prev_tensor = torch.tensor(cu_prev, device=device, dtype=torch.int32)
cu_seqlens_q_next_tensor = torch.tensor(cu_next, device=device, dtype=torch.int32)
total_q_prev_tokens = cu_prev[-1]
total_q_next_tokens = cu_next[-1]
max_seqlen_q_prev = max(actual_seq_q_prev_list) if actual_seq_q_prev_list else 0
max_seqlen_q_next = max(actual_seq_q_next_list) if actual_seq_q_next_list else 0
total_seq_lens = sum(extend_seqs_len)
# Cheap invariants: metadata must be a valid permutation spec.
# - split_list has bs * cp_segment_num pieces (all blocks, all seqs).
# - zigzag_index has 2 * bs entries (this rank's prev + next per seq).
# - cp_reverse_index has bs * cp_segment_num entries (reorders the
# full allgathered stream back to per-seq-original order).
assert len(split_list) == bs * cp_segment_num
assert sum(split_list) == total_seq_lens
assert len(zigzag_index) == 2 * bs
assert len(cp_reverse_index) == bs * cp_segment_num
assert sorted(cp_reverse_index) == list(range(bs * cp_segment_num))
assert sum(per_rank_actual_token) == total_seq_lens
return ContextParallelMetadata(
split_list=split_list,
zigzag_index=zigzag_index,
cp_reverse_index=cp_reverse_index,
reverse_split_len=reverse_split_len,
per_rank_actual_token=per_rank_actual_token,
max_rank_len=max_rank_len,
kv_len_prev_tensor=kv_len_prev_tensor,
kv_len_next_tensor=kv_len_next_tensor,
actual_seq_q_prev_tensor=actual_seq_q_prev_tensor,
actual_seq_q_next_tensor=actual_seq_q_next_tensor,
cu_seqlens_q_prev_tensor=cu_seqlens_q_prev_tensor,
cu_seqlens_q_next_tensor=cu_seqlens_q_next_tensor,
total_q_prev_tokens=total_q_prev_tokens,
total_q_next_tokens=total_q_next_tokens,
max_seqlen_q_prev=max_seqlen_q_prev,
max_seqlen_q_next=max_seqlen_q_next,
kv_len_prev_list=kv_len_prev_list,
kv_len_next_list=kv_len_next_list,
actual_seq_q_prev_list=actual_seq_q_prev_list,
actual_seq_q_next_list=actual_seq_q_next_list,
total_seq_lens=total_seq_lens,
bs=bs,
)
_deprecated_platform_cp()
python/sglang/srt/managers/scheduler_pp_mixin.py
+3 -13
View File
@@ -553,10 +553,6 @@ class SchedulerPPMixin:
def init_pp_loop_state(self: Scheduler):
self.pp_loop_size: int = self.ps.pp_size + get_parallel().pp_async_batch_depth
# In CP mode, attention weights are duplicated, eliminating the need for the attention TP all-gather operation.
self.require_attn_tp_allgather = (
not get_parallel().enable_dsa_prefill_context_parallel
)
self.mbs = [None] * self.pp_loop_size
self.last_mbs = [None] * self.pp_loop_size
self.running_mbs = [
@@ -817,9 +813,7 @@ class SchedulerPPMixin:
p2p_work.extend(
self.pp_group.send_tensor_dict(
tensor_dict=tensor_dict,
all_gather_group=(
self.attn_tp_group if self.require_attn_tp_allgather else None
),
all_gather_group=(self.attn_tp_group),
async_send=async_send,
)
)
@@ -864,9 +858,7 @@ class SchedulerPPMixin:
pp_proxy_tensors = PPProxyTensors(
self._pp_recv_typed_dict(
expected_kind="proxy",
all_gather_group=(
self.attn_tp_group if self.require_attn_tp_allgather else None
),
all_gather_group=(self.attn_tp_group),
)
)
return pp_proxy_tensors
@@ -876,9 +868,7 @@ class SchedulerPPMixin:
) -> Dict[str, torch.Tensor]:
return self._pp_recv_typed_dict(
expected_kind="output",
all_gather_group=(
self.attn_tp_group if self.require_attn_tp_allgather else None
),
all_gather_group=(self.attn_tp_group),
)
def _pp_make_skip_output_result(
python/sglang/srt/model_executor/forward_batch_info.py
-14
View File
@@ -1331,10 +1331,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
def prepare_mlp_sync_batch(self, model_runner: ModelRunner):
from sglang.srt.batch_overlap.two_batch_overlap import TboForwardBatchPreparer
# Local imports: module-level CP helper imports here are circular (#27014).
from sglang.srt.layers.cp.padding import get_cp_padding_align_size
from sglang.srt.layers.cp.utils import enable_cp_v2
assert self.global_num_tokens_cpu is not None
assert self.global_num_tokens_for_logprob_cpu is not None
@@ -1348,16 +1344,6 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
# there is no reduce-scatter in LM logprob, so we do not need to adjust the padded length for logprob
global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_tp_size)
# make sure that each rank has the same number of tokens to do collective communication.
# Zigzag (in-seq-split) CP pads to 2 * attn_cp_size for load balance; other CP modes
# pad to attn_cp_size; CP off pads nothing (extra padding breaks EAGLE/MTP draft
# prefill with NaN draft logits, see #23269).
# FIXME(kpham-sgl): revisit so draft prefill-extend tolerates padded dummy tokens.
if not enable_cp_v2():
cp_align_size = get_cp_padding_align_size()
for i in range(sync_group_size):
global_num_tokens[i] = ceil_align(global_num_tokens[i], cp_align_size)
dp_padding_mode = DpPaddingMode.get_dp_padding_mode(
self.is_extend_in_batch, global_num_tokens
)
python/sglang/srt/models/deepseek_common/attention_backend_handler.py
+2 -7
View File
@@ -1,7 +1,6 @@
from sglang.srt.environ import envs
from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
from sglang.srt.layers.cp.utils import enable_cp_v2, is_cp_v2_active
from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp
from sglang.srt.layers.cp.utils import is_cp_v2_active
from sglang.srt.model_executor.forward_context import get_attn_backend
from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
is_in_breakable_cuda_graph,
@@ -115,10 +114,7 @@ def _handle_attention_backend(attn, forward_batch, backend_name):
# Strategy CP gathers latent KV in the backend's absorbed MLA path;
# normal MHA would write rank-local KV against full cache locations.
# Protected platform CP retains its model-side materialization path.
if is_cp_v2_active(forward_batch) or (
not enable_cp_v2() and mla_use_prefill_cp(forward_batch)
):
if is_cp_v2_active(forward_batch):
return _dispatch_mla_subtype(attn, forward_batch)
sum_extend_prefix_lens = _get_sum_extend_prefix_lens(forward_batch)
@@ -236,7 +232,6 @@ def _can_use_triton_dense_fp8_prefill(attn, forward_batch) -> bool:
and attn.v_head_dim == 128
and attn.kv_lora_rank == 512
and not get_parallel().dcp_enabled
and not mla_use_prefill_cp(forward_batch)
and forward_batch.forward_mode.is_extend_without_speculative()
and prefix_lens is not None
and any(prefix_lens)
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py
-7
View File
@@ -17,7 +17,6 @@ from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.attention.dsa.utils import is_graph_dsa_split_op_surface
from sglang.srt.layers.attention.dsa_backend import prepare_kv_for_attention
from sglang.srt.layers.communicator import get_attn_tp_context
from sglang.srt.layers.cp.utils import enable_cp_v2
from sglang.srt.layers.dcp import (
all_gather_kv_cache_for_mla_extend,
all_gather_q_for_mla_decode,
@@ -26,7 +25,6 @@ from sglang.srt.layers.dcp import (
)
from sglang.srt.layers.logits_processor import get_in_autotune_dummy_run
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.lora.deepseek_mla_correction import (
apply_q_correction as apply_kv_b_lora_q_correction,
)
@@ -626,11 +624,6 @@ class DeepseekMLAForwardMixin:
defer_materialization=fuse_rope_for_trtllm_mla,
)
if not enable_cp_v2() and mla_use_prefill_cp(forward_batch):
k_nope, k_pe = self.rebuild_cp_kv_cache(
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 get_parallel().dcp_enabled:
if is_dcp_mla_decode_phase(forward_batch):
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py
-30
View File
@@ -20,9 +20,7 @@ from sglang.kernels.ops.quantization.fp8_kernel import (
)
from sglang.srt.environ import envs
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
from sglang.srt.layers.communicator import get_attn_tp_context
from sglang.srt.layers.cp.utils import is_cp_v2_active
from sglang.srt.layers.dcp import (
all_gather_kv_cache_for_mla_extend,
all_gather_q_for_mla_decode,
@@ -35,7 +33,6 @@ from sglang.srt.layers.quantization.fp8_utils import (
materialize_bpreshuffle_fp8_scale_tuple,
view_aiter_fused_rms_transposed_fp8_scale_tuple,
)
from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp
from sglang.srt.lora.deepseek_mla_correction import (
apply_q_correction as apply_kv_b_lora_q_correction,
)
@@ -54,7 +51,6 @@ from sglang.srt.models.deepseek_common.attention_forward_methods.forward_mla imp
_select_local_dcp_heads_for_autotune,
is_dcp_mla_decode_phase,
is_mla_dcp_lse_base_on_e,
should_defer_dsa_cp_kv_gather,
)
from sglang.srt.models.deepseek_common.utils import (
FORWARD_ABSORB_CORE_ATTENTION_BACKENDS,
@@ -617,32 +613,6 @@ class DeepseekMLARocmForwardMixin:
):
q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
dsa_prefill_cp = dsa_use_prefill_cp(forward_batch)
mla_prefill_cp = mla_use_prefill_cp(forward_batch)
defer_kv_gather_until_after_rope = should_defer_dsa_cp_kv_gather(
dsa_prefill_cp=dsa_prefill_cp,
fuse_rope_for_trtllm_mla=fuse_rope_for_trtllm_mla,
)
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,
)
# all_gather q_pe, q_nope_out,take tp8 as an example， q_pe [B, H, ROPE_DIM], q_nope_out [B, H, NOPE_DIM] gathered to [B, H * dcp_world_size, ROPE_DIM] [B, H * dcp_world_size, NOPE_DIM] for decode batch, and all gather k_pe, k_nope for extend batch.
if get_parallel().dcp_enabled:
if is_dcp_mla_decode_phase(forward_batch):
python/sglang/srt/models/deepseek_nextn.py
-95
View File
@@ -28,28 +28,12 @@ from sglang.kernels.ops.layernorm.fused_eh_norm import fused_eh_norm
from sglang.srt.distributed import get_pp_group
from sglang.srt.environ import envs
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.layers.attention.dsa.utils import (
can_dsa_cp_split,
dsa_use_prefill_cp,
is_dsa_enable_prefill_cp,
is_dsa_prefill_cp_round_robin_split,
)
from sglang.srt.layers.attention.index_topk_share import IndexTopKShareState
from sglang.srt.layers.cp.utils import enable_cp_v2
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.quantization import Fp8Config
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.utils.cp_utils import (
can_cp_split,
cp_all_gather_rerange_output,
cp_split_and_rebuild_data,
cp_split_and_rebuild_position,
is_mla_prefill_cp_enabled,
mla_use_prefill_cp,
prepare_context_parallel_metadata,
)
from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
@@ -62,34 +46,6 @@ from sglang.srt.models.utils import WeightsMapper
from sglang.srt.runtime_context import get_model, get_parallel, get_spec
from sglang.srt.utils import BumpAllocator, add_prefix, is_cuda, is_npu
def _gather_dsa_topk_indices_for_cp(
topk_indices: torch.Tensor,
local_num_tokens: int,
cp_size: int,
forward_batch: ForwardBatch,
stream,
) -> torch.Tensor:
if (
is_dsa_prefill_cp_round_robin_split()
and topk_indices.shape[0] < local_num_tokens
):
pad_rows = local_num_tokens - topk_indices.shape[0]
topk_indices = torch.cat(
[
topk_indices,
topk_indices.new_full((pad_rows, topk_indices.shape[1]), -1),
],
dim=0,
)
return cp_all_gather_rerange_output(
topk_indices,
cp_size,
forward_batch,
stream,
)
logger = logging.getLogger(__name__)
@@ -260,13 +216,6 @@ class DeepseekModelNextN(nn.Module):
else:
hidden_states = self.eh_proj(eh_input)
# Protected platforms retain their model-side token split.
use_platform_cp = not enable_cp_v2() and (
dsa_use_prefill_cp(forward_batch) or mla_use_prefill_cp(forward_batch)
)
if use_platform_cp:
hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
positions = cp_split_and_rebuild_position(forward_batch, positions)
residual = None
index_topk_share = IndexTopKShareState.from_mtp_carry(forward_batch)
with get_global_expert_distribution_recorder().disable_this_region():
@@ -284,22 +233,6 @@ class DeepseekModelNextN(nn.Module):
else:
hidden_states = self.shared_head.norm(hidden_states)
if use_platform_cp:
local_num_tokens = hidden_states.shape[0]
hidden_states = cp_all_gather_rerange_output(
hidden_states,
get_parallel().attn_cp_size,
forward_batch,
torch.cuda.current_stream(),
)
if index_topk_share.should_publish and topk_indices is not None:
topk_indices = _gather_dsa_topk_indices_for_cp(
topk_indices,
local_num_tokens,
get_parallel().attn_cp_size,
forward_batch,
torch.cuda.current_stream(),
)
index_topk_share.update(topk_indices)
index_topk_share.publish()
finally:
@@ -371,34 +304,6 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
positions: torch.Tensor,
forward_batch: ForwardBatch,
) -> torch.Tensor:
if not enable_cp_v2():
if is_dsa_enable_prefill_cp():
if can_dsa_cp_split(
len(input_ids),
get_parallel().attn_cp_size,
self.model.decoder.self_attn.use_dsa,
forward_batch,
):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len(input_ids),
get_parallel().attn_cp_rank,
get_parallel().attn_cp_size,
forward_batch.seq_lens_cpu.tolist(),
extend_seqs_len=forward_batch.extend_seq_lens_cpu,
)
elif (
is_mla_prefill_cp_enabled() and not self.model.decoder.self_attn.use_dsa
):
if can_cp_split(
len(input_ids), get_parallel().attn_cp_size, forward_batch
):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len(input_ids),
get_parallel().attn_cp_rank,
get_parallel().attn_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
-90
View File
@@ -64,11 +64,6 @@ from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.amx_utils import PackWeightMethod
from sglang.srt.layers.attention.dsa.dsa_indexer import Indexer
from sglang.srt.layers.attention.dsa.dsa_indexer_kpool import IndexerKPool
from sglang.srt.layers.attention.dsa.utils import (
can_dsa_cp_split,
dsa_use_prefill_cp,
is_dsa_enable_prefill_cp,
)
from sglang.srt.layers.attention.index_topk_share import IndexTopKShareState
from sglang.srt.layers.aux_hidden_states import (
AuxHiddenStateAccumulator,
@@ -85,7 +80,6 @@ from sglang.srt.layers.communicator_dsa_cp import (
maybe_prefetch_next_full_attention_kv,
)
from sglang.srt.layers.cp.cp_decode_attn_tp import get_cp_decode_attn_tp_ctx
from sglang.srt.layers.cp.utils import enable_cp_v2
from sglang.srt.layers.dcp.planner import (
prepare_decode_context_parallel_metadata,
)
@@ -135,15 +129,6 @@ from sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe import (
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.rotary_embedding import get_rope_wrapper
from sglang.srt.layers.utils import PPMissingLayer
from sglang.srt.layers.utils.cp_utils import (
can_cp_split,
cp_all_gather_rerange_output,
cp_split_and_rebuild_data,
cp_split_and_rebuild_position,
is_prefill_context_parallel_enabled,
mla_use_prefill_cp,
prepare_context_parallel_metadata,
)
from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
@@ -531,14 +516,6 @@ class MoEGate(nn.Module):
if get_exec().deterministic.enable_deterministic_inference:
return F.linear(hidden_states, self.weight, None)
if (
not enable_cp_v2()
and not self.is_deepseek_v4
and forward_batch is not None
and (dsa_use_prefill_cp(forward_batch) or mla_use_prefill_cp(forward_batch))
):
return F.linear(hidden_states, self.weight, None)
if hidden_states.shape[0] <= self.tiny_router_gemm_max_tokens:
logits = tiny_gemm_bf16(
hidden_states,
@@ -2266,20 +2243,6 @@ class DeepseekV2AttentionMLA(
q = self.q_b_proj(q_lora)[0]
return q.view(-1, self.num_local_heads, self.qk_head_dim)
def rebuild_cp_kv_cache(self, latent_cache, forward_batch, k_nope, k_pe):
# Retained for the platform MLA paths.
latent_cache[..., : self.kv_lora_rank] = k_nope.squeeze(1)
latent_cache[..., self.kv_lora_rank :] = k_pe.squeeze(1)
latent_cache_output = cp_all_gather_rerange_output(
latent_cache.contiguous(),
get_parallel().attn_cp_size,
forward_batch,
torch.cuda.current_stream(),
)
k_nope = latent_cache_output[..., : self.kv_lora_rank].unsqueeze(1)
k_pe = latent_cache_output[..., self.kv_lora_rank :].unsqueeze(1)
return k_nope, k_pe
@staticmethod
def _get_q_b_proj_quant_config(quant_config):
if envs.SGLANG_NVFP4_CKPT_FP8_GEMM_IN_ATTN.get():
@@ -2809,15 +2772,6 @@ class DeepseekV2Model(nn.Module):
else None
)
# HIP/NPU/MUSA retain their model-side CP boundary.
use_platform_cp = not enable_cp_v2() and (
dsa_use_prefill_cp(forward_batch) or mla_use_prefill_cp(forward_batch)
)
if use_platform_cp:
if self.pp_group.is_first_rank:
hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
positions = cp_split_and_rebuild_position(forward_batch, positions)
# llama_4_scaling: for supporting Mistral-Large-3 model
# Compute llama 4 scaling once per forward pass if enabled
llama_4_scaling: Optional[torch.Tensor] = None
@@ -2917,14 +2871,6 @@ class DeepseekV2Model(nn.Module):
else:
hidden_states, _ = self.norm(hidden_states, residual)
if self.pp_group.is_last_rank and use_platform_cp:
# allgather + rerrange
hidden_states = cp_all_gather_rerange_output(
hidden_states,
get_parallel().attn_cp_size,
forward_batch,
torch.cuda.current_stream(),
)
if len(aux_hidden_states) == 0:
return hidden_states
return hidden_states, aux_hidden_states.finalize()
@@ -3085,42 +3031,6 @@ 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:
len_input_ids = input_ids.shape[0]
elif input_embeds is not None:
len_input_ids = input_embeds.shape[0]
else:
len_input_ids = pp_proxy_tensors["hidden_states"].shape[0]
if not enable_cp_v2():
if is_dsa_enable_prefill_cp():
if can_dsa_cp_split(
len_input_ids,
get_parallel().attn_cp_size,
self.use_dsa,
forward_batch,
):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len_input_ids,
get_parallel().attn_cp_rank,
get_parallel().attn_cp_size,
forward_batch.seq_lens_cpu.tolist(),
extend_seqs_len=forward_batch.extend_seq_lens_cpu,
)
elif is_prefill_context_parallel_enabled() and not self.use_dsa:
if can_cp_split(
len_input_ids, get_parallel().attn_cp_size, forward_batch
):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len_input_ids,
get_parallel().attn_cp_rank,
get_parallel().attn_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(
input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors
python/sglang/srt/models/deepseek_v4.py
+11 -300
View File
@@ -55,10 +55,8 @@ from sglang.srt.hardware_backend.npu.utils import (
use_npu_arch35_mxfp8_wo_a,
)
from sglang.srt.layers.attention.dsa.utils import (
can_dsa_cp_split,
dsa_use_prefill_cp,
is_dsa_enable_prefill_cp,
is_dsa_prefill_cp_round_robin_split,
)
from sglang.srt.layers.attention.dsv4.compressor import Compressor
from sglang.srt.layers.attention.dsv4.indexer import C4Indexer
@@ -70,13 +68,9 @@ from sglang.srt.layers.communicator_dsa_cp import (
from sglang.srt.layers.cp.cp_decode_attn_tp import get_cp_decode_attn_tp_ctx
from sglang.srt.layers.cp.utils import (
cp_materialize_global_token_order,
enable_cp_v2,
is_cp_v2_active,
)
from sglang.srt.layers.dp_attention import (
_tbo_event,
attn_cp_overlap_all_gather_into_tensor,
attn_cp_overlap_reduce_scatter_tensor,
attn_tp_all_gather,
attn_tp_all_reduce,
dp_gather_partial,
@@ -109,15 +103,6 @@ from sglang.srt.layers.quantization.fp8_utils import (
)
from sglang.srt.layers.rotary_embedding import get_rope_wrapper
from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
from sglang.srt.layers.utils.cp_utils import (
cp_all_gather_rerange_finish,
cp_all_gather_rerange_launch,
cp_all_gather_rerange_output,
cp_round_robin_input_ids,
cp_split_and_rebuild_data,
cp_split_and_rebuild_position,
prepare_context_parallel_metadata,
)
from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
from sglang.srt.mem_cache.memory_pool import RadixAttention
from sglang.srt.model_executor.cuda_graph_config import (
@@ -172,12 +157,6 @@ from sglang.srt.runtime_context import (
get_parallel,
get_platform,
)
if not _is_hip:
from sglang.srt.layers.utils.cp_utils import (
prepare_context_parallel_metadata,
)
from sglang.srt.utils import (
LazyValue,
add_prefix,
@@ -666,9 +645,6 @@ class MqaAttentionBase(nn.Module):
if attn_tp_rank is None or attn_tp_size is None:
attn_tp_rank = get_parallel().attn_tp_rank
attn_tp_size = get_parallel().attn_tp_size
if self.dsa_enable_prefill_cp:
self.cp_size = get_parallel().attn_cp_size
attn_tp_rank, attn_tp_size = 0, 1
self.attn_tp_rank: int = attn_tp_rank
self.attn_tp_size: int = attn_tp_size
@@ -1415,14 +1391,6 @@ class MQALayer(MqaAttentionBase):
x_quant=None,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
x_linear = x_quant if x_quant is not None else x
# kv_score depends only on x, so its CP all-gather can start before the
# projections and be collected inside forward_core_compressor below --
# the projections are what hides it. No-op unless the CP+TBO path armed
# _cp_prefetch_comm_stream.
if _is_hip and self.compressor is not None:
self.compressor.prelaunch_kv_score(x, forward_batch)
if self.indexer is not None:
self.indexer.compressor.prelaunch_kv_score(x, forward_batch)
if self.fuse_wqa_wkv:
qkv_a, _ = self.wqkv_a(x_linear)
@@ -1433,7 +1401,6 @@ class MQALayer(MqaAttentionBase):
use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch)
kv: Optional[torch.Tensor]
kv_handle = None
from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import (
is_unified_kv_triton,
@@ -1582,29 +1549,6 @@ class MQALayer(MqaAttentionBase):
# unified_kv prefill: keep bf16 kv; the backend writes
# the ring AFTER attention (2-source path).
kv = self._compute_kv_bf16(x_linear, positions, qkv_a=qkv_a)
# HIP/ROCm-only: the unified_kv 2-source prefill path is exclusive
# to DeepseekV4HipRadixBackend. Guard with _is_hip so this CP
# all-gather never enters the NVIDIA (DeepseekV4AttnBackend) path.
if use_cp and _is_hip:
# unified_kv + DSA CP: the 2-source prefill path needs the
# FULL current-chunk KV (extend source + ring write), so
# all-gather the per-rank bf16 KV across the CP group.
comm_stream = getattr(
forward_batch, "_cp_prefetch_comm_stream", None
)
if comm_stream is not None:
# kv is not read again until this function returns, so the
# indexer + compressor below can run while it gathers.
kv_handle = cp_all_gather_rerange_launch(
kv, self.cp_size, comm_stream, ("kv", self.layer_id)
)
kv = None
else:
kv = cp_materialize_global_token_order(
kv.contiguous(),
forward_batch,
torch.cuda.current_stream(),
)
elif use_cp:
# NSA CP: keep bf16 kv around for the cross-rank all-gather, then
# write to the FlashMLA cache after gather.
@@ -1642,9 +1586,6 @@ class MQALayer(MqaAttentionBase):
self.compressor,
)
if _is_hip and kv_handle is not None:
kv = cp_all_gather_rerange_finish(kv_handle)
return q, kv
def forward(
@@ -2856,67 +2797,6 @@ class DeepseekV4DecoderLayer(nn.Module):
hidden = hidden + shared_local[:n]
state.hidden_states_mlp_output = hidden
def _cp_tbo_launch(self, state, x, key, out_rows, collective):
assert _is_hip, "CP+TBO MoE overlap is HIP-only"
x = x.contiguous()
sub = state.tbo_subbatch_index
out = get_tbo_persistent_buffer(
(key, sub), out_rows, x.shape[1], x.dtype, x.device
)
comm = get_dp_tbo_comm_stream()
comm.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(comm):
collective(out, x)
event = _tbo_event((key, sub))
event.record(comm)
return out, event, x
def op_cp_gather_a(self, state):
local = state.pop("hidden_states_mlp_input")
out, event, keepalive = self._cp_tbo_launch(
state,
local,
"cpgh",
local.shape[0] * get_parallel().attn_cp_size,
attn_cp_overlap_all_gather_into_tensor,
)
state.global_hidden = out
state.cp_gather_event = event
state.cp_gather_keepalive = keepalive
def op_cp_gather_b(self, state):
torch.cuda.current_stream().wait_event(state.pop("cp_gather_event"))
state.pop("cp_gather_keepalive")
def op_cp_moe(self, state):
fb = state.forward_batch
global_ids = fb._cp_moe_input_ids
with get_forward().scoped(mlp_reduce_scatter=True):
state.global_expert_out = self.mlp(
state.pop("global_hidden"),
fb,
input_ids=global_ids,
input_ids_global=global_ids,
)
def op_cp_combine_a(self, state):
global_out = state.pop("global_expert_out")
out, event, keepalive = self._cp_tbo_launch(
state,
global_out,
"cplo",
global_out.shape[0] // get_parallel().attn_cp_size,
attn_cp_overlap_reduce_scatter_tensor,
)
state.local_out = out
state.cp_combine_event = event
state.cp_combine_keepalive = keepalive
def op_cp_combine_b(self, state):
torch.cuda.current_stream().wait_event(state.pop("cp_combine_event"))
state.pop("cp_combine_keepalive")
state.hidden_states_mlp_output = state.pop("local_out")
class DeepseekV4Model(nn.Module):
fall_back_to_pt_during_load = False
@@ -2992,12 +2872,9 @@ class DeepseekV4Model(nn.Module):
self.hc_head_scale,
) = make_hc_head_params(hc_mult, config.hidden_size)
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
self.use_fused_mhc_post_pre = (
is_cross_layer_mhc_fusion_enabled() or _is_fused_mhc_post_pre_enabled_xpu()
)
if self.dsa_enable_prefill_cp:
self.cp_size = get_parallel().attn_cp_size
self.dspark_layers_to_capture: Optional[List[int]] = None
@@ -3040,44 +2917,22 @@ class DeepseekV4Model(nn.Module):
hc_eps=self.hc_eps,
)
def _cp_children_splittable(self, forward_batch: ForwardBatch) -> bool:
children = forward_batch.tbo_children
if not children:
return False
cp_size = get_parallel().attn_cp_size
for child in children:
if child.batch_size <= 0 or child.extend_seq_lens_cpu is None:
return False
if sum(child.extend_seq_lens_cpu) < cp_size:
return False
return True
def _can_run_tbo(self, forward_batch: ForwardBatch) -> bool:
"""DSV4 prefill-only two-batch-overlap gate.
TBO batch prep (tbo_split_seq_index / tbo_children) is populated
model-agnostically when --enable-two-batch-overlap is set and the
DP-attention preparer allows it (mori `normal` mode permits prefill
TBO). We additionally restrict to: prefill (EXTEND), single PP, and a
path the DSV4 op strategy implements -- the non-CP path everywhere, plus
the round-robin DSA prefill CP path on HIP.
TBO). We additionally restrict to prefill (EXTEND), single PP, and
non-CP paths supported by the DSV4 op strategy.
"""
from sglang.srt.layers.moe import is_tbo_enabled
if dsa_use_prefill_cp(forward_batch):
path_ok = (
_is_hip
and not is_cp_v2_active(forward_batch)
and is_dsa_prefill_cp_round_robin_split()
and get_moe_a2a_backend().is_none()
and self._cp_children_splittable(forward_batch)
)
else:
path_ok = (
not _is_hip
or not get_moe_a2a_backend().is_none()
or get_parallel().attn_dp_size > 1
)
path_ok = not dsa_use_prefill_cp(forward_batch) and (
not _is_hip
or not get_moe_a2a_backend().is_none()
or get_parallel().attn_dp_size > 1
)
return (
is_tbo_enabled()
and forward_batch.can_run_tbo
@@ -3103,13 +2958,6 @@ class DeepseekV4Model(nn.Module):
_model_forward_tbo_merge_outputs,
)
if _is_hip and dsa_use_prefill_cp(forward_batch):
return self._forward_layers_tbo_cp(
positions=positions,
hidden_states=hidden_states,
forward_batch=forward_batch,
)
layers = [self.layers[i] for i in range(self.start_layer, self.end_layer)]
operations_strategy = OperationsStrategy.init_new_tbo(
layers, forward_batch.global_forward_mode
@@ -3186,97 +3034,6 @@ class DeepseekV4Model(nn.Module):
)
return hidden_states
def _setup_child_cp_metadata(self, child: ForwardBatch, child_backend) -> None:
cp_rank = get_parallel().attn_cp_rank
cp_size = get_parallel().attn_cp_size
child.attn_cp_metadata = prepare_context_parallel_metadata(
len(child.input_ids),
cp_rank,
cp_size,
child.seq_lens_cpu.tolist(),
extend_seqs_len=child.extend_seq_lens_cpu,
)
if is_dsa_prefill_cp_round_robin_split():
metadata = child_backend.forward_metadata
core_meta = metadata.core_attn_metadata
core_meta.apply_cp_reindex()
core_meta.init_flashmla_related(is_prefill=True)
if metadata.indexer_metadata is not None:
metadata.indexer_metadata = child_backend.init_forward_metadata_indexer(
core_meta
)
def _forward_layers_tbo_cp(
self,
positions: torch.Tensor,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
) -> torch.Tensor:
assert _is_hip, "CP+TBO prefill path is HIP-only"
from sglang.srt.batch_overlap.operations import execute_overlapped_operations
from sglang.srt.batch_overlap.operations_strategy import OperationsStrategy
from sglang.srt.batch_overlap.two_batch_overlap import (
_model_forward_filter_inputs,
_model_forward_tbo_merge_outputs,
)
original_len = hidden_states.shape[0]
cp_size = get_parallel().attn_cp_size
layers = [self.layers[i] for i in range(self.start_layer, self.end_layer)]
operations_strategy = OperationsStrategy.init_new_tbo(
layers, forward_batch.global_forward_mode, use_cp=True
)
attn_backend = get_attn_backend()
children = forward_batch.tbo_children
# Attention-side CP gathers run two-phase (launch early on the comm
# stream / collect right before their consumer). Only the MoE
# collectives are splittable across a YieldOperation, so without this the
# ~2.5 attention-side collectives per layer would stay on the compute
# stream and defeat most of TBO's overlap.
prefetch_comm_stream = get_dp_tbo_comm_stream()
inputs_arr = []
for idx, child in enumerate(children):
child_inputs = _model_forward_filter_inputs(
hidden_states=hidden_states,
residual=None,
positions=positions,
output_forward_batch=child,
tbo_subbatch_index=idx,
)
self._setup_child_cp_metadata(child, attn_backend.children[idx])
if self.pp_group.is_first_rank:
child_inputs["hidden_states"] = cp_split_and_rebuild_data(
child, child_inputs["hidden_states"]
)
child_inputs["positions"] = cp_split_and_rebuild_position(
child, child_inputs["positions"]
)
child._cp_moe_input_ids = cp_round_robin_input_ids(child.input_ids)
child._cp_prefetch_comm_stream = prefetch_comm_stream
inputs_arr.append(child_inputs)
outputs_arr = execute_overlapped_operations(
inputs_arr=inputs_arr,
operations_arr=[operations_strategy.operations] * 2,
delta_stages=[0, operations_strategy.tbo_delta_stages],
)
if self.pp_group.is_last_rank:
for idx, child in enumerate(children):
outputs_arr[idx]["hidden_states"] = cp_all_gather_rerange_output(
outputs_arr[idx]["hidden_states"],
cp_size,
child,
torch.cuda.current_stream(),
)
hidden_states, _ = _model_forward_tbo_merge_outputs(
outputs_arr[0], outputs_arr[1], original_len
)
return hidden_states
def forward(
self,
input_ids: torch.Tensor,
@@ -3321,19 +3078,14 @@ class DeepseekV4Model(nn.Module):
# DSpark aux capture needs the per-layer eager loop (TBO's overlapped
# execution cannot expose per-layer completed hidden states), so skip
# TBO when capturing -- a perf-only downgrade, not a correctness one.
run_tbo = self._can_run_tbo(forward_batch) and not capture_dspark
use_platform_cp = not enable_cp_v2() and dsa_use_prefill_cp(forward_batch)
if use_platform_cp and not run_tbo:
if self.pp_group.is_first_rank:
hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
positions = cp_split_and_rebuild_position(forward_batch, positions)
input_ids = cp_round_robin_input_ids(input_ids)
input_ids_global = input_ids
# Reset Compressor's per-step freqs_cis cache from any previous step.
for _attr in ("freqs_cis_c4", "freqs_cis_c128"):
if hasattr(forward_batch, _attr):
delattr(forward_batch, _attr)
run_tbo = self._can_run_tbo(forward_batch) and not capture_dspark
if _is_npu and not run_tbo:
# Rope cos/sin for the whole forward: one bf16 gather per rope
# config on the current stream, before the layer loop forks the
@@ -3347,6 +3099,7 @@ class DeepseekV4Model(nn.Module):
forward_batch,
positions,
)
if run_tbo:
# Two-batch-overlap prefill (EP / mori). Cross-layer mHC fusion is
# disabled here (each layer self-contained), so no trailing hc_post.
@@ -3391,24 +3144,6 @@ class DeepseekV4Model(nn.Module):
hidden_states, prev_residual, prev_post, prev_comb
)
# CP all-gather only on the last PP rank; PP IPC carries CP-split tensors.
if self.pp_group.is_last_rank and use_platform_cp and not run_tbo:
stream = torch.cuda.current_stream()
hidden_states = cp_all_gather_rerange_output(
hidden_states,
self.cp_size,
forward_batch,
stream,
)
# Gather DSpark aux tensors on the same CP token split.
if capture_dspark:
dspark_aux_hidden_states = [
cp_all_gather_rerange_output(
aux, self.cp_size, forward_batch, stream
)
for aux in dspark_aux_hidden_states
]
if not self.pp_group.is_last_rank:
# Flatten 3D mHC tensor for PP IPC.
return PPProxyTensors({"hidden_states": hidden_states.flatten(1)})
@@ -3482,11 +3217,6 @@ class DeepseekV4ForCausalLM(nn.Module):
self.start_layer = self.model.start_layer
self.end_layer = self.model.end_layer
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if self.dsa_enable_prefill_cp:
self.cp_rank = get_parallel().attn_cp_rank
self.cp_size = get_parallel().attn_cp_size
# update_weights_from_disk/_tensor/_distributed re-enter load_weights
# mid-serving (RL refit sends many partial batches); the prewarm and
# its barrier must only run on the first (startup) load.
@@ -3553,25 +3283,6 @@ class DeepseekV4ForCausalLM(nn.Module):
input_embeds: Optional[torch.Tensor] = None,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> torch.Tensor:
if not enable_cp_v2() and self.dsa_enable_prefill_cp:
if can_dsa_cp_split(len(input_ids), self.cp_size, True, 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 is_dsa_prefill_cp_round_robin_split():
attn_backend = get_attn_backend()
metadata = attn_backend.forward_metadata
core_meta = metadata.core_attn_metadata
core_meta.apply_cp_reindex()
core_meta.init_flashmla_related(is_prefill=True)
if metadata.indexer_metadata is not None:
metadata.indexer_metadata = (
attn_backend.init_forward_metadata_indexer(core_meta)
)
with get_attn_tp_context().maybe_input_scattered(forward_batch):
hidden_states = self.model.forward(
python/sglang/srt/models/deepseek_v4_nextn.py
-64
View File
@@ -8,15 +8,6 @@ from transformers import PretrainedConfig
from sglang.srt.distributed import get_pp_group
from sglang.srt.hardware_backend.npu.dsv4.dsv4_rope import prime_rope_cos_sin
from sglang.srt.layers.attention.dsa.utils import (
can_dsa_cp_split,
dsa_use_prefill_cp,
is_dsa_enable_prefill_cp,
is_dsa_prefill_cp_round_robin_split,
)
from sglang.srt.layers.cp.utils import (
enable_cp_v2,
)
from sglang.srt.layers.dp_attention import (
dp_gather_replicate,
get_global_dp_buffer_len,
@@ -28,19 +19,11 @@ from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig
from sglang.srt.layers.utils.cp_utils import (
cp_all_gather_rerange_output,
cp_round_robin_input_ids,
cp_split_and_rebuild_data,
cp_split_and_rebuild_position,
prepare_context_parallel_metadata,
)
from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.forward_context import get_attn_backend
from sglang.srt.models.deepseek_v4 import (
DeepseekV4DecoderLayer,
DeepseekV4ForCausalLM,
@@ -114,12 +97,6 @@ class DeepseekV4ModelNextN(nn.Module):
compress_ratio_override=COMPRESS_RATIO_NEXTN_LAYER,
)
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if self.dsa_enable_prefill_cp:
self.cp_size = get_parallel().attn_cp_size
else:
self.cp_size = None
self.shared_head = nn.Module()
self.shared_head.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
@@ -148,7 +125,6 @@ class DeepseekV4ModelNextN(nn.Module):
forward_batch: ForwardBatch,
input_embeds: torch.Tensor = None,
) -> torch.Tensor:
use_platform_cp = not enable_cp_v2() and dsa_use_prefill_cp(forward_batch)
if input_embeds is None:
hidden_states = self.embed_tokens(input_ids)
else:
@@ -184,12 +160,6 @@ class DeepseekV4ModelNextN(nn.Module):
else:
input_ids_global = getattr(forward_batch, "input_ids_global", input_ids)
if use_platform_cp:
hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
positions = cp_split_and_rebuild_position(forward_batch, positions)
input_ids = cp_round_robin_input_ids(input_ids)
input_ids_global = input_ids
if _is_npu:
# Same per-forward rope prime as DeepseekV4Model.forward: the
# decoder layer reads the memoized gather instead of re-gathering.
@@ -207,14 +177,6 @@ class DeepseekV4ModelNextN(nn.Module):
# deferred fused hc_post state.
hidden_states = self.decoder.hc_post(hidden_states, residual, post, comb)
if use_platform_cp:
hidden_states = cp_all_gather_rerange_output(
hidden_states,
self.cp_size,
forward_batch,
torch.cuda.current_stream(),
)
pre_hc_head = hidden_states.flatten(1)
hidden_states = self.hc_head(
@@ -238,13 +200,6 @@ class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM):
self.pp_group = get_pp_group()
self.quant_config = quant_config
self.determine_num_fused_shared_experts()
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
if self.dsa_enable_prefill_cp:
self.cp_rank = get_parallel().attn_cp_rank
self.cp_size = get_parallel().attn_cp_size
else:
self.cp_rank = None
self.cp_size = None
self.model = DeepseekV4ModelNextN(
config, quant_config, prefix=add_prefix("model", prefix)
@@ -265,25 +220,6 @@ class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM):
positions: torch.Tensor,
forward_batch: ForwardBatch,
) -> torch.Tensor:
if self.dsa_enable_prefill_cp and not enable_cp_v2():
if can_dsa_cp_split(len(input_ids), self.cp_size, True, 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 is_dsa_prefill_cp_round_robin_split():
attn_backend = get_attn_backend()
metadata = attn_backend.forward_metadata
core_meta = metadata.core_attn_metadata
core_meta.apply_cp_reindex()
core_meta.init_flashmla_related(is_prefill=True)
if metadata.indexer_metadata is not None:
metadata.indexer_metadata = (
attn_backend.init_forward_metadata_indexer(core_meta)
)
hidden_states, pre_hc_head = self.model(input_ids, positions, forward_batch)
return self.logits_processor(
python/sglang/srt/models/qwen2_moe.py
-32
View File
@@ -48,7 +48,6 @@ from sglang.srt.layers.communicator import (
LayerScatterModes,
ScatterMode,
)
from sglang.srt.layers.cp.utils import enable_cp_v2
from sglang.srt.layers.dp_attention import (
is_dp_attention_enabled,
)
@@ -77,12 +76,6 @@ from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.rotary_embedding import get_rope
from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
from sglang.srt.layers.utils.cp_utils import (
cp_all_gather_rerange_output,
cp_split_and_rebuild_data,
cp_split_and_rebuild_position,
is_prefill_context_parallel_enabled,
)
from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
VocabParallelEmbedding,
@@ -1064,7 +1057,6 @@ class Qwen2MoeModel(nn.Module):
self.pp_group = get_pp_group()
self.moe_dp_size = get_parallel().moe_dp_size
self.attn_cp_size = get_parallel().attn_cp_size
if self.pp_group.is_first_rank:
self.embed_tokens = VocabParallelEmbedding(
@@ -1130,16 +1122,6 @@ class Qwen2MoeModel(nn.Module):
hidden_states = pp_proxy_tensors["hidden_states"]
residual = pp_proxy_tensors["residual"]
if (
is_prefill_context_parallel_enabled()
and not enable_cp_v2()
and forward_batch.forward_mode.is_context_parallel_extend()
and forward_batch.attn_cp_metadata is not None
):
if self.pp_group.is_first_rank:
hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
positions = cp_split_and_rebuild_position(forward_batch, positions)
aux_hidden_states = []
if forward_batch.can_run_tbo:
hidden_states, residual = model_forward_maybe_tbo(
@@ -1196,20 +1178,6 @@ class Qwen2MoeModel(nn.Module):
else:
hidden_states, _ = self.norm(hidden_states, residual)
if (
self.pp_group.is_last_rank
and not enable_cp_v2()
and is_prefill_context_parallel_enabled()
and forward_batch.forward_mode.is_context_parallel_extend()
and forward_batch.attn_cp_metadata is not None
):
hidden_states = cp_all_gather_rerange_output(
hidden_states,
self.attn_cp_size,
forward_batch,
torch.cuda.current_stream(),
)
if len(aux_hidden_states) == 0:
return hidden_states
python/sglang/srt/models/qwen3_moe.py
-16
View File
@@ -34,7 +34,6 @@ from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_r
from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes
from sglang.srt.layers.cp.utils import enable_cp_v2
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import (
QKVParallelLinear,
@@ -58,11 +57,6 @@ from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.rotary_embedding import MRotaryEmbedding, get_rope
from sglang.srt.layers.utils import get_layer_id
from sglang.srt.layers.utils.cp_utils import (
can_cp_split,
is_prefill_context_parallel_enabled,
prepare_context_parallel_metadata,
)
from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_loader.weight_utils import default_weight_loader
@@ -975,7 +969,6 @@ class Qwen3MoeForCausalLM(nn.Module):
)
self.attn_cp_size = get_parallel().attn_cp_size
self.attn_cp_rank = get_parallel().attn_cp_rank
self.moe_dp_size = get_parallel().moe_dp_size
assert self.attn_cp_size % self.moe_dp_size == 0, (
@@ -995,15 +988,6 @@ class Qwen3MoeForCausalLM(nn.Module):
input_embeds: torch.Tensor = None,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> torch.Tensor:
if is_prefill_context_parallel_enabled() and not enable_cp_v2():
if can_cp_split(len(input_ids), self.attn_cp_size, forward_batch):
forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
len(input_ids),
self.attn_cp_rank,
self.attn_cp_size,
forward_batch.seq_lens_cpu.tolist(),
extend_seqs_len=forward_batch.extend_seq_lens_cpu,
)
hidden_states = self.model(
input_ids,
python/sglang/srt/models/utils.py
+2 -6
View File
@@ -37,7 +37,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
from sglang.srt.model_executor.runner import get_is_capture_mode
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.runtime_context import get_exec, get_parallel
from sglang.srt.runtime_context import get_exec
from sglang.srt.utils import get_current_device_stream_fast, is_cuda, is_hip
from sglang.srt.utils.custom_op import register_custom_op
@@ -298,11 +298,7 @@ def enable_fused_set_kv_buffer(forward_batch: ForwardBatch):
and not isinstance(pool, SWAKVPool)
and not is_cp_v2_active(forward_batch)
and getattr(forward_batch, "dcp_kv_mask", None) is None
) or (
_is_hip
and not get_parallel().enable_prefill_context_parallel
and getattr(forward_batch, "dcp_kv_mask", None) is None
)
) or (_is_hip and getattr(forward_batch, "dcp_kv_mask", None) is None)
def create_fused_set_kv_buffer_arg(
python/sglang/srt/server_args.py
-24
View File
@@ -596,13 +596,6 @@ class ServerArgs:
dest="cuda_graph_max_bs_prefill",
help="Deprecated alias for --cuda-graph-max-bs-prefill.",
)
parser.add_argument(
"--enable-nsa-prefill-context-parallel",
dest="enable_dsa_prefill_context_parallel",
action=DeprecatedStoreTrueAction,
new_flag="--enable-prefill-cp",
help="[Deprecated] Use --enable-prefill-cp instead.",
)
parser.add_argument(
"--enable-gdn-replayssm-spec",
dest="enable_linear_replayssm_spec",
@@ -610,23 +603,6 @@ class ServerArgs:
new_flag="--enable-linear-replayssm-spec",
help="[Deprecated] Use --enable-linear-replayssm-spec instead.",
)
parser.add_argument(
"--enable-prefill-context-parallel",
dest="enable_prefill_context_parallel",
action=DeprecatedStoreTrueAction,
new_flag="--enable-prefill-cp",
help="[Deprecated] Use --enable-prefill-cp instead.",
)
parser.add_argument(
"--nsa-prefill-cp-mode",
dest="dsa_prefill_cp_mode",
action=DeprecatedAliasStoreAction,
new_flag="--cp-strategy",
type=str,
default=argparse.SUPPRESS,
choices=["in-seq-split", "round-robin-split"],
help="[Deprecated] Use --cp-strategy instead.",
)
parser.add_argument(
"--enable-flashinfer-allreduce-fusion",
action="store_true",
test/registered/amd/test_deepseek_v4_pro_fp4_cp.py
+7 -1
View File
@@ -23,7 +23,10 @@ from sglang.test.test_utils import (
)
register_amd_ci(
est_time=5400, suite="nightly-amd-8-gpu-mi35x-deepseek-v4-pro", nightly=True
est_time=5400,
suite="nightly-amd-8-gpu-mi35x-deepseek-v4-pro",
nightly=True,
disabled="Prefill CP on HIP/NPU/MUSA is deprecated; CP support will be refactored soon.",
)
DEEPSEEK_V4_PRO_FP4_MODEL_PATH = os.environ.get(
@@ -49,6 +52,9 @@ FP4_ENV_VARS = {
}
@unittest.skip(
"Prefill CP on HIP/NPU/MUSA is deprecated; CP support will be refactored soon."
)
class TestDeepseekV4ProFp4CPInterleave(CustomTestCase):
"""DeepSeek-V4-Pro FP4 unified_kv prefill CP, interleave (round-robin-split), tp=8."""
test/registered/amd/test_deepseek_v4_pro_fp4_cp_tbo.py
+7 -1
View File
@@ -36,7 +36,10 @@ from sglang.test.test_utils import (
)
register_amd_ci(
est_time=5400, suite="nightly-amd-8-gpu-mi35x-deepseek-v4-pro", nightly=True
est_time=5400,
suite="nightly-amd-8-gpu-mi35x-deepseek-v4-pro",
nightly=True,
disabled="Prefill CP on HIP/NPU/MUSA is deprecated; CP support will be refactored soon.",
)
DEEPSEEK_V4_PRO_FP4_MODEL_PATH = os.environ.get(
@@ -63,6 +66,9 @@ FP4_ENV_VARS = {
}
@unittest.skip(
"Prefill CP on HIP/NPU/MUSA is deprecated; CP support will be refactored soon."
)
class TestDeepseekV4ProFp4CPInterleaveTbo(CustomTestCase):
"""DeepSeek-V4-Pro FP4 unified_kv prefill CP (round-robin-split) + TBO, tp=8."""
test/registered/cp/test_cp_strategy_unit.py
+1 -2
View File
@@ -99,9 +99,8 @@ class TestCPStrategyUnit(CustomTestCase):
self.assertTrue(is_cp_enabled())
self.assertTrue(is_interleave())
def test_hip_dsa_cp_uses_protected_legacy_runtime_flag(self):
def test_hip_dsa_cp_is_disabled(self):
parallel = SimpleNamespace(
enable_dsa_prefill_context_parallel=False,
attn_cp_size=2,
)
model_config = SimpleNamespace(hf_config=SimpleNamespace())
test/registered/npu/accuracy/glm5_1/test_npu_glm5_1_w4a8_1p1d_32p_in64k_out1k_50ms_aime26.py
+4 -1
View File
@@ -13,7 +13,7 @@ register_npu_ci(
est_time=4800,
suite="",
nightly=True,
disabled="accuracy testcase",
disabled="Prefill CP on HIP/NPU/MUSA is deprecated; CP support will be refactored soon.",
)
GLM_5_1_PD_SEP_PREFILL_ENVS = {
@@ -164,6 +164,9 @@ GLM_5_1_PD_SEP_MODEL_CONFIG = {
}
@unittest.skip(
"Prefill CP on HIP/NPU/MUSA is deprecated; CP support will be refactored soon."
)
class TestNPUGLM5_1_W4A8_PD_SEP_AIME2026(TestNpuAccuracyMultiNodePdSepTestCaseBase):
"""Test NPU accuracy for GLM-5.1-w4a8 PD separation on AIME2026"""
test/registered/npu/llm_models/test_npu_qwen3_30b_attn_cp.py
+9 -1
View File
@@ -6,9 +6,17 @@ from sglang.test.ascend.test_ascend_utils import QWEN3_30B_A3B_WEIGHTS_PATH
from sglang.test.ci.ci_register import register_npu_ci
from sglang.test.test_utils import CustomTestCase
register_npu_ci(est_time=500, suite="full-4-npu-a3", nightly=True)
register_npu_ci(
est_time=500,
suite="full-4-npu-a3",
nightly=True,
disabled="Prefill CP on HIP/NPU/MUSA is deprecated; CP support will be refactored soon.",
)
@unittest.skip(
"Prefill CP on HIP/NPU/MUSA is deprecated; CP support will be refactored soon."
)
class TestQwen330BAttnCP(GSM8KAscendMixin, CustomTestCase):
"""GSM8K accuracy test for Qwen3-30B-A3B mixed deployment on 4 NPUs.
test/registered/npu/performance/glm5_1/test_npu_glm5_1_w4a8_1p1d_32p_in64k_out1k_50ms.py
+4 -1
View File
@@ -13,7 +13,7 @@ register_npu_ci(
est_time=3600,
suite="",
nightly=True,
disabled="performance testcase",
disabled="Prefill CP on HIP/NPU/MUSA is deprecated; CP support will be refactored soon.",
)
GLM_5_1_PD_SEP_PREFILL_ENVS = {
@@ -172,6 +172,9 @@ GLM_5_1_PD_SEP_MODEL_CONFIG = {
}
@unittest.skip(
"Prefill CP on HIP/NPU/MUSA is deprecated; CP support will be refactored soon."
)
class TestNPUGLM5_1_W4A8_PD_SEP_In3k5_Out1k5(TestNpuPerfMultiNodePdSepTestCaseBase):
"""Test NPU performance for GLM-5.1-w4a8 PD separation 4 nodes in3k5 out1k5"""
test/registered/unit/models/test_deepseek_mla_dispatch.py
-6
View File
@@ -17,7 +17,6 @@ from unittest import mock
from sglang.srt.layers.cp import base as cp_base
from sglang.srt.layers.cp import utils as cp_utils
from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy
from sglang.srt.layers.utils import cp_utils as platform_cp_utils
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.models.deepseek_common import attention_backend_handler as abh
from sglang.srt.models.deepseek_common.attention_forward_methods.forward_methods import (
@@ -114,11 +113,6 @@ class TestCPMLADispatch(CustomTestCase):
mock.patch.object(abh, "_is_hip", False),
mock.patch.object(cp_utils, "enable_cp_v2", return_value=True),
mock.patch.object(cp_base, "_STRATEGY", ZigzagCPStrategy(cp_size=4)),
mock.patch.object(
platform_cp_utils,
"get_parallel",
return_value=SimpleNamespace(enable_prefill_context_parallel=False),
),
):
for prefix in (0, 32):
for capacity in (0, 8192):
test/registered/unit/models/test_deepseek_nextn_mm_embed.py
-16
View File
@@ -92,14 +92,6 @@ class TestDeepseekNextNMmEmbed(CustomTestCase):
object.__setattr__(model, "embed_tokens", mock_embed)
with (
patch(
"sglang.srt.models.deepseek_nextn.dsa_use_prefill_cp",
return_value=False,
),
patch(
"sglang.srt.models.deepseek_nextn.mla_use_prefill_cp",
return_value=False,
),
patch(
"sglang.srt.models.deepseek_nextn.fused_eh_norm",
side_effect=lambda h, p, ew, hw, eps: torch.cat(
@@ -157,14 +149,6 @@ class TestDeepseekNextNMmEmbed(CustomTestCase):
embed_calls = mock_embed.call_args_list
with (
patch(
"sglang.srt.models.deepseek_nextn.dsa_use_prefill_cp",
return_value=False,
),
patch(
"sglang.srt.models.deepseek_nextn.mla_use_prefill_cp",
return_value=False,
),
patch(
"sglang.srt.models.deepseek_nextn.fused_eh_norm",
side_effect=lambda h, p, ew, hw, eps: torch.cat(
test/registered/unit/server_args/test_platform_prefill_cp_deprecation.py
+84
View File
@@ -0,0 +1,84 @@
"""Reject deprecated platform CP before model loading or topology setup."""
import unittest
from sglang.srt.arg_groups.parallel_hook import (
handle_context_parallelism,
validate_prefill_cp_platform,
)
from sglang.srt.runtime_context import override_platform
from sglang.srt.server_args import ServerArgs
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
class TestPlatformPrefillCPDeprecation(CustomTestCase):
def test_platform_cp_rejected_before_model_lookup(self):
for platform in ("is_hip", "is_npu", "is_musa"):
facts = dict(is_hip=False, is_npu=False, is_musa=False)
facts[platform] = True
for strategy in (None, "zigzag", "interleave"):
with self.subTest(platform=platform, strategy=strategy):
with override_platform(**facts):
args = ServerArgs(
model_path="missing-model-must-not-be-loaded",
enable_prefill_cp=True,
cp_strategy=strategy,
)
with self.assertRaisesRegex(ValueError, "deprecated.*refactor"):
validate_prefill_cp_platform(args)
def test_context_parallel_handler_rejects_before_model_lookup(self):
for platform in ("is_hip", "is_npu", "is_musa"):
facts = dict(is_hip=False, is_npu=False, is_musa=False)
facts[platform] = True
with self.subTest(platform=platform), override_platform(**facts):
args = ServerArgs(
model_path="missing-model-must-not-be-loaded",
enable_prefill_cp=True,
cp_strategy="interleave",
)
with self.assertRaisesRegex(ValueError, "deprecated.*refactor"):
handle_context_parallelism(args)
def test_resolution_rejects_even_dummy_models(self):
for platform in ("is_hip", "is_npu", "is_musa"):
facts = dict(is_hip=False, is_npu=False, is_musa=False)
facts[platform] = True
for model_path in ("dummy", "none", "missing-model-must-not-be-loaded"):
with self.subTest(platform=platform, model_path=model_path):
with override_platform(**facts):
args = ServerArgs(
model_path=model_path,
enable_prefill_cp=True,
cp_strategy="interleave",
)
with self.assertRaisesRegex(ValueError, "deprecated.*refactor"):
args.resolve_once()
def test_non_cp_and_decode_cp_are_not_rejected(self):
for platform in ("is_hip", "is_npu", "is_musa"):
facts = dict(is_hip=False, is_npu=False, is_musa=False)
facts[platform] = True
for dcp_size in (1, 2):
with self.subTest(platform=platform, dcp_size=dcp_size):
with override_platform(**facts):
args = ServerArgs(model_path="dummy", dcp_size=dcp_size)
validate_prefill_cp_platform(args)
@override_platform(is_hip=False, is_npu=False, is_musa=False)
def test_generic_cp_is_not_rejected_or_modified(self):
for strategy in ("zigzag", "interleave"):
with self.subTest(strategy=strategy):
args = ServerArgs(
model_path="dummy", enable_prefill_cp=True, cp_strategy=strategy
)
validate_prefill_cp_platform(args)
self.assertTrue(args.enable_prefill_cp)
self.assertEqual(args.cp_strategy, strategy)
if __name__ == "__main__":
unittest.main()
test/registered/unit/server_args/test_server_args.py
+3 -99
View File
@@ -49,8 +49,6 @@ from sglang.srt.arg_groups.overrides import (
from sglang.srt.arg_groups.parallel_hook import (
handle_context_parallelism,
handle_data_parallelism,
handle_legacy_cp_runtime_compatibility,
handle_platform_cp_compatibility,
)
from sglang.srt.arg_groups.pd_disaggregation_hook import handle_pd_disaggregation
from sglang.srt.arg_groups.serving_hook import (
@@ -1026,13 +1024,9 @@ class TestContextParallelServerArgs(CustomTestCase):
def _new_cp_args(self, **overrides):
server_args = object.__new__(ServerArgs)
defaults = dict(
enable_prefill_context_parallel=False,
enable_dsa_prefill_context_parallel=False,
enable_prefill_cp=False,
cp_strategy=None,
model_path="instance://127.0.0.1:8000/dummy",
dsa_prefill_cp_mode="round-robin-split",
prefill_cp_mode="in-seq-split",
attn_cp_size=1,
tp_size=1,
dp_size=1,
@@ -1075,52 +1069,11 @@ class TestContextParallelServerArgs(CustomTestCase):
with self.assertRaisesRegex(ValueError, "DeepSeek V3.2.*interleave"):
handle_context_parallelism(server_args)
@override_platform(is_hip=False, is_npu=False, is_musa=False)
def test_generic_canonical_cp_does_not_enable_platform_runtime_fields(self):
cases = (
(
"zigzag_mla_or_gqa",
"zigzag",
"fa3",
),
(
"interleave_dsa",
"interleave",
"dsa",
),
)
for name, strategy, backend in cases:
with self.subTest(name=name):
server_args = self._new_cp_args(
enable_prefill_cp=True,
cp_strategy=strategy,
attention_backend=backend,
)
handle_platform_cp_compatibility(server_args)
handle_legacy_cp_runtime_compatibility(server_args)
self.assertFalse(
resolution_result(server_args, "enable_prefill_context_parallel")
)
self.assertFalse(
resolution_result(
server_args, "enable_dsa_prefill_context_parallel"
)
)
@override_platform(is_hip=False, is_npu=False, is_musa=False)
def test_non_platform_legacy_prefill_cp_is_rejected(self):
server_args = ServerArgs(
model_path="instance://127.0.0.1:8000/dummy",
enable_prefill_context_parallel=True,
)
with self.assertRaisesRegex(ValueError, "protected HIP, Ascend NPU, or MUSA"):
handle_platform_cp_compatibility(server_args)
def test_generic_v1_cp_options_are_not_public_cli(self):
removed_options = (
("--enable-prefill-context-parallel", []),
("--enable-nsa-prefill-context-parallel", []),
("--nsa-prefill-cp-mode", ["round-robin-split"]),
("--enable-dsa-prefill-context-parallel", []),
("--dsa-prefill-cp-mode", ["round-robin-split"]),
("--prefill-cp-mode", ["in-seq-split"]),
@@ -1130,55 +1083,6 @@ class TestContextParallelServerArgs(CustomTestCase):
with self.subTest(option=option), self.assertRaises(SystemExit):
self.parser.parse_args(["--model", "dummy", option, *values])
def test_npu_cp_compatibility_options_remain_public_cli(self):
args = self.parser.parse_args(
[
"--model",
"dummy",
"--enable-prefill-context-parallel",
"--enable-nsa-prefill-context-parallel",
"--nsa-prefill-cp-mode",
"round-robin-split",
]
)
self.assertTrue(resolution_result(args, "enable_prefill_context_parallel"))
self.assertTrue(resolution_result(args, "enable_dsa_prefill_context_parallel"))
self.assertEqual(
resolution_result(args, "dsa_prefill_cp_mode"), "round-robin-split"
)
def test_platform_interleave_cp_mirrors_to_dsa_runtime_aliases(self):
for platform in ("is_hip", "is_npu", "is_musa"):
facts = dict(is_hip=False, is_npu=False, is_musa=False)
facts[platform] = True
with self.subTest(platform=platform), override_platform(**facts):
server_args = self._new_cp_args(
enable_prefill_cp=True,
cp_strategy="interleave",
attention_backend="dsa",
)
handle_legacy_cp_runtime_compatibility(server_args)
handle_context_parallelism(server_args)
self.assertTrue(
resolution_result(
server_args, "enable_dsa_prefill_context_parallel"
)
)
self.assertFalse(
resolution_result(server_args, "enable_prefill_context_parallel")
)
self.assertEqual(
resolution_result(server_args, "dsa_prefill_cp_mode"),
"round-robin-split",
)
self.assertEqual(
resolution_result(server_args, "prefill_cp_mode"),
"round-robin-split",
)
def test_context_parallel_handler_initializes_cp_strategy(self):
server_args = self._new_cp_args(
enable_prefill_cp=True,
test/registered/unit/test_runtime_context.py
+1 -1
View File
@@ -925,7 +925,7 @@ class TestForwardFlags(_IsolatedServerArgs):
@torch.compile(fullgraph=True, backend="eager", dynamic=False)
def probe(x):
par = get_parallel()
if par.enable_prefill_context_parallel:
if par.enable_prefill_cp:
x = x + 1
if par.moe_dense_tp_size == 1:
x = x + 2