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

python/sglang/srt/layers/attention/qsa/sparse_attn.py

@@ -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, :]
+    # 64-bit element offsets: slot * heads * dim exceeds int32 once the pool holds
-    dst = (pack_start + cols)[:, None] * heads * dim + head * dim + dims[None, :]
+    # more than 2^31 / (heads * dim) tokens (~4.2M for 2 x 256), which an FP8 pool
-    mask = valid[:, None] & (dims[None, :] < dim)
+    # on one GPU does reach.
-    tl.store(out_k + dst, tl.load(k + src, mask=mask, other=0.0), mask=mask)
+    src = slots.to(tl.int64)[:, None] * heads * dim + head * dim + dims[None, :]
-    tl.store(out_v + dst, tl.load(v + src, mask=mask, other=0.0), mask=mask)
+    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

@@ -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

@@ -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"]))