minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-17 18:36:22 -07:00
committed by GitHub
parent 1981464ba4
commit 9888b7b42b
3 changed files with 27 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/triton_ops/trtllm_mha_page_table.py
+8
View File
@@ -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
+18 -38
View File
@@ -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