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, :]
-    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

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