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

[CP V1 Deprecation 1/5] Migrate tests to strategy-based prefill CP (#36222)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-08-24 20:15:34 -07:00
committed by GitHub
parent a0f52eca98
commit 833be86c15
20 changed files with 37 additions and 558 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/dsv4/test_b200_flash.py
+3 -3
View File
@@ -87,9 +87,9 @@ class TestB200FlashCP(DSV4FlashAime25TestBase):
"4",
"--moe-a2a-backend",
"deepep",
"--enable-dsa-prefill-context-parallel",
"--dsa-prefill-cp-mode",
"round-robin-split",
"--enable-prefill-cp",
"--cp-strategy",
"interleave",
"--chunked-prefill-size",
"16384",
"--mem-fraction-static",
test/manual/dsv4/test_b200_pro.py
+3 -3
View File
@@ -101,9 +101,9 @@ class TestB200ProCP(DSV4ProAime25TestBase):
"8",
"--moe-a2a-backend",
"deepep",
"--enable-dsa-prefill-context-parallel",
"--dsa-prefill-cp-mode",
"round-robin-split",
"--enable-prefill-cp",
"--cp-strategy",
"interleave",
"--chunked-prefill-size",
"16384",
"--mem-fraction-static",
test/manual/dsv4/test_b300_flash.py
+3 -3
View File
@@ -89,9 +89,9 @@ class TestB300FlashCP(DSV4FlashAime25TestBase):
"4",
"--moe-a2a-backend",
"deepep",
"--enable-dsa-prefill-context-parallel",
"--dsa-prefill-cp-mode",
"round-robin-split",
"--enable-prefill-cp",
"--cp-strategy",
"interleave",
"--chunked-prefill-size",
"16384",
"--mem-fraction-static",
test/manual/dsv4/test_b300_pro.py
+3 -3
View File
@@ -103,9 +103,9 @@ class TestB300ProCP(DSV4ProAime25TestBase):
"8",
"--moe-a2a-backend",
"deepep",
"--enable-dsa-prefill-context-parallel",
"--dsa-prefill-cp-mode",
"round-robin-split",
"--enable-prefill-cp",
"--cp-strategy",
"interleave",
"--chunked-prefill-size",
"16384",
"--mem-fraction-static",
test/manual/dsv4/test_gb300_flash.py
+3 -3
View File
@@ -87,9 +87,9 @@ class TestGB300FlashCP(DSV4FlashAime25TestBase):
"4",
"--moe-a2a-backend",
"deepep",
"--enable-dsa-prefill-context-parallel",
"--dsa-prefill-cp-mode",
"round-robin-split",
"--enable-prefill-cp",
"--cp-strategy",
"interleave",
"--chunked-prefill-size",
"16384",
"--mem-fraction-static",
test/manual/dsv4/test_gb300_pro.py
+3 -3
View File
@@ -103,9 +103,9 @@ class TestGB300ProCP(DSV4ProAime25TestBase):
"4",
"--moe-a2a-backend",
"deepep",
"--enable-dsa-prefill-context-parallel",
"--dsa-prefill-cp-mode",
"round-robin-split",
"--enable-prefill-cp",
"--cp-strategy",
"interleave",
"--chunked-prefill-size",
"16384",
"--mem-fraction-static",
test/manual/dsv4/test_h200_fp8_flash.py
+3 -3
View File
@@ -98,9 +98,9 @@ class TestH200Fp8FlashCP(DSV4FlashAime25TestBase):
"4",
"--moe-a2a-backend",
"deepep",
"--enable-dsa-prefill-context-parallel",
"--dsa-prefill-cp-mode",
"round-robin-split",
"--enable-prefill-cp",
"--cp-strategy",
"interleave",
"--chunked-prefill-size",
"16384",
"--mem-fraction-static",
test/manual/test_kv_events.py
+3 -1
View File
@@ -329,7 +329,9 @@ class TestKvEvents(CustomTestCase):
2,
"--moe-dp-size",
2,
"--enable-prefill-context-parallel",
"--enable-prefill-cp",
"--cp-strategy",
"zigzag",
"--trust-remote-code",
"--max-total-tokens",
4096,
test/manual/test_qwen3_235b.py
+2 -1
View File
@@ -80,7 +80,8 @@ class TestQwen3235BFP8(unittest.TestCase):
"--moe-dp-size=2",
"--attn-cp-size=2",
"--ep-size=4",
"--enable-prefill-context-parallel",
"--enable-prefill-cp",
"--cp-strategy=zigzag",
]
MTP_ARGS = [
test/registered/cp/test_cp_strategy_unit.py
+2 -13
View File
@@ -89,7 +89,7 @@ class TestCPStrategyUnit(CustomTestCase):
self.assertFalse(is_interleave())
self.assertEqual(get_cp_strategy_kind(), ContextParallelStrategyKind.ZIGZAG)
def test_get_cp_strategy_is_initialized_under_cp_v1_and_cp_v2(self):
def test_get_cp_strategy_is_initialized_under_cp_v2(self):
init_cp_strategy(
SimpleNamespace(
enable_prefill_cp=True,
@@ -99,17 +99,12 @@ class TestCPStrategyUnit(CustomTestCase):
)
with patch(
"sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=False
"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_interleave())
with patch(
"sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=True
):
self.assertIsNotNone(get_cp_strategy())
class TestPrefillCPBCGReplay(CustomTestCase):
def tearDown(self):
@@ -312,12 +307,6 @@ class TestCPZigzagStrategy(CustomTestCase):
extend_seq_lens_cpu=[7],
)
with patch(
"sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=False
):
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
):
test/registered/cp/test_deepseek_v3_cp_single_node.py
+4 -2
View File
@@ -37,7 +37,9 @@ class TestDeepseekV3CPInSeqSplit(CustomTestCase):
"8",
"--dp",
"2",
"--enable-prefill-context-parallel",
"--enable-prefill-cp",
"--cp-strategy",
"zigzag",
"--attention-backend",
"fa3",
"--mem-frac",
@@ -79,7 +81,7 @@ class TestDeepseekV3CPInSeqSplit(CustomTestCase):
if is_in_ci():
write_github_step_summary(
f"### test_a_gsm8k (deepseek-v3-mla-cp-in-seq-split)\n"
f"### test_a_gsm8k (deepseek-v3-mla-cp-zigzag)\n"
f'{metrics["score"]=:.3f}\n'
)
self.assertGreater(metrics["score"], GSM8K_ACCURACY_THRESHOLD)
test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py
-4
View File
@@ -29,7 +29,6 @@ SERVER_LAUNCH_TIMEOUT = 3600
DEEPEP_CONFIG = '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}'
_DEEPEP_ENV = {
"SGLANG_ENABLE_CP_V2": "1",
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
# The draft-extend graph pool costs ~4.5 GB here (DeepEP MoE workspace is
# captured at full dispatch capacity), which starves the eager prefill
@@ -39,7 +38,6 @@ _DEEPEP_ENV = {
}
_MEGAMOE_ENV = {
"SGLANG_ENABLE_CP_V2": "1",
"SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK": "8320",
}
@@ -188,7 +186,6 @@ class TestDSV4FlashFP4B200Balanced_CP_NonDeepEP(
"--moe-runner-backend", # for fp4 checkpoint
"flashinfer_mxfp4",
],
env={"SGLANG_ENABLE_CP_V2": "1"},
)
@classmethod
@@ -232,7 +229,6 @@ class TestDSV4FlashFP4B200_CP_DSpark(
"--moe-runner-backend", # for fp4 checkpoint
"flashinfer_mxfp4",
],
env={"SGLANG_ENABLE_CP_V2": "1"},
)
@classmethod
test/registered/cp/test_dsa_prefill_cp.py
-1
View File
@@ -54,7 +54,6 @@ class TestDSACPV2Interleave(CustomTestCase):
cls.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
other_args=other_args,
env={"SGLANG_ENABLE_CP_V2": "1"},
)
@classmethod
test/registered/cp/test_dsa_prefill_cp_legacy.py
-157
View File
@@ -1,157 +0,0 @@
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
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,
is_in_ci,
popen_launch_server,
write_github_step_summary,
)
register_cuda_ci(est_time=616, stage="extra-b", runner_config="8-gpu-h200")
DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2"
class TestDeepseekV32CPInSeqSplit(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEEPSEEK_V32_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--tp",
"8",
"--enable-dp-attention",
"--dp",
"2",
"--attn-cp-size",
"4",
"--enable-dsa-prefill-context-parallel",
"--dsa-prefill-cp-mode",
"in-seq-split",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--mem-frac",
"0.85",
"--cuda-graph-max-bs-decode",
"32",
"--max-running-requests",
"32",
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 64}',
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=500,
num_threads=32,
num_shots=20,
)
metrics = run_eval(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_a_gsm8k (deepseek-v32-cp-in-seq-split)\n"
f'{metrics["score"]=:.3f}\n'
)
self.assertGreater(metrics["score"], 0.935)
class TestDeepseekV32CPRoundRobinSplit(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEEPSEEK_V32_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--tp",
"8",
"--attn-cp-size",
"8",
"--enable-dsa-prefill-context-parallel",
"--dsa-prefill-cp-mode",
"round-robin-split",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--mem-frac",
"0.85",
"--cuda-graph-max-bs-decode",
"32",
"--max-running-requests",
"32",
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 64}',
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
env={"SGLANG_ENABLE_CP_V2": "0"},
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=500,
num_threads=32,
num_shots=20,
)
metrics = run_eval(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_a_gsm8k (deepseek-v32-cp-in-seq-split)\n"
f'{metrics["score"]=:.3f}\n'
)
self.assertGreater(metrics["score"], 0.935)
if __name__ == "__main__":
unittest.main()
test/registered/cp/test_gqa_prefill_cp.py
-2
View File
@@ -48,7 +48,6 @@ class TestGQACP2TP2EP2(CustomTestCase):
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 64}',
],
env={"SGLANG_ENABLE_CP_V2": "1"},
)
@classmethod
@@ -107,7 +106,6 @@ class TestGQACPTP2CP2EP4(CustomTestCase):
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 64}',
],
env={"SGLANG_ENABLE_CP_V2": "1"},
)
@classmethod
test/registered/hicache/test_hicache_storage_mooncake_backend.py
+2 -1
View File
@@ -279,7 +279,8 @@ class TestMooncakeBackendQwen330BCP2(
"--tp-size": 2,
"--moe-dp-size": 2,
"--attn-cp-size": 2,
"--enable-prefill-context-parallel": True,
"--enable-prefill-cp": True,
"--cp-strategy": "zigzag",
"--trust-remote-code": True,
"--cuda-graph-max-bs-decode": 32,
"--max-running-requests": 32,
test/registered/kernels/ops/attention/test_cp_prefix_len_fa3_parity.py
-143
View File
@@ -1,143 +0,0 @@
"""
FA3 parity test for `prepare_context_parallel_metadata`.
Drives the real function and feeds its `kv_len_prev/next_tensor` into FA3
via `flash_attn_with_kvcache`. Compares per-rank CP output against a
full-sequence FA3 reference computed over the unpadded `(prefix + extend)`
KV. Any discrepancy indicates the metadata function emitted wrong
`cache_seqlens` for at least one rank.
"""
import unittest
from unittest.mock import patch
import torch
from sglang.srt.layers.utils.cp_utils import prepare_context_parallel_metadata
from sglang.srt.utils.common import ceil_align
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=5, stage="extra-a", runner_config="1-gpu-large")
_DSA_UTILS = "sglang.srt.layers.attention.dsa.utils"
_DEVICE = "cuda"
_DTYPE = torch.bfloat16
_HEAD_NUM = 8
_HEAD_DIM = 128
_SCALE = _HEAD_DIM**-0.5
class TestCPPrefixLenFA3Parity(CustomTestCase):
"""Per-rank FA3 output under CP must match a full-sequence reference."""
def _run_parity(self, prefix_len: int, extend_len: int, cp_size: int):
from sgl_kernel.flash_attn import flash_attn_with_kvcache
torch.manual_seed(extend_len * 1_000_003 + prefix_len * 101 + cp_size)
padded_extend = ceil_align(extend_len, cp_size)
pad = padded_extend - extend_len
self.assertGreaterEqual(
padded_extend,
2 * cp_size,
"runtime `can_cp_split` would skip this case; pick a larger extend",
)
# Reference: one full-sequence FA3 call over the unpadded KV.
q_full = torch.randn(
extend_len, _HEAD_NUM, _HEAD_DIM, device=_DEVICE, dtype=_DTYPE
)
k_full = torch.randn(
prefix_len + extend_len, _HEAD_NUM, _HEAD_DIM, device=_DEVICE, dtype=_DTYPE
)
v_full = torch.randn(
prefix_len + extend_len, _HEAD_NUM, _HEAD_DIM, device=_DEVICE, dtype=_DTYPE
)
ref = flash_attn_with_kvcache(
q=q_full.unsqueeze(0),
k_cache=k_full.unsqueeze(0),
v_cache=v_full.unsqueeze(0),
cache_seqlens=torch.tensor(
[k_full.shape[0]], dtype=torch.int32, device=_DEVICE
),
softmax_scale=_SCALE,
causal=True,
).squeeze(0)
# CP path sees tensors padded to `ceil_align(extend, cp_size)`,
# matching what `prepare_mlp_sync_batch` does in production.
zeros = torch.zeros(pad, _HEAD_NUM, _HEAD_DIM, device=_DEVICE, dtype=_DTYPE)
q_padded = torch.cat([q_full, zeros], dim=0)
k_padded = torch.cat([k_full, zeros], dim=0)
v_padded = torch.cat([v_full, zeros], dim=0)
seqs_len = [prefix_len + extend_len]
extend_lens = [extend_len]
def _call_meta(rank: int):
return prepare_context_parallel_metadata(
padded_extend, rank, cp_size, seqs_len, extend_seqs_len=extend_lens
)
# Exercise the non-DSA branch; the DSA branch uses a separate
# `prefix_len` pathway re-added by `_get_topk_ragged_with_cp`.
with (
patch(f"{_DSA_UTILS}.is_dsa_enable_prefill_cp", return_value=False),
patch(
f"{_DSA_UTILS}.is_dsa_prefill_cp_round_robin_split",
return_value=False,
),
):
meta0 = _call_meta(0)
cp_segment_num = 2 * cp_size
blocks_q = list(torch.split(q_padded, meta0.split_list, dim=0))
outs = [None] * cp_segment_num
for rank in range(cp_size):
meta = meta0 if rank == 0 else _call_meta(rank)
for idx, cs_tensor in (
(rank, meta.kv_len_prev_tensor),
(cp_size * 2 - rank - 1, meta.kv_len_next_tensor),
):
if meta0.split_list[idx] == 0:
outs[idx] = torch.empty(
0, _HEAD_NUM, _HEAD_DIM, device=_DEVICE, dtype=_DTYPE
)
continue
outs[idx] = flash_attn_with_kvcache(
q=blocks_q[idx].unsqueeze(0),
k_cache=k_padded.unsqueeze(0),
v_cache=v_padded.unsqueeze(0),
cache_seqlens=cs_tensor,
softmax_scale=_SCALE,
causal=True,
).squeeze(0)
cp_out = torch.cat(outs, dim=0)
err = (cp_out[:extend_len].float() - ref.float()).abs().max().item()
self.assertLess(
err,
1e-2,
f"CP output diverges from full-sequence FA3 reference by "
f"max_err={err:.5f} "
f"(prefix_len={prefix_len}, extend_len={extend_len}, "
f"cp_size={cp_size}, pad={pad})",
)
def test_cp2_prefix1_extend3(self):
"""cp_size=2, prefix_len=1, extend_len=3 (pad=1)."""
self._run_parity(prefix_len=1, extend_len=3, cp_size=2)
def test_cp4_prefix1_extend7(self):
"""cp_size=4, prefix_len=1, extend_len=7 (pad=1)."""
self._run_parity(prefix_len=1, extend_len=7, cp_size=4)
def test_cp8_prefix1_extend17(self):
"""cp_size=8, prefix_len=1, extend_len=17 (pad=7)."""
self._run_parity(prefix_len=1, extend_len=17, cp_size=8)
if __name__ == "__main__":
unittest.main()
test/registered/kernels/ops/attention/test_mla_cp_fa3_parity.py
-210
View File
@@ -1,210 +0,0 @@
"""FA3 numerical parity for MLA prefill CP.
Verifies the rank-local zigzag-split FA3 path (``_mla_cp_attn`` +
``cp_attn_forward_extend`` in ``flashattention_backend.py``) matches a
single non-CP ``flash_attn_with_kvcache`` over the full sequence.
Single-process, single-layer, pre-populated paged KV cache. Requires
FA3 ver=3 (Hopper+).
"""
import math
import sys
from types import SimpleNamespace
import pytest
import torch
from sglang.srt.layers.utils.cp_utils import (
ContextParallelMetadata,
cp_attn_forward_extend,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, stage="extra-a", runner_config="1-gpu-large")
if not torch.cuda.is_available():
pytest.skip(reason="CUDA required for FA3", allow_module_level=True)
_cap = torch.cuda.get_device_capability(0)
if _cap[0] < 9:
pytest.skip(
reason=f"FA3 ver=3 requires Hopper (sm90+); got sm{_cap[0]}{_cap[1]}",
allow_module_level=True,
)
try:
from sgl_kernel.flash_attn import flash_attn_with_kvcache
except ImportError as e:
pytest.skip(
reason=f"sgl_kernel.flash_attn unavailable: {e}",
allow_module_level=True,
)
DEVICE = torch.device("cuda")
DTYPE = torch.bfloat16
# Default shape is DeepSeek V3/R1 TP=8 MLA: 16 heads, v=512, rope=64.
NUM_HEADS = 16
V_HEAD_DIM = 512
QK_ROPE_HEAD_DIM = 64
PAGE_SIZE = 1
def _build_cache_and_q(seq_len):
"""Pre-populated paged KV cache + full-sequence q.
Pre-population mirrors upstream ``rebuild_cp_kv_cache``, which all-gathers
rank-local KV into the global pool before the attention call, so each
rank's FA3 invocation sees the same fully-populated cache.
"""
num_pages = (seq_len + PAGE_SIZE - 1) // PAGE_SIZE
c_kv_cache = torch.randn(
num_pages, PAGE_SIZE, 1, V_HEAD_DIM, dtype=DTYPE, device=DEVICE
)
k_rope_cache = torch.randn(
num_pages, PAGE_SIZE, 1, QK_ROPE_HEAD_DIM, dtype=DTYPE, device=DEVICE
)
q_nope = torch.randn(seq_len, NUM_HEADS, V_HEAD_DIM, dtype=DTYPE, device=DEVICE)
q_rope = torch.randn(
seq_len, NUM_HEADS, QK_ROPE_HEAD_DIM, dtype=DTYPE, device=DEVICE
)
page_table = torch.arange(num_pages, dtype=torch.int32, device=DEVICE).unsqueeze(0)
return c_kv_cache, k_rope_cache, q_nope, q_rope, page_table
def _full_seq_attn(
seq_len, q_nope, q_rope, c_kv_cache, k_rope_cache, page_table, softmax_scale
):
"""Non-CP reference: single flash_attn_with_kvcache over the full seq."""
return flash_attn_with_kvcache(
q=q_rope,
qv=q_nope,
k_cache=k_rope_cache,
v_cache=c_kv_cache,
page_table=page_table,
cache_seqlens=torch.tensor([seq_len], dtype=torch.int32, device=DEVICE),
cu_seqlens_q=torch.tensor([0, seq_len], dtype=torch.int32, device=DEVICE),
cu_seqlens_k_new=None,
max_seqlen_q=seq_len,
softmax_scale=softmax_scale,
causal=True,
ver=3,
)
def _cp_attn_for_rank(
rank,
cp_size,
block_size,
q_nope,
q_rope,
c_kv_cache,
k_rope_cache,
page_table,
softmax_scale,
):
"""Run the rank-local CP closure from ``flashattention_backend.py``.
Zigzag layout: rank r gets blocks [r, num_blocks - 1 - r] where
num_blocks = cp_size * 2. kv_len for each half is the cumulative KV
extent through the end of that block.
"""
num_blocks = cp_size * 2
b_prev, b_next = rank, num_blocks - 1 - rank
prev_slice = slice(b_prev * block_size, (b_prev + 1) * block_size)
next_slice = slice(b_next * block_size, (b_next + 1) * block_size)
q_nope_local = torch.cat([q_nope[prev_slice], q_nope[next_slice]], dim=0)
q_rope_local = torch.cat([q_rope[prev_slice], q_rope[next_slice]], dim=0)
q_fused = torch.cat([q_nope_local, q_rope_local], dim=-1)
cp_meta = ContextParallelMetadata(
kv_len_prev_tensor=torch.tensor(
[(b_prev + 1) * block_size], dtype=torch.int32, device=DEVICE
),
kv_len_next_tensor=torch.tensor(
[(b_next + 1) * block_size], dtype=torch.int32, device=DEVICE
),
cu_seqlens_q_prev_tensor=torch.tensor(
[0, block_size], dtype=torch.int32, device=DEVICE
),
cu_seqlens_q_next_tensor=torch.tensor(
[0, block_size], dtype=torch.int32, device=DEVICE
),
max_seqlen_q_prev=block_size,
max_seqlen_q_next=block_size,
total_q_prev_tokens=block_size,
)
fb = SimpleNamespace(attn_cp_metadata=cp_meta)
def _mla_cp_attn(q_chunk, cu_seqlens_q_cp, cache_seqlens_cp, max_seqlen_q_cp):
q_nope_chunk = q_chunk[..., :V_HEAD_DIM]
q_rope_chunk = q_chunk[..., V_HEAD_DIM:]
return flash_attn_with_kvcache(
q=q_rope_chunk,
qv=q_nope_chunk,
k_cache=k_rope_cache,
v_cache=c_kv_cache,
page_table=page_table,
cache_seqlens=cache_seqlens_cp,
cu_seqlens_q=cu_seqlens_q_cp,
cu_seqlens_k_new=None,
max_seqlen_q=max_seqlen_q_cp,
softmax_scale=softmax_scale,
causal=True,
ver=3,
)
local_out = cp_attn_forward_extend(fb, q_fused, DEVICE, _mla_cp_attn)
return local_out, prev_slice, next_slice
@pytest.mark.parametrize(
"cp_size, block_size",
[
(2, 64), # DSv3 TP=8 baseline
(2, 128), # longer per-block seq
(4, 32), # multi-rank zigzag: rank r gets blocks [r, 7-r]
],
)
def test_cp_parity(cp_size, block_size):
torch.manual_seed(0)
seq_len = block_size * cp_size * 2
softmax_scale = 1.0 / math.sqrt(V_HEAD_DIM + QK_ROPE_HEAD_DIM)
c_kv_cache, k_rope_cache, q_nope, q_rope, page_table = _build_cache_and_q(seq_len)
ref_out = _full_seq_attn(
seq_len, q_nope, q_rope, c_kv_cache, k_rope_cache, page_table, softmax_scale
)
for rank in range(cp_size):
local_out, prev_slice, next_slice = _cp_attn_for_rank(
rank,
cp_size,
block_size,
q_nope,
q_rope,
c_kv_cache,
k_rope_cache,
page_table,
softmax_scale,
)
torch.testing.assert_close(
local_out[:block_size],
ref_out[prev_slice],
rtol=1e-3,
atol=5e-3,
msg=f"rank={rank} prev-half mismatch",
)
torch.testing.assert_close(
local_out[block_size:],
ref_out[next_slice],
rtol=1e-3,
atol=5e-3,
msg=f"rank={rank} next-half mismatch",
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))
test/registered/pp/test_pp_parallel_compat.py
-1
View File
@@ -90,7 +90,6 @@ class TestQwen3MoePPxCP(_Qwen3MoePPCompatMixin, CustomTestCase):
"--cp-strategy",
"zigzag",
]
server_env = {"SGLANG_ENABLE_CP_V2": "1"}
class TestQwen3MoePPxDP(_Qwen3MoePPCompatMixin, CustomTestCase):
test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_cp.py
+3 -1
View File
@@ -39,7 +39,9 @@ class TestUnifiedQwen3HiCacheCP(UnifiedRadixTreeTestMixin, CustomTestCase):
"4",
"--attn-cp-size",
"2",
"--enable-prefill-context-parallel",
"--enable-prefill-cp",
"--cp-strategy",
"zigzag",
"--mem-fraction-static",
"0.8",
"--cuda-graph-max-bs-decode",