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

Support batch size > 1 when enable CP (#23269)

Browse Source 
Co-authored-by: Shunkang <182541032+Shunkangz@users.noreply.github.co>
Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
Shunkangz
2026-05-27 14:11:17 -07:00
committed by GitHub
co-authored by Shunkang Khoa Pham Baizhou Zhang
parent ddf0627254
commit 19663aafcd
13 changed files with 263 additions and 300 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/cp/test_qwen3_30b.py
-132
View File
@@ -1,132 +0,0 @@
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=261, stage="extra-b", runner_config="4-gpu-h100")
QWEN3_30B_MODEL_PATH = "Qwen/Qwen3-30B-A3B-FP8"
GSM8K_BASELINE_ACCURACY = 0.85
class TestQwen330B(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",
"--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}',
],
)
@classmethod
def tearDownClass(cls):
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 TestQwen330BCP(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",
"--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}',
],
)
@classmethod
def tearDownClass(cls):
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()
test/registered/kernels/test_cp_prefix_len_fa3_parity.py
+1 -1
View File
@@ -77,7 +77,7 @@ class TestCPPrefixLenFA3Parity(CustomTestCase):
def _call_meta(rank: int):
return prepare_context_parallel_metadata(
padded_extend, rank, cp_size, seqs_len, extend_lens=extend_lens
padded_extend, rank, cp_size, seqs_len, extend_seqs_len=extend_lens
)
# Exercise the non-DSA branch; the DSA branch uses a separate
test/registered/kernels/test_mla_cp_fa3_parity.py
+9 -2
View File
@@ -126,8 +126,15 @@ def _cp_attn_for_rank(
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,
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)