diff --git a/python/sglang/srt/layers/attention/triton_ops/trtllm_mha_page_table.py b/python/sglang/srt/layers/attention/triton_ops/trtllm_mha_page_table.py index 7a7de6fc0..fdf41a811 100644 --- a/python/sglang/srt/layers/attention/triton_ops/trtllm_mha_page_table.py +++ b/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)) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 5ce1a4302..c5a834741 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/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 - # (_fill_page_table_device) sizes to the static max_num_pages and bounds - # the actual writes by cache_seqlens, so no runtime host max is needed. + # The device-side build (_fill_page_table_device) sizes to the static + # max_num_pages and bounds the actual writes by cache_seqlens, so no + # 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,39 +491,24 @@ 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( - 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( - num_tokens_per_bs, - bs * num_tokens_per_bs + 1, - num_tokens_per_bs, - dtype=torch.int32, - device=metadata.cu_seqlens_q.device, - ) + 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(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( + num_tokens_per_bs, + bs * num_tokens_per_bs + 1, + num_tokens_per_bs, + dtype=torch.int32, + 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 diff --git a/test/registered/attention/test_trtllm_mha_page_table.py b/test/registered/attention/test_trtllm_mha_page_table.py index 44b5d8d11..14d2de417 100644 --- a/test/registered/attention/test_trtllm_mha_page_table.py +++ b/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)