From 29345f08ec0c25a1a7f235be60722c1594f0adde Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Tue, 16 Jun 2026 02:30:01 -0700 Subject: [PATCH 1/5] Implement CP v2 zigzag for Qwen3 --- .../attention/flashattention_backend.py | 51 ++-- python/sglang/srt/layers/cp/base.py | 8 +- python/sglang/srt/layers/cp/utils.py | 85 +++++- python/sglang/srt/layers/cp/zigzag.py | 267 +++++++++++++++++- .../sglang/srt/model_executor/model_runner.py | 69 ++++- python/sglang/srt/models/utils.py | 3 + test/registered/cp/test_cp_strategy_unit.py | 237 ++++++++++++++-- test/registered/cp/test_qwen3_30b.py | 61 ++++ 8 files changed, 716 insertions(+), 65 deletions(-) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index d5f724b9bd20..6843a8c1f7fb 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -13,6 +13,7 @@ prepare_swa_spec_page_table_triton, ) from sglang.srt.layers.attention.utils import assert_buffer_fits +from sglang.srt.layers.cp.base import get_cp_strategy from sglang.srt.layers.radix_attention import AttentionType from sglang.srt.layers.utils.cp_utils import ( cp_allgather_and_save_kv_cache, @@ -811,18 +812,27 @@ def forward_extend( elif is_cp_mode: # Dense-MHA CP: k, v are still rank-local; backend # all-gathers and writes to the per-rank pool. - cp_allgather_and_save_kv_cache( - forward_batch, - layer, - k, - v, - self.attn_cp_size, - swa_loc=( - self.forward_metadata.swa_out_cache_loc - if self.use_sliding_window_kv_pool - else None - ), + swa_loc = ( + self.forward_metadata.swa_out_cache_loc + if self.use_sliding_window_kv_pool + else None ) + cp_strategy = get_cp_strategy() + if cp_strategy is not None: + forward_batch.cp_swa_out_cache_loc = swa_loc + try: + cp_strategy.materialize_full_kv(forward_batch, layer, k, v) + finally: + forward_batch.cp_swa_out_cache_loc = None + else: + cp_allgather_and_save_kv_cache( + forward_batch, + layer, + k, + v, + self.attn_cp_size, + swa_loc=swa_loc, + ) else: self.token_to_kv_pool.set_kv_buffer( layer, @@ -967,12 +977,19 @@ def _fa_cp_attn( **kwargs, ) - result = cp_attn_forward_extend( - forward_batch, - q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim), - self.device, - _fa_cp_attn, - ) + q_cp = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) + cp_strategy = get_cp_strategy() + if cp_strategy is not None: + result = cp_strategy.run_attention( + q_cp, forward_batch, self.device, _fa_cp_attn + ) + else: + result = cp_attn_forward_extend( + forward_batch, + q_cp, + self.device, + _fa_cp_attn, + ) 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 diff --git a/python/sglang/srt/layers/cp/base.py b/python/sglang/srt/layers/cp/base.py index 57a7dcd15867..961c9bc6f97b 100644 --- a/python/sglang/srt/layers/cp/base.py +++ b/python/sglang/srt/layers/cp/base.py @@ -258,18 +258,22 @@ def _get_cp_strategy() -> Optional[ContextParallelStrategy]: def get_cp_strategy() -> Optional[ContextParallelStrategy]: """Return the configured CP strategy for runtime dispatch.""" + from sglang.srt.environ import envs + + if not envs.SGLANG_ENABLE_CP_V2.get(): + return None return _get_cp_strategy() def get_cp_strategy_kind() -> ContextParallelStrategyKind: - strategy = _get_cp_strategy() + strategy = get_cp_strategy() if strategy is None: return ContextParallelStrategyKind.NONE return strategy.kind def is_cp_enabled() -> bool: - return _get_cp_strategy() is not None + return get_cp_strategy() is not None def is_zigzag() -> bool: diff --git a/python/sglang/srt/layers/cp/utils.py b/python/sglang/srt/layers/cp/utils.py index a8a8582db4b2..447465fc5b5b 100644 --- a/python/sglang/srt/layers/cp/utils.py +++ b/python/sglang/srt/layers/cp/utils.py @@ -12,13 +12,17 @@ # limitations under the License. # ============================================================================== -"""Public import facade for context parallel strategy helpers.""" +"""Public import facade and runtime helpers for context parallel strategies.""" + +from contextlib import contextmanager +from typing import Any, Dict, Optional from sglang.srt.layers.cp.base import ( BaseContextParallelMetadata, ContextParallelStrategy, ContextParallelStrategyKind, CPAttentionBackendKind, + get_cp_strategy, ) from sglang.srt.layers.cp.interleave import ( InterleaveContextParallelMetadata, @@ -30,6 +34,81 @@ ZigzagCPStrategy, ) + +def maybe_prepare_cp_forward(forward_batch) -> None: + """Build CP-v2 metadata for a prefill batch when the strategy can apply.""" + strategy = get_cp_strategy() + if strategy is None: + return + + input_ids = getattr(forward_batch, "input_ids", None) + if input_ids is None: + return + + num_tokens = len(input_ids) + if not strategy.can_apply(num_tokens, forward_batch): + return + + seq_lens_cpu = _to_int_list(getattr(forward_batch, "seq_lens_cpu", None)) + extend_lens_cpu = _to_int_list(getattr(forward_batch, "extend_seq_lens_cpu", None)) + forward_batch.attn_cp_metadata = strategy.build_metadata( + num_tokens=num_tokens, + seqs_len=seq_lens_cpu, + extend_seqs_len=extend_lens_cpu, + ) + forward_batch.cp_v2_active = True + + +def maybe_cp_split_before_forward( + model: Any, + forward_batch, + kwargs: Dict[str, Any], +) -> Optional[Any]: + """Shard embeddings and positions for CP-v2 model-runner forwarding.""" + strategy = get_cp_strategy() + if ( + strategy is None + or getattr(forward_batch, "attn_cp_metadata", None) is None + or not getattr(forward_batch, "cp_v2_active", False) + ): + return None + + input_embeds = kwargs.get("input_embeds") + if input_embeds is None: + embed_layer = model.get_input_embeddings() + input_embeds = embed_layer(forward_batch.input_ids) + kwargs["input_embeds"] = strategy.shard_hidden_states(input_embeds, forward_batch) + return strategy.shard_position_ids(forward_batch.positions, forward_batch) + + +@contextmanager +def disable_legacy_cp_during_cp_v2(forward_batch): + """Prevent model-local legacy CP hooks from double-sharding CP-v2 inputs.""" + if not getattr(forward_batch, "cp_v2_active", False): + yield + return + + from sglang.srt.server_args import get_global_server_args + + server_args = get_global_server_args() + old_enable_prefill_context_parallel = server_args.enable_prefill_context_parallel + server_args.enable_prefill_context_parallel = False + try: + yield + finally: + server_args.enable_prefill_context_parallel = ( + old_enable_prefill_context_parallel + ) + + +def _to_int_list(values) -> Optional[list[int]]: + if values is None: + return None + if hasattr(values, "tolist"): + values = values.tolist() + return [int(x) for x in values] + + __all__ = [ "BaseContextParallelMetadata", "CPAttentionBackendKind", @@ -40,4 +119,8 @@ "InterleaveContextParallelMetadata", "ZigzagCPStrategy", "ZigzagContextParallelMetadata", + "disable_legacy_cp_during_cp_v2", + "get_cp_strategy", + "maybe_cp_split_before_forward", + "maybe_prepare_cp_forward", ] diff --git a/python/sglang/srt/layers/cp/zigzag.py b/python/sglang/srt/layers/cp/zigzag.py index c25a75898659..e156a51c0de0 100644 --- a/python/sglang/srt/layers/cp/zigzag.py +++ b/python/sglang/srt/layers/cp/zigzag.py @@ -30,15 +30,29 @@ from __future__ import annotations +from contextlib import nullcontext from dataclasses import dataclass +from itertools import accumulate from typing import Any, List, Optional +import torch +import torch.nn.functional as F + +from sglang.srt.distributed.device_communicators.pynccl_allocator import ( + use_symmetric_memory, +) from sglang.srt.layers.cp.base import ( BaseContextParallelMetadata, ContextParallelStrategy, ContextParallelStrategyKind, CPAttentionBackendKind, ) +from sglang.srt.layers.dp_attention import ( + get_attention_cp_group, + is_allocation_symmetric, +) +from sglang.srt.mem_cache.memory_pool import KVWriteLoc +from sglang.srt.model_executor.forward_context import get_token_to_kv_pool @dataclass @@ -85,7 +99,13 @@ def can_apply(self, num_tokens: int, forward_batch) -> bool: if self.cp_size <= 1 or num_tokens < self.cp_size * 2: return False forward_mode = getattr(forward_batch, "forward_mode", None) - return forward_mode is None or forward_mode.is_context_parallel_extend() + if forward_mode is not None and not forward_mode.is_context_parallel_extend(): + return False + + extend_lens = getattr(forward_batch, "extend_seq_lens_cpu", None) + if extend_lens is None: + return True + return all(int(length) >= self.cp_size * 2 for length in extend_lens) def build_metadata( self, @@ -93,32 +113,183 @@ def build_metadata( seqs_len: Optional[List[int]], extend_seqs_len: Optional[List[int]] = None, ) -> ZigzagContextParallelMetadata: + if extend_seqs_len is None: + extend_seqs_len = seqs_len or [num_tokens] + extend_seqs_len = [int(x) for x in extend_seqs_len] + + pad_len = int(num_tokens) - 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 = self.cp_size * 2 + if seqs_len is not None and len(seqs_len) == bs: + prefix_offsets = [ + max(int(seqs_len[i]) - extend_seqs_len[i], 0) for i in range(bs) + ] + else: + prefix_offsets = [0] * bs + + per_seq_block_sizes: List[List[int]] = [] + split_list: List[int] = [] + for length in extend_seqs_len: + base = length // cp_segment_num + rem = length % cp_segment_num + block_sizes = [ + base + 1 if block_id < rem else base + for block_id in range(cp_segment_num) + ] + per_seq_block_sizes.append(block_sizes) + split_list.extend(block_sizes) + + per_rank_actual_token = [] + for rank in range(self.cp_size): + per_rank_actual_token.append( + sum( + block_sizes[rank] + block_sizes[cp_segment_num - 1 - rank] + for block_sizes in per_seq_block_sizes + ) + ) + max_rank_len = [max(per_rank_actual_token)] * self.cp_size + + cp_rank = self.cp_rank + 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, + ) + ) + + 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, + ) + ) + ) + + reverse_split_len: List[int] = [] + for rank in range(self.cp_size): + for batch_id in range(bs): + reverse_split_len.append(per_seq_block_sizes[batch_id][rank]) + for batch_id in range(bs): + reverse_split_len.append( + per_seq_block_sizes[batch_id][cp_segment_num - 1 - rank] + ) + + 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 batch_id, block_sizes in enumerate(per_seq_block_sizes): + kv_len_prev_list.append( + prefix_offsets[batch_id] + sum(block_sizes[: cp_rank + 1]) + ) + kv_len_next_list.append( + prefix_offsets[batch_id] + sum(block_sizes[: cp_segment_num - cp_rank]) + ) + actual_seq_q_prev_list.append(block_sizes[cp_rank]) + actual_seq_q_next_list.append(block_sizes[cp_segment_num - cp_rank - 1]) + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + cu_prev = [0] + list(accumulate(actual_seq_q_prev_list)) + cu_next = [0] + list(accumulate(actual_seq_q_next_list)) + + total_seq_lens = sum(extend_seqs_len) + 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 ZigzagContextParallelMetadata( - total_seq_lens=sum(extend_seqs_len or seqs_len or [num_tokens]), - bs=len(extend_seqs_len or seqs_len or [num_tokens]), + 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=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_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 + ), + 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, ) def shard_hidden_states(self, x: Any, forward_batch) -> Any: - raise NotImplementedError( - "Zigzag hidden-state sharding will land in a follow-up PR" + chunks = torch.split(x, forward_batch.attn_cp_metadata.split_list, dim=0) + return torch.cat( + [chunks[i] for i in forward_batch.attn_cp_metadata.zigzag_index], dim=0 ) def shard_position_ids(self, positions: Any, forward_batch) -> Any: - raise NotImplementedError( - "Zigzag position-id sharding will land in a follow-up PR" + chunks = torch.split( + positions, forward_batch.attn_cp_metadata.split_list, dim=-1 + ) + return torch.cat( + [chunks[i] for i in forward_batch.attn_cp_metadata.zigzag_index], dim=-1 ) def gather_hidden_states( self, x: Any, forward_batch, stream: Optional[Any] = None ) -> Any: - raise NotImplementedError( - "Zigzag hidden-state gather will land in a follow-up PR" + gathered = self._all_gather_reorganized(x, forward_batch, stream) + chunks = torch.split( + gathered, forward_batch.attn_cp_metadata.reverse_split_len, dim=0 + ) + return torch.cat( + [chunks[i] for i in forward_batch.attn_cp_metadata.cp_reverse_index], dim=0 ) def gather_kv_cache( self, x: Any, forward_batch, stream: Optional[Any] = None ) -> Any: - raise NotImplementedError("Zigzag KV gather will land in a follow-up PR") + gathered = self._all_gather_reorganized(x, forward_batch, stream) + chunks = torch.split( + gathered, forward_batch.attn_cp_metadata.reverse_split_len, dim=0 + ) + return torch.cat( + [chunks[i] for i in forward_batch.attn_cp_metadata.cp_reverse_index], dim=0 + ) def run_attention( self, @@ -128,11 +299,79 @@ def run_attention( attn_fn, attention_backend: CPAttentionBackendKind = CPAttentionBackendKind.FLASH_ATTENTION, ) -> Any: - raise NotImplementedError( - "Zigzag attention dispatch will land in a follow-up PR" + del device + if attention_backend != CPAttentionBackendKind.FLASH_ATTENTION: + raise NotImplementedError( + f"{self.name} CP does not support {attention_backend=}" + ) + + meta = forward_batch.attn_cp_metadata + q_prev = q[: meta.total_q_prev_tokens] + q_next = q[meta.total_q_prev_tokens :] + + result_prev = attn_fn( + q_prev, + meta.cu_seqlens_q_prev_tensor, + meta.kv_len_prev_tensor, + meta.max_seqlen_q_prev, + ) + result_next = attn_fn( + q_next, + meta.cu_seqlens_q_next_tensor, + meta.kv_len_next_tensor, + meta.max_seqlen_q_next, ) + return torch.cat([result_prev, result_next], dim=0) def materialize_full_kv(self, forward_batch, layer: Any, k: Any, v: Any) -> None: - raise NotImplementedError( - "Zigzag KV materialization will land in a follow-up PR" + cache_loc = ( + forward_batch.out_cache_loc + if not layer.is_cross_attention + else forward_batch.encoder_out_cache_loc + ) + key_cache_full = self.gather_kv_cache( + k.contiguous(), forward_batch, torch.cuda.current_stream() + ) + value_cache_full = self.gather_kv_cache( + v.contiguous(), forward_batch, torch.cuda.current_stream() + ) + get_token_to_kv_pool().set_kv_buffer( + layer, + KVWriteLoc(cache_loc, getattr(forward_batch, "cp_swa_out_cache_loc", None)), + key_cache_full, + value_cache_full, + layer.k_scale, + layer.v_scale, + ) + + def _all_gather_reorganized(self, x: torch.Tensor, forward_batch, stream): + meta = forward_batch.attn_cp_metadata + max_len = meta.max_rank_len[0] + pad_size = max_len - x.shape[0] + if pad_size > 0: + padding = [0, 0] * (x.ndim - 1) + [0, pad_size] + x = F.pad(x, padding, mode="constant", value=0) + + group = get_attention_cp_group() + ctx = ( + use_symmetric_memory(group, disabled=not is_allocation_symmetric()) + if x.is_cuda + else nullcontext() + ) + with ctx: + gathered = torch.empty( + max_len * self.cp_size, + *x.shape[1:], + device=x.device, + dtype=x.dtype, + ) + group.cp_all_gather_into_tensor_async(gathered, x, stream) + + chunks = torch.split(gathered, meta.max_rank_len, dim=0) + return torch.cat( + [ + chunks[rank][:per_rank_len] + for rank, per_rank_len in enumerate(meta.per_rank_actual_token) + ], + dim=0, ) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index e9ce8579ed73..1fbe7a19c1e5 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -117,6 +117,12 @@ ) from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp from sglang.srt.layers.attention.tbo_backend import TboAttnBackend +from sglang.srt.layers.cp.utils import ( + disable_legacy_cp_during_cp_v2, + get_cp_strategy, + maybe_cp_split_before_forward, + maybe_prepare_cp_forward, +) from sglang.srt.layers.dp_attention import ( DpPaddingMode, get_attention_tp_group, @@ -3308,6 +3314,8 @@ def forward_extend( self.prefill_cuda_graph_runner is not None and self.prefill_cuda_graph_runner.can_run(forward_batch) ) + if get_cp_strategy() is not None: + can_run_graph = False if can_run_graph: # TODO: device_timer.wrap is too broad here — it also includes # replay_prepare time. Move timing into the prefill cuda graph @@ -3331,6 +3339,10 @@ def forward_extend( # e.g. Moss-VL's prefill cross-attention custom mask. self.model.prepare_forward_batch(forward_batch) self.attn_backend.init_forward_metadata(forward_batch) + maybe_prepare_cp_forward(forward_batch) + cp_v2_positions = maybe_cp_split_before_forward( + self.model, forward_batch, kwargs + ) ctx = ( self.device_timer.wrap(metadata={"category": "extend"}) @@ -3363,14 +3375,59 @@ def forward_extend( **kwargs, ) else: - ret = self.model.forward( - forward_batch.input_ids, - forward_batch.positions, - forward_batch, - **kwargs, - ) + with disable_legacy_cp_during_cp_v2(forward_batch): + if cp_v2_positions is not None: + ret = self._forward_extend_cp_v2( + forward_batch, + cp_v2_positions, + kwargs, + ) + else: + ret = self.model.forward( + forward_batch.input_ids, + forward_batch.positions, + forward_batch, + **kwargs, + ) return (ret, can_run_graph) + def _forward_extend_cp_v2( + self, + forward_batch: ForwardBatch, + cp_v2_positions: torch.Tensor, + kwargs: dict, + ): + cp_strategy = get_cp_strategy() + assert cp_strategy is not None + + pp_proxy_tensors = kwargs.get("pp_proxy_tensors") + hidden_states = self.model.model( + forward_batch.input_ids, + cp_v2_positions, + forward_batch, + input_embeds=kwargs.get("input_embeds"), + pp_proxy_tensors=pp_proxy_tensors, + ) + + aux_hidden_states = None + if getattr(self.model, "capture_aux_hidden_states", False): + hidden_states, aux_hidden_states = hidden_states + + if self.model.pp_group.is_last_rank: + hidden_states = cp_strategy.gather_hidden_states( + hidden_states, + forward_batch, + torch.cuda.current_stream(), + ) + return self.model.logits_processor( + forward_batch.input_ids, + hidden_states, + self.model.lm_head, + forward_batch, + aux_hidden_states, + ) + return hidden_states + def forward_idle( self, forward_batch: ForwardBatch, pp_proxy_tensors=None ) -> Union[LogitsProcessorOutput, PPProxyTensors]: diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index 9465701c64de..0348eced18d5 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -287,6 +287,9 @@ def enable_fused_set_kv_buffer(forward_batch: ForwardBatch): `.view(-> 4D NHD)` reshape, and let the rotary forward pass `flash_layout=False`. See `create_fused_set_kv_buffer_arg` below. """ + if getattr(forward_batch, "cp_v2_active", False): + return False + pool = get_token_to_kv_pool() return ( _is_cuda diff --git a/test/registered/cp/test_cp_strategy_unit.py b/test/registered/cp/test_cp_strategy_unit.py index 8f306b5eaf92..f5ae3aa9861a 100644 --- a/test/registered/cp/test_cp_strategy_unit.py +++ b/test/registered/cp/test_cp_strategy_unit.py @@ -2,6 +2,8 @@ from types import SimpleNamespace from unittest.mock import patch +import torch + from sglang.srt.layers.cp.base import ( ContextParallelStrategyKind, get_cp_strategy, @@ -11,16 +13,81 @@ is_interleave, is_zigzag, ) +from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy 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") +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class _ExtendMode: + def is_context_parallel_extend(self): + return True + + +class _FakeCPGroup: + def __init__(self, all_rank_tensors): + self.all_rank_tensors = all_rank_tensors + + def cp_all_gather_into_tensor_async(self, output, input_tensor, stream): + del input_tensor, stream + torch.cat(self.all_rank_tensors, dim=0, out=output) class TestCPStrategyUnit(CustomTestCase): def tearDown(self): init_cp_strategy(SimpleNamespace(enable_prefill_cp=False)) + def _init_zigzag(self, cp_size=4): + init_cp_strategy( + SimpleNamespace( + enable_prefill_cp=True, + cp_strategy="zigzag", + attn_cp_size=cp_size, + attention_backend="fa3", + ) + ) + + def _metadata_for_rank(self, rank, cp_size=4, num_tokens=19): + strategy = ZigzagCPStrategy(cp_size=cp_size) + with patch( + "sglang.srt.layers.dp_attention.get_attention_cp_rank", + return_value=rank, + ): + return strategy.build_metadata( + num_tokens=num_tokens, + seqs_len=[11, 13], + extend_seqs_len=[9, 10], + ) + + def _forward_batch(self, metadata): + return SimpleNamespace( + forward_mode=_ExtendMode(), + extend_seq_lens_cpu=[9, 10], + attn_cp_metadata=metadata, + ) + + def _padded_rank_tensors(self, x, cp_size=4): + per_rank = [] + metas = [] + for rank in range(cp_size): + metadata = self._metadata_for_rank(rank, cp_size=cp_size) + metas.append(metadata) + fb = self._forward_batch(metadata) + with patch( + "sglang.srt.layers.dp_attention.get_attention_cp_rank", + return_value=rank, + ): + local = ZigzagCPStrategy(cp_size=cp_size).shard_hidden_states(x, fb) + pad = metadata.max_rank_len[0] - local.shape[0] + if pad: + local = torch.nn.functional.pad( + local, + [0, 0] * (local.ndim - 1) + [0, pad], + ) + per_rank.append(local) + return metas, per_rank + def test_strategy_kind_maps_cli_values(self): self.assertEqual(ContextParallelStrategyKind.NONE.value, 0) self.assertEqual( @@ -34,40 +101,160 @@ def test_strategy_kind_maps_cli_values(self): self.assertEqual(ContextParallelStrategyKind.ZIGZAG.cli_value, "zigzag") self.assertEqual(ContextParallelStrategyKind.INTERLEAVE.cli_value, "interleave") - def test_init_cp_strategy_binds_zigzag_strategy(self): - init_cp_strategy( - SimpleNamespace( - enable_prefill_cp=True, - cp_strategy="zigzag", - attn_cp_size=4, - ) - ) + def test_cp_v2_strategy_is_only_active_when_env_enabled(self): + self._init_zigzag() - self.assertTrue(is_cp_enabled()) - self.assertTrue(is_zigzag()) - self.assertFalse(is_interleave()) - self.assertEqual(get_cp_strategy_kind(), ContextParallelStrategyKind.ZIGZAG) + with patch( + "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=False + ): + self.assertIsNone(get_cp_strategy()) + self.assertFalse(is_cp_enabled()) - def test_get_cp_strategy_is_initialized_under_cp_v1_and_cp_v2(self): - init_cp_strategy( - SimpleNamespace( - enable_prefill_cp=True, - cp_strategy="interleave", - attn_cp_size=4, + with patch( + "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=True + ): + self.assertIsNotNone(get_cp_strategy()) + self.assertTrue(is_cp_enabled()) + self.assertTrue(is_zigzag()) + self.assertFalse(is_interleave()) + self.assertEqual(get_cp_strategy_kind(), ContextParallelStrategyKind.ZIGZAG) + + def test_zigzag_metadata_for_batched_sequences(self): + metadata = self._metadata_for_rank(rank=2) + + self.assertEqual(metadata.bs, 2) + self.assertEqual(metadata.total_seq_lens, 19) + self.assertEqual( + metadata.split_list, + [2, 1, 1, 1, 1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 1, 1], + ) + self.assertEqual(metadata.zigzag_index, [2, 10, 5, 13]) + self.assertEqual(metadata.per_rank_actual_token, [6, 5, 4, 4]) + self.assertEqual(metadata.max_rank_len, [6, 6, 6, 6]) + self.assertEqual( + metadata.reverse_split_len, + [2, 2, 1, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], + ) + self.assertEqual( + metadata.cp_reverse_index, + [0, 4, 8, 12, 14, 10, 6, 2, 1, 5, 9, 13, 15, 11, 7, 3], + ) + self.assertEqual(metadata.kv_len_prev_list, [6, 8]) + self.assertEqual(metadata.kv_len_next_list, [9, 11]) + self.assertEqual(metadata.actual_seq_q_prev_list, [1, 1]) + self.assertEqual(metadata.actual_seq_q_next_list, [1, 1]) + self.assertEqual(metadata.total_q_prev_tokens, 2) + self.assertEqual(metadata.total_q_next_tokens, 2) + self.assertTrue( + torch.equal( + metadata.cu_seqlens_q_prev_tensor.cpu(), torch.tensor([0, 1, 2]) + ) + ) + self.assertTrue( + torch.equal( + metadata.cu_seqlens_q_next_tensor.cpu(), torch.tensor([0, 1, 2]) ) ) + def test_zigzag_shards_hidden_states_and_position_ids(self): + metadata = self._metadata_for_rank(rank=2) + fb = self._forward_batch(metadata) + x = torch.arange(19 * 2).view(19, 2) + positions = torch.arange(19) + with patch( - "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=False + "sglang.srt.layers.dp_attention.get_attention_cp_rank", + return_value=2, ): - self.assertIsNotNone(get_cp_strategy()) - self.assertTrue(is_cp_enabled()) - self.assertTrue(is_interleave()) + strategy = ZigzagCPStrategy(cp_size=4) + local_x = strategy.shard_hidden_states(x, fb) + local_positions = strategy.shard_position_ids(positions, fb) + + expected_indices = [3, 13, 6, 16] + self.assertTrue(torch.equal(local_x, x[expected_indices])) + self.assertTrue(torch.equal(local_positions, positions[expected_indices])) + + def test_zigzag_gathers_hidden_states_to_original_order(self): + cp_size = 4 + x = torch.arange(19 * 2).view(19, 2) + metas, padded_rank_tensors = self._padded_rank_tensors(x, cp_size=cp_size) + rank = 2 + local_x = padded_rank_tensors[rank][: metas[rank].per_rank_actual_token[rank]] + fb = self._forward_batch(metas[rank]) + with ( + patch( + "sglang.srt.layers.cp.zigzag.get_attention_cp_group", + return_value=_FakeCPGroup(padded_rank_tensors), + ), + patch( + "sglang.srt.distributed.device_communicators.pynccl_allocator.use_symmetric_memory", + return_value=torch.no_grad(), + ), + ): + gathered = ZigzagCPStrategy(cp_size=cp_size).gather_hidden_states( + local_x, fb, stream=None + ) + + self.assertTrue(torch.equal(gathered, x)) + + def test_zigzag_gathers_kv_cache_to_original_order(self): + cp_size = 4 + kv = torch.arange(19 * 2 * 3).view(19, 2, 3) + metas, padded_rank_tensors = self._padded_rank_tensors(kv, cp_size=cp_size) + rank = 1 + local_kv = padded_rank_tensors[rank][: metas[rank].per_rank_actual_token[rank]] + fb = self._forward_batch(metas[rank]) + + with ( + patch( + "sglang.srt.layers.cp.zigzag.get_attention_cp_group", + return_value=_FakeCPGroup(padded_rank_tensors), + ), + patch( + "sglang.srt.distributed.device_communicators.pynccl_allocator.use_symmetric_memory", + return_value=torch.no_grad(), + ), + ): + gathered = ZigzagCPStrategy(cp_size=cp_size).gather_kv_cache( + local_kv, fb, stream=None + ) + + self.assertTrue(torch.equal(gathered, kv)) + + def test_zigzag_attention_dispatch_runs_prev_then_next(self): with patch( - "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=True + "sglang.srt.layers.dp_attention.get_attention_cp_rank", + return_value=0, ): - self.assertIsNotNone(get_cp_strategy()) + metadata = ZigzagCPStrategy(cp_size=2).build_metadata( + num_tokens=8, + seqs_len=[8], + extend_seqs_len=[8], + ) + fb = SimpleNamespace(attn_cp_metadata=metadata) + q = torch.arange(4 * 2).view(4, 2) + calls = [] + + def attn_fn(q_chunk, cu_seqlens_q, cache_seqlens, max_seqlen_q): + calls.append( + ( + q_chunk.clone(), + cu_seqlens_q.clone(), + cache_seqlens.clone(), + max_seqlen_q, + ) + ) + return q_chunk + 100 + + out = ZigzagCPStrategy(cp_size=2).run_attention( + q, fb, device=torch.device("cpu"), attn_fn=attn_fn + ) + + self.assertEqual(len(calls), 2) + self.assertTrue(torch.equal(calls[0][0], q[:2])) + self.assertTrue(torch.equal(calls[1][0], q[2:])) + self.assertTrue(torch.equal(out, q + 100)) if __name__ == "__main__": diff --git a/test/registered/cp/test_qwen3_30b.py b/test/registered/cp/test_qwen3_30b.py index e21dafb5c3c9..1cccc14adcd5 100644 --- a/test/registered/cp/test_qwen3_30b.py +++ b/test/registered/cp/test_qwen3_30b.py @@ -128,5 +128,66 @@ def test_gsm8k(self): self.assertGreaterEqual(metrics["score"], GSM8K_BASELINE_ACCURACY) +class TestQwen330BCPV2ZigzagDeepEP(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = QWEN3_30B_MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp-size", + "4", + "--ep", + "4", + "--attn-cp-size", + "4", + "--enable-prefill-cp", + "--cp-strategy", + "zigzag", + "--moe-a2a-backend", + "deepep", + "--attention-backend", + "fa3", + "--cuda-graph-max-bs", + "32", + "--max-running-requests", + "32", + "--trust-remote-code", + "--disable-piecewise-cuda-graph", + "--model-loader-extra-config", + '{"enable_multithread_load": true, "num_threads": 64}', + ], + env={"SGLANG_ENABLE_CP_V2": "1"}, + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + model=self.model, + eval_name="gsm8k", + num_shots=5, + num_examples=200, + max_tokens=16000, + num_threads=128, + repeat=1, + temperature=0.6, + top_p=0.95, + top_k=20, + base_url=self.base_url, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreaterEqual(metrics["score"], GSM8K_BASELINE_ACCURACY) + + if __name__ == "__main__": unittest.main() From f6d620139ad3a4cf97bc5adc2a70bed4bdf798d7 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Wed, 17 Jun 2026 21:05:32 -0700 Subject: [PATCH 2/5] Address CP v2 zigzag review comments --- .../attention/flashattention_backend.py | 27 +- python/sglang/srt/layers/cp/base.py | 12 +- python/sglang/srt/layers/cp/interleave.py | 4 +- python/sglang/srt/layers/cp/utils.py | 116 +++-- python/sglang/srt/layers/cp/zigzag.py | 26 +- .../sglang/srt/model_executor/model_runner.py | 127 +++-- python/sglang/srt/models/qwen2_moe.py | 3 + python/sglang/srt/models/qwen3_moe.py | 3 +- python/sglang/srt/models/utils.py | 3 - python/sglang/srt/server_args.py | 8 + test/registered/cp/test_cp_strategy_unit.py | 464 +++++++++++++----- ...est_qwen3_30b.py => test_gqa_preill_cp.py} | 27 +- 12 files changed, 527 insertions(+), 293 deletions(-) rename test/registered/cp/{test_qwen3_30b.py => test_gqa_preill_cp.py} (88%) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 6843a8c1f7fb..0a50f0d60d04 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -13,7 +13,8 @@ prepare_swa_spec_page_table_triton, ) from sglang.srt.layers.attention.utils import assert_buffer_fits -from sglang.srt.layers.cp.base import get_cp_strategy +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.radix_attention import AttentionType from sglang.srt.layers.utils.cp_utils import ( cp_allgather_and_save_kv_cache, @@ -817,13 +818,12 @@ def forward_extend( if self.use_sliding_window_kv_pool else None ) - cp_strategy = get_cp_strategy() - if cp_strategy is not None: - forward_batch.cp_swa_out_cache_loc = swa_loc - try: - cp_strategy.materialize_full_kv(forward_batch, layer, k, v) - finally: - forward_batch.cp_swa_out_cache_loc = 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, @@ -978,10 +978,15 @@ def _fa_cp_attn( ) q_cp = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim) - cp_strategy = get_cp_strategy() - if cp_strategy is not None: + 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 + q_cp, + forward_batch, + self.device, + _fa_cp_attn, + attention_backend=CPAttentionBackendKind.FLASH_ATTENTION, ) else: result = cp_attn_forward_extend( diff --git a/python/sglang/srt/layers/cp/base.py b/python/sglang/srt/layers/cp/base.py index 961c9bc6f97b..c11be84f9e59 100644 --- a/python/sglang/srt/layers/cp/base.py +++ b/python/sglang/srt/layers/cp/base.py @@ -184,6 +184,7 @@ def materialize_full_kv( layer: Any, k: Any, v: Any, + swa_loc: Optional[Any] = None, ) -> None: """Write full-layout K/V to the backend cache if needed.""" @@ -234,7 +235,7 @@ def init_cp_strategy(server_args: ServerArgs) -> None: ) -def _get_cp_strategy() -> Optional[ContextParallelStrategy]: +def get_cp_strategy() -> Optional[ContextParallelStrategy]: """Return the configured strategy, initializing lazily on first call. Subprocesses re-import this module with ``_STRATEGY = None`` and never @@ -256,15 +257,6 @@ def _get_cp_strategy() -> Optional[ContextParallelStrategy]: return _STRATEGY -def get_cp_strategy() -> Optional[ContextParallelStrategy]: - """Return the configured CP strategy for runtime dispatch.""" - from sglang.srt.environ import envs - - if not envs.SGLANG_ENABLE_CP_V2.get(): - return None - return _get_cp_strategy() - - def get_cp_strategy_kind() -> ContextParallelStrategyKind: strategy = get_cp_strategy() if strategy is None: diff --git a/python/sglang/srt/layers/cp/interleave.py b/python/sglang/srt/layers/cp/interleave.py index 7ce9db8968df..7db762ab9cb8 100644 --- a/python/sglang/srt/layers/cp/interleave.py +++ b/python/sglang/srt/layers/cp/interleave.py @@ -99,7 +99,9 @@ def run_attention( "Interleave attention dispatch will land in a follow-up PR" ) - def materialize_full_kv(self, forward_batch, layer: Any, k: Any, v: Any) -> None: + def materialize_full_kv( + self, forward_batch, layer: Any, k: Any, v: Any, swa_loc: Optional[Any] = None + ) -> None: raise NotImplementedError( "Interleave KV materialization will land in a follow-up PR" ) diff --git a/python/sglang/srt/layers/cp/utils.py b/python/sglang/srt/layers/cp/utils.py index 447465fc5b5b..e7561d99a0a8 100644 --- a/python/sglang/srt/layers/cp/utils.py +++ b/python/sglang/srt/layers/cp/utils.py @@ -14,8 +14,7 @@ """Public import facade and runtime helpers for context parallel strategies.""" -from contextlib import contextmanager -from typing import Any, Dict, Optional +from typing import Any, Optional, Tuple from sglang.srt.layers.cp.base import ( BaseContextParallelMetadata, @@ -34,20 +33,49 @@ ZigzagCPStrategy, ) +CP_V2_DEFAULT_MODEL_CLASSES = frozenset( + { + "Qwen3MoeForCausalLM", + } +) + + +def enable_cp_v2() -> bool: + """Return whether the CP-v2 path is enabled for this process.""" + from sglang.srt.environ import envs + + return bool(envs.SGLANG_ENABLE_CP_V2.get()) + + +def is_cp_v2_active(forward_batch) -> bool: + """Return whether the current forward batch is running through CP-v2.""" + if not enable_cp_v2(): + return False + forward_mode = getattr(forward_batch, "forward_mode", None) + if forward_mode is None or not forward_mode.is_context_parallel_extend(): + return False -def maybe_prepare_cp_forward(forward_batch) -> None: - """Build CP-v2 metadata for a prefill batch when the strategy can apply.""" strategy = get_cp_strategy() if strategy is None: - return + return False input_ids = getattr(forward_batch, "input_ids", None) if input_ids is None: - return + return False + + return strategy.can_apply(len(input_ids), forward_batch) + + +def prepare_cp_forward(forward_batch) -> None: + """Build CP-v2 metadata for an active context-parallel prefill batch.""" + strategy = get_cp_strategy() + assert strategy is not None + + input_ids = getattr(forward_batch, "input_ids", None) + assert input_ids is not None num_tokens = len(input_ids) - if not strategy.can_apply(num_tokens, forward_batch): - return + assert strategy.can_apply(num_tokens, forward_batch) seq_lens_cpu = _to_int_list(getattr(forward_batch, "seq_lens_cpu", None)) extend_lens_cpu = _to_int_list(getattr(forward_batch, "extend_seq_lens_cpu", None)) @@ -56,49 +84,40 @@ def maybe_prepare_cp_forward(forward_batch) -> None: seqs_len=seq_lens_cpu, extend_seqs_len=extend_lens_cpu, ) - forward_batch.cp_v2_active = True -def maybe_cp_split_before_forward( - model: Any, +def cp_split_before_forward( + complete_hidden_states: Any, + complete_position_ids: Any, forward_batch, - kwargs: Dict[str, Any], -) -> Optional[Any]: +) -> Tuple[Optional[Any], Optional[Any]]: """Shard embeddings and positions for CP-v2 model-runner forwarding.""" strategy = get_cp_strategy() - if ( - strategy is None - or getattr(forward_batch, "attn_cp_metadata", None) is None - or not getattr(forward_batch, "cp_v2_active", False) - ): - return None + assert strategy is not None + assert complete_hidden_states is not None + assert getattr(forward_batch, "attn_cp_metadata", None) is not None + return ( + strategy.shard_hidden_states(complete_hidden_states, forward_batch), + strategy.shard_position_ids(complete_position_ids, forward_batch), + ) + - input_embeds = kwargs.get("input_embeds") - if input_embeds is None: - embed_layer = model.get_input_embeddings() - input_embeds = embed_layer(forward_batch.input_ids) - kwargs["input_embeds"] = strategy.shard_hidden_states(input_embeds, forward_batch) - return strategy.shard_position_ids(forward_batch.positions, forward_batch) - - -@contextmanager -def disable_legacy_cp_during_cp_v2(forward_batch): - """Prevent model-local legacy CP hooks from double-sharding CP-v2 inputs.""" - if not getattr(forward_batch, "cp_v2_active", False): - yield - return - - from sglang.srt.server_args import get_global_server_args - - server_args = get_global_server_args() - old_enable_prefill_context_parallel = server_args.enable_prefill_context_parallel - server_args.enable_prefill_context_parallel = False - try: - yield - finally: - server_args.enable_prefill_context_parallel = ( - old_enable_prefill_context_parallel +def maybe_cp_gather_after_forward(x: Any, forward_batch, stream: Optional[Any] = None): + """Gather CP-v2 hidden states at the model boundary when this batch is active.""" + if not is_cp_v2_active(forward_batch): + return x + + strategy = get_cp_strategy() + assert strategy is not None + + if isinstance(x, tuple): + hidden_states, *rest = x + hidden_states = strategy.gather_hidden_states( + hidden_states, forward_batch, stream ) + return (hidden_states, *rest) + + return strategy.gather_hidden_states(x, forward_batch, stream) def _to_int_list(values) -> Optional[list[int]]: @@ -119,8 +138,11 @@ def _to_int_list(values) -> Optional[list[int]]: "InterleaveContextParallelMetadata", "ZigzagCPStrategy", "ZigzagContextParallelMetadata", - "disable_legacy_cp_during_cp_v2", + "CP_V2_DEFAULT_MODEL_CLASSES", + "enable_cp_v2", "get_cp_strategy", - "maybe_cp_split_before_forward", - "maybe_prepare_cp_forward", + "is_cp_v2_active", + "maybe_cp_gather_after_forward", + "cp_split_before_forward", + "prepare_cp_forward", ] diff --git a/python/sglang/srt/layers/cp/zigzag.py b/python/sglang/srt/layers/cp/zigzag.py index e156a51c0de0..3e5af00ba776 100644 --- a/python/sglang/srt/layers/cp/zigzag.py +++ b/python/sglang/srt/layers/cp/zigzag.py @@ -133,6 +133,8 @@ def build_metadata( else: prefix_offsets = [0] * bs + # TODO: move these per-request layout/index computations to a Triton + # kernel if Python-side metadata construction becomes a bottleneck. per_seq_block_sizes: List[List[int]] = [] split_list: List[int] = [] for length in extend_seqs_len: @@ -202,7 +204,12 @@ def build_metadata( actual_seq_q_prev_list.append(block_sizes[cp_rank]) actual_seq_q_next_list.append(block_sizes[cp_segment_num - cp_rank - 1]) - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + from sglang.srt.server_args import get_global_server_args + + try: + device = torch.device(get_global_server_args().device) + except Exception: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") cu_prev = [0] + list(accumulate(actual_seq_q_prev_list)) cu_next = [0] + list(accumulate(actual_seq_q_next_list)) @@ -291,6 +298,9 @@ def gather_kv_cache( [chunks[i] for i in forward_batch.attn_cp_metadata.cp_reverse_index], dim=0 ) + def get_supported_attention_backend(self): + return [CPAttentionBackendKind.FLASH_ATTENTION] + def run_attention( self, q: Any, @@ -299,11 +309,9 @@ def run_attention( attn_fn, attention_backend: CPAttentionBackendKind = CPAttentionBackendKind.FLASH_ATTENTION, ) -> Any: - del device - if attention_backend != CPAttentionBackendKind.FLASH_ATTENTION: - raise NotImplementedError( - f"{self.name} CP does not support {attention_backend=}" - ) + assert ( + attention_backend in self.get_supported_attention_backend() + ), f"{self.name} CP does not support {attention_backend=}" meta = forward_batch.attn_cp_metadata q_prev = q[: meta.total_q_prev_tokens] @@ -323,7 +331,9 @@ def run_attention( ) return torch.cat([result_prev, result_next], dim=0) - def materialize_full_kv(self, forward_batch, layer: Any, k: Any, v: Any) -> None: + def materialize_full_kv( + self, forward_batch, layer: Any, k: Any, v: Any, swa_loc: Optional[Any] = None + ) -> None: cache_loc = ( forward_batch.out_cache_loc if not layer.is_cross_attention @@ -337,7 +347,7 @@ def materialize_full_kv(self, forward_batch, layer: Any, k: Any, v: Any) -> None ) get_token_to_kv_pool().set_kv_buffer( layer, - KVWriteLoc(cache_loc, getattr(forward_batch, "cp_swa_out_cache_loc", None)), + KVWriteLoc(cache_loc, swa_loc), key_cache_full, value_cache_full, layer.k_scale, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 1fbe7a19c1e5..d2c0e8061c0d 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -118,10 +118,11 @@ from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp from sglang.srt.layers.attention.tbo_backend import TboAttnBackend from sglang.srt.layers.cp.utils import ( - disable_legacy_cp_during_cp_v2, + cp_split_before_forward, get_cp_strategy, - maybe_cp_split_before_forward, - maybe_prepare_cp_forward, + is_cp_v2_active, + maybe_cp_gather_after_forward, + prepare_cp_forward, ) from sglang.srt.layers.dp_attention import ( DpPaddingMode, @@ -3339,10 +3340,63 @@ def forward_extend( # e.g. Moss-VL's prefill cross-attention custom mask. self.model.prepare_forward_batch(forward_batch) self.attn_backend.init_forward_metadata(forward_batch) - maybe_prepare_cp_forward(forward_batch) - cp_v2_positions = maybe_cp_split_before_forward( - self.model, forward_batch, kwargs - ) + cp_v2_active = is_cp_v2_active(forward_batch) + forward_positions = forward_batch.positions + if cp_v2_active: + prepare_cp_forward(forward_batch) + complete_hidden_states = kwargs.get("input_embeds") + if complete_hidden_states is None: + embed_layer = self.model.get_input_embeddings() + complete_hidden_states = embed_layer(forward_batch.input_ids) + sharded_hidden_states, sharded_positions = cp_split_before_forward( + complete_hidden_states, + forward_batch.positions, + forward_batch, + ) + kwargs["input_embeds"] = sharded_hidden_states + forward_positions = sharded_positions + + def _forward_model(): + if not cp_v2_active: + return self.model.forward( + forward_batch.input_ids, + forward_positions, + forward_batch, + **kwargs, + ) + + hidden_states = self.model.model( + forward_batch.input_ids, + forward_positions, + forward_batch, + input_embeds=kwargs.get("input_embeds"), + pp_proxy_tensors=kwargs.get("pp_proxy_tensors"), + ) + + aux_hidden_states = None + capture_aux_hidden_states = getattr( + self.model, "capture_aux_hidden_states", False + ) + if capture_aux_hidden_states: + hidden_states, aux_hidden_states = hidden_states + + if self.model.pp_group.is_last_rank: + hidden_states = maybe_cp_gather_after_forward( + hidden_states, + forward_batch, + torch.cuda.current_stream(), + ) + return self.model.logits_processor( + forward_batch.input_ids, + hidden_states, + self.model.lm_head, + forward_batch, + aux_hidden_states, + ) + + if capture_aux_hidden_states: + return hidden_states, aux_hidden_states + return hidden_states ctx = ( self.device_timer.wrap(metadata={"category": "extend"}) @@ -3368,66 +3422,11 @@ def forward_extend( dsa_indexers=self.dsa_indexers, ), ): - ret = self.model.forward( - forward_batch.input_ids, - forward_batch.positions, - forward_batch, - **kwargs, - ) + ret = _forward_model() else: - with disable_legacy_cp_during_cp_v2(forward_batch): - if cp_v2_positions is not None: - ret = self._forward_extend_cp_v2( - forward_batch, - cp_v2_positions, - kwargs, - ) - else: - ret = self.model.forward( - forward_batch.input_ids, - forward_batch.positions, - forward_batch, - **kwargs, - ) + ret = _forward_model() return (ret, can_run_graph) - def _forward_extend_cp_v2( - self, - forward_batch: ForwardBatch, - cp_v2_positions: torch.Tensor, - kwargs: dict, - ): - cp_strategy = get_cp_strategy() - assert cp_strategy is not None - - pp_proxy_tensors = kwargs.get("pp_proxy_tensors") - hidden_states = self.model.model( - forward_batch.input_ids, - cp_v2_positions, - forward_batch, - input_embeds=kwargs.get("input_embeds"), - pp_proxy_tensors=pp_proxy_tensors, - ) - - aux_hidden_states = None - if getattr(self.model, "capture_aux_hidden_states", False): - hidden_states, aux_hidden_states = hidden_states - - if self.model.pp_group.is_last_rank: - hidden_states = cp_strategy.gather_hidden_states( - hidden_states, - forward_batch, - torch.cuda.current_stream(), - ) - return self.model.logits_processor( - forward_batch.input_ids, - hidden_states, - self.model.lm_head, - forward_batch, - aux_hidden_states, - ) - return hidden_states - def forward_idle( self, forward_batch: ForwardBatch, pp_proxy_tensors=None ) -> Union[LogitsProcessorOutput, PPProxyTensors]: diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 324e8678d8c7..54aa3eddbe19 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -48,6 +48,7 @@ LayerScatterModes, ScatterMode, ) +from sglang.srt.layers.cp.utils import is_cp_v2_active from sglang.srt.layers.dp_attention import ( get_attention_tp_rank, get_attention_tp_size, @@ -896,6 +897,7 @@ def forward( if ( is_prefill_context_parallel_enabled() + and not is_cp_v2_active(forward_batch) and forward_batch.forward_mode.is_context_parallel_extend() and forward_batch.attn_cp_metadata is not None ): @@ -951,6 +953,7 @@ def forward( if ( self.pp_group.is_last_rank + and not is_cp_v2_active(forward_batch) 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 5c6c76bf13fa..ab2a92db5b4b 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -40,6 +40,7 @@ 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.dp_attention import get_attention_tp_rank, get_attention_tp_size from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( @@ -994,7 +995,7 @@ def forward( input_embeds: torch.Tensor = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> torch.Tensor: - if is_prefill_context_parallel_enabled(): + if is_prefill_context_parallel_enabled() and not is_cp_v2_active(forward_batch): 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 0348eced18d5..9465701c64de 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -287,9 +287,6 @@ def enable_fused_set_kv_buffer(forward_batch: ForwardBatch): `.view(-> 4D NHD)` reshape, and let the rotary forward pass `flash_layout=False`. See `create_fused_set_kv_buffer_arg` below. """ - if getattr(forward_batch, "cp_v2_active", False): - return False - pool = get_token_to_kv_pool() return ( _is_cuda diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 28cc109de5f1..9e29767ade6f 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1965,6 +1965,14 @@ def _handle_model_specific_adjustments(self): hf_config = model_config.hf_config model_arch = hf_config.architectures[0] + from sglang.srt.layers.cp.utils import CP_V2_DEFAULT_MODEL_CLASSES + + if ( + model_arch in CP_V2_DEFAULT_MODEL_CLASSES + and not envs.SGLANG_ENABLE_CP_V2.is_set() + ): + envs.SGLANG_ENABLE_CP_V2.set(True) + _hybrid_spec = get_linear_attn_spec_by_arch(model_arch) if _hybrid_spec is not None and _hybrid_spec.uses_mamba_radix_cache: self._handle_mamba_radix_cache( diff --git a/test/registered/cp/test_cp_strategy_unit.py b/test/registered/cp/test_cp_strategy_unit.py index f5ae3aa9861a..1b0ae8e6bf9e 100644 --- a/test/registered/cp/test_cp_strategy_unit.py +++ b/test/registered/cp/test_cp_strategy_unit.py @@ -13,6 +13,11 @@ is_interleave, is_zigzag, ) +from sglang.srt.layers.cp.utils import ( + cp_split_before_forward, + enable_cp_v2, + is_cp_v2_active, +) from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -48,190 +53,377 @@ def _init_zigzag(self, cp_size=4): ) ) - def _metadata_for_rank(self, rank, cp_size=4, num_tokens=19): + def test_strategy_kind_maps_cli_values(self): + self.assertEqual(ContextParallelStrategyKind.NONE.value, 0) + self.assertEqual( + ContextParallelStrategyKind.from_string("zigzag"), + ContextParallelStrategyKind.ZIGZAG, + ) + self.assertEqual( + ContextParallelStrategyKind.from_string("interleave"), + ContextParallelStrategyKind.INTERLEAVE, + ) + self.assertEqual(ContextParallelStrategyKind.ZIGZAG.cli_value, "zigzag") + self.assertEqual(ContextParallelStrategyKind.INTERLEAVE.cli_value, "interleave") + + def test_get_cp_strategy_is_initialized_under_cp_v1_and_cp_v2(self): + self._init_zigzag() + + for enable_v2 in (False, True): + with patch( + "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", + return_value=enable_v2, + ): + self.assertIsNotNone(get_cp_strategy()) + self.assertTrue(is_cp_enabled()) + self.assertTrue(is_zigzag()) + self.assertFalse(is_interleave()) + self.assertEqual( + get_cp_strategy_kind(), ContextParallelStrategyKind.ZIGZAG + ) + + +class TestCPZigzagStrategy(CustomTestCase): + def setUp(self): + init_cp_strategy( + SimpleNamespace( + enable_prefill_cp=True, + cp_strategy="zigzag", + attn_cp_size=4, + attention_backend="fa3", + ) + ) + + def tearDown(self): + init_cp_strategy(SimpleNamespace(enable_prefill_cp=False)) + + def _metadata_for_rank(self, rank, *, cp_size, seq_lens, extend_seq_lens): strategy = ZigzagCPStrategy(cp_size=cp_size) with patch( "sglang.srt.layers.dp_attention.get_attention_cp_rank", return_value=rank, ): return strategy.build_metadata( - num_tokens=num_tokens, - seqs_len=[11, 13], - extend_seqs_len=[9, 10], + num_tokens=sum(extend_seq_lens), + seqs_len=seq_lens, + extend_seqs_len=extend_seq_lens, ) - def _forward_batch(self, metadata): + def _forward_batch(self, metadata, extend_seq_lens): return SimpleNamespace( + input_ids=torch.arange(sum(extend_seq_lens)), forward_mode=_ExtendMode(), - extend_seq_lens_cpu=[9, 10], + extend_seq_lens_cpu=extend_seq_lens, attn_cp_metadata=metadata, ) - def _padded_rank_tensors(self, x, cp_size=4): - per_rank = [] - metas = [] - for rank in range(cp_size): - metadata = self._metadata_for_rank(rank, cp_size=cp_size) - metas.append(metadata) - fb = self._forward_batch(metadata) - with patch( - "sglang.srt.layers.dp_attention.get_attention_cp_rank", - return_value=rank, - ): - local = ZigzagCPStrategy(cp_size=cp_size).shard_hidden_states(x, fb) - pad = metadata.max_rank_len[0] - local.shape[0] - if pad: - local = torch.nn.functional.pad( - local, - [0, 0] * (local.ndim - 1) + [0, pad], - ) - per_rank.append(local) - return metas, per_rank - - def test_strategy_kind_maps_cli_values(self): - self.assertEqual(ContextParallelStrategyKind.NONE.value, 0) - self.assertEqual( - ContextParallelStrategyKind.from_string("zigzag"), - ContextParallelStrategyKind.ZIGZAG, + def test_enable_cp_v2_and_is_cp_v2_active(self): + active_batch = SimpleNamespace( + input_ids=torch.arange(8), + forward_mode=_ExtendMode(), + extend_seq_lens_cpu=[8], ) - self.assertEqual( - ContextParallelStrategyKind.from_string("interleave"), - ContextParallelStrategyKind.INTERLEAVE, + inactive_batch = SimpleNamespace( + input_ids=torch.arange(7), + forward_mode=_ExtendMode(), + extend_seq_lens_cpu=[7], ) - self.assertEqual(ContextParallelStrategyKind.ZIGZAG.cli_value, "zigzag") - self.assertEqual(ContextParallelStrategyKind.INTERLEAVE.cli_value, "interleave") - - def test_cp_v2_strategy_is_only_active_when_env_enabled(self): - self._init_zigzag() with patch( "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=False ): - self.assertIsNone(get_cp_strategy()) - self.assertFalse(is_cp_enabled()) + self.assertFalse(enable_cp_v2()) + self.assertFalse(is_cp_v2_active(active_batch)) with patch( "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=True ): - self.assertIsNotNone(get_cp_strategy()) - self.assertTrue(is_cp_enabled()) - self.assertTrue(is_zigzag()) - self.assertFalse(is_interleave()) - self.assertEqual(get_cp_strategy_kind(), ContextParallelStrategyKind.ZIGZAG) + self.assertTrue(enable_cp_v2()) + self.assertTrue(is_cp_v2_active(active_batch)) + self.assertFalse(is_cp_v2_active(inactive_batch)) + + def _expected_metadata(self, *, rank, cp_size, seq_lens, extend_seq_lens): + bs = len(extend_seq_lens) + cp_segment_num = cp_size * 2 + prefix_offsets = [ + max(int(seq_lens[i]) - int(extend_seq_lens[i]), 0) for i in range(bs) + ] + + per_seq_block_sizes = [] + split_list = [] + for length in extend_seq_lens: + base = length // cp_segment_num + rem = length % cp_segment_num + block_sizes = [ + base + 1 if block_id < rem else base + for block_id in range(cp_segment_num) + ] + per_seq_block_sizes.append(block_sizes) + split_list.extend(block_sizes) + + per_rank_actual_token = [ + sum( + block_sizes[rank_id] + block_sizes[cp_segment_num - 1 - rank_id] + for block_sizes in per_seq_block_sizes + ) + for rank_id in range(cp_size) + ] + max_rank_len = [max(per_rank_actual_token)] * cp_size - def test_zigzag_metadata_for_batched_sequences(self): - metadata = self._metadata_for_rank(rank=2) + zigzag_index = list(range(rank, rank + bs * cp_segment_num, cp_segment_num)) + zigzag_index += list( + range(cp_segment_num - rank - 1, bs * cp_segment_num, cp_segment_num) + ) + + cp_reverse_index = [] + 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)) + ) + + reverse_split_len = [] + for rank_id in range(cp_size): + for batch_id in range(bs): + reverse_split_len.append(per_seq_block_sizes[batch_id][rank_id]) + for batch_id in range(bs): + reverse_split_len.append( + per_seq_block_sizes[batch_id][cp_segment_num - 1 - rank_id] + ) - self.assertEqual(metadata.bs, 2) - self.assertEqual(metadata.total_seq_lens, 19) + kv_len_prev_list = [] + kv_len_next_list = [] + actual_seq_q_prev_list = [] + actual_seq_q_next_list = [] + for batch_id, block_sizes in enumerate(per_seq_block_sizes): + kv_len_prev_list.append( + prefix_offsets[batch_id] + sum(block_sizes[: rank + 1]) + ) + kv_len_next_list.append( + prefix_offsets[batch_id] + sum(block_sizes[: cp_segment_num - rank]) + ) + actual_seq_q_prev_list.append(block_sizes[rank]) + actual_seq_q_next_list.append(block_sizes[cp_segment_num - rank - 1]) + + return { + "bs": bs, + "total_seq_lens": sum(extend_seq_lens), + "split_list": split_list, + "zigzag_index": zigzag_index, + "per_rank_actual_token": per_rank_actual_token, + "max_rank_len": max_rank_len, + "reverse_split_len": reverse_split_len, + "cp_reverse_index": cp_reverse_index, + "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, + } + + def _assert_metadata_matches(self, metadata, expected): + self.assertEqual(metadata.bs, expected["bs"]) + self.assertEqual(metadata.total_seq_lens, expected["total_seq_lens"]) + self.assertEqual(metadata.split_list, expected["split_list"]) + self.assertEqual(metadata.zigzag_index, expected["zigzag_index"]) self.assertEqual( - metadata.split_list, - [2, 1, 1, 1, 1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 1, 1], + metadata.per_rank_actual_token, expected["per_rank_actual_token"] ) - self.assertEqual(metadata.zigzag_index, [2, 10, 5, 13]) - self.assertEqual(metadata.per_rank_actual_token, [6, 5, 4, 4]) - self.assertEqual(metadata.max_rank_len, [6, 6, 6, 6]) + self.assertEqual(metadata.max_rank_len, expected["max_rank_len"]) + self.assertEqual(metadata.reverse_split_len, expected["reverse_split_len"]) + self.assertEqual(metadata.cp_reverse_index, expected["cp_reverse_index"]) + self.assertEqual(metadata.kv_len_prev_list, expected["kv_len_prev_list"]) + self.assertEqual(metadata.kv_len_next_list, expected["kv_len_next_list"]) self.assertEqual( - metadata.reverse_split_len, - [2, 2, 1, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], + metadata.actual_seq_q_prev_list, expected["actual_seq_q_prev_list"] ) self.assertEqual( - metadata.cp_reverse_index, - [0, 4, 8, 12, 14, 10, 6, 2, 1, 5, 9, 13, 15, 11, 7, 3], + metadata.actual_seq_q_next_list, expected["actual_seq_q_next_list"] ) - self.assertEqual(metadata.kv_len_prev_list, [6, 8]) - self.assertEqual(metadata.kv_len_next_list, [9, 11]) - self.assertEqual(metadata.actual_seq_q_prev_list, [1, 1]) - self.assertEqual(metadata.actual_seq_q_next_list, [1, 1]) - self.assertEqual(metadata.total_q_prev_tokens, 2) - self.assertEqual(metadata.total_q_next_tokens, 2) - self.assertTrue( - torch.equal( - metadata.cu_seqlens_q_prev_tensor.cpu(), torch.tensor([0, 1, 2]) - ) + self.assertEqual( + metadata.cu_seqlens_q_prev_tensor.cpu().tolist(), + [0] + + list( + torch.tensor(expected["actual_seq_q_prev_list"]).cumsum(dim=0).tolist() + ), ) - self.assertTrue( - torch.equal( - metadata.cu_seqlens_q_next_tensor.cpu(), torch.tensor([0, 1, 2]) - ) + self.assertEqual( + metadata.cu_seqlens_q_next_tensor.cpu().tolist(), + [0] + + list( + torch.tensor(expected["actual_seq_q_next_list"]).cumsum(dim=0).tolist() + ), ) + def _padded_rank_tensors(self, x, *, cp_size, seq_lens, extend_seq_lens): + per_rank = [] + metas = [] + for rank in range(cp_size): + metadata = self._metadata_for_rank( + rank, + cp_size=cp_size, + seq_lens=seq_lens, + extend_seq_lens=extend_seq_lens, + ) + metas.append(metadata) + fb = self._forward_batch(metadata, extend_seq_lens) + with patch( + "sglang.srt.layers.dp_attention.get_attention_cp_rank", + return_value=rank, + ): + local = ZigzagCPStrategy(cp_size=cp_size).shard_hidden_states(x, fb) + pad = metadata.max_rank_len[0] - local.shape[0] + if pad: + local = torch.nn.functional.pad( + local, + [0, 0] * (local.ndim - 1) + [0, pad], + ) + per_rank.append(local) + return metas, per_rank + + def test_zigzag_metadata_for_batched_sequences(self): + cases = [ + (4, [11, 13], [9, 10]), + (2, [8], [8]), + (4, [100000, 200000, 80], [100000, 200000, 64]), + (4, [100005, 200011, 25], [100000, 200000, 16]), + ] + + for cp_size, seq_lens, extend_seq_lens in cases: + for rank in range(cp_size): + with self.subTest( + cp_size=cp_size, + rank=rank, + seq_lens=seq_lens, + extend_seq_lens=extend_seq_lens, + ): + metadata = self._metadata_for_rank( + rank, + cp_size=cp_size, + seq_lens=seq_lens, + extend_seq_lens=extend_seq_lens, + ) + expected = self._expected_metadata( + rank=rank, + cp_size=cp_size, + seq_lens=seq_lens, + extend_seq_lens=extend_seq_lens, + ) + self._assert_metadata_matches(metadata, expected) + def test_zigzag_shards_hidden_states_and_position_ids(self): - metadata = self._metadata_for_rank(rank=2) - fb = self._forward_batch(metadata) - x = torch.arange(19 * 2).view(19, 2) - positions = torch.arange(19) + cp_size = 4 + seq_lens = [11, 13] + extend_seq_lens = [9, 10] + x = torch.arange(sum(extend_seq_lens) * 2).view(sum(extend_seq_lens), 2) + positions = torch.arange(sum(extend_seq_lens)) + + for rank in range(cp_size): + metadata = self._metadata_for_rank( + rank, + cp_size=cp_size, + seq_lens=seq_lens, + extend_seq_lens=extend_seq_lens, + ) + fb = self._forward_batch(metadata, extend_seq_lens) + strategy = ZigzagCPStrategy(cp_size=cp_size) + chunks = torch.split(x, metadata.split_list, dim=0) + position_chunks = torch.split(positions, metadata.split_list, dim=-1) + expected_x = torch.cat([chunks[i] for i in metadata.zigzag_index], dim=0) + expected_positions = torch.cat( + [position_chunks[i] for i in metadata.zigzag_index], dim=-1 + ) - with patch( - "sglang.srt.layers.dp_attention.get_attention_cp_rank", - return_value=2, - ): - strategy = ZigzagCPStrategy(cp_size=4) local_x = strategy.shard_hidden_states(x, fb) local_positions = strategy.shard_position_ids(positions, fb) + helper_x, helper_positions = cp_split_before_forward( + x, + positions, + fb, + ) - expected_indices = [3, 13, 6, 16] - self.assertTrue(torch.equal(local_x, x[expected_indices])) - self.assertTrue(torch.equal(local_positions, positions[expected_indices])) + self.assertTrue(torch.equal(local_x, expected_x)) + self.assertTrue(torch.equal(local_positions, expected_positions)) + self.assertTrue(torch.equal(helper_x, expected_x)) + self.assertTrue(torch.equal(helper_positions, expected_positions)) def test_zigzag_gathers_hidden_states_to_original_order(self): cp_size = 4 - x = torch.arange(19 * 2).view(19, 2) - metas, padded_rank_tensors = self._padded_rank_tensors(x, cp_size=cp_size) - rank = 2 - local_x = padded_rank_tensors[rank][: metas[rank].per_rank_actual_token[rank]] - fb = self._forward_batch(metas[rank]) - - with ( - patch( - "sglang.srt.layers.cp.zigzag.get_attention_cp_group", - return_value=_FakeCPGroup(padded_rank_tensors), - ), - patch( - "sglang.srt.distributed.device_communicators.pynccl_allocator.use_symmetric_memory", - return_value=torch.no_grad(), - ), - ): - gathered = ZigzagCPStrategy(cp_size=cp_size).gather_hidden_states( - local_x, fb, stream=None - ) + seq_lens = [11, 13] + extend_seq_lens = [9, 10] + x = torch.arange(sum(extend_seq_lens) * 2).view(sum(extend_seq_lens), 2) + metas, padded_rank_tensors = self._padded_rank_tensors( + x, + cp_size=cp_size, + seq_lens=seq_lens, + extend_seq_lens=extend_seq_lens, + ) - self.assertTrue(torch.equal(gathered, x)) + for rank in range(cp_size): + local_x = padded_rank_tensors[rank][ + : metas[rank].per_rank_actual_token[rank] + ] + fb = self._forward_batch(metas[rank], extend_seq_lens) + + with ( + patch( + "sglang.srt.layers.cp.zigzag.get_attention_cp_group", + return_value=_FakeCPGroup(padded_rank_tensors), + ), + patch( + "sglang.srt.distributed.device_communicators.pynccl_allocator.use_symmetric_memory", + return_value=torch.no_grad(), + ), + ): + gathered = ZigzagCPStrategy(cp_size=cp_size).gather_hidden_states( + local_x, fb, stream=None + ) + + self.assertTrue(torch.equal(gathered, x)) def test_zigzag_gathers_kv_cache_to_original_order(self): cp_size = 4 - kv = torch.arange(19 * 2 * 3).view(19, 2, 3) - metas, padded_rank_tensors = self._padded_rank_tensors(kv, cp_size=cp_size) - rank = 1 - local_kv = padded_rank_tensors[rank][: metas[rank].per_rank_actual_token[rank]] - fb = self._forward_batch(metas[rank]) - - with ( - patch( - "sglang.srt.layers.cp.zigzag.get_attention_cp_group", - return_value=_FakeCPGroup(padded_rank_tensors), - ), - patch( - "sglang.srt.distributed.device_communicators.pynccl_allocator.use_symmetric_memory", - return_value=torch.no_grad(), - ), - ): - gathered = ZigzagCPStrategy(cp_size=cp_size).gather_kv_cache( - local_kv, fb, stream=None - ) + seq_lens = [11, 13] + extend_seq_lens = [9, 10] + kv = torch.arange(sum(extend_seq_lens) * 2 * 3).view(sum(extend_seq_lens), 2, 3) + metas, padded_rank_tensors = self._padded_rank_tensors( + kv, + cp_size=cp_size, + seq_lens=seq_lens, + extend_seq_lens=extend_seq_lens, + ) + + for rank in range(cp_size): + local_kv = padded_rank_tensors[rank][ + : metas[rank].per_rank_actual_token[rank] + ] + fb = self._forward_batch(metas[rank], extend_seq_lens) + + with ( + patch( + "sglang.srt.layers.cp.zigzag.get_attention_cp_group", + return_value=_FakeCPGroup(padded_rank_tensors), + ), + patch( + "sglang.srt.distributed.device_communicators.pynccl_allocator.use_symmetric_memory", + return_value=torch.no_grad(), + ), + ): + gathered = ZigzagCPStrategy(cp_size=cp_size).gather_kv_cache( + local_kv, fb, stream=None + ) - self.assertTrue(torch.equal(gathered, kv)) + self.assertTrue(torch.equal(gathered, kv)) def test_zigzag_attention_dispatch_runs_prev_then_next(self): - with patch( - "sglang.srt.layers.dp_attention.get_attention_cp_rank", - return_value=0, - ): - metadata = ZigzagCPStrategy(cp_size=2).build_metadata( - num_tokens=8, - seqs_len=[8], - extend_seqs_len=[8], - ) + cp_size = 2 + seq_lens = [8] + extend_seq_lens = [8] + metadata = self._metadata_for_rank( + 0, + cp_size=cp_size, + seq_lens=seq_lens, + extend_seq_lens=extend_seq_lens, + ) fb = SimpleNamespace(attn_cp_metadata=metadata) q = torch.arange(4 * 2).view(4, 2) calls = [] @@ -247,7 +439,7 @@ def attn_fn(q_chunk, cu_seqlens_q, cache_seqlens, max_seqlen_q): ) return q_chunk + 100 - out = ZigzagCPStrategy(cp_size=2).run_attention( + out = ZigzagCPStrategy(cp_size=cp_size).run_attention( q, fb, device=torch.device("cpu"), attn_fn=attn_fn ) diff --git a/test/registered/cp/test_qwen3_30b.py b/test/registered/cp/test_gqa_preill_cp.py similarity index 88% rename from test/registered/cp/test_qwen3_30b.py rename to test/registered/cp/test_gqa_preill_cp.py index 1cccc14adcd5..e54f728f0d86 100644 --- a/test/registered/cp/test_qwen3_30b.py +++ b/test/registered/cp/test_gqa_preill_cp.py @@ -11,17 +11,17 @@ popen_launch_server, ) -register_cuda_ci(est_time=261, stage="extra-b", runner_config="4-gpu-h100") +register_cuda_ci(est_time=400, stage="extra-b", runner_config="4-gpu-h100") -QWEN3_30B_MODEL_PATH = "Qwen/Qwen3-30B-A3B-FP8" +GQA_MODEL_PATH = "Qwen/Qwen3-30B-A3B-FP8" -GSM8K_BASELINE_ACCURACY = 0.85 +GSM8K_BASELINE_ACCURACY = 0.93 -class TestQwen330B(CustomTestCase): +class TestGQACP2TP2EP2(CustomTestCase): @classmethod def setUpClass(cls): - cls.model = QWEN3_30B_MODEL_PATH + cls.model = GQA_MODEL_PATH cls.base_url = DEFAULT_URL_FOR_TEST cls.process = popen_launch_server( cls.model, @@ -46,11 +46,13 @@ def setUpClass(cls): "--model-loader-extra-config", '{"enable_multithread_load": true, "num_threads": 64}', ], + env={"SGLANG_ENABLE_CP_V2": "0"}, ) @classmethod def tearDownClass(cls): - kill_process_tree(cls.process.pid) + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) def test_gsm8k(self): args = SimpleNamespace( @@ -73,10 +75,10 @@ def test_gsm8k(self): self.assertGreaterEqual(metrics["score"], GSM8K_BASELINE_ACCURACY) -class TestQwen330BCP(CustomTestCase): +class TestGQACPTP2CP2EP4(CustomTestCase): @classmethod def setUpClass(cls): - cls.model = QWEN3_30B_MODEL_PATH + cls.model = GQA_MODEL_PATH cls.base_url = DEFAULT_URL_FOR_TEST cls.process = popen_launch_server( cls.model, @@ -101,11 +103,13 @@ def setUpClass(cls): "--model-loader-extra-config", '{"enable_multithread_load": true, "num_threads": 64}', ], + env={"SGLANG_ENABLE_CP_V2": "0"}, ) @classmethod def tearDownClass(cls): - kill_process_tree(cls.process.pid) + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) def test_gsm8k(self): args = SimpleNamespace( @@ -128,10 +132,10 @@ def test_gsm8k(self): self.assertGreaterEqual(metrics["score"], GSM8K_BASELINE_ACCURACY) -class TestQwen330BCPV2ZigzagDeepEP(CustomTestCase): +class TestGQACPCP4EP4(CustomTestCase): @classmethod def setUpClass(cls): - cls.model = QWEN3_30B_MODEL_PATH + cls.model = GQA_MODEL_PATH cls.base_url = DEFAULT_URL_FOR_TEST cls.process = popen_launch_server( cls.model, @@ -160,7 +164,6 @@ def setUpClass(cls): "--model-loader-extra-config", '{"enable_multithread_load": true, "num_threads": 64}', ], - env={"SGLANG_ENABLE_CP_V2": "1"}, ) @classmethod From 003008639382089ae2f9304ef08e24ab45b2dfb3 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Thu, 18 Jun 2026 01:31:59 -0700 Subject: [PATCH 3/5] Address CP v2 review comments --- python/sglang/srt/layers/cp/utils.py | 17 +-- .../sglang/srt/model_executor/model_runner.py | 97 +++++++------ python/sglang/srt/server_args.py | 19 +-- test/registered/cp/test_cp_strategy_unit.py | 80 ++++++----- .../cp/test_gqa_prefill_cp_legacy.py | 136 ++++++++++++++++++ test/registered/cp/test_gqa_preill_cp.py | 6 +- 6 files changed, 250 insertions(+), 105 deletions(-) create mode 100644 test/registered/cp/test_gqa_prefill_cp_legacy.py diff --git a/python/sglang/srt/layers/cp/utils.py b/python/sglang/srt/layers/cp/utils.py index e7561d99a0a8..21f980987a0f 100644 --- a/python/sglang/srt/layers/cp/utils.py +++ b/python/sglang/srt/layers/cp/utils.py @@ -68,14 +68,10 @@ def is_cp_v2_active(forward_batch) -> bool: 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) strategy = get_cp_strategy() assert strategy is not None - - input_ids = getattr(forward_batch, "input_ids", None) - assert input_ids is not None - - num_tokens = len(input_ids) - assert strategy.can_apply(num_tokens, forward_batch) + num_tokens = len(forward_batch.input_ids) seq_lens_cpu = _to_int_list(getattr(forward_batch, "seq_lens_cpu", None)) extend_lens_cpu = _to_int_list(getattr(forward_batch, "extend_seq_lens_cpu", None)) @@ -92,6 +88,7 @@ def cp_split_before_forward( forward_batch, ) -> Tuple[Optional[Any], Optional[Any]]: """Shard embeddings and positions for CP-v2 model-runner forwarding.""" + assert is_cp_v2_active(forward_batch) strategy = get_cp_strategy() assert strategy is not None assert complete_hidden_states is not None @@ -102,11 +99,9 @@ def cp_split_before_forward( ) -def maybe_cp_gather_after_forward(x: Any, forward_batch, stream: Optional[Any] = None): +def cp_gather_after_forward(x: Any, forward_batch, stream: Optional[Any] = None): """Gather CP-v2 hidden states at the model boundary when this batch is active.""" - if not is_cp_v2_active(forward_batch): - return x - + assert is_cp_v2_active(forward_batch) strategy = get_cp_strategy() assert strategy is not None @@ -142,7 +137,7 @@ def _to_int_list(values) -> Optional[list[int]]: "enable_cp_v2", "get_cp_strategy", "is_cp_v2_active", - "maybe_cp_gather_after_forward", + "cp_gather_after_forward", "cp_split_before_forward", "prepare_cp_forward", ] diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index f6b922e2d039..24b2ae4bb86a 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -124,10 +124,10 @@ from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp from sglang.srt.layers.attention.tbo_backend import TboAttnBackend from sglang.srt.layers.cp.utils import ( + cp_gather_after_forward, cp_split_before_forward, get_cp_strategy, is_cp_v2_active, - maybe_cp_gather_after_forward, prepare_cp_forward, ) from sglang.srt.layers.dp_attention import ( @@ -3415,55 +3415,17 @@ def forward_extend( kwargs["input_embeds"] = sharded_hidden_states forward_positions = sharded_positions - def _forward_model(): - if not cp_v2_active: - return self.model.forward( - forward_batch.input_ids, - forward_positions, - forward_batch, - **kwargs, - ) - - hidden_states = self.model.model( - forward_batch.input_ids, - forward_positions, - forward_batch, - input_embeds=kwargs.get("input_embeds"), - pp_proxy_tensors=kwargs.get("pp_proxy_tensors"), - ) - - aux_hidden_states = None - capture_aux_hidden_states = getattr( - self.model, "capture_aux_hidden_states", False - ) - if capture_aux_hidden_states: - hidden_states, aux_hidden_states = hidden_states - - if self.model.pp_group.is_last_rank: - hidden_states = maybe_cp_gather_after_forward( - hidden_states, - forward_batch, - torch.cuda.current_stream(), - ) - return self.model.logits_processor( - forward_batch.input_ids, - hidden_states, - self.model.lm_head, - forward_batch, - aux_hidden_states, - ) - - if capture_aux_hidden_states: - return hidden_states, aux_hidden_states - return hidden_states - ctx = ( self.device_timer.wrap(metadata={"category": "extend"}) if self.device_timer else contextlib.nullcontext() ) with ctx: - if _is_hip and self.prefill_cuda_graph_runner is not None: + if ( + _is_hip + and self.prefill_cuda_graph_runner is not None + and not cp_v2_active + ): # AMD/HIP: when PCG is enabled but the batch exceeds max captured # size, run eagerly under enable_tc_piecewise_cuda_graph() and # set_tc_piecewise_forward_context() so that (a) Dynamo guards on @@ -3481,9 +3443,52 @@ def _forward_model(): dsa_indexers=self.dsa_indexers, ), ): - ret = _forward_model() + ret = self.model.forward( + forward_batch.input_ids, + forward_positions, + forward_batch, + **kwargs, + ) + elif cp_v2_active: + hidden_states = self.model.model( + forward_batch.input_ids, + forward_positions, + forward_batch, + input_embeds=kwargs.get("input_embeds"), + pp_proxy_tensors=kwargs.get("pp_proxy_tensors"), + ) + + aux_hidden_states = None + capture_aux_hidden_states = getattr( + self.model, "capture_aux_hidden_states", False + ) + if capture_aux_hidden_states: + hidden_states, aux_hidden_states = hidden_states + + if self.model.pp_group.is_last_rank: + hidden_states = cp_gather_after_forward( + hidden_states, + forward_batch, + torch.cuda.current_stream(), + ) + ret = self.model.logits_processor( + forward_batch.input_ids, + hidden_states, + self.model.lm_head, + forward_batch, + aux_hidden_states, + ) + elif capture_aux_hidden_states: + ret = hidden_states, aux_hidden_states + else: + ret = hidden_states else: - ret = _forward_model() + ret = self.model.forward( + forward_batch.input_ids, + forward_positions, + forward_batch, + **kwargs, + ) return (ret, can_run_graph) def forward_idle( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index f98475b4cf43..3e5d29b3379b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1972,14 +1972,6 @@ def _handle_model_specific_adjustments(self): hf_config = model_config.hf_config model_arch = hf_config.architectures[0] - from sglang.srt.layers.cp.utils import CP_V2_DEFAULT_MODEL_CLASSES - - if ( - model_arch in CP_V2_DEFAULT_MODEL_CLASSES - and not envs.SGLANG_ENABLE_CP_V2.is_set() - ): - envs.SGLANG_ENABLE_CP_V2.set(True) - _hybrid_spec = get_linear_attn_spec_by_arch(model_arch) if _hybrid_spec is not None and _hybrid_spec.uses_mamba_radix_cache: self._handle_mamba_radix_cache( @@ -3504,6 +3496,17 @@ def _handle_legacy_cp_arguments(self): self.prefill_cp_mode = mode def _handle_context_parallelism(self): + if parse_connector_type(self.model_path) != ConnectorType.INSTANCE: + from sglang.srt.layers.cp.utils import CP_V2_DEFAULT_MODEL_CLASSES + + model_config = self.get_model_config() + model_arch = model_config.hf_config.architectures[0] + if ( + model_arch in CP_V2_DEFAULT_MODEL_CLASSES + and not envs.SGLANG_ENABLE_CP_V2.is_set() + ): + envs.SGLANG_ENABLE_CP_V2.set(True) + if self.enable_prefill_cp and self.cp_strategy is None: raise ValueError( "--cp-strategy must be set when --enable-prefill-cp is enabled." diff --git a/test/registered/cp/test_cp_strategy_unit.py b/test/registered/cp/test_cp_strategy_unit.py index 1b0ae8e6bf9e..270ab9feb749 100644 --- a/test/registered/cp/test_cp_strategy_unit.py +++ b/test/registered/cp/test_cp_strategy_unit.py @@ -19,6 +19,7 @@ is_cp_v2_active, ) from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy +from sglang.srt.runtime_context import get_parallel from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -43,16 +44,6 @@ class TestCPStrategyUnit(CustomTestCase): def tearDown(self): init_cp_strategy(SimpleNamespace(enable_prefill_cp=False)) - def _init_zigzag(self, cp_size=4): - init_cp_strategy( - SimpleNamespace( - enable_prefill_cp=True, - cp_strategy="zigzag", - attn_cp_size=cp_size, - attention_backend="fa3", - ) - ) - def test_strategy_kind_maps_cli_values(self): self.assertEqual(ContextParallelStrategyKind.NONE.value, 0) self.assertEqual( @@ -66,21 +57,40 @@ def test_strategy_kind_maps_cli_values(self): self.assertEqual(ContextParallelStrategyKind.ZIGZAG.cli_value, "zigzag") self.assertEqual(ContextParallelStrategyKind.INTERLEAVE.cli_value, "interleave") + def test_init_cp_strategy_binds_zigzag_strategy(self): + init_cp_strategy( + SimpleNamespace( + enable_prefill_cp=True, + cp_strategy="zigzag", + attn_cp_size=4, + ) + ) + + self.assertTrue(is_cp_enabled()) + self.assertTrue(is_zigzag()) + self.assertFalse(is_interleave()) + self.assertEqual(get_cp_strategy_kind(), ContextParallelStrategyKind.ZIGZAG) + def test_get_cp_strategy_is_initialized_under_cp_v1_and_cp_v2(self): - self._init_zigzag() + init_cp_strategy( + SimpleNamespace( + enable_prefill_cp=True, + cp_strategy="interleave", + attn_cp_size=4, + ) + ) - for enable_v2 in (False, True): - with patch( - "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", - return_value=enable_v2, - ): - self.assertIsNotNone(get_cp_strategy()) - self.assertTrue(is_cp_enabled()) - self.assertTrue(is_zigzag()) - self.assertFalse(is_interleave()) - self.assertEqual( - get_cp_strategy_kind(), ContextParallelStrategyKind.ZIGZAG - ) + with patch( + "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=False + ): + self.assertIsNotNone(get_cp_strategy()) + self.assertTrue(is_cp_enabled()) + self.assertTrue(is_interleave()) + + with patch( + "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=True + ): + self.assertIsNotNone(get_cp_strategy()) class TestCPZigzagStrategy(CustomTestCase): @@ -99,10 +109,7 @@ def tearDown(self): def _metadata_for_rank(self, rank, *, cp_size, seq_lens, extend_seq_lens): strategy = ZigzagCPStrategy(cp_size=cp_size) - with patch( - "sglang.srt.layers.dp_attention.get_attention_cp_rank", - return_value=rank, - ): + with get_parallel().override(attn_cp_rank=rank): return strategy.build_metadata( num_tokens=sum(extend_seq_lens), seqs_len=seq_lens, @@ -266,11 +273,7 @@ def _padded_rank_tensors(self, x, *, cp_size, seq_lens, extend_seq_lens): ) metas.append(metadata) fb = self._forward_batch(metadata, extend_seq_lens) - with patch( - "sglang.srt.layers.dp_attention.get_attention_cp_rank", - return_value=rank, - ): - local = ZigzagCPStrategy(cp_size=cp_size).shard_hidden_states(x, fb) + local = ZigzagCPStrategy(cp_size=cp_size).shard_hidden_states(x, fb) pad = metadata.max_rank_len[0] - local.shape[0] if pad: local = torch.nn.functional.pad( @@ -335,11 +338,14 @@ def test_zigzag_shards_hidden_states_and_position_ids(self): local_x = strategy.shard_hidden_states(x, fb) local_positions = strategy.shard_position_ids(positions, fb) - helper_x, helper_positions = cp_split_before_forward( - x, - positions, - fb, - ) + with patch( + "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=True + ): + helper_x, helper_positions = cp_split_before_forward( + x, + positions, + fb, + ) self.assertTrue(torch.equal(local_x, expected_x)) self.assertTrue(torch.equal(local_positions, expected_positions)) diff --git a/test/registered/cp/test_gqa_prefill_cp_legacy.py b/test/registered/cp/test_gqa_prefill_cp_legacy.py new file mode 100644 index 000000000000..925e54184892 --- /dev/null +++ b/test/registered/cp/test_gqa_prefill_cp_legacy.py @@ -0,0 +1,136 @@ +import unittest +from types import SimpleNamespace + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + kill_process_tree, + popen_launch_server, +) + +register_cuda_ci(est_time=260, stage="extra-b", runner_config="4-gpu-h100") + +GQA_MODEL_PATH = "Qwen/Qwen3-30B-A3B-FP8" + +GSM8K_BASELINE_ACCURACY = 0.93 + + +class TestGQACP2TP2EP2(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = GQA_MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp-size", + "4", + "--moe-dp-size", + "2", + "--ep-size", + "2", + "--attn-cp-size", + "2", + "--enable-prefill-context-parallel", + "--cuda-graph-max-bs", + "32", + "--max-running-requests", + "32", + "--trust-remote-code", + "--disable-piecewise-cuda-graph", + "--model-loader-extra-config", + '{"enable_multithread_load": true, "num_threads": 64}', + ], + env={"SGLANG_ENABLE_CP_V2": "0"}, + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + model=self.model, + eval_name="gsm8k", + num_shots=5, + num_examples=200, + max_tokens=16000, + num_threads=128, + repeat=1, + temperature=0.6, + top_p=0.95, + top_k=20, + base_url=self.base_url, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreaterEqual(metrics["score"], GSM8K_BASELINE_ACCURACY) + + +class TestGQACPTP2CP2EP4(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = GQA_MODEL_PATH + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp-size", + "4", + "--moe-dp-size", + "1", + "--ep-size", + "4", + "--attn-cp-size", + "2", + "--enable-prefill-context-parallel", + "--cuda-graph-max-bs", + "32", + "--max-running-requests", + "32", + "--trust-remote-code", + "--disable-piecewise-cuda-graph", + "--model-loader-extra-config", + '{"enable_multithread_load": true, "num_threads": 64}', + ], + env={"SGLANG_ENABLE_CP_V2": "0"}, + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + model=self.model, + eval_name="gsm8k", + num_shots=5, + num_examples=200, + max_tokens=16000, + num_threads=128, + repeat=1, + temperature=0.6, + top_p=0.95, + top_k=20, + base_url=self.base_url, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreaterEqual(metrics["score"], GSM8K_BASELINE_ACCURACY) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/cp/test_gqa_preill_cp.py b/test/registered/cp/test_gqa_preill_cp.py index e54f728f0d86..c1cb25829edf 100644 --- a/test/registered/cp/test_gqa_preill_cp.py +++ b/test/registered/cp/test_gqa_preill_cp.py @@ -11,7 +11,7 @@ popen_launch_server, ) -register_cuda_ci(est_time=400, stage="extra-b", runner_config="4-gpu-h100") +register_cuda_ci(est_time=500, stage="extra-b", runner_config="4-gpu-h100") GQA_MODEL_PATH = "Qwen/Qwen3-30B-A3B-FP8" @@ -46,7 +46,7 @@ def setUpClass(cls): "--model-loader-extra-config", '{"enable_multithread_load": true, "num_threads": 64}', ], - env={"SGLANG_ENABLE_CP_V2": "0"}, + env={"SGLANG_ENABLE_CP_V2": "1"}, ) @classmethod @@ -103,7 +103,7 @@ def setUpClass(cls): "--model-loader-extra-config", '{"enable_multithread_load": true, "num_threads": 64}', ], - env={"SGLANG_ENABLE_CP_V2": "0"}, + env={"SGLANG_ENABLE_CP_V2": "1"}, ) @classmethod From 02c77b7f70e520f9c0c00245b214dd79e0c4ab32 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Thu, 18 Jun 2026 12:58:17 -0700 Subject: [PATCH 4/5] Fix server args CP unit test --- test/registered/unit/server_args/test_server_args.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 92a954a7c10d..53f5b8908c6f 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -159,6 +159,7 @@ def _new_cp_args(self, **overrides): 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, From 58819f81f36b0e55a7314c14b99823c4a3ed371b Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Thu, 18 Jun 2026 14:53:56 -0700 Subject: [PATCH 5/5] Use CP v2 flags in GQA CP test --- test/registered/cp/test_gqa_preill_cp.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/test/registered/cp/test_gqa_preill_cp.py b/test/registered/cp/test_gqa_preill_cp.py index c1cb25829edf..d19847ab58e6 100644 --- a/test/registered/cp/test_gqa_preill_cp.py +++ b/test/registered/cp/test_gqa_preill_cp.py @@ -36,7 +36,9 @@ def setUpClass(cls): "2", "--attn-cp-size", "2", - "--enable-prefill-context-parallel", + "--enable-prefill-cp", + "--cp-strategy", + "zigzag", "--cuda-graph-max-bs", "32", "--max-running-requests", @@ -93,7 +95,9 @@ def setUpClass(cls): "4", "--attn-cp-size", "2", - "--enable-prefill-context-parallel", + "--enable-prefill-cp", + "--cp-strategy", + "zigzag", "--cuda-graph-max-bs", "32", "--max-running-requests",