From ed183d45acfb0f6d1a2b1cd6b6a34a579e1a3ad7 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Tue, 8 Sep 2026 16:03:36 -0700 Subject: [PATCH] [CP V1 Deprecation 4/5] Canonicalize prefill CP API names (#36229) --- .../sglang/kernels/ops/attention/__init__.py | 2 +- .../kernels/ops/attention/dsa/cp_split.py | 4 +- .../ops/attention/dsv4/metadata_kernel.py | 2 +- python/sglang/kernels/ops/layernorm/mhc.py | 4 +- .../layers/attention/deepseek_v4_backend.py | 16 +++--- .../srt/layers/attention/dsa/dsa_indexer.py | 6 +- .../sglang/srt/layers/attention/dsa/utils.py | 20 ++++--- .../srt/layers/attention/dsa_backend.py | 10 ++-- .../attention/flashattention_backend.py | 4 +- .../srt/layers/attention/index_topk_share.py | 4 +- .../layers/attention/trtllm_mha_backend.py | 20 +++---- python/sglang/srt/layers/communicator.py | 17 +++--- .../sglang/srt/layers/communicator_dsa_cp.py | 6 +- python/sglang/srt/layers/cp/base.py | 4 +- python/sglang/srt/layers/cp/bcg.py | 7 +-- .../sglang/srt/layers/cp/cp_decode_attn_tp.py | 4 +- python/sglang/srt/layers/cp/interleave.py | 4 +- python/sglang/srt/layers/cp/utils.py | 56 ++++++++----------- python/sglang/srt/layers/dcp/metadata.py | 4 +- .../srt/model_executor/forward_batch_info.py | 4 +- .../sglang/srt/model_executor/model_runner.py | 10 ++-- .../runner/decode_cuda_graph_runner.py | 6 +- .../srt/model_executor/runner/eager_runner.py | 18 +++--- .../runner/prefill_cuda_graph_runner.py | 28 +++++----- .../attention_backend_handler.py | 4 +- python/sglang/srt/models/utils.py | 4 +- test/registered/cp/test_cp_strategy_unit.py | 34 ++++------- .../cp/test_deepseek_v4_flash_fp4_b200_cp.py | 2 +- test/registered/cp/test_dsa_prefill_cp.py | 4 +- .../kernels/ops/layernorm/test_mhc_kernels.py | 2 +- .../test_layer_scatter_modes_cp_dense_mlp.py | 2 +- .../test_prefill_cuda_graph_runner.py | 2 +- .../unit/models/test_deepseek_mla_dispatch.py | 2 - 33 files changed, 143 insertions(+), 173 deletions(-) diff --git a/python/sglang/kernels/ops/attention/__init__.py b/python/sglang/kernels/ops/attention/__init__.py index 42e51effc..3c9e6d30b 100644 --- a/python/sglang/kernels/ops/attention/__init__.py +++ b/python/sglang/kernels/ops/attention/__init__.py @@ -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"), diff --git a/python/sglang/kernels/ops/attention/dsa/cp_split.py b/python/sglang/kernels/ops/attention/dsa/cp_split.py index d20f3aca3..d2b181dff 100644 --- a/python/sglang/kernels/ops/attention/dsa/cp_split.py +++ b/python/sglang/kernels/ops/attention/dsa/cp_split.py @@ -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, diff --git a/python/sglang/kernels/ops/attention/dsv4/metadata_kernel.py b/python/sglang/kernels/ops/attention/dsv4/metadata_kernel.py index d477a0238..9025187f9 100644 --- a/python/sglang/kernels/ops/attention/dsv4/metadata_kernel.py +++ b/python/sglang/kernels/ops/attention/dsv4/metadata_kernel.py @@ -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] diff --git a/python/sglang/kernels/ops/layernorm/mhc.py b/python/sglang/kernels/ops/layernorm/mhc.py index 4784ebf0d..42446c837 100644 --- a/python/sglang/kernels/ops/layernorm/mhc.py +++ b/python/sglang/kernels/ops/layernorm/mhc.py @@ -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) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index b398df93f..5827504d4 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -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 diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index 285f66342..6678e4b3a 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -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 diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py index d31d8e6f9..db99b5b20 100644 --- a/python/sglang/srt/layers/attention/dsa/utils.py +++ b/python/sglang/srt/layers/attention/dsa/utils.py @@ -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 diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 22c35d366..0a090c271 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -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) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 1b57ecc7a..79d100be9 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -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 diff --git a/python/sglang/srt/layers/attention/index_topk_share.py b/python/sglang/srt/layers/attention/index_topk_share.py index 3731a9f2b..7b56b487c 100644 --- a/python/sglang/srt/layers/attention/index_topk_share.py +++ b/python/sglang/srt/layers/attention/index_topk_share.py @@ -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 diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 357239c61..59291270f 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -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( diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index fcdeb296c..9e07ac3e4 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -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 diff --git a/python/sglang/srt/layers/communicator_dsa_cp.py b/python/sglang/srt/layers/communicator_dsa_cp.py index 26e91eff3..b4696e499 100644 --- a/python/sglang/srt/layers/communicator_dsa_cp.py +++ b/python/sglang/srt/layers/communicator_dsa_cp.py @@ -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 diff --git a/python/sglang/srt/layers/cp/base.py b/python/sglang/srt/layers/cp/base.py index 6328151fa..111f47781 100644 --- a/python/sglang/srt/layers/cp/base.py +++ b/python/sglang/srt/layers/cp/base.py @@ -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) diff --git a/python/sglang/srt/layers/cp/bcg.py b/python/sglang/srt/layers/cp/bcg.py index 1adabb21d..6631c9d48 100644 --- a/python/sglang/srt/layers/cp/bcg.py +++ b/python/sglang/srt/layers/cp/bcg.py @@ -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( diff --git a/python/sglang/srt/layers/cp/cp_decode_attn_tp.py b/python/sglang/srt/layers/cp/cp_decode_attn_tp.py index b894bfa40..d4b57aafd 100644 --- a/python/sglang/srt/layers/cp/cp_decode_attn_tp.py +++ b/python/sglang/srt/layers/cp/cp_decode_attn_tp.py @@ -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) diff --git a/python/sglang/srt/layers/cp/interleave.py b/python/sglang/srt/layers/cp/interleave.py index 60c6f2e38..157a4ce7d 100644 --- a/python/sglang/srt/layers/cp/interleave.py +++ b/python/sglang/srt/layers/cp/interleave.py @@ -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 diff --git a/python/sglang/srt/layers/cp/utils.py b/python/sglang/srt/layers/cp/utils.py index 82da04255..d234987af 100644 --- a/python/sglang/srt/layers/cp/utils.py +++ b/python/sglang/srt/layers/cp/utils.py @@ -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", diff --git a/python/sglang/srt/layers/dcp/metadata.py b/python/sglang/srt/layers/dcp/metadata.py index f9614148e..9ec2ff33e 100644 --- a/python/sglang/srt/layers/dcp/metadata.py +++ b/python/sglang/srt/layers/dcp/metadata.py @@ -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 diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index a9a2aa8df..84c449422 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -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 diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index b843fe27a..7b5b75fb7 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -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() ), ) diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index cef7ac162..e7abb66e0 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -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() diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index 1bc6b3e8d..54a01852d 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -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 diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 6bb8f30fb..4ef9f69ac 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -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, diff --git a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py index ca677aecc..8543c3fdf 100644 --- a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py +++ b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py @@ -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) diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index 5bca8eb81..4c8ade007 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -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) diff --git a/test/registered/cp/test_cp_strategy_unit.py b/test/registered/cp/test_cp_strategy_unit.py index 644d2dc19..1bc6a537d 100644 --- a/test/registered/cp/test_cp_strategy_unit.py +++ b/test/registered/cp/test_cp_strategy_unit.py @@ -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) diff --git a/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py b/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py index 33e676a5c..812053168 100644 --- a/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py +++ b/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py @@ -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 diff --git a/test/registered/cp/test_dsa_prefill_cp.py b/test/registered/cp/test_dsa_prefill_cp.py index 9dace61a8..7660ca29b 100644 --- a/test/registered/cp/test_dsa_prefill_cp.py +++ b/test/registered/cp/test_dsa_prefill_cp.py @@ -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) diff --git a/test/registered/kernels/ops/layernorm/test_mhc_kernels.py b/test/registered/kernels/ops/layernorm/test_mhc_kernels.py index f284cf84d..3311b34f2 100644 --- a/test/registered/kernels/ops/layernorm/test_mhc_kernels.py +++ b/test/registered/kernels/ops/layernorm/test_mhc_kernels.py @@ -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 diff --git a/test/registered/unit/layers/test_layer_scatter_modes_cp_dense_mlp.py b/test/registered/unit/layers/test_layer_scatter_modes_cp_dense_mlp.py index b725e75bc..3a112232e 100644 --- a/test/registered/unit/layers/test_layer_scatter_modes_cp_dense_mlp.py +++ b/test/registered/unit/layers/test_layer_scatter_modes_cp_dense_mlp.py @@ -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( diff --git a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py index 07a2c5f05..613656dc1 100644 --- a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py +++ b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner.py @@ -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 diff --git a/test/registered/unit/models/test_deepseek_mla_dispatch.py b/test/registered/unit/models/test_deepseek_mla_dispatch.py index 560d11088..dbd144857 100644 --- a/test/registered/unit/models/test_deepseek_mla_dispatch.py +++ b/test/registered/unit/models/test_deepseek_mla_dispatch.py @@ -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):