[dflash] fa3/fa4: device-side page table; drop seq_lens_cpu D2H sync (#29343)

test/registered/attention/test_trtllm_mha_page_table.py

@@ -6,6 +6,11 @@ it never reads a runtime max (no D2H sync). This test checks the device build is
 bit-identical to the legacy gather for the columns each request uses, for both
 the full page table and the SWA-translated page table, across context lengths,
 page sizes, and batch sizes.
+
+It also pins the invariant that lets the no-host-max (GPU-only) path hand the
+kernel a static ``max_num_pages``-wide buffer: every column past a request's
+page count must be left untouched, i.e. the kernel bounds its writes by the
+device-side ``cache_seqlens`` alone.
 """
 
 import unittest
@@ -20,7 +25,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
 # Triton kernel unit test for the trtllm_mha device-side page-table build.
-register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
+register_cuda_ci(est_time=14, stage="base-b", runner_config="1-gpu-small")
 
 
 def _build_page_table_reference(
@@ -143,6 +148,82 @@ class TestTrtllmMhaPageTable(CustomTestCase):
                         max_ctx, page_size, num_reqs=max(64, bs), bs=bs, swa=True
                     )
 
+    def _run_self_guard_case(self, max_context_len, page_size, bs, swa=False):
+        """Short sequences against a full static buffer -- the GPU-only shape.
+
+        Pre-fill the page table with a sentinel and run the kernel with the
+        static ``max_num_pages`` width (no host max to tighten it). Used columns
+        must hold the right block ids; every tail column must keep the sentinel,
+        proving the kernel never writes past the device-side ``cache_seqlens``.
+        """
+        torch.manual_seed(0)
+        dev = "cuda"
+        num_reqs = max(64, bs)
+        max_num_pages = (max_context_len + page_size - 1) // page_size
+        n_slots = num_reqs * max_context_len
+        req_to_token = torch.randint(
+            0, n_slots, (num_reqs, max_context_len), dtype=torch.int32, device=dev
+        )
+        req_pool_indices = torch.randperm(num_reqs, device=dev)[:bs].to(torch.int32)
+        # Cap lengths well below max_context_len so most tail columns stay unused.
+        hi = max(2, max_context_len // 8)
+        cache_seqlens = torch.randint(1, hi + 1, (bs,), dtype=torch.int32, device=dev)
+        full_to_swa = (
+            torch.randint(0, n_slots, (n_slots,), dtype=torch.int32, device=dev)
+            if swa
+            else None
+        )
+
+        SENTINEL = -1
+        page_table = torch.full(
+            (bs, max_num_pages), SENTINEL, dtype=torch.int32, device=dev
+        )
+        swa_page_table = (
+            torch.full((bs, max_num_pages), SENTINEL, dtype=torch.int32, device=dev)
+            if swa
+            else None
+        )
+        build_trtllm_mha_page_table(
+            req_to_token=req_to_token,
+            req_pool_indices=req_pool_indices,
+            cache_seqlens=cache_seqlens,
+            page_table=page_table,
+            page_size=page_size,
+            swa_page_table=swa_page_table,
+            full_to_swa=full_to_swa,
+        )
+        pt_ref, swa_ref = _build_page_table_reference(
+            req_to_token, req_pool_indices, cache_seqlens, page_size, full_to_swa
+        )
+
+        tag = f"max_ctx={max_context_len} page_size={page_size} bs={bs} swa={swa}"
+        for i in range(bs):
+            npages = (int(cache_seqlens[i].item()) + page_size - 1) // page_size
+            self.assertTrue(
+                torch.equal(page_table[i, :npages], pt_ref[i, :npages]),
+                f"used-column mismatch req={i} {tag}",
+            )
+            self.assertTrue(
+                torch.all(page_table[i, npages:] == SENTINEL),
+                f"kernel wrote past cache_seqlens req={i} npages={npages} {tag}",
+            )
+            if swa:
+                self.assertTrue(
+                    torch.equal(swa_page_table[i, :npages], swa_ref[i, :npages]),
+                    f"swa used-column mismatch req={i} {tag}",
+                )
+                self.assertTrue(
+                    torch.all(swa_page_table[i, npages:] == SENTINEL),
+                    f"swa wrote past cache_seqlens req={i} npages={npages} {tag}",
+                )
+
+    def test_writes_bounded_by_cache_seqlens(self):
+        for max_ctx in (4096, 131072):
+            for page_size in (1, 64, 256):
+                for bs in (1, 8):
+                    self._run_self_guard_case(max_ctx, page_size, bs=bs)
+                    self._run_self_guard_case(max_ctx, page_size, bs=bs, swa=True)
+
 
 if __name__ == "__main__":
     unittest.main()