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

[CP] 1/N: Support MLA Prefill Context Parallel (#23292)

Browse Source 
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Khoa Pham
2026-05-23 03:07:09 -07:00
committed by GitHub
co-authored by Claude Opus 4.7
parent 81cd338fcc
commit b0ce16d0c5
21 changed files with 900 additions and 161 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/cp/test_deepseek_v3_cp_single_node.py
+89
View File
@@ -0,0 +1,89 @@
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=500, stage="extra-b", runner_config="deepep-8-gpu-h200")
DEEPSEEK_V3_MODEL_PATH = "deepseek-ai/DeepSeek-V3-0324"
# Matches the non-CP DSv3 production baseline in
# ``test_deepseek_v3_basic.py`` / ``test_deepseek_v3_mtp.py``. Pinning
# MLA CP to the same threshold makes this test double as a regression
# gate against the known production accuracy.
GSM8K_ACCURACY_THRESHOLD = 0.935
class TestDeepseekV3CPInSeqSplit(CustomTestCase):
"""tp=8, dp=2, attn-cp=4 — DP attention + DeepEP MoE + MLA CP."""
@classmethod
def setUpClass(cls):
cls.model = DEEPSEEK_V3_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--tp",
"8",
"--dp",
"2",
"--enable-prefill-context-parallel",
"--attention-backend",
"fa3",
"--mem-frac",
"0.7",
"--cuda-graph-max-bs",
"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 * 5,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
# "test_a_" prefix pins alphabetical first-run ordering so this
# warms up the server before any follow-up sibling test methods.
def test_a_gsm8k(self):
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-v3-mla-cp-in-seq-split)\n"
f'{metrics["score"]=:.3f}\n'
)
self.assertGreater(metrics["score"], GSM8K_ACCURACY_THRESHOLD)
if __name__ == "__main__":
unittest.main()
test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py
+85
View File
@@ -0,0 +1,85 @@
"""B200 extra CI: DeepSeek-V4-Flash FP4 with attn-CP (DSA prefill CP).
Balanced recipe (TP=4, DeepEP, EAGLE) plus --attn-cp-size=4 with the
DSA prefill-CP round-robin-split mode. Split out of
models_e2e/test_deepseek_v4_flash_fp4_b200.py so the `cp` group covers
all context-parallel tests.
Registry: extra-b-test-4-gpu-b200 (label-gated extra CI, 4x B200)
"""
import unittest
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.basic_decode_correctness_kit import BasicDecodeCorrectnessMixin
from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
from sglang.test.test_utils import (
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
try_cached_model,
)
register_cuda_ci(est_time=235, stage="extra-b", runner_config="4-gpu-b200")
MODEL = "deepseek-ai/DeepSeek-V4-Flash"
SERVER_LAUNCH_TIMEOUT = 3600
DEEPEP_CONFIG = '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}'
_DEEPEP_ENV = {
"SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
}
class TestDSV4FlashFP4B200Balanced_CP(
BasicDecodeCorrectnessMixin,
GSM8KMixin,
CustomTestCase,
):
"""Balanced recipe: TP=4, DP=4, DeepEP, EAGLE (1-step spec)."""
gsm8k_accuracy_thres = 0.93
@classmethod
def setUpClass(cls):
cls.model = try_cached_model(MODEL)
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--attn-cp-size",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
"--enable-dsa-prefill-context-parallel",
"--dsa-prefill-cp-mode",
"round-robin-split",
"--deepep-config",
DEEPEP_CONFIG,
],
env=_DEEPEP_ENV,
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
if __name__ == "__main__":
unittest.main()
test/registered/4-gpu-models/test_qwen3_30b.py → test/registered/cp/test_qwen3_30b.py
+1 -1
View File
@@ -11,7 +11,7 @@ from sglang.test.test_utils import (
popen_launch_server,
)
register_cuda_ci(est_time=261, stage="base-c", runner_config="4-gpu-h100")
register_cuda_ci(est_time=261, stage="extra-b", runner_config="4-gpu-h100")
QWEN3_30B_MODEL_PATH = "Qwen/Qwen3-30B-A3B-FP8"
test/registered/kernels/test_cp_prefix_len_fa3_parity.py
+143
View File
@@ -0,0 +1,143 @@
"""
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_lens=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/test_mla_cp_fa3_parity.py
+203
View File
@@ -0,0 +1,203 @@
"""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
),
actual_seq_q_prev=block_size,
actual_seq_q_next=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/models_e2e/test_deepseek_v4_flash_fp4_b200.py
+1 -50
View File
@@ -20,7 +20,7 @@ from sglang.test.test_utils import (
try_cached_model,
)
register_cuda_ci(est_time=700, stage="base-c", runner_config="dsv4-4-gpu-b200")
register_cuda_ci(est_time=465, stage="base-c", runner_config="dsv4-4-gpu-b200")
MODEL = "deepseek-ai/DeepSeek-V4-Flash"
SERVER_LAUNCH_TIMEOUT = 3600
@@ -156,54 +156,5 @@ class TestDSV4FlashFP4NonMTPB200(
kill_process_tree(cls.process.pid)
class TestDSV4FlashFP4B200Balanced_CP(
BasicDecodeCorrectnessMixin,
GSM8KMixin,
CustomTestCase,
):
"""Balanced recipe: TP=4, DP=4, DeepEP, EAGLE (1-step spec)."""
gsm8k_accuracy_thres = 0.93
@classmethod
def setUpClass(cls):
cls.model = try_cached_model(MODEL)
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
other_args=[
"--trust-remote-code",
"--tp",
"4",
"--attn-cp-size",
"4",
"--enable-dp-attention",
"--moe-a2a-backend",
"deepep",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"1",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"2",
"--enable-dsa-prefill-context-parallel",
"--dsa-prefill-cp-mode",
"round-robin-split",
"--deepep-config",
DEEPEP_CONFIG,
],
env=_DEEPEP_ENV,
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
if __name__ == "__main__":
unittest.main()