[misc] Trim dead code in trtllm_mha page-table backend; reuse eager page-table buffer (#28578)

python/sglang/srt/layers/attention/triton_ops/trtllm_mha_page_table.py

@@ -19,6 +19,7 @@ import triton
 import triton.language as tl
 
 # Tokens covered per CTA along the page-block (grid axis-1) dimension.
+# Must be a multiple of page_size (asserted in build_trtllm_mha_page_table).
 _MHA_KV_INDEX_BLOCK_TOKENS = 4096
 # Triton kernels can only read module globals that are tl.constexpr instances.
 _MHA_KV_INDEX_BLOCK_TOKENS_TL = tl.constexpr(_MHA_KV_INDEX_BLOCK_TOKENS)
@@ -55,6 +56,10 @@ def create_trtllm_mha_kv_indices_triton(
     from ``req_to_token`` and converts it to a block id (``slot // PAGE_SIZE``).
     Programs past the request's page count are guarded out, so the work (and the
     DRAM traffic) is bounded by the device-side ``seq_lens`` — no host max needed.
+
+    The SWA lookup assumes valid (``>= 0``) slots, unlike
+    ``translate_loc_from_full_to_swa``'s ``-1`` sentinel handling; page-boundary
+    reads stay within ``seq_len``, so slots are always valid here.
     """
     PAGES_PER_BLOCK: tl.constexpr = _MHA_KV_INDEX_BLOCK_TOKENS_TL // PAGE_SIZE
     pid_req = tl.program_id(0)
@@ -109,6 +114,9 @@ def build_trtllm_mha_page_table(
     assert has_swa == (
         swa_page_table is not None
     ), "full_to_swa and swa_page_table must be provided together"
+    assert (
+        _MHA_KV_INDEX_BLOCK_TOKENS % page_size == 0
+    ), f"page_size={page_size} must divide _MHA_KV_INDEX_BLOCK_TOKENS={_MHA_KV_INDEX_BLOCK_TOKENS}"
     bs, num_pages = page_table.shape
     create_trtllm_mha_kv_indices_triton[
         (bs, get_num_mha_kv_index_blocks(num_pages, page_size))

python/sglang/srt/layers/attention/trtllm_mha_backend.py

@@ -54,8 +54,6 @@ class TRTLLMMHAMetadata:
     cache_seqlens_int32: torch.Tensor = None
     # Maximum sequence length for query
     max_seq_len_q: int = 1
-    # Maximum sequence length for key
-    max_seq_len_k: int = 0
     # Cumulative sequence lengths for `query
     cu_seqlens_q: torch.Tensor = None
     # Cumulative sequence lengths for key
@@ -257,7 +255,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         kv_indices_buf: Optional[torch.Tensor] = None,
     ):
         """Initialize CUDA graph state for TRTLLM MHA."""
-        max_num_pages = (self.max_context_len + self.page_size - 1) // self.page_size
+        max_num_pages = self.max_num_pages
         self.decode_cuda_graph_metadata = {
             "cache_seqlens": torch.zeros(max_bs, dtype=torch.int32, device=self.device),
             "page_table": torch.zeros(
@@ -454,9 +452,9 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         """
         seq_lens = seq_lens[:bs]
         req_pool_indices = req_pool_indices[:bs]
-        # max_seq_len_k is the page-table width upper bound; the device-side build
+        # The device-side build (_fill_page_table_device) sizes to the static
-        # (_fill_page_table_device) sizes to the static max_num_pages and bounds
+        # max_num_pages and bounds the actual writes by cache_seqlens, so no
-        # the actual writes by cache_seqlens, so no runtime host max is needed.
+        # runtime host max is needed.
         metadata = None
         if forward_mode.is_decode_or_idle():
             if spec_info is not None:
@@ -474,7 +472,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                 metadata = self.decode_cuda_graph_metadata[bs]
                 metadata.cache_seqlens_int32.copy_(seq_lens)
 
-            metadata.max_seq_len_k = self.max_context_len
             metadata.cu_seqlens_k[1:].copy_(
                 torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32)
             )
@@ -485,7 +482,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
             # Here we only support topk = 1 for now.
             metadata = self.target_verify_metadata[bs]
             metadata.cache_seqlens_int32.copy_(seq_lens + metadata.max_seq_len_q)
-            metadata.max_seq_len_k = self.max_context_len
             metadata.cu_seqlens_k[1:].copy_(
                 torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32)
             )
@@ -495,18 +491,14 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         elif forward_mode.is_draft_extend_v2():
             metadata = self.draft_extend_metadata[bs]
             metadata.cache_seqlens_int32.copy_(seq_lens)
-            metadata.max_seq_len_k = self.max_context_len
             metadata.cu_seqlens_k[1:].copy_(
                 torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32)
             )
-            if forward_mode.is_draft_extend_v2():
             num_tokens_per_bs = spec_info.num_tokens_per_req
             if num_tokens_per_bs <= 0:
                 # Capture uses a synthetic EagleDraftExtendInput; infer the
                 # fixed V2 stride from the capture buffer when it is unset.
-                    num_tokens_per_bs = int(
+                num_tokens_per_bs = int(spec_info.num_accept_tokens[:bs].max().item())
-                        spec_info.num_accept_tokens[:bs].max().item()
-                    )
             metadata.max_seq_len_q = num_tokens_per_bs
             metadata.cu_seqlens_q[1:].copy_(
                 torch.arange(
@@ -517,17 +509,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
                     device=metadata.cu_seqlens_q.device,
                 )
             )
-            else:
-                extend_lens = spec_info.num_accept_tokens[:bs]
-                if spec_info.num_accept_tokens_cpu:
-                    metadata.max_seq_len_q = max(spec_info.num_accept_tokens_cpu)
-                else:
-                    metadata.max_seq_len_q = 1
-
-                metadata.cu_seqlens_q[1:].copy_(
-                    torch.cumsum(extend_lens, dim=0, dtype=torch.int32)
-                )
-
             self._fill_page_table_device(
                 metadata, req_pool_indices, metadata.cache_seqlens_int32
             )
@@ -692,7 +673,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
             else:
                 metadata.cu_seqlens_q = metadata.cu_seqlens_k
 
-        metadata.max_seq_len_k = self.max_context_len
         has_swa = self._swa_kv_pool is not None
         metadata.page_table = torch.empty(
             (batch_size, self.max_num_pages), dtype=torch.int32, device=device

test/registered/attention/test_trtllm_mha_page_table.py

@@ -28,7 +28,6 @@ def _build_page_table_reference(
     req_pool_indices: torch.Tensor,
     cache_seqlens: torch.Tensor,
     page_size: int,
-    max_num_pages: int,
     full_to_swa: Optional[torch.Tensor] = None,
 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
     """Reference impl: host-side strided gather, then // page_size.
@@ -113,7 +112,6 @@ class TestTrtllmMhaPageTable(CustomTestCase):
             req_pool_indices,
             cache_seqlens,
             page_size,
-            max_num_pages,
             full_to_swa=full_to_swa,
         )
 
@@ -133,7 +131,7 @@ class TestTrtllmMhaPageTable(CustomTestCase):
 
     def test_matches_reference_gather(self):
         for max_ctx in (2048, 4096, 131072):
-            for page_size in (1, 32, 64, 128):
+            for page_size in (1, 32, 64, 128, 256):
                 for bs in (1, 7, 32):
                     self._run_case(max_ctx, page_size, num_reqs=max(64, bs), bs=bs)