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

[AMD] Parallelize aiter spec-decode KV index building over token blocks (#37659)

Browse Source 
Co-authored-by: Zijie Chen <300606707+zijiecode@users.noreply.github.com>
Co-authored-by: jacky.cheng <yichiche@amd.com>
This commit is contained in:
zijiec
2026-09-09 17:53:53 -07:00
committed by GitHub
co-authored by Zijie Chen jacky.cheng
parent 880d6fa64d
commit 92d831d3d7
5 changed files with 361 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/utils.py
+3
View File
@@ -48,6 +48,9 @@ from sglang.kernels.ops.kvcache.kv_indices import (
from sglang.kernels.ops.kvcache.kv_indices import (
get_num_page_per_block_flashmla as get_num_page_per_block_flashmla,
)
from sglang.kernels.ops.kvcache.kv_indices import (
kv_indices_num_token_blocks as kv_indices_num_token_blocks,
)
from sglang.kernels.ops.kvcache.rope_cache import (
fused_qk_rope_reshape_and_cache as fused_qk_rope_reshape_and_cache,
)
python/sglang/kernels/ops/kvcache/kv_indices.py
+31 -1
View File
@@ -4,6 +4,19 @@ import triton.language as tl
_FLASHMLA_CREATE_KV_BLOCK_SIZE = 4096
FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON = tl.constexpr(_FLASHMLA_CREATE_KV_BLOCK_SIZE)
# Token-block parallelism for the index-copy kernels below: aim for about
# _TARGET_PROGRAMS programs in total, one extra block per
# _MIN_TOKENS_PER_BLOCK of table width at most, and fall back to the
# historical single-block grid when the base grid is already wide.
_MIN_TOKENS_PER_BLOCK = 8192
_TARGET_PROGRAMS = 512
def kv_indices_num_token_blocks(table_width: int, base_programs: int) -> int:
cap = (table_width + _MIN_TOKENS_PER_BLOCK - 1) // _MIN_TOKENS_PER_BLOCK
want = _TARGET_PROGRAMS // max(1, base_programs)
return max(1, min(cap, want))
@triton.jit
def create_flashinfer_kv_indices_triton(
@@ -19,6 +32,7 @@ def create_flashinfer_kv_indices_triton(
# (a recompile every few decode steps at small page sizes).
req_to_token_ptr_stride,
ENTRY_PAGE_SIZE: tl.constexpr = 1,
TOKEN_BLOCK_PARALLEL: tl.constexpr = False,
):
"""Gather per-request token ids into a flat CSR kv_indices stream.
@@ -28,9 +42,21 @@ def create_flashinfer_kv_indices_triton(
read table (entries already kernel-facing page ids); token ids are rebuilt
as ``token = entry * ps + pos % ps``, exact because converting an id keeps
its offset inside the page.
``TOKEN_BLOCK_PARALLEL`` (default False): launched on a 2D grid
``(batch, num_blocks)``, the programs of a request stride over its copy
loop together instead of one program crawling the whole context serially
(which bottlenecks long-context spec decode, where this kernel runs every
iteration). With the default, the kernel is the historical
one-program-per-request loop and 1D launch sites are unaffected.
"""
BLOCK_SIZE: tl.constexpr = 512
pid = tl.program_id(axis=0)
if TOKEN_BLOCK_PARALLEL:
blk = tl.program_id(axis=1)
num_blk = tl.num_programs(axis=1)
else:
blk = 0
num_blk = 1
# find the req pool idx, this is for batch to token
req_pool_index = tl.load(req_pool_indices_ptr + pid).to(tl.int64)
@@ -44,7 +70,11 @@ def create_flashinfer_kv_indices_triton(
kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32)
num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE)
for i in range(num_loop):
if TOKEN_BLOCK_PARALLEL:
# Blocks with no copy work exit early.
if blk >= num_loop:
return
for i in range(blk, num_loop, num_blk):
# index into req_to_token_ptr needs to be int64
offset = tl.arange(0, BLOCK_SIZE).to(tl.int64) + i * BLOCK_SIZE
mask = offset < kv_end - kv_start
python/sglang/kernels/ops/speculative/cache_locs.py
+93 -42
View File
@@ -67,13 +67,30 @@ def generate_draft_decode_kv_indices(
iter_upper: tl.constexpr,
num_tokens_upper: tl.constexpr,
page_size: tl.constexpr,
NUM_STEPS: tl.constexpr = 0,
):
BLOCK_SIZE: tl.constexpr = 128
iters = tl.program_id(axis=0)
# Optional token-block parallelism (NUM_STEPS > 0): the first grid axis
# packs (draft step, token block) as ``step + NUM_STEPS * block``,
# spreading the per-request index copy below over many programs instead
# of one program crawling the whole context serially (which bottlenecks
# long-context spec decode, where this kernel runs every iteration).
# NUM_STEPS == 0 (default) is the historical one-program-per-step kernel:
# the same 128-wide copy loop, in the same order, with the token-block
# branches folded away at compile time.
BLOCK_SIZE: tl.constexpr = 128 if NUM_STEPS == 0 else 512
pid0 = tl.program_id(axis=0)
bid = tl.program_id(axis=1)
topk_id = tl.program_id(axis=2)
num_steps = tl.num_programs(axis=0)
if NUM_STEPS == 0:
iters = pid0
num_steps = tl.num_programs(axis=0)
blk = 0
else:
iters = pid0 % NUM_STEPS
blk = pid0 // NUM_STEPS
num_steps = NUM_STEPS
num_blk = tl.num_programs(axis=0) // NUM_STEPS
num_seqs = tl.num_programs(axis=1)
topk = tl.num_programs(axis=2)
@@ -81,56 +98,90 @@ def generate_draft_decode_kv_indices(
kv_indptr += kv_indptr_stride * iters
iters += 1
load_offset = tl.arange(0, bs_upper)
seq_lens = tl.load(paged_kernel_lens + load_offset, mask=load_offset < bid, other=0)
seq_len = tl.load(paged_kernel_lens + bid)
cum_seq_len = tl.sum(seq_lens)
if NUM_STEPS == 0:
load_offset = tl.arange(0, bs_upper)
seq_lens = tl.load(
paged_kernel_lens + load_offset, mask=load_offset < bid, other=0
)
seq_len = tl.load(paged_kernel_lens + bid)
cum_seq_len = tl.sum(seq_lens)
else:
seq_len = tl.load(paged_kernel_lens + bid)
num_loop = tl.cdiv(seq_len, BLOCK_SIZE)
# Blocks with no copy work exit before the O(bs) prefix-sum below;
# block 0 always continues (it owns the extension and kv_indptr).
if blk >= num_loop and blk > 0:
return
load_offset = tl.arange(0, bs_upper)
seq_lens = tl.load(
paged_kernel_lens + load_offset, mask=load_offset < bid, other=0
)
cum_seq_len = tl.sum(seq_lens)
# Update kv_indices
kv_offset = cum_seq_len * topk + bid * iters * topk + topk_id * (seq_len + iters)
kv_ptr = kv_indices + kv_offset
token_pool_ptr = req_to_token + tl.load(req_pool_indices + bid) * pool_len
kv_offset = tl.arange(0, BLOCK_SIZE)
num_loop = tl.cdiv(seq_len, BLOCK_SIZE)
for _ in range(num_loop):
mask = kv_offset < seq_len
data = tl.load(token_pool_ptr + kv_offset, mask=mask)
tl.store(kv_ptr + kv_offset, data, mask=mask)
kv_offset += BLOCK_SIZE
extend_offset = tl.arange(0, iter_upper)
if page_size == 1 or topk == 1:
extend_data = tl.load(
token_pool_ptr + seq_len + topk_id * num_steps + tl.arange(0, iter_upper),
mask=extend_offset < iters,
)
if NUM_STEPS == 0:
kv_offset = tl.arange(0, BLOCK_SIZE)
num_loop = tl.cdiv(seq_len, BLOCK_SIZE)
for _ in range(num_loop):
mask = kv_offset < seq_len
data = tl.load(token_pool_ptr + kv_offset, mask=mask)
tl.store(kv_ptr + kv_offset, data, mask=mask)
kv_offset += BLOCK_SIZE
else:
prefix_len = seq_len
last_page_len = prefix_len % page_size
num_new_pages_per_topk = (
last_page_len + num_steps + page_size - 1
) // page_size
prefix_base = seq_len // page_size * page_size
start = (
prefix_base + topk_id * num_new_pages_per_topk * page_size + last_page_len
)
extend_data = tl.load(
token_pool_ptr + start + extend_offset,
for i in range(blk, num_loop, num_blk):
tok_off = i * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = tok_off < seq_len
data = tl.load(token_pool_ptr + tok_off, mask=mask)
tl.store(kv_ptr + tok_off, data, mask=mask)
# Extension entries and kv_indptr belong to token block 0 alone; other
# blocks neither compute nor store them.
if blk == 0:
extend_offset = tl.arange(0, iter_upper)
if page_size == 1 or topk == 1:
extend_data = tl.load(
token_pool_ptr
+ seq_len
+ topk_id * num_steps
+ tl.arange(0, iter_upper),
mask=extend_offset < iters,
)
else:
prefix_len = seq_len
last_page_len = prefix_len % page_size
num_new_pages_per_topk = (
last_page_len + num_steps + page_size - 1
) // page_size
prefix_base = seq_len // page_size * page_size
start = (
prefix_base
+ topk_id * num_new_pages_per_topk * page_size
+ last_page_len
)
extend_data = tl.load(
token_pool_ptr + start + extend_offset,
mask=extend_offset < iters,
)
tl.store(
kv_ptr + seq_len + extend_offset,
extend_data,
mask=extend_offset < iters,
)
tl.store(kv_ptr + seq_len + extend_offset, extend_data, mask=extend_offset < iters)
# Update kv_indptr
bs_offset = tl.arange(0, num_tokens_upper)
# Update kv_indptr
bs_offset = tl.arange(0, num_tokens_upper)
zid = bid * topk + topk_id
if zid == 0:
zid = num_seqs * topk
positions = tl.load(positions + bs_offset, mask=bs_offset < zid, other=0)
base = tl.sum(positions)
tl.store(kv_indptr + zid, base + zid * iters)
zid = bid * topk + topk_id
if zid == 0:
zid = num_seqs * topk
pos_vals = tl.load(positions + bs_offset, mask=bs_offset < zid, other=0)
base = tl.sum(pos_vals)
tl.store(kv_indptr + zid, base + zid * iters)
@triton.jit
python/sglang/srt/layers/attention/aiter_backend.py
+39 -6
View File
@@ -23,6 +23,7 @@ from sglang.kernels.ops.attention.utils import (
create_flashinfer_kv_indices_triton,
create_flashmla_kv_indices_triton,
get_num_kv_index_blocks_flashmla,
kv_indices_num_token_blocks,
)
from sglang.kernels.ops.kvcache.aiter_unified_attention import (
scatter_ragged_to_page_table_kernel,
@@ -124,6 +125,13 @@ fast_mode = False
intra_batch_mode = True if _use_mla_ps_kernel else False
# Token-block parallel KV-index building is enabled only where it pays:
# the speculative-decoding paths (target_verify / draft_extend / draft
# decode) of long-context servers. Everything else keeps the historical
# one-program-per-request launch.
_KV_INDEX_BLOCKS_MIN_CONTEXT = 32768
class WrapperDispatch(Enum):
SLIDING_WINDOW = auto()
CROSS_ATTENTION = auto()
@@ -1170,6 +1178,11 @@ class AiterAttnBackend(AttentionBackend):
)
return output[:, : layer.tp_q_head_num, :] if head_pad else output
def _kv_index_blocks(self, bs: int) -> int:
if self.max_context_len < _KV_INDEX_BLOCKS_MIN_CONTEXT:
return 1
return kv_indices_num_token_blocks(self.req_to_token.shape[1], bs)
def init_forward_metadata_out_graph(
self,
forward_batch: ForwardBatch,
@@ -1374,7 +1387,8 @@ class AiterAttnBackend(AttentionBackend):
forward_batch.seq_lens_sum, device
)
create_flashinfer_kv_indices_triton[(bs,)](
num_token_blocks = self._kv_index_blocks(bs)
create_flashinfer_kv_indices_triton[(bs, num_token_blocks)](
self.req_to_token,
forward_batch.req_pool_indices,
forward_batch.seq_lens,
@@ -1382,6 +1396,7 @@ class AiterAttnBackend(AttentionBackend):
None,
kv_indices,
self.req_to_token.stride(0),
TOKEN_BLOCK_PARALLEL=num_token_blocks > 1,
)
if _use_mla_ps_kernel:
@@ -1467,7 +1482,8 @@ class AiterAttnBackend(AttentionBackend):
kv_lens_sum,
device,
)
create_flashinfer_kv_indices_triton[(bs,)](
num_token_blocks = self._kv_index_blocks(bs)
create_flashinfer_kv_indices_triton[(bs, num_token_blocks)](
self.req_to_token,
forward_batch.req_pool_indices,
kv_lens,
@@ -1475,6 +1491,7 @@ class AiterAttnBackend(AttentionBackend):
None,
kv_indices,
self.req_to_token.stride(0),
TOKEN_BLOCK_PARALLEL=num_token_blocks > 1,
)
# if self.kv_cache_dtype == fp8_dtype:
@@ -1564,7 +1581,8 @@ class AiterAttnBackend(AttentionBackend):
kv_indices = torch.empty(
kv_indptr[-1], dtype=torch.int64, device=self.device
)
create_flashinfer_kv_indices_triton[(bs,)](
num_token_blocks = self._kv_index_blocks(bs)
create_flashinfer_kv_indices_triton[(bs, num_token_blocks)](
self.req_to_token,
forward_batch.req_pool_indices,
forward_batch.seq_lens,
@@ -1572,6 +1590,7 @@ class AiterAttnBackend(AttentionBackend):
None,
kv_indices,
self.req_to_token.stride(0),
TOKEN_BLOCK_PARALLEL=num_token_blocks > 1,
)
custom_mask = spec_info.custom_mask
@@ -2018,7 +2037,8 @@ class AiterAttnBackend(AttentionBackend):
bs=bs,
seq_lens_sum=seq_lens_sum,
)
create_flashinfer_kv_indices_triton[(bs,)](
num_token_blocks = self._kv_index_blocks(bs)
create_flashinfer_kv_indices_triton[(bs, num_token_blocks)](
self.req_to_token,
req_pool_indices,
kv_lens,
@@ -2026,6 +2046,7 @@ class AiterAttnBackend(AttentionBackend):
None,
kv_indices,
self.req_to_token.stride(0),
TOKEN_BLOCK_PARALLEL=num_token_blocks > 1,
)
kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs]
@@ -2144,7 +2165,8 @@ class AiterAttnBackend(AttentionBackend):
kv_indptr = self.kv_indptr[: bs + 1]
kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
kv_indices = self.cuda_graph_kv_indices
create_flashinfer_kv_indices_triton[(bs,)](
num_token_blocks = self._kv_index_blocks(bs)
create_flashinfer_kv_indices_triton[(bs, num_token_blocks)](
self.req_to_token,
req_pool_indices,
seq_lens,
@@ -2152,6 +2174,7 @@ class AiterAttnBackend(AttentionBackend):
None,
kv_indices,
self.req_to_token.stride(0),
TOKEN_BLOCK_PARALLEL=num_token_blocks > 1,
)
kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs]
@@ -3544,8 +3567,15 @@ class AiterMultiStepDraftBackend:
bs = self.topk * num_seqs
seq_lens_sum = forward_batch.seq_lens_sum
num_token_blocks = (
kv_indices_num_token_blocks(
self.pool_len, self.speculative_num_steps * num_seqs * self.topk
)
if self.max_context_len >= _KV_INDEX_BLOCKS_MIN_CONTEXT
else 1
)
self.generate_draft_decode_kv_indices[
(self.speculative_num_steps, num_seqs, self.topk)
(self.speculative_num_steps * num_token_blocks, num_seqs, self.topk)
](
forward_batch.req_pool_indices,
self.req_to_token_pool.req_to_token,
@@ -3560,6 +3590,9 @@ class AiterMultiStepDraftBackend:
triton.next_power_of_2(self.speculative_num_steps),
triton.next_power_of_2(bs),
self.page_size,
# A single token block is the historical launch; NUM_STEPS=0 keeps
# its 128-wide program instead of the token-block specialization.
NUM_STEPS=self.speculative_num_steps if num_token_blocks > 1 else 0,
)
for i in range(self.speculative_num_steps - 1):