[DSA] Wire DeepGEMM sparse MQA logits into the DeepSeek V4.1 indexer (#56254)
Signed-off-by: Jared Wen <w13431838023@gmail.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Jared Wen
co-authored by
Claude Fable 5.1
Claude Fable 5.1
parent
f2aad6aa70
commit
1257512305
@@ -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))
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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 = (
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
),
|
||||
|
||||
Reference in New Issue
Block a user