Fix DSV4 prefill large Triton recompilation idle across context lengths (#30255)

Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>

python/sglang/srt/layers/attention/dsv4/metadata_kernel.py

@@ -5,7 +5,7 @@ import triton
 import triton.language as tl
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["bs", "c128_cur_max_seq_len"])
 def _init_compressed_attn_metadata_kernel(
     seq_lens_ptr,
     positions_ptr,
@@ -22,8 +22,7 @@ def _init_compressed_attn_metadata_kernel(
     c128_page_indices_ptr,
     bs,
     max_pages,
-    page_size: tl.constexpr,
-    c128_max_seq_len: tl.constexpr,
+    c128_cur_max_seq_len,
     c128_page_size: tl.constexpr,
     BLOCK_SIZE: tl.constexpr,
     COMPUTE_PAGE_INDICES: tl.constexpr,
@@ -59,10 +58,10 @@ def _init_compressed_attn_metadata_kernel(
     tl.store(c128_seq_lens_clamp1_ptr + batch_id, c128_seq_lens_clamp1)
 
     if COMPUTE_PAGE_INDICES:
-        page_indices_base = batch_id * c128_max_seq_len
-        for block_start in range(0, c128_max_seq_len, BLOCK_SIZE):
+        page_indices_base = batch_id * c128_cur_max_seq_len
+        for block_start in tl.range(0, c128_cur_max_seq_len, BLOCK_SIZE):
             offsets = block_start + tl.arange(0, BLOCK_SIZE)
-            mask = offsets < c128_max_seq_len
+            mask = offsets < c128_cur_max_seq_len
 
             page_idx = offsets // c128_page_size
             offset_in_page = offsets % c128_page_size
@@ -121,18 +120,20 @@ def _init_compressed_attn_metadata_triton(
         assert (
             page_table is not None
         ), "page_table required when compute_page_indices=True"
-        assert page_size > 0, "page_size required when compute_page_indices=True"
+        assert (
+            page_size >= 128 and page_size % 128 == 0
+        ), "page_size must be a multiple of 128 when compute_page_indices=True"
         max_pages = page_table.shape[1]
         c128_page_size = page_size // 128
-        c128_max_seq_len = c128_page_size * max_pages
+        c128_cur_max_seq_len = c128_page_size * max_pages
         c128_page_indices = torch.empty(
-            bs, c128_max_seq_len, dtype=torch.int32, device=device
+            bs, c128_cur_max_seq_len, dtype=torch.int32, device=device
         )
         BLOCK_SIZE = triton.next_power_of_2(max(c128_page_size, 64))
     else:
         max_pages = 0
         c128_page_size = 1
-        c128_max_seq_len = 0
+        c128_cur_max_seq_len = 0
         c128_page_indices = None
         BLOCK_SIZE = 64
         if page_table is None:
@@ -159,8 +160,7 @@ def _init_compressed_attn_metadata_triton(
         ),
         bs,
         max_pages,
-        page_size if page_size > 0 else 128,
-        c128_max_seq_len,
+        c128_cur_max_seq_len,
         c128_page_size,
         BLOCK_SIZE,
         compute_page_indices,

python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py

@@ -135,7 +135,10 @@ def combine_topk_swa_indices(
     assert gather_lens.dtype == torch.int32
     assert compressed_base.dtype == torch.int32
     assert swa_base.dtype == torch.int32
-    assert compress_ratio >= 1, "COMPRESS_RATIO must be >= 1 (use TOP_K=0 for SWA-only)"
+    assert compress_ratio >= 1, "compress_ratio must be >= 1 (use topk=0 for SWA-only)"
+    assert (
+        topk_indices.shape[-1] >= topk
+    ), f"topk_indices width {topk_indices.shape[-1]} must be >= topk {topk}"
 
     num_tokens = topk_indices.shape[0]
     num_reqs = seq_lens.shape[0]
@@ -172,7 +175,7 @@ def combine_topk_swa_indices(
         gather_lens,
         compressed_base,
         swa_base,
-        TOP_K=topk,
+        top_k=topk,
         COMPRESS_RATIO=compress_ratio,
         WINDOW_SIZE=window_size,
         PADDED_TOP_K=triton.next_power_of_2(topk_indices.shape[-1]),
@@ -279,7 +282,7 @@ def _build_swa_token_ids_kernel(
         tl.store(out_ptr + out_off + i, swa_id)
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["top_k"])
 def _combine_topk_swa_indices_kernel(
     combined_indices_ptr,
     combined_indices_stride,
@@ -291,7 +294,7 @@ def _combine_topk_swa_indices_kernel(
     gather_lens_ptr,
     compressed_base_ptr,
     swa_base_ptr,
-    TOP_K: tl.constexpr,
+    top_k,
     COMPRESS_RATIO: tl.constexpr,
     WINDOW_SIZE: tl.constexpr,
     PADDED_TOP_K: tl.constexpr,
@@ -321,8 +324,8 @@ def _combine_topk_swa_indices_kernel(
         pos = start_pos + token_idx_in_query
         # Both the C4 indexer and the C128 metadata builder emit
         # min((pos+1)//compress_ratio, topk_tokens) valid entries. Caller
-        # passes TOP_K=0 for SWA-only layers to zero this out.
-        topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K)
+        # passes top_k=0 for SWA-only layers to zero this out.
+        topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, top_k)
         swa_len = tl.minimum(pos + 1, WINDOW_SIZE)
 
         combined_row = token_idx.to(tl.int64) * combined_indices_stride