diff --git a/python/sglang/srt/layers/attention/qsa/sparse_attn.py b/python/sglang/srt/layers/attention/qsa/sparse_attn.py index e07e1ae6f..601945f05 100644 --- a/python/sglang/srt/layers/attention/qsa/sparse_attn.py +++ b/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, ) diff --git a/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py b/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py index 4f8eaefef..54b28ceff 100644 --- a/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py +++ b/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( diff --git a/test/registered/kernel/qsa/test_qsa_strided_zero_fill.py b/test/registered/kernel/qsa/test_qsa_strided_zero_fill.py new file mode 100644 index 000000000..9e8422884 --- /dev/null +++ b/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"]))