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

fix(qsa): make the paged sparse-decode gather memory-safe (zero-fill scratch, int64 offsets, dequant FP8 on gather) (#38851)

Browse Source
This commit is contained in:
YAMY
2026-09-11 15:47:25 -07:00
committed by GitHub
parent ec30f19e4a
commit f963d7a27c
3 changed files with 265 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/qsa/sparse_attn.py
+68 -7
View File
@@ -388,8 +388,10 @@ def _compact_kv(
dim: tl.constexpr,
req_stride: tl.constexpr,
idx_stride: tl.constexpr,
pad_cols,
BLOCK_TOPK: tl.constexpr,
BLOCK_D: tl.constexpr,
ZERO_FILL: tl.constexpr,
):
batch, head, block = tl.program_id(0), tl.program_id(1), tl.program_id(2)
cols = block * BLOCK_TOPK + tl.arange(0, BLOCK_TOPK)
@@ -405,11 +407,39 @@ def _compact_kv(
mask=valid,
other=0,
)
src = slots[:, None] * heads * dim + head * dim + dims[None, :]
dst = (pack_start + cols)[:, None] * heads * dim + head * dim + dims[None, :]
mask = valid[:, None] & (dims[None, :] < dim)
tl.store(out_k + dst, tl.load(k + src, mask=mask, other=0.0), mask=mask)
tl.store(out_v + dst, tl.load(v + src, mask=mask, other=0.0), mask=mask)
# 64-bit element offsets: slot * heads * dim exceeds int32 once the pool holds
# more than 2^31 / (heads * dim) tokens (~4.2M for 2 x 256), which an FP8 pool
# on one GPU does reach.
src = slots.to(tl.int64)[:, None] * heads * dim + head * dim + dims[None, :]
dst = (
(pack_start + cols).to(tl.int64)[:, None] * heads * dim
+ head * dim
+ dims[None, :]
)
load_mask = valid[:, None] & (dims[None, :] < dim)
if ZERO_FILL:
# Strided (page-aligned) packing: the paged decode kernel reads whole pages,
# so every slot in [valid_count, pad_cols) must hold zeros, never stale bytes.
# `valid_count` here is the row's page-aligned stride, not its valid count, so
# the store covers the full region while the load stays limited to valid rows.
store_mask = (cols < pad_cols)[:, None] & (dims[None, :] < dim)
else:
store_mask = load_mask
# Dequantize while gathering: the scratch is allocated in the query dtype, so an
# FP8 pool is read as fp8 and stored as bf16. The QSA backend writes the pool
# without per-tensor k/v scales (see set_kv_buffer calls in
# qwen_sparse_attn_backend.py), so no scale is applied here either.
out_dtype = out_k.dtype.element_ty
tl.store(
out_k + dst,
tl.load(k + src, mask=load_mask, other=0.0).to(out_dtype),
mask=store_mask,
)
tl.store(
out_v + dst,
tl.load(v + src, mask=load_mask, other=0.0).to(out_dtype),
mask=store_mask,
)
def qwen_sparse_valid_counts_triton(seq_lens, indices, counts, batch, topk):
@@ -426,11 +456,40 @@ def qwen_sparse_valid_counts_triton(seq_lens, indices, counts, batch, topk):
def qwen_sparse_kv_extraction_compact_triton(
k, v, req_to_token, req_indices, indices, seq_lens, cu_k, out_k, out_v, batch, topk
k,
v,
req_to_token,
req_indices,
indices,
seq_lens,
cu_k,
out_k,
out_v,
batch,
topk,
zero_fill_cols: int = 0,
):
"""Gather the selected K/V rows into ``out_k``/``out_v``.
``zero_fill_cols`` > 0 selects the strided (page-aligned) layout used by the paged
decode kernel: row ``b`` owns ``[cu_k[b], cu_k[b] + zero_fill_cols)`` and every slot
past its valid rows is zero-filled. Paged kernels read whole pages and multiply the
masked probabilities into V, so stale or uninitialized bytes there (NaN/Inf bit
patterns) would otherwise leak into the output. ``0`` keeps the compact layout for
the varlen fallback, whose rows are packed back-to-back.
``out_k``/``out_v`` may use a wider dtype than the pool (bf16 scratch for an FP8
pool); rows are converted while gathering.
Both layouts assume the valid entries of each ``indices`` row are contiguous at
the front (``expand_qsa_block_indices`` sorts them that way): ``valid_count`` is a
count, not a mask, so a ``-1`` in the middle of a row would shift the packing.
"""
_, heads, dim = k.shape
block_topk = 16
_compact_kv[(batch, heads, triton.cdiv(topk, block_topk))](
zero_fill = zero_fill_cols > 0
num_cols = zero_fill_cols if zero_fill else topk
_compact_kv[(batch, heads, triton.cdiv(num_cols, block_topk))](
k,
v,
req_to_token,
@@ -445,8 +504,10 @@ def qwen_sparse_kv_extraction_compact_triton(
dim,
req_to_token.stride(0),
indices.stride(0),
num_cols,
BLOCK_TOPK=block_topk,
BLOCK_D=triton.next_power_of_2(dim),
ZERO_FILL=zero_fill,
num_warps=8,
)
python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py
+5 -2
View File
@@ -1454,11 +1454,13 @@ class QwenSparseAttnBackend(AttentionBackend):
batch, pages_per_row, page, device
)
capacity_rows = self._cuda_graph_max_tokens if metadata.is_cuda_graph else batch
# Gather into the query dtype: an FP8 pool is dequantized on the way in, so the
# paged kernel always runs the bf16 q + bf16 KV path.
packed_k, packed_v = self._get_fa2_scratch(
max(capacity_rows, batch) * stride,
k_buffer.shape[1],
k_buffer.shape[2],
k_buffer.dtype,
q.dtype,
k_buffer.device,
)
qwen_sparse_kv_extraction_compact_triton(
@@ -1477,6 +1479,7 @@ class QwenSparseAttnBackend(AttentionBackend):
packed_v,
batch,
topk,
zero_fill_cols=stride,
)
num_kv_heads = k_buffer.shape[1]
head_dim = k_buffer.shape[2]
@@ -1587,7 +1590,7 @@ class QwenSparseAttnBackend(AttentionBackend):
scratch_capacity,
k_buffer.shape[1],
k_buffer.shape[2],
k_buffer.dtype,
q.dtype,
k_buffer.device,
)
qwen_sparse_kv_extraction_compact_triton(
test/registered/kernel/qsa/test_qsa_strided_zero_fill.py
+192
View File
@@ -0,0 +1,192 @@
"""Regression test for the QSA strided sparse-decode scratch zero-fill.
Poison the packed scratch with NaN, gather with the strided layout used by
`_forward_trtllm_sparse`, and require that (a) valid rows are copied exactly and
(b) every slot in [valid_count, stride) is zero, so the paged decode kernel can never
multiply masked probabilities into stale NaN/Inf bytes. Also checks the compact
(FA2 fallback) layout is unchanged. Intended for test/registered/kernel/qsa/.
"""
import sys
import pytest
import torch
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=120, stage="base-b-kernel-unit", runner_config="1-gpu-large")
from sglang.srt.layers.attention.qsa.sparse_attn import (
qwen_sparse_fa2_cu_seqlens_triton,
qwen_sparse_kv_extraction_compact_triton,
)
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn])
def test_strided_gather_zero_fills_tail(dtype):
if not torch.cuda.is_available():
pytest.skip("CUDA required")
torch.manual_seed(0)
device = torch.device("cuda")
batch, topk, page, heads, dim = 3, 2051, 64, 2, 256
pages_per_row = (topk + page - 1) // page
stride = pages_per_row * page
pool_rows = 8192
k_pool = torch.randn(pool_rows, heads, dim, device=device, dtype=torch.bfloat16).to(
dtype
)
v_pool = torch.randn(pool_rows, heads, dim, device=device, dtype=torch.bfloat16).to(
dtype
)
seq_lens = torch.tensor([733, 109, 2500], device=device, dtype=torch.int32)
req_to_token = (
torch.randperm(pool_rows, device=device)[: batch * 2600]
.reshape(batch, 2600)
.to(torch.int32)
)
req_indices = torch.arange(batch, device=device, dtype=torch.int32)
# top-k rows: the first min(seq_len, topk) logical positions, then -1 padding
indices = torch.full((batch, topk), -1, device=device, dtype=torch.int32)
for b in range(batch):
n = min(int(seq_lens[b]), topk)
indices[b, :n] = torch.arange(n, device=device, dtype=torch.int32)
cu_strided = torch.arange(batch + 1, device=device, dtype=torch.int32) * stride
# the scratch is always in the compute dtype (bf16); an FP8 pool is dequantized on the way in
packed_k = torch.full(
(batch * stride, heads, dim), float("nan"), device=device, dtype=torch.bfloat16
)
packed_v = packed_k.clone()
qwen_sparse_kv_extraction_compact_triton(
k_pool,
v_pool,
req_to_token,
req_indices,
indices,
seq_lens,
cu_strided,
packed_k,
packed_v,
batch,
topk,
zero_fill_cols=stride,
)
pk, pv = (
packed_k.float().view(batch, stride, heads, dim),
packed_v.float().view(batch, stride, heads, dim),
)
assert torch.isfinite(pk).all() and torch.isfinite(pv).all()
for b in range(batch):
n = min(int(seq_lens[b]), topk)
slots = req_to_token[b, :n].long()
torch.testing.assert_close(pk[b, :n], k_pool[slots].to(torch.bfloat16).float())
torch.testing.assert_close(pv[b, :n], v_pool[slots].to(torch.bfloat16).float())
assert (pk[b, n:] == 0).all() and (pv[b, n:] == 0).all()
def test_compact_gather_unchanged():
if not torch.cuda.is_available():
pytest.skip("CUDA required")
torch.manual_seed(0)
device = torch.device("cuda")
batch, topk, heads, dim = 2, 2051, 2, 256
k_pool = torch.randn(4096, heads, dim, device=device, dtype=torch.bfloat16)
v_pool = torch.randn(4096, heads, dim, device=device, dtype=torch.bfloat16)
seq_lens = torch.tensor([300, 50], device=device, dtype=torch.int32)
req_to_token = torch.arange(batch * 512, device=device, dtype=torch.int32).reshape(
batch, 512
)
req_indices = torch.arange(batch, device=device, dtype=torch.int32)
indices = torch.full((batch, topk), -1, device=device, dtype=torch.int32)
for b in range(batch):
indices[b, : int(seq_lens[b])] = torch.arange(
int(seq_lens[b]), device=device, dtype=torch.int32
)
counts = torch.empty(batch, device=device, dtype=torch.int32)
cu_k = torch.empty(batch + 1, device=device, dtype=torch.int32)
qwen_sparse_fa2_cu_seqlens_triton(seq_lens, indices, counts, cu_k, batch, topk)
assert cu_k.tolist() == [0, 300, 350]
packed_k = torch.full(
(batch * topk, heads, dim), float("nan"), device=device, dtype=torch.bfloat16
)
packed_v = packed_k.clone()
qwen_sparse_kv_extraction_compact_triton(
k_pool,
v_pool,
req_to_token,
req_indices,
indices,
seq_lens,
cu_k,
packed_k,
packed_v,
batch,
topk,
)
torch.testing.assert_close(packed_k[:300], k_pool[req_to_token[0, :300].long()])
torch.testing.assert_close(packed_k[300:350], k_pool[req_to_token[1, :50].long()])
# compact layout leaves the region past the packed rows untouched (still NaN)
assert torch.isnan(packed_k[350:]).all()
def test_strided_gather_addresses_pool_beyond_int32_elements():
"""Slots past 2^31 / (heads * dim) must be addressed with 64-bit offsets.
An FP8 KV pool on one GB300 holds ~7.6M tokens for Qwen3.8-Flash-Next (2 kv heads x 256),
so slot indices above 4,194,304 occur in production; int32 element offsets wrap there.
"""
if not torch.cuda.is_available():
pytest.skip("CUDA required")
if torch.cuda.get_device_properties(0).total_memory < 6 * 1024**3:
pytest.skip("needs ~2.5 GB of device memory")
torch.manual_seed(0)
device = torch.device("cuda")
heads, dim = 2, 256
threshold = (1 << 31) // (heads * dim) # 4,194,304
pool_rows = threshold + 4096
k_pool = torch.zeros(
pool_rows, heads, dim, device=device, dtype=torch.float8_e4m3fn
)
v_pool = torch.zeros(
pool_rows, heads, dim, device=device, dtype=torch.float8_e4m3fn
)
hi = torch.arange(threshold + 64, threshold + 64 + 300, device=device)
k_pool[hi] = torch.randn(300, heads, dim, device=device, dtype=torch.bfloat16).to(
torch.float8_e4m3fn
)
v_pool[hi] = torch.randn(300, heads, dim, device=device, dtype=torch.bfloat16).to(
torch.float8_e4m3fn
)
batch, topk, page = 1, 2051, 64
stride = ((topk + page - 1) // page) * page
seq_lens = torch.tensor([300], device=device, dtype=torch.int32)
req_to_token = torch.zeros(batch, 512, device=device, dtype=torch.int32)
req_to_token[0, :300] = hi.to(torch.int32)
indices = torch.full((batch, topk), -1, device=device, dtype=torch.int32)
indices[0, :300] = torch.arange(300, device=device, dtype=torch.int32)
cu_strided = torch.arange(batch + 1, device=device, dtype=torch.int32) * stride
packed_k = torch.full(
(batch * stride, heads, dim), float("nan"), device=device, dtype=torch.bfloat16
)
packed_v = packed_k.clone()
qwen_sparse_kv_extraction_compact_triton(
k_pool,
v_pool,
req_to_token,
torch.zeros(1, device=device, dtype=torch.int32),
indices,
seq_lens,
cu_strided,
packed_k,
packed_v,
batch,
topk,
zero_fill_cols=stride,
)
torch.testing.assert_close(packed_k[:300], k_pool[hi].to(torch.bfloat16))
torch.testing.assert_close(packed_v[:300], v_pool[hi].to(torch.bfloat16))
assert (packed_k[300:] == 0).all() and (packed_v[300:] == 0).all()
if __name__ == "__main__":
sys.exit(pytest.main([__file__, "-v"]))