diff --git a/tests/kernels/test_fused_indexer_q_rope_quant.py b/tests/kernels/test_fused_indexer_q_rope_quant.py index 8770882ddf..7464c90769 100644 --- a/tests/kernels/test_fused_indexer_q_rope_quant.py +++ b/tests/kernels/test_fused_indexer_q_rope_quant.py @@ -560,3 +560,55 @@ def test_indexer_k_store_roundtrips_through_rocm_gather(block_size, compress_rat torch.testing.assert_close( k_scale[emitted], expected_scale[emitted], rtol=0, atol=0 ) + + +@pytest.mark.parametrize("use_cutedsl", [False, True]) +@pytest.mark.skipif( + not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ), + reason="MXFP4 indexer cache requires an SM100-family GPU", +) +@torch.inference_mode() +def test_fused_indexer_q_rope_quant_writes_bf16_weights(use_cutedsl): + """The MXFP4 path can emit the per-head weights in bf16 (what DeepGEMM's + sparse MQA-logits kernels take) instead of fp32; the values are the fp32 + result rounded once, so the scoring kernel needs no cast.""" + if use_cutedsl and not has_cutedsl(): + pytest.skip("cutedsl (cutlass) not installed") + + device = "cuda" + torch.manual_seed(0) + num_tokens, n_head = 257, 32 + q = torch.randn(num_tokens, n_head, HEAD_DIM, dtype=torch.bfloat16, device=device) + positions = torch.randint( + 0, MAX_POS, (num_tokens,), dtype=torch.int64, device=device + ) + cos_sin_cache = torch.randn(MAX_POS, ROPE_DIM, dtype=torch.float32, device=device) + weights = torch.randn(num_tokens, n_head, dtype=torch.bfloat16, device=device) + softmax_scale, head_scale = HEAD_DIM**-0.5, 0.125 # pow2 head_scale, see above + + _, weights_ref = _reference( + positions, q, cos_sin_cache, weights, softmax_scale, head_scale, n_head, True + ) + cutedsl_patch = ( + mock.patch( + "vllm.models.deepseek_v4.common.ops.fused_indexer_q.has_cutedsl", + return_value=False, + ) + if not use_cutedsl + else contextlib.nullcontext() + ) + with cutedsl_patch: + _, weights_fused = fused_indexer_q_rope_quant( + positions, + q.clone(), + cos_sin_cache, + weights, + softmax_scale, + head_scale, + use_fp4=True, + weights_out_dtype=torch.bfloat16, + ) + assert weights_fused.dtype == torch.bfloat16 + assert torch.equal(weights_fused, weights_ref.to(torch.bfloat16)) diff --git a/tests/model_executor/layers/test_mla_short_prefill_indexer.py b/tests/model_executor/layers/test_mla_short_prefill_indexer.py index f8e7a89cca..41a06d0546 100644 --- a/tests/model_executor/layers/test_mla_short_prefill_indexer.py +++ b/tests/model_executor/layers/test_mla_short_prefill_indexer.py @@ -450,3 +450,547 @@ def test_candidate_kernels_preserve_packed_bounds_and_padding( keep = valid & (block >= 0) & ((cols - ks[:, None]) // block_size == block) reference = torch.where(keep, 3.0, -torch.inf) torch.testing.assert_close(logits, reference, rtol=0, atol=0) + + +# --- DeepGEMM sparse-MQA indexer path (V4.1 two-level candidate filtering) --- + + +def _make_candidates(lens: torch.Tensor, num_candidates: int, cbk: int): + """Production-faithful candidates: top min(avail, K) blocks per row — + full coverage for lens <= K*cbk, a fixed-size subset beyond that.""" + avail = ((lens.cpu() + cbk - 1) // cbk).tolist() + rows = [] + for a in avail: + n = min(a, num_candidates) + picks = torch.randperm(a)[:n].sort().values[:num_candidates].tolist() + rows.append(picks + [-1] * (num_candidates - len(picks))) + return torch.tensor(rows, dtype=torch.int32) + + +def _candidate_token_mask( + candidates: torch.Tensor, cbk: int, total: int, ks: torch.Tensor | None = None +): + """[rows, total] bool mask of the candidate-covered token positions. + + ``ks`` converts request-local candidate blocks to absolute positions + (prefill's packed workspace); pass None when already request-local. + """ + rows = candidates.shape[0] + device = candidates.device + # Sentinel block keeps -1 padding from scattering onto real columns. + width = max(total, (candidates.shape[1] + 1) * cbk) + if ks is not None: + width += int(ks.max()) + cbk + in_cand = torch.zeros(rows, width, dtype=torch.bool, device=device) + block_id = torch.where(candidates >= 0, candidates.long(), candidates.shape[1]) + tok = block_id.unsqueeze(2) * cbk + torch.arange(cbk, device=device) + if ks is not None: + tok = tok + ks.unsqueeze(1).unsqueeze(2) + keep = (candidates >= 0).unsqueeze(2).expand(-1, -1, cbk) + in_cand.scatter_(1, tok.flatten(1), keep.flatten(1)) + return in_cand[:, :total] + + +def _reference_topk(masked: torch.Tensor, topk: int) -> torch.Tensor: + """Plain ``torch.topk`` over -inf-masked logits; slots past a row's valid + count are -1 (the indexer's "no token" sentinel).""" + values, indices = masked.topk(topk, dim=1) + return torch.where(values > float("-inf"), indices.int(), -1) + + +def _assert_same_selection( + logits: torch.Tensor, + got: torch.Tensor, + expected: torch.Tensor, + row_starts: torch.Tensor | None = None, +) -> None: + """Both top-k index sets must pick the same multiset of logit values per + row (-1 padding must match too). Indices themselves may differ: the + top-k kernels break ties among equal values in nondeterministic order.""" + for r in range(logits.shape[0]): + base = int(row_starts[r]) if row_starts is not None else 0 + g, e = got[r][got[r] >= 0], expected[r][expected[r] >= 0] + assert g.numel() == e.numel(), f"row {r}: {g.numel()} vs {e.numel()} selected" + vg = logits[r][(base + g).long()].sort().values + ve = logits[r][(base + e).long()].sort().values + assert torch.equal(vg, ve), f"row {r}: selected values differ" + + +def _sparse_scratch(rows: int, num_sparse: int, topk: int) -> dict: + """Caller-owned buffers the sparse path writes into (the metadata builder + owns these in production).""" + i32 = dict(dtype=torch.int32, device="cuda") + return dict( + sparse_indices=torch.zeros(rows, num_sparse, **i32), + end=torch.zeros(rows, **i32), + col_indices=torch.zeros(rows, topk, **i32), + ) + + +def _expand_candidates_reference( + candidates: torch.Tensor, ks: torch.Tensor, ke: torch.Tensor, cbk: int, sbk: int +) -> tuple[torch.Tensor, torch.Tensor]: + """PyTorch reference for `candidate_blocks_to_sparse_indices`: expand + each candidate block into ``cbk // sbk`` sparse blocks, keep those inside + ``[0, ke - ks)``, sort ascending, anchor at ``ks // sbk``, pad by + repeating the last valid block, and count the valid sparse columns.""" + ratio = cbk // sbk + rows, num_candidates = candidates.shape + num_sparse = num_candidates * ratio + big = torch.iinfo(torch.int32).max + offsets = torch.arange(ratio, device=candidates.device, dtype=torch.int32) + blocks = (candidates.unsqueeze(-1) * ratio + offsets).flatten(1) + max_block = ((ke - ks).clamp(min=0) + sbk - 1) // sbk + valid = (blocks >= 0) & (blocks < max_block.unsqueeze(1)) + sorted_blocks, _ = torch.where(valid, blocks, big).sort(dim=1) + n_valid = (sorted_blocks != big).sum(dim=1, dtype=torch.int32) + base = (ks // sbk).unsqueeze(1) + kernel_blocks = torch.where(sorted_blocks != big, sorted_blocks, 0) + base + last_slot = (n_valid - 1).clamp(min=0).unsqueeze(1).long() + last_block = kernel_blocks.gather(1, last_slot) + pad = torch.arange(num_sparse, device=candidates.device) >= n_valid.unsqueeze(1) + sparse_indices = torch.where(pad, last_block, kernel_blocks) + last_start = last_block.squeeze(1) * sbk + ks % sbk + last_fill = (ke - last_start).clamp(min=1, max=sbk) + end = torch.where(n_valid > 0, (n_valid - 1) * sbk + last_fill, 0) + return sparse_indices.int(), end.int() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel") +def test_candidate_blocks_to_sparse_indices_math(): + """Candidate->sparse-block expansion: ks anchoring, range filtering, + repeat-last padding, and the per-row valid column count (``end``).""" + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + candidate_blocks_to_sparse_indices, + ) + + # Rows cover: exactly one block, an unaligned ks with a partial tail + # block, an empty row, duplicates / -1 padding / out-of-range garbage + # (warmup runs touch uninitialized buffers), and int32-overflow junk. + ks = torch.tensor([64, 13, 100, 8, 0], dtype=torch.int32, device="cuda") + ke = torch.tensor([72, 30, 100, 40, 1000], dtype=torch.int32, device="cuda") + candidates = torch.tensor( + [ + [0, -1, -1, -1], + [2, 0, -1, -1], + [3, -1, -1, -1], + [1, 1, 7, 0], + [2**30, -(2**30), 0, -1], + ], + dtype=torch.int32, + device="cuda", + ) + indices, end = candidate_blocks_to_sparse_indices( + candidates, ks, ke, candidate_block_size=8, sparse_block_kv=8 + ) + assert indices.tolist() == [ + [8, 8, 8, 8], + [1, 3, 3, 3], + [12, 12, 12, 12], + [1, 2, 2, 2], + [0, 0, 0, 0], + ] + # end: row 1's last block covers [29, 37) but ke = 30 -> one valid token; + # row 3 keeps the duplicated block (production candidates are unique). + assert end.tolist() == [8, 9, 0, 24, 8] + + # Wider candidate blocks expand into multiple sparse blocks. + idx2, end2 = candidate_blocks_to_sparse_indices( + torch.tensor([[1, 0]], dtype=torch.int32, device="cuda"), + torch.tensor([0], dtype=torch.int32, device="cuda"), + torch.tensor([48], dtype=torch.int32, device="cuda"), + candidate_block_size=16, + sparse_block_kv=8, + ) + assert idx2.tolist() == [[0, 1, 2, 3]] + assert end2.tolist() == [32] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton kernel") +@pytest.mark.parametrize("cbk,sbk", [(8, 8), (16, 8), (16, 16)]) +def test_candidate_blocks_to_sparse_indices_matches_reference(cbk: int, sbk: int): + """The Triton expansion equals the PyTorch reference on production-shaped + random inputs, including strided (column-sliced) candidate rows and + strided output buffers.""" + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + candidate_blocks_to_sparse_indices, + ) + + torch.manual_seed(0) + rows, num_candidates = 41, 2048 + ratio = cbk // sbk + lens = torch.randint(1, 40000, (rows,), dtype=torch.int32) + lens[0], lens[1] = 1, 3 + ks = torch.zeros(rows, dtype=torch.int32) + ks[1:] = lens.cumsum(0)[:-1] + ks, ke = ks.cuda(), (ks + lens).cuda() + candidates = _make_candidates(lens, num_candidates, cbk).cuda() + # Row 2 gets out-of-range garbage past the padding, as a CUDA-graph + # warm-up on uninitialized buffers would. + candidates[2, -3:] = torch.tensor([2**30, -(2**29), 10**6], dtype=torch.int32) + + ref_indices, ref_end = _expand_candidates_reference(candidates, ks, ke, cbk, sbk) + + cand_buf = torch.full((rows, num_candidates + 64), -7, dtype=torch.int32).cuda() + cand_buf[:, :num_candidates] = candidates + si_buf = torch.zeros(rows, num_candidates * ratio + 32, dtype=torch.int32).cuda() + end_buf = torch.zeros(rows * 2, dtype=torch.int32, device="cuda") + ks_buf = torch.zeros(rows * 2, dtype=torch.int32, device="cuda") + ks_buf[::2] = ks + got_indices, got_end = candidate_blocks_to_sparse_indices( + cand_buf[:, :num_candidates], + ks_buf[::2], + ke, + cbk, + sbk, + out=(si_buf[:, : num_candidates * ratio], end_buf[::2]), + ) + assert torch.equal(got_indices, ref_indices) + assert torch.equal(got_end, ref_end) + assert torch.equal( + cand_buf[:, num_candidates:], torch.full_like(cand_buf[:, num_candidates:], -7) + ) + assert not si_buf[:, num_candidates * ratio :].any() + assert not end_buf[1::2].any() + + +def _skip_unless_sm100_sparse_kernels(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10: + pytest.skip("SM100 required") + deep_gemm = pytest.importorskip("deep_gemm") + if not hasattr(deep_gemm, "fp8_fp4_sparse_mqa_logits"): + pytest.skip("DeepGEMM >= 2.8 required for sparse MQA logits") + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + has_deep_select, + ) + + if not has_deep_select(): + pytest.skip("DeepSelect extension (vllm._deepselect_C) required") + return deep_gemm + + +def _quant_fp4(x: torch.Tensor): + from deep_gemm.utils import per_token_cast_to_fp4 + + return per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + + +@pytest.mark.parametrize("capture", [False, True]) +def test_sparse_mqa_logits_prefill_matches_dense_masked(capture: bool): + """The sparse-MQA candidate path selects exactly the same logit values as + ``torch.topk`` over dense bf16 logits masked to the candidates (the two + DeepGEMM kernels agree bitwise on candidate columns), incl. under CUDA + graph capture, and a later layer reusing the step's expansion/schedule + selects the same values.""" + deep_gemm = _skip_unless_sm100_sparse_kernels() + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + sparse_mqa_logits_prefill_chunk, + ) + + torch.manual_seed(0) + rows, num_heads, head_dim = 37, 32, 128 + cbk, topk, num_candidates = 8, 512, 2048 # V4.1-Flash production values + lens = torch.randint(1, 30000, (rows,), dtype=torch.int32) + lens[0], lens[1], lens[2] = 1, 5, 2048 # single token, sub-block, long + ks = torch.zeros(rows, dtype=torch.int32) + ks[1:] = lens.cumsum(0)[:-1] + ks, ke, lens = ks.cuda(), (ks + lens).cuda(), lens.cuda() + total = int(lens.sum()) + q_fp, q_sf = _quant_fp4( + torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + ) + k_fp, k_sf = _quant_fp4( + torch.randn(total, head_dim, device="cuda", dtype=torch.bfloat16) + ) + weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16) + candidates = _make_candidates(lens, num_candidates, cbk).cuda() + + # Reference: dense bf16 logits masked to the candidates, then torch.topk. + dense = deep_gemm.fp8_fp4_mqa_logits( + ( + q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2), + q_sf.view(rows, num_heads), + ), + (k_fp.view(torch.int8), k_sf.view(total)), + weights, + ks, + ke, + clean_logits=False, + logits_dtype=torch.bfloat16, + ) + cols = torch.arange(total, device="cuda") + in_range = (cols - ks.unsqueeze(1) >= 0) & ( + cols - ks.unsqueeze(1) < lens.unsqueeze(1) + ) + masked = dense.masked_fill( + ~(in_range & _candidate_token_mask(candidates, cbk, total, ks)), + float("-inf"), + ).float() + ref_abs = _reference_topk(masked, topk) # packed-workspace positions + ref = torch.where(ref_abs >= 0, ref_abs - ks.unsqueeze(1), -1) + + out = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda") + scratch = _sparse_scratch(rows, num_candidates, topk) + + def run_sparse(kernel_metadata=None): + return sparse_mqa_logits_prefill_chunk( + q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2), + q_sf.view(rows, num_heads), + k_fp.view(torch.int8), + k_sf.view(total), + weights, + ks, + ke, + candidates, + cbk, + cbk, + topk, + out, + kernel_metadata=kernel_metadata, + **scratch, + ) + + if capture: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run_sparse() # warm up JIT + the DeepGEMM workspace pre-capture + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run_sparse() + out.fill_(-1) + graph.replay() + else: + kernel_metadata = run_sparse() + _assert_same_selection(dense, out, ref, row_starts=ks) + out.fill_(-1) + run_sparse(kernel_metadata) + + _assert_same_selection(dense, out, ref, row_starts=ks) + + +def _build_fp4_paged_cache(req_lens: torch.Tensor, page_kv: int, head_dim: int): + """MXFP4 paged cache in DeepGEMM's fused layout: per page, all tokens' + fp4 data first, then their packed UE8M0 scales.""" + device = req_lens.device + pages_per = (req_lens + page_kv - 1) // page_kv + num_pages, max_pages = int(pages_per.sum()), int(pages_per.max()) + kv_cache = torch.zeros( + num_pages, page_kv, 1, head_dim // 2 + 4, dtype=torch.uint8, device=device + ) + block_table = torch.zeros( + len(req_lens), max_pages, dtype=torch.int32, device=device + ) + page_pool = torch.randperm(num_pages, device=device) + offset = 0 + for b in range(len(req_lens)): + n = int(pages_per[b]) + block_table[b, :n] = page_pool[offset : offset + n].to(torch.int32) + offset += n + k_fp, k_sf = _quant_fp4( + torch.randn(int(req_lens[b]), head_dim, device=device, dtype=torch.bfloat16) + ) + rows68 = torch.zeros(n * page_kv, 68, dtype=torch.uint8, device=device) + rows68[: int(req_lens[b]), : head_dim // 2] = k_fp.view(torch.uint8) + rows68[: int(req_lens[b]), head_dim // 2 :] = k_sf.view(torch.uint8).view(-1, 4) + by_page = rows68.view(n, page_kv, 68) + split = torch.cat( + [by_page[..., :64].reshape(n, -1), by_page[..., 64:].reshape(n, -1)], dim=1 + ).view(n, page_kv, 1, head_dim // 2 + 4) + kv_cache[block_table[b, :n].long()] = split + return kv_cache, block_table + + +@pytest.mark.parametrize("next_n", [1, 2]) +@pytest.mark.parametrize("capture", [False, True]) +def test_sparse_mqa_logits_paged_matches_dense_masked(next_n: int, capture: bool): + """Paged sparse-MQA decode selects exactly the values ``torch.topk`` picks + from the dense bf16 masked logits, incl. MTP rows (next_n > 1 flattened to + one query row each), CUDA graph capture/replay and schedule reuse by a + later layer.""" + deep_gemm = _skip_unless_sm100_sparse_kernels() + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + sparse_mqa_logits_paged_decode, + ) + + torch.manual_seed(0) + batch, num_heads, head_dim, page_kv = 9, 32, 128, 128 + cbk, topk, num_candidates = 8, 512, 2048 + rows = batch * next_n + req_lens = torch.randint(1, 30000, (batch,), dtype=torch.int32, device="cuda") + req_lens[0], req_lens[1] = 1, 30000 # tiny row + a filtered (subset) row + # Row (b, j) sees L_b - next_n + j + 1 tokens (native MTP layout). + row_lens = ( + (req_lens.unsqueeze(1) - next_n + torch.arange(1, next_n + 1, device="cuda")) + .clamp(min=0) + .to(torch.int32) + ) + kv_cache, block_table = _build_fp4_paged_cache(req_lens, page_kv, head_dim) + assert kv_cache.stride(0) % 512 == 0 # the paged sparse kernel requires it + + q_fp, q_sf = _quant_fp4( + torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + ) + q = q_fp.view(torch.int8).view(batch, next_n, num_heads, head_dim // 2) + q_scale = q_sf.view(batch, next_n, num_heads) + weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16) + candidates = _make_candidates(row_lens.flatten(), num_candidates, cbk).cuda() + + def reference_topk(): + context_lens_2d = row_lens.view(batch, next_n) + sched = deep_gemm.get_paged_mqa_logits_metadata( + context_lens_2d, page_kv, deep_gemm.get_num_sms(), None + ) + logits = deep_gemm.fp8_fp4_paged_mqa_logits( + (q, q_scale), + kv_cache, + weights, + context_lens_2d, + block_table, + sched, + int(req_lens.max()), + False, + torch.bfloat16, + None, + ) + lens_flat = row_lens.flatten() + cols = torch.arange(logits.shape[1], device="cuda") + masked = logits.masked_fill( + ~( + (cols.unsqueeze(0) < lens_flat.unsqueeze(1)) + & _candidate_token_mask(candidates, cbk, logits.shape[1]) + ), + float("-inf"), + ).float() + return _reference_topk(masked, topk), logits + + expected, dense_logits = reference_topk() + out = torch.full((rows, topk), -1, dtype=torch.int32, device="cuda") + + # The metadata builder flattens every query to one row: per-row context + # lengths, an expanded block table and a row -> request map. + context_lens = row_lens.flatten().contiguous() + row_block_table = block_table.repeat_interleave(next_n, dim=0) + row_indices = torch.arange(batch, dtype=torch.int32, device="cuda") + row_indices = row_indices.repeat_interleave(next_n) + row_ks = torch.zeros(rows, dtype=torch.int32, device="cuda") + scratch = _sparse_scratch(rows, num_candidates, topk) + + def run_sparse(o: torch.Tensor, kernel_metadata=None): + return sparse_mqa_logits_paged_decode( + q.view(rows, 1, num_heads, head_dim // 2), + q_scale.view(rows, 1, num_heads), + kv_cache, + weights, + context_lens, + row_block_table, + row_indices, + candidates, + cbk, + cbk, + topk, + o, + row_ks=row_ks, + kernel_metadata=kernel_metadata, + **scratch, + ) + + if capture: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run_sparse(out) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run_sparse(out) + out.fill_(-1) + graph.replay() + else: + kernel_metadata = run_sparse(out) + _assert_same_selection(dense_logits, out, expected) + out.fill_(-1) + run_sparse(out, kernel_metadata) + + _assert_same_selection(dense_logits, out, expected) + + +def test_sparse_mqa_logits_accuracy_vs_torch(): + """Sparse-kernel numerics vs a PyTorch fp32 reference on the dequantized + MXFP4 inputs.""" + _skip_unless_sm100_sparse_kernels() + from deep_gemm.utils import cast_back_from_fp4 + + from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + candidate_blocks_to_sparse_indices, + ) + from vllm.utils.deep_gemm import ( + fp8_fp4_sparse_mqa_logits, + get_sparse_mqa_logits_metadata, + ) + + torch.manual_seed(0) + rows, num_heads, head_dim = 17, 32, 128 + sbk, num_candidates = 8, 2048 + lens = torch.randint(1, 1500, (rows,), dtype=torch.int32) + lens[0] = 1 + ks = torch.zeros(rows, dtype=torch.int32) + ks[1:] = lens.cumsum(0)[:-1] + ks, ke, lens = ks.cuda(), (ks + lens).cuda(), lens.cuda() + total = int(lens.sum()) + q_fp, q_sf = _quant_fp4( + torch.randn(rows * num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + ) + k_fp, k_sf = _quant_fp4( + torch.randn(total, head_dim, device="cuda", dtype=torch.bfloat16) + ) + weights = torch.randn(rows, num_heads, device="cuda", dtype=torch.bfloat16) + candidates = _make_candidates(lens, num_candidates, sbk).cuda() + + sparse_idx, end = candidate_blocks_to_sparse_indices(candidates, ks, ke, sbk, sbk) + metadata = get_sparse_mqa_logits_metadata( + ks, ke, total, sparse_idx, torch.int8, sbk + ) + logits = fp8_fp4_sparse_mqa_logits( + ( + q_fp.view(torch.int8).view(rows, num_heads, head_dim // 2), + q_sf.view(rows, num_heads), + ), + (k_fp.view(torch.int8), k_sf.view(total)), + weights, + metadata, + sparse_idx.shape[1], + sbk, + ) + + # PyTorch fp32 reference on the dequantized inputs. + q_deq = cast_back_from_fp4(q_fp, q_sf, gran_k=32, use_packed_ue8m0=True) + k_deq = cast_back_from_fp4(k_fp, k_sf, gran_k=32, use_packed_ue8m0=True) + dots = torch.einsum( + "rhd,td->rht", + q_deq.view(rows, num_heads, head_dim).float(), + k_deq.float(), + ) + ref = (dots.relu() * weights.float().unsqueeze(-1)).sum(1) # [rows, total] + + # Compare on valid sparse columns (absolute packed-workspace positions). + width = sparse_idx.shape[1] * sbk + cols = torch.arange(width, device="cuda") + valid = cols.unsqueeze(0) < end.unsqueeze(1).long() + pos = ( + sparse_idx.repeat_interleave(sbk, dim=1).long() * sbk + + (ks % sbk).unsqueeze(1) + + (cols % sbk).unsqueeze(0) + ).clamp(min=0, max=total - 1) + abs_err = (logits.float() - ref.gather(1, pos)).abs()[valid] + row_scale = ref.abs().amax(dim=1) + norm_err = ((logits.float() - ref.gather(1, pos)).abs() / row_scale.unsqueeze(1))[ + valid + ] + print( + f"sparse vs torch fp32: max_abs={abs_err.max().item():.4f}; " + f"rowmax-normalized p50={norm_err.median().item():.5f} " + f"p99={norm_err.quantile(0.99).item():.5f} max={norm_err.max().item():.5f}" + ) + # bf16-level numerics: far below the value gaps that change top-k selection. + assert norm_err.quantile(0.99).item() < 1e-2 + assert norm_err.max().item() < 2e-2 diff --git a/vllm/config/attention.py b/vllm/config/attention.py index 129fbcbb08..3646677728 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -86,6 +86,15 @@ class AttentionConfig: indexer). Quantized formats (fp8, mxfp4, nvfp4) require indexer kernel support in the backend.""" + indexer_sparse_logits: bool = False + """DeepSeek V4.1 two-level indexer: score only the candidate blocks with + DeepGEMM's sparse MQA-logits kernels instead of computing dense logits over + the whole context and masking them. Requires `indexer_kv_dtype="mxfp4"`, + an SM100-class GPU, DeepGEMM >= 2.8 and the DeepSelect top-k extension + (the top-k runs on the kernels' bf16 logits). The sparse path costs + O(candidate blocks) per query regardless of context length, so it pays off + for long contexts (roughly 32K tokens and beyond) and is slower below.""" + hisparse_config: HiSparseConfig | None = None """HiSparse host-resident KV configuration. Setting this enables experimental Model Runner V2-only HiSparse sparse-MLA decode hot-buffering. It is inferred diff --git a/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py new file mode 100644 index 0000000000..839c581583 --- /dev/null +++ b/vllm/model_executor/kernels/attention/dsa/sparse_mqa_logits.py @@ -0,0 +1,469 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepGEMM sparse MQA-logits path for the DeepSeek V4.1 two-level indexer. + +The candidate-source indexer publishes the top candidate blocks (in units of +``candidate_block_size`` compressed positions). Instead of scoring every KV +position and masking the dense logits down to those blocks, this path scores +only the candidate tokens with DeepGEMM's sparse MQA logits kernels +(``fp8_fp4_(paged_)sparse_mqa_logits``, DeepGEMM >= 2.8, SM100 only), runs +DeepSelect's top-k directly on the bf16 logits they produce, and remaps the +selected columns back to request-local positions. + +Only the MXFP4 indexer cache is supported: the kernels require UE8M0-packed +(granularity-32) Q/KV scales, which is exactly the MXFP4 cache layout. The +FP8 indexer cache stores one fp32 scale per quant block instead. + +The functions here take caller-owned output buffers (``sparse_indices``, +``end``, ``col_indices``) so the hot path allocates nothing; the indexer +metadata builder owns them. +""" + +import importlib.util + +import torch + +from vllm.model_executor.layers.indexer_topk import ( + deep_select_topk, + get_deep_select_stride_requirement, +) +from vllm.platforms import current_platform +from vllm.triton_utils import tl, triton +from vllm.utils.deep_gemm import ( + fp8_fp4_paged_sparse_mqa_logits, + fp8_fp4_sparse_mqa_logits, + get_paged_sparse_mqa_logits_metadata, + get_sparse_mqa_logits_metadata, +) + +# Sparse-block sizes (in tokens) supported by the DeepGEMM kernels. +SPARSE_BLOCK_KV_CHOICES = (16, 8) + + +def has_deep_select() -> bool: + """Whether the DeepSelect top-k extension is built and this GPU runs it.""" + return ( + importlib.util.find_spec("vllm._deepselect_C") is not None + and current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ) + + +def check_deep_select_layout(num_sparse_cols: int, topk_tokens: int) -> None: + """Raise unless DeepSelect accepts the sparse logits layout. + + The sparse kernels return contiguous bf16 ``[rows, num_sparse_cols]`` + logits and the top-k writes int32 ``[rows, topk_tokens]`` indices; + DeepSelect needs both row strides aligned (1024B / 32B on SM100). + """ + in_align, out_align = get_deep_select_stride_requirement() + logits_row_bytes = num_sparse_cols * torch.bfloat16.itemsize + if num_sparse_cols >= 2**23 or logits_row_bytes % in_align != 0: + raise ValueError( + f"DeepSelect needs the sparse logits row ({num_sparse_cols} bf16 " + f"columns, {logits_row_bytes} B) to be {in_align}B-aligned." + ) + if topk_tokens > 4096 or (topk_tokens * torch.int32.itemsize) % out_align != 0: + raise ValueError( + f"DeepSelect needs topk <= 4096 with a {out_align}B-aligned int32 " + f"output row, got topk={topk_tokens}." + ) + + +_INT32_BIG = tl.constexpr(2**31 - 1) + + +@triton.jit +def _expand_candidates_kernel( + cand_ptr, + cand_stride, + ks_ptr, + ks_stride, + ke_ptr, + ke_stride, + si_ptr, + si_stride, + end_ptr, + end_stride, + K: tl.constexpr, + S: tl.constexpr, + RATIO: tl.constexpr, + SBK: tl.constexpr, +): + """One program per row: expand + filter + sort candidate blocks. + + Each row's valid sparse blocks come out sorted ascending (no dedup pass: + production candidates are unique, and the kernel tolerates repeats), + padded by repeating the last valid block — DeepGEMM's padding convention. + ``end`` counts valid sparse-token columns; valid columns form a prefix + because block ids are sorted. + """ + row = tl.program_id(0) + ks = tl.load(ks_ptr + row * ks_stride) + ke = tl.load(ke_ptr + row * ke_stride) + cols = tl.arange(0, K) + cand = tl.load(cand_ptr + row * cand_stride + cols) # K is pow2, no mask + if RATIO == 2: + cand = tl.interleave(cand * RATIO, cand * RATIO + 1) + + # Keep only blocks inside the row's KV range [0, ke - ks); the lower + # bound is free since candidates are non-negative. + kv_len = tl.maximum(ke - ks, 0) + max_block = (kv_len + SBK - 1) // SBK + valid = (cand >= 0) & (cand < max_block) + vals = tl.where(valid, cand, _INT32_BIG) + vals = tl.sort(vals) + + # Anchor at the row's ks rounded down to a sparse-block boundary + # (identity for the paged path, where ks == 0). + base = ks // SBK + si = tl.where(vals != _INT32_BIG, vals, 0) + base + n_valid = tl.sum(valid.to(tl.int32)) + last_rel = tl.max(tl.where(valid, cand, -1)) + pad_val = tl.where(n_valid > 0, last_rel + base, base) + tl.store( + si_ptr + row * si_stride + tl.arange(0, S), + tl.where(tl.arange(0, S) < n_valid, si, pad_val), + ) + + # Valid sparse-token column count; only the last block can be partial. + # Absolute start of the last block: (base + last_rel) * SBK + ks % SBK, + # which is exactly last_rel * SBK + ks. + last_start = last_rel * SBK + ks + fill = tl.minimum(tl.maximum(ke - last_start, 1), SBK) + end = tl.where(n_valid > 0, (n_valid - 1) * SBK + fill, 0) + tl.store(end_ptr + row * end_stride, end) + + +@triton.jit +def _sparse_topk_remap_kernel( + col_ptr, + col_stride, + si_ptr, + si_stride, + ks_ptr, + ks_stride, + out_ptr, + out_stride, + k, + SBK: tl.constexpr, + K_POW2: tl.constexpr, +): + """Remap sparse top-k columns to request-local positions. + + Column ``j * SBK + o`` scores the token at + ``(si[row, j] - ks // SBK) * SBK + o``; -1 columns stay -1. + """ + row = tl.program_id(0) + ks_base = tl.load(ks_ptr + row * ks_stride) // SBK + cols = tl.arange(0, K_POW2) + c = tl.load(col_ptr + row * col_stride + cols, mask=cols < k, other=-1) + valid = c >= 0 + cc = tl.where(valid, c, 0) + blocks = tl.load(si_ptr + row * si_stride + cc // SBK) + pos = (blocks - ks_base) * SBK + cc % SBK + tl.store(out_ptr + row * out_stride + cols, tl.where(valid, pos, -1), mask=cols < k) + + +def pick_sparse_block_kv(candidate_block_size: int) -> int: + """Largest sparse-block size (tokens) dividing the candidate block size.""" + for sbk in SPARSE_BLOCK_KV_CHOICES: + if candidate_block_size % sbk == 0: + return sbk + raise ValueError( + f"candidate_block_size={candidate_block_size} is not a multiple of " + f"any supported sparse_block_kv {SPARSE_BLOCK_KV_CHOICES}." + ) + + +def candidate_blocks_to_sparse_indices( + candidate_blocks: torch.Tensor, + row_ks: torch.Tensor, + row_ke: torch.Tensor, + candidate_block_size: int, + sparse_block_kv: int, + out: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Expand request-local candidate blocks into DeepGEMM sparse block ids. + + Args: + candidate_blocks: [rows, K] int32 request-local candidate block ids + (-1 padded), in units of ``candidate_block_size`` positions. K + must be a power of two (the in-kernel sort); rows may be strided. + row_ks/row_ke: [rows] int32 per-row K range; bounds are in the same + (packed-workspace) coordinates the sparse kernel iterates over. + Pass zeros for the paged path, whose blocks are context-relative. + candidate_block_size: Positions per candidate block. + sparse_block_kv: Positions per sparse block (8 or 16). + out: Optional ``(sparse_indices, end)`` buffers to write into, + shaped [rows, K * ratio] and [rows], both int32. + + Returns: + sparse_indices: [rows, K * ratio] int32 block ids for DeepGEMM, + sorted ascending per row with the valid prefix first (unique in + production, where candidates are unique) and the remaining slots + repeating the last valid block (the kernel's padding convention). + With the unaligned-ks variant, block ``i`` starts at + ``i * sparse_block_kv + ks % sparse_block_kv``. + end: [rows] int32 count of valid sparse-token columns; valid columns + always form a prefix because the block ids are sorted. + """ + rows, num_candidates = candidate_blocks.shape + ratio = candidate_block_size // sparse_block_kv + assert ratio * sparse_block_kv == candidate_block_size and ratio in (1, 2), ( + candidate_block_size, + sparse_block_kv, + ) + assert num_candidates & (num_candidates - 1) == 0, num_candidates + num_sparse = num_candidates * ratio + if out is None: + device = candidate_blocks.device + out = ( + torch.empty(rows, num_sparse, dtype=torch.int32, device=device), + torch.empty(rows, dtype=torch.int32, device=device), + ) + sparse_indices, end = out + assert sparse_indices.shape == (rows, num_sparse) and end.shape == (rows,) + assert candidate_blocks.stride(1) == 1 and sparse_indices.stride(1) == 1 + for t in (candidate_blocks, row_ks, row_ke, sparse_indices, end): + assert t.dtype == torch.int32, t.dtype + + _expand_candidates_kernel[(rows,)]( + candidate_blocks, + candidate_blocks.stride(0), + row_ks, + row_ks.stride(0), + row_ke, + row_ke.stride(0), + sparse_indices, + sparse_indices.stride(0), + end, + end.stride(0), + K=num_candidates, + S=num_sparse, + RATIO=ratio, + SBK=sparse_block_kv, + num_warps=16, + ) + return sparse_indices, end + + +def sparse_topk_remap( + logits: torch.Tensor, + sparse_indices: torch.Tensor, + end: torch.Tensor, + row_ks: torch.Tensor, + sparse_block_kv: int, + topk_tokens: int, + topk_indices: torch.Tensor, + *, + col_indices: torch.Tensor, +) -> None: + """Row top-k over sparse logits, remapped to request-local positions. + + Column ``j * sparse_block_kv + o`` of ``logits`` scores the token at + request-local position + ``(sparse_indices[row, j] - ks // sparse_block_kv) * sparse_block_kv + o``. + + DeepSelect consumes the bf16 logits as they come out of the sparse + kernels; the valid columns form a prefix (block ids are sorted), so the + per-row ``end`` bounds drive it directly. Slots beyond a row's valid + count come back as -1. + + Args: + col_indices: [rows, topk_tokens] int32 scratch for the sparse-column + top-k result (row stride 32B-aligned for DeepSelect). + """ + rows, width = logits.shape + assert logits.dtype == torch.bfloat16 and width >= topk_tokens + assert end.dtype == torch.int32 and end.shape == (rows,) + cols = col_indices[:rows] + deep_select_topk(logits, topk_tokens, end=end, output_idx=cols) + _sparse_topk_remap_kernel[(rows,)]( + cols, + cols.stride(0), + sparse_indices, + sparse_indices.stride(0), + row_ks, + row_ks.stride(0), + topk_indices, + topk_indices.stride(0), + topk_tokens, + SBK=sparse_block_kv, + K_POW2=triton.next_power_of_2(topk_tokens), + ) + + +def sparse_mqa_logits_prefill_chunk( + q: torch.Tensor, + q_scale: torch.Tensor, + k_quant: torch.Tensor, + k_scale: torch.Tensor, + weights: torch.Tensor, + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + candidate_blocks: torch.Tensor, + candidate_block_size: int, + sparse_block_kv: int, + topk_tokens: int, + topk_indices: torch.Tensor, + *, + sparse_indices: torch.Tensor, + end: torch.Tensor, + col_indices: torch.Tensor, + kernel_metadata: torch.Tensor | None = None, +) -> torch.Tensor: + """Sparse-MQA logits + top-k for one prefill chunk (packed KV workspace). + + Args: + q: [rows, H, D] packed Q values (MXFP4 viewed as int8). + q_scale: [rows, H] int32 packed UE8M0 Q scales. + k_quant: [total_kv, D] packed K workspace (MXFP4 viewed as int8). + k_scale: [total_kv] int32 packed UE8M0 K scales. + weights: [rows, H] bf16 per-head weights; the sparse kernels take + bf16 and do not fold the Q scale in. + cu_seqlen_ks/cu_seqlen_ke: [rows] int32 per-token K bounds in the + packed workspace. + candidate_blocks: [rows, K] int32 request-local candidate block ids. + topk_indices: [rows, topk_tokens] output buffer. + sparse_indices/end/col_indices: Caller-owned scratch, see + `candidate_blocks_to_sparse_indices` and `sparse_topk_remap`. + kernel_metadata: DeepGEMM schedule from a previous call with the same + candidates and bounds (i.e. another indexer layer in the same + step). When given, the candidate expansion is skipped and + ``sparse_indices``/``end`` are assumed to be up to date. + + Returns: + The DeepGEMM schedule metadata used, for reuse by later layers. + """ + assert weights.dtype == torch.bfloat16, weights.dtype + if kernel_metadata is None: + candidate_blocks_to_sparse_indices( + candidate_blocks, + cu_seqlen_ks, + cu_seqlen_ke, + candidate_block_size, + sparse_block_kv, + out=(sparse_indices, end), + ) + kernel_metadata = get_sparse_mqa_logits_metadata( + cu_seqlen_ks, + cu_seqlen_ke, + k_quant.shape[0], + sparse_indices, + q.dtype, + sparse_block_kv, + ) + logits = fp8_fp4_sparse_mqa_logits( + (q, q_scale), + (k_quant, k_scale), + weights, + kernel_metadata, + sparse_indices.shape[1], + sparse_block_kv, + ) + sparse_topk_remap( + logits, + sparse_indices, + end, + cu_seqlen_ks, + sparse_block_kv, + topk_tokens, + topk_indices, + col_indices=col_indices, + ) + return kernel_metadata + + +def sparse_mqa_logits_paged_decode( + q: torch.Tensor, + q_scale: torch.Tensor, + kv_cache: torch.Tensor, + weights: torch.Tensor, + context_lens: torch.Tensor, + block_table: torch.Tensor, + row_indices: torch.Tensor, + candidate_blocks: torch.Tensor, + candidate_block_size: int, + sparse_block_kv: int, + topk_tokens: int, + topk_indices: torch.Tensor, + *, + row_ks: torch.Tensor, + sparse_indices: torch.Tensor, + end: torch.Tensor, + col_indices: torch.Tensor, + kernel_metadata: torch.Tensor | None = None, +) -> torch.Tensor: + """Sparse-MQA logits + top-k for paged decode rows. + + Every query is one row (spec-decode rows are flattened by the metadata + builder), so ``block_table`` is already per row. + + Args: + q: [rows, 1, H, D] packed Q values (MXFP4 viewed as int8). + q_scale: [rows, 1, H] int32 packed UE8M0 Q scales. + kv_cache: 4D [num_pages, page_kv, 1, head_bytes] uint8 fused cache + view; the page stride must be 512B-aligned. + weights: [rows, H] bf16 per-head weights. + context_lens: [rows] int32 per-row (compressed) context lengths. + block_table: [rows, P] int32 per-row page table, ``stride(-1) == 1``. + row_indices: [rows] int32 row -> request map (pairing only affects + scheduling). + candidate_blocks: [rows, K] int32 candidate block ids. + topk_indices: [rows, topk_tokens] output buffer. + row_ks: [rows] int32 zeros (paged blocks are context-relative). + sparse_indices/end/col_indices: Caller-owned scratch. + kernel_metadata: See `sparse_mqa_logits_prefill_chunk`. + + Returns: + The DeepGEMM schedule metadata used, for reuse by later layers. + """ + rows = q.shape[0] + page_kv = kv_cache.shape[1] + assert weights.dtype == torch.bfloat16, weights.dtype + assert kv_cache.stride(0) % 512 == 0, ( + "paged cache page stride must be 512B-aligned: " + f"shape={tuple(kv_cache.shape)} stride={kv_cache.stride()}" + ) + assert page_kv % sparse_block_kv == 0, (page_kv, sparse_block_kv) + assert block_table.shape[0] == rows and block_table.stride(-1) == 1 + assert context_lens.shape == (rows,) and context_lens.is_contiguous() + + if kernel_metadata is None: + candidate_blocks_to_sparse_indices( + candidate_blocks, + row_ks, + context_lens, + candidate_block_size, + sparse_block_kv, + out=(sparse_indices, end), + ) + kernel_metadata = get_paged_sparse_mqa_logits_metadata( + context_lens, + block_table, + row_indices, + page_kv, + sparse_indices, + q.dtype, + sparse_block_kv, + ) + logits = fp8_fp4_paged_sparse_mqa_logits( + (q, q_scale), + kv_cache, + weights, + kernel_metadata, + sparse_indices.shape[1], + sparse_block_kv, + ) + sparse_topk_remap( + logits, + sparse_indices, + end, + row_ks, + sparse_block_kv, + topk_tokens, + topk_indices, + col_indices=col_indices, + ) + return kernel_metadata diff --git a/vllm/model_executor/layers/sparse_mqa_indexer.py b/vllm/model_executor/layers/sparse_mqa_indexer.py new file mode 100644 index 0000000000..78db1b3df5 --- /dev/null +++ b/vllm/model_executor/layers/sparse_mqa_indexer.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sparse Attention Indexer that scores only the candidate blocks. + +DeepSeek V4.1 two-level selection: the candidate-source indexer publishes +the top candidate blocks and later indexers pick their top-k inside them. +`SparseAttnIndexer` does that by computing dense logits over the whole +context and masking; this layer instead calls DeepGEMM's sparse MQA-logits +kernels on the candidate blocks only, so the work is O(candidate blocks) +instead of O(context). It requires the `DeepseekV41SparseIndexerBackend` +metadata (see ``AttentionConfig.indexer_sparse_logits``). +""" + +import torch +from torch import nn + +import vllm.envs as envs +from vllm import _custom_ops as ops +from vllm.forward_context import get_forward_context +from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + has_deep_select, + sparse_mqa_logits_paged_decode, + sparse_mqa_logits_prefill_chunk, +) +from vllm.model_executor.layers.sparse_attn_indexer import ( + _gather_workspace_shapes, + kv_cache_as_quant_view, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import has_deep_gemm_sparse_mqa +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerPrefillChunkMetadata, +) +from vllm.v1.attention.backends.mla.sparse_indexer import ( + DeepseekV41SparseIndexerMetadata, +) +from vllm.v1.worker.workspace import current_workspace_manager + + +def _prefill_k_workspaces( + total_seq_lens: int, head_dim: int +) -> tuple[torch.Tensor, torch.Tensor]: + """The packed MXFP4 K-gather workspace ``(values, scales)``, shared with + the dense indexer layers of the same model.""" + values_spec, scales_spec = _gather_workspace_shapes( + total_seq_lens, head_dim, current_platform.fp8_dtype(), use_fp4_cache=True + ) + k_quant, k_scale = current_workspace_manager().get_simultaneous( + values_spec, scales_spec + ) + return k_quant, k_scale + + +def _gather_prefill_chunk_k( + kv_cache: torch.Tensor, + k_quant_full: torch.Tensor, + k_scale_full: torch.Tensor, + chunk: DeepseekV32IndexerPrefillChunkMetadata, +) -> tuple[torch.Tensor, torch.Tensor]: + """Gather one prefill chunk's paged K into the packed workspace.""" + assert chunk.local_cu_seq_lens is not None + k_quant = k_quant_full[: chunk.max_local_total_seq_lens] + k_scale = k_scale_full[: chunk.max_local_total_seq_lens] + if not chunk.skip_kv_gather and chunk.local_total_seq_lens > 0: + ops.cp_gather_indexer_k_quant_cache( + kv_cache, + k_quant, + k_scale, + chunk.block_table, + chunk.local_cu_seq_lens, + ) + return k_quant, k_scale + + +class SparseMQAIndexer(nn.Module): + """Candidate-consuming indexer on DeepGEMM's sparse MQA-logits kernels. + + Only valid for indexer layers that read candidate blocks with the MXFP4 + indexer cache on SM100. The K cache is written by the model before this + runs; ``forward`` takes the same arguments as `SparseAttnIndexer` so the + attention layer can call either. + """ + + weights_dtype = torch.bfloat16 + """Per-head weights dtype the sparse kernels take. The fused Q RoPE-quant + kernel writes it directly so no cast runs per step.""" + + def __init__( + self, + k_cache, + topk_tokens: int, + head_dim: int, + max_total_seq_len: int, + topk_indices_buffer: torch.Tensor, + candidate_blocks: torch.Tensor, + candidate_block_size: int, + ): + super().__init__() + if not ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + and has_deep_gemm_sparse_mqa() + ): + raise ValueError( + "SparseMQAIndexer requires an SM100-class GPU and DeepGEMM >= 2.8." + ) + if not has_deep_select(): + raise ValueError( + "SparseMQAIndexer requires the DeepSelect top-k extension " + "(vllm._deepselect_C)." + ) + self.k_cache = k_cache + self.topk_tokens = topk_tokens + self.head_dim = head_dim + self.max_total_seq_len = max_total_seq_len + self.topk_indices_buffer = topk_indices_buffer + self.candidate_blocks = candidate_blocks + self.candidate_block_size = candidate_block_size + + def _reserve_workspaces(self, device: torch.device) -> None: + """Profiling run: claim the K-gather workspace and the peak sparse + logits allocation so the memory estimate covers them.""" + _prefill_k_workspaces(self.max_total_seq_len, self.head_dim) + max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 + torch.empty(max_logits_bytes, dtype=torch.uint8, device=device) + + def forward( + self, + hidden_states: torch.Tensor, + q_quant: tuple[torch.Tensor, torch.Tensor], + k: torch.Tensor | None, + weights: torch.Tensor, + ) -> torch.Tensor: + attn_metadata = get_forward_context().attn_metadata + if not isinstance(attn_metadata, dict): + self._reserve_workspaces(hidden_states.device) + return self.topk_indices_buffer + + metadata = attn_metadata[self.k_cache.prefix] + assert isinstance(metadata, DeepseekV41SparseIndexerMetadata), ( + "SparseMQAIndexer needs DeepseekV41SparseIndexerBackend metadata" + ) + assert k is None, "the model writes the indexer K cache" + q_values, q_scale = q_quant + kv_cache = self.k_cache.kv_cache + topk_tokens = self.topk_tokens + topk_indices_buffer = self.topk_indices_buffer + topk_indices_buffer[: hidden_states.shape[0]] = -1 + sparse_block_kv = metadata.sparse_block_kv + + if metadata.num_prefills > 0: + prefill = metadata.prefill + sparse_chunks = metadata.sparse_prefill + assert prefill is not None and sparse_chunks is not None + k_quant_full, k_scale_full = _prefill_k_workspaces( + self.max_total_seq_len, self.head_dim + ) + for chunk, rows in zip(prefill.chunks, sparse_chunks): + k_quant, k_scale = _gather_prefill_chunk_k( + kv_cache, k_quant_full, k_scale_full, chunk + ) + if chunk.local_total_seq_lens == 0: + continue # buffer already holds -1 + start, end = chunk.token_start, chunk.token_end + rows.kernel_metadata = sparse_mqa_logits_prefill_chunk( + q_values[start:end].view(torch.int8), + q_scale[start:end], + k_quant.view(torch.int8), + k_scale.view(torch.int32).squeeze(-1), + weights[start:end], + chunk.cu_seqlen_ks, + chunk.cu_seqlen_ke, + self.candidate_blocks[start:end], + self.candidate_block_size, + sparse_block_kv, + topk_tokens, + topk_indices_buffer[start:end, :topk_tokens], + sparse_indices=rows.sparse_indices, + end=rows.end, + col_indices=rows.col_indices, + kernel_metadata=rows.kernel_metadata, + ) + + if metadata.num_decodes > 0: + decode_rows = metadata.sparse_decode + assert decode_rows is not None and decode_rows.block_table is not None + assert decode_rows.row_indices is not None + num_rows = decode_rows.row_ke.shape[0] + if num_rows > 0: + decode_rows.kernel_metadata = sparse_mqa_logits_paged_decode( + q_values[:num_rows].view(torch.int8).unsqueeze(1), + q_scale[:num_rows].unsqueeze(1), + kv_cache_as_quant_view(kv_cache, self.head_dim, use_fp4_cache=True), + weights[:num_rows], + decode_rows.row_ke, + decode_rows.block_table, + decode_rows.row_indices, + self.candidate_blocks[:num_rows], + self.candidate_block_size, + sparse_block_kv, + topk_tokens, + topk_indices_buffer[:num_rows, :topk_tokens], + row_ks=decode_rows.row_ks, + sparse_indices=decode_rows.sparse_indices, + end=decode_rows.end, + col_indices=decode_rows.col_indices, + kernel_metadata=decode_rows.kernel_metadata, + ) + return topk_indices_buffer diff --git a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py index cf6e7235f9..669186edbb 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py +++ b/vllm/models/deepseek_v4/common/ops/fused_indexer_q.py @@ -309,12 +309,22 @@ class FusedIndexerQRopeQuantTritonKernel( ) +def _indexer_weights_out_dtypes(vllm_config: Any) -> tuple[torch.dtype, ...]: + """Weights dtypes the model's indexer layers ask for: fp32 for the dense + scoring kernels, plus bf16 when the DeepSeek V4.1 sparse-logits indexer + (`SparseMQAIndexer`) is enabled.""" + if vllm_config.attention_config.indexer_sparse_logits: + return (torch.float32, torch.bfloat16) + return (torch.float32,) + + class FusedIndexerQRopeMxFp4TritonKernel( VllmTritonJitKernel["FusedIndexerQRopeMxFp4TritonKernel.CompileKey"] ): @dataclass(frozen=True) class CompileKey: dtype: torch.dtype + weights_out_dtype: torch.dtype num_heads: int index_q_half_rot_dim: int index_q_head_dim: int @@ -434,12 +444,14 @@ class FusedIndexerQRopeMxFp4TritonKernel( self, *, dtype: torch.dtype, + weights_out_dtype: torch.dtype, num_heads: int, head_dim: int, rope_dim: int, ) -> CompileKey: return self.CompileKey( dtype=dtype, + weights_out_dtype=weights_out_dtype, num_heads=num_heads, index_q_half_rot_dim=rope_dim // 2, index_q_head_dim=head_dim, @@ -456,6 +468,7 @@ class FusedIndexerQRopeMxFp4TritonKernel( return self._trace_dispatch(self.dispatch)( dtype=vllm_config.model_config.dtype, + weights_out_dtype=_indexer_weights_out_dtypes(vllm_config), num_heads=num_heads, head_dim=head_dim, rope_dim=rope_dim, @@ -491,7 +504,7 @@ class FusedIndexerQRopeMxFp4TritonKernel( shape=(1, compile_key.num_heads, scale_head_dim), ), index_weights_out=TritonWarmupTensor( - torch.float32, + compile_key.weights_out_dtype, shape=(1, compile_key.num_heads), ), ) @@ -540,12 +553,17 @@ def fused_indexer_q_rope_quant( index_weights_softmax_scale: float, index_weights_head_scale: float, use_fp4: bool = False, + weights_out_dtype: torch.dtype = torch.float32, ) -> tuple[ torch.Tensor | tuple[torch.Tensor, torch.Tensor], torch.Tensor, ]: """Fused RoPE + quantize Q for the sparse indexer. + ``weights_out_dtype`` is the dtype the downstream scoring kernel takes: + fp32 for the dense MQA-logits kernels, bf16 for DeepGEMM's sparse + MQA-logits kernels (CUDA MXFP4 path only). + Weight-fold semantics (important — the two paths differ): FP8 path (use_fp4=False, default): @@ -577,7 +595,7 @@ def fused_indexer_q_rope_quant( num_index_q_heads = index_q.shape[1] index_q_head_dim = index_q.shape[2] - index_weights_out = torch.empty_like(index_weights, dtype=torch.float32) + index_weights_out = torch.empty_like(index_weights, dtype=weights_out_dtype) if use_fp4: assert index_q_head_dim % MXFP4_BLOCK_SIZE == 0, ( @@ -613,6 +631,7 @@ def fused_indexer_q_rope_quant( weights_out=index_weights_out, ) elif current_platform.is_xpu(): + assert weights_out_dtype == torch.float32, weights_out_dtype torch.ops.vllm.xpu_deepseek_fused_indexer_q_rope_mxfp4( index_q, positions, @@ -647,6 +666,7 @@ def fused_indexer_q_rope_quant( index_q_scale.view(torch.int32).squeeze(-1), ), index_weights_out + assert weights_out_dtype == torch.float32, weights_out_dtype fp8_dtype = current_platform.fp8_dtype() use_fnuz = fp8_dtype == torch.float8_e4m3fnuz fp8_max = 224.0 if use_fnuz else 448.0 diff --git a/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py b/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py index 8052178a03..f67a77e254 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py +++ b/vllm/models/deepseek_v4/nvidia/ops/fused_indexer_q_cutedsl.py @@ -165,6 +165,7 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"] rope_dim: int num_heads: int cos_sin_dtype: type[cutlass.Numeric] + weights_out_dtype: type[cutlass.Numeric] coarsen: int @staticmethod @@ -177,6 +178,7 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"] # process multiple heads at the same time to armotize RoPE load costs coarsen = compile_key.coarsen tb_size = IndexerQMxFp4Kernel.tb_size + weights_out_dtype = compile_key.weights_out_dtype assert num_heads % coarsen == 0 # later we will use 32B load = 16 BF16 elems @@ -287,7 +289,7 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"] weight_head_id = global_tid % num_heads weights_out[weight_token_id, weight_head_id] = ( weights[weight_token_id, weight_head_id].to(Float32) * scale - ) + ).to(weights_out_dtype) @cute.jit def host_entrypoint( @@ -320,11 +322,13 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"] self, *, cos_sin_dtype: type[cutlass.Numeric], + weights_out_dtype: type[cutlass.Numeric], **compile_key_fields: int, ) -> CompileKey: return self.CompileKey( **compile_key_fields, cos_sin_dtype=cos_sin_dtype, + weights_out_dtype=weights_out_dtype, ) def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: @@ -340,11 +344,17 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"] if num_heads <= 0 or head_dim <= 0 or rope_dim <= 0: return [] + # The DeepSeek V4.1 sparse-logits indexer takes bf16 weights; the + # dense scoring kernels take fp32. + weights_out_dtypes: tuple[type[cutlass.Numeric], ...] = (Float32,) + if vllm_config.attention_config.indexer_sparse_logits: + weights_out_dtypes = (Float32, BFloat16) return self._trace_dispatch(self.dispatch)( head_dim=head_dim, rope_dim=rope_dim, num_heads=num_heads, cos_sin_dtype=Float32, + weights_out_dtype=weights_out_dtypes, coarsen=(1, 4), ) @@ -380,7 +390,9 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"] divisibility=4, ) weights_out = make_fake_tensor( - Float32, (num_tokens, compile_key.num_heads), divisibility=4 + compile_key.weights_out_dtype, + (num_tokens, compile_key.num_heads), + divisibility=4, ) return ( positions, @@ -413,6 +425,7 @@ class IndexerQMxFp4Kernel(VllmCuTeDSLJitKernel["IndexerQMxFp4Kernel.CompileKey"] rope_dim=cos_sin_cache.shape[-1], num_heads=num_heads, cos_sin_dtype=torch_to_cute_dtype(cos_sin_cache.dtype), + weights_out_dtype=torch_to_cute_dtype(weights_out.dtype), coarsen=1 if num_tokens < 512 else 4, ) launch_args = ( diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index fa26351681..134328dbe9 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -4,6 +4,7 @@ DeepseekV4 MLA Attention Layer """ +import math from abc import ABC, abstractmethod from collections.abc import Callable from functools import cached_property @@ -25,6 +26,7 @@ from vllm.model_executor.layers.linear import ( RowParallelLinear, ) from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.model_executor.layers.sparse_mqa_indexer import SparseMQAIndexer from vllm.models.common.ops import fused_q_kv_rmsnorm from vllm.models.deepseek_v41.common.ops import ( MXFP4_BLOCK_SIZE, @@ -63,6 +65,9 @@ from vllm.v1.attention.backends.mla.indexer import ( dsa_indexer_uses_fp4, get_max_prefill_buffer_size, ) +from vllm.v1.attention.backends.mla.sparse_indexer import ( + DeepseekV41SparseIndexerBackend, +) from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache from vllm.v1.kv_cache_interface import ( KVCacheSpec, @@ -997,7 +1002,9 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): self.cache_config = cache_config self.dtype = dtype self.compress_ratio = compress_ratio - compilation_config = get_current_vllm_config().compilation_config + vllm_config = get_current_vllm_config() + self.sparse_logits = vllm_config.attention_config.indexer_sparse_logits + compilation_config = vllm_config.compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self @@ -1010,6 +1017,9 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): # head_dim already carries the fp8 scale padding # tokens_per_state=1 for V3.2, >1 for DeepseekV4; same cache layout. uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" + page_alignment = ( + 576 if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record() else 512 + ) return MLAAttentionSpec( block_size=self.cache_config.block_size, num_kv_heads=1, @@ -1020,16 +1030,21 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): # (the block stride is the sum of every page in the group), so it # has to carry the same padding: the packed record's TMA stride, # else 512B for FlashInfer sparse (#44577). - alignment=( - 576 - if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record() - else 512 + alignment=page_alignment, + # DeepGEMM's paged sparse MQA logits address pages by the block + # stride, which packs every layer's page in this layout, and need + # it 512B-aligned; the main KV record read from the same blocks + # needs its own TMA stride, so keep both. + block_stride_alignment=( + math.lcm(512, page_alignment) if self.sparse_logits else None ), ) def forward(self): ... def get_attn_backend(self) -> type[AttentionBackend]: + if self.sparse_logits: + return DeepseekV41SparseIndexerBackend return DeepseekV41IndexerBackend @@ -1129,21 +1144,52 @@ class DeepseekV4Indexer(nn.Module): self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps) self.k_cache = k_cache - self.indexer_op = SparseAttnIndexer( - self.k_cache, - self.quant_block_size, - self.scale_fmt, - self.topk_tokens, - self.head_dim, - self.max_model_len, - self.max_total_seq_len, - self.topk_indices_buffer, - skip_k_cache_insert=True, - use_fp4_cache=self.use_fp4_kv, - compress_ratio=self.compress_ratio, - candidate_blocks=candidate_block_buffer, - candidate_block_size=candidate_block_size, - candidate_write=candidate_write, + # Candidate consumers can score only the candidate blocks (opt-in); + # the candidate source and pre-candidate indexers stay dense. + use_sparse_logits = ( + vllm_config.attention_config.indexer_sparse_logits + and candidate_block_buffer is not None + and not candidate_write + ) + self.indexer_op: SparseAttnIndexer | SparseMQAIndexer + if use_sparse_logits: + if not self.use_fp4_kv: + raise ValueError( + "attention_config.indexer_sparse_logits requires the MXFP4 " + "indexer cache (attention_config.indexer_kv_dtype='mxfp4')." + ) + assert candidate_block_buffer is not None + assert topk_indices_buffer is not None + self.indexer_op = SparseMQAIndexer( + self.k_cache, + self.topk_tokens, + self.head_dim, + self.max_total_seq_len, + topk_indices_buffer, + candidate_block_buffer, + candidate_block_size, + ) + else: + self.indexer_op = SparseAttnIndexer( + self.k_cache, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + skip_k_cache_insert=True, + use_fp4_cache=self.use_fp4_kv, + compress_ratio=self.compress_ratio, + candidate_blocks=candidate_block_buffer, + candidate_block_size=candidate_block_size, + candidate_write=candidate_write, + ) + # The fused Q kernel writes the per-head weights in the dtype the + # scoring kernels take, so no cast runs per step. + self.indexer_weights_dtype = ( + SparseMQAIndexer.weights_dtype if use_sparse_logits else torch.float32 ) def _produce_k( @@ -1231,6 +1277,7 @@ class DeepseekV4Indexer(nn.Module): self.softmax_scale, self.n_head**-0.5, use_fp4=self.use_fp4_kv, + weights_out_dtype=self.indexer_weights_dtype, ) if isinstance(q_quant, tuple): q, q_scale = q_quant diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 5d05c29a84..91e315afb9 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -158,6 +158,10 @@ _grouped_fp4_impl: Callable[..., Any] | None = None _fp8_fp4_mqa_logits_impl: Callable[..., Any] | None = None _fp8_fp4_paged_mqa_logits_impl: Callable[..., Any] | None = None _get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None +_fp8_fp4_sparse_mqa_logits_impl: Callable[..., Any] | None = None +_fp8_fp4_paged_sparse_mqa_logits_impl: Callable[..., Any] | None = None +_get_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None +_get_paged_sparse_mqa_logits_metadata_impl: Callable[..., Any] | None = None _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None _mega_mhc_impl: Callable[..., Any] | None = None _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None @@ -232,6 +236,9 @@ def _lazy_init() -> None: global _grouped_impl, _grouped_masked_impl, _grouped_fp4_impl global _fp8_fp4_mqa_logits_impl, _fp8_fp4_paged_mqa_logits_impl global _get_paged_mqa_logits_metadata_impl + global _fp8_fp4_sparse_mqa_logits_impl, _fp8_fp4_paged_sparse_mqa_logits_impl + global _get_sparse_mqa_logits_metadata_impl + global _get_paged_sparse_mqa_logits_metadata_impl global _tf32_hc_prenorm_gemm_impl, _mega_mhc_impl global _get_mn_major_tma_aligned_tensor_impl global _get_mk_alignment_for_contiguous_layout_impl @@ -295,6 +302,17 @@ def _lazy_init() -> None: _get_paged_mqa_logits_metadata_impl = getattr( _dg, "get_paged_mqa_logits_metadata", None ) + # Sparse-indexer kernels (DeepGEMM >= 2.8, SM100 only). + _fp8_fp4_sparse_mqa_logits_impl = getattr(_dg, "fp8_fp4_sparse_mqa_logits", None) + _fp8_fp4_paged_sparse_mqa_logits_impl = getattr( + _dg, "fp8_fp4_paged_sparse_mqa_logits", None + ) + _get_sparse_mqa_logits_metadata_impl = getattr( + _dg, "get_sparse_mqa_logits_metadata", None + ) + _get_paged_sparse_mqa_logits_metadata_impl = getattr( + _dg, "get_paged_sparse_mqa_logits_metadata", None + ) _tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None) _mega_mhc_impl = getattr(_dg, "mega_mhc", None) _get_mn_major_tma_aligned_tensor_impl = getattr( @@ -681,6 +699,171 @@ def fp8_fp4_paged_mqa_logits( ) +def has_deep_gemm_sparse_mqa() -> bool: + """Whether the installed DeepGEMM provides the sparse-indexer kernels + (``fp8_fp4_(paged_)sparse_mqa_logits``, added in DeepGEMM 2.8, SM100-only).""" + _lazy_init() + return ( + _fp8_fp4_sparse_mqa_logits_impl is not None + and _fp8_fp4_paged_sparse_mqa_logits_impl is not None + and _get_sparse_mqa_logits_metadata_impl is not None + and _get_paged_sparse_mqa_logits_metadata_impl is not None + ) + + +def get_sparse_mqa_logits_metadata( + cu_seqlen_ks: torch.Tensor, + cu_seqlen_ke: torch.Tensor, + num_kv_tokens: int, + sparse_kv_block_indices: torch.Tensor, + qk_dtype: torch.dtype, + sparse_block_kv: int, +) -> torch.Tensor: + """Build scheduling metadata for `fp8_fp4_sparse_mqa_logits`. + + Always uses the unaligned-ks variant, which degenerates to the aligned + variant when ``cu_seqlen_ks % sparse_block_kv == 0``, so callers never + need a host-side sync to pick one. + + Args: + cu_seqlen_ks/cu_seqlen_ke: Per-row K range bounds in the packed KV + workspace, shape [num_q_tokens], dtype int32. + num_kv_tokens: Total KV tokens in the packed workspace. + sparse_kv_block_indices: Per-row candidate block ids, shape + [num_q_tokens, num_max_sparse_blocks], dtype int32. Each row's + valid prefix must be sorted ascending (repeats are tolerated but + wasteful); pad by repeating the last valid block. Block ``i`` + covers tokens + ``[i * sparse_block_kv + ks % sparse_block_kv, ...)``. + qk_dtype: dtype of the packed Q values (``torch.int8`` for MXFP4, + ``torch.float8_e4m3fn`` for FP8). + sparse_block_kv: Tokens per sparse block, 8 or 16. + """ + _lazy_init() + if _get_sparse_mqa_logits_metadata_impl is None: + return _missing() + return _get_sparse_mqa_logits_metadata_impl( + cu_seqlen_ks, + cu_seqlen_ke, + num_kv_tokens, + sparse_kv_block_indices, + qk_dtype, + sparse_block_kv, + use_unaligned_ks=True, + ) + + +def get_paged_sparse_mqa_logits_metadata( + context_lens: torch.Tensor, + block_table: torch.Tensor, + indices: torch.Tensor, + page_kv: int, + sparse_kv_block_indices: torch.Tensor, + qk_dtype: torch.dtype, + sparse_block_kv: int, +) -> torch.Tensor: + """Build scheduling metadata for `fp8_fp4_paged_sparse_mqa_logits`. + + Args: + context_lens: Per-row context length, shape [num_q_tokens], int32. + block_table: Per-row page table, shape [num_q_tokens, max_pages], + int32; rows of queries paired with the same request must be + identical (expand it per query row before calling). + indices: Request index for each query row, shape [num_q_tokens]. + page_kv: Tokens per KV page; must be a multiple of sparse_block_kv. + sparse_kv_block_indices: Per-row logical (context-relative) candidate + block ids, [num_q_tokens, num_max_sparse_blocks], int32. + qk_dtype: dtype of the packed Q values. + sparse_block_kv: Tokens per sparse block, 8 or 16. + """ + _lazy_init() + if _get_paged_sparse_mqa_logits_metadata_impl is None: + return _missing() + return _get_paged_sparse_mqa_logits_metadata_impl( + context_lens, + block_table, + indices, + page_kv, + sparse_kv_block_indices, + qk_dtype, + sparse_block_kv, + ) + + +def fp8_fp4_sparse_mqa_logits( + q: tuple[torch.Tensor, torch.Tensor], + kv: tuple[torch.Tensor, torch.Tensor], + weights: torch.Tensor, + metadata: torch.Tensor, + num_max_sparse_blocks: int, + sparse_block_kv: int, +) -> torch.Tensor: + """Compute MQA logits only at the candidate blocks (prefill path). + + Args: + q: ``(q_values, q_scale)``. q_values is [M, H, D] (packed FP4 viewed + as int8, or FP8); q_scale is the packed UE8M0 scale tensor [M, H] + int32 (mandatory, unlike the dense kernels). + kv: ``(kv_values, kv_scale)`` — the packed KV workspace with UE8M0 + int32 scales. + weights: [M, H] ``torch.bfloat16`` (the Q scale is NOT folded in). + metadata: From `get_sparse_mqa_logits_metadata`. + num_max_sparse_blocks: Width of each sparse-index row. + sparse_block_kv: Tokens per sparse block, 8 or 16. + + Returns: + bf16 logits of shape [M, num_max_sparse_blocks * sparse_block_kv]; + column ``j * sparse_block_kv + o`` scores the token at + ``sparse_kv_block_indices[row, j] * sparse_block_kv + ks % sbk + o``. + """ + _lazy_init() + if _fp8_fp4_sparse_mqa_logits_impl is None: + return _missing() + return _fp8_fp4_sparse_mqa_logits_impl( + q, + kv, + weights, + metadata, + num_max_sparse_blocks, + sparse_block_kv, + use_unaligned_ks=True, + ) + + +def fp8_fp4_paged_sparse_mqa_logits( + q: tuple[torch.Tensor, torch.Tensor], + kv_cache: torch.Tensor, + weights: torch.Tensor, + metadata: torch.Tensor, + num_max_sparse_blocks: int, + sparse_block_kv: int, +) -> torch.Tensor: + """Compute MQA logits only at the candidate blocks (paged decode path). + + Args: + q: ``(q_values, q_scale)``; q_values is [num_q_tokens, 1, H, D]. + kv_cache: Fused paged cache [num_pages, page_kv, 1, D' + 4] uint8, + page stride 512B-aligned. + weights: [num_q_tokens, H] ``torch.bfloat16``. + metadata: From `get_paged_sparse_mqa_logits_metadata`. + + Returns: + bf16 logits of shape + [num_q_tokens, num_max_sparse_blocks * sparse_block_kv]. + """ + _lazy_init() + if _fp8_fp4_paged_sparse_mqa_logits_impl is None: + return _missing() + return _fp8_fp4_paged_sparse_mqa_logits_impl( + q, + kv_cache, + weights, + metadata, + num_max_sparse_blocks, + sparse_block_kv, + ) + + def tf32_hc_prenorm_gemm( x: torch.Tensor, fn: torch.Tensor, diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 837dbca817..53206afd93 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -1116,6 +1116,11 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): for request_slice, query_slice in chunk_specs ] + def _prefill_split_seq_lens(self, seq_lens_cpu: torch.Tensor) -> torch.Tensor: + """Per-request KV lengths the prefill chunker budgets logits with; + subclasses whose logits rows are wider than the context override.""" + return seq_lens_cpu + @staticmethod def _split_indexer_prefill_chunks( compressed_seq_lens_cpu: torch.Tensor, @@ -1278,7 +1283,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): ) else: chunk_specs = self._split_indexer_prefill_chunks( - compressed_seq_lens_cpu[num_decodes:], + self._prefill_split_seq_lens(compressed_seq_lens_cpu[num_decodes:]), prefill_query_lens_cpu, self.max_prefill_buffer_size, max_logits_bytes, diff --git a/vllm/v1/attention/backends/mla/sparse_indexer.py b/vllm/v1/attention/backends/mla/sparse_indexer.py new file mode 100644 index 0000000000..c0b445534d --- /dev/null +++ b/vllm/v1/attention/backends/mla/sparse_indexer.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepSeek V4.1 indexer backend for the DeepGEMM sparse MQA-logits path. + +Candidate-consuming indexer layers score only the candidate blocks published +by the candidate-source layer (``fp8_fp4_(paged_)sparse_mqa_logits``) instead +of dense logits over the whole context. This backend extends the dense +indexer metadata with the per-row scratch that path needs, so the layer op +allocates nothing per step, and lets all consumer layers of a step share one +candidate expansion and DeepGEMM schedule (they all read the same candidate +blocks). Opt in with ``AttentionConfig.indexer_sparse_logits``. +""" + +from dataclasses import dataclass, fields + +import torch + +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( + check_deep_select_layout, + has_deep_select, + pick_sparse_block_kv, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import has_deep_gemm_sparse_mqa +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadata, + DeepseekV32IndexerMetadataBuilder, + DeepseekV41IndexerBackend, + dsa_indexer_uses_fp4, +) +from vllm.v1.kv_cache_interface import AttentionSpec + +logger = init_logger(__name__) + + +@dataclass +class SparseMQARowsMetadata: + """Per-row state for one sparse-logits call: the decode rows, or one + prefill chunk. Buffers are builder-owned views sized to the rows. + + ``kernel_metadata`` (the DeepGEMM schedule) is filled by the first + consumer indexer layer of the step together with ``sparse_indices`` / + ``end``; later consumer layers reuse all three. + """ + + row_ks: torch.Tensor + """[rows] int32 K-range start per row; zeros for paged decode.""" + row_ke: torch.Tensor + """[rows] int32 K-range end (prefill) or compressed context length.""" + sparse_indices: torch.Tensor + """[rows, num_sparse_blocks] int32 DeepGEMM sparse block ids.""" + end: torch.Tensor + """[rows] int32 valid sparse-column count per row.""" + col_indices: torch.Tensor + """[rows, topk] int32 scratch for the sparse-column top-k.""" + block_table: torch.Tensor | None = None + """Decode only: [rows, pages] per-row page table.""" + row_indices: torch.Tensor | None = None + """Decode only: [rows] row -> request index.""" + kernel_metadata: torch.Tensor | None = None + + +@dataclass +class DeepseekV41SparseIndexerMetadata(DeepseekV32IndexerMetadata): + sparse_block_kv: int = 0 + sparse_decode: SparseMQARowsMetadata | None = None + sparse_prefill: list[SparseMQARowsMetadata] | None = None + """Parallel to ``prefill.chunks``.""" + + +class DeepseekV41SparseIndexerBackend(DeepseekV41IndexerBackend): + @staticmethod + def get_name() -> str: + return "DEEPSEEK_V41_SPARSE_INDEXER" + + @staticmethod + def get_builder_cls() -> type["DeepseekV41SparseIndexerMetadataBuilder"]: + return DeepseekV41SparseIndexerMetadataBuilder + + +class DeepseekV41SparseIndexerMetadataBuilder(DeepseekV32IndexerMetadataBuilder): + """Dense indexer metadata plus the sparse-logits row scratch. + + Every requirement of the sparse kernels is checked here, at engine start, + rather than falling back to the dense path per step. + """ + + def __init__( + self, + kv_cache_spec: AttentionSpec, + layer_names: list[str], + vllm_config: VllmConfig, + device: torch.device, + **kwargs, + ) -> None: + super().__init__(kv_cache_spec, layer_names, vllm_config, device, **kwargs) + hf_config = vllm_config.model_config.hf_text_config + candidate_block_size = getattr(hf_config, "candidate_block_size", 0) + candidate_topk_blocks = getattr(hf_config, "candidate_topk_blocks", 0) + topk_tokens = getattr(hf_config, "index_topk", 0) + prefix = "attention_config.indexer_sparse_logits" + if candidate_block_size <= 0 or candidate_topk_blocks <= 0: + raise ValueError( + f"{prefix} requires a two-level (candidate block) indexer model " + "such as DeepSeek-V4.1." + ) + if not dsa_indexer_uses_fp4(vllm_config): + raise ValueError( + f"{prefix} requires the MXFP4 indexer cache " + "(attention_config.indexer_kv_dtype='mxfp4'): the sparse kernels " + "take UE8M0-packed granularity-32 scales." + ) + if not ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + ): + raise ValueError(f"{prefix} requires an SM100-class GPU.") + if not has_deep_gemm_sparse_mqa(): + raise ValueError( + f"{prefix} requires DeepGEMM >= 2.8 " + "(fp8_fp4_sparse_mqa_logits not found)." + ) + if self.dcp_world_size > 1 or self.use_pcp: + raise NotImplementedError( + f"{prefix} is not supported with decode/prefill context parallel." + ) + if not self.supports_varlen: + raise ValueError( + f"{prefix} requires the varlen paged MQA logits decode layout." + ) + # The top-k runs on the kernels' bf16 logits with DeepSelect; there is + # no fp32 fallback, so the other top-k backends cannot be honored. + topk_backend = vllm_config.kernel_config.sparse_indexer_topk_backend + if not has_deep_select(): + raise ValueError( + f"{prefix} requires the DeepSelect top-k extension " + "(vllm._deepselect_C) for this GPU." + ) + if topk_backend not in ("auto", "deep_select"): + raise ValueError( + f"{prefix} uses DeepSelect for the top-k; " + f"kernel_config.sparse_indexer_topk_backend='{topk_backend}' " + "is not supported with it." + ) + self.sparse_block_kv = pick_sparse_block_kv(candidate_block_size) + if kv_cache_spec.block_size % self.sparse_block_kv != 0: + raise ValueError( + f"{prefix}: indexer cache pages ({kv_cache_spec.block_size} " + f"tokens) must be a multiple of {self.sparse_block_kv} tokens." + ) + self.num_sparse_blocks = candidate_topk_blocks * ( + candidate_block_size // self.sparse_block_kv + ) + num_sparse_cols = self.num_sparse_blocks * self.sparse_block_kv + check_deep_select_layout(num_sparse_cols, topk_tokens) + # Sparse logits are bf16 [rows, num_sparse_cols]; make the prefill + # chunker (which budgets fp32 [rows, seq_len]) size chunks as if every + # row were this wide. + self.min_split_seq_len = (num_sparse_cols * 2 + 3) // 4 + + # Rows are indexed by batch token (decode rows first, then prefill + # chunks by token range), like ``arange_buffer``. + max_rows = self.arange_buffer.shape[0] + int32 = dict(dtype=torch.int32, device=device) + self.sparse_indices_buffer = torch.zeros( + (max_rows, self.num_sparse_blocks), **int32 + ) + self.end_buffer = torch.zeros(max_rows, **int32) + self.col_indices_buffer = torch.zeros((max_rows, topk_tokens), **int32) + self.zero_row_ks = torch.zeros(max_rows, **int32) + logger.info_once( + "DeepSeek V4.1 indexer: candidate consumers score only the candidate " + "blocks with DeepGEMM sparse MQA logits (%d sparse blocks of %d " + "tokens per row).", + self.num_sparse_blocks, + self.sparse_block_kv, + ) + + def _prefill_split_seq_lens(self, seq_lens_cpu: torch.Tensor) -> torch.Tensor: + return seq_lens_cpu.clamp(min=self.min_split_seq_len) + + def _rows(self, start: int, end: int, row_ks, row_ke, **kwargs): + return SparseMQARowsMetadata( + row_ks=row_ks, + row_ke=row_ke, + sparse_indices=self.sparse_indices_buffer[start:end], + end=self.end_buffer[start:end], + col_indices=self.col_indices_buffer[start:end], + **kwargs, + ) + + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + ) -> DeepseekV41SparseIndexerMetadata: + base = super().build(common_prefix_len, common_attn_metadata, fast_build) + metadata = DeepseekV41SparseIndexerMetadata( + **{f.name: getattr(base, f.name) for f in fields(base)}, + sparse_block_kv=self.sparse_block_kv, + ) + if base.decode is not None: + decode = base.decode + rows = decode.seq_lens.shape[0] + # The varlen layout flattens every query to one row, so the base + # decode tensors are already per row. + assert decode.seq_lens.shape[1] == 1 and not decode.requires_padding + assert decode.block_table.shape[0] == rows + metadata.sparse_decode = self._rows( + 0, + rows, + row_ks=self.zero_row_ks[:rows], + row_ke=decode.seq_lens.view(-1), + block_table=decode.block_table, + row_indices=( + decode.indices + if decode.indices is not None + else self.arange_buffer[:rows] + ), + ) + if base.prefill is not None: + metadata.sparse_prefill = [ + self._rows( + chunk.token_start, + chunk.token_end, + row_ks=chunk.cu_seqlen_ks, + row_ke=chunk.cu_seqlen_ke, + ) + for chunk in base.prefill.chunks + ] + return metadata diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 28478b69f5..1c6e102c16 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1608,6 +1608,15 @@ def _get_kv_cache_bytes_per_block( ] if hot_page_sizes: bytes_per_block = round_up(bytes_per_block, math.lcm(*hot_page_sizes)) + stride_alignments = [ + spec.block_stride_alignment + for group in kv_cache_groups + for layer_name in group.layer_names + if isinstance(spec := _get_per_layer_spec(group, layer_name), MLAAttentionSpec) + and spec.block_stride_alignment + ] + if stride_alignments: + bytes_per_block = round_up(bytes_per_block, math.lcm(*stride_alignments)) return bytes_per_block diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 8cd867eb9c..7e91300b2f 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -653,6 +653,12 @@ class MLAAttentionSpec(FullAttentionSpec): is_index_group_leader: bool = False storage_block_size: int | None = None """Token width used to view storage when it differs from the kernel block.""" + block_stride_alignment: int | None = None + """Required alignment, in bytes, of the distance between consecutive + blocks of this cache. In block-major layouts that distance is the whole + block (all layers' pages), so the allocator rounds the block up to it. + DeepGEMM's paged sparse MQA-logits kernels address pages as + ``base + page * stride`` and need it 512B-aligned.""" # Group capability enabled when any member flattens a non-causal query block # into decode rows. Runtime metadata still selects causal vs. non-causal mode. non_causal_multi_token_decode: bool = False @@ -674,6 +680,7 @@ class MLAAttentionSpec(FullAttentionSpec): cache_role_set = {spec.cache_role for spec in specs} index_group_leader_set = {spec.is_index_group_leader for spec in specs} storage_block_size_set = set(spec.storage_block_size for spec in specs) + block_stride_alignment_set = {spec.block_stride_alignment for spec in specs} assert ( len(cache_dtype_str_set) == 1 and len(tokens_per_state_set) == 1 @@ -681,10 +688,11 @@ class MLAAttentionSpec(FullAttentionSpec): and len(cache_role_set) == 1 and len(index_group_leader_set) == 1 and len(storage_block_size_set) == 1 + and len(block_stride_alignment_set) == 1 ), ( "All attention layers in the same KV cache group must use the same " "quantization method, tokens per state, model version, cache role, " - "index-sharing role, and storage block size." + "index-sharing role, storage block size and block stride alignment." ) merged_spec = cls( block_size=specs[0].block_size, @@ -701,6 +709,7 @@ class MLAAttentionSpec(FullAttentionSpec): cache_role=cache_role_set.pop(), is_index_group_leader=index_group_leader_set.pop(), storage_block_size=storage_block_size_set.pop(), + block_stride_alignment=block_stride_alignment_set.pop(), non_causal_multi_token_decode=any( spec.non_causal_multi_token_decode for spec in specs ),