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

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>

python/sglang/kernels/ops/attention/dsa/transform_index.py

@@ -116,6 +116,50 @@ def transform_index_page_table_decode_kernel(
     tl.store(result_ptr + offset, -1, mask=~mask)
 
 
+@triton.jit
+def transform_index_page_table_decode_tiled_kernel(
+    page_table_ptr: torch.Tensor,
+    topk_indices_ptr: torch.Tensor,
+    result_ptr: torch.Tensor,
+    page_table_row_stride: tl.constexpr,
+    topk_indices_stride_0: tl.constexpr,
+    topk_indices_stride_1: tl.constexpr,
+    result_stride_0: tl.constexpr,
+    result_stride_1: tl.constexpr,
+    TOPK: tl.constexpr,
+    BLOCK_TOPK: tl.constexpr,
+):
+    """Width-generic form of the kernel above.
+
+    The 2048 variant folds the row stride into a compile-time TOPK and covers a
+    whole row with one unmasked `tl.arange`, which needs TOPK to be a power of
+    two. k-pool widths are not: `index_topk + index_kpool - 1` is 2051 for
+    GLM-5.3-Flash. Tile the row instead and carry the strides explicitly.
+    """
+    req_id = tl.program_id(0)
+    topk_offsets = tl.program_id(1) * BLOCK_TOPK + tl.arange(0, BLOCK_TOPK)
+    in_row = topk_offsets < TOPK
+
+    loaded_topk_indices = tl.load(
+        topk_indices_ptr
+        + req_id * topk_indices_stride_0
+        + topk_offsets * topk_indices_stride_1,
+        mask=in_row,
+        other=-1,
+    )
+    selected = in_row & (loaded_topk_indices >= 0)
+    loaded_kv_indices = tl.load(
+        page_table_ptr + req_id * page_table_row_stride + loaded_topk_indices,
+        mask=selected,
+        other=-1,
+    )
+    tl.store(
+        result_ptr + req_id * result_stride_0 + topk_offsets * result_stride_1,
+        loaded_kv_indices,
+        mask=in_row,
+    )
+
+
 # Expanded EAGLE page tables are contiguous, so their row stride changes with
 # the exact context length. Treating it as constexpr creates one cubin per
 # observed length and grows the loaded-module set in long-lived processes.
@@ -194,18 +238,37 @@ def transform_index_page_table_decode_fast(
     """
     assert page_size == 1
     assert page_table.shape[0] == topk_indices.shape[0]
-    assert topk_indices.shape[1] == 2048
     qo_len = topk_indices.shape[0]
+    topk = topk_indices.shape[1]
     if result is None:
         result = torch.empty_like(topk_indices, dtype=torch.int32)
-    # Launch triton kernel
+    if topk == 2048:
-    grid = (qo_len,)
+        # Keep the single-program path for the unpooled width, which covers a
-    transform_index_page_table_decode_kernel[grid](
+        # whole row per program with no masking.
+        transform_index_page_table_decode_kernel[(qo_len,)](
+            page_table,
+            topk_indices,
+            result,
+            page_size,
+            page_table_row_stride=page_table.stride(0),
+        )
+        return result
+
+    block_topk = 256
+    transform_index_page_table_decode_tiled_kernel[
+        (qo_len, triton.cdiv(topk, block_topk))
+    ](
         page_table,
         topk_indices,
         result,
-        page_size,
+        page_table.stride(0),
-        page_table_row_stride=page_table.stride(0),
+        topk_indices.stride(0),
+        topk_indices.stride(1),
+        result.stride(0),
+        result.stride(1),
+        TOPK=topk,
+        BLOCK_TOPK=block_topk,
+        num_warps=4,
     )
     return result
 

test/registered/kernels/ops/attention/test_dsa_transform_index.py

@@ -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()