diff --git a/python/sglang/kernels/ops/attention/dsa/transform_index.py b/python/sglang/kernels/ops/attention/dsa/transform_index.py index ae5abf381..506add470 100644 --- a/python/sglang/kernels/ops/attention/dsa/transform_index.py +++ b/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 - grid = (qo_len,) - transform_index_page_table_decode_kernel[grid]( + if topk == 2048: + # Keep the single-program path for the unpooled width, which covers a + # 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_row_stride=page_table.stride(0), + 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 diff --git a/test/registered/kernels/ops/attention/test_dsa_transform_index.py b/test/registered/kernels/ops/attention/test_dsa_transform_index.py index 4598bbe38..231540540 100644 --- a/test/registered/kernels/ops/attention/test_dsa_transform_index.py +++ b/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()