[Fix][DSA] Bound prefill Triton specializations for page-table stride (#37093)

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

@@ -116,13 +116,16 @@ def transform_index_page_table_decode_kernel(
     tl.store(result_ptr + offset, -1, mask=~mask)
 
 
-@triton.jit
+# 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.
+@triton.jit(do_not_specialize=["page_table_stride_0"])
 def transform_index_page_table_prefill_kernel(
     page_table_ptr: torch.Tensor,
     topk_indices_ptr: torch.Tensor,
     cu_seqlens_q_ptr: torch.Tensor,
     result_ptr: torch.Tensor,
-    page_table_stride_0: tl.constexpr,
+    page_table_stride_0,
     page_table_stride_1: tl.constexpr,
     topk_indices_stride_0: tl.constexpr,
     topk_indices_stride_1: tl.constexpr,

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

@@ -190,6 +190,36 @@ class TestDSATransformIndex(CustomTestCase):
                 cu_seqlens_q=cu_seqlens_q,
             )
 
+    def test_prefill_page_table_row_stride_is_not_specialized(self):
+        kernel = transform_index_module.transform_index_page_table_prefill_kernel
+        stride_param = next(
+            param for param in kernel.params if param.name == "page_table_stride_0"
+        )
+
+        self.assertFalse(stride_param.is_constexpr)
+        self.assertTrue(stride_param.do_not_specialize)
+
+    def test_prefill_dynamic_page_table_row_strides(self):
+        context_lengths = (4096, 4160, 4224)
+        kernel = transform_index_module.transform_index_page_table_prefill_kernel
+
+        self._check_case(
+            [2, 1],
+            context_lengths[0],
+            page_table_is_expanded=True,
+        )
+        kernel_cache = kernel.device_caches[torch.cuda.current_device()][0]
+        specialization_count = len(kernel_cache)
+
+        for context_length in context_lengths[1:]:
+            with self.subTest(context_length=context_length):
+                self._check_case(
+                    [2, 1],
+                    context_length,
+                    page_table_is_expanded=True,
+                )
+                self.assertEqual(len(kernel_cache), specialization_count)
+
     def test_mixed_lengths_padding_and_empty_batch(self):
         self._check_case(
             [0, 3, 1, 0, 4],