[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
-    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