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

support GLM-5.2 MTP index sharing with prefill CP (#30992)

Browse Source 
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
This commit is contained in:
Yuxuan Zhang
2026-07-13 21:11:27 -07:00
committed by GitHub
co-authored by Baizhou Zhang
parent 1b4176cc46
commit 7e229e2a81
7 changed files with 529 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/cp/test_glm52_cp_index_share.py
+63
View File
@@ -0,0 +1,63 @@
import unittest
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
register_cuda_ci(est_time=400, stage="extra-b", runner_config="4-gpu-b200")
GLM52_NVFP4_MODEL_PATH = "nvidia/GLM-5.2-NVFP4"
class TestGLM52CPInterleave(GSM8KMixin, CustomTestCase):
gsm8k_accuracy_thres = 0.935
gsm8k_num_examples = 500
gsm8k_num_threads = 32
gsm8k_num_shots = 20
gsm8k_accept_length_thres = 3
@classmethod
def setUpClass(cls):
cls.model = GLM52_NVFP4_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",
"4",
"--attn-cp-size",
"4",
"--enable-prefill-cp",
"--cp-strategy",
"interleave",
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--mem-frac",
"0.85",
"--model-loader-extra-config",
'{"enable_multithread_load": true, "num_threads": 64}',
],
)
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
if __name__ == "__main__":
unittest.main()
test/registered/kernels/test_dsa_transform_index.py
+240
View File
@@ -0,0 +1,240 @@
import unittest
from unittest.mock import patch
import torch
import sglang.srt.layers.attention.dsa.transform_index as transform_index_module
from sglang.srt.layers.attention.dsa.transform_index import (
transform_index_page_table_decode_fast,
transform_index_page_table_prefill_fast,
)
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=60, stage="base-b", runner_config="1-gpu-large")
TOPK = 2048
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for this test.")
class TestDSATransformIndex(CustomTestCase):
def setUp(self):
super().setUp()
self.device = torch.device("cuda")
def tearDown(self):
torch.cuda.empty_cache()
super().tearDown()
def _make_page_table(self, rows: int, context_length: int) -> torch.Tensor:
columns = torch.arange(context_length, dtype=torch.int32, device=self.device)
row_bias = (
torch.arange(rows, dtype=torch.int32, device=self.device).unsqueeze(1) * 17
)
return columns.unsqueeze(0) + row_bias
def _make_topk(self, rows: int, context_length: int) -> torch.Tensor:
topk = (
torch.arange(TOPK, dtype=torch.int64, device=self.device)
.remainder(context_length)
.repeat(rows, 1)
)
if rows > 0:
topk[:, 0] = 0
topk[:, 1] = context_length - 1
topk[:, 257::257] = -1
return topk
def _expected(
self,
page_table: torch.Tensor,
topk_indices: torch.Tensor,
extend_lens_cpu: list[int],
output_num_tokens: int,
page_table_is_expanded: bool,
) -> torch.Tensor:
real_num_tokens = sum(extend_lens_cpu)
expected = torch.full(
(output_num_tokens, TOPK),
-1,
dtype=torch.int32,
device=self.device,
)
if real_num_tokens == 0:
return expected
if page_table_is_expanded:
source_rows = page_table[:real_num_tokens]
else:
request_ids = torch.repeat_interleave(
torch.arange(
len(extend_lens_cpu), dtype=torch.int64, device=self.device
),
torch.tensor(extend_lens_cpu, dtype=torch.int64, device=self.device),
)
source_rows = page_table[request_ids]
real_topk = topk_indices[:real_num_tokens]
torch.gather(
source_rows,
dim=1,
index=real_topk.clamp(min=0),
out=expected[:real_num_tokens],
)
expected[:real_num_tokens][real_topk < 0] = -1
return expected
def _check_decode_case(
self,
batch_size: int,
context_length: int,
*,
zero_row_stride: bool = False,
provide_result: bool = False,
) -> None:
if zero_row_stride:
page_table = self._make_page_table(1, context_length).expand(batch_size, -1)
else:
page_table = self._make_page_table(batch_size, context_length)
topk_indices = self._make_topk(batch_size, context_length)
expected = torch.empty(
(batch_size, TOPK), dtype=torch.int32, device=self.device
)
torch.gather(
page_table,
dim=1,
index=topk_indices.clamp(min=0),
out=expected,
)
expected[topk_indices < 0] = -1
result = torch.empty_like(expected) if provide_result else None
actual = transform_index_page_table_decode_fast(
page_table=page_table,
topk_indices=topk_indices,
result=result,
)
torch.cuda.synchronize()
if result is not None:
self.assertIs(actual, result)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
def _check_case(
self,
extend_lens_cpu: list[int],
context_length: int,
*,
page_table_is_expanded: bool,
topk_padding: int = 0,
output_padding: int = 0,
) -> None:
real_num_tokens = sum(extend_lens_cpu)
page_table_rows = (
real_num_tokens if page_table_is_expanded else len(extend_lens_cpu)
)
topk_num_tokens = real_num_tokens + topk_padding
output_num_tokens = topk_num_tokens + output_padding
page_table = self._make_page_table(page_table_rows, context_length)
topk_indices = self._make_topk(topk_num_tokens, context_length)
expected = self._expected(
page_table,
topk_indices,
extend_lens_cpu,
output_num_tokens,
page_table_is_expanded,
)
actual = transform_index_page_table_prefill_fast(
page_table=page_table,
topk_indices=topk_indices,
extend_lens_cpu=extend_lens_cpu,
output_num_tokens=output_num_tokens,
page_table_is_expanded=page_table_is_expanded,
)
torch.cuda.synchronize()
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
def test_prefill_uses_dedicated_kernel(self):
extend_lens_cpu = [2, 1]
context_length = 4096
page_table = self._make_page_table(len(extend_lens_cpu), context_length)
topk_indices = self._make_topk(sum(extend_lens_cpu), context_length)
with patch.object(
transform_index_module,
"transform_index_page_table_decode_fast",
side_effect=AssertionError("prefill must not launch decode per request"),
):
transform_index_page_table_prefill_fast(
page_table=page_table,
topk_indices=topk_indices,
extend_lens_cpu=extend_lens_cpu,
)
def test_prefill_uses_device_cu_seqlens(self):
extend_lens_cpu = [2, 1]
context_length = 4096
page_table = self._make_page_table(len(extend_lens_cpu), context_length)
topk_indices = self._make_topk(sum(extend_lens_cpu), context_length)
cu_seqlens_q = torch.tensor([0, 2, 3], dtype=torch.int32, device=self.device)
with patch.object(
transform_index_module.torch,
"tensor",
side_effect=AssertionError("must reuse device-side metadata"),
):
transform_index_page_table_prefill_fast(
page_table=page_table,
topk_indices=topk_indices,
extend_lens_cpu=extend_lens_cpu,
cu_seqlens_q=cu_seqlens_q,
)
def test_mixed_lengths_padding_and_empty_batch(self):
self._check_case(
[0, 3, 1, 0, 4],
8192,
page_table_is_expanded=False,
topk_padding=5,
output_padding=7,
)
self._check_case(
[0, 0],
16,
page_table_is_expanded=False,
output_padding=8,
)
def test_large_batch_size(self):
self._check_case(
[1] * 8192,
4096,
page_table_is_expanded=False,
)
def test_large_context_lengths(self):
for context_length, page_table_is_expanded in (
(640_000, True),
(1_000_000, False),
):
with self.subTest(
context_length=context_length,
page_table_is_expanded=page_table_is_expanded,
):
self._check_case(
[2, 1],
context_length,
page_table_is_expanded=page_table_is_expanded,
)
def test_decode_fast_correctness_and_strides(self):
self._check_decode_case(17, 8192, provide_result=True)
self._check_decode_case(17, 8192, zero_row_stride=True)
def test_decode_fast_extreme_shapes(self):
self._check_decode_case(8192, 4096)
self._check_decode_case(2, 1_000_000)
if __name__ == "__main__":
unittest.main()