[ROCm][DSv4.1] Paged MXFP4 sparse indexer on aiter's MQA-logits kernel (#58671)
Signed-off-by: Mehmet Cagri Kaymak <mehmet.kaymak@amd.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Shanshan Shen <467638484@qq.com>
This commit is contained in:
Mehmet Cagri
co-authored by
Claude Opus 5.5
Shanshan Shen
Claude Opus 5.5
Shanshan Shen
parent
e05095c9e1
commit
63b9931b73
@@ -0,0 +1,599 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""DeepSeek V4.1's ROCm MXFP4 sparse indexer against a torch reference.
|
||||
|
||||
The indexer K cache is a view into vLLM's block-major KV pool: its pages are
|
||||
a block stride apart with other layers' pages in between, so the tests put
|
||||
it there. The reference reads the keys back from a second cache the same
|
||||
writer fills in natural order, so it shares none of the preshuffled
|
||||
addressing it checks.
|
||||
"""
|
||||
|
||||
import itertools
|
||||
import types
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.is_rocm():
|
||||
pytest.skip("ROCm-only", allow_module_level=True)
|
||||
|
||||
from vllm.config import CUDAGraphMode
|
||||
from vllm.models.deepseek_v41.amd.rocm import (
|
||||
rocm_mxfp4_indexer_k_store,
|
||||
rocm_mxfp4_indexer_q_quant,
|
||||
)
|
||||
from vllm.platforms.rocm import get_cdna_version
|
||||
from vllm.utils.math_utils import cdiv
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
from vllm.v1.attention.backends.mla.indexer import (
|
||||
DeepseekV32IndexerPrefillChunkMetadata,
|
||||
)
|
||||
from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import (
|
||||
RocmMxfp4IndexerMetadata,
|
||||
native_decode,
|
||||
plan_gather_launches,
|
||||
plan_prefill_chunks,
|
||||
)
|
||||
from vllm.v1.attention.ops import rocm_paged_mxfp4_indexer as ops
|
||||
from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager
|
||||
|
||||
if get_cdna_version() != 4:
|
||||
pytest.skip("CDNA4-only", allow_module_level=True)
|
||||
|
||||
HEADS, HEAD_DIM = 32, 128
|
||||
WIDTH = HEAD_DIM // 2 + HEAD_DIM // 32
|
||||
# V4.1's 8-token blocks, with the pool (2048 blocks) and top-k (512) scaled
|
||||
# down so short test contexts still select.
|
||||
CAND_BLOCK, POOL_BLOCKS, TOPK = 8, 16, 32
|
||||
MAX_LEN = 4096
|
||||
DEVICE = "cuda"
|
||||
|
||||
_E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0])
|
||||
E2M1 = torch.cat([_E2M1, -_E2M1])
|
||||
|
||||
|
||||
def _dequant(packed, e8m0):
|
||||
"""[..., D/2] e2m1 pairs and [..., D/32] e8m0 to [..., D] fp32."""
|
||||
table = E2M1.to(packed.device)
|
||||
vals = torch.stack([table[(packed & 0xF).long()], table[(packed >> 4).long()]], -1)
|
||||
scale = torch.pow(2.0, e8m0.float() - 127)
|
||||
return (vals.flatten(-2).unflatten(-1, (-1, 32)) * scale[..., None]).flatten(-2)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _workspace():
|
||||
init_workspace_manager(torch.device(DEVICE))
|
||||
yield
|
||||
reset_workspace_manager()
|
||||
|
||||
|
||||
class _Case:
|
||||
"""Requests with their keys written into a block-major pool (ratio-2 and
|
||||
ratio-1 pages side by side in every block) and into natural-order caches
|
||||
the reference reads. The bf16 inputs are kept for the end-to-end test."""
|
||||
|
||||
def __init__(self, seq_lens, block):
|
||||
from aiter.ops.triton.fusions.k_norm_rope_mxfp4_cache import (
|
||||
k_norm_rope_mxfp4_cache,
|
||||
)
|
||||
|
||||
self.seq_lens = seq_lens
|
||||
num_blocks = sum(cdiv(n, block) for n in seq_lens) + 3
|
||||
perm = torch.randperm(num_blocks)
|
||||
width = max(cdiv(n, block) for n in seq_lens)
|
||||
self.block_table = torch.zeros(len(seq_lens), width, dtype=torch.int32)
|
||||
used = 0
|
||||
for req, n in enumerate(seq_lens):
|
||||
self.block_table[req, : cdiv(n, block)] = perm[used : used + cdiv(n, block)]
|
||||
used += cdiv(n, block)
|
||||
self.block_table = self.block_table.to(DEVICE)
|
||||
|
||||
def align(n):
|
||||
return cdiv(n, 576) * 576
|
||||
|
||||
page2, page1 = block // 2 * WIDTH, block * WIDTH
|
||||
stride = align(page2) + align(page1) + 576
|
||||
self.pool = torch.zeros(num_blocks * stride, dtype=torch.uint8, device=DEVICE)
|
||||
self.page_bytes = torch.zeros_like(self.pool, dtype=torch.bool)
|
||||
for i in range(num_blocks):
|
||||
base = i * stride
|
||||
self.page_bytes[base : base + page2] = True
|
||||
base += align(page2)
|
||||
self.page_bytes[base : base + page1] = True
|
||||
self.cache = {
|
||||
2: torch.as_strided(
|
||||
self.pool, (num_blocks, block // 2, WIDTH), (stride, WIDTH, 1), 0
|
||||
),
|
||||
1: torch.as_strided(
|
||||
self.pool,
|
||||
(num_blocks, block, WIDTH),
|
||||
(stride, WIDTH, 1),
|
||||
align(page2),
|
||||
),
|
||||
}
|
||||
self.offsets = [0]
|
||||
for n in seq_lens:
|
||||
self.offsets.append(self.offsets[-1] + n)
|
||||
# the same keys in natural order, one entry per page
|
||||
self.natural = {
|
||||
r: torch.zeros(self.offsets[-1], 1, WIDTH, dtype=torch.uint8, device=DEVICE)
|
||||
for r in (1, 2)
|
||||
}
|
||||
self.cos_sin = cos_sin = torch.randn(MAX_LEN, 64, device=DEVICE)
|
||||
self.norm = norm = torch.rand(HEAD_DIM, device=DEVICE) + 0.5
|
||||
self.k_pre = {}
|
||||
for req, n in enumerate(seq_lens):
|
||||
if n == 0:
|
||||
continue
|
||||
k_pre = torch.randn(n, HEAD_DIM, device=DEVICE).to(torch.bfloat16)
|
||||
self.k_pre[req] = k_pre
|
||||
pos = torch.arange(n, device=DEVICE)
|
||||
for ratio, cache in self.cache.items():
|
||||
comp = pos // ratio
|
||||
boundary = (pos + 1) % ratio == 0
|
||||
entries = cache.shape[1]
|
||||
page = self.block_table[req, comp // entries].long()
|
||||
slot = torch.where(boundary, page * entries + comp % entries, -1)
|
||||
nat = torch.where(boundary, self.offsets[req] + comp, -1)
|
||||
rocm_mxfp4_indexer_k_store(
|
||||
k_pre,
|
||||
pos,
|
||||
cos_sin,
|
||||
norm,
|
||||
1e-6,
|
||||
cache,
|
||||
slot,
|
||||
ratio,
|
||||
True,
|
||||
num_heads=HEADS,
|
||||
)
|
||||
k_norm_rope_mxfp4_cache(
|
||||
k_pre, pos, cos_sin, norm, 1e-6, self.natural[ratio], nat, ratio
|
||||
)
|
||||
|
||||
def keys(self, ratio, req, n):
|
||||
nat = self.natural[ratio][self.offsets[req] : self.offsets[req] + n, 0]
|
||||
return _dequant(nat[:, : HEAD_DIM // 2], nat[:, HEAD_DIM // 2 :])
|
||||
|
||||
|
||||
def _queries(num_rows):
|
||||
q = torch.randint(0, 256, (num_rows, HEADS, HEAD_DIM // 2), dtype=torch.uint8)
|
||||
scales = torch.randint(122, 130, (num_rows, HEADS, HEAD_DIM // 32))
|
||||
q, scales = q.to(DEVICE), scales.to(torch.uint8).to(DEVICE)
|
||||
weights = torch.randn(num_rows, HEADS, device=DEVICE)
|
||||
# the fused Q kernel returns the scales as one int32 per head
|
||||
return q, scales.view(torch.int32).squeeze(-1), weights, _dequant(q, scales)
|
||||
|
||||
|
||||
def _scores(q, w, keys):
|
||||
return (torch.relu(q @ keys.T) * w[:, None]).sum(0)
|
||||
|
||||
|
||||
def _assert_topk(row, scores, k):
|
||||
"""``row`` is a top-k of the entries of ``scores`` above -inf, up to the
|
||||
fp32 reduce order."""
|
||||
valid = scores > float("-inf")
|
||||
want = min(k, int(valid.sum()))
|
||||
sel = row[row >= 0].long()
|
||||
assert sel.numel() == want and sel.unique().numel() == want
|
||||
if want:
|
||||
assert bool(valid[sel].all())
|
||||
kth = scores[valid].topk(want).values[-1]
|
||||
finite = scores[torch.isfinite(scores)]
|
||||
tol = 1e-4 * float(finite.abs().max()) if finite.numel() else 0.0
|
||||
assert bool((scores[sel] >= kth - tol).all())
|
||||
|
||||
|
||||
def _block_maxima(scores):
|
||||
end = scores.numel()
|
||||
if end == 0:
|
||||
return scores
|
||||
padded = torch.full((cdiv(end, CAND_BLOCK) * CAND_BLOCK,), float("-inf"))
|
||||
padded[:end] = scores.cpu()
|
||||
out = padded.view(-1, CAND_BLOCK).amax(1)
|
||||
out[(end - 1) // CAND_BLOCK] = float("inf")
|
||||
return out.to(scores.device)
|
||||
|
||||
|
||||
def _forward_context(monkeypatch, metadata):
|
||||
context = types.SimpleNamespace(
|
||||
attn_metadata={"indexer": metadata},
|
||||
cudagraph_runtime_mode=CUDAGraphMode.NONE,
|
||||
)
|
||||
monkeypatch.setattr(ops, "get_forward_context", lambda: context)
|
||||
|
||||
|
||||
def _decode_metadata(case, rows, ratio, query_lens):
|
||||
lens = torch.tensor([(p + 1) // ratio for _, p in rows], dtype=torch.int32)
|
||||
lens = lens.to(DEVICE)
|
||||
block_table = case.block_table[[req for req, _ in rows]].contiguous()
|
||||
entries = case.cache[ratio].shape[1]
|
||||
next_n = max(query_lens)
|
||||
context_lens = (torch.tensor(case.seq_lens, dtype=torch.int32) // ratio).to(DEVICE)
|
||||
native = native_decode(lens, block_table, query_lens, next_n, context_lens)
|
||||
schedule = torch.empty(
|
||||
ops.rocm_mxfp4_decode_schedule_words(HEADS, HEAD_DIM, entries, next_n),
|
||||
dtype=torch.int32,
|
||||
device=DEVICE,
|
||||
)
|
||||
return RocmMxfp4IndexerMetadata(
|
||||
seq_lens=None,
|
||||
max_seq_len=max(case.seq_lens),
|
||||
slot_mapping=None,
|
||||
num_decodes=len(case.seq_lens),
|
||||
num_decode_tokens=len(rows),
|
||||
num_prefills=0,
|
||||
num_prefill_tokens=0,
|
||||
decode=types.SimpleNamespace(block_table=block_table, seq_lens=lens[:, None]),
|
||||
decode_row_lens=lens,
|
||||
decode_block_ends=(lens + CAND_BLOCK - 1) // CAND_BLOCK,
|
||||
decode_native=native,
|
||||
decode_schedule=ops.build_rocm_mxfp4_decode_schedule(
|
||||
lens, HEADS, HEAD_DIM, entries, schedule, MAX_LEN // ratio, native
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _prefill_metadata(
|
||||
case, rows, ratio, chunk_bounds, query_start_loc, gather_rows, gather
|
||||
):
|
||||
chunks = []
|
||||
for req_lo, req_hi, t0, t1 in chunk_bounds:
|
||||
ends = torch.tensor(
|
||||
[(rows[t][1] + 1) // ratio for t in range(t0, t1)], dtype=torch.int32
|
||||
).to(DEVICE)
|
||||
chunks.append(
|
||||
DeepseekV32IndexerPrefillChunkMetadata(
|
||||
block_table=case.block_table[req_lo:req_hi],
|
||||
cu_seqlen_ks=torch.zeros_like(ends),
|
||||
cu_seqlen_ke=ends,
|
||||
cu_seq_lens=torch.zeros(1, dtype=torch.int32, device=DEVICE),
|
||||
token_to_seq=torch.zeros(1, dtype=torch.int32, device=DEVICE),
|
||||
total_seq_lens=1,
|
||||
token_start=t0,
|
||||
token_end=t1,
|
||||
num_reqs=req_hi - req_lo,
|
||||
)
|
||||
)
|
||||
seq_lens = torch.tensor(case.seq_lens)
|
||||
context_lens = (seq_lens // ratio).int().to(DEVICE)
|
||||
plans = plan_prefill_chunks(
|
||||
chunks,
|
||||
query_start_loc,
|
||||
seq_lens,
|
||||
context_lens,
|
||||
ratio,
|
||||
CAND_BLOCK,
|
||||
0.0 if gather else None,
|
||||
torch.tensor(query_start_loc, dtype=torch.int32, device=DEVICE),
|
||||
)
|
||||
return RocmMxfp4IndexerMetadata(
|
||||
seq_lens=None,
|
||||
max_seq_len=max(case.seq_lens),
|
||||
slot_mapping=None,
|
||||
num_decodes=0,
|
||||
num_decode_tokens=0,
|
||||
num_prefills=len(case.seq_lens),
|
||||
num_prefill_tokens=len(rows),
|
||||
prefill=types.SimpleNamespace(chunks=chunks),
|
||||
prefill_plans=plans,
|
||||
gather_launches=plan_gather_launches(chunks, plans, gather_rows),
|
||||
)
|
||||
|
||||
|
||||
# The kernel prefers 128-token pages.
|
||||
BLOCKS = pytest.mark.parametrize("block", [128, 64])
|
||||
|
||||
|
||||
@BLOCKS
|
||||
def test_k_store_stays_inside_its_pages(block):
|
||||
"""The writer touches only its pages of the block-major pool. What they hold
|
||||
is checked by the layer tests, which read them back through aiter's kernel."""
|
||||
set_random_seed(0)
|
||||
case = _Case([700, 333], block)
|
||||
assert int(case.pool[~case.page_bytes].count_nonzero()) == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"rows, width, k",
|
||||
# vLLM's kernel below 128 rows, aiter's from 128 rows and on every k=2048
|
||||
[(8, 4096, 512), (128, 4096, 512), (8, 16384, 2048)],
|
||||
)
|
||||
def test_topk_matches_torch_on_both_kernels(rows, width, k):
|
||||
set_random_seed(0)
|
||||
logits = torch.randn(rows, width, device=DEVICE)
|
||||
lengths = torch.randint(1, width + 1, (rows,), dtype=torch.int32, device=DEVICE)
|
||||
lengths[0] = k // 2 # shorter than k, so the rest of the row pads with -1
|
||||
out = torch.full((rows, k), 7, dtype=torch.int32, device=DEVICE)
|
||||
ops._topk(logits, lengths, out, k)
|
||||
for r in range(rows):
|
||||
n = min(k, int(lengths[r]))
|
||||
want = torch.topk(logits[r, : int(lengths[r])], n).indices.sort().values
|
||||
got = out[r]
|
||||
assert torch.equal(got[got >= 0].sort().values, want.int())
|
||||
assert int((got == -1).sum()) == k - n
|
||||
|
||||
|
||||
def _run_layers(monkeypatch, case, rows, metadata):
|
||||
"""Ratio-2 dense layer, the candidate source, and a consumer both ways;
|
||||
each checked against the torch reference."""
|
||||
num_rows = len(rows)
|
||||
q, q_scale, weights, q_ref = _queries(num_rows)
|
||||
hidden = torch.empty(num_rows, 1, device=DEVICE)
|
||||
|
||||
def topk_buffer():
|
||||
return torch.full((num_rows + 2, TOPK), 7, dtype=torch.int32, device=DEVICE)
|
||||
|
||||
_forward_context(monkeypatch, metadata(2))
|
||||
out = topk_buffer()
|
||||
ops.rocm_mxfp4_sparse_attn_indexer(
|
||||
hidden,
|
||||
"indexer",
|
||||
case.cache[2],
|
||||
q,
|
||||
q_scale,
|
||||
weights,
|
||||
TOPK,
|
||||
HEAD_DIM,
|
||||
MAX_LEN // 2,
|
||||
out,
|
||||
compress_ratio=2,
|
||||
)
|
||||
for i, (req, pos) in enumerate(rows):
|
||||
end = (pos + 1) // 2
|
||||
keys = case.keys(2, req, end)
|
||||
_assert_topk(out[i], _scores(q_ref[i], weights[i], keys), TOPK)
|
||||
|
||||
source = metadata(1)
|
||||
_forward_context(monkeypatch, source)
|
||||
pool = torch.full((num_rows + 2, POOL_BLOCKS), 5, dtype=torch.int32, device=DEVICE)
|
||||
out = topk_buffer()
|
||||
ops.rocm_mxfp4_sparse_attn_indexer(
|
||||
hidden,
|
||||
"indexer",
|
||||
case.cache[1],
|
||||
q,
|
||||
q_scale,
|
||||
weights,
|
||||
TOPK,
|
||||
HEAD_DIM,
|
||||
MAX_LEN,
|
||||
out,
|
||||
compress_ratio=1,
|
||||
candidate_blocks=pool,
|
||||
candidate_block_size=CAND_BLOCK,
|
||||
candidate_write=True,
|
||||
)
|
||||
allowed = []
|
||||
for i, (req, pos) in enumerate(rows):
|
||||
scores = _scores(q_ref[i], weights[i], case.keys(1, req, pos + 1))
|
||||
_assert_topk(out[i], scores, TOPK)
|
||||
_assert_topk(pool[i], _block_maxima(scores), POOL_BLOCKS)
|
||||
mask = torch.zeros_like(scores, dtype=torch.bool)
|
||||
for block in pool[i][pool[i] >= 0].tolist():
|
||||
mask[block * CAND_BLOCK : (block + 1) * CAND_BLOCK] = True
|
||||
allowed.append(torch.where(mask, scores, float("-inf")))
|
||||
|
||||
pa = ops._aiter()
|
||||
build = pa.build_candidate_gather
|
||||
builds = []
|
||||
|
||||
def counting_build(*args, **kwargs):
|
||||
builds.append(1)
|
||||
return build(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(pa, "build_candidate_gather", counting_build)
|
||||
for gather in (True, False):
|
||||
consumer = metadata(1, gather)
|
||||
consumer.decode_use_gather = gather
|
||||
_forward_context(monkeypatch, consumer)
|
||||
builds.clear()
|
||||
# Two consumer layers of one step: the second reuses the first's pools.
|
||||
for _ in range(2):
|
||||
out = topk_buffer()
|
||||
ops.rocm_mxfp4_sparse_mqa_indexer(
|
||||
hidden,
|
||||
"indexer",
|
||||
case.cache[1],
|
||||
q,
|
||||
q_scale,
|
||||
weights,
|
||||
TOPK,
|
||||
HEAD_DIM,
|
||||
MAX_LEN,
|
||||
out,
|
||||
1,
|
||||
pool,
|
||||
CAND_BLOCK,
|
||||
POOL_BLOCKS * CAND_BLOCK,
|
||||
)
|
||||
for i in range(num_rows):
|
||||
_assert_topk(out[i], allowed[i], TOPK)
|
||||
if gather:
|
||||
launches = len(consumer.gather_launches) + (consumer.decode is not None)
|
||||
assert len(builds) == launches
|
||||
|
||||
|
||||
@BLOCKS
|
||||
@pytest.mark.parametrize(
|
||||
"query_lens",
|
||||
# Uniform steps launch the dense layers on next_n-row sequences, also with
|
||||
# cudagraph padding (query length 0) after them; a ragged step keeps a row
|
||||
# per token.
|
||||
[[6, 6, 6, 6], [2, 2, 2], [2, 2, 2, 0], [2, 1, 2]],
|
||||
ids=["native6", "native2", "padded", "ragged"],
|
||||
)
|
||||
def test_decode_layers_match_reference(monkeypatch, block, query_lens):
|
||||
set_random_seed(0)
|
||||
next_n = max(query_lens)
|
||||
case = _Case([q and n for q, n in zip(query_lens, [900, 333, 610, 1200])], block)
|
||||
# a padding request has next_n rows of its own and no context
|
||||
rows = [
|
||||
(req, n - q + j if q else -1)
|
||||
for req, (n, q) in enumerate(zip(case.seq_lens, query_lens))
|
||||
for j in range(q or next_n)
|
||||
]
|
||||
_run_layers(
|
||||
monkeypatch,
|
||||
case,
|
||||
rows,
|
||||
lambda r, gather=False: _decode_metadata(case, rows, r, query_lens),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"seq_lens,new_tokens,chunks",
|
||||
[
|
||||
# one request past a cached prefix, then one sliced into two chunks
|
||||
([760, 500], [300, 500], [(0, 1, 0, 300), (1, 2, 300, 560), (1, 2, 560, 800)]),
|
||||
# three requests in one chunk, the first two with equal rows
|
||||
([600, 450, 900], [200, 200, 120], [(0, 3, 0, 520)]),
|
||||
# four requests with four different query lengths in one chunk
|
||||
([900, 600, 450, 700], [150, 90, 200, 60], [(0, 4, 0, 500)]),
|
||||
],
|
||||
ids=["sliced", "batched", "ragged"],
|
||||
)
|
||||
# The consumers rejoin a request's rows across chunks, then cut them at most
|
||||
# gather_rows a launch.
|
||||
@pytest.mark.parametrize("gather_rows", [1 << 20, 128], ids=["joined", "cut"])
|
||||
@BLOCKS
|
||||
def test_prefill_layers_match_reference(
|
||||
monkeypatch, seq_lens, new_tokens, chunks, gather_rows, block
|
||||
):
|
||||
set_random_seed(0)
|
||||
case = _Case(seq_lens, block)
|
||||
rows = [
|
||||
(req, n - q + i)
|
||||
for req, (n, q) in enumerate(zip(seq_lens, new_tokens))
|
||||
for i in range(q)
|
||||
]
|
||||
query_start_loc = [0, *itertools.accumulate(new_tokens)]
|
||||
_run_layers(
|
||||
monkeypatch,
|
||||
case,
|
||||
rows,
|
||||
lambda r, gather=True: _prefill_metadata(
|
||||
case, rows, r, chunks, query_start_loc, gather_rows, gather
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _rope_gptj(x, cos_sin, positions):
|
||||
"""GPT-J RoPE on the last 64 dims, as the indexer rotates Q and K: bf16 in
|
||||
and out."""
|
||||
x = x.bfloat16().float()
|
||||
cos, sin = cos_sin[positions].chunk(2, dim=-1)
|
||||
if x.dim() == 3:
|
||||
cos, sin = cos[:, None], sin[:, None]
|
||||
rot = x[..., -64:]
|
||||
even, odd = rot[..., 0::2].clone(), rot[..., 1::2].clone()
|
||||
rot[..., 0::2] = even * cos - odd * sin
|
||||
rot[..., 1::2] = odd * cos + even * sin
|
||||
return x.bfloat16()
|
||||
|
||||
|
||||
def _mxfp4_round_trip(x):
|
||||
"""MXFP4 round trip of x, rounded as the writers do: a scale of
|
||||
2^ceil(log2(amax / 6)) per 32 values, then the nearest e2m1 value, ties to
|
||||
even."""
|
||||
blocks = x.float().unflatten(-1, (-1, 32))
|
||||
amax = blocks.abs().amax(-1, keepdim=True).clamp(min=2**-126)
|
||||
scale = torch.exp2(torch.ceil(torch.log2(amax / 6)))
|
||||
y = (blocks / scale).clamp(-6, 6)
|
||||
# Even codes first: argmin keeps the first of two equal distances.
|
||||
grid = torch.cat([_E2M1[0::2], _E2M1[1::2]]).to(x.device)
|
||||
idx = (y.abs()[..., None] - grid).abs().argmin(-1)
|
||||
return (grid[idx] * y.sign() * scale).flatten(-2)
|
||||
|
||||
|
||||
def _torch_keys(case, ratio, req):
|
||||
"""A request's compressed keys in torch: k_norm of each group's last
|
||||
token, rotated at the group's first position, through MXFP4."""
|
||||
k = case.k_pre[req][ratio - 1 :: ratio].float()
|
||||
k = k * torch.rsqrt(k.pow(2).mean(-1, keepdim=True) + 1e-6) * case.norm
|
||||
pos = torch.arange(k.shape[0], device=DEVICE) * ratio
|
||||
return _mxfp4_round_trip(_rope_gptj(k, case.cos_sin, pos))
|
||||
|
||||
|
||||
def _topk_indices(q, weights, keys, ends, k):
|
||||
"""Top-k of sum_h relu(q_h . key) * weights_h over each row's [0, end)."""
|
||||
scores = (torch.einsum("thd,cd->thc", q, keys).relu() * weights[..., None]).sum(1)
|
||||
cols = torch.arange(keys.shape[0], device=DEVICE)
|
||||
return scores.masked_fill(cols >= ends[:, None], float("-inf")).topk(k).indices
|
||||
|
||||
|
||||
def _recall(selected, reference):
|
||||
"""Per row, the fraction of the reference top-k that ``selected`` picked."""
|
||||
return (selected[:, :, None] == reference[:, None, :]).any(1).float().mean(1)
|
||||
|
||||
|
||||
@BLOCKS
|
||||
@pytest.mark.parametrize("ratio", [2, 1])
|
||||
@pytest.mark.parametrize("step", ["decode", "prefill"])
|
||||
def test_selection_from_bf16_matches_torch(monkeypatch, step, ratio, block):
|
||||
"""End to end from bf16 K and Q: vLLM's MXFP4 K store and Q quant, the
|
||||
paged pool, metadata from the planning functions the builder calls, and
|
||||
the indexer op pick what the same MXFP4 math picks in torch."""
|
||||
set_random_seed(0)
|
||||
if step == "decode":
|
||||
case = _Case([3000, 2600, 1900, 3500, 1200, 2200, 3900, 1600], block)
|
||||
rows = [(req, n - 1) for req, n in enumerate(case.seq_lens)]
|
||||
metadata = _decode_metadata(case, rows, ratio, [1] * len(rows))
|
||||
else:
|
||||
new_tokens = [300, 200, 250]
|
||||
case = _Case([1500, 900, 2000], block)
|
||||
rows = [
|
||||
(req, n - q + i)
|
||||
for req, (n, q) in enumerate(zip(case.seq_lens, new_tokens))
|
||||
for i in range(q)
|
||||
]
|
||||
metadata = _prefill_metadata(
|
||||
case,
|
||||
rows,
|
||||
ratio,
|
||||
[(0, len(new_tokens), 0, len(rows))],
|
||||
[0, *itertools.accumulate(new_tokens)],
|
||||
1 << 20,
|
||||
gather=False,
|
||||
)
|
||||
num_rows = len(rows)
|
||||
q_bf16 = torch.randn(num_rows, HEADS, HEAD_DIM, device=DEVICE).to(torch.bfloat16)
|
||||
w = torch.randn(num_rows, HEADS, device=DEVICE)
|
||||
pos = torch.tensor([p for _, p in rows], device=DEVICE)
|
||||
scales = (HEAD_DIM**-0.5, HEADS**-0.5)
|
||||
|
||||
(q, q_scale), weights = rocm_mxfp4_indexer_q_quant(
|
||||
pos, q_bf16, case.cos_sin, w, *scales
|
||||
)
|
||||
_forward_context(monkeypatch, metadata)
|
||||
out = torch.full((num_rows, TOPK), 7, dtype=torch.int32, device=DEVICE)
|
||||
ops.rocm_mxfp4_sparse_attn_indexer(
|
||||
torch.empty(num_rows, 1, device=DEVICE),
|
||||
"indexer",
|
||||
case.cache[ratio],
|
||||
q,
|
||||
q_scale,
|
||||
weights,
|
||||
TOPK,
|
||||
HEAD_DIM,
|
||||
MAX_LEN // ratio,
|
||||
out,
|
||||
compress_ratio=ratio,
|
||||
)
|
||||
ends = (pos + 1) // ratio
|
||||
assert bool(((out >= 0) & (out < ends[:, None])).all())
|
||||
assert all(row.unique().numel() == TOPK for row in out)
|
||||
|
||||
q_ref = _mxfp4_round_trip(_rope_gptj(q_bf16, case.cos_sin, pos))
|
||||
w_ref = w * scales[0] * scales[1]
|
||||
agree = []
|
||||
for req in range(len(case.seq_lens)):
|
||||
mine = torch.tensor([r == req for r, _ in rows], device=DEVICE)
|
||||
ref = _topk_indices(
|
||||
q_ref[mine], w_ref[mine], _torch_keys(case, ratio, req), ends[mine], TOPK
|
||||
)
|
||||
agree.append(_recall(out[mine], ref))
|
||||
assert torch.cat(agree).mean().item() >= 0.99
|
||||
@@ -0,0 +1,162 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""The ROCm MXFP4 indexer reads the paged cache in place, so each prefill
|
||||
chunk is split back into its requests' rows, which one launch finds through
|
||||
the chunk's query_start_loc. A chunk is whole requests or one query slice of a
|
||||
long request, and both have to come back with the right rows and the right
|
||||
block-table rows.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.is_rocm():
|
||||
pytest.skip("ROCm-only", allow_module_level=True)
|
||||
|
||||
from vllm.v1.attention.backends.mla.indexer import (
|
||||
DeepseekV32IndexerPrefillChunkMetadata,
|
||||
)
|
||||
from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import (
|
||||
native_decode,
|
||||
plan_gather_launches,
|
||||
plan_prefill_chunks,
|
||||
split_prefill_chunks,
|
||||
)
|
||||
from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import rocm_mxfp4_consumer_rows
|
||||
|
||||
|
||||
def _chunk(block_table, token_start, token_end, row_ends):
|
||||
ends = torch.tensor(row_ends, dtype=torch.int32)
|
||||
return DeepseekV32IndexerPrefillChunkMetadata(
|
||||
block_table=block_table,
|
||||
cu_seqlen_ks=torch.full_like(ends, 7),
|
||||
cu_seqlen_ke=ends + 7,
|
||||
cu_seq_lens=torch.zeros(1, dtype=torch.int32),
|
||||
token_to_seq=torch.zeros(1, dtype=torch.int32),
|
||||
total_seq_lens=1,
|
||||
token_start=token_start,
|
||||
token_end=token_end,
|
||||
num_reqs=block_table.shape[0],
|
||||
)
|
||||
|
||||
|
||||
# A decode (1 token), then prefills of 3, 3, 5 and 6 query tokens; the last
|
||||
# one is sliced into two chunks.
|
||||
QUERY_START_LOC = [0, 1, 4, 7, 12, 18]
|
||||
SEQ_LENS = torch.tensor([40, 10, 12, 21, 64])
|
||||
BLOCK_TABLE = torch.arange(20, dtype=torch.int32).view(5, 4)
|
||||
CHUNKS = [
|
||||
_chunk(BLOCK_TABLE[1:4], 1, 12, list(range(11))),
|
||||
_chunk(BLOCK_TABLE[4:5], 12, 15, [29, 30, 30]),
|
||||
_chunk(BLOCK_TABLE[4:5], 15, 18, [31, 31, 32]),
|
||||
]
|
||||
|
||||
|
||||
def _plans(min_gather_width, block=8):
|
||||
return plan_prefill_chunks(
|
||||
CHUNKS,
|
||||
QUERY_START_LOC,
|
||||
SEQ_LENS,
|
||||
SEQ_LENS.int() // 2,
|
||||
2,
|
||||
block,
|
||||
min_gather_width,
|
||||
torch.tensor(QUERY_START_LOC, dtype=torch.int32),
|
||||
)
|
||||
|
||||
|
||||
def test_chunks_split_into_requests():
|
||||
plans = _plans(20.0)
|
||||
|
||||
assert [p.requests for p in plans] == [
|
||||
[(0, 3, 0), (3, 6, 1), (6, 11, 2)],
|
||||
[(0, 3, 0)],
|
||||
[(0, 3, 0)],
|
||||
]
|
||||
# every chunk launches once, with offsets local to the chunk
|
||||
assert [p.query_start_loc.tolist() for p in plans] == [
|
||||
[0, 3, 6, 11],
|
||||
[0, 3],
|
||||
[0, 3],
|
||||
]
|
||||
# the logits are as wide as the longest compressed context in the chunk
|
||||
assert [p.width for p in plans] == [10, 32, 32]
|
||||
assert [p.context_lens.tolist() for p in plans] == [[5, 6, 10], [32], [32]]
|
||||
torch.testing.assert_close(plans[2].row_ends, torch.tensor([31, 31, 32]).int())
|
||||
torch.testing.assert_close(plans[2].block_ends, torch.tensor([4, 4, 4]).int())
|
||||
assert [p.use_gather for p in plans] == [False, True, True]
|
||||
assert [p.first_request for p in plans] == [1, 4, 4]
|
||||
# no candidate blocks, as in the model-neutral builder
|
||||
for plan in _plans(None, block=0):
|
||||
assert plan.block_ends is None and not plan.use_gather
|
||||
|
||||
|
||||
def test_launches_follow_the_logits_budget(monkeypatch):
|
||||
"""Each launch's fp32 logits fit in VLLM_SPARSE_INDEXER_MAX_LOGITS_MB. A dense
|
||||
chunk's logits are [rows, longest context in the chunk]: K is read in place,
|
||||
not gathered, so requests share a launch until their rows no longer fit, and
|
||||
a request that does not fit alone is split on its query rows. A consumer's
|
||||
logits are [rows, pool]."""
|
||||
budget = 4 * 1000 * 30
|
||||
contexts, queries = torch.tensor([1000, 600, 200]), torch.tensor([10, 10, 10])
|
||||
assert split_prefill_chunks(contexts, queries, budget, 2) == [
|
||||
(slice(2, 5), slice(0, 30))
|
||||
]
|
||||
# the third request would take the chunk past the budget
|
||||
queries = torch.tensor([10, 10, 11])
|
||||
assert split_prefill_chunks(contexts, queries, budget) == [
|
||||
(slice(0, 2), slice(0, 20)),
|
||||
(slice(2, 3), slice(0, 11)),
|
||||
]
|
||||
assert split_prefill_chunks(torch.tensor([1000]), torch.tensor([100]), budget) == [
|
||||
(slice(0, 1), slice(lo, min(lo + 30, 100))) for lo in range(0, 100, 30)
|
||||
]
|
||||
# the default 512 MiB holds 128 rows of a 1M context, or 8192 of a 16K pool
|
||||
assert split_prefill_chunks(
|
||||
torch.tensor([1 << 20]), torch.tensor([200]), 512 << 20
|
||||
) == [(slice(0, 1), slice(0, 128)), (slice(0, 1), slice(128, 200))]
|
||||
monkeypatch.setenv("VLLM_SPARSE_INDEXER_MAX_LOGITS_MB", "512")
|
||||
assert rocm_mxfp4_consumer_rows(16384) == 8192
|
||||
|
||||
|
||||
def test_gather_launches_rejoin_a_requests_rows():
|
||||
"""The consumers' logits are [rows, pool] whatever the context, so a
|
||||
request's rows go back together across the chunks that cut them, and a
|
||||
launch never mixes requests."""
|
||||
# the first chunk is too narrow to gather, the last request's slices rejoin
|
||||
(joined,) = plan_gather_launches(CHUNKS, _plans(20.0), 100)
|
||||
assert (joined.token_start, joined.token_end) == (12, 18)
|
||||
# every chunk gathers: the first chunk's three requests launch apart
|
||||
launches = plan_gather_launches(CHUNKS, _plans(0.0), 100)
|
||||
assert [(g.token_start, g.token_end) for g in launches] == [
|
||||
(1, 4),
|
||||
(4, 7),
|
||||
(7, 12),
|
||||
(12, 18),
|
||||
]
|
||||
assert plan_gather_launches(CHUNKS, _plans(None), 100) == []
|
||||
|
||||
|
||||
def test_decode_launches_native_on_uniform_steps():
|
||||
"""A decode step goes to the kernel as next_n-row sequences only when its
|
||||
shape says so: each request has next_n rows, and cudagraph padding (query
|
||||
length 0) comes only after them. A FULL graph captured on a uniform batch
|
||||
then replays the same launch on a padded one."""
|
||||
rows = torch.zeros(8, dtype=torch.int32)
|
||||
row_block_table = torch.arange(32, dtype=torch.int32).view(8, 4)
|
||||
context_lens = torch.tensor([9, 4, 7, 0, 5], dtype=torch.int32)
|
||||
|
||||
native = native_decode(rows, row_block_table, [2, 2, 2, 0], 2, context_lens)
|
||||
assert native is not None and native.next_n == 2
|
||||
assert native.context_lens.tolist() == [9, 4, 7, 0]
|
||||
# every request's first flattened row, as a view
|
||||
assert native.block_table.tolist() == row_block_table[::2].tolist()
|
||||
assert native.block_table.data_ptr() == row_block_table.data_ptr()
|
||||
|
||||
assert native_decode(rows, row_block_table, [2, 0, 2, 2], 2, context_lens) is None
|
||||
# a ragged step, rows padded past the requests, and no speculation
|
||||
assert native_decode(rows[:5], row_block_table, [2, 1, 2], 2, context_lens) is None
|
||||
assert native_decode(rows, row_block_table, [2, 2, 2], 2, context_lens) is None
|
||||
assert native_decode(rows[:4], row_block_table, [1] * 4, 1, context_lens) is None
|
||||
@@ -93,7 +93,10 @@ class AttentionConfig:
|
||||
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."""
|
||||
for long contexts (roughly 32K tokens and beyond) and is slower below.
|
||||
On ROCm gfx950 it runs aiter's paged MXFP4 MQA-logits kernel over the
|
||||
candidate pool instead, with fp32 logits, and keeps the masked dense walk
|
||||
for steps whose contexts are too short for the pool to pay."""
|
||||
|
||||
hisparse_config: HiSparseConfig | None = None
|
||||
"""HiSparse host-resident KV configuration. Setting this enables experimental
|
||||
|
||||
@@ -0,0 +1,116 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Sparse attention indexer layers on aiter's paged MXFP4 kernels, for ROCm.
|
||||
|
||||
With ``indexer_kv_dtype="mxfp4"`` on gfx950, the model writes the indexer K
|
||||
cache in the order aiter's MQA-logits kernel reads (see
|
||||
`rocm_paged_mxfp4_cache_layout`), and the kernel scores the cache in place.
|
||||
Requires `RocmMxfp4IndexerMetadataBuilder` metadata.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm.config import get_current_vllm_config
|
||||
from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
|
||||
from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import (
|
||||
rocm_mxfp4_sparse_attn_indexer,
|
||||
rocm_mxfp4_sparse_mqa_indexer,
|
||||
)
|
||||
|
||||
|
||||
class RocmSparseAttnIndexer(SparseAttnIndexer):
|
||||
"""`SparseAttnIndexer` whose HIP path walks the paged MXFP4 cache with
|
||||
aiter's kernel. A two-level indexer's source layer also publishes its
|
||||
candidate pool here."""
|
||||
|
||||
def forward_hip(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
|
||||
k: torch.Tensor | None,
|
||||
weights: torch.Tensor,
|
||||
):
|
||||
assert isinstance(q_quant, tuple) and self.skip_k_cache_insert, (
|
||||
"the ROCm MXFP4 indexer takes (values, scales) Q and a K cache "
|
||||
"the model already wrote"
|
||||
)
|
||||
q_values, q_scale = q_quant
|
||||
return rocm_mxfp4_sparse_attn_indexer(
|
||||
hidden_states,
|
||||
self.k_cache.prefix,
|
||||
self.k_cache.kv_cache,
|
||||
q_values,
|
||||
q_scale,
|
||||
weights,
|
||||
self.topk_tokens,
|
||||
self.head_dim,
|
||||
self.max_model_len,
|
||||
self.topk_indices_buffer,
|
||||
compress_ratio=self.compress_ratio,
|
||||
candidate_blocks=self.candidate_blocks,
|
||||
candidate_block_size=self.candidate_block_size,
|
||||
candidate_write=self.candidate_write,
|
||||
)
|
||||
|
||||
|
||||
class RocmSparseMQAIndexer(nn.Module):
|
||||
"""DeepSeek-V4.1's candidate-consuming indexer on aiter's paged MXFP4
|
||||
MQA-logits kernel: the consumers score only the candidate pool the source
|
||||
layer published, read straight from the paged cache.
|
||||
|
||||
Built and called as `SparseMQAIndexer` is, so the model can take either.
|
||||
"""
|
||||
|
||||
weights_dtype = torch.float32
|
||||
"""The scorer's weights and logits stay fp32: gfx950 has no packed bf16
|
||||
arithmetic, so its head reduce costs the same in either dtype."""
|
||||
|
||||
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__()
|
||||
self.k_cache = k_cache
|
||||
self.topk_tokens = topk_tokens
|
||||
self.head_dim = head_dim
|
||||
# aiter's kernel bounds a row by the model length in compressed
|
||||
# positions, not by the prefill buffer SparseMQAIndexer sizes.
|
||||
model_config = get_current_vllm_config().model_config
|
||||
self.max_model_len = model_config.max_model_len // k_cache.compress_ratio
|
||||
self.topk_indices_buffer = topk_indices_buffer
|
||||
self.candidate_blocks = candidate_blocks
|
||||
self.candidate_block_size = candidate_block_size
|
||||
self.num_candidate_cols = candidate_blocks.shape[1] * candidate_block_size
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
q_quant: tuple[torch.Tensor, torch.Tensor],
|
||||
k: torch.Tensor | None,
|
||||
weights: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
assert k is None, "the model writes the indexer K cache"
|
||||
q_values, q_scale = q_quant
|
||||
return rocm_mxfp4_sparse_mqa_indexer(
|
||||
hidden_states,
|
||||
self.k_cache.prefix,
|
||||
self.k_cache.kv_cache,
|
||||
q_values,
|
||||
q_scale,
|
||||
weights,
|
||||
self.topk_tokens,
|
||||
self.head_dim,
|
||||
self.max_model_len,
|
||||
self.topk_indices_buffer,
|
||||
self.k_cache.compress_ratio,
|
||||
self.candidate_blocks,
|
||||
self.candidate_block_size,
|
||||
self.num_candidate_cols,
|
||||
)
|
||||
@@ -2,8 +2,9 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import functools
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import cast
|
||||
from typing import Any, cast
|
||||
|
||||
import torch
|
||||
|
||||
@@ -15,8 +16,14 @@ from vllm.forward_context import get_forward_context
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.fusion.quant_activation import QuantizedActivation
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import kMxfp8Dynamic
|
||||
from vllm.model_executor.layers.rocm_paged_mxfp4_indexer import (
|
||||
RocmSparseAttnIndexer,
|
||||
RocmSparseMQAIndexer,
|
||||
)
|
||||
from vllm.models.deepseek_v41.attention import (
|
||||
DeepseekV4Attention,
|
||||
DeepseekV4Indexer,
|
||||
DeepseekV4IndexerCache,
|
||||
_replace_layer_index,
|
||||
)
|
||||
from vllm.models.deepseek_v41.common.ops import dequantize_and_gather_k_cache
|
||||
@@ -32,6 +39,7 @@ from vllm.triton_utils import tl, triton
|
||||
from vllm.utils.multi_stream_utils import execute_in_parallel
|
||||
from vllm.v1.attention.backend import (
|
||||
CommonAttentionMetadata,
|
||||
MultipleOf,
|
||||
)
|
||||
from vllm.v1.attention.backends.mla.sparse_swa import (
|
||||
DeepseekSparseSWABackend,
|
||||
@@ -47,6 +55,9 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
|
||||
rocm_sparse_attn_decode,
|
||||
rocm_sparse_attn_prefill,
|
||||
)
|
||||
from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import (
|
||||
rocm_paged_mxfp4_cache_layout,
|
||||
)
|
||||
from vllm.v1.worker.workspace import current_workspace_manager
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -517,11 +528,141 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu
|
||||
)
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _aiter_indexer_cache_ops() -> tuple[Callable[..., None], Callable[..., tuple]]:
|
||||
"""The aiter indexer key writer and query quantizer."""
|
||||
from aiter.ops.triton.fusions.k_norm_rope_mxfp4_cache import (
|
||||
k_norm_rope_mxfp4_cache,
|
||||
)
|
||||
from aiter.ops.triton.rope.q_rope_mxfp4_quant import q_rope_mxfp4_quant
|
||||
|
||||
return k_norm_rope_mxfp4_cache, q_rope_mxfp4_quant
|
||||
|
||||
|
||||
def rocm_mxfp4_indexer_k_store(
|
||||
k_pre: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
rms_norm_weight: torch.Tensor,
|
||||
rms_norm_eps: float,
|
||||
k_cache: torch.Tensor,
|
||||
kv_slot_mapping: torch.Tensor,
|
||||
compress_ratio: int,
|
||||
use_fp4_cache: bool,
|
||||
*,
|
||||
num_heads: int,
|
||||
) -> None:
|
||||
"""`indexer_k_norm_rope_store` for the ROCm MXFP4 cache: aiter's cache op
|
||||
writes the key in the order its MQA-logits kernel reads with ``num_heads``
|
||||
query heads."""
|
||||
assert use_fp4_cache, "the ROCm indexer cache op writes MXFP4 only"
|
||||
layout = rocm_paged_mxfp4_cache_layout(num_heads, k_pre.shape[1], k_cache.shape[1])
|
||||
k_norm_rope_mxfp4_cache, _ = _aiter_indexer_cache_ops()
|
||||
k_norm_rope_mxfp4_cache(
|
||||
k_pre,
|
||||
positions,
|
||||
cos_sin_cache,
|
||||
rms_norm_weight,
|
||||
rms_norm_eps,
|
||||
k_cache,
|
||||
kv_slot_mapping,
|
||||
compress_ratio,
|
||||
shuffle=layout,
|
||||
)
|
||||
|
||||
|
||||
def rocm_mxfp4_indexer_q_quant(
|
||||
positions: torch.Tensor,
|
||||
index_q: torch.Tensor,
|
||||
index_q_cos_sin_cache: torch.Tensor,
|
||||
index_weights: torch.Tensor,
|
||||
index_weights_softmax_scale: float,
|
||||
index_weights_head_scale: float,
|
||||
use_fp4: bool = True,
|
||||
weights_out_dtype: torch.dtype = torch.float32,
|
||||
) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor]:
|
||||
"""`fused_indexer_q_rope_quant` for the ROCm MXFP4 indexer, on aiter's op:
|
||||
((packed [T, H, D // 2], e8m0 as one int32 per head [T, H]), fp32
|
||||
weights)."""
|
||||
assert use_fp4 and weights_out_dtype == torch.float32, (
|
||||
"the ROCm indexer quantizes Q to MXFP4 and scores with fp32 weights"
|
||||
)
|
||||
_, q_rope_mxfp4_quant = _aiter_indexer_cache_ops()
|
||||
q_packed, q_scale, weights_out = q_rope_mxfp4_quant(
|
||||
index_q,
|
||||
positions,
|
||||
index_q_cos_sin_cache,
|
||||
index_weights,
|
||||
index_weights_softmax_scale * index_weights_head_scale,
|
||||
)
|
||||
return (q_packed, q_scale.view(torch.int32).squeeze(-1)), weights_out
|
||||
|
||||
|
||||
class DeepseekV41RocmMxfp4Indexer(DeepseekV4Indexer):
|
||||
"""The indexer on aiter's paged MXFP4 cache: its K store and Q quant write
|
||||
in the order aiter's MQA-logits kernel reads, and its layers score with it.
|
||||
``_produce_k`` and ``forward_q`` mirror ``DeepseekV4Indexer``'s except for
|
||||
those two calls."""
|
||||
|
||||
mqa_cls = RocmSparseMQAIndexer
|
||||
attn_cls = RocmSparseAttnIndexer
|
||||
|
||||
def _produce_k(
|
||||
self,
|
||||
latent: torch.Tensor | None,
|
||||
positions: torch.Tensor,
|
||||
rotary_emb: torch.nn.Module,
|
||||
) -> None:
|
||||
attn_metadata = get_forward_context().attn_metadata
|
||||
if not isinstance(attn_metadata, dict) or latent is None:
|
||||
return
|
||||
assert self.owns_k
|
||||
indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix])
|
||||
k_pre, _ = self.wk(latent)
|
||||
rocm_mxfp4_indexer_k_store(
|
||||
k_pre,
|
||||
positions,
|
||||
rotary_emb.cos_sin_cache,
|
||||
self.k_norm.weight,
|
||||
self.k_norm.variance_epsilon,
|
||||
self.k_cache.kv_cache,
|
||||
indexer_metadata.slot_mapping,
|
||||
self.compress_ratio,
|
||||
self.use_fp4_kv,
|
||||
num_heads=self.n_head,
|
||||
)
|
||||
|
||||
def forward_q(
|
||||
self,
|
||||
qr: torch.Tensor | QuantizedActivation,
|
||||
qr_scale: torch.Tensor | None,
|
||||
indexer_weights: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
rotary_emb: torch.nn.Module,
|
||||
) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]:
|
||||
q = self._wq_b_proj(qr, qr_scale).view(-1, self.n_head, self.head_dim)
|
||||
(q, q_scale), weights = rocm_mxfp4_indexer_q_quant(
|
||||
positions,
|
||||
q,
|
||||
rotary_emb.cos_sin_cache,
|
||||
indexer_weights,
|
||||
self.softmax_scale,
|
||||
self.n_head**-0.5,
|
||||
use_fp4=self.use_fp4_kv,
|
||||
weights_out_dtype=self.indexer_weights_dtype,
|
||||
)
|
||||
return q, q_scale, weights
|
||||
|
||||
|
||||
class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4SparseMLABackend):
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "ROCM_FLASHMLA_SPARSE_DSV4"
|
||||
|
||||
@staticmethod
|
||||
def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
|
||||
return [64, 128]
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type[DeepseekV4SparseMLAMetadataBuilder]:
|
||||
return DeepseekV4SparseMLAMetadataBuilder
|
||||
@@ -540,6 +681,13 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention):
|
||||
swa_backend_cls = DeepseekV41ROCMAiterSparseSWABackend
|
||||
_use_aiter_sparse_mla = False
|
||||
|
||||
def _indexer_cls(
|
||||
self, k_cache: DeepseekV4IndexerCache | None
|
||||
) -> type[DeepseekV4Indexer]:
|
||||
if k_cache is not None and k_cache.rocm_mxfp4:
|
||||
return DeepseekV41RocmMxfp4Indexer
|
||||
return super()._indexer_cls(k_cache)
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
vllm_config = args[0] if args else kwargs["vllm_config"]
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
@@ -35,6 +35,9 @@ from vllm.models.deepseek_v41.common.ops import (
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.model_executor.kernels.linear.cute_dsl.gemm_rs_ar import GemmRsAr
|
||||
from vllm.model_executor.layers.rocm_paged_mxfp4_indexer import (
|
||||
RocmSparseMQAIndexer,
|
||||
)
|
||||
from vllm.v1.attention.backends.mla.sparse_swa import (
|
||||
DeepseekSparseSWAMetadata,
|
||||
)
|
||||
@@ -261,6 +264,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
|
||||
"""Return whether this instance stores fp8 KV in fp8_ds_mla layout."""
|
||||
return self.use_fp8_ds_mla_layout
|
||||
|
||||
def _indexer_cls(
|
||||
self, k_cache: "DeepseekV4IndexerCache | None"
|
||||
) -> type["DeepseekV4Indexer"]:
|
||||
"""The indexer for this layer's index K cache. A platform subclass
|
||||
returns one that writes and scores its own cache layout."""
|
||||
return DeepseekV4Indexer
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
@@ -449,7 +459,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
|
||||
)
|
||||
is_candidate_source = layer_id == self.candidate_source_layer
|
||||
uses_candidates = 0 <= self.candidate_source_layer < layer_id
|
||||
self.indexer = DeepseekV4Indexer(
|
||||
self.indexer = self._indexer_cls(index_k_cache)(
|
||||
vllm_config,
|
||||
config=config,
|
||||
hidden_size=self.hidden_size,
|
||||
@@ -1152,6 +1162,10 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
|
||||
self.compress_ratio = compress_ratio
|
||||
vllm_config = get_current_vllm_config()
|
||||
self.sparse_logits = vllm_config.attention_config.indexer_sparse_logits
|
||||
# aiter's paged MXFP4 kernels write this cache and read it in place.
|
||||
self.rocm_mxfp4 = current_platform.is_rocm() and dsa_indexer_uses_fp4(
|
||||
vllm_config
|
||||
)
|
||||
compilation_config = vllm_config.compilation_config
|
||||
if prefix in compilation_config.static_forward_context:
|
||||
raise ValueError(f"Duplicate layer name: {prefix}")
|
||||
@@ -1190,13 +1204,21 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
|
||||
# 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
|
||||
math.lcm(512, page_alignment)
|
||||
if self.sparse_logits and not self.rocm_mxfp4
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
def forward(self): ...
|
||||
|
||||
def get_attn_backend(self) -> type[AttentionBackend]:
|
||||
if self.rocm_mxfp4:
|
||||
from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import (
|
||||
DeepseekV41RocmMxfp4IndexerBackend,
|
||||
)
|
||||
|
||||
return DeepseekV41RocmMxfp4IndexerBackend
|
||||
if self.sparse_logits:
|
||||
return DeepseekV41SparseIndexerBackend
|
||||
return DeepseekV41IndexerBackend
|
||||
@@ -1218,6 +1240,12 @@ class DeepseekV4Indexer(nn.Module):
|
||||
later indexers mask their scores to those blocks before their own top-k.
|
||||
"""
|
||||
|
||||
# The scoring layers; a platform subclass swaps in its own.
|
||||
mqa_cls: ClassVar[type["SparseMQAIndexer | RocmSparseMQAIndexer"]] = (
|
||||
SparseMQAIndexer
|
||||
)
|
||||
attn_cls: ClassVar[type[SparseAttnIndexer]] = SparseAttnIndexer
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vllm_config: VllmConfig,
|
||||
@@ -1305,7 +1333,7 @@ class DeepseekV4Indexer(nn.Module):
|
||||
and candidate_block_buffer is not None
|
||||
and not candidate_write
|
||||
)
|
||||
self.indexer_op: SparseAttnIndexer | SparseMQAIndexer
|
||||
self.indexer_op: SparseAttnIndexer | SparseMQAIndexer | RocmSparseMQAIndexer
|
||||
if use_sparse_logits:
|
||||
if not self.use_fp4_kv:
|
||||
raise ValueError(
|
||||
@@ -1314,7 +1342,7 @@ class DeepseekV4Indexer(nn.Module):
|
||||
)
|
||||
assert candidate_block_buffer is not None
|
||||
assert topk_indices_buffer is not None
|
||||
self.indexer_op = SparseMQAIndexer(
|
||||
self.indexer_op = self.mqa_cls(
|
||||
self.k_cache,
|
||||
self.topk_tokens,
|
||||
self.head_dim,
|
||||
@@ -1324,7 +1352,7 @@ class DeepseekV4Indexer(nn.Module):
|
||||
candidate_block_size,
|
||||
)
|
||||
else:
|
||||
self.indexer_op = SparseAttnIndexer(
|
||||
self.indexer_op = self.attn_cls(
|
||||
self.k_cache,
|
||||
self.quant_block_size,
|
||||
self.scale_fmt,
|
||||
@@ -1343,7 +1371,7 @@ class DeepseekV4Indexer(nn.Module):
|
||||
# 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
|
||||
self.mqa_cls.weights_dtype if use_sparse_logits else torch.float32
|
||||
)
|
||||
|
||||
def _produce_k(
|
||||
|
||||
@@ -51,7 +51,7 @@ from vllm.v1.kv_cache_interface import (
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# The DSA indexer K cache is always quantized; "auto" means fp8 (V3.2 layout)
|
||||
# and mxfp4 is the opt-in Blackwell path.
|
||||
# and mxfp4 is the opt-in Blackwell and gfx950 path.
|
||||
DSA_INDEXER_KV_DTYPES = ("fp8", "mxfp4")
|
||||
|
||||
|
||||
@@ -64,6 +64,28 @@ def dsa_indexer_uses_fp4(vllm_config: VllmConfig) -> bool:
|
||||
f"sparse indexer (expected one of {DSA_INDEXER_KV_DTYPES})."
|
||||
)
|
||||
use_fp4 = kv_dtype == "mxfp4"
|
||||
if use_fp4 and current_platform.is_rocm():
|
||||
from vllm._aiter_ops import rocm_aiter_ops
|
||||
from vllm.platforms.rocm import get_cdna_version
|
||||
|
||||
# Only DeepSeek-V4.1 is wired to the ROCm MXFP4 cache; other DSA models
|
||||
# would silently keep their FP8 one.
|
||||
model_config = vllm_config.model_config
|
||||
if model_config is None or model_config.hf_config.model_type != "deepseek_v41":
|
||||
raise ValueError(
|
||||
"indexer_kv_dtype='mxfp4' on ROCm is only supported for "
|
||||
"DeepSeek-V4.1-Flash."
|
||||
)
|
||||
if get_cdna_version() != 4:
|
||||
raise ValueError(
|
||||
"indexer_kv_dtype='mxfp4' on ROCm requires CDNA4 (MI350X/MI355X)."
|
||||
)
|
||||
if not rocm_aiter_ops.is_enabled():
|
||||
raise ValueError(
|
||||
"indexer_kv_dtype='mxfp4' on ROCm runs on aiter's kernels; enable "
|
||||
"aiter with VLLM_ROCM_USE_AITER=1."
|
||||
)
|
||||
return True
|
||||
if use_fp4 and not current_platform.is_device_capability_family(100):
|
||||
raise ValueError(
|
||||
"indexer_kv_dtype='mxfp4' requires Blackwell datacenter GPUs "
|
||||
|
||||
@@ -0,0 +1,513 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Indexer metadata for the ROCm paged MXFP4 path.
|
||||
|
||||
Extends the dense indexer metadata with what aiter's paged MXFP4 MQA-logits
|
||||
kernel needs to read the cache in place: each prefill chunk as its requests,
|
||||
packed through query_start_loc into one launch, and the decode rows as
|
||||
next_n-row sequences on uniform steps. `RocmMxfp4IndexerMetadataBuilder` plans
|
||||
this for any model's indexer, and a two-level indexer's candidate blocks once
|
||||
a subclass sets their size. DeepSeek-V4.1's builder does, and with
|
||||
``AttentionConfig.indexer_sparse_logits`` adds the per-step state the
|
||||
candidate consumers share.
|
||||
"""
|
||||
|
||||
import bisect
|
||||
from dataclasses import dataclass, field, fields
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.logger import init_logger
|
||||
from vllm.v1.attention.backend import (
|
||||
AttentionCGSupport,
|
||||
CommonAttentionMetadata,
|
||||
MultipleOf,
|
||||
)
|
||||
from vllm.v1.attention.backends.mla.indexer import (
|
||||
DeepseekV32IndexerMetadata,
|
||||
DeepseekV32IndexerMetadataBuilder,
|
||||
DeepseekV32IndexerPrefillChunkMetadata,
|
||||
DeepseekV41IndexerBackend,
|
||||
)
|
||||
from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import (
|
||||
build_rocm_mxfp4_decode_schedule,
|
||||
rocm_mxfp4_consumer_rows,
|
||||
rocm_mxfp4_decode_schedule_words,
|
||||
rocm_paged_mxfp4_cache_layout,
|
||||
)
|
||||
from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# The consumers walk the pool instead of the context once the context is this
|
||||
# many pools long. Set from kernel timings, not end-to-end runs: the gather
|
||||
# loads more per key than the dense walk, so it only wins on a large cut.
|
||||
_GATHER_MIN_CUT_DECODE = 3.0
|
||||
_GATHER_MIN_CUT_PREFILL = 1.5
|
||||
|
||||
|
||||
@dataclass
|
||||
class RocmMxfp4PrefillPlan:
|
||||
"""One prefill chunk as the requests the kernel launches on, all in one
|
||||
launch: the kernel finds each row's request through query_start_loc, and
|
||||
a request's rows share its KV tiles."""
|
||||
|
||||
requests: list[tuple[int, int, int]]
|
||||
"""(first row, end row, request) per request with rows in the chunk; rows
|
||||
count from ``chunk.token_start``, requests index ``chunk.block_table``."""
|
||||
width: int
|
||||
"""Logits columns: an upper bound on the chunk's compressed contexts."""
|
||||
row_ends: torch.Tensor
|
||||
"""[rows] int32 exclusive compressed key bound of each row."""
|
||||
context_lens: torch.Tensor
|
||||
"""[chunk.num_reqs] int32 compressed context of each request."""
|
||||
query_start_loc: torch.Tensor
|
||||
"""[chunk.num_reqs + 1] int32 row offsets of the chunk's requests, local to
|
||||
the chunk."""
|
||||
first_request: int
|
||||
"""The step's index of the chunk's first request."""
|
||||
block_ends: torch.Tensor | None = None
|
||||
"""[rows] int32 candidate blocks each row sees, for the source's pool."""
|
||||
use_gather: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class RocmMxfp4GatherLaunch:
|
||||
"""Query rows [token_start, token_end) of one request, which the candidate
|
||||
consumers score against the pool in one launch."""
|
||||
|
||||
token_start: int
|
||||
token_end: int
|
||||
block_table: torch.Tensor
|
||||
"""[1, max_blocks] int32, the request's block table row."""
|
||||
context_len: torch.Tensor
|
||||
"""[1] int32 compressed context of the request."""
|
||||
row_ends: torch.Tensor
|
||||
"""[rows] int32 exclusive compressed key bound of each row."""
|
||||
pool: tuple[dict, torch.Tensor] | None = None
|
||||
"""The resolved candidate lists, built by the first consumer layer and
|
||||
reused by the rest."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class RocmMxfp4NativeDecode:
|
||||
"""A decode step whose requests all have next_n query rows, as the
|
||||
kernel's sequences: a workgroup then walks a KV tile once for up to
|
||||
next_n rows instead of once per row."""
|
||||
|
||||
next_n: int
|
||||
context_lens: torch.Tensor
|
||||
"""[requests] int32 compressed context of each request."""
|
||||
block_table: torch.Tensor
|
||||
"""[requests, max_blocks] int32, each request's first flattened row."""
|
||||
|
||||
|
||||
def native_decode(
|
||||
row_lens: torch.Tensor,
|
||||
row_block_table: torch.Tensor,
|
||||
query_lens: list[int],
|
||||
next_n: int,
|
||||
context_lens: torch.Tensor,
|
||||
) -> RocmMxfp4NativeDecode | None:
|
||||
"""The step as next_n-row sequences, or None when a request has fewer rows.
|
||||
|
||||
``row_lens`` and ``row_block_table`` are the flattened rows' bounds and
|
||||
block table, ``query_lens`` the decode requests' query lengths. Trailing
|
||||
empty requests are cudagraph padding: they keep their next_n rows, with no
|
||||
context. Only the step's shape decides, so a FULL graph captured on a
|
||||
uniform batch replays the same launch on a padded one.
|
||||
"""
|
||||
num_reqs = len(query_lens)
|
||||
num_full = next((i for i, n in enumerate(query_lens) if n != next_n), num_reqs)
|
||||
if (
|
||||
next_n <= 1
|
||||
or row_lens.shape[0] != num_reqs * next_n
|
||||
or any(query_lens[num_full:])
|
||||
):
|
||||
return None
|
||||
return RocmMxfp4NativeDecode(
|
||||
next_n, context_lens[:num_reqs], row_block_table[::next_n]
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RocmMxfp4IndexerMetadata(DeepseekV32IndexerMetadata):
|
||||
prefill_plans: list[RocmMxfp4PrefillPlan] = field(default_factory=list)
|
||||
"""Parallel to ``prefill.chunks``."""
|
||||
decode_row_lens: torch.Tensor | None = None
|
||||
"""[rows] int32 compressed context of each decode query row."""
|
||||
decode_block_ends: torch.Tensor | None = None
|
||||
"""[rows] int32 candidate blocks each decode row sees, for the source's
|
||||
pool."""
|
||||
decode_native: RocmMxfp4NativeDecode | None = None
|
||||
"""The same rows as next_n-row sequences, for the dense launches."""
|
||||
decode_schedule: torch.Tensor | None = None
|
||||
"""The dense decode launches' work schedule, built once for the step."""
|
||||
decode_use_gather: bool = False
|
||||
"""Whether the candidate consumers gather for the decode rows."""
|
||||
decode_gather: tuple[dict, torch.Tensor] | None = None
|
||||
"""The pool the first consumer resolved for the decode rows."""
|
||||
gather_launches: list[RocmMxfp4GatherLaunch] = field(default_factory=list)
|
||||
"""The prefill rows of the chunks that gather, as the consumers launch
|
||||
them."""
|
||||
|
||||
|
||||
def plan_prefill_chunks(
|
||||
chunks: list[DeepseekV32IndexerPrefillChunkMetadata],
|
||||
query_start_loc: list[int],
|
||||
seq_lens_cpu: torch.Tensor,
|
||||
context_lens: torch.Tensor,
|
||||
compress_ratio: int,
|
||||
candidate_block_size: int,
|
||||
min_gather_width: float | None,
|
||||
query_start_loc_device: torch.Tensor,
|
||||
) -> list[RocmMxfp4PrefillPlan]:
|
||||
"""Split each chunk into its requests.
|
||||
|
||||
``query_start_loc`` and ``seq_lens_cpu`` are the step's host copies (the
|
||||
latter an upper bound), ``context_lens`` the device compressed lengths.
|
||||
A chunk covers whole requests, or a query slice of one request; either
|
||||
way it launches once, packed through the chunk's part of the device
|
||||
``query_start_loc``.
|
||||
"""
|
||||
block = candidate_block_size
|
||||
plans = []
|
||||
for chunk in chunks:
|
||||
t0, t1 = chunk.token_start, chunk.token_end
|
||||
first = bisect.bisect_right(query_start_loc, t0) - 1
|
||||
requests = []
|
||||
for req in range(chunk.num_reqs):
|
||||
lo = max(query_start_loc[first + req], t0) - t0
|
||||
hi = min(query_start_loc[first + req + 1], t1) - t0
|
||||
if hi > lo:
|
||||
requests.append((lo, hi, req))
|
||||
reqs = slice(first, first + chunk.num_reqs)
|
||||
packed = torch.clamp(query_start_loc_device[first : reqs.stop + 1], t0, t1)
|
||||
packed -= t0
|
||||
width = int(seq_lens_cpu[reqs].max()) // compress_ratio
|
||||
row_ends = chunk.cu_seqlen_ke - chunk.cu_seqlen_ks
|
||||
plans.append(
|
||||
RocmMxfp4PrefillPlan(
|
||||
requests=requests,
|
||||
width=width,
|
||||
row_ends=row_ends,
|
||||
context_lens=context_lens[reqs],
|
||||
query_start_loc=packed,
|
||||
first_request=first,
|
||||
block_ends=(row_ends + block - 1) // block if block else None,
|
||||
use_gather=min_gather_width is not None and width >= min_gather_width,
|
||||
)
|
||||
)
|
||||
return plans
|
||||
|
||||
|
||||
def split_prefill_chunks(
|
||||
context_lens: torch.Tensor,
|
||||
query_lens: torch.Tensor,
|
||||
max_logits_bytes: int,
|
||||
request_offset: int = 0,
|
||||
) -> list[tuple[slice, slice]]:
|
||||
"""(request slice, query slice) chunks, as the base chunker returns them.
|
||||
|
||||
Nothing gathers K here, so the requests' contexts do not add up: a chunk's
|
||||
logits are [rows, widest compressed context] fp32, at most
|
||||
``max_logits_bytes``. A request too wide to launch whole is cut on its
|
||||
query rows.
|
||||
"""
|
||||
|
||||
def max_rows(width: int) -> int:
|
||||
return max(1, max_logits_bytes // (4 * max(width, 1)))
|
||||
|
||||
chunks: list[tuple[slice, slice]] = []
|
||||
end = 0
|
||||
while end < len(context_lens):
|
||||
start, rows, width = end, 0, 0
|
||||
while end < len(context_lens):
|
||||
q = int(query_lens[end])
|
||||
w = max(width, int(context_lens[end]))
|
||||
if rows and rows + q > max_rows(w):
|
||||
break
|
||||
rows, width = rows + q, w
|
||||
end += 1
|
||||
step = max_rows(width)
|
||||
reqs = slice(start + request_offset, end + request_offset)
|
||||
chunks.extend(
|
||||
(reqs, slice(lo, min(lo + step, rows))) for lo in range(0, rows, step)
|
||||
)
|
||||
return chunks
|
||||
|
||||
|
||||
def plan_gather_launches(
|
||||
chunks: list[DeepseekV32IndexerPrefillChunkMetadata],
|
||||
plans: list[RocmMxfp4PrefillPlan],
|
||||
max_rows: int,
|
||||
) -> list[RocmMxfp4GatherLaunch]:
|
||||
"""The candidate consumers' launches over the chunks that gather.
|
||||
|
||||
A consumer's logits are [rows, pool] however long the context, so the
|
||||
dense split, sized for the context, would launch it far more often than
|
||||
its memory needs. Each request's rows go back together across the chunks
|
||||
that cut them, then out again at most ``max_rows`` a launch.
|
||||
"""
|
||||
# Per request: (first token, [(plan, first row, end row)]), and the chunk
|
||||
# and request index its block table and context come from.
|
||||
runs: list[tuple[int, list, DeepseekV32IndexerPrefillChunkMetadata, int]] = []
|
||||
last: tuple[int, int] | None = None
|
||||
for chunk, plan in zip(chunks, plans):
|
||||
if not plan.use_gather:
|
||||
last = None
|
||||
continue
|
||||
for lo, hi, req in plan.requests:
|
||||
request, t_lo = plan.first_request + req, chunk.token_start + lo
|
||||
if last != (request, t_lo):
|
||||
runs.append((t_lo, [], chunk, req))
|
||||
runs[-1][1].append((plan, lo, hi))
|
||||
last = (request, chunk.token_start + hi)
|
||||
launches = []
|
||||
for token_start, pieces, chunk, req in runs:
|
||||
ends = [plan.row_ends[lo:hi] for plan, lo, hi in pieces]
|
||||
row_ends = ends[0] if len(ends) == 1 else torch.cat(ends)
|
||||
block_table = chunk.block_table[req : req + 1]
|
||||
context_len = pieces[0][0].context_lens[req : req + 1]
|
||||
for lo in range(0, row_ends.shape[0], max_rows):
|
||||
hi = min(lo + max_rows, row_ends.shape[0])
|
||||
launches.append(
|
||||
RocmMxfp4GatherLaunch(
|
||||
token_start + lo,
|
||||
token_start + hi,
|
||||
block_table,
|
||||
context_len,
|
||||
row_ends[lo:hi],
|
||||
)
|
||||
)
|
||||
return launches
|
||||
|
||||
|
||||
class DeepseekV41RocmMxfp4IndexerBackend(DeepseekV41IndexerBackend):
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "DEEPSEEK_V41_ROCM_MXFP4_INDEXER"
|
||||
|
||||
@staticmethod
|
||||
def get_supported_kernel_block_sizes(
|
||||
kv_cache_spec: KVCacheSpec | None = None,
|
||||
) -> list[int | MultipleOf]:
|
||||
# --block-size 128 preferred.
|
||||
return [64, 128]
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type["DeepseekV41RocmMxfp4IndexerMetadataBuilder"]:
|
||||
return DeepseekV41RocmMxfp4IndexerMetadataBuilder
|
||||
|
||||
|
||||
class RocmMxfp4IndexerMetadataBuilder(DeepseekV32IndexerMetadataBuilder):
|
||||
"""Dense indexer metadata plus the in-place paged launches' layout.
|
||||
|
||||
Requirements are checked here, at engine start.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def get_cudagraph_support(
|
||||
cls,
|
||||
vllm_config: VllmConfig,
|
||||
kv_cache_spec: KVCacheSpec,
|
||||
) -> AttentionCGSupport:
|
||||
# Decode rows are flattened into persistent buffers, and whether a
|
||||
# step launches them as next_n-row sequences depends on its shape
|
||||
# alone (see native_decode).
|
||||
return AttentionCGSupport.ALWAYS
|
||||
|
||||
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)
|
||||
name = "the ROCm MXFP4 indexer"
|
||||
if self.dcp_world_size > 1 or self.use_pcp:
|
||||
raise NotImplementedError(
|
||||
f"{name} does not support decode/prefill context parallel."
|
||||
)
|
||||
hf_config = vllm_config.model_config.hf_text_config
|
||||
num_heads, head_dim = hf_config.index_n_heads, hf_config.index_head_dim
|
||||
self.num_heads, self.head_dim = num_heads, head_dim
|
||||
self.page_entries = kv_cache_spec.block_size // self.compress_ratio
|
||||
self.cache_layout = rocm_paged_mxfp4_cache_layout(
|
||||
num_heads, head_dim, self.page_entries
|
||||
)
|
||||
# Ragged steps take a row per query token, as do the candidate
|
||||
# consumers' gathers. Uniform steps also launch the dense layers on
|
||||
# next_n-row sequences, see native_decode.
|
||||
self.use_flattening = True
|
||||
# A two-level (candidate block) indexer's builder sets these.
|
||||
self.candidate_block_size = 0
|
||||
self.min_gather_width: float | None = None
|
||||
|
||||
self.decode_block_ends_buffer = torch.zeros(
|
||||
self.arange_buffer.shape[0], dtype=torch.int32, device=device
|
||||
)
|
||||
self.context_lens_buffer = torch.zeros(
|
||||
self.arange_buffer.shape[0], dtype=torch.int32, device=device
|
||||
)
|
||||
spec_config = vllm_config.speculative_config
|
||||
# Adaptive verification replays a FULL graph on drafts reallocated
|
||||
# across requests, so a step's shape does not fix its query lengths.
|
||||
adaptive = spec_config is not None and spec_config.enable_adaptive_verification
|
||||
self.native_next_n = 1 if adaptive else self.num_speculative_tokens + 1
|
||||
# The indexer's logits width, which the decode schedule sizes slices by.
|
||||
max_model_len = vllm_config.model_config.max_model_len
|
||||
self.logits_width = max_model_len // self.compress_ratio
|
||||
# Every dense layer of this group reads the same rows through the same
|
||||
# page geometry, so the step builds their schedule once, not per layer.
|
||||
self.decode_schedule_buffer = torch.empty(
|
||||
rocm_mxfp4_decode_schedule_words(
|
||||
num_heads, head_dim, self.page_entries, self.native_next_n
|
||||
),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
def _split_indexer_prefill_chunks( # type: ignore[override]
|
||||
self,
|
||||
compressed_seq_lens_cpu: torch.Tensor,
|
||||
prefill_query_lens_cpu: torch.Tensor,
|
||||
workspace_size: int,
|
||||
max_logits_bytes: int,
|
||||
request_offset: int = 0,
|
||||
) -> list[tuple[slice, slice]]:
|
||||
return split_prefill_chunks(
|
||||
compressed_seq_lens_cpu,
|
||||
prefill_query_lens_cpu,
|
||||
max_logits_bytes,
|
||||
request_offset,
|
||||
)
|
||||
|
||||
def build(
|
||||
self,
|
||||
common_prefix_len: int,
|
||||
common_attn_metadata: CommonAttentionMetadata,
|
||||
fast_build: bool = False,
|
||||
) -> RocmMxfp4IndexerMetadata:
|
||||
base = super().build(common_prefix_len, common_attn_metadata, fast_build)
|
||||
metadata = RocmMxfp4IndexerMetadata(
|
||||
**{f.name: getattr(base, f.name) for f in fields(base)}
|
||||
)
|
||||
cm = common_attn_metadata
|
||||
context_lens = self.context_lens_buffer[: cm.num_reqs]
|
||||
torch.floor_divide(
|
||||
cm.seq_lens[: cm.num_reqs], self.compress_ratio, out=context_lens
|
||||
)
|
||||
if base.decode is not None:
|
||||
lengths = base.decode.seq_lens.view(-1)
|
||||
metadata.decode_row_lens = lengths
|
||||
block = self.candidate_block_size
|
||||
if block:
|
||||
ends = self.decode_block_ends_buffer[: lengths.shape[0]]
|
||||
torch.add(lengths, block - 1, out=ends)
|
||||
ends.floor_divide_(block)
|
||||
metadata.decode_block_ends = ends
|
||||
query_start_loc = cm.query_start_loc_cpu[: base.num_decodes + 1]
|
||||
metadata.decode_native = native_decode(
|
||||
lengths,
|
||||
base.decode.block_table,
|
||||
torch.diff(query_start_loc).tolist(),
|
||||
self.native_next_n,
|
||||
context_lens,
|
||||
)
|
||||
metadata.decode_schedule = build_rocm_mxfp4_decode_schedule(
|
||||
lengths,
|
||||
self.num_heads,
|
||||
self.head_dim,
|
||||
self.page_entries,
|
||||
self.decode_schedule_buffer,
|
||||
self.logits_width,
|
||||
metadata.decode_native,
|
||||
)
|
||||
if base.prefill is not None:
|
||||
assert cm.seq_lens_cpu_upper_bound is not None
|
||||
metadata.prefill_plans = plan_prefill_chunks(
|
||||
base.prefill.chunks,
|
||||
cm.query_start_loc_cpu.tolist(),
|
||||
cm.seq_lens_cpu_upper_bound,
|
||||
context_lens,
|
||||
self.compress_ratio,
|
||||
self.candidate_block_size,
|
||||
self.min_gather_width,
|
||||
cm.query_start_loc,
|
||||
)
|
||||
return metadata
|
||||
|
||||
|
||||
class DeepseekV41RocmMxfp4IndexerMetadataBuilder(RocmMxfp4IndexerMetadataBuilder):
|
||||
"""DeepSeek-V4.1's two-level indexer: sets the candidate block size the
|
||||
base plans with and, with ``indexer_sparse_logits``, adds the consumers'
|
||||
gathers."""
|
||||
|
||||
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
|
||||
self.candidate_block_size = getattr(hf_config, "candidate_block_size", 0)
|
||||
self.num_candidate_cols = 0
|
||||
if vllm_config.attention_config.indexer_sparse_logits:
|
||||
prefix = "attention_config.indexer_sparse_logits"
|
||||
block = self.candidate_block_size
|
||||
topk_blocks = getattr(hf_config, "candidate_topk_blocks", 0)
|
||||
layout = self.cache_layout
|
||||
if block <= 0 or topk_blocks <= 0:
|
||||
raise ValueError(
|
||||
f"{prefix} requires a two-level (candidate block) indexer "
|
||||
"model such as DeepSeek-V4.1."
|
||||
)
|
||||
if topk_blocks & (topk_blocks - 1):
|
||||
raise ValueError(
|
||||
f"{prefix} on ROCm resolves the candidate pool with a sort "
|
||||
f"that needs a power-of-two block count, not {topk_blocks}."
|
||||
)
|
||||
if self.page_entries % block or block > layout.n_per_tile:
|
||||
raise ValueError(
|
||||
f"{prefix} on ROCm needs {block}-entry candidate blocks to "
|
||||
f"tile a {self.page_entries}-entry indexer page and fit in "
|
||||
f"one {layout.n_per_tile}-entry shuffle group."
|
||||
)
|
||||
self.num_candidate_cols = topk_blocks * block
|
||||
self.min_gather_width = self.num_candidate_cols * _GATHER_MIN_CUT_PREFILL
|
||||
self.gather_rows = rocm_mxfp4_consumer_rows(self.num_candidate_cols)
|
||||
logger.info_once(
|
||||
"DeepSeek V4.1 indexer: candidate consumers score the %d "
|
||||
"candidate positions with aiter's paged MXFP4 gather once the "
|
||||
"context is %.1fx (decode) / %.1fx (prefill) as long.",
|
||||
self.num_candidate_cols,
|
||||
_GATHER_MIN_CUT_DECODE,
|
||||
_GATHER_MIN_CUT_PREFILL,
|
||||
)
|
||||
|
||||
def _gather_pays(self, context: int, cut: float) -> bool:
|
||||
return self.num_candidate_cols > 0 and context >= cut * self.num_candidate_cols
|
||||
|
||||
def build(
|
||||
self,
|
||||
common_prefix_len: int,
|
||||
common_attn_metadata: CommonAttentionMetadata,
|
||||
fast_build: bool = False,
|
||||
) -> RocmMxfp4IndexerMetadata:
|
||||
metadata = super().build(common_prefix_len, common_attn_metadata, fast_build)
|
||||
if metadata.decode is not None:
|
||||
metadata.decode_use_gather = self._gather_pays(
|
||||
metadata.max_seq_len // self.compress_ratio, _GATHER_MIN_CUT_DECODE
|
||||
)
|
||||
if metadata.prefill is not None and self.num_candidate_cols:
|
||||
metadata.gather_launches = plan_gather_launches(
|
||||
metadata.prefill.chunks, metadata.prefill_plans, self.gather_rows
|
||||
)
|
||||
return metadata
|
||||
@@ -0,0 +1,676 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Sparse attention indexer on aiter's paged MXFP4 MQA-logits kernel.
|
||||
|
||||
gfx950 only. The kernel reads the preshuffled paged indexer K cache in place,
|
||||
so no layer gathers K into a contiguous buffer, prefill included. Any model
|
||||
whose indexer supplies the inputs `rocm_mxfp4_sparse_attn_indexer` documents
|
||||
can run the dense path. DeepSeek-V4.1's two-level indexer adds to it: the
|
||||
candidate source takes its block maxima from the same walk that writes its
|
||||
logits, and the candidate consumers can walk the candidate pool instead of
|
||||
the whole context: the pool is resolved once per step and shared by all of
|
||||
them.
|
||||
"""
|
||||
|
||||
import functools
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, NamedTuple
|
||||
|
||||
import torch
|
||||
|
||||
import vllm.envs as envs
|
||||
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
|
||||
from vllm.config import CUDAGraphMode
|
||||
from vllm.forward_context import get_forward_context
|
||||
from vllm.model_executor.kernels.attention.dsa.candidate_blocks import (
|
||||
apply_candidate_mask,
|
||||
)
|
||||
from vllm.triton_utils import tl, triton
|
||||
from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
|
||||
_apply_candidate_mask_strided,
|
||||
_max_decode_logits_rows,
|
||||
)
|
||||
from vllm.v1.worker.workspace import current_workspace_manager
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.v1.attention.backends.mla.indexer import (
|
||||
DeepseekV32IndexerPrefillChunkMetadata,
|
||||
)
|
||||
from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import (
|
||||
RocmMxfp4GatherLaunch,
|
||||
RocmMxfp4IndexerMetadata,
|
||||
RocmMxfp4NativeDecode,
|
||||
RocmMxfp4PrefillPlan,
|
||||
)
|
||||
|
||||
MXFP4_BLOCK_SIZE = 32
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _aiter():
|
||||
"""The aiter module with the paged MXFP4 MQA-logits launcher, its schedule
|
||||
and cache_format."""
|
||||
from aiter.ops.triton.attention import pa_mqa_logits_mxfp4
|
||||
|
||||
return pa_mqa_logits_mxfp4
|
||||
|
||||
|
||||
@functools.cache
|
||||
def _aiter_topk() -> Callable[..., None]:
|
||||
"""The aiter per-row top-k."""
|
||||
from aiter.ops.topk import top_k_per_row_decode
|
||||
|
||||
return top_k_per_row_decode
|
||||
|
||||
|
||||
class RocmPagedMxfp4CacheLayout(NamedTuple):
|
||||
"""The byte order of an indexer K page, as the kernel reads it, and the
|
||||
shuffle pattern aiter's K cache op writes it with.
|
||||
|
||||
A page is cut into runs of ``n_per_tile`` tokens. Inside a run, each token's
|
||||
packed values are split into ``d_per_tile``-byte chunks along the head
|
||||
dimension, and the run stores chunk 0 of every token, then chunk 1, and so
|
||||
on: ``[chunk, token, byte]``. Its e8m0 scales are stored as
|
||||
``[scale % scale_lanes, token, scale // scale_lanes]``.
|
||||
"""
|
||||
|
||||
n_per_tile: int
|
||||
d_per_tile: int
|
||||
scale_lanes: int
|
||||
|
||||
|
||||
# The scale order the kernel reads: aiter's mode 1, a lane's scales innermost.
|
||||
# cache_format does not report the order or the lane split yet.
|
||||
_SCALE_ORDER = 1
|
||||
_WAVE_SIZE = 64
|
||||
|
||||
|
||||
@functools.cache
|
||||
def rocm_paged_mxfp4_cache_layout(
|
||||
num_heads: int, head_dim: int, page_entries: int
|
||||
) -> RocmPagedMxfp4CacheLayout:
|
||||
"""The K page layout from aiter's ``cache_format``, the only place vLLM
|
||||
reads it, so a change there cannot reach the writer unchecked.
|
||||
|
||||
Raises:
|
||||
ValueError: aiter rejects the page size, or describes an order this
|
||||
layout cannot express.
|
||||
|
||||
"""
|
||||
fmt = _aiter().cache_format(num_heads, head_dim, page_entries)
|
||||
try:
|
||||
n_per_tile = int(fmt["n_per_tile"])
|
||||
d_per_tile = int(fmt["d_per_tile"])
|
||||
scale_order = fmt.get("scale_mode", _SCALE_ORDER)
|
||||
scale_lanes = int(fmt.get("scale_lanes", _WAVE_SIZE // n_per_tile))
|
||||
except (KeyError, TypeError, ValueError, ZeroDivisionError) as e:
|
||||
raise ValueError(f"aiter's cache_format() returned {fmt!r}") from e
|
||||
if (
|
||||
scale_order != _SCALE_ORDER
|
||||
or min(n_per_tile, d_per_tile, scale_lanes) <= 0
|
||||
or page_entries % n_per_tile
|
||||
or (head_dim // 2) % d_per_tile
|
||||
or (head_dim // MXFP4_BLOCK_SIZE) % scale_lanes
|
||||
):
|
||||
raise ValueError(
|
||||
f"aiter's cache_format() returned {fmt!r} for {num_heads} heads of "
|
||||
f"{head_dim} in {page_entries}-entry pages, a K page order "
|
||||
"RocmPagedMxfp4CacheLayout cannot express"
|
||||
)
|
||||
return RocmPagedMxfp4CacheLayout(n_per_tile, d_per_tile, scale_lanes)
|
||||
|
||||
|
||||
def rocm_mxfp4_decode_schedule_words(
|
||||
num_heads: int, head_dim: int, page_entries: int, next_n: int = 1
|
||||
) -> int:
|
||||
"""int32 words of the largest schedule a decode step can build, flattened or
|
||||
not. The slice cap can take it past target_wgs, up to SCHED_SLOT_CAP."""
|
||||
words = 4 * _aiter().SCHED_SLOT_CAP
|
||||
for rows in {1, next_n}:
|
||||
config = _aiter().select_config(num_heads, head_dim, rows, page_entries)
|
||||
words = max(words, 4 * config["target_wgs"])
|
||||
return words
|
||||
|
||||
|
||||
def build_rocm_mxfp4_decode_schedule(
|
||||
row_lens: torch.Tensor,
|
||||
num_heads: int,
|
||||
head_dim: int,
|
||||
page_entries: int,
|
||||
out: torch.Tensor,
|
||||
logits_width: int,
|
||||
native: "RocmMxfp4NativeDecode | None" = None,
|
||||
) -> torch.Tensor | None:
|
||||
"""Work descriptors that even out a decode step, or None where the static
|
||||
grid already fills the machine. Depends only on the rows' lengths, the
|
||||
cache geometry and the logits width, so one serves every layer of a group.
|
||||
The width sizes the slices, capped the way the static grid caps them."""
|
||||
if native is None:
|
||||
return _aiter().build_schedule(
|
||||
row_lens,
|
||||
1,
|
||||
num_heads,
|
||||
head_dim,
|
||||
page_entries,
|
||||
out=out,
|
||||
max_model_len=logits_width,
|
||||
)
|
||||
return _aiter().build_schedule(
|
||||
native.context_lens,
|
||||
native.next_n,
|
||||
num_heads,
|
||||
head_dim,
|
||||
page_entries,
|
||||
out=out,
|
||||
row_ends=row_lens,
|
||||
max_model_len=logits_width,
|
||||
)
|
||||
|
||||
|
||||
def _kv_view(kv_cache: torch.Tensor, head_dim: int) -> torch.Tensor:
|
||||
"""The indexer cache as [pages, entries, 1, bytes]. The page stride is the
|
||||
block-major pool's, not the page's own size."""
|
||||
num_pages, entries, width = kv_cache.shape
|
||||
assert kv_cache.dtype == torch.uint8
|
||||
assert width == head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE, width
|
||||
return torch.as_strided(
|
||||
kv_cache, (num_pages, entries, 1, width), (kv_cache.stride(0), width, width, 1)
|
||||
)
|
||||
|
||||
|
||||
def rocm_mxfp4_consumer_rows(num_candidate_cols: int) -> int:
|
||||
"""Query rows per candidate-consumer launch. Its logits are [rows, pool]
|
||||
fp32 whatever the context, so the logits budget alone sizes it."""
|
||||
budget = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
|
||||
return max(1, budget // (4 * num_candidate_cols))
|
||||
|
||||
|
||||
def reserve_rocm_mxfp4_indexer_workspace(
|
||||
hidden_states: torch.Tensor,
|
||||
logits_width: int,
|
||||
candidate_block_size: int = 0,
|
||||
gather_block_size: int = 0,
|
||||
num_candidate_cols: int = 0,
|
||||
) -> None:
|
||||
"""Profiling run: claim the decode logits workspace and the peak prefill
|
||||
logits, block scores included when the layer writes them, candidate lists
|
||||
when it gathers."""
|
||||
rows = _max_decode_logits_rows(hidden_states.shape[0])
|
||||
specs = [((rows, logits_width), torch.float32)]
|
||||
budget = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
|
||||
if candidate_block_size:
|
||||
nblocks = triton.cdiv(logits_width, candidate_block_size)
|
||||
specs.append(((rows, nblocks), torch.float32))
|
||||
budget += budget // candidate_block_size
|
||||
if gather_block_size:
|
||||
# aiter allocates the candidate lists: each gathered row, 8 B per
|
||||
# candidate block (int32 slot and position). The first consumer builds
|
||||
# them and they live until the step ends, across the later layers, so
|
||||
# one reservation is held for the rest of the forward.
|
||||
held = get_forward_context().additional_kwargs
|
||||
if "rocm_mxfp4_candidate_lists" not in held:
|
||||
held["rocm_mxfp4_candidate_lists"] = torch.empty(
|
||||
hidden_states.shape[0] * (num_candidate_cols // gather_block_size) * 8,
|
||||
dtype=torch.uint8,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
current_workspace_manager().get_simultaneous(*specs)
|
||||
torch.empty(budget, dtype=torch.uint8, device=hidden_states.device)
|
||||
|
||||
|
||||
def _topk(
|
||||
logits: torch.Tensor, lengths: torch.Tensor, out: torch.Tensor, k: int
|
||||
) -> None:
|
||||
"""Top-k over each row's [0, lengths[row]); -1 pads rows shorter than k.
|
||||
|
||||
aiter's kernel is faster on the candidate top-k and from 128 rows on; below
|
||||
that, a wide decode step is faster on vLLM's. Shape alone decides, so a FULL
|
||||
graph replays the same choice at every context length.
|
||||
"""
|
||||
rows, width = logits.shape
|
||||
if rows == 0:
|
||||
return
|
||||
if width >= k and (k >= 2048 or rows >= 128):
|
||||
_aiter_topk()(
|
||||
logits,
|
||||
1,
|
||||
lengths.reshape(-1),
|
||||
out,
|
||||
rows,
|
||||
logits.stride(0),
|
||||
logits.stride(1),
|
||||
k=k,
|
||||
)
|
||||
else:
|
||||
torch.ops._C.top_k_per_row_decode(
|
||||
logits, 1, lengths, out, rows, logits.stride(0), logits.stride(1), k
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _remap_compact_topk_kernel(
|
||||
idx_ptr,
|
||||
idx_stride,
|
||||
pos_ptr,
|
||||
pos_stride,
|
||||
K: tl.constexpr,
|
||||
BLOCK: tl.constexpr,
|
||||
PADDED_K: tl.constexpr,
|
||||
):
|
||||
row = tl.program_id(0).to(tl.int64)
|
||||
cols = tl.arange(0, PADDED_K)
|
||||
live = cols < K
|
||||
slot = tl.load(idx_ptr + row * idx_stride + cols, mask=live, other=-1)
|
||||
hit = slot >= 0
|
||||
slot = tl.where(hit, slot, 0)
|
||||
start = tl.load(
|
||||
pos_ptr + row * pos_stride + slot // BLOCK, mask=live & hit, other=0
|
||||
)
|
||||
tl.store(
|
||||
idx_ptr + row * idx_stride + cols,
|
||||
tl.where(hit, start + slot % BLOCK, -1).to(tl.int32),
|
||||
mask=live,
|
||||
)
|
||||
|
||||
|
||||
def _remap_compact_topk(
|
||||
indices: torch.Tensor, positions: torch.Tensor, block: int
|
||||
) -> None:
|
||||
"""Candidate slots to context positions, in place: slot j sits in pool
|
||||
block j // block, which starts at positions[j // block]."""
|
||||
rows, k = indices.shape
|
||||
if rows == 0:
|
||||
return
|
||||
_remap_compact_topk_kernel[(rows,)](
|
||||
indices,
|
||||
indices.stride(0),
|
||||
positions,
|
||||
positions.stride(0),
|
||||
K=k,
|
||||
BLOCK=block,
|
||||
PADDED_K=triton.next_power_of_2(k),
|
||||
num_warps=4,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Layer:
|
||||
"""One indexer layer's inputs and outputs for this step."""
|
||||
|
||||
metadata: "RocmMxfp4IndexerMetadata"
|
||||
kv: torch.Tensor
|
||||
q: torch.Tensor
|
||||
q_scale: torch.Tensor
|
||||
weights: torch.Tensor
|
||||
topk_buffer: torch.Tensor
|
||||
topk_tokens: int
|
||||
compress_ratio: int
|
||||
candidates: torch.Tensor | None
|
||||
block: int
|
||||
full_graph: bool
|
||||
|
||||
@property
|
||||
def num_heads(self) -> int:
|
||||
return self.q.shape[1]
|
||||
|
||||
@property
|
||||
def head_dim(self) -> int:
|
||||
return self.q.shape[2] * 2
|
||||
|
||||
def rows(self, lo: int, hi: int, seqs: int = 1) -> tuple[torch.Tensor, ...]:
|
||||
"""Q, its scales and the weights of token rows [lo, hi), as ``seqs``
|
||||
sequences of next_n rows each."""
|
||||
return (
|
||||
self.q[lo:hi].view(seqs, -1, *self.q.shape[1:]),
|
||||
self.q_scale[lo:hi].view(seqs, -1, *self.q_scale.shape[1:]),
|
||||
self.weights[lo:hi],
|
||||
)
|
||||
|
||||
def decode_rows(self, n: int) -> tuple[torch.Tensor, ...]:
|
||||
"""The first n token rows, one sequence each."""
|
||||
return (
|
||||
self.q[:n].unsqueeze(1),
|
||||
self.q_scale[:n].unsqueeze(1),
|
||||
self.weights[:n],
|
||||
)
|
||||
|
||||
|
||||
def _layer(
|
||||
k_cache_prefix: str,
|
||||
kv_cache: torch.Tensor,
|
||||
q_values: torch.Tensor,
|
||||
q_scale: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
topk_indices_buffer: torch.Tensor,
|
||||
topk_tokens: int,
|
||||
head_dim: int,
|
||||
compress_ratio: int,
|
||||
candidate_blocks: torch.Tensor | None,
|
||||
candidate_block_size: int,
|
||||
) -> _Layer:
|
||||
from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import (
|
||||
RocmMxfp4IndexerMetadata,
|
||||
)
|
||||
|
||||
forward_context = get_forward_context()
|
||||
attn_metadata = forward_context.attn_metadata
|
||||
assert isinstance(attn_metadata, dict)
|
||||
metadata = attn_metadata[k_cache_prefix]
|
||||
assert isinstance(metadata, RocmMxfp4IndexerMetadata), (
|
||||
"the ROCm MXFP4 indexer needs RocmMxfp4IndexerMetadataBuilder metadata"
|
||||
)
|
||||
kv = _kv_view(kv_cache, head_dim)
|
||||
num_heads = q_values.shape[1]
|
||||
return _Layer(
|
||||
metadata=metadata,
|
||||
kv=kv,
|
||||
q=q_values,
|
||||
# The fused Q kernel returns one int32 per head; the kernel reads the
|
||||
# four ue8m0 bytes behind it.
|
||||
q_scale=q_scale.view(torch.uint8).view(q_scale.shape[0], num_heads, -1),
|
||||
weights=weights,
|
||||
topk_buffer=topk_indices_buffer,
|
||||
topk_tokens=topk_tokens,
|
||||
compress_ratio=compress_ratio,
|
||||
candidates=candidate_blocks,
|
||||
block=candidate_block_size,
|
||||
full_graph=forward_context.cudagraph_runtime_mode == CUDAGraphMode.FULL,
|
||||
)
|
||||
|
||||
|
||||
def _block_scores(layer: _Layer, scores: torch.Tensor | None) -> dict:
|
||||
if scores is None:
|
||||
return {}
|
||||
return dict(
|
||||
calc_block_scores=True, block_scores=scores, candidate_block_size=layer.block
|
||||
)
|
||||
|
||||
|
||||
def _dense_prefill(
|
||||
layer: _Layer,
|
||||
chunk: "DeepseekV32IndexerPrefillChunkMetadata",
|
||||
plan: "RocmMxfp4PrefillPlan",
|
||||
candidate_write: bool,
|
||||
) -> None:
|
||||
pa = _aiter()
|
||||
t0, t1 = chunk.token_start, chunk.token_end
|
||||
logits = layer.q.new_empty((t1 - t0, plan.width), dtype=torch.float32)
|
||||
scores = None
|
||||
if candidate_write:
|
||||
nblocks = triton.cdiv(plan.width, layer.block)
|
||||
scores = layer.q.new_empty((t1 - t0, nblocks), dtype=torch.float32)
|
||||
# One launch for the chunk: the kernel finds each row's request from
|
||||
# query_start_loc on its own grid.
|
||||
pa.paged_mxfp4_mqa_logits(
|
||||
layer.q[t0:t1],
|
||||
layer.q_scale[t0:t1],
|
||||
layer.kv,
|
||||
layer.weights[t0:t1],
|
||||
plan.context_lens,
|
||||
chunk.block_table,
|
||||
plan.width,
|
||||
out_logits=logits,
|
||||
clean_logits=False,
|
||||
row_ends=plan.row_ends,
|
||||
query_start_loc=plan.query_start_loc,
|
||||
**_block_scores(layer, scores),
|
||||
)
|
||||
candidates = None if layer.candidates is None else layer.candidates[t0:t1]
|
||||
if scores is not None:
|
||||
assert candidates is not None and plan.block_ends is not None
|
||||
_topk(scores, plan.block_ends, candidates, candidates.shape[1])
|
||||
elif candidates is not None:
|
||||
apply_candidate_mask(logits, None, plan.row_ends, candidates, layer.block)
|
||||
_topk(
|
||||
logits,
|
||||
plan.row_ends,
|
||||
layer.topk_buffer[t0:t1, : layer.topk_tokens],
|
||||
layer.topk_tokens,
|
||||
)
|
||||
|
||||
|
||||
def _dense_decode(layer: _Layer, logits_width: int, candidate_write: bool) -> None:
|
||||
metadata = layer.metadata
|
||||
lengths = metadata.decode_row_lens
|
||||
assert metadata.decode is not None and lengths is not None
|
||||
rows = lengths.shape[0]
|
||||
specs = [((rows, logits_width), torch.float32)]
|
||||
if candidate_write:
|
||||
specs.append(((rows, triton.cdiv(logits_width, layer.block)), torch.float32))
|
||||
logits, *scores = current_workspace_manager().get_simultaneous(*specs)
|
||||
native = metadata.decode_native
|
||||
if native is not None:
|
||||
# A request's next_n rows go in as one sequence, so a workgroup walks
|
||||
# each KV tile once for all of them.
|
||||
q, q_scale, weights = layer.rows(0, rows, native.context_lens.shape[0])
|
||||
context_lens, block_table, row_ends = (
|
||||
native.context_lens,
|
||||
native.block_table,
|
||||
lengths,
|
||||
)
|
||||
else:
|
||||
q, q_scale, weights = layer.decode_rows(rows)
|
||||
context_lens, block_table, row_ends = lengths, metadata.decode.block_table, None
|
||||
assert block_table.shape[0] == rows, "one block-table row per query row"
|
||||
_aiter().paged_mxfp4_mqa_logits(
|
||||
q,
|
||||
q_scale,
|
||||
layer.kv,
|
||||
weights,
|
||||
context_lens,
|
||||
block_table,
|
||||
logits_width,
|
||||
out_logits=logits,
|
||||
clean_logits=False,
|
||||
row_ends=row_ends,
|
||||
schedule=metadata.decode_schedule,
|
||||
**_block_scores(layer, scores[0] if scores else None),
|
||||
)
|
||||
candidates = None if layer.candidates is None else layer.candidates[:rows]
|
||||
if scores:
|
||||
assert candidates is not None and metadata.decode_block_ends is not None
|
||||
_topk(scores[0], metadata.decode_block_ends, candidates, candidates.shape[1])
|
||||
elif candidates is not None:
|
||||
_apply_candidate_mask_strided(logits, None, lengths, candidates, layer.block)
|
||||
_topk(
|
||||
logits,
|
||||
lengths,
|
||||
layer.topk_buffer[:rows, : layer.topk_tokens],
|
||||
layer.topk_tokens,
|
||||
)
|
||||
|
||||
|
||||
def _gather_prefill(
|
||||
layer: _Layer, launch: "RocmMxfp4GatherLaunch", num_cols: int
|
||||
) -> None:
|
||||
pa = _aiter()
|
||||
assert layer.candidates is not None
|
||||
t0, t1 = launch.token_start, launch.token_end
|
||||
if launch.pool is None:
|
||||
launch.pool = pa.build_candidate_gather(
|
||||
layer.candidates[t0:t1],
|
||||
launch.row_ends,
|
||||
launch.block_table.expand(t1 - t0, -1),
|
||||
layer.kv,
|
||||
layer.num_heads,
|
||||
layer.head_dim,
|
||||
layer.block,
|
||||
)
|
||||
gather, slot_ends = launch.pool
|
||||
compact = layer.q.new_empty((t1 - t0, num_cols), dtype=torch.float32)
|
||||
q, q_scale, weights = layer.rows(t0, t1)
|
||||
pa.paged_mxfp4_mqa_logits(
|
||||
q,
|
||||
q_scale,
|
||||
layer.kv,
|
||||
weights,
|
||||
launch.context_len,
|
||||
launch.block_table,
|
||||
num_cols,
|
||||
out_logits=compact,
|
||||
clean_logits=False,
|
||||
row_ends=slot_ends,
|
||||
use_gather=True,
|
||||
candidates=gather,
|
||||
)
|
||||
out = layer.topk_buffer[t0:t1, : layer.topk_tokens]
|
||||
_topk(compact, slot_ends, out, layer.topk_tokens)
|
||||
_remap_compact_topk(out, gather["positions"], layer.block)
|
||||
|
||||
|
||||
def _gather_decode(layer: _Layer, num_cols: int) -> None:
|
||||
pa = _aiter()
|
||||
assert layer.candidates is not None
|
||||
metadata = layer.metadata
|
||||
lengths = metadata.decode_row_lens
|
||||
assert metadata.decode is not None and lengths is not None
|
||||
block_table = metadata.decode.block_table
|
||||
rows = lengths.shape[0]
|
||||
if metadata.decode_gather is None:
|
||||
metadata.decode_gather = pa.build_candidate_gather(
|
||||
layer.candidates[:rows],
|
||||
lengths,
|
||||
block_table,
|
||||
layer.kv,
|
||||
layer.num_heads,
|
||||
layer.head_dim,
|
||||
layer.block,
|
||||
)
|
||||
gather, slot_ends = metadata.decode_gather
|
||||
(compact,) = current_workspace_manager().get_simultaneous(
|
||||
((rows, num_cols), torch.float32)
|
||||
)
|
||||
q, q_scale, weights = layer.decode_rows(rows)
|
||||
pa.paged_mxfp4_mqa_logits(
|
||||
q,
|
||||
q_scale,
|
||||
layer.kv,
|
||||
weights,
|
||||
lengths,
|
||||
block_table,
|
||||
num_cols,
|
||||
out_logits=compact,
|
||||
clean_logits=False,
|
||||
row_ends=slot_ends,
|
||||
use_gather=True,
|
||||
candidates=gather,
|
||||
)
|
||||
out = layer.topk_buffer[:rows, : layer.topk_tokens]
|
||||
_topk(compact, slot_ends, out, layer.topk_tokens)
|
||||
_remap_compact_topk(out, gather["positions"], layer.block)
|
||||
|
||||
|
||||
@eager_break_during_capture
|
||||
def rocm_mxfp4_sparse_attn_indexer(
|
||||
hidden_states: torch.Tensor,
|
||||
k_cache_prefix: str,
|
||||
kv_cache: torch.Tensor,
|
||||
q_values: torch.Tensor,
|
||||
q_scale: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
topk_tokens: int,
|
||||
head_dim: int,
|
||||
max_model_len: int,
|
||||
topk_indices_buffer: torch.Tensor,
|
||||
compress_ratio: int = 1,
|
||||
candidate_blocks: torch.Tensor | None = None,
|
||||
candidate_block_size: int = 0,
|
||||
candidate_write: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""Dense indexer: every layer that scores the whole context, the
|
||||
candidate source included. With ``candidate_blocks`` and not
|
||||
``candidate_write`` the scores are masked to the pool first, as the
|
||||
shared path does.
|
||||
|
||||
The model supplies ``q_values`` [T, H, D // 2] packed e2m1, ``q_scale``
|
||||
holding each head's D // 32 ue8m0 bytes (one int32 per head at D = 128),
|
||||
fp32 ``weights`` with the softmax and head scales folded in, and
|
||||
``kv_cache`` pages in `rocm_paged_mxfp4_cache_layout` order for H heads.
|
||||
"""
|
||||
if not isinstance(get_forward_context().attn_metadata, dict):
|
||||
reserve_rocm_mxfp4_indexer_workspace(
|
||||
hidden_states,
|
||||
max_model_len,
|
||||
candidate_block_size if candidate_write else 0,
|
||||
)
|
||||
return topk_indices_buffer
|
||||
layer = _layer(
|
||||
k_cache_prefix,
|
||||
kv_cache,
|
||||
q_values,
|
||||
q_scale,
|
||||
weights,
|
||||
topk_indices_buffer,
|
||||
topk_tokens,
|
||||
head_dim,
|
||||
compress_ratio,
|
||||
candidate_blocks,
|
||||
candidate_block_size,
|
||||
)
|
||||
metadata = layer.metadata
|
||||
topk_indices_buffer[: hidden_states.shape[0]] = -1
|
||||
if metadata.prefill is not None:
|
||||
for chunk, plan in zip(metadata.prefill.chunks, metadata.prefill_plans):
|
||||
_dense_prefill(layer, chunk, plan, candidate_write)
|
||||
if metadata.decode is not None:
|
||||
_dense_decode(layer, max_model_len, candidate_write)
|
||||
return topk_indices_buffer
|
||||
|
||||
|
||||
@eager_break_during_capture
|
||||
def rocm_mxfp4_sparse_mqa_indexer(
|
||||
hidden_states: torch.Tensor,
|
||||
k_cache_prefix: str,
|
||||
kv_cache: torch.Tensor,
|
||||
q_values: torch.Tensor,
|
||||
q_scale: torch.Tensor,
|
||||
weights: torch.Tensor,
|
||||
topk_tokens: int,
|
||||
head_dim: int,
|
||||
max_model_len: int,
|
||||
topk_indices_buffer: torch.Tensor,
|
||||
compress_ratio: int,
|
||||
candidate_blocks: torch.Tensor,
|
||||
candidate_block_size: int,
|
||||
num_candidate_cols: int,
|
||||
) -> torch.Tensor:
|
||||
"""Candidate consumer: score only the source's pool where the builder's
|
||||
length gate says it pays, else the dense walk masked to the pool."""
|
||||
if not isinstance(get_forward_context().attn_metadata, dict):
|
||||
reserve_rocm_mxfp4_indexer_workspace(
|
||||
hidden_states,
|
||||
max(max_model_len, num_candidate_cols),
|
||||
gather_block_size=candidate_block_size,
|
||||
num_candidate_cols=num_candidate_cols,
|
||||
)
|
||||
return topk_indices_buffer
|
||||
layer = _layer(
|
||||
k_cache_prefix,
|
||||
kv_cache,
|
||||
q_values,
|
||||
q_scale,
|
||||
weights,
|
||||
topk_indices_buffer,
|
||||
topk_tokens,
|
||||
head_dim,
|
||||
compress_ratio,
|
||||
candidate_blocks,
|
||||
candidate_block_size,
|
||||
)
|
||||
metadata = layer.metadata
|
||||
topk_indices_buffer[: hidden_states.shape[0]] = -1
|
||||
if metadata.prefill is not None:
|
||||
for chunk, plan in zip(metadata.prefill.chunks, metadata.prefill_plans):
|
||||
if not plan.use_gather:
|
||||
_dense_prefill(layer, chunk, plan, candidate_write=False)
|
||||
for launch in metadata.gather_launches:
|
||||
_gather_prefill(layer, launch, num_candidate_cols)
|
||||
if metadata.decode is not None:
|
||||
# A FULL graph cannot follow the per-step gate; the gather is the side
|
||||
# that stays flat in the context length.
|
||||
if metadata.decode_use_gather or layer.full_graph:
|
||||
_gather_decode(layer, num_candidate_cols)
|
||||
else:
|
||||
_dense_decode(layer, max_model_len, candidate_write=False)
|
||||
return topk_indices_buffer
|
||||
Reference in New Issue
Block a user