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