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

[trtllm_mha] Fuse cuda-graph metadata rebuild into one triton kernel (#29843)

Browse Source 
Co-authored-by: pranjalssh <pranjalssh@fb.com>
This commit is contained in:
Pranjal Shankhdhar
2026-07-03 22:24:28 -07:00
committed by GitHub
co-authored by pranjalssh
parent 4028304579
commit ad744c6c6b
9 changed files with 787 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/hybrid_attn_backend.py
+4
View File
@@ -59,6 +59,10 @@ class HybridAttnBackend(AttentionBackend):
backend = self._select_backend(forward_batch.forward_mode)
backend.init_forward_metadata_out_graph(forward_batch, in_capture=in_capture)
def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch):
backend = self._select_backend(forward_batch.forward_mode)
backend.init_forward_metadata_in_graph(forward_batch)
def init_forward_metadata(self, forward_batch: ForwardBatch):
backend = self._select_backend(forward_batch.forward_mode)
backend.init_forward_metadata(forward_batch)
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+4
View File
@@ -838,6 +838,10 @@ class HybridLinearAttnBackend(AttentionBackend):
forward_batch, in_capture=in_capture
)
def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch):
for attn_backend in self.attn_backend_list:
attn_backend.init_forward_metadata_in_graph(forward_batch)
def init_forward_metadata(self, forward_batch: ForwardBatch):
if forward_batch.forward_mode.is_draft_extend_v2():
# DRAFT_EXTEND_V2 runs only full-attn layers in the draft model; skip
python/sglang/srt/layers/attention/triton_ops/trtllm_mha_graph_metadata.py
+198
View File
@@ -0,0 +1,198 @@
"""Fused CUDA-graph metadata update for the TRTLLM MHA backend.
`TRTLLMHAAttnBackend._apply_cuda_graph_metadata` used to rebuild the
page table(s) and seqlen buffers with ~25 small aten ops per graph
replay (index gathers, floor_divide, cumsum, dtype casts, copies).
On some CPUs that is ~0.7-1.0 ms of pure host dispatch, repeated 4x
per decode step (2 draft-decode steps + target-verify + draft-extend)
on every TP rank. The resulting per-rank CPU jitter skews the
cudaGraphLaunch across ranks and is paid as spin time inside the first
custom all-reduce of every replayed graph.
This kernel performs the whole update in ONE launch:
- cache_seqlens[i] = seq_lens[i] + seqlen_offset (int32)
- cu_seqlens_k[1:] = cumsum(cache_seqlens) (int32)
- cu_seqlens_q[1:] = cumsum(qlens) or arange*q_stride (optional)
- page_table[i, p] = req_to_token[req_pool_indices[i],
p * page_size] // page_size
- swa_page_table = full_to_swa_mapping[token] // page_size (optional)
- swa_out_cache_loc = full_to_swa_mapping[out_cache_loc], zero padded
(optional)
"""
import triton
import triton.language as tl
# cu_seqlens_q handling inside the fused kernel
Q_MODE_NONE = 0 # cu_seqlens_q is preset (decode / target-verify)
Q_MODE_CUMSUM = 1 # cu_seqlens_q[1:] = cumsum(qlens) (draft-extend)
Q_MODE_STRIDED = 2 # cu_seqlens_q[1:] = arange*q_stride (draft-extend v2)
@triton.jit
def update_trtllm_mha_graph_metadata_kernel(
# inputs
req_pool_indices_ptr, # [bs] int
seq_lens_ptr, # [bs] int
req_to_token_ptr, # [pool_size, req_to_token_stride] int32
swa_mapping_ptr, # [full_size + page_size + 1] int64, or None
out_cache_loc_ptr, # [num_out_tokens] int64, or None
qlens_ptr, # [bs] int, or None (Q_MODE_CUMSUM only)
# outputs
cache_seqlens_ptr, # [bs] int32
cu_seqlens_k_ptr, # [bs + 1] int32
cu_seqlens_q_ptr, # [bs + 1] int32, or None
page_table_ptr, # [bs, page_table_stride] int32
swa_page_table_ptr, # [bs, swa_page_table_stride] int32, or None
swa_out_cache_loc_ptr, # [swa_out_len] int64, or None
# scalars
bs,
seqlen_offset, # added to seq_lens for cache_seqlens / cu_seqlens_k
max_seq_pages, # page-table columns to (re)write per row
q_stride, # Q_MODE_STRIDED stride
num_out_tokens, # valid prefix of out_cache_loc
swa_out_len, # full swa_out_cache_loc length (zero-padded tail)
req_to_token_stride,
page_table_stride,
swa_page_table_stride,
# constexpr
PAGE_SIZE: tl.constexpr,
HAS_SWA: tl.constexpr,
HAS_SWA_OUT: tl.constexpr,
Q_MODE: tl.constexpr,
PAGE_BLOCK: tl.constexpr,
BS_BLOCK: tl.constexpr,
):
pid = tl.program_id(axis=0)
if pid < bs:
# One program per batch row: cache_seqlens + page table row(s).
req_pool_index = tl.load(req_pool_indices_ptr + pid).to(tl.int64)
seqlen = (tl.load(seq_lens_ptr + pid) + seqlen_offset).to(tl.int32)
tl.store(cache_seqlens_ptr + pid, seqlen)
row_in = req_to_token_ptr + req_pool_index * req_to_token_stride
row_out = page_table_ptr + pid.to(tl.int64) * page_table_stride
if HAS_SWA:
swa_row_out = swa_page_table_ptr + pid.to(tl.int64) * swa_page_table_stride
for i in range(tl.cdiv(max_seq_pages, PAGE_BLOCK)):
page_idx = i * PAGE_BLOCK + tl.arange(0, PAGE_BLOCK)
mask = page_idx < max_seq_pages
token = tl.load(
row_in + page_idx.to(tl.int64) * PAGE_SIZE, mask=mask, other=0
)
tl.store(row_out + page_idx, token // PAGE_SIZE, mask=mask)
if HAS_SWA:
token64 = token.to(tl.int64)
# Real req_to_token slots are >=0; the token>=0 guard + other=-1 mirror
# the swa_out_cache_loc -1 sentinel (uniform handling, no wrap).
swa_token = tl.load(
swa_mapping_ptr + token64, mask=mask & (token64 >= 0), other=-1
)
swa_page = tl.where(swa_token < 0, -1, swa_token // PAGE_SIZE)
tl.store(swa_row_out + page_idx, swa_page.to(tl.int32), mask=mask)
elif pid == bs:
# Single program: cu_seqlens_k (+ optional cu_seqlens_q) cumsum.
offs = tl.arange(0, BS_BLOCK)
mask = offs < bs
seqlens = (tl.load(seq_lens_ptr + offs, mask=mask, other=0)).to(tl.int32)
seqlens = tl.where(mask, seqlens + seqlen_offset, 0)
tl.store(cu_seqlens_k_ptr + 1 + offs, tl.cumsum(seqlens, axis=0), mask=mask)
if Q_MODE == 1: # Q_MODE_CUMSUM
qlens = tl.load(qlens_ptr + offs, mask=mask, other=0).to(tl.int32)
qlens = tl.where(mask, qlens, 0)
tl.store(cu_seqlens_q_ptr + 1 + offs, tl.cumsum(qlens, axis=0), mask=mask)
if Q_MODE == 2: # Q_MODE_STRIDED
tl.store(
cu_seqlens_q_ptr + 1 + offs,
((offs + 1) * q_stride).to(tl.int32),
mask=mask,
)
else:
# Remaining programs: swa_out_cache_loc translate + zero padding.
if HAS_SWA_OUT:
out_idx = (pid - bs - 1) * PAGE_BLOCK + tl.arange(0, PAGE_BLOCK)
in_range = out_idx < swa_out_len
is_real = in_range & (out_idx < num_out_tokens)
loc = tl.load(out_cache_loc_ptr + out_idx, mask=is_real, other=0)
translated = tl.load(
swa_mapping_ptr + loc, mask=is_real & (loc >= 0), other=0
)
translated = tl.where(is_real & (loc < 0), -1, translated)
tl.store(swa_out_cache_loc_ptr + out_idx, translated, mask=in_range)
def update_trtllm_mha_graph_metadata(
*,
req_pool_indices,
seq_lens,
req_to_token,
cache_seqlens,
cu_seqlens_k,
page_table,
bs: int,
seqlen_offset: int,
max_seq_pages: int,
page_size: int,
swa_mapping=None,
swa_page_table=None,
out_cache_loc=None,
swa_out_cache_loc=None,
cu_seqlens_q=None,
qlens=None,
q_stride: int = 0,
q_mode: int = Q_MODE_NONE,
):
"""Launch the fused metadata update (one kernel for the whole replay init)."""
if bs == 0:
return
# Launch-block width: page-table columns each program writes per iteration
# (also the swa_out_cache_loc tile width). 512 keeps the per-program working
# set small enough to stay off the register-pressure / occupancy cliff while
# being wide enough to cover the static page-table width in few iterations.
PAGE_BLOCK = 512
has_swa = swa_page_table is not None
has_swa_out = swa_out_cache_loc is not None
swa_out_len = swa_out_cache_loc.shape[0] if has_swa_out else 0
if has_swa_out and out_cache_loc is not None:
num_out_tokens = min(swa_out_len, out_cache_loc.shape[0])
else:
num_out_tokens = 0
if num_out_tokens == 0:
# All loads are masked out; pass a valid dummy pointer for codegen.
out_cache_loc = swa_out_cache_loc
grid_extra = triton.cdiv(swa_out_len, PAGE_BLOCK) if has_swa_out else 0
grid = (bs + 1 + grid_extra,)
update_trtllm_mha_graph_metadata_kernel[grid](
req_pool_indices,
seq_lens,
req_to_token,
swa_mapping,
out_cache_loc,
qlens,
cache_seqlens,
cu_seqlens_k,
cu_seqlens_q,
page_table,
swa_page_table,
swa_out_cache_loc,
bs,
seqlen_offset,
max_seq_pages,
q_stride,
num_out_tokens,
swa_out_len,
req_to_token.stride(0),
page_table.stride(0),
swa_page_table.stride(0) if has_swa else 0,
PAGE_SIZE=page_size,
HAS_SWA=has_swa,
HAS_SWA_OUT=has_swa_out,
Q_MODE=q_mode,
PAGE_BLOCK=PAGE_BLOCK,
BS_BLOCK=triton.next_power_of_2(bs),
)
python/sglang/srt/layers/attention/trtllm_mha_backend.py
+111 -77
View File
@@ -19,6 +19,11 @@ 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_graph_metadata import (
Q_MODE_NONE,
Q_MODE_STRIDED,
update_trtllm_mha_graph_metadata,
)
from sglang.srt.layers.attention.triton_ops.trtllm_mha_page_table import (
build_trtllm_mha_page_table,
)
@@ -136,6 +141,16 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
# SWA hybrid models split the KV cache into full and SWA pools with
# separate index spaces; SWA layers need a translated page_table.
self._swa_kv_pool: Optional[SWAKVPool] = self._resolve_swa_kv_pool(model_runner)
# Raw full->swa index mapping tensor for the fused cuda-graph
# metadata kernel (gather + // page_size happen on device).
if self._swa_kv_pool is not None:
self._swa_full_to_swa_mapping = self._swa_kv_pool.full_to_swa_index_mapping
assert self._swa_full_to_swa_mapping is not None, (
"SWA pool must register full_to_swa_index_mapping before "
"TRTLLMHAAttnBackend is constructed"
)
else:
self._swa_full_to_swa_mapping = None
# 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
@@ -307,8 +322,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
}
# SWA write-target buffer; bound as a [:num_tokens] view in
# _build_cuda_graph_metadata, refilled before each replay in
# init_forward_metadata_out_graph.
# _build_cuda_graph_metadata and refilled by the fused metadata kernel.
self.cuda_graph_swa_out_cache_loc = (
torch.zeros(max_num_tokens, dtype=torch.int64, device=self.device)
if self.use_sliding_window_kv_pool
@@ -471,7 +485,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
)
self.draft_extend_metadata[bs] = metadata
# Bind the SWA write-target buffer slice (refilled at replay).
# Bind the SWA write-target buffer slice (refilled by in-graph metadata).
if self.use_sliding_window_kv_pool:
metadata.swa_out_cache_loc = self.cuda_graph_swa_out_cache_loc[:num_tokens]
@@ -484,73 +498,86 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
seq_lens: torch.Tensor,
forward_mode: ForwardMode,
spec_info: Optional[SpecInput],
out_cache_loc: Optional[torch.Tensor] = None,
):
"""Shared capture+replay body for the cuda-graph init path.
Public entry: :py:meth:`init_forward_metadata_out_graph`.
One fused triton kernel (update_trtllm_mha_graph_metadata) rebuilds
cache_seqlens, cu_seqlens_k/q, the page table(s), and swa_out_cache_loc.
The previous aten-op implementation issued ~25 host dispatches per graph
replay, whose per-rank jitter was paid as spin time inside the first
all-reduce of every replayed graph.
The page table is rewritten to the static ``max_num_pages`` width (the
same upper bound ``_fill_page_table_device`` uses); the kernel
bounds the actual KV reads by the on-device ``cache_seqlens``, so no
runtime host max / seq_lens_cpu D2H sync is needed.
Public entry: :py:meth:`init_forward_metadata_in_graph`.
"""
seq_lens = seq_lens[:bs]
req_pool_indices = req_pool_indices[:bs]
# 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
seqlen_offset = 0
cu_seqlens_q = None
qlens = None
q_stride = 0
q_mode = Q_MODE_NONE
if forward_mode.is_decode_or_idle():
if spec_info is not None:
# Draft Decode
# Here we only support topk = 1 for now.
metadata = self.decode_cuda_graph_metadata[bs]
metadata.cache_seqlens_int32 = self.decode_cuda_graph_metadata[
"cache_seqlens"
][:bs]
metadata.cache_seqlens_int32.copy_(
seq_lens + self.speculative_step_id + 1
)
seqlen_offset = self.speculative_step_id + 1
else:
# Normal Decode
metadata = self.decode_cuda_graph_metadata[bs]
metadata.cache_seqlens_int32.copy_(seq_lens)
metadata.cu_seqlens_k[1:].copy_(
torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32)
)
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.cu_seqlens_k[1:].copy_(
torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32)
)
self._fill_page_table_device(
metadata, req_pool_indices, metadata.cache_seqlens_int32
)
seqlen_offset = metadata.max_seq_len_q
elif forward_mode.is_draft_extend_v2():
metadata = self.draft_extend_metadata[bs]
metadata.cache_seqlens_int32.copy_(seq_lens)
metadata.cu_seqlens_k[1:].copy_(
torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32)
)
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,
)
)
self._fill_page_table_device(
metadata, req_pool_indices, metadata.cache_seqlens_int32
# Static per-request query width, fixed by the captured graph shape.
# Do not inspect replay-time tensors here; this body is recorded into
# the CUDA graph.
num_tokens_per_bs = metadata.max_seq_len_q
cu_seqlens_q = metadata.cu_seqlens_q
q_stride = num_tokens_per_bs
q_mode = Q_MODE_STRIDED
else:
raise ValueError(
"TRTLLM-MHA CUDA graph metadata build got an unsupported forward "
f"mode: {forward_mode}"
)
assert metadata is not None
# Static upper-bound page-table width (see docstring); the kernel
# bounds real KV reads by cache_seqlens, so this is a fixed loop
# bound only — never a host max / seq_lens_cpu D2H sync.
max_seq_pages = self.max_num_pages
update_trtllm_mha_graph_metadata(
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
req_to_token=self.req_to_token,
cache_seqlens=metadata.cache_seqlens_int32,
cu_seqlens_k=metadata.cu_seqlens_k,
page_table=metadata.page_table,
bs=bs,
seqlen_offset=seqlen_offset,
max_seq_pages=max_seq_pages,
page_size=self.page_size,
swa_mapping=self._swa_full_to_swa_mapping,
swa_page_table=metadata.swa_page_table,
out_cache_loc=out_cache_loc,
swa_out_cache_loc=metadata.swa_out_cache_loc,
cu_seqlens_q=cu_seqlens_q,
qlens=qlens,
q_stride=q_stride,
q_mode=q_mode,
)
self.forward_metadata = metadata
def update_verify_buffers_to_fill_after_draft(
@@ -598,44 +625,36 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
in_capture: bool = False,
):
bs = forward_batch.batch_size
req_pool_indices = forward_batch.req_pool_indices
seq_lens = forward_batch.seq_lens
encoder_lens = forward_batch.encoder_lens
forward_mode = forward_batch.forward_mode
spec_info = forward_batch.spec_info
if in_capture:
num_tokens = forward_batch.positions.numel()
self._build_cuda_graph_metadata(
bs, num_tokens, forward_mode, spec_info, seq_lens.device
)
self._apply_cuda_graph_metadata(
bs=bs,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
forward_mode=forward_mode,
spec_info=spec_info,
)
else:
self._apply_cuda_graph_metadata(
bs=bs,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
forward_mode=forward_mode,
spec_info=spec_info,
bs, num_tokens, forward_mode, spec_info, forward_batch.seq_lens.device
)
# Refill the SWA write-target buffer from the live out_cache_loc before
# replay (the per-bs metadata holds a view bound in _build).
if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None:
n = forward_batch.out_cache_loc.shape[0]
self.cuda_graph_swa_out_cache_loc[n:].zero_()
self.cuda_graph_swa_out_cache_loc[:n].copy_(
self.token_to_kv_pool.translate_loc_from_full_to_swa(
forward_batch.out_cache_loc
)
if forward_mode.is_decode_or_idle():
self.forward_metadata = self.decode_cuda_graph_metadata[bs]
elif forward_mode.is_target_verify():
self.forward_metadata = self.target_verify_metadata[bs]
elif forward_mode.is_draft_extend_v2():
self.forward_metadata = self.draft_extend_metadata[bs]
else:
raise ValueError(
f"Invalid forward mode: {forward_mode=} for CUDA Graph replay."
)
def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch):
self._apply_cuda_graph_metadata(
bs=forward_batch.batch_size,
req_pool_indices=forward_batch.req_pool_indices,
seq_lens=forward_batch.seq_lens,
forward_mode=forward_batch.forward_mode,
spec_info=forward_batch.spec_info,
out_cache_loc=forward_batch.out_cache_loc,
)
def init_forward_metadata(self, forward_batch: ForwardBatch):
"""Initialize the metadata for a forward pass."""
@@ -980,3 +999,18 @@ class TRTLLMHAAttnMultiStepDraftBackend(FlashInferMultiStepDraftBackend):
self.attn_backends[i].init_forward_metadata_out_graph(
inner_fb, in_capture=in_capture
)
def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None:
from sglang.srt.model_executor.forward_batch_info import build_inner_fb_view
assert forward_batch.spec_info is not None
assert forward_batch.spec_info.is_draft_input()
inner_fb = build_inner_fb_view(
forward_batch,
bs=forward_batch.batch_size,
forward_mode=ForwardMode.DECODE,
encoder_lens=forward_batch.encoder_lens,
)
for i in range(self.speculative_num_steps - 1):
self.attn_backends[i].init_forward_metadata_in_graph(inner_fb)
python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
+2
View File
@@ -385,6 +385,8 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
)
def run_once():
self.draft_extend_attn_backend.init_forward_metadata_in_graph(forward_batch)
# Clean intermediate result cache for DP attention
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
set_dp_buffer_len(
python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py
+9
View File
@@ -294,6 +294,15 @@ class FrozenKVMTPCudaGraphRunner(DecodeCudaGraphRunner):
)
def run_once():
# Record the metadata rebuild against the committed target-prefix
# geometry (spec_info nulled → plain target-length decode), matching
# every other frozen-KV metadata init. Without the view, backends
# that key seqlen offsets off spec_info (trtllm_mha's draft-decode
# branch adds speculative_step_id + 1) bake a +1 offset into the
# captured graph and replay reads one extra, never-written KV slot.
with self.frozen_kv_mtp_worker._frozen_kv_target_view(forward_batch):
self.draft_attn_backend.init_forward_metadata_in_graph(forward_batch)
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
set_dp_buffer_len(
global_dp_buffer_len,
python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py
+2
View File
@@ -337,6 +337,8 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
attn_backend = self.eagle_worker.draft_extend_attn_backend_list[self.step]
def run_once():
attn_backend.init_forward_metadata_in_graph(forward_batch)
# Clean intermediate result cache for DP attention
forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
set_dp_buffer_len(