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

[CP V1 Deprecation 4/5] Canonicalize prefill CP API names (#36229)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-09-08 16:03:36 -07:00
committed by GitHub
parent 559c7fa75b
commit ed183d45ac
33 changed files with 143 additions and 173 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/__init__.py
+1 -1
View File
@@ -181,7 +181,7 @@ for _mod, _fn in [
("dsa.transform_index", "transform_index_page_table_prefill"),
("dsa.transform_index", "transform_index_page_table_decode"),
("dsa.transform_index", "prepare_trtllm_nope_sparse_metadata"),
("dsa.cp_split", "dsa_cp_round_robin_split_q_seqs_kernel"),
("dsa.cp_split", "dsa_cp_interleave_q_seqs_kernel"),
("dsv4.fp4_indexer", "quantize_fp4_indexer_tensor"),
("dsv4.fp4_indexer", "store_fp4_index_k_cache"),
("dsv4.rms_normalize_hip", "rms_normalize_triton"),
python/sglang/kernels/ops/attention/dsa/cp_split.py
+2 -2
View File
@@ -1,4 +1,4 @@
"""Round-robin CP q-sequence split kernel for DSA prefill.
"""Interleave CP q-sequence split kernel for DSA prefill.
Migrated from ``sglang.srt.layers.attention.dsa.utils`` (RFC #29630, Phase 2.5).
"""
@@ -8,7 +8,7 @@ import triton.language as tl
@triton.jit
def dsa_cp_round_robin_split_q_seqs_kernel(
def dsa_cp_interleave_q_seqs_kernel(
in_seqs_ptr,
out_seqs_ptr,
bs_idx_ptr,
python/sglang/kernels/ops/attention/dsv4/metadata_kernel.py
+1 -1
View File
@@ -106,7 +106,7 @@ def _init_compressed_attn_metadata_triton(
Optional[torch.Tensor],
]:
bs = seq_lens.shape[0]
# CP-v2 may add padding rows to the attention metadata, but those rows have
# CP may add padding rows to the attention metadata, but those rows have
# no cache-write locations. Keep the write buffers unpadded and mask those
# rows in the kernel.
num_write_tokens = raw_out_loc.shape[0]
python/sglang/kernels/ops/layernorm/mhc.py
+2 -2
View File
@@ -15,7 +15,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
)
from sglang.srt.distributed.parallel_state import get_tp_group
from sglang.srt.environ import envs
from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_round_robin_split
from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_interleave
from sglang.srt.layers.dp_attention import is_allocation_symmetric
from sglang.srt.layers.utils.common import strict_contiguous
from sglang.srt.utils.common import is_gfx1250_supported
@@ -1240,7 +1240,7 @@ def mhc_post(
post_layer_mix: torch.Tensor,
comb_res_mix: torch.Tensor,
) -> torch.Tensor:
if is_dsa_prefill_cp_round_robin_split():
if is_dsa_prefill_cp_interleave():
x = strict_contiguous(x)
residual = strict_contiguous(residual)
post_layer_mix = strict_contiguous(post_layer_mix)
python/sglang/srt/layers/attention/deepseek_v4_backend.py
+8 -8
View File
@@ -66,7 +66,7 @@ from sglang.srt.layers.attention.verify_mask import (
VerifyMask,
maybe_create_verify_mask,
)
from sglang.srt.layers.cp.utils import is_cp_v2_active
from sglang.srt.layers.cp.utils import is_cp_active
from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.runtime_context import (
@@ -290,7 +290,7 @@ class DSV4AttnMetadata:
def init_compression_metadata(self, num_tokens: Optional[int] = None) -> None:
assert self.page_table.dim() == 2
# CP-v2 pads causal metadata for per-rank partitioning, while cache-write
# CP pads causal metadata for per-rank partitioning, while cache-write
# locations remain one-per-logical-token. num_tokens tracks that unpadded
# length; legacy paths use the metadata length.
if num_tokens is None:
@@ -746,8 +746,8 @@ class DeepseekV4AttnBackend(
forward_batch: Optional[ForwardBatch] = None,
) -> DSV4Metadata:
padded_num_tokens = out_cache_loc.shape[0]
cp_v2_active = forward_batch is not None and is_cp_v2_active(forward_batch)
if cp_v2_active:
cp_active = forward_batch is not None and is_cp_active(forward_batch)
if cp_active:
cp_metadata = forward_batch.attn_cp_metadata
assert cp_metadata is not None
padded_num_tokens = sum(cp_metadata.per_rank_actual_token)
@@ -771,9 +771,9 @@ class DeepseekV4AttnBackend(
need_compress=need_compress,
is_prefill=True,
dspark_block_size=dspark_block_size,
num_tokens=num_tokens if cp_v2_active else None,
num_tokens=num_tokens if cp_active else None,
)
if cp_v2_active:
if cp_active:
core_attn_metadata.apply_cp_reindex(num_tokens=num_tokens)
core_attn_metadata.init_flashmla_related(is_prefill=True)
indexer_metadata = (
@@ -1364,7 +1364,7 @@ class DeepseekV4AttnBackend(
) -> None:
# Sparse prefill otherwise reads req_to_token/full_to_swa lazily in its
# first layer. DFLASH/DSPARK have no later prefill draft-extend reader;
# CP-v2 shards the query layout that this global snapshot assumes.
# CP shards the query layout that this global snapshot assumes.
metadata = self.forward_metadata
if isinstance(metadata, DSV4Metadata):
metadata.prefill_shared_reads_snapshotted = False
@@ -1372,7 +1372,7 @@ class DeepseekV4AttnBackend(
envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.get()
and forward_batch.forward_mode == ForwardMode.EXTEND
and self.model_runner.spec_algorithm.is_dflash_family()
and not is_cp_v2_active(forward_batch)
and not is_cp_active(forward_batch)
)
if not snapshot_shared_prefill_reads:
return
python/sglang/srt/layers/attention/dsa/dsa_indexer.py
+3 -3
View File
@@ -117,7 +117,7 @@ from sglang.srt.distributed import (
from sglang.srt.distributed.parallel_state import get_pp_group
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.cp.base import get_cp_strategy
from sglang.srt.layers.cp.utils import is_cp_v2_active
from sglang.srt.layers.cp.utils import is_cp_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
@@ -529,7 +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 is_cp_v2_active(forward_batch):
elif self.alt_stream is not None and is_cp_active(forward_batch):
key = self._maybe_rotate(key)
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
@@ -548,7 +548,7 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
key = self._maybe_rotate(key)
# allgather+rerrange
if is_cp_v2_active(forward_batch):
if is_cp_active(forward_batch):
key = get_cp_strategy().materialize_full_indexer_k_cache(key, forward_batch)
return query, key, weights_raw
python/sglang/srt/layers/attention/dsa/utils.py
+11 -9
View File
@@ -128,10 +128,14 @@ def is_dsa_enable_prefill_cp():
return is_deepseek_dsa(hf_config) or is_deepseek_v4(hf_config)
def is_dsa_prefill_cp_round_robin_split():
def is_dsa_prefill_cp_interleave():
return is_dsa_enable_prefill_cp() and get_parallel().cp_strategy == "interleave"
# Retain the name imported by the unchanged HIP radix attention backend.
is_dsa_prefill_cp_round_robin_split = is_dsa_prefill_cp_interleave
# Structural surface where the graph DSA split-op dispatch (DSA indexer) and the
# MLA BMM-into-attention fusion apply: a non-speculative extend (prefill) running
# inside a piecewise/breakable CUDA graph. Both fusions are now on by default on
@@ -145,13 +149,13 @@ def is_graph_dsa_split_op_surface(forward_batch: "ForwardBatch") -> bool:
)
def can_dsa_prefill_cp_round_robin_split(forward_batch: "ForwardBatch"):
def can_dsa_prefill_cp_interleave(forward_batch: "ForwardBatch"):
if not forward_batch.forward_mode.is_context_parallel_extend():
return False
cp_size = get_parallel().attn_cp_size
seq_len = sum(forward_batch.extend_seq_lens_cpu)
return (
is_dsa_prefill_cp_round_robin_split()
is_dsa_prefill_cp_interleave()
and seq_len > 0
and seq_len >= cp_size
and cp_size > 1
@@ -161,10 +165,10 @@ def can_dsa_prefill_cp_round_robin_split(forward_batch: "ForwardBatch"):
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.utils import is_cp_v2_active
from sglang.srt.layers.cp.utils import is_cp_active
# CP-v2 already pads each rank-local shard to its physical size
if is_cp_v2_active(forward_batch):
if is_cp_active(forward_batch):
return forward_batch.attn_cp_metadata.per_rank_actual_token[
get_parallel().attn_cp_rank
]
@@ -185,16 +189,14 @@ def cal_padded_tokens(forward_batch: "ForwardBatch"):
tokens = global_num_tokens[get_parallel().attn_dp_rank]
else:
tokens = global_num_tokens[0]
if can_dsa_prefill_cp_round_robin_split(forward_batch):
if can_dsa_prefill_cp_interleave(forward_batch):
tokens = ceil_div(tokens, attn_cp_size)
return tokens
def pad_dsa_cache_seqlens(forward_batch: "ForwardBatch", dsa_cache_seqlens):
attn_cp_size = get_parallel().attn_cp_size
needs_cp_pad = attn_cp_size > 1 and can_dsa_prefill_cp_round_robin_split(
forward_batch
)
needs_cp_pad = attn_cp_size > 1 and can_dsa_prefill_cp_interleave(forward_batch)
needs_dp_pad = forward_batch.global_num_tokens_cpu is not None
if not needs_cp_pad and not needs_dp_pad:
return dsa_cache_seqlens
python/sglang/srt/layers/attention/dsa_backend.py
+5 -5
View File
@@ -74,7 +74,7 @@ from sglang.srt.layers.attention.dsa.kpool_plan import (
KPoolWritePlan,
)
from sglang.srt.layers.attention.dsa.utils import (
can_dsa_prefill_cp_round_robin_split,
can_dsa_prefill_cp_interleave,
compute_dsa_seqlens,
dsa_use_prefill_cp,
is_dsa_enable_prefill_cp,
@@ -86,7 +86,7 @@ from sglang.srt.layers.attention.trtllm_mla_backend import (
make_persistent_multi_ctas_kv_counter_buffer,
)
from sglang.srt.layers.cp.base import get_cp_strategy
from sglang.srt.layers.cp.utils import is_cp_v2_active
from sglang.srt.layers.cp.utils import is_cp_active
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 (
@@ -128,7 +128,7 @@ def prepare_kv_for_attention(
if (
defer_materialization
or not dsa_use_prefill_cp(forward_batch)
or not is_cp_v2_active(forward_batch)
or not is_cp_active(forward_batch)
):
return k_nope, k_pe
strategy = get_cp_strategy()
@@ -1005,7 +1005,7 @@ class DeepseekSparseAttnBackend(
)
kpool_inputs.full_seqlens_expanded = seqlens_expanded
if can_dsa_prefill_cp_round_robin_split(forward_batch):
if can_dsa_prefill_cp_interleave(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 = (
@@ -1222,7 +1222,7 @@ class DeepseekSparseAttnBackend(
ke = torch.cat(ke_list, dim=0)
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)
assert can_dsa_prefill_cp_interleave(forward_batch)
split_per_token = get_cp_strategy().shard_local_tokens
ks = split_per_token(ks)
ke = split_per_token(ke)
python/sglang/srt/layers/attention/flashattention_backend.py
+2 -2
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 is_cp_v2_active
from sglang.srt.layers.cp.utils import is_cp_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
@@ -1243,7 +1243,7 @@ class FlashAttentionBackend(AttentionBackend):
):
if score_mod is not None and self.fa_impl_ver != 4:
raise RuntimeError("score_mod is only supported by the FA4 backend.")
cp_active = is_cp_v2_active(forward_batch)
cp_active = is_cp_active(forward_batch)
if k is not None:
assert v is not None
python/sglang/srt/layers/attention/index_topk_share.py
+2 -2
View File
@@ -3,7 +3,7 @@ from __future__ import annotations
from contextlib import contextmanager
from typing import TYPE_CHECKING, Iterator, Optional
from sglang.srt.layers.cp.utils import cp_gather_after_forward, is_cp_v2_active
from sglang.srt.layers.cp.utils import cp_gather_after_forward, is_cp_active
if TYPE_CHECKING:
import torch
@@ -54,7 +54,7 @@ class IndexTopKShareState:
if (
topk_indices is not None
and self.should_publish
and is_cp_v2_active(self._forward_batch)
and is_cp_active(self._forward_batch)
):
topk_indices = cp_gather_after_forward(topk_indices, self._forward_batch)
self._topk_indices = topk_indices
python/sglang/srt/layers/attention/trtllm_mha_backend.py
+10 -10
View File
@@ -37,7 +37,7 @@ from sglang.srt.layers.attention.trtllm_mla_backend import (
make_persistent_multi_ctas_kv_counter_buffer,
)
from sglang.srt.layers.cp.base import CPAttentionBackendKind, get_cp_strategy
from sglang.srt.layers.cp.utils import is_cp_v2_active
from sglang.srt.layers.cp.utils import is_cp_active
from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import (
KVCacheAttentionAccessKind,
)
@@ -90,7 +90,7 @@ class TRTLLMMHAMetadata:
page_table: torch.Tensor = None
# Page table for SWA layers (translated from full pool indices to SWA pool indices)
swa_page_table: torch.Tensor = None
# CP-v2 zigzag treats prev/next halves as a synthetic 2 * batch_size batch.
# CP zigzag treats prev/next halves as a synthetic 2 * batch_size batch.
zigzag_page_table: torch.Tensor = None
zigzag_swa_page_table: torch.Tensor = None
# full->SWA translated out_cache_loc (SWA KV-store write target)
@@ -409,7 +409,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
forward_batch: ForwardBatch,
) -> None:
"""Duplicate request rows once for the combined prev-then-next CP launch."""
if not is_cp_v2_active(forward_batch):
if not is_cp_active(forward_batch):
return
# TODO: Avoid materializing duplicated page tables to reduce zigzag CP
@@ -818,7 +818,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
) -> bool:
"""Check if we should use the fused FP8 KV cache write path."""
return (
not is_cp_v2_active(forward_batch)
not is_cp_active(forward_batch)
and save_kv_cache
and k is not None
and self.data_type == torch.float8_e4m3fn
@@ -1334,9 +1334,9 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
)
cache_loc = forward_batch.out_cache_loc
cp_v2_active = is_cp_v2_active(forward_batch)
cp_active = is_cp_active(forward_batch)
# The fused path writes rank-local K/V directly to cache. CP-v2 needs
# The fused path writes rank-local K/V directly to cache. CP needs
# the strategy to gather K/V into full logical token order first.
use_fused_fp8_path = self._should_use_fused_fp8_path(
save_kv_cache, k, forward_batch
@@ -1353,7 +1353,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
v = None
else:
if save_kv_cache and k is not None:
if cp_v2_active:
if cp_active:
cp_strategy = get_cp_strategy()
assert cp_strategy is not None
cp_strategy.materialize_full_kv(
@@ -1480,8 +1480,8 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
sinks=attention_sink,
q_len_per_req=self.forward_metadata.max_seq_len_q,
)
elif self.use_fmha_v2 and not cp_v2_active:
# CP-v2 must go through cp_strategy.run_attention (per-shard
elif self.use_fmha_v2 and not cp_active:
# CP must go through cp_strategy.run_attention (per-shard
# masking); the plain-causal fmha_v2 call below would be wrong.
paged_kv = torch.stack([k_cache, v_cache], dim=1)
o = flashinfer.prefill.trtllm_fmha_v2_prefill(
@@ -1544,7 +1544,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
out_dtype=self.q_data_type,
)
if cp_v2_active:
if cp_active:
cp_strategy = get_cp_strategy()
assert cp_strategy is not None
o = cp_strategy.run_attention(
python/sglang/srt/layers/communicator.py
+7 -10
View File
@@ -38,8 +38,8 @@ from sglang.srt.layers.attention.dsa.utils import (
)
from sglang.srt.layers.aux_hidden_states import AuxHiddenStateAccumulator
from sglang.srt.layers.cp.utils import (
is_mla_prefill_cp_enabled,
mla_use_prefill_cp,
is_mla_cp_active,
is_mla_cp_enabled,
)
from sglang.srt.layers.dp_attention import (
attn_tp_all_gather_into_tensor,
@@ -238,7 +238,7 @@ class ScatterMode(Enum):
@staticmethod
def model_input_output():
"""The scatter mode for model forward pass input and output data"""
if is_dsa_enable_prefill_cp() or is_mla_prefill_cp_enabled():
if is_dsa_enable_prefill_cp() or is_mla_cp_enabled():
return ScatterMode.SCATTERED
return ScatterMode.TP_ATTN_FULL
@@ -432,7 +432,7 @@ class LayerScatterModes:
return ScatterMode.SCATTERED
# DSA CP and MLA CP both don't support MOE_FULL yet; fall back to FULL.
if is_enable_moe_cp_allgather() and not (
is_dsa_enable_prefill_cp() or is_mla_prefill_cp_enabled()
is_dsa_enable_prefill_cp() or is_mla_cp_enabled()
):
return ScatterMode.MOE_FULL
return ScatterMode.FULL
@@ -444,7 +444,7 @@ class LayerScatterModes:
# first or the all-reduce sums different tokens' partial outputs.
# MLA/DSA CP models do this in DSACPLayerCommunicator instead.
if _generic_prefill_cp_shards_tokens() and not (
is_dsa_enable_prefill_cp() or is_mla_prefill_cp_enabled()
is_dsa_enable_prefill_cp() or is_mla_cp_enabled()
):
return ScatterMode.MOE_FULL
return ScatterMode.FULL
@@ -487,11 +487,8 @@ def enable_moe_dense_fully_dp():
def _generic_prefill_cp_shards_tokens() -> bool:
"""Whether the strategy prefill CP path shards prefill tokens across CP ranks."""
# Local import: module-level CP helper imports here are circular (#27014).
from sglang.srt.layers.cp.utils import enable_cp_v2
parallel = get_parallel()
return parallel.attn_cp_size > 1 and parallel.enable_prefill_cp and enable_cp_v2()
return parallel.attn_cp_size > 1 and parallel.enable_prefill_cp
def enable_dwdp():
@@ -918,7 +915,7 @@ class LayerCommunicator:
return True
if forward_batch.dp_padding_mode.is_max_len():
return True
if dsa_use_prefill_cp(forward_batch) or mla_use_prefill_cp(forward_batch):
if dsa_use_prefill_cp(forward_batch) or is_mla_cp_active(forward_batch):
return True
if get_attn_tp_context().input_scattered and not self.is_last_layer:
return True
python/sglang/srt/layers/communicator_dsa_cp.py
+3 -3
View File
@@ -30,7 +30,7 @@ from sglang.srt.layers.communicator import (
LayerScatterModes,
ScatterMode,
)
from sglang.srt.layers.cp.utils import mla_use_prefill_cp
from sglang.srt.layers.cp.utils import is_mla_cp_active
from sglang.srt.layers.dp_attention import (
attn_cp_all_gather_into_tensor,
attn_cp_reduce_scatter_tensor,
@@ -189,7 +189,7 @@ class DSACPCommunicateWithAllReduceAndLayerNormFn(
hidden_states, residual = layernorm(hidden_states, residual)
# for prefill: attn tp scattered -> full
# for decode: attn tp full -> full
if dsa_use_prefill_cp(forward_batch) or mla_use_prefill_cp(forward_batch):
if dsa_use_prefill_cp(forward_batch) or is_mla_cp_active(forward_batch):
hidden_states = dsa_cp_gather_hidden_states(hidden_states)
return hidden_states, residual
@@ -234,6 +234,6 @@ class DSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn):
):
# for prefill: full -> attn tp scattered
# for decode: full -> attn tp full
if dsa_use_prefill_cp(forward_batch) or mla_use_prefill_cp(forward_batch):
if dsa_use_prefill_cp(forward_batch) or is_mla_cp_active(forward_batch):
hidden_states = dsa_cp_reduce_scatter_hidden_states(hidden_states)
return hidden_states, residual
python/sglang/srt/layers/cp/base.py
+2 -2
View File
@@ -180,9 +180,9 @@ class ContextParallelStrategy(ABC):
positions: Any,
input_embeds: Optional[Any] = None,
) -> Optional[Any]:
"""Shard model inputs before model.forward in CP-v2 paths."""
"""Shard model inputs before model.forward in CP paths."""
if input_ids is not None:
forward_batch.cp_v2_input_ids = self.shard_hidden_states(
forward_batch.cp_input_ids = self.shard_hidden_states(
input_ids, forward_batch
)
forward_batch.positions = self.shard_position_ids(positions, forward_batch)
python/sglang/srt/layers/cp/bcg.py
+3 -4
View File
@@ -31,7 +31,6 @@ from sglang.srt.layers.cp.padding import get_cp_padding_align_size
from sglang.srt.layers.cp.utils import (
cp_gather_after_forward,
cp_split_before_forward,
enable_cp_v2,
prepare_cp_forward,
)
from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy
@@ -60,9 +59,9 @@ def supports_prefill_cp_bcg(server_args: ServerArgs) -> bool:
)
def enable_cp_v2_bcg_capture(server_args: ServerArgs) -> bool:
"""Return whether CP-v2 breakable prefill capture is enabled."""
return enable_cp_v2() and supports_prefill_cp_bcg(server_args)
def enable_cp_bcg_capture(server_args: ServerArgs) -> bool:
"""Return whether CP breakable prefill capture is enabled."""
return supports_prefill_cp_bcg(server_args)
def filter_prefill_cp_bcg_capture_num_tokens(
python/sglang/srt/layers/cp/cp_decode_attn_tp.py
+2 -2
View File
@@ -13,7 +13,7 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
import torch
from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
from sglang.srt.layers.cp.utils import is_cp_v2_active
from sglang.srt.layers.cp.utils import is_cp_active
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.runtime_context import get_parallel
@@ -74,7 +74,7 @@ class CpDecodeAttnTpContext:
return
# Skip during prefill context parallel (needs all heads); apply on every
# other forward, which includes decode.
self.use_decode_attn_tp = not is_cp_v2_active(
self.use_decode_attn_tp = not is_cp_active(
forward_batch
) and not dsa_use_prefill_cp(forward_batch)
python/sglang/srt/layers/cp/interleave.py
+2 -2
View File
@@ -132,7 +132,7 @@ class InterleaveCPStrategy(ContextParallelStrategy):
):
"""Build device outputs in the shared kernel to keep the split graph-safe."""
from sglang.kernels.ops.attention.dsa.cp_split import (
dsa_cp_round_robin_split_q_seqs_kernel,
dsa_cp_interleave_q_seqs_kernel,
)
cp_size = self.cp_size
@@ -154,7 +154,7 @@ class InterleaveCPStrategy(ContextParallelStrategy):
bs_idx = torch.empty(
(len(bs_idx_cpu),), device=extend_seqs.device, dtype=torch.int32
)
dsa_cp_round_robin_split_q_seqs_kernel[(1,)](
dsa_cp_interleave_q_seqs_kernel[(1,)](
extend_seqs, q_lens, bs_idx, len(extend_seqs), cp_size, cp_rank
)
return q_lens_cpu, q_lens, bs_idx_cpu, bs_idx
python/sglang/srt/layers/cp/utils.py
+23 -33
View File
@@ -118,17 +118,8 @@ def get_layer_owner(local_layer_idx: int, shard_size: int, total_layers: int) ->
)
def enable_cp_v2() -> bool:
"""Return whether the strategy-based generic prefill CP path is available."""
from sglang.srt.utils import is_hip, is_musa, is_npu
return not (is_hip() or is_npu() or is_musa())
def is_cp_v2_active(forward_batch) -> bool:
"""Return whether the current forward batch is running through CP-v2."""
if not enable_cp_v2():
return False
def is_cp_active(forward_batch) -> bool:
"""Return whether the current forward batch is running through CP."""
forward_mode = getattr(forward_batch, "forward_mode", None)
if forward_mode is None or not forward_mode.is_context_parallel_extend():
return False
@@ -144,19 +135,19 @@ def is_cp_v2_active(forward_batch) -> bool:
return strategy.can_apply(len(input_ids), forward_batch)
def is_mla_prefill_cp_enabled() -> bool:
def is_mla_cp_enabled() -> bool:
"""Return whether prefill CP is configured for an MLA attention backend."""
return enable_cp_v2() and is_cp_enabled() and uses_mla_backend()
return is_cp_enabled() and uses_mla_backend()
def mla_use_prefill_cp(forward_batch) -> bool:
def is_mla_cp_active(forward_batch) -> bool:
"""Return whether this MLA forward batch is using prefill CP."""
return is_mla_prefill_cp_enabled() and is_cp_v2_active(forward_batch)
return is_mla_cp_enabled() and is_cp_active(forward_batch)
def prepare_cp_forward(forward_batch) -> None:
"""Build CP-v2 metadata for an active context-parallel prefill batch."""
assert is_cp_v2_active(forward_batch)
"""Build CP metadata for an active context-parallel prefill batch."""
assert is_cp_active(forward_batch)
strategy = get_cp_strategy()
assert strategy is not None
@@ -191,8 +182,8 @@ def cp_split_before_forward(
complete_position_ids: Any,
forward_batch,
) -> Tuple[Optional[Any], Optional[Any]]:
"""Shard embeddings and positions for CP-v2 model-runner forwarding."""
assert is_cp_v2_active(forward_batch)
"""Shard embeddings and positions for CP model-runner forwarding."""
assert is_cp_active(forward_batch)
assert complete_hidden_states is not None
assert getattr(forward_batch, "attn_cp_metadata", None) is not None
return (
@@ -202,7 +193,7 @@ def cp_split_before_forward(
def cp_shard_hidden_states(complete_hidden_states: Any, forward_batch):
assert is_cp_v2_active(forward_batch)
assert is_cp_active(forward_batch)
strategy = get_cp_strategy()
assert strategy is not None
assert complete_hidden_states is not None
@@ -211,7 +202,7 @@ def cp_shard_hidden_states(complete_hidden_states: Any, forward_batch):
def cp_shard_position_ids(complete_position_ids: Any, forward_batch):
assert is_cp_v2_active(forward_batch)
assert is_cp_active(forward_batch)
strategy = get_cp_strategy()
assert strategy is not None
assert complete_position_ids is not None
@@ -219,8 +210,8 @@ def cp_shard_position_ids(complete_position_ids: Any, forward_batch):
return strategy.shard_position_ids(complete_position_ids, forward_batch)
def cp_round_robin_input_ids_v2(input_ids: Any, forward_batch):
assert is_cp_v2_active(forward_batch)
def cp_interleave_input_ids(input_ids: Any, forward_batch):
assert is_cp_active(forward_batch)
if not get_moe_a2a_backend().is_none():
return cp_shard_hidden_states(input_ids, forward_batch)
@@ -231,8 +222,8 @@ def cp_round_robin_input_ids_v2(input_ids: Any, forward_batch):
def cp_gather_after_forward(x: Any, forward_batch, stream: Optional[Any] = None):
"""Gather CP-v2 hidden states at the model boundary when this batch is active."""
assert is_cp_v2_active(forward_batch)
"""Gather CP hidden states at the model boundary when this batch is active."""
assert is_cp_active(forward_batch)
strategy = get_cp_strategy()
assert strategy is not None
@@ -257,7 +248,7 @@ def cp_materialize_global_token_order(
x: Any, forward_batch, stream: Optional[Any] = None
):
"""Materialize a CP tensor in the global logical token order."""
assert is_cp_v2_active(forward_batch)
assert is_cp_active(forward_batch)
strategy = get_cp_strategy()
assert strategy is not None
return strategy.gather_kv_cache(x, forward_batch, stream)
@@ -271,7 +262,7 @@ def cp_shard_model_inputs(
complete_input_ids: Optional[Any] = None,
):
"""Restore the shared batch so logits processing keeps full-batch metadata."""
assert is_cp_v2_active(forward_batch)
assert is_cp_active(forward_batch)
sharded_hidden_states = cp_shard_hidden_states(
complete_hidden_states, forward_batch
)
@@ -285,7 +276,7 @@ def cp_shard_model_inputs(
had_input_ids_global = hasattr(forward_batch, "input_ids_global")
input_ids_global_backup = getattr(forward_batch, "input_ids_global", None)
if complete_input_ids is not None:
forward_batch.input_ids_global = cp_round_robin_input_ids_v2(
forward_batch.input_ids_global = cp_interleave_input_ids(
complete_input_ids, forward_batch
)
@@ -330,14 +321,13 @@ __all__ = [
"InterleaveContextParallelMetadata",
"ZigzagCPStrategy",
"ZigzagContextParallelMetadata",
"enable_cp_v2",
"get_cp_strategy",
"is_cp_v2_active",
"is_mla_prefill_cp_enabled",
"mla_use_prefill_cp",
"is_cp_active",
"is_mla_cp_enabled",
"is_mla_cp_active",
"cp_gather_after_forward",
"cp_materialize_global_token_order",
"cp_round_robin_input_ids_v2",
"cp_interleave_input_ids",
"cp_shard_hidden_states",
"cp_shard_model_inputs",
"cp_shard_position_ids",
python/sglang/srt/layers/dcp/metadata.py
+2 -2
View File
@@ -23,10 +23,10 @@ import torch
# NOTE: This is intentionally a standalone dataclass, NOT a subclass of
# layers.cp.base.BaseContextParallelMetadata. It is preserved verbatim from #14194
# and is stored in its own ForwardBatch field (attn_dcp_metadata), separate from the
# prefill-CP attn_cp_metadata, so it never participates in the CP-v2 build_metadata
# prefill-CP attn_cp_metadata, so it never participates in the CP build_metadata
# contract today. Whether the decode metadata should re-parent onto
# BaseContextParallelMetadata is deferred to P2 (DecodeContextParallelStrategy); decide
# it there rather than coupling this relocation to the CP-v2 ABC.
# it there rather than coupling this relocation to the CP ABC.
@dataclass
class DecodeContextParallelMetadata:
# For decode context parallel
python/sglang/srt/model_executor/forward_batch_info.py
+2 -2
View File
@@ -305,12 +305,12 @@ def compute_local_num_token_non_padded_cpu(
def prefill_graph_tolerates_sum_len() -> bool:
"""Whether MegaMoE may replay prefill graphs with local shapes."""
from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
from sglang.srt.layers.cp.utils import is_mla_prefill_cp_enabled
from sglang.srt.layers.cp.utils import is_mla_cp_enabled
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
if not get_moe_a2a_backend().is_megamoe():
return False
return not (is_dsa_enable_prefill_cp() or is_mla_prefill_cp_enabled())
return not (is_dsa_enable_prefill_cp() or is_mla_cp_enabled())
@dataclass
python/sglang/srt/model_executor/model_runner.py
+5 -5
View File
@@ -76,8 +76,8 @@ from sglang.srt.layers import deep_gemm_wrapper, model_parallel
from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
from sglang.srt.layers.cp.utils import (
get_cp_strategy,
is_cp_v2_active,
is_mla_prefill_cp_enabled,
is_cp_active,
is_mla_cp_enabled,
)
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.sampler import create_sampler
@@ -266,8 +266,8 @@ def _prefill_cuda_graph_allows_context_parallel(
) -> bool:
"""Allow CP only through a runner that captured the validated CP body."""
return get_cp_strategy() is None or (
bool(getattr(prefill_runner, "enable_cp_v2_bcg_capture", False))
and is_cp_v2_active(forward_batch)
bool(getattr(prefill_runner, "enable_cp_bcg_capture", False))
and is_cp_active(forward_batch)
)
@@ -1544,7 +1544,7 @@ class ModelRunner:
sharded=(
forward_batch.attn_tp_sequence_sharded
and not is_dsa_enable_prefill_cp()
and not is_mla_prefill_cp_enabled()
and not is_mla_cp_enabled()
),
)
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+2 -4
View File
@@ -49,7 +49,7 @@ from sglang.srt.layers.attention.base_attn_backend import (
SharedReadEnds,
)
from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
from sglang.srt.layers.cp.utils import is_mla_prefill_cp_enabled
from sglang.srt.layers.cp.utils import is_mla_cp_enabled
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
set_dp_buffer_len,
@@ -293,9 +293,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
# #18233) uses the plain LayerCommunicator with an attn_tp-replicated
# layout and is intentionally excluded so the attn_tp-local
# num_token_non_padded adjustment still runs for it.
self.enable_prefill_cp = (
is_dsa_enable_prefill_cp() or is_mla_prefill_cp_enabled()
)
self.enable_prefill_cp = is_dsa_enable_prefill_cp() or is_mla_cp_enabled()
self.deepep_adapter = DeepEPCudaGraphRunnerAdapter()
python/sglang/srt/model_executor/runner/eager_runner.py
+9 -9
View File
@@ -28,7 +28,7 @@ from sglang.srt.layers.cp.utils import (
cp_gather_after_forward,
cp_shard_model_inputs,
get_cp_strategy,
is_cp_v2_active,
is_cp_active,
prepare_cp_forward,
)
from sglang.srt.layers.pooler import EmbeddingPoolerOutput
@@ -281,8 +281,8 @@ class EagerRunner(BaseRunner):
if not self.enable_pdmux:
forward_batch = self.load_batch(forward_batch, pp_proxy_tensors)
cp_v2_active = is_cp_v2_active(forward_batch)
if cp_v2_active:
cp_active = is_cp_active(forward_batch)
if cp_active:
prepare_cp_forward(forward_batch)
# Target verify can arrive with ``forward_metadata_ready`` set by an
@@ -293,7 +293,7 @@ class EagerRunner(BaseRunner):
# directly from the live ``spec_info`` tensors.
if (
forward_batch.needs_forward_metadata_init()
or cp_v2_active
or cp_active
or forward_batch.forward_mode.is_target_verify()
):
if model_runner.ps.attn_dcp_size > 1 and hasattr(
@@ -330,7 +330,7 @@ class EagerRunner(BaseRunner):
torch.get_device_module(model_runner.device),
)
if not cp_v2_active:
if not cp_active:
forward_batch.attn_cp_metadata = None
category = (
@@ -344,7 +344,7 @@ class EagerRunner(BaseRunner):
_is_hip
and pcg_runner is not None
and not isinstance(pcg_runner, EagerRunner)
and not cp_v2_active
and not cp_active
):
# HIP PCG eager fallback: enter the PCG context so Dynamo guards
# and PCG-specific MoE/attention paths stay consistent.
@@ -366,8 +366,8 @@ class EagerRunner(BaseRunner):
forward_batch,
**kwargs,
)
elif cp_v2_active:
ret = self._execute_extend_cp_v2(forward_batch, kwargs)
elif cp_active:
ret = self._execute_extend_cp(forward_batch, kwargs)
else:
ret = model_runner.model.forward(
forward_batch.input_ids,
@@ -377,7 +377,7 @@ class EagerRunner(BaseRunner):
)
return ret
def _execute_extend_cp_v2(
def _execute_extend_cp(
self, forward_batch: ForwardBatch, kwargs: dict
) -> Union[LogitsProcessorOutput, PPProxyTensors]:
"""CP extend: shard inputs at the model boundary, run the body on the
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+13 -15
View File
@@ -58,15 +58,15 @@ from sglang.srt.layers.cp.bcg import (
PrefillCPBCGInput,
)
from sglang.srt.layers.cp.bcg import (
enable_cp_v2_bcg_capture as should_enable_cp_v2_bcg_capture,
enable_cp_bcg_capture as should_enable_cp_bcg_capture,
)
from sglang.srt.layers.cp.bcg import (
execute_prefill_cp_bcg,
filter_prefill_cp_bcg_capture_num_tokens,
)
from sglang.srt.layers.cp.utils import (
is_cp_v2_active,
is_mla_prefill_cp_enabled,
is_cp_active,
is_mla_cp_enabled,
)
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
@@ -380,9 +380,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
or self.prefill_backend_name == Backend.FULL
),
require_gathered_buffer=require_gathered_buffer(),
enable_prefill_cp=(
is_dsa_enable_prefill_cp() or is_mla_prefill_cp_enabled()
),
enable_prefill_cp=(is_dsa_enable_prefill_cp() or is_mla_cp_enabled()),
attn_tp_sharded_fn=self.model_runner.attn_tp_sequence_sharded,
source=self.buffers,
)
@@ -415,7 +413,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
self._is_full_backend = False
# Same ordering requirement: capture_prepare reads this.
self._capture_lora = False
self.enable_cp_v2_bcg_capture = False
self.enable_cp_bcg_capture = False
self.prefill_cp_bcg_input: Optional[PrefillCPBCGInput] = None
# TcPiecewise does its compile pass during backend construction.
# Wrap only that path with the prefill CUDA graph failure hint.
@@ -513,10 +511,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
}
server_args = model_runner.server_args
self.enable_cp_v2_bcg_capture = isinstance(
self.enable_cp_bcg_capture = isinstance(
self.backend, BreakableCudaGraphBackend
) and should_enable_cp_v2_bcg_capture(server_args)
if self.enable_cp_v2_bcg_capture:
) and should_enable_cp_bcg_capture(server_args)
if self.enable_cp_bcg_capture:
self.capture_num_tokens = filter_prefill_cp_bcg_capture_num_tokens(
self.capture_num_tokens, server_args
)
@@ -1232,7 +1230,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
),
):
return False
if getattr(self, "enable_cp_v2_bcg_capture", False) and is_cp_v2_active(
if getattr(self, "enable_cp_bcg_capture", False) and is_cp_active(
forward_batch
):
assert self.prefill_cp_bcg_input is not None
@@ -1454,7 +1452,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
"""
num_tokens = size
forward_batch, attn_backend = self.capture_prepare(num_tokens)
if self.enable_cp_v2_bcg_capture:
if self.enable_cp_bcg_capture:
assert self.prefill_cp_bcg_input is not None
self.prefill_cp_bcg_input.prepare(
self,
@@ -1532,7 +1530,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
"""
num_tokens = len(forward_batch.input_ids)
static_num_tokens = self._pad_to_bucket(num_tokens, self.capture_num_tokens)
if getattr(self, "enable_cp_v2_bcg_capture", False) and is_cp_v2_active(
if getattr(self, "enable_cp_bcg_capture", False) and is_cp_active(
forward_batch
):
assert self.prefill_cp_bcg_input is not None
@@ -1746,7 +1744,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
)
metadata_forward_batch = forward_batch
if self.enable_cp_v2_bcg_capture:
if self.enable_cp_bcg_capture:
assert self.prefill_cp_bcg_input is not None
self.prefill_cp_bcg_input.prepare(
self,
@@ -1914,7 +1912,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
self.model_runner, forward_batch, self.device_module
)
if self.enable_cp_v2_bcg_capture:
if self.enable_cp_bcg_capture:
output = execute_prefill_cp_bcg(
self,
forward_batch,
python/sglang/srt/models/deepseek_common/attention_backend_handler.py
+2 -2
View File
@@ -1,6 +1,6 @@
from sglang.srt.environ import envs
from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
from sglang.srt.layers.cp.utils import is_cp_v2_active
from sglang.srt.layers.cp.utils import is_cp_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,
@@ -114,7 +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.
if is_cp_v2_active(forward_batch):
if is_cp_active(forward_batch):
return _dispatch_mla_subtype(attn, forward_batch)
sum_extend_prefix_lens = _get_sum_extend_prefix_lens(forward_batch)
python/sglang/srt/models/utils.py
+2 -2
View File
@@ -30,7 +30,7 @@ from sglang.kernels.ops.layernorm.norm import (
fused_inplace_qknorm,
)
from sglang.srt.environ import envs
from sglang.srt.layers.cp.utils import is_cp_v2_active
from sglang.srt.layers.cp.utils import is_cp_active
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -296,7 +296,7 @@ def enable_fused_set_kv_buffer(forward_batch: ForwardBatch):
_is_cuda
and pool.dtype == torch.bfloat16
and not isinstance(pool, SWAKVPool)
and not is_cp_v2_active(forward_batch)
and not is_cp_active(forward_batch)
and getattr(forward_batch, "dcp_kv_mask", None) is None
) or (_is_hip and getattr(forward_batch, "dcp_kv_mask", None) is None)
test/registered/cp/test_cp_strategy_unit.py
+11 -23
View File
@@ -24,8 +24,7 @@ from sglang.srt.layers.cp.padding import (
)
from sglang.srt.layers.cp.utils import (
cp_split_before_forward,
enable_cp_v2,
is_cp_v2_active,
is_cp_active,
prepare_cp_forward,
)
from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy
@@ -88,7 +87,7 @@ class TestCPStrategyUnit(CustomTestCase):
self.assertFalse(is_interleave())
self.assertEqual(get_cp_strategy_kind(), ContextParallelStrategyKind.ZIGZAG)
def test_get_cp_strategy_is_initialized_under_cp_v2(self):
def test_get_cp_strategy_is_initialized_under_cp(self):
init_cp_strategy(
enable_prefill_cp=True,
cp_size=4,
@@ -122,16 +121,6 @@ class TestCPStrategyUnit(CustomTestCase):
):
self.assertFalse(is_dsa_enable_prefill_cp())
@patch("sglang.srt.utils.is_npu", return_value=False)
@patch("sglang.srt.utils.is_hip", return_value=True)
def test_hip_keeps_strategy_cp_disabled(self, _mock_is_hip, _mock_is_npu):
self.assertFalse(enable_cp_v2())
@patch("sglang.srt.utils.is_npu", return_value=True)
@patch("sglang.srt.utils.is_hip", return_value=False)
def test_npu_keeps_strategy_cp_disabled(self, _mock_is_hip, _mock_is_npu):
self.assertFalse(enable_cp_v2())
class TestPrefillCPBCGReplay(CustomTestCase):
def tearDown(self):
@@ -147,7 +136,7 @@ class TestPrefillCPBCGReplay(CustomTestCase):
runner.capture_hidden_mode = CaptureHiddenMode.NULL
runner.capture_num_tokens = [2048, 2304]
runner.max_num_tokens = 2304
runner.enable_cp_v2_bcg_capture = True
runner.enable_cp_bcg_capture = True
return runner
def _make_forward_batch(self):
@@ -305,7 +294,7 @@ class TestCPZigzagStrategy(CustomTestCase):
attn_cp_metadata=metadata,
)
def test_enable_cp_v2_and_is_cp_v2_active(self):
def test_is_cp_active(self):
active_batch = SimpleNamespace(
input_ids=torch.arange(8),
forward_mode=_ExtendMode(),
@@ -318,9 +307,8 @@ class TestCPZigzagStrategy(CustomTestCase):
)
with patch.dict("os.environ", {"SGLANG_ENABLE_CP_V2": "0"}):
self.assertTrue(enable_cp_v2())
self.assertTrue(is_cp_v2_active(active_batch))
self.assertFalse(is_cp_v2_active(inactive_batch))
self.assertTrue(is_cp_active(active_batch))
self.assertFalse(is_cp_active(inactive_batch))
def _expected_metadata(self, *, rank, cp_size, seq_lens, extend_seq_lens):
bs = len(extend_seq_lens)
@@ -832,7 +820,7 @@ class TestCPInterleaveStrategy(CustomTestCase):
def _rank_tensors(self, x, *, cp_size, seq_lens, extend_seq_lens):
per_rank = []
metas = []
with self._patch_legacy_round_robin_mode():
with self._patch_legacy_interleave_mode():
for rank in range(cp_size):
metadata = self._metadata_for_rank(
rank,
@@ -848,9 +836,9 @@ class TestCPInterleaveStrategy(CustomTestCase):
return metas, per_rank
@contextmanager
def _patch_legacy_round_robin_mode(self):
def _patch_legacy_interleave_mode(self):
with patch(
"sglang.srt.layers.attention.dsa.utils.is_dsa_prefill_cp_round_robin_split",
"sglang.srt.layers.attention.dsa.utils.is_dsa_prefill_cp_interleave",
return_value=True,
):
yield
@@ -875,7 +863,7 @@ class TestCPInterleaveStrategy(CustomTestCase):
return_value=torch.no_grad(),
),
patch(
"sglang.srt.layers.attention.dsa.utils.is_dsa_prefill_cp_round_robin_split",
"sglang.srt.layers.attention.dsa.utils.is_dsa_prefill_cp_interleave",
return_value=True,
),
)
@@ -960,7 +948,7 @@ class TestCPInterleaveStrategy(CustomTestCase):
attn_cp_rank=rank,
attn_cp_size=cp_size,
),
self._patch_legacy_round_robin_mode(),
self._patch_legacy_interleave_mode(),
):
local_x = strategy.shard_hidden_states(x, fb)
local_positions = strategy.shard_position_ids(positions, fb)
test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py
+1 -1
View File
@@ -203,7 +203,7 @@ class TestDSV4FlashFP4B200_CP_DSpark(
GSM8KMixin,
CustomTestCase,
):
"""DSPARK speculation + prefill CP (interleave, CP_V2, attn_cp=tp)."""
"""DSPARK speculation + prefill CP (interleave, CP, attn_cp=tp)."""
gsm8k_accuracy_thres = 0.90
test/registered/cp/test_dsa_prefill_cp.py
+2 -2
View File
@@ -18,7 +18,7 @@ GLM52_MODEL_PATH = "zai-org/GLM-5.2-FP8"
SERVER_LAUNCH_TIMEOUT = max(DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, 1800)
class TestDSACPV2Interleave(CustomTestCase):
class TestDSACPInterleave(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = GLM52_MODEL_PATH
@@ -78,7 +78,7 @@ class TestDSACPV2Interleave(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f'### test_a_gsm8k (dsa-cp-v2-interleave)\n{metrics["score"]=:.3f}\n'
f'### test_a_gsm8k (dsa-cp-interleave)\n{metrics["score"]=:.3f}\n'
)
self.assertGreater(metrics["score"], 0.935)
test/registered/kernels/ops/layernorm/test_mhc_kernels.py
+1 -1
View File
@@ -19,7 +19,7 @@ def test_mhc_fused_post_pre_matches_unfused(
if not torch.cuda.is_available():
pytest.skip("CUDA is required for TileLang mHC kernels")
monkeypatch.setattr(mhc, "is_dsa_prefill_cp_round_robin_split", lambda: False)
monkeypatch.setattr(mhc, "is_dsa_prefill_cp_interleave", lambda: False)
# This is a single-process kernel unit test with no TP group initialized.
# mhc_pre / mhc_fused_post_pre allocate the MoE input in the symmetric-memory
# pool via use_symmetric_memory(get_tp_group(), ...); bypass that path so the
test/registered/unit/layers/test_layer_scatter_modes_cp_dense_mlp.py
+1 -1
View File
@@ -18,7 +18,7 @@ class TestDenseMlpScatterModeUnderPrefillCP(CustomTestCase):
with (
patch.object(comm, "_generic_prefill_cp_shards_tokens", return_value=True),
patch.object(comm, "is_dsa_enable_prefill_cp", return_value=False),
patch.object(comm, "is_mla_prefill_cp_enabled", return_value=False),
patch.object(comm, "is_mla_cp_enabled", return_value=False),
patch.object(comm, "enable_moe_dense_fully_dp", return_value=False),
):
modes = LayerScatterModes.init_new(
test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py
+1 -1
View File
@@ -213,7 +213,7 @@ class TestPrefillCudaGraphRunnerChunkedPrefix(CustomTestCase):
runner.capture_num_tokens = [4]
runner.buffer_registry = _FakeBatchRegistry()
runner.model_runner = SimpleNamespace(attn_tp_sequence_sharded=lambda _: False)
runner.enable_cp_v2_bcg_capture = False
runner.enable_cp_bcg_capture = False
runner._is_full_backend = False
runner.backend = SimpleNamespace()
runner.has_mha_companion_layers = False
test/registered/unit/models/test_deepseek_mla_dispatch.py
-2
View File
@@ -15,7 +15,6 @@ from types import SimpleNamespace
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.model_executor.forward_batch_info import ForwardMode
from sglang.srt.models.deepseek_common import attention_backend_handler as abh
@@ -111,7 +110,6 @@ class TestCPMLADispatch(CustomTestCase):
)
with (
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)),
):
for prefix in (0, 32):