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

[AMD] [GLM-5.3-Flash Day 0] Support non-2048 top-k widths in the DSA page-table transform (#39340)

Browse Source 
Co-authored-by: Thomas Wang <thomawan@amd.com>
Co-authored-by: Kevin Mi <45493463+kevin-mii@users.noreply.github.com>
Co-authored-by: Kevin Mi <mikevin920@yahoo.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Jacob0226
2026-09-21 19:50:17 -07:00
committed by GitHub
co-authored by Thomas Wang Kevin Mi Kevin Mi Cursor
parent 5f9c6b9eb0
commit 90cf471723
2 changed files with 107 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/ops/attention/test_dsa_transform_index.py
+38 -7
View File
@@ -9,12 +9,17 @@ from sglang.kernels.ops.attention.dsa.transform_index import (
transform_index_page_table_prefill_fast,
transform_index_page_table_prefill_ref,
)
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(est_time=9, stage="base-b-kernel-unit", runner_config="1-gpu-large")
register_amd_ci(est_time=9, suite="stage-b-test-1-gpu-small-amd-mi35x")
TOPK = 2048
# k-pool appends up to index_kpool - 1 open-tail tokens to index_topk, so the
# width the indexer hands over is not a power of two: 2048 + 4 - 1 for
# GLM-5.3-Flash. See get_dsa_mtp_topk_width() in srt/configs/model_config.py.
KPOOL_TOPK = 2051
@unittest.skipUnless(torch.cuda.is_available(), "CUDA is required for this test.")
@@ -34,9 +39,11 @@ class TestDSATransformIndex(CustomTestCase):
)
return columns.unsqueeze(0) + row_bias
def _make_topk(self, rows: int, context_length: int) -> torch.Tensor:
def _make_topk(
self, rows: int, context_length: int, topk_width: int = TOPK
) -> torch.Tensor:
topk = (
torch.arange(TOPK, dtype=torch.int64, device=self.device)
torch.arange(topk_width, dtype=torch.int64, device=self.device)
.remainder(context_length)
.repeat(rows, 1)
)
@@ -53,10 +60,11 @@ class TestDSATransformIndex(CustomTestCase):
extend_lens_cpu: list[int],
output_num_tokens: int,
page_table_is_expanded: bool,
topk_width: int = TOPK,
) -> torch.Tensor:
real_num_tokens = sum(extend_lens_cpu)
expected = torch.full(
(output_num_tokens, TOPK),
(output_num_tokens, topk_width),
-1,
dtype=torch.int32,
device=self.device,
@@ -92,14 +100,15 @@ class TestDSATransformIndex(CustomTestCase):
*,
zero_row_stride: bool = False,
provide_result: bool = False,
topk_width: int = TOPK,
) -> 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)
topk_indices = self._make_topk(batch_size, context_length, topk_width)
expected = torch.empty(
(batch_size, TOPK), dtype=torch.int32, device=self.device
(batch_size, topk_width), dtype=torch.int32, device=self.device
)
torch.gather(
page_table,
@@ -128,6 +137,7 @@ class TestDSATransformIndex(CustomTestCase):
page_table_is_expanded: bool,
topk_padding: int = 0,
output_padding: int = 0,
topk_width: int = TOPK,
) -> None:
real_num_tokens = sum(extend_lens_cpu)
page_table_rows = (
@@ -136,13 +146,14 @@ class TestDSATransformIndex(CustomTestCase):
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)
topk_indices = self._make_topk(topk_num_tokens, context_length, topk_width)
expected = self._expected(
page_table,
topk_indices,
extend_lens_cpu,
output_num_tokens,
page_table_is_expanded,
topk_width,
)
actual = transform_index_page_table_prefill_fast(
@@ -303,6 +314,26 @@ class TestDSATransformIndex(CustomTestCase):
self._check_decode_case(8192, 4096)
self._check_decode_case(2, 1_000_000)
def test_decode_fast_kpool_widths(self):
# 2051 is the GLM-5.3-Flash k-pool width; the others cover a partial
# trailing tile and a width below one tile.
for topk_width in (KPOOL_TOPK, 515, 257):
with self.subTest(topk_width=topk_width):
self._check_decode_case(17, 8192, topk_width=topk_width)
self._check_decode_case(
17, 8192, topk_width=topk_width, provide_result=True
)
def test_prefill_kpool_widths(self):
for topk_width in (KPOOL_TOPK, 515, 257):
with self.subTest(topk_width=topk_width):
self._check_case(
[2, 1],
4096,
page_table_is_expanded=False,
topk_width=topk_width,
)
if __name__ == "__main__":
unittest.main()