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

[attn backend] Make seq_lens_cpu optional in trtllm_mha backend (#28106)

Browse Source 
Co-authored-by: jonnykong <jonnykong@fb.com>
This commit is contained in:
Jonny Kong
2026-06-17 16:24:19 -07:00
committed by GitHub
co-authored by jonnykong
parent e4fd613def
commit cd60c4edd0
3 changed files with 360 additions and 125 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
+126
View File
@@ -0,0 +1,126 @@
"""Device-side page-table builder for the trtllm_mha attention backend.
trtllm_mha builds its block (page) table from the global ``req_to_token`` pool.
Doing it with a host-max PyTorch gather forces a ``seq_lens.max().item()`` D2H
sync (the CPU must know the page-table width before launching). This kernel
instead derives the per-request page count from the device-side ``seq_lens``
tensor, so the build is sync-free: the grid/buffer use the static
``max_num_pages`` upper bound, while each program self-guards on the real length.
The kernel is MHA-owned (no dependency on the MLA kv-index kernels) and also
emits the SWA-translated block table in the same pass via the full->SWA lookup
table, so SWA hybrid models stay sync-free too.
"""
from typing import Optional
import torch
import triton
import triton.language as tl
# Tokens covered per CTA along the page-block (grid axis-1) dimension.
_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)
def get_num_mha_kv_index_blocks(num_pages: int, page_size: int) -> int:
"""Grid axis-1 size: number of page-block CTAs spanning the widest sequence.
``num_pages`` is the per-row width of the page-table buffer (the static
``max_num_pages`` upper bound). One CTA handles ``_MHA_KV_INDEX_BLOCK_TOKENS
// page_size`` pages.
"""
pages_per_block = _MHA_KV_INDEX_BLOCK_TOKENS // page_size
return (num_pages + pages_per_block - 1) // pages_per_block
@triton.jit
def create_trtllm_mha_kv_indices_triton(
req_to_token_ptr, # [max_reqs, max_context_len], int32
req_pool_indices_ptr, # [bs]
seq_lens_ptr, # [bs], per-request KV length in tokens
full_to_swa_ptr, # full->SWA token-slot lookup table, or dummy when not SWA
page_table_ptr, # [bs, num_pages] int32 block ids (output)
swa_page_table_ptr, # [bs, num_pages] int32 SWA block ids (output), or dummy
req_to_token_stride: tl.constexpr,
page_table_stride: tl.constexpr,
PAGE_SIZE: tl.constexpr,
HAS_SWA: tl.constexpr,
):
"""Fill ``page_table_ptr`` (and ``swa_page_table_ptr`` when ``HAS_SWA``).
Program ``(pid_req, pid_blk)`` writes the block ids of request ``pid_req`` for
the page-block ``pid_blk``. It reads the KV token slot at each page boundary
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.
"""
PAGES_PER_BLOCK: tl.constexpr = _MHA_KV_INDEX_BLOCK_TOKENS_TL // PAGE_SIZE
pid_req = tl.program_id(0)
pid_blk = tl.program_id(1)
seq_len = tl.load(seq_lens_ptr + pid_req)
num_pages = tl.cdiv(seq_len, PAGE_SIZE)
num_page_blocks = tl.cdiv(seq_len, _MHA_KV_INDEX_BLOCK_TOKENS_TL)
if pid_blk >= num_page_blocks:
return
req_pool_index = tl.load(req_pool_indices_ptr + pid_req)
page_idx = tl.arange(0, PAGES_PER_BLOCK) + pid_blk * PAGES_PER_BLOCK
token_pos = page_idx.to(tl.int64) * PAGE_SIZE
mask = page_idx < num_pages
slot = tl.load(
req_to_token_ptr
+ req_pool_index.to(tl.int64) * req_to_token_stride
+ token_pos,
mask=mask,
)
out_off = pid_req * page_table_stride + page_idx
tl.store(page_table_ptr + out_off, (slot // PAGE_SIZE).to(tl.int32), mask=mask)
if HAS_SWA:
swa_slot = tl.load(full_to_swa_ptr + slot.to(tl.int64), mask=mask)
tl.store(
swa_page_table_ptr + out_off,
(swa_slot // PAGE_SIZE).to(tl.int32),
mask=mask,
)
def build_trtllm_mha_page_table(
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
cache_seqlens: torch.Tensor,
page_table: torch.Tensor,
page_size: int,
swa_page_table: Optional[torch.Tensor] = None,
full_to_swa: Optional[torch.Tensor] = None,
) -> None:
"""Fill ``page_table`` (and ``swa_page_table`` when SWA) on-device, no D2H sync.
Computes the launch grid from the static page-table width and dispatches
``create_trtllm_mha_kv_indices_triton``. ``page_table`` (and, for SWA models,
``swa_page_table``) are written in place; the caller owns the buffers so the
cuda-graph path can reuse its pre-allocated tensors. SWA is enabled iff
``full_to_swa`` is provided, which then also requires ``swa_page_table``.
"""
has_swa = full_to_swa is not None
assert has_swa == (
swa_page_table is not None
), "full_to_swa and swa_page_table must be provided together"
bs, num_pages = page_table.shape
create_trtllm_mha_kv_indices_triton[
(bs, get_num_mha_kv_index_blocks(num_pages, page_size))
](
req_to_token,
req_pool_indices,
cache_seqlens,
full_to_swa,
page_table,
swa_page_table,
req_to_token.stride(0),
page_table.stride(0),
PAGE_SIZE=page_size,
HAS_SWA=has_swa,
)
python/sglang/srt/layers/attention/trtllm_mha_backend.py
+84 -125
View File
@@ -19,6 +19,9 @@ from sglang.srt.layers.attention.flashinfer_backend import (
from sglang.srt.layers.attention.triton_ops.trtllm_fp8_kv_kernel import (
fused_fp8_set_kv_buffer,
)
from sglang.srt.layers.attention.triton_ops.trtllm_mha_page_table import (
build_trtllm_mha_page_table,
)
from sglang.srt.layers.attention.utils import canonicalize_stride
from sglang.srt.mem_cache.memory_pool import KVWriteLoc
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
@@ -68,6 +71,11 @@ class TRTLLMMHAMetadata:
class TRTLLMHAAttnBackend(FlashInferAttnBackend):
"""TRTLLM MHA attention kernel from flashinfer."""
# Build the page table on-device from seq_lens (incl. the SWA-translated table
# via the full->SWA lookup; see _fill_page_table_device), so we never need the
# seq_lens_cpu D2H sync; opt out of it, matching trtllm_mla / triton.
needs_cpu_seq_lens: bool = False
def __init__(
self,
model_runner: ModelRunner,
@@ -131,6 +139,13 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
# separate index spaces; SWA layers need a translated page_table.
self._swa_kv_pool: Optional[SWAKVPool] = self._resolve_swa_kv_pool(model_runner)
# Static page-table width (upper bound). The CUDA-graph path builds the
# page table on-device sized to this constant, so it never reads a runtime
# max. See _fill_page_table_device.
self.max_num_pages = (
self.max_context_len + self.page_size - 1
) // self.page_size
# Forward metadata
self.forward_metadata: Optional[TRTLLMMHAMetadata] = None
@@ -166,20 +181,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
kvcache = allocator.get_kvcache()
return kvcache if isinstance(kvcache, SWAKVPool) else None
def _maybe_translate_swa(
self, token_indices: torch.Tensor
) -> Optional[torch.Tensor]:
"""Translate full-pool token indices to SWA-pool indices, or return None."""
if self._swa_kv_pool is None:
return None
shape = token_indices.shape
# trtllm-gen SWA attention kernels require int32 page indices.
return (
self._swa_kv_pool.translate_loc_from_full_to_swa(token_indices.reshape(-1))
.reshape(shape)
.to(torch.int32)
)
def _alloc_swa_page_table(
self, max_bs: int, max_num_pages: int
) -> Optional[torch.Tensor]:
@@ -188,17 +189,33 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
return None
return torch.zeros(max_bs, max_num_pages, dtype=torch.int32, device=self.device)
def _copy_swa_page_table(
def _fill_page_table_device(
self,
metadata: TRTLLMMHAMetadata,
page_indices: torch.Tensor,
num_pages: int,
req_pool_indices: torch.Tensor,
cache_seqlens: torch.Tensor,
):
"""Translate and copy SWA page indices into metadata. No-op for non-SWA."""
if metadata.swa_page_table is None:
return
swa_indices = self._maybe_translate_swa(page_indices)
metadata.swa_page_table[:, :num_pages].copy_(swa_indices // self.page_size)
"""Build the page table on-device from per-request KV lengths (no sync).
Fills ``metadata.page_table`` (a [bs, max_num_pages] buffer) in place with
block ids derived from ``cache_seqlens`` (a GPU tensor); for SWA models it
also fills ``metadata.swa_page_table`` via the full->SWA lookup. The Triton
kernel self-guards per request on the device-side length, so the grid and
buffer width use the static ``max_num_pages`` upper bound while the actual
writes stay bounded by ``cache_seqlens`` — no host-side max / D2H sync.
"""
has_swa = self._swa_kv_pool is not None
build_trtllm_mha_page_table(
req_to_token=self.req_to_token,
req_pool_indices=req_pool_indices,
cache_seqlens=cache_seqlens,
page_table=metadata.page_table,
page_size=self.page_size,
swa_page_table=metadata.swa_page_table if has_swa else None,
full_to_swa=(
self._swa_kv_pool.full_to_swa_index_mapping if has_swa else None
),
)
def _get_layer_cache_loc(
self,
@@ -250,9 +267,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
device=self.device,
),
"swa_page_table": self._alloc_swa_page_table(max_bs, max_num_pages),
"strided_indices": torch.arange(
0, self.max_context_len, self.page_size, device=self.device
),
}
# SWA write-target buffer; bound as a [:num_tokens] view in
@@ -305,9 +319,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
device=self.device,
),
"swa_page_table": self._alloc_swa_page_table(max_bs, max_num_pages),
"strided_indices": torch.arange(
0, self.max_context_len, self.page_size, device=self.device
),
}
self.draft_extend_metadata = {
@@ -329,9 +340,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
device=self.device,
),
"swa_page_table": self._alloc_swa_page_table(max_bs, max_num_pages),
"strided_indices": torch.arange(
0, self.max_context_len, self.page_size, device=self.device
),
}
def _build_cuda_graph_metadata(
@@ -439,15 +447,16 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
seq_lens: torch.Tensor,
forward_mode: ForwardMode,
spec_info: Optional[SpecInput],
seq_lens_cpu: Optional[torch.Tensor],
):
"""Shared capture+replay body for the cuda-graph init path.
Public entry: :py:meth:`init_forward_metadata_out_graph`.
"""
seq_lens = seq_lens[:bs]
seq_lens_cpu = seq_lens_cpu[: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.
metadata = None
if forward_mode.is_decode_or_idle():
if spec_info is not None:
@@ -460,58 +469,33 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
metadata.cache_seqlens_int32.copy_(
seq_lens + self.speculative_step_id + 1
)
metadata.max_seq_len_k = seq_lens.max().item() + (
self.speculative_step_id + 1
)
max_seq_pages = (
metadata.max_seq_len_k + self.page_size - 1
) // self.page_size
else:
# Normal Decode
metadata = self.decode_cuda_graph_metadata[bs]
max_len = seq_lens_cpu.max().item()
max_seq_pages = (max_len + self.page_size - 1) // self.page_size
metadata.max_seq_len_k = max_len
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)
)
page_indices = self.req_to_token[
req_pool_indices[:, None],
self.decode_cuda_graph_metadata["strided_indices"][:max_seq_pages][
None, :
],
]
metadata.page_table[:, :max_seq_pages].copy_(page_indices // self.page_size)
self._copy_swa_page_table(metadata, page_indices, max_seq_pages)
self._fill_page_table_device(
metadata, req_pool_indices, metadata.cache_seqlens_int32
)
elif forward_mode.is_target_verify():
# 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 = seq_lens_cpu.max().item() + metadata.max_seq_len_q
max_len = seq_lens_cpu.max().item()
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)
)
max_seq_pages = (
metadata.max_seq_len_k + self.page_size - 1
) // self.page_size
page_indices = self.req_to_token[
req_pool_indices[:, None],
self.decode_cuda_graph_metadata["strided_indices"][:max_seq_pages],
]
metadata.page_table[:, :max_seq_pages].copy_(page_indices // self.page_size)
self._copy_swa_page_table(metadata, page_indices, max_seq_pages)
self._fill_page_table_device(
metadata, req_pool_indices, metadata.cache_seqlens_int32
)
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 = seq_lens_cpu.max().item()
max_len = seq_lens_cpu.max().item()
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)
)
@@ -544,15 +528,9 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
torch.cumsum(extend_lens, dim=0, dtype=torch.int32)
)
max_seq_pages = (
metadata.max_seq_len_k + self.page_size - 1
) // self.page_size
page_indices = self.req_to_token[
req_pool_indices[:, None],
self.draft_extend_metadata["strided_indices"][:max_seq_pages],
]
metadata.page_table[:, :max_seq_pages].copy_(page_indices // self.page_size)
self._copy_swa_page_table(metadata, page_indices, max_seq_pages)
self._fill_page_table_device(
metadata, req_pool_indices, metadata.cache_seqlens_int32
)
self.forward_metadata = metadata
def update_verify_buffers_to_fill_after_draft(
@@ -608,7 +586,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
if in_capture:
num_tokens = forward_batch.positions.numel()
seq_lens_cpu = seq_lens.cpu()
self._build_cuda_graph_metadata(
bs, num_tokens, forward_mode, spec_info, seq_lens.device
)
@@ -618,7 +595,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
seq_lens=seq_lens,
forward_mode=forward_mode,
spec_info=spec_info,
seq_lens_cpu=seq_lens_cpu,
)
else:
self._apply_cuda_graph_metadata(
@@ -627,7 +603,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
seq_lens=seq_lens,
forward_mode=forward_mode,
spec_info=spec_info,
seq_lens_cpu=forward_batch.seq_lens_cpu,
)
# Refill the SWA write-target buffer from the live out_cache_loc before
@@ -656,9 +631,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
metadata.cache_seqlens_int32 = (
seqlens_in_batch + (self.speculative_step_id + 1)
).to(torch.int32)
metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item() + (
self.speculative_step_id + 1
)
metadata.cu_seqlens_q = torch.arange(
0, batch_size + 1, dtype=torch.int32, device=device
)
@@ -668,22 +640,15 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
),
(1, 0),
)
metadata.page_table = self.req_to_token_pool.req_to_token[
forward_batch.req_pool_indices, : metadata.max_seq_len_k
]
else:
# Normal Decode
metadata.cache_seqlens_int32 = seqlens_in_batch.to(torch.int32)
metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item()
metadata.cu_seqlens_q = torch.arange(
0, batch_size + 1, dtype=torch.int32, device=device
)
metadata.cu_seqlens_k = torch.nn.functional.pad(
torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)
)
metadata.page_table = self.req_to_token_pool.req_to_token[
forward_batch.req_pool_indices, : metadata.max_seq_len_k
]
elif forward_batch.forward_mode.is_target_verify():
# Only support topk = 1 for now.
tokens_per_req = forward_batch.input_ids.shape[0] // batch_size
@@ -691,9 +656,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
torch.int32
)
metadata.max_seq_len_q = tokens_per_req
metadata.max_seq_len_k = (
forward_batch.seq_lens_cpu.max().item() + tokens_per_req
)
metadata.cu_seqlens_q = torch.arange(
0,
batch_size * tokens_per_req + 1,
@@ -705,40 +667,46 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32),
(1, 0),
)
metadata.page_table = self.req_to_token_pool.req_to_token[
forward_batch.req_pool_indices, : metadata.max_seq_len_k
]
else:
metadata.cache_seqlens_int32 = seqlens_in_batch.to(torch.int32)
metadata.max_seq_len_k = forward_batch.seq_lens_cpu.max().item()
metadata.cu_seqlens_k = torch.nn.functional.pad(
torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)
)
metadata.page_table = self.req_to_token_pool.req_to_token[
forward_batch.req_pool_indices, : metadata.max_seq_len_k
]
# Query-side max length, sourced from the host-resident extend lengths
# (sync-free); for plain prefill these equal the full seq lens.
# NOTE: in piecewise CUDA graph warmup, extend_seq_lens_cpu is a torch.Tensor;
# Python's max() returns a 0-d tensor, but flashinfer expects an int.
max_q = max(forward_batch.extend_seq_lens_cpu)
metadata.max_seq_len_q = (
int(max_q.item()) if isinstance(max_q, torch.Tensor) else int(max_q)
)
if (
any(forward_batch.extend_prefix_lens_cpu)
or forward_batch.forward_mode.is_draft_extend_v2()
):
forward_batch.extend_prefix_lens_cpu is not None
and any(forward_batch.extend_prefix_lens_cpu)
) or forward_batch.forward_mode.is_draft_extend_v2():
extend_seq_lens = forward_batch.extend_seq_lens
# NOTE: in piecewise CUDA graph warmup, extend_seq_lens_cpu is a torch.Tensor;
# Python's max() returns a 0-d tensor, but flashinfer expects an int.
max_q = max(forward_batch.extend_seq_lens_cpu)
metadata.max_seq_len_q = (
int(max_q.item()) if isinstance(max_q, torch.Tensor) else int(max_q)
)
metadata.cu_seqlens_q = torch.nn.functional.pad(
torch.cumsum(extend_seq_lens, dim=0, dtype=torch.int32), (1, 0)
)
else:
metadata.max_seq_len_q = metadata.max_seq_len_k
metadata.cu_seqlens_q = metadata.cu_seqlens_k
# Compute SWA page table (None for non-SWA models)
metadata.swa_page_table = self._maybe_translate_swa(metadata.page_table)
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
)
metadata.swa_page_table = (
torch.empty(
(batch_size, self.max_num_pages), dtype=torch.int32, device=device
)
if has_swa
else None
)
self._fill_page_table_device(
metadata, forward_batch.req_pool_indices, metadata.cache_seqlens_int32
)
# int64 scatter index (unlike the int32 read page table above).
if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None:
@@ -748,19 +716,6 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
)
)
# Convert the page tables to a strided format
if self.page_size > 1:
self.strided_indices = torch.arange(
0, metadata.page_table.shape[1], self.page_size, device=self.device
)
metadata.page_table = (
metadata.page_table[:, self.strided_indices] // self.page_size
)
if metadata.swa_page_table is not None:
metadata.swa_page_table = (
metadata.swa_page_table[:, self.strided_indices] // self.page_size
)
self.forward_metadata = metadata
def forward_decode(
@@ -970,6 +925,10 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
class TRTLLMHAAttnMultiStepDraftBackend(FlashInferMultiStepDraftBackend):
"""Multi-step TRTLLM MHA attention kernel used by EAGLE."""
# Per-step backends build the page table on-device (sync-free); mirror that so
# decide_needs_cpu_seq_lens sees a consistent target + draft value.
needs_cpu_seq_lens: bool = False
def __init__(
self, model_runner: ModelRunner, topk: int, speculative_num_steps: int
):