From b6c31b155c2b788c7c0e56ab91e03d9016b73356 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Sun, 6 Sep 2026 21:53:53 -0700 Subject: [PATCH] [CP V1 Deprecation 3/5] Remove generic prefill CP v1 runtime (#36228) --- .../sglang/srt/arg_groups/deepseek_v4_hook.py | 2 +- python/sglang/srt/arg_groups/kv_cache_hook.py | 6 +- python/sglang/srt/arg_groups/parallel_hook.py | 11 +- python/sglang/srt/arg_groups/pipeline.py | 3 +- .../srt/layers/attention/dsa/dsa_indexer.py | 240 ++---------------- .../sglang/srt/layers/attention/dsa/utils.py | 11 +- .../srt/layers/attention/dsa_backend.py | 42 ++- .../attention/flashattention_backend.py | 105 +++----- .../srt/layers/attention/index_topk_share.py | 15 +- python/sglang/srt/layers/communicator.py | 8 +- .../sglang/srt/layers/communicator_dsa_cp.py | 10 +- python/sglang/srt/layers/cp/padding.py | 5 +- python/sglang/srt/layers/cp/utils.py | 50 +++- python/sglang/srt/layers/utils/cp_utils.py | 53 +--- python/sglang/srt/managers/schedule_policy.py | 10 - .../srt/model_executor/forward_batch_info.py | 6 +- .../sglang/srt/model_executor/model_runner.py | 4 +- .../runner/decode_cuda_graph_runner.py | 2 +- .../srt/model_executor/runner/eager_runner.py | 11 +- .../runner/prefill_cuda_graph_runner.py | 6 +- .../attention_backend_handler.py | 11 +- .../attention_forward_methods/forward_mla.py | 42 +-- python/sglang/srt/models/deepseek_nextn.py | 73 ++---- python/sglang/srt/models/deepseek_v2.py | 184 +++++--------- python/sglang/srt/models/deepseek_v4.py | 31 +-- python/sglang/srt/models/deepseek_v4_nextn.py | 23 +- python/sglang/srt/models/qwen2_moe.py | 6 +- python/sglang/srt/models/qwen3_moe.py | 4 +- python/sglang/srt/models/utils.py | 8 +- .../attention_methods/dsa_attention.py | 3 +- test/registered/cp/test_cp_strategy_unit.py | 12 +- .../layers/attention/test_index_topk_share.py | 5 +- .../unit/models/test_deepseek_mla_dispatch.py | 55 ++++ .../unit/server_args/test_server_args.py | 84 +++--- 34 files changed, 400 insertions(+), 741 deletions(-) diff --git a/python/sglang/srt/arg_groups/deepseek_v4_hook.py b/python/sglang/srt/arg_groups/deepseek_v4_hook.py index 42f515ca4..ce8c7d5d1 100644 --- a/python/sglang/srt/arg_groups/deepseek_v4_hook.py +++ b/python/sglang/srt/arg_groups/deepseek_v4_hook.py @@ -168,7 +168,7 @@ 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: + 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( diff --git a/python/sglang/srt/arg_groups/kv_cache_hook.py b/python/sglang/srt/arg_groups/kv_cache_hook.py index 2a1093325..157a4fde5 100644 --- a/python/sglang/srt/arg_groups/kv_cache_hook.py +++ b/python/sglang/srt/arg_groups/kv_cache_hook.py @@ -487,9 +487,9 @@ def validate_prefill_only_disable_kv_cache_args(server_args: Any): "radix cache indexes KV pool slots that no longer hold real data." ) - # Context-parallel prefill stages K/V through cp_allgather_and_save_kv_cache, - # which writes to the pool via set_kv_buffer. NoOpMHATokenToKVPool intentionally - # raises on writes, so the engine would boot fine but fail on the first request. + # Context-parallel prefill writes K/V to the pool via set_kv_buffer. + # NoOpMHATokenToKVPool intentionally raises on writes, so the engine would + # boot fine but fail on the first request. if resolved_view(server_args).attn_cp_size > 1: raise ValueError( "--prefill-only-disable-kv-cache is incompatible with --attn-cp-size > 1: " diff --git a/python/sglang/srt/arg_groups/parallel_hook.py b/python/sglang/srt/arg_groups/parallel_hook.py index c293c193d..6a6598de8 100644 --- a/python/sglang/srt/arg_groups/parallel_hook.py +++ b/python/sglang/srt/arg_groups/parallel_hook.py @@ -39,7 +39,7 @@ def handle_context_parallelism(server_args: Any): cfg.enable_prefill_cp and model_arch == "DeepseekV32ForCausalLM" and cfg.cp_strategy == "zigzag" - and not (platform.is_hip or platform.is_npu) + 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 " @@ -563,7 +563,7 @@ def handle_eplb_and_dispatch(server_args: Any): 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 + 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 @@ -571,7 +571,7 @@ def handle_platform_cp_compatibility(server_args: Any): ): raise ValueError( "Legacy prefill context-parallel options are supported only " - "by protected HIP or Ascend NPU paths. Use " + "by protected HIP, Ascend NPU, or MUSA paths. Use " "--enable-prefill-cp with --cp-strategy." ) return @@ -603,7 +603,10 @@ def handle_platform_cp_compatibility(server_args: Any): def handle_legacy_cp_runtime_compatibility(server_args: Any): - """Project canonical CP settings for runtime consumers removed by PR3.""" + """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: diff --git a/python/sglang/srt/arg_groups/pipeline.py b/python/sglang/srt/arg_groups/pipeline.py index 74ed99b05..daef9af83 100644 --- a/python/sglang/srt/arg_groups/pipeline.py +++ b/python/sglang/srt/arg_groups/pipeline.py @@ -287,8 +287,7 @@ def run_resolution_pipeline(server_args: Any) -> None: # Normalize load balancing defaults. handle_load_balance_method(server_args) - # The old runtime distinguishes DSA from other CP paths through legacy - # fields, so project only after attention_backend has been resolved. + # Protected runtimes still consume platform CP fields after backend selection. handle_legacy_cp_runtime_compatibility(server_args) # Handle context parallelism. diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index b53130894..6718a37a8 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -2,7 +2,7 @@ from __future__ import annotations import contextlib import logging -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union import torch from einops import rearrange @@ -29,7 +29,6 @@ from sglang.srt.layers.attention.dsa.paged_mqa_logits_backend import ( from sglang.srt.layers.attention.dsa.utils import ( aiter_can_use_preshuffle_paged_mqa, is_dsa_enable_prefill_cp, - is_dsa_prefill_cp_in_seq_split, is_graph_dsa_split_op_surface, ) from sglang.srt.layers.layernorm import LayerNorm, RMSNorm @@ -1418,167 +1417,6 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp): return None return raw_topk_result - def _get_topk_ragged_with_cp( - self, - forward_batch: ForwardBatch, - layer_id: int, - q_fp8: torch.Tensor, - weights: torch.Tensor, - metadata: BaseIndexerMetadata, - kv_len: int, - actual_seq_q: int, - cp_index: List[Tuple[int, int, int]] = None, - ) -> torch.Tensor: - assert not _is_in_piecewise_or_breakable_cuda_graph(), ( - "DSA context parallel (_get_topk_ragged_with_cp) not supported under " - "piecewise/breakable CUDA graph" - ) - if TYPE_CHECKING: - assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool) - - page_size = get_token_to_kv_pool().page_size - if _is_xpu: - assert page_size in ( - 64, - 128, - ), f"XPU DSA requires page_size 64 or 128, got {page_size}" - else: - assert page_size == 64, "only support page size 64" - assert len(weights.shape) == 3 - weights = weights.squeeze(-1) - k_fp8_list = [] - k_scale_list = [] - ks_list = [] - ke_offset_list = [] - offset = 0 - actual_seq_q_list = [] - batch_idx_list = [] - - block_tables = metadata.get_page_table_64() - - assert ( - forward_batch.seq_lens_cpu is not None - and forward_batch.extend_seq_lens_cpu is not None - ) - if cp_index is not None: - # TODO Multi-batch support has accuracy issues - for batch_idx, start_seq_position, end_seq_position in cp_index: - pre_chunk_offset = ( - forward_batch.seq_lens_cpu[batch_idx].item() - - forward_batch.extend_seq_lens_cpu[batch_idx] - ) - start_seq_position += pre_chunk_offset - end_seq_position += pre_chunk_offset - if offset == 0 and batch_idx != 0: - offset += forward_batch.extend_seq_lens_cpu[batch_idx - 1] - k_fp8 = get_token_to_kv_pool().get_index_k_continuous( - layer_id, - end_seq_position, - block_tables[batch_idx], - ) - k_scale = get_token_to_kv_pool().get_index_k_scale_continuous( - layer_id, - end_seq_position, - block_tables[batch_idx], - ) - - extend_seq_len = end_seq_position - start_seq_position - ks = torch.full( - (extend_seq_len,), offset, dtype=torch.int32, device="cuda" - ) - k_fp8_list.append(k_fp8) - k_scale_list.append(k_scale) - ks_list.append(ks) - ke_offset = torch.arange( - start_seq_position + 1, - end_seq_position + 1, - dtype=torch.int32, - device="cuda", - ) - ke_offset_list.append(ke_offset) - actual_seq_q = torch.tensor( - [extend_seq_len], dtype=torch.int32, device="cuda" - ) - actual_seq_q_list.append(actual_seq_q) - batch_idx_list.append(batch_idx) - - k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn) - k_scale = torch.cat(k_scale_list, dim=0).view(torch.float32).squeeze(-1) - kv_fp8 = (k_fp8, k_scale) - ks = torch.cat(ks_list, dim=0) - ke_offset = torch.cat(ke_offset_list, dim=0) - ke = ks + ke_offset - actual_seq_q = torch.cat(actual_seq_q_list, dim=0) - with self._with_real_sm_count(): - q_padded, w_padded, _ = self._pad_heads_for_deep_gemm(q_fp8, weights) - logits = deep_gemm.fp8_mqa_logits( - q_padded, - kv_fp8, - w_padded, - ks, - ke, - clean_logits=False, - ) - topk_result = metadata.topk_transform( - logits, - self.index_topk, - ks=ks, - cu_seqlens_q=actual_seq_q, - ke_offset=ke_offset, - batch_idx_list=batch_idx_list, - ) - else: - kv_len = ( - forward_batch.seq_lens_cpu[0].item() - - forward_batch.extend_seq_lens_cpu[0] - + kv_len - ) - k_fp8 = get_token_to_kv_pool().get_index_k_continuous( - layer_id, - kv_len, - block_tables[0], - ) - k_scale = get_token_to_kv_pool().get_index_k_scale_continuous( - layer_id, - kv_len, - block_tables[0], - ) - - k_fp8 = k_fp8.view(torch.float8_e4m3fn) - k_scale = k_scale.view(torch.float32).squeeze(-1) - kv_fp8 = (k_fp8, k_scale) - ks = torch.full((actual_seq_q,), offset, dtype=torch.int32, device="cuda") - ke_offset = torch.arange( - (kv_len - actual_seq_q) + 1, - kv_len + 1, - dtype=torch.int32, - device="cuda", - ) - ke = ks + ke_offset - - with self._with_real_sm_count(): - q_padded, w_padded, _ = self._pad_heads_for_deep_gemm(q_fp8, weights) - logits = deep_gemm.fp8_mqa_logits( - q_padded, - kv_fp8, - w_padded, - ks, - ke, - clean_logits=False, - ) - actual_seq_q = torch.tensor([actual_seq_q], dtype=torch.int32).to( - device="cuda", non_blocking=True - ) - topk_result = metadata.topk_transform( - logits, - self.index_topk, - ks=ks, - cu_seqlens_q=actual_seq_q, - ke_offset=ke_offset, - ) - - return topk_result - def _store_index_k_cache( self, forward_batch: ForwardBatch, @@ -1955,67 +1793,21 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp): forward_batch, layer_id, q_fp8, weights, metadata ) else: - if ( - forward_batch.attn_cp_metadata is not None - and is_dsa_prefill_cp_in_seq_split() - ): - kv_len_prev = forward_batch.attn_cp_metadata.kv_len_prev_list[0] - kv_len_next = forward_batch.attn_cp_metadata.kv_len_next_list[0] - actual_seq_q_prev = ( - forward_batch.attn_cp_metadata.actual_seq_q_prev_list[0] - ) - actual_seq_q_next = ( - forward_batch.attn_cp_metadata.actual_seq_q_next_list[0] - ) - - # TODO support mutil-batch - # cp_batch_seq_index_prev = forward_batch.attn_cp_metadata["cp_batch_seq_index_prev"] - # cp_batch_seq_index_next = forward_batch.attn_cp_metadata["cp_batch_seq_index_next"] - # TODO prev, next, combined into a single call - q_fp8_prev, q_fp8_next = torch.split( - q_fp8, (q_fp8.shape[0] + 1) // 2, dim=0 - ) - weights_prev, weights_next = torch.split( - weights, (weights.shape[0] + 1) // 2, dim=0 - ) - topk_result_prev = self._get_topk_ragged_with_cp( - forward_batch, - layer_id, - q_fp8_prev, - weights_prev, - metadata, - kv_len_prev, - actual_seq_q_prev, - ) - - topk_result_next = self._get_topk_ragged_with_cp( - forward_batch, - layer_id, - q_fp8_next, - weights_next, - metadata, - kv_len_next, - actual_seq_q_next, - ) - topk_result = torch.cat([topk_result_prev, topk_result_next], dim=0) - topk_result = _broadcast_indexer_topk_from_rank0(topk_result) - return maybe_capture_indexer_topk(layer_id, topk_result) - else: - # In-graph (PCG/BCG) non-CP prefill is handled earlier by the - # graph DSA split-op dispatch, so only the eager path reaches - # here. - assert not in_piecewise_or_breakable_cuda_graph, ( - "Internal error: in-graph DSA prefill must go through the " - "graph DSA split-op dispatch" - ) - topk_result = self._get_topk_ragged( - enable_dual_stream, - forward_batch, - layer_id, - q_fp8, - weights, - metadata, - ) + # In-graph (PCG/BCG) non-CP prefill is handled earlier by the + # graph DSA split-op dispatch, so only the eager path reaches + # here. + assert not in_piecewise_or_breakable_cuda_graph, ( + "Internal error: in-graph DSA prefill must go through the " + "graph DSA split-op dispatch" + ) + topk_result = self._get_topk_ragged( + enable_dual_stream, + forward_batch, + layer_id, + q_fp8, + weights, + metadata, + ) else: raise NotImplementedError("DSA indexer only supports CUDA, HIP, and NPU") topk_result = _broadcast_indexer_topk_from_rank0(topk_result) diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py index 7c2c2d0f4..8e077654f 100644 --- a/python/sglang/srt/layers/attention/dsa/utils.py +++ b/python/sglang/srt/layers/attention/dsa/utils.py @@ -18,7 +18,7 @@ from sglang.srt.runtime_context import ( get_parallel, process_model_config, ) -from sglang.srt.utils import get_bool_env_var, is_cuda, is_hip, is_npu +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 @@ -115,7 +115,7 @@ 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(): + if is_hip() or is_npu() or is_musa(): return get_parallel().enable_dsa_prefill_context_parallel # Generic prefill CP derives activation from the runtime topology and model @@ -128,13 +128,6 @@ def is_dsa_enable_prefill_cp(): return is_deepseek_dsa(hf_config) or is_deepseek_v4(hf_config) -def is_dsa_prefill_cp_in_seq_split(): - return ( - is_dsa_enable_prefill_cp() - and get_parallel().dsa_prefill_cp_mode == "in-seq-split" - ) - - def is_dsa_prefill_cp_round_robin_split(): return ( 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 bb403461b..939c45f75 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -80,7 +80,6 @@ from sglang.srt.layers.attention.dsa.utils import ( dsa_cp_round_robin_split_q_seqs, dsa_use_prefill_cp, is_dsa_enable_prefill_cp, - is_dsa_prefill_cp_in_seq_split, pad_dsa_cache_seqlens, should_use_dsa_fused_topk, ) @@ -141,6 +140,31 @@ def _all_gather_dsa_trtllm_fp8_kv( return kv.split((kv_lora_rank, qk_rope_head_dim), dim=-1) +def prepare_kv_for_attention( + attn_mla, + forward_batch: ForwardBatch, + k_nope: torch.Tensor, + k_pe: torch.Tensor, + *, + defer_materialization: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + """Materialize KV needed before attention for the active layout.""" + if ( + defer_materialization + or not dsa_use_prefill_cp(forward_batch) + or not is_cp_v2_active(forward_batch) + ): + return k_nope, k_pe + 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, + ) + + def materialize_full_kv_cp( attn_mla, forward_batch: ForwardBatch, @@ -148,13 +172,18 @@ def materialize_full_kv_cp( 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): - return get_cp_strategy().materialize_full_mla_kv( + 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) @@ -3583,15 +3612,6 @@ class DeepseekSparseAttnBackend( block_tables = page_table_1.unsqueeze(1) seq_lens = metadata.cache_seqlens_int32 if seq_lens is None else seq_lens - if ( - dsa_use_prefill_cp(forward_batch) - and is_dsa_prefill_cp_in_seq_split() - and forward_batch.attn_cp_metadata is not None - ): - cp_meta = forward_batch.attn_cp_metadata - seq_chunks = list(torch.split(seq_lens, cp_meta.split_list, dim=0)) - seq_lens = torch.cat([seq_chunks[i] for i in cp_meta.zigzag_index], dim=0) - out = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla( query=q, kv_cache=kv, diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 7034956f1..8d487ede0 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -24,12 +24,8 @@ from sglang.srt.configs.model_config import AttentionArch from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask from sglang.srt.layers.cp.base import CPAttentionBackendKind, get_cp_strategy -from sglang.srt.layers.cp.utils import is_cp_v2_active +from sglang.srt.layers.cp.utils import enable_cp_v2, is_cp_v2_active from sglang.srt.layers.radix_attention import AttentionType -from sglang.srt.layers.utils.cp_utils import ( - cp_allgather_and_save_kv_cache, - cp_attn_forward_extend, -) from sglang.srt.mem_cache.memory_pool import KVWriteLoc from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode @@ -1012,7 +1008,7 @@ class FlashAttentionBackend(AttentionBackend): # (req_to_token is zero-init) and outputs for padding queries are # discarded downstream. if ( - not is_cp_v2_active(forward_batch) + 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 @@ -1266,11 +1262,7 @@ class FlashAttentionBackend(AttentionBackend): ): if score_mod is not None and self.fa_impl_ver != 4: raise RuntimeError("score_mod is only supported by the FA4 backend.") - is_cp_mode = ( - forward_batch.forward_mode.is_context_parallel_extend() - and forward_batch.attn_cp_metadata is not None - and self.attn_cp_size > 1 - ) + cp_active = is_cp_v2_active(forward_batch) if k is not None: assert v is not None @@ -1282,25 +1274,20 @@ class FlashAttentionBackend(AttentionBackend): else forward_batch.encoder_out_cache_loc ) if self.use_mla: - if is_cp_v2_active(forward_batch): - # CP-v2: k/k_rope are rank-local; the strategy gathers - # the latent to full sequence and writes it. + if cp_active: cp_strategy = get_cp_strategy() assert cp_strategy is not None cp_strategy.materialize_full_mla_kv( forward_batch, layer, k, k_rope ) else: - # CP-v1: k/k_rope arrive full-sequence (rebuild_cp_kv_cache - # ran upstream); rank-local when CP is off. out_cache_loc is - # never zigzag-split, so the write lands in the right slots. self.token_to_kv_pool.set_mla_kv_buffer( layer, cache_loc, k, k_rope, ) - elif is_cp_mode: + elif cp_active: # Dense-MHA CP: k, v are still rank-local; backend # all-gathers and writes to the per-rank pool. swa_loc = ( @@ -1308,21 +1295,11 @@ class FlashAttentionBackend(AttentionBackend): if self.use_sliding_window_kv_pool else None ) - if is_cp_v2_active(forward_batch): - cp_strategy = get_cp_strategy() - assert cp_strategy is not None - cp_strategy.materialize_full_kv( - forward_batch, layer, k, v, swa_loc=swa_loc - ) - else: - cp_allgather_and_save_kv_cache( - forward_batch, - layer, - k, - v, - self.attn_cp_size, - swa_loc=swa_loc, - ) + cp_strategy = get_cp_strategy() + assert cp_strategy is not None + cp_strategy.materialize_full_kv( + forward_batch, layer, k, v, swa_loc=swa_loc + ) else: k_scale = k_descale if self.kv_cache_is_mxfp8 else layer.k_scale v_scale = v_descale if self.kv_cache_is_mxfp8 else layer.v_scale @@ -1459,11 +1436,7 @@ class FlashAttentionBackend(AttentionBackend): cu_seqlens_k = metadata.encoder_cu_seqlens_k window_size = (-1, -1) - if ( - forward_batch.forward_mode.is_context_parallel_extend() - and forward_batch.attn_cp_metadata is not None - and self.attn_cp_size > 1 - ): + if cp_active: def _fa_cp_attn( q_chunk, cu_seqlens_q_cp, cache_seqlens_cp, max_seqlen_q_cp @@ -1488,23 +1461,15 @@ class FlashAttentionBackend(AttentionBackend): ) q_cp = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) - if is_cp_v2_active(forward_batch): - cp_strategy = get_cp_strategy() - assert cp_strategy is not None - result = cp_strategy.run_attention( - q_cp, - forward_batch, - self.device, - _fa_cp_attn, - attention_backend=CPAttentionBackendKind.FLASH_ATTENTION, - ) - else: - result = cp_attn_forward_extend( - forward_batch, - q_cp, - self.device, - _fa_cp_attn, - ) + cp_strategy = get_cp_strategy() + assert cp_strategy is not None + result = cp_strategy.run_attention( + q_cp, + forward_batch, + self.device, + _fa_cp_attn, + attention_backend=CPAttentionBackendKind.FLASH_ATTENTION, + ) elif self.fa_skip_kv_cache: # Embedding mode: skip KV cache read and use raw K/V tensors # directly via flash_attn_varlen_func. The KV cache write is @@ -1731,16 +1696,15 @@ class FlashAttentionBackend(AttentionBackend): q_nope = q_all[:, :, : layer.v_head_dim] q_rope = q_all[:, :, layer.v_head_dim :] - if is_cp_mode: + if cp_active: # MLA CP: q is rank-local zigzag-split; run the # absorbed-MLA kernel twice (prev/next halves) against - # the full latent KV pool (which rebuild_cp_kv_cache - # populated upstream) via cp_attn_forward_extend. + # the full latent KV pool through the selected strategy. # Concat q_nope + q_rope along dim=-1 so the wrapper's # chunk(2, dim=0) keeps their alignment; split back # inside the closure. assert not use_cascade_attn, ( - "Cascade attention under MLA CP is not supported in v1." + "Cascade attention under MLA CP is not supported." ) q_fused = torch.cat([q_nope, q_rope], dim=-1) @@ -1773,20 +1737,15 @@ class FlashAttentionBackend(AttentionBackend): ver=self.fa_impl_ver, ) - if is_cp_v2_active(forward_batch): - cp_strategy = get_cp_strategy() - assert cp_strategy is not None - o = cp_strategy.run_attention( - q_fused, - forward_batch, - self.device, - _mla_cp_attn, - attention_backend=CPAttentionBackendKind.FLASH_ATTENTION, - ) - else: - o = cp_attn_forward_extend( - forward_batch, q_fused, self.device, _mla_cp_attn - ) + cp_strategy = get_cp_strategy() + assert cp_strategy is not None + o = cp_strategy.run_attention( + q_fused, + forward_batch, + self.device, + _mla_cp_attn, + attention_backend=CPAttentionBackendKind.FLASH_ATTENTION, + ) else: result = flash_attn_with_kvcache( q=q_rope, diff --git a/python/sglang/srt/layers/attention/index_topk_share.py b/python/sglang/srt/layers/attention/index_topk_share.py index e7500d7da..3731a9f2b 100644 --- a/python/sglang/srt/layers/attention/index_topk_share.py +++ b/python/sglang/srt/layers/attention/index_topk_share.py @@ -3,6 +3,8 @@ from __future__ import annotations from contextlib import contextmanager from typing import TYPE_CHECKING, Iterator, Optional +from sglang.srt.layers.cp.utils import cp_gather_after_forward, is_cp_v2_active + if TYPE_CHECKING: import torch @@ -33,8 +35,11 @@ class IndexTopKShareState: @property def _seed_buf(self) -> Optional[torch.Tensor]: - if self._forward_batch.forward_mode.is_extend(include_draft_extend_v2=True): - return self._forward_batch.spec_info.dsa_seed_topk_capture + spec_info = self._forward_batch.spec_info + if spec_info is not None and self._forward_batch.forward_mode.is_extend( + include_draft_extend_v2=True + ): + return spec_info.dsa_seed_topk_capture return None @property @@ -46,6 +51,12 @@ class IndexTopKShareState: return self._topk_indices def update(self, topk_indices: Optional[torch.Tensor]) -> None: + if ( + topk_indices is not None + and self.should_publish + and is_cp_v2_active(self._forward_batch) + ): + topk_indices = cp_gather_after_forward(topk_indices, self._forward_batch) self._topk_indices = topk_indices def publish(self) -> None: diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 21ed8c17b..fcdeb296c 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -37,6 +37,10 @@ from sglang.srt.layers.attention.dsa.utils import ( is_dsa_enable_prefill_cp, ) from sglang.srt.layers.aux_hidden_states import AuxHiddenStateAccumulator +from sglang.srt.layers.cp.utils import ( + is_mla_prefill_cp_enabled, + mla_use_prefill_cp, +) from sglang.srt.layers.dp_attention import ( attn_tp_all_gather_into_tensor, attn_tp_reduce_scatter_tensor, @@ -64,10 +68,6 @@ from sglang.srt.layers.quantization.fp8_utils import ( _use_aiter_bpreshuffle_gfx95, materialize_bpreshuffle_fp8_scale_tuple, ) -from sglang.srt.layers.utils.cp_utils import ( - is_mla_prefill_cp_enabled, - mla_use_prefill_cp, -) from sglang.srt.model_executor.cuda_graph_config import ( Backend, Phase, diff --git a/python/sglang/srt/layers/communicator_dsa_cp.py b/python/sglang/srt/layers/communicator_dsa_cp.py index ee661b90f..26e91eff3 100644 --- a/python/sglang/srt/layers/communicator_dsa_cp.py +++ b/python/sglang/srt/layers/communicator_dsa_cp.py @@ -20,7 +20,6 @@ import torch from sglang.srt.layers.attention.dsa.utils import ( dsa_use_prefill_cp, - is_dsa_enable_prefill_cp, ) from sglang.srt.layers.communicator import ( CommunicateContext, @@ -31,24 +30,17 @@ from sglang.srt.layers.communicator import ( LayerScatterModes, ScatterMode, ) +from sglang.srt.layers.cp.utils import mla_use_prefill_cp from sglang.srt.layers.dp_attention import ( attn_cp_all_gather_into_tensor, attn_cp_reduce_scatter_tensor, get_local_dp_buffer, ) -from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp 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.runtime_context import get_parallel -def dsa_enable_prefill_cp(): - # After using cp, the communication mode of this part changes. - # The three parts of prepare_attn, prepare_mlp, and postprocess_layer - # no longer require additional communication for reduce, scatter, etc. - return is_dsa_enable_prefill_cp() - - def maybe_prefetch_next_full_attention_kv( forward_batch: ForwardBatch, next_full_attention_layer_id: Optional[int], diff --git a/python/sglang/srt/layers/cp/padding.py b/python/sglang/srt/layers/cp/padding.py index b482fe9ed..8c2921506 100644 --- a/python/sglang/srt/layers/cp/padding.py +++ b/python/sglang/srt/layers/cp/padding.py @@ -23,11 +23,10 @@ from sglang.srt.runtime_context import get_parallel def get_cp_padding_align_size() -> int: """Return the token-count alignment required by the active CP strategy.""" - from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_in_seq_split - from sglang.srt.layers.utils.cp_utils import is_prefill_cp_in_seq_split + from sglang.srt.layers.cp.base import is_zigzag attn_cp_size = get_parallel().attn_cp_size - if is_prefill_cp_in_seq_split() or is_dsa_prefill_cp_in_seq_split(): + if is_zigzag(): return attn_cp_size * 2 return attn_cp_size diff --git a/python/sglang/srt/layers/cp/utils.py b/python/sglang/srt/layers/cp/utils.py index 21f46a0e4..6be2d6285 100644 --- a/python/sglang/srt/layers/cp/utils.py +++ b/python/sglang/srt/layers/cp/utils.py @@ -23,6 +23,7 @@ from sglang.srt.layers.cp.base import ( ContextParallelStrategyKind, CPAttentionBackendKind, get_cp_strategy, + is_cp_enabled, ) from sglang.srt.layers.cp.interleave import ( InterleaveContextParallelMetadata, @@ -35,7 +36,7 @@ from sglang.srt.layers.cp.zigzag import ( ZigzagCPStrategy, ) from sglang.srt.layers.moe.utils import get_moe_a2a_backend -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_parallel, uses_mla_backend if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunner @@ -119,9 +120,9 @@ def get_layer_owner(local_layer_idx: int, shard_size: int, total_layers: int) -> def enable_cp_v2() -> bool: """Return whether the strategy-based generic prefill CP path is available.""" - from sglang.srt.utils import is_hip, is_npu + from sglang.srt.utils import is_hip, is_musa, is_npu - return not (is_hip() or is_npu()) + return not (is_hip() or is_npu() or is_musa()) def is_cp_v2_active(forward_batch) -> bool: @@ -143,6 +144,24 @@ def is_cp_v2_active(forward_batch) -> bool: return strategy.can_apply(len(input_ids), forward_batch) +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() + + +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() + ) + + def prepare_cp_forward(forward_batch) -> None: """Build CP-v2 metadata for an active context-parallel prefill batch.""" assert is_cp_v2_active(forward_batch) @@ -251,8 +270,8 @@ def cp_materialize_global_token_order( assert strategy is not None return strategy.gather_kv_cache(x, forward_batch, stream) - # TODO(hzh0425): Keep the legacy gather temporarily for CP-v1 compatibility. Remove it - # with the follow-up CP-v1 cleanup. + # 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( @@ -265,6 +284,7 @@ def cp_shard_model_inputs( complete_hidden_states: Any, complete_position_ids: Any, forward_batch, + complete_input_ids: Optional[Any] = None, ): """Restore the shared batch so logits processing keeps full-batch metadata.""" assert is_cp_v2_active(forward_batch) @@ -272,6 +292,18 @@ def cp_shard_model_inputs( complete_hidden_states, forward_batch ) sharded_positions = cp_shard_position_ids(complete_position_ids, forward_batch) + model_input_ids = ( + cp_shard_hidden_states(complete_input_ids, forward_batch) + if complete_input_ids is not None + else None + ) + + had_input_ids_global = hasattr(forward_batch, "input_ids_global") + input_ids_global_backup = getattr(forward_batch, "input_ids_global", None) + if complete_input_ids is not None: + forward_batch.input_ids_global = cp_round_robin_input_ids_v2( + complete_input_ids, forward_batch + ) spec_info = getattr(forward_batch, "spec_info", None) spec_hidden_states = getattr(spec_info, "hidden_states", None) @@ -286,10 +318,14 @@ def cp_shard_model_inputs( ) try: - yield sharded_hidden_states, sharded_positions + yield sharded_hidden_states, sharded_positions, model_input_ids finally: if spec_hidden_states_backup is not None: spec_info.hidden_states = spec_hidden_states_backup + if had_input_ids_global: + forward_batch.input_ids_global = input_ids_global_backup + elif hasattr(forward_batch, "input_ids_global"): + delattr(forward_batch, "input_ids_global") def _to_int_list(values) -> Optional[list[int]]: @@ -313,6 +349,8 @@ __all__ = [ "enable_cp_v2", "get_cp_strategy", "is_cp_v2_active", + "is_mla_prefill_cp_enabled", + "mla_use_prefill_cp", "cp_gather_after_forward", "cp_materialize_global_token_order", "cp_round_robin_input_ids_v2", diff --git a/python/sglang/srt/layers/utils/cp_utils.py b/python/sglang/srt/layers/utils/cp_utils.py index 0171952cf..3ac89faa1 100644 --- a/python/sglang/srt/layers/utils/cp_utils.py +++ b/python/sglang/srt/layers/utils/cp_utils.py @@ -1,6 +1,8 @@ +"""Legacy prefill CP helpers retained for HIP, NPU, and MUSA callers.""" + from dataclasses import dataclass from itertools import accumulate -from typing import Callable, List +from typing import List import torch import torch.nn.functional as F @@ -66,13 +68,6 @@ def is_prefill_context_parallel_enabled(): return get_parallel().enable_prefill_context_parallel -def is_prefill_cp_in_seq_split(): - return ( - is_prefill_context_parallel_enabled() - and get_parallel().prefill_cp_mode == "in-seq-split" - ) - - def is_mla_prefill_cp_enabled() -> bool: return get_parallel().enable_prefill_context_parallel and uses_mla_backend() @@ -477,48 +472,6 @@ def cp_allgather_and_save_kv_cache(forward_batch, layer, k, v, cp_size, swa_loc= ) -def cp_attn_forward_extend( - forward_batch, - q: torch.Tensor, - 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. - - For bs > 1, q is laid out as [all_prev_tokens_across_seqs, - all_next_tokens_across_seqs]; the split point is total_q_prev_tokens. - cu_seqlens_q_prev/next tensors have shape [bs+1] and carry the - per-sequence boundaries through FlashAttention's variable-length API. - - 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[: cp_meta.total_q_prev_tokens] - q_next = q[cp_meta.total_q_prev_tokens :] - - result_prev = attn_fn( - q_prev, - cp_meta.cu_seqlens_q_prev_tensor, - cp_meta.kv_len_prev_tensor, - cp_meta.max_seqlen_q_prev, - ) - result_next = attn_fn( - q_next, - cp_meta.cu_seqlens_q_next_tensor, - cp_meta.kv_len_next_tensor, - cp_meta.max_seqlen_q_next, - ) - - return torch.concat([result_prev, result_next], dim=0) - - def prepare_context_parallel_metadata( kv_len, cp_rank, diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 023071a67..3f9e5dd0f 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -39,8 +39,6 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union import torch from sglang.srt.dllm.config import DllmConfig -from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_in_seq_split -from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled from sglang.srt.managers.schedule_batch import ( Req, ScheduleBatch, @@ -582,9 +580,7 @@ class PrefillAdder: self.priority_scheduling_preemption_threshold = ( priority_scheduling_preemption_threshold ) - self.dsa_prefill_cp_in_seq_split = is_dsa_prefill_cp_in_seq_split() self.max_running_requests = max_running_requests - self.prefill_context_parallel_enabled = is_prefill_context_parallel_enabled() self.prefill_max_requests = prefill_max_requests self.prefill_delayer_single_pass = prefill_delayer_single_pass self.max_prefill_bs = max_prefill_bs @@ -1200,12 +1196,6 @@ class PrefillAdder: def add_one_req( self, req: Req, has_chunked_req: bool, truncation_align_size: Optional[int] ): - # TODO support cp with multiple requests - # Enabling context parallelism currently presents precision issues; - # therefore, the prefill-batch setting is temporarily set to 1. - if (self.dsa_prefill_cp_in_seq_split) and len(self.can_run_list) >= 1: - return AddReqResult.OTHER - if (x := self.prefill_max_requests) is not None and len(self.can_run_list) >= x: return AddReqResult.OTHER diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 5872abf93..2191348fa 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -67,9 +67,9 @@ from sglang.srt.utils import ( from sglang.srt.utils.common import ceil_align, is_pin_memory_available if TYPE_CHECKING: + from sglang.srt.layers.cp.base import BaseContextParallelMetadata from sglang.srt.layers.dcp.metadata import DecodeContextParallelMetadata from sglang.srt.layers.logits_processor import LogitsProcessorOutput - from sglang.srt.layers.utils.cp_utils import ContextParallelMetadata from sglang.srt.managers.schedule_batch import MultimodalInputs, ScheduleBatch from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo @@ -304,8 +304,8 @@ def compute_local_num_token_non_padded_cpu( def prefill_graph_tolerates_sum_len() -> bool: """Whether MegaMoE may replay prefill graphs with local shapes.""" from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp + from sglang.srt.layers.cp.utils import is_mla_prefill_cp_enabled from sglang.srt.layers.moe.utils import get_moe_a2a_backend - from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled if not get_moe_a2a_backend().is_megamoe(): return False @@ -598,7 +598,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): tbo_padded_len: Optional[int] = None tbo_children: Optional[List[ForwardBatch]] = None - attn_cp_metadata: Optional[ContextParallelMetadata] = None + attn_cp_metadata: Optional[BaseContextParallelMetadata] = None # For decode context parallel. # NOTE: DecodeContextParallelMetadata is imported under TYPE_CHECKING only (see the diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index ea7f75da7..b843fe27a 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -77,10 +77,10 @@ from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp from sglang.srt.layers.cp.utils import ( get_cp_strategy, is_cp_v2_active, + is_mla_prefill_cp_enabled, ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.sampler import create_sampler -from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled from sglang.srt.lora.lora_manager import LoRAManager, init_lora_cuda_graph_moe_buffers from sglang.srt.lora.lora_registry import LoRARef from sglang.srt.managers.schedule_batch import sanity_check_mm_pad_shift_value @@ -264,7 +264,7 @@ class SamplingPrewarmResult: def _prefill_cuda_graph_allows_context_parallel( prefill_runner, forward_batch: ForwardBatch ) -> bool: - """Allow CP only through a runner that captured the validated CP-v2 body.""" + """Allow CP only through a runner that captured the validated CP body.""" return get_cp_strategy() is None or ( bool(getattr(prefill_runner, "enable_cp_v2_bcg_capture", False)) and is_cp_v2_active(forward_batch) diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index b4cbebedf..47b5dbdaf 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -49,13 +49,13 @@ from sglang.srt.layers.attention.base_attn_backend import ( SharedReadEnds, ) from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp +from sglang.srt.layers.cp.utils import is_mla_prefill_cp_enabled from sglang.srt.layers.dp_attention import ( DpPaddingMode, set_dp_buffer_len, set_is_extend_in_batch, ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput -from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled from sglang.srt.model_executor.cuda_graph_buffer_registry import ( CudaGraphBufferRegistry, build_decode_registry, diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index 690c6e423..1bc6b3e8d 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -380,7 +380,7 @@ class EagerRunner(BaseRunner): def _execute_extend_cp_v2( self, forward_batch: ForwardBatch, kwargs: dict ) -> Union[LogitsProcessorOutput, PPProxyTensors]: - """CP-v2 extend: shard inputs at the model boundary, run the body on the + """CP extend: shard inputs at the model boundary, run the body on the rank-local slice, then gather hidden states before the logits step. """ model = self.model_runner.model @@ -389,13 +389,16 @@ class EagerRunner(BaseRunner): if input_embeds is None: input_embeds = model.get_input_embeddings()(forward_batch.input_ids) with cp_shard_model_inputs( - input_embeds, forward_batch.positions, forward_batch - ) as (sharded_input_embeds, sharded_positions): + input_embeds, + forward_batch.positions, + forward_batch, + forward_batch.input_ids, + ) as (sharded_input_embeds, sharded_positions, model_input_ids): model_kwargs = {"input_embeds": sharded_input_embeds} if (pp_proxy_tensors := kwargs.get("pp_proxy_tensors")) is not None: model_kwargs["pp_proxy_tensors"] = pp_proxy_tensors hidden_states = model.model( - forward_batch.input_ids, + model_input_ids, sharded_positions, forward_batch, **model_kwargs, diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index d04a57f23..d70853be8 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -64,7 +64,10 @@ from sglang.srt.layers.cp.bcg import ( execute_prefill_cp_bcg, filter_prefill_cp_bcg_capture_num_tokens, ) -from sglang.srt.layers.cp.utils import is_cp_v2_active +from sglang.srt.layers.cp.utils import ( + is_cp_v2_active, + is_mla_prefill_cp_enabled, +) from sglang.srt.layers.dp_attention import ( DpPaddingMode, set_dp_buffer_len, @@ -72,7 +75,6 @@ from sglang.srt.layers.dp_attention import ( ) from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.pooler import EmbeddingPoolerOutput -from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled from sglang.srt.model_executor.cuda_graph_buffer_registry import ( CudaGraphBufferRegistry, build_prefill_registry, 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 821455aee..f0d7f32c0 100644 --- a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py +++ b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py @@ -1,5 +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.model_executor.forward_context import get_attn_backend from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( @@ -111,10 +112,12 @@ def _handle_attention_backend(attn, forward_batch, backend_name): if is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph(): return AttnForwardMethod.MLA - # MLA prefill CP forces absorbed MLA regardless of prefix length: the - # CP path gathers latent KV via rebuild_cp_kv_cache and feeds the - # backend's absorbed-MLA kernel. - if mla_use_prefill_cp(forward_batch): + # 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) + ): return _dispatch_mla_subtype(attn, forward_batch) sum_extend_prefix_lens = _get_sum_extend_prefix_lens(forward_batch) diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index a5e89882b..b4fd29fe0 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 @@ -14,12 +14,10 @@ from sglang.kernels.ops.quantization.fp8_kernel import ( from sglang.srt.compilation.compilation_config import register_split_op 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, - is_graph_dsa_split_op_surface, -) +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 is_cp_v2_active +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, @@ -115,6 +113,7 @@ def should_defer_dsa_cp_kv_gather( dsa_prefill_cp: bool, fuse_rope_for_trtllm_mla: bool, ) -> bool: + """Compatibility predicate imported by the unchanged ROCm MLA path.""" return dsa_prefill_cp and fuse_rope_for_trtllm_mla @@ -266,10 +265,6 @@ class DeepseekMLAForwardMixin: return None if get_parallel().dcp_enabled: return None - # Context-parallel prefill reshuffles the KV side; keep the handshake - # out of those paths. - if dsa_use_prefill_cp(forward_batch) or mla_use_prefill_cp(forward_batch): - return None # Kernel shape constraints (tl.arange / tl.dot / block tiling). K # (qk_nope_head_dim) needs only K % 16 == 0 and K <= 256: power-of-2 # K (DeepSeek 128) takes the kernel's preload-once path, other K @@ -623,30 +618,17 @@ class DeepseekMLAForwardMixin: num_tokens, self.num_local_heads, self.kv_lora_rank, q_nope.device ) - 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, + k_nope, k_pe = prepare_kv_for_attention( + self, + forward_batch, + k_nope, + k_pe, + defer_materialization=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). + 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, + 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. diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index b3ca28468..95582e079 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -25,7 +25,6 @@ from torch import nn from transformers import PretrainedConfig from sglang.kernels.ops.layernorm.fused_eh_norm import fused_eh_norm -from sglang.srt.configs.model_config import is_deepseek_dsa 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 @@ -36,7 +35,7 @@ from sglang.srt.layers.attention.dsa.utils import ( is_dsa_prefill_cp_round_robin_split, ) from sglang.srt.layers.attention.index_topk_share import IndexTopKShareState -from sglang.srt.layers.cp.utils import cp_gather_after_forward, is_cp_v2_active +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 @@ -166,14 +165,6 @@ class DeepseekModelNextN(nn.Module): layer_name = "layers." + str(config.num_hidden_layers) self.quant_config = quant_config - self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() - self.mla_enable_prefill_cp = ( - is_mla_prefill_cp_enabled() and not is_deepseek_dsa(config) - ) - if self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp: - self.cp_size = get_parallel().attn_cp_size - else: - self.cp_size = None self.decoder = DeepseekV2DecoderLayer( config, 0, @@ -183,8 +174,6 @@ class DeepseekModelNextN(nn.Module): prefix=add_prefix(layer_name, prefix), alt_stream=self.alt_stream, skip_rope=config.qk_rope_head_dim == 0, - dsa_enable_prefill_cp=self.dsa_enable_prefill_cp, - mla_enable_prefill_cp=self.mla_enable_prefill_cp, ) self.shared_head = nn.Module() @@ -271,13 +260,11 @@ class DeepseekModelNextN(nn.Module): else: hidden_states = self.eh_proj(eh_input) - # CP-v2 shards/gathers hidden states at the eager-runner boundary. - cp_v2_active = is_cp_v2_active(forward_batch) - use_cp_v1 = ( - dsa_use_prefill_cp(forward_batch, self.dsa_enable_prefill_cp) - or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp) - ) and not cp_v2_active - if use_cp_v1: + # 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 @@ -297,11 +284,11 @@ class DeepseekModelNextN(nn.Module): else: hidden_states = self.shared_head.norm(hidden_states) - if use_cp_v1: + if use_platform_cp: local_num_tokens = hidden_states.shape[0] hidden_states = cp_all_gather_rerange_output( hidden_states, - self.cp_size, + get_parallel().attn_cp_size, forward_batch, torch.cuda.current_stream(), ) @@ -309,16 +296,10 @@ class DeepseekModelNextN(nn.Module): topk_indices = _gather_dsa_topk_indices_for_cp( topk_indices, local_num_tokens, - self.cp_size, + get_parallel().attn_cp_size, forward_batch, torch.cuda.current_stream(), ) - elif ( - cp_v2_active - and index_topk_share.should_publish - and topk_indices is not None - ): - topk_indices = cp_gather_after_forward(topk_indices, forward_batch) index_topk_share.update(topk_indices) index_topk_share.publish() finally: @@ -369,16 +350,6 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM): # if not set, model load will be broken in DeepseekV3ForCausalLM load_weights() self.pp_group = get_pp_group() self.determine_num_fused_shared_experts() - self.use_dsa = is_deepseek_dsa(config) - self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() - self.mla_enable_prefill_cp = is_mla_prefill_cp_enabled() and not self.use_dsa - if self.dsa_enable_prefill_cp or self.mla_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 - nextn_quant_config = self._resolve_nextn_quant_config(config, quant_config) self.model = DeepseekModelNextN( @@ -400,25 +371,31 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM): positions: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: - # TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2 - if not is_cp_v2_active(forward_batch): - if self.dsa_enable_prefill_cp: + if not enable_cp_v2(): + if is_dsa_enable_prefill_cp(): if can_dsa_cp_split( - len(input_ids), self.cp_size, self.use_dsa, forward_batch + 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), - self.cp_rank, - self.cp_size, + 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 self.mla_enable_prefill_cp: - if can_cp_split(len(input_ids), self.cp_size, forward_batch): + 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), - self.cp_rank, - self.cp_size, + 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, ) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 7d0a46848..e4f54b59b 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -85,7 +85,7 @@ 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 is_cp_v2_active +from sglang.srt.layers.cp.utils import enable_cp_v2 from sglang.srt.layers.dcp.planner import ( prepare_decode_context_parallel_metadata, ) @@ -470,14 +470,10 @@ class MoEGate(nn.Module): config, quant_config, prefix: str = "", - is_nextn: bool = False, is_hash_moe: bool = False, is_deepseek_v4: bool = False, - dsa_enable_prefill_cp: bool = False, - mla_enable_prefill_cp: bool = False, ): super().__init__() - self.is_nextn = is_nextn self.is_deepseek_v4 = is_deepseek_v4 self.weight = nn.Parameter( torch.empty( @@ -509,9 +505,6 @@ class MoEGate(nn.Module): self.e_score_correction_bias = None if _is_cpu and _is_cpu_amx_available: self.quant_method = PackWeightMethod(weight_names=["weight"]) - self.use_dsa = is_deepseek_dsa(config) - self.dsa_enable_prefill_cp = dsa_enable_prefill_cp - self.mla_enable_prefill_cp = mla_enable_prefill_cp self.tiny_router_gemm_max_tokens = tiny_router_gemm_max_tokens( num_experts=config.n_routed_experts, hidden_size=config.hidden_size, @@ -539,36 +532,29 @@ class MoEGate(nn.Module): return F.linear(hidden_states, self.weight, None) if ( - not self.is_deepseek_v4 + not enable_cp_v2() + and not self.is_deepseek_v4 and forward_batch is not None - and ( - dsa_use_prefill_cp(forward_batch, self.dsa_enable_prefill_cp) - or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp) - ) + and (dsa_use_prefill_cp(forward_batch) or mla_use_prefill_cp(forward_batch)) ): - if _is_cuda: - from sglang.kernels.ops.attention.dsv4 import linear_bf16_fp32 - - return linear_bf16_fp32(hidden_states, self.weight) 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, + self.weight, + out_dtype=torch.float32, + max_m=self.tiny_router_gemm_max_tokens, + ) + elif _use_aiter: + logits = aiter_dsv3_router_gemm(hidden_states, self.weight) + elif not _is_cuda: + logits = F.linear(hidden_states, self.weight, None) else: - if hidden_states.shape[0] <= self.tiny_router_gemm_max_tokens: - logits = tiny_gemm_bf16( - hidden_states, - self.weight, - out_dtype=torch.float32, - max_m=self.tiny_router_gemm_max_tokens, - ) + # cuBLAS bf16 x bf16 -> fp32 GEMM (torch.mm's out_dtype kwarg is CUDA-only) + from sglang.kernels.ops.attention.dsv4 import linear_bf16_fp32 - elif _use_aiter: - logits = aiter_dsv3_router_gemm(hidden_states, self.weight) - elif not _is_cuda: - logits = F.linear(hidden_states, self.weight, None) - else: - # cuBLAS bf16 x bf16 -> fp32 GEMM (torch.mm's out_dtype kwarg is CUDA-only) - from sglang.kernels.ops.attention.dsv4 import linear_bf16_fp32 - - logits = linear_bf16_fp32(hidden_states, self.weight) + logits = linear_bf16_fp32(hidden_states, self.weight) return logits @@ -583,8 +569,6 @@ class DeepseekV2MoE(nn.Module): alt_stream: Optional[torch.cuda.Stream] = None, is_nextn: bool = False, is_deepseek_v4: bool = False, - dsa_enable_prefill_cp: bool = False, - mla_enable_prefill_cp: bool = False, ): super().__init__() self.tp_size = get_parallel().tp_size @@ -643,11 +627,8 @@ class DeepseekV2MoE(nn.Module): config=config, quant_config=quant_config, prefix=add_prefix("gate", prefix), - is_nextn=is_nextn, is_hash_moe=self.is_hash, is_deepseek_v4=is_deepseek_v4, - dsa_enable_prefill_cp=dsa_enable_prefill_cp, - mla_enable_prefill_cp=mla_enable_prefill_cp, ) # scaling factor for fused shared experts on AMD-platform. @@ -1771,8 +1752,6 @@ class DeepseekV2AttentionMLA( alt_stream: Optional[torch.cuda.Stream] = None, skip_rope: bool = False, is_nextn: bool = False, - dsa_enable_prefill_cp: bool = False, - mla_enable_prefill_cp: bool = False, ) -> None: super().__init__() self.layer_id = layer_id @@ -1788,15 +1767,6 @@ class DeepseekV2AttentionMLA( attn_tp_rank = get_parallel().attn_tp_rank attn_tp_size = get_parallel().attn_tp_size self.use_dsa = is_deepseek_dsa(config) - self.dsa_enable_prefill_cp = dsa_enable_prefill_cp - self.mla_enable_prefill_cp = mla_enable_prefill_cp - if self.dsa_enable_prefill_cp: - assert self.use_dsa, "CP currently only supports deepseek v3.2 model" - # cp reuses the attn_tp comm group but needs to duplicate the weights; - # store cp_size whenever either CP flavor is active so rebuild_cp_kv_cache - # and the FA3 MLA wrapper can reach it on the dense MLA path too. - if self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp: - self.cp_size = get_parallel().attn_cp_size self.num_heads = num_heads assert num_heads % attn_tp_size == 0 self.num_local_heads = num_heads // attn_tp_size @@ -2291,12 +2261,12 @@ class DeepseekV2AttentionMLA( 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): - # support allgather+rerrange + # 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(), - self.cp_size, + get_parallel().attn_cp_size, forward_batch, torch.cuda.current_stream(), ) @@ -2327,8 +2297,6 @@ class DeepseekV2DecoderLayer(nn.Module): prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, skip_rope: bool = False, - dsa_enable_prefill_cp: bool = False, - mla_enable_prefill_cp: bool = False, ) -> None: super().__init__() self.hidden_size = config.hidden_size @@ -2345,8 +2313,6 @@ class DeepseekV2DecoderLayer(nn.Module): self.speculative_algorithm = SpeculativeAlgorithm.from_string( get_spec().speculative_algorithm ) - self.dsa_enable_prefill_cp = dsa_enable_prefill_cp - self.mla_enable_prefill_cp = mla_enable_prefill_cp self.layer_id = layer_id self.is_nextn = is_nextn if is_nextn and getattr(config, "mla_nope", False): @@ -2374,8 +2340,6 @@ class DeepseekV2DecoderLayer(nn.Module): alt_stream=alt_stream, skip_rope=skip_rope, is_nextn=is_nextn, - dsa_enable_prefill_cp=dsa_enable_prefill_cp, - mla_enable_prefill_cp=mla_enable_prefill_cp, ) if not hasattr(config, "q_lora_rank") and envs.SGLANG_USE_AG_AFTER_QLORA.get(): raise ValueError( @@ -2402,8 +2366,6 @@ class DeepseekV2DecoderLayer(nn.Module): layer_id=self.layer_id, alt_stream=alt_stream, is_nextn=is_nextn, - dsa_enable_prefill_cp=dsa_enable_prefill_cp, - mla_enable_prefill_cp=mla_enable_prefill_cp, ) else: if enable_moe_dense_fully_dp(): @@ -2428,31 +2390,21 @@ class DeepseekV2DecoderLayer(nn.Module): self._gfx95_quant_format = self._detect_gfx95_quant_format() - if self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp: - # DSACPLayerCommunicator is flavor-agnostic; its internal gates - # read both dsa_use_prefill_cp and mla_use_prefill_cp. The rename - # to CPLayerCommunicator is deferred to a cleanup PR. - self.layer_communicator = DSACPLayerCommunicator( - layer_scatter_modes=self.layer_scatter_modes, - input_layernorm=self.input_layernorm, - post_attention_layernorm=self.post_attention_layernorm, - allow_reduce_scatter=True, - is_last_layer=( - is_nextn or (self.layer_id == self.config.num_hidden_layers - 1) - ), - qkv_latent_func=self.self_attn.prepare_qkv_latent, - ) - else: - self.layer_communicator = LayerCommunicator( - layer_scatter_modes=self.layer_scatter_modes, - input_layernorm=self.input_layernorm, - post_attention_layernorm=self.post_attention_layernorm, - allow_reduce_scatter=True, - is_last_layer=( - is_nextn or (self.layer_id == self.config.num_hidden_layers - 1) - ), - qkv_latent_func=self.self_attn.prepare_qkv_latent, - ) + communicator_cls = ( + DSACPLayerCommunicator + if get_parallel().enable_prefill_cp + else LayerCommunicator + ) + self.layer_communicator = communicator_cls( + layer_scatter_modes=self.layer_scatter_modes, + input_layernorm=self.input_layernorm, + post_attention_layernorm=self.post_attention_layernorm, + allow_reduce_scatter=True, + is_last_layer=( + is_nextn or (self.layer_id == self.config.num_hidden_layers - 1) + ), + qkv_latent_func=self.self_attn.prepare_qkv_latent, + ) def _detect_gfx95_quant_format(self) -> str: if not _is_gfx95_supported: @@ -2577,10 +2529,7 @@ class DeepseekV2DecoderLayer(nn.Module): gemm_output_zero_allocator, ) - if ( - not (self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp) - and fuse_mlp_allreduce - ): + if fuse_mlp_allreduce: hidden_states._sglang_needs_allreduce_fusion = True if not fuse_mlp_allreduce: @@ -2666,14 +2615,6 @@ class DeepseekV2Model(nn.Module): self.vocab_size = config.vocab_size self.first_k_dense_replace = config.first_k_dense_replace self.pp_group = get_pp_group() - self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() - self.mla_enable_prefill_cp = ( - is_prefill_context_parallel_enabled() and not self.use_dsa - ) - if self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp: - self.cp_size = get_parallel().attn_cp_size - else: - self.cp_size = None if self.pp_group.is_first_rank: self.embed_tokens = VocabParallelEmbedding( @@ -2704,8 +2645,6 @@ class DeepseekV2Model(nn.Module): prefix=prefix, alt_stream=self.alt_stream, skip_rope=config.qk_rope_head_dim == 0, - dsa_enable_prefill_cp=self.dsa_enable_prefill_cp, - mla_enable_prefill_cp=self.mla_enable_prefill_cp, ), pp_rank=self.pp_group.rank_in_group, pp_size=self.pp_group.world_size, @@ -2864,13 +2803,11 @@ class DeepseekV2Model(nn.Module): else None ) - # CP-v2 shards/gathers at the eager-runner boundary instead. - use_cp_v1 = ( - dsa_use_prefill_cp(forward_batch, self.dsa_enable_prefill_cp) - or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp) - ) and not is_cp_v2_active(forward_batch) - - if use_cp_v1: + # 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) @@ -2974,11 +2911,11 @@ class DeepseekV2Model(nn.Module): else: hidden_states, _ = self.norm(hidden_states, residual) - if self.pp_group.is_last_rank and use_cp_v1: + if self.pp_group.is_last_rank and use_platform_cp: # allgather + rerrange hidden_states = cp_all_gather_rerange_output( hidden_states, - self.cp_size, + get_parallel().attn_cp_size, forward_batch, torch.cuda.current_stream(), ) @@ -3051,16 +2988,6 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin): ) self.capture_aux_hidden_states = False - self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() - self.mla_enable_prefill_cp = ( - is_prefill_context_parallel_enabled() and not is_deepseek_dsa(config) - ) - if self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp: - self.cp_rank = get_parallel().attn_cp_rank - self.cp_size = get_parallel().attn_cp_size - else: - self.cp_rank = self.cp_size = None - q_lora_rank = config.q_lora_rank if hasattr(config, "q_lora_rank") else None get_attn_tp_context().init_context(q_lora_rank, is_deepseek_dsa(config)) @@ -3161,24 +3088,29 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin): len_input_ids = input_embeds.shape[0] else: len_input_ids = pp_proxy_tensors["hidden_states"].shape[0] - if not is_cp_v2_active(forward_batch): - if self.dsa_enable_prefill_cp: + if not enable_cp_v2(): + if is_dsa_enable_prefill_cp(): if can_dsa_cp_split( - len_input_ids, self.cp_size, self.use_dsa, forward_batch + 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, - self.cp_rank, - self.cp_size, + 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 self.mla_enable_prefill_cp: - if can_cp_split(len_input_ids, self.cp_size, forward_batch): + 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, - self.cp_rank, - self.cp_size, + 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, ) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 768dcfdc0..fe65ef542 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -62,7 +62,7 @@ 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, - cp_round_robin_input_ids_v2, + enable_cp_v2, is_cp_v2_active, ) from sglang.srt.layers.dp_attention import ( @@ -3223,8 +3223,6 @@ class DeepseekV4Model(nn.Module): input_embeds: Optional[torch.Tensor], pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> Union[torch.Tensor, PPProxyTensors]: - cp_v2_active = is_cp_v2_active(forward_batch) - use_prefill_cp = dsa_use_prefill_cp(forward_batch) if self.pp_group.is_first_rank: if input_embeds is None: hidden_states = self.embed_tokens(input_ids) @@ -3254,7 +3252,7 @@ class DeepseekV4Model(nn.Module): ) input_ids_global = input_ids_global.squeeze(-1) else: - input_ids_global = input_ids + input_ids_global = getattr(forward_batch, "input_ids_global", input_ids) capture_dspark = self.dspark_layers_to_capture is not None dspark_aux_hidden_states: List[torch.Tensor] = [] @@ -3262,16 +3260,12 @@ class DeepseekV4Model(nn.Module): # 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 - if use_prefill_cp and not run_tbo: - if cp_v2_active: - input_ids = cp_round_robin_input_ids_v2(input_ids, forward_batch) - else: - 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) + 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. @@ -3323,12 +3317,7 @@ class DeepseekV4Model(nn.Module): ) # CP all-gather only on the last PP rank; PP IPC carries CP-split tensors. - if ( - self.pp_group.is_last_rank - and use_prefill_cp - and not cp_v2_active - and not run_tbo - ): + 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, @@ -3489,7 +3478,7 @@ class DeepseekV4ForCausalLM(nn.Module): input_embeds: Optional[torch.Tensor] = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> torch.Tensor: - if self.dsa_enable_prefill_cp: + 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), diff --git a/python/sglang/srt/models/deepseek_v4_nextn.py b/python/sglang/srt/models/deepseek_v4_nextn.py index df937a1ba..c9403fbfd 100644 --- a/python/sglang/srt/models/deepseek_v4_nextn.py +++ b/python/sglang/srt/models/deepseek_v4_nextn.py @@ -14,8 +14,7 @@ from sglang.srt.layers.attention.dsa.utils import ( is_dsa_prefill_cp_round_robin_split, ) from sglang.srt.layers.cp.utils import ( - cp_round_robin_input_ids_v2, - is_cp_v2_active, + enable_cp_v2, ) from sglang.srt.layers.dp_attention import ( dp_gather_replicate, @@ -144,8 +143,7 @@ class DeepseekV4ModelNextN(nn.Module): forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, ) -> torch.Tensor: - cp_v2_active = is_cp_v2_active(forward_batch) - use_prefill_cp = dsa_use_prefill_cp(forward_batch) + 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: @@ -179,15 +177,12 @@ class DeepseekV4ModelNextN(nn.Module): ) input_ids_global = input_ids_global.squeeze(-1) else: - input_ids_global = input_ids + input_ids_global = getattr(forward_batch, "input_ids_global", input_ids) - if use_prefill_cp: - if cp_v2_active: - input_ids = cp_round_robin_input_ids_v2(input_ids, forward_batch) - else: - 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) + 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 hidden_states, residual, post, comb = self.decoder( @@ -202,7 +197,7 @@ class DeepseekV4ModelNextN(nn.Module): # deferred fused hc_post state. hidden_states = self.decoder.hc_post(hidden_states, residual, post, comb) - if use_prefill_cp and not cp_v2_active: + if use_platform_cp: hidden_states = cp_all_gather_rerange_output( hidden_states, self.cp_size, @@ -260,7 +255,7 @@ class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM): positions: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: - if self.dsa_enable_prefill_cp and not is_cp_v2_active(forward_batch): + 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), diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index ea6cf62f4..b79a98376 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -48,7 +48,7 @@ from sglang.srt.layers.communicator import ( LayerScatterModes, ScatterMode, ) -from sglang.srt.layers.cp.utils import is_cp_v2_active +from sglang.srt.layers.cp.utils import enable_cp_v2 from sglang.srt.layers.dp_attention import ( is_dp_attention_enabled, ) @@ -1132,7 +1132,7 @@ class Qwen2MoeModel(nn.Module): if ( is_prefill_context_parallel_enabled() - and not is_cp_v2_active(forward_batch) + and not enable_cp_v2() and forward_batch.forward_mode.is_context_parallel_extend() and forward_batch.attn_cp_metadata is not None ): @@ -1198,7 +1198,7 @@ class Qwen2MoeModel(nn.Module): if ( self.pp_group.is_last_rank - and not is_cp_v2_active(forward_batch) + 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 diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index 0937e36e1..1d578e9fc 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -34,7 +34,7 @@ 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 is_cp_v2_active +from sglang.srt.layers.cp.utils import enable_cp_v2 from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( QKVParallelLinear, @@ -995,7 +995,7 @@ 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 is_cp_v2_active(forward_batch): + 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), diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index 27c7addd1..bf6559a77 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -30,14 +30,14 @@ from sglang.kernels.ops.layernorm.norm import ( fused_inplace_qknorm, ) from sglang.srt.environ import envs +from sglang.srt.layers.cp.utils import is_cp_v2_active from sglang.srt.layers.radix_attention import RadixAttention -from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool 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 +from sglang.srt.runtime_context import get_exec, get_parallel from sglang.srt.utils import get_current_device_stream_fast, is_cuda, is_hip from sglang.srt.utils.custom_op import register_custom_op @@ -296,11 +296,11 @@ def enable_fused_set_kv_buffer(forward_batch: ForwardBatch): _is_cuda and pool.dtype == torch.bfloat16 and not isinstance(pool, SWAKVPool) - and not is_prefill_context_parallel_enabled() + and not is_cp_v2_active(forward_batch) and getattr(forward_batch, "dcp_kv_mask", None) is None ) or ( _is_hip - and not is_prefill_context_parallel_enabled() + and not get_parallel().enable_prefill_context_parallel and getattr(forward_batch, "dcp_kv_mask", None) is None ) diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py index c42545dda..a69346934 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py @@ -342,12 +342,11 @@ class DSAMockModelRunner(ModelRunner): dllm_algorithm_config=None, dp_size=1, dsa_decode_backend=dsa_decode_backend, - dsa_prefill_cp_mode="round-robin-split", dsa_prefill_backend=dsa_prefill_backend, device=device, enable_deterministic_inference=False, enable_dp_attention=False, - enable_dsa_prefill_context_parallel=False, + enable_prefill_cp=False, enable_mis=False, is_embedding=False, kv_cache_dtype="auto", diff --git a/test/registered/cp/test_cp_strategy_unit.py b/test/registered/cp/test_cp_strategy_unit.py index cd41f9367..977d599f4 100644 --- a/test/registered/cp/test_cp_strategy_unit.py +++ b/test/registered/cp/test_cp_strategy_unit.py @@ -597,17 +597,7 @@ class TestCPZigzagStrategy(CustomTestCase): max_rank_len=[7, 7], ) - with ( - get_parallel().override(attn_cp_size=cp_size), - patch( - "sglang.srt.layers.utils.cp_utils.is_prefill_cp_in_seq_split", - return_value=True, - ), - patch( - "sglang.srt.layers.attention.dsa.utils.is_dsa_prefill_cp_in_seq_split", - return_value=False, - ), - ): + with get_parallel().override(attn_cp_size=cp_size): align_size = get_cp_padding_align_size() pad_logical_token_to_physical(metadata) diff --git a/test/registered/unit/layers/attention/test_index_topk_share.py b/test/registered/unit/layers/attention/test_index_topk_share.py index aa298bef3..f92a2bb34 100644 --- a/test/registered/unit/layers/attention/test_index_topk_share.py +++ b/test/registered/unit/layers/attention/test_index_topk_share.py @@ -5,6 +5,7 @@ import pytest import torch from sglang.srt.layers.attention.index_topk_share import IndexTopKShareState +from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=6, suite="base-a-test-cpu") @@ -15,9 +16,7 @@ def _batch( ) -> SimpleNamespace: return SimpleNamespace( reuse_dsa_topk_indices=reuse, - forward_mode=SimpleNamespace( - is_extend=lambda include_draft_extend_v2: is_extend - ), + forward_mode=ForwardMode.DRAFT_EXTEND_V2 if is_extend else ForwardMode.DECODE, spec_info=SimpleNamespace( dsa_topk_indices=carried, dsa_seed_topk_capture=seed_buf, diff --git a/test/registered/unit/models/test_deepseek_mla_dispatch.py b/test/registered/unit/models/test_deepseek_mla_dispatch.py index 9c491f2ef..3d87545e8 100644 --- a/test/registered/unit/models/test_deepseek_mla_dispatch.py +++ b/test/registered/unit/models/test_deepseek_mla_dispatch.py @@ -14,6 +14,11 @@ import unittest from types import SimpleNamespace from unittest import mock +from sglang.srt.layers.cp import base as cp_base +from sglang.srt.layers.cp import utils as cp_utils +from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy +from sglang.srt.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 ( AttnForwardMethod, @@ -95,5 +100,55 @@ class TestResolveRocmForwardMethod(CustomTestCase): self.assertEqual(abh.resolve_rocm_forward_method(method), method) +class TestCPMLADispatch(CustomTestCase): + def test_strategy_cp_uses_absorbed_mla_without_legacy_flags(self): + # Normal MHA writes rank-local KV against full out_cache_loc before + # the CP backend can gather it. Both one-shot and chunked MHA must + # therefore be bypassed for an active strategy-based CP batch. + attn = SimpleNamespace( + chunked_prefix_cache_threshold=0, + disable_chunked_prefix_cache=False, + flashinfer_mla_disable_ragged=False, + ) + with ( + mock.patch.object(abh, "_is_hip", False), + mock.patch.object(cp_utils, "enable_cp_v2", return_value=True), + mock.patch.object(cp_base, "_STRATEGY", ZigzagCPStrategy(cp_size=4)), + 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): + for num_tokens in (1, 3952): + with self.subTest( + prefix=prefix, capacity=capacity, num_tokens=num_tokens + ): + batch = SimpleNamespace( + forward_mode=ForwardMode.EXTEND, + input_ids=range(num_tokens), + attn_cp_metadata=None, + extend_prefix_lens_cpu=[prefix], + extend_seq_lens_cpu=[num_tokens], + seq_lens_cpu=[prefix + num_tokens], + get_max_chunk_capacity=lambda: capacity, + ) + expected = ( + AttnForwardMethod.MLA + if num_tokens == 3952 + else ( + AttnForwardMethod.MHA_ONE_SHOT + if capacity == 8192 + else AttnForwardMethod.MHA_CHUNKED_KV + ) + ) + self.assertEqual( + abh._handle_attention_backend(attn, batch, "fa3"), + expected, + ) + + 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 7e2ed387a..48bc0a15a 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -1075,28 +1075,22 @@ class TestContextParallelServerArgs(CustomTestCase): with self.assertRaisesRegex(ValueError, "DeepSeek V3.2.*interleave"): handle_context_parallelism(server_args) - @override_platform(is_hip=False, is_npu=False) - def test_generic_canonical_cp_mirrors_to_transitional_runtime_fields(self): + @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", - True, - False, - "in-seq-split", ), ( "interleave_dsa", "interleave", "dsa", - False, - True, - "round-robin-split", ), ) - for name, strategy, backend, expect_generic, expect_dsa, mode in cases: + for name, strategy, backend in cases: with self.subTest(name=name): server_args = self._new_cp_args( enable_prefill_cp=True, @@ -1105,6 +1099,7 @@ class TestContextParallelServerArgs(CustomTestCase): ) handle_platform_cp_compatibility(server_args) + handle_legacy_cp_runtime_compatibility(server_args) self.assertFalse( resolution_result(server_args, "enable_prefill_context_parallel") @@ -1115,32 +1110,13 @@ class TestContextParallelServerArgs(CustomTestCase): ) ) - handle_legacy_cp_runtime_compatibility(server_args) - - self.assertEqual( - resolution_result(server_args, "enable_prefill_context_parallel"), - expect_generic, - ) - self.assertEqual( - resolution_result( - server_args, "enable_dsa_prefill_context_parallel" - ), - expect_dsa, - ) - self.assertEqual( - resolution_result(server_args, "dsa_prefill_cp_mode"), mode - ) - self.assertEqual( - resolution_result(server_args, "prefill_cp_mode"), mode - ) - - @override_platform(is_hip=False, is_npu=False) + @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, "HIP or Ascend NPU"): + 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): @@ -1172,28 +1148,36 @@ class TestContextParallelServerArgs(CustomTestCase): resolution_result(args, "dsa_prefill_cp_mode"), "round-robin-split" ) - def test_canonical_interleave_cp_mirrors_to_dsa_runtime_aliases(self): - server_args = self._new_cp_args( - enable_prefill_cp=True, - cp_strategy="interleave", - attention_backend="dsa", - ) + 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) + 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" - ) + 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(