| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """DeepSeek V4.1's ROCm MXFP4 sparse indexer against a torch reference. |
| 4 | |
| 5 | The indexer K cache is a view into vLLM's block-major KV pool: its pages are |
| 6 | a block stride apart with other layers' pages in between, so the tests put |
| 7 | it there. The reference reads the keys back from a second cache the same |
| 8 | writer fills in natural order, so it shares none of the preshuffled |
| 9 | addressing it checks. |
| 10 | """ |
| 11 | |
| 12 | import itertools |
| 13 | import types |
| 14 | |
| 15 | import pytest |
| 16 | import torch |
| 17 | |
| 18 | from vllm.platforms import current_platform |
| 19 | |
| 20 | if not current_platform.is_rocm(): |
| 21 | pytest.skip("ROCm-only", allow_module_level=True) |
| 22 | |
| 23 | from vllm.config import CUDAGraphMode |
| 24 | from vllm.models.deepseek_v41.amd.rocm import ( |
| 25 | rocm_mxfp4_indexer_k_store, |
| 26 | rocm_mxfp4_indexer_q_quant, |
| 27 | ) |
| 28 | from vllm.platforms.rocm import get_cdna_version |
| 29 | from vllm.utils.math_utils import cdiv |
| 30 | from vllm.utils.torch_utils import set_random_seed |
| 31 | from vllm.v1.attention.backends.mla.indexer import ( |
| 32 | DeepseekV32IndexerPrefillChunkMetadata, |
| 33 | ) |
| 34 | from vllm.v1.attention.backends.mla.rocm_paged_mxfp4_indexer import ( |
| 35 | RocmMxfp4IndexerMetadata, |
| 36 | native_decode, |
| 37 | plan_gather_launches, |
| 38 | plan_prefill_chunks, |
| 39 | ) |
| 40 | from vllm.v1.attention.ops import rocm_paged_mxfp4_indexer as ops |
| 41 | from vllm.v1.worker.workspace import init_workspace_manager, reset_workspace_manager |
| 42 | |
| 43 | if get_cdna_version() != 4: |
| 44 | pytest.skip("CDNA4-only", allow_module_level=True) |
| 45 | |
| 46 | HEADS, HEAD_DIM = 32, 128 |
| 47 | WIDTH = HEAD_DIM // 2 + HEAD_DIM // 32 |
| 48 | # V4.1's 8-token blocks, with the pool (2048 blocks) and top-k (512) scaled |
| 49 | # down so short test contexts still select. |
| 50 | CAND_BLOCK, POOL_BLOCKS, TOPK = 8, 16, 32 |
| 51 | MAX_LEN = 4096 |
| 52 | DEVICE = "cuda" |
| 53 | |
| 54 | _E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) |
| 55 | E2M1 = torch.cat([_E2M1, -_E2M1]) |
| 56 | |
| 57 | |
| 58 | def _dequant(packed, e8m0): |
| 59 | """[..., D/2] e2m1 pairs and [..., D/32] e8m0 to [..., D] fp32.""" |
| 60 | table = E2M1.to(packed.device) |
| 61 | vals = torch.stack([table[(packed & 0xF).long()], table[(packed >> 4).long()]], -1) |
| 62 | scale = torch.pow(2.0, e8m0.float() - 127) |
| 63 | return (vals.flatten(-2).unflatten(-1, (-1, 32)) * scale[..., None]).flatten(-2) |
| 64 | |
| 65 | |
| 66 | @pytest.fixture(autouse=True) |
| 67 | def _workspace(): |
| 68 | init_workspace_manager(torch.device(DEVICE)) |
| 69 | yield |
| 70 | reset_workspace_manager() |
| 71 | |
| 72 | |
| 73 | class _Case: |
| 74 | """Requests with their keys written into a block-major pool (ratio-2 and |
| 75 | ratio-1 pages side by side in every block) and into natural-order caches |
| 76 | the reference reads. The bf16 inputs are kept for the end-to-end test.""" |
| 77 | |
| 78 | def __init__(self, seq_lens, block): |
| 79 | from aiter.ops.triton.fusions.k_norm_rope_mxfp4_cache import ( |
| 80 | k_norm_rope_mxfp4_cache, |
| 81 | ) |
| 82 | |
| 83 | self.seq_lens = seq_lens |
| 84 | num_blocks = sum(cdiv(n, block) for n in seq_lens) + 3 |
| 85 | perm = torch.randperm(num_blocks) |
| 86 | width = max(cdiv(n, block) for n in seq_lens) |
| 87 | self.block_table = torch.zeros(len(seq_lens), width, dtype=torch.int32) |
| 88 | used = 0 |
| 89 | for req, n in enumerate(seq_lens): |
| 90 | self.block_table[req, : cdiv(n, block)] = perm[used : used + cdiv(n, block)] |
| 91 | used += cdiv(n, block) |
| 92 | self.block_table = self.block_table.to(DEVICE) |
| 93 | |
| 94 | def align(n): |
| 95 | return cdiv(n, 576) * 576 |
| 96 | |
| 97 | page2, page1 = block // 2 * WIDTH, block * WIDTH |
| 98 | stride = align(page2) + align(page1) + 576 |
| 99 | self.pool = torch.zeros(num_blocks * stride, dtype=torch.uint8, device=DEVICE) |
| 100 | self.page_bytes = torch.zeros_like(self.pool, dtype=torch.bool) |
| 101 | for i in range(num_blocks): |
| 102 | base = i * stride |
| 103 | self.page_bytes[base : base + page2] = True |
| 104 | base += align(page2) |
| 105 | self.page_bytes[base : base + page1] = True |
| 106 | self.cache = { |
| 107 | 2: torch.as_strided( |
| 108 | self.pool, (num_blocks, block // 2, WIDTH), (stride, WIDTH, 1), 0 |
| 109 | ), |
| 110 | 1: torch.as_strided( |
| 111 | self.pool, |
| 112 | (num_blocks, block, WIDTH), |
| 113 | (stride, WIDTH, 1), |
| 114 | align(page2), |
| 115 | ), |
| 116 | } |
| 117 | self.offsets = [0] |
| 118 | for n in seq_lens: |
| 119 | self.offsets.append(self.offsets[-1] + n) |
| 120 | # the same keys in natural order, one entry per page |
| 121 | self.natural = { |
| 122 | r: torch.zeros(self.offsets[-1], 1, WIDTH, dtype=torch.uint8, device=DEVICE) |
| 123 | for r in (1, 2) |
| 124 | } |
| 125 | self.cos_sin = cos_sin = torch.randn(MAX_LEN, 64, device=DEVICE) |
| 126 | self.norm = norm = torch.rand(HEAD_DIM, device=DEVICE) + 0.5 |
| 127 | self.k_pre = {} |
| 128 | for req, n in enumerate(seq_lens): |
| 129 | if n == 0: |
| 130 | continue |
| 131 | k_pre = torch.randn(n, HEAD_DIM, device=DEVICE).to(torch.bfloat16) |
| 132 | self.k_pre[req] = k_pre |
| 133 | pos = torch.arange(n, device=DEVICE) |
| 134 | for ratio, cache in self.cache.items(): |
| 135 | comp = pos // ratio |
| 136 | boundary = (pos + 1) % ratio == 0 |
| 137 | entries = cache.shape[1] |
| 138 | page = self.block_table[req, comp // entries].long() |
| 139 | slot = torch.where(boundary, page * entries + comp % entries, -1) |
| 140 | nat = torch.where(boundary, self.offsets[req] + comp, -1) |
| 141 | rocm_mxfp4_indexer_k_store( |
| 142 | k_pre, |
| 143 | pos, |
| 144 | cos_sin, |
| 145 | norm, |
| 146 | 1e-6, |
| 147 | cache, |
| 148 | slot, |
| 149 | ratio, |
| 150 | True, |
| 151 | num_heads=HEADS, |
| 152 | ) |
| 153 | k_norm_rope_mxfp4_cache( |
| 154 | k_pre, pos, cos_sin, norm, 1e-6, self.natural[ratio], nat, ratio |
| 155 | ) |
| 156 | |
| 157 | def keys(self, ratio, req, n): |
| 158 | nat = self.natural[ratio][self.offsets[req] : self.offsets[req] + n, 0] |
| 159 | return _dequant(nat[:, : HEAD_DIM // 2], nat[:, HEAD_DIM // 2 :]) |
| 160 | |
| 161 | |
| 162 | def _queries(num_rows): |
| 163 | q = torch.randint(0, 256, (num_rows, HEADS, HEAD_DIM // 2), dtype=torch.uint8) |
| 164 | scales = torch.randint(122, 130, (num_rows, HEADS, HEAD_DIM // 32)) |
| 165 | q, scales = q.to(DEVICE), scales.to(torch.uint8).to(DEVICE) |
| 166 | weights = torch.randn(num_rows, HEADS, device=DEVICE) |
| 167 | # the fused Q kernel returns the scales as one int32 per head |
| 168 | return q, scales.view(torch.int32).squeeze(-1), weights, _dequant(q, scales) |
| 169 | |
| 170 | |
| 171 | def _scores(q, w, keys): |
| 172 | return (torch.relu(q @ keys.T) * w[:, None]).sum(0) |
| 173 | |
| 174 | |
| 175 | def _assert_topk(row, scores, k): |
| 176 | """``row`` is a top-k of the entries of ``scores`` above -inf, up to the |
| 177 | fp32 reduce order.""" |
| 178 | valid = scores > float("-inf") |
| 179 | want = min(k, int(valid.sum())) |
| 180 | sel = row[row >= 0].long() |
| 181 | assert sel.numel() == want and sel.unique().numel() == want |
| 182 | if want: |
| 183 | assert bool(valid[sel].all()) |
| 184 | kth = scores[valid].topk(want).values[-1] |
| 185 | finite = scores[torch.isfinite(scores)] |
| 186 | tol = 1e-4 * float(finite.abs().max()) if finite.numel() else 0.0 |
| 187 | assert bool((scores[sel] >= kth - tol).all()) |
| 188 | |
| 189 | |
| 190 | def _block_maxima(scores): |
| 191 | end = scores.numel() |
| 192 | if end == 0: |
| 193 | return scores |
| 194 | padded = torch.full((cdiv(end, CAND_BLOCK) * CAND_BLOCK,), float("-inf")) |
| 195 | padded[:end] = scores.cpu() |
| 196 | out = padded.view(-1, CAND_BLOCK).amax(1) |
| 197 | out[(end - 1) // CAND_BLOCK] = float("inf") |
| 198 | return out.to(scores.device) |
| 199 | |
| 200 | |
| 201 | def _forward_context(monkeypatch, metadata): |
| 202 | context = types.SimpleNamespace( |
| 203 | attn_metadata={"indexer": metadata}, |
| 204 | cudagraph_runtime_mode=CUDAGraphMode.NONE, |
| 205 | ) |
| 206 | monkeypatch.setattr(ops, "get_forward_context", lambda: context) |
| 207 | |
| 208 | |
| 209 | def _decode_metadata(case, rows, ratio, query_lens): |
| 210 | lens = torch.tensor([(p + 1) // ratio for _, p in rows], dtype=torch.int32) |
| 211 | lens = lens.to(DEVICE) |
| 212 | block_table = case.block_table[[req for req, _ in rows]].contiguous() |
| 213 | entries = case.cache[ratio].shape[1] |
| 214 | next_n = max(query_lens) |
| 215 | context_lens = (torch.tensor(case.seq_lens, dtype=torch.int32) // ratio).to(DEVICE) |
| 216 | native = native_decode(lens, block_table, query_lens, next_n, context_lens) |
| 217 | schedule = torch.empty( |
| 218 | ops.rocm_mxfp4_decode_schedule_words(HEADS, HEAD_DIM, entries, next_n), |
| 219 | dtype=torch.int32, |
| 220 | device=DEVICE, |
| 221 | ) |
| 222 | return RocmMxfp4IndexerMetadata( |
| 223 | seq_lens=None, |
| 224 | max_seq_len=max(case.seq_lens), |
| 225 | slot_mapping=None, |
| 226 | num_decodes=len(case.seq_lens), |
| 227 | num_decode_tokens=len(rows), |
| 228 | num_prefills=0, |
| 229 | num_prefill_tokens=0, |
| 230 | decode=types.SimpleNamespace(block_table=block_table, seq_lens=lens[:, None]), |
| 231 | decode_row_lens=lens, |
| 232 | decode_block_ends=(lens + CAND_BLOCK - 1) // CAND_BLOCK, |
| 233 | decode_native=native, |
| 234 | decode_schedule=ops.build_rocm_mxfp4_decode_schedule( |
| 235 | lens, HEADS, HEAD_DIM, entries, schedule, MAX_LEN // ratio, native |
| 236 | ), |
| 237 | ) |
| 238 | |
| 239 | |
| 240 | def _prefill_metadata( |
| 241 | case, rows, ratio, chunk_bounds, query_start_loc, gather_rows, gather |
| 242 | ): |
| 243 | chunks = [] |
| 244 | for req_lo, req_hi, t0, t1 in chunk_bounds: |
| 245 | ends = torch.tensor( |
| 246 | [(rows[t][1] + 1) // ratio for t in range(t0, t1)], dtype=torch.int32 |
| 247 | ).to(DEVICE) |
| 248 | chunks.append( |
| 249 | DeepseekV32IndexerPrefillChunkMetadata( |
| 250 | block_table=case.block_table[req_lo:req_hi], |
| 251 | cu_seqlen_ks=torch.zeros_like(ends), |
| 252 | cu_seqlen_ke=ends, |
| 253 | cu_seq_lens=torch.zeros(1, dtype=torch.int32, device=DEVICE), |
| 254 | token_to_seq=torch.zeros(1, dtype=torch.int32, device=DEVICE), |
| 255 | total_seq_lens=1, |
| 256 | token_start=t0, |
| 257 | token_end=t1, |
| 258 | num_reqs=req_hi - req_lo, |
| 259 | ) |
| 260 | ) |
| 261 | seq_lens = torch.tensor(case.seq_lens) |
| 262 | context_lens = (seq_lens // ratio).int().to(DEVICE) |
| 263 | plans = plan_prefill_chunks( |
| 264 | chunks, |
| 265 | query_start_loc, |
| 266 | seq_lens, |
| 267 | context_lens, |
| 268 | ratio, |
| 269 | CAND_BLOCK, |
| 270 | 0.0 if gather else None, |
| 271 | torch.tensor(query_start_loc, dtype=torch.int32, device=DEVICE), |
| 272 | ) |
| 273 | return RocmMxfp4IndexerMetadata( |
| 274 | seq_lens=None, |
| 275 | max_seq_len=max(case.seq_lens), |
| 276 | slot_mapping=None, |
| 277 | num_decodes=0, |
| 278 | num_decode_tokens=0, |
| 279 | num_prefills=len(case.seq_lens), |
| 280 | num_prefill_tokens=len(rows), |
| 281 | prefill=types.SimpleNamespace(chunks=chunks), |
| 282 | prefill_plans=plans, |
| 283 | gather_launches=plan_gather_launches(chunks, plans, gather_rows), |
| 284 | ) |
| 285 | |
| 286 | |
| 287 | # The kernel prefers 128-token pages. |
| 288 | BLOCKS = pytest.mark.parametrize("block", [128, 64]) |
| 289 | |
| 290 | |
| 291 | @BLOCKS |
| 292 | def test_k_store_stays_inside_its_pages(block): |
| 293 | """The writer touches only its pages of the block-major pool. What they hold |
| 294 | is checked by the layer tests, which read them back through aiter's kernel.""" |
| 295 | set_random_seed(0) |
| 296 | case = _Case([700, 333], block) |
| 297 | assert int(case.pool[~case.page_bytes].count_nonzero()) == 0 |
| 298 | |
| 299 | |
| 300 | @pytest.mark.parametrize( |
| 301 | "rows, width, k", |
| 302 | # vLLM's kernel below 128 rows, aiter's from 128 rows and on every k=2048 |
| 303 | [(8, 4096, 512), (128, 4096, 512), (8, 16384, 2048)], |
| 304 | ) |
| 305 | def test_topk_matches_torch_on_both_kernels(rows, width, k): |
| 306 | set_random_seed(0) |
| 307 | logits = torch.randn(rows, width, device=DEVICE) |
| 308 | lengths = torch.randint(1, width + 1, (rows,), dtype=torch.int32, device=DEVICE) |
| 309 | lengths[0] = k // 2 # shorter than k, so the rest of the row pads with -1 |
| 310 | out = torch.full((rows, k), 7, dtype=torch.int32, device=DEVICE) |
| 311 | ops._topk(logits, lengths, out, k) |
| 312 | for r in range(rows): |
| 313 | n = min(k, int(lengths[r])) |
| 314 | want = torch.topk(logits[r, : int(lengths[r])], n).indices.sort().values |
| 315 | got = out[r] |
| 316 | assert torch.equal(got[got >= 0].sort().values, want.int()) |
| 317 | assert int((got == -1).sum()) == k - n |
| 318 | |
| 319 | |
| 320 | def _run_layers(monkeypatch, case, rows, metadata): |
| 321 | """Ratio-2 dense layer, the candidate source, and a consumer both ways; |
| 322 | each checked against the torch reference.""" |
| 323 | num_rows = len(rows) |
| 324 | q, q_scale, weights, q_ref = _queries(num_rows) |
| 325 | hidden = torch.empty(num_rows, 1, device=DEVICE) |
| 326 | |
| 327 | def topk_buffer(): |
| 328 | return torch.full((num_rows + 2, TOPK), 7, dtype=torch.int32, device=DEVICE) |
| 329 | |
| 330 | _forward_context(monkeypatch, metadata(2)) |
| 331 | out = topk_buffer() |
| 332 | ops.rocm_mxfp4_sparse_attn_indexer( |
| 333 | hidden, |
| 334 | "indexer", |
| 335 | case.cache[2], |
| 336 | q, |
| 337 | q_scale, |
| 338 | weights, |
| 339 | TOPK, |
| 340 | HEAD_DIM, |
| 341 | MAX_LEN // 2, |
| 342 | out, |
| 343 | compress_ratio=2, |
| 344 | ) |
| 345 | for i, (req, pos) in enumerate(rows): |
| 346 | end = (pos + 1) // 2 |
| 347 | keys = case.keys(2, req, end) |
| 348 | _assert_topk(out[i], _scores(q_ref[i], weights[i], keys), TOPK) |
| 349 | |
| 350 | source = metadata(1) |
| 351 | _forward_context(monkeypatch, source) |
| 352 | pool = torch.full((num_rows + 2, POOL_BLOCKS), 5, dtype=torch.int32, device=DEVICE) |
| 353 | out = topk_buffer() |
| 354 | ops.rocm_mxfp4_sparse_attn_indexer( |
| 355 | hidden, |
| 356 | "indexer", |
| 357 | case.cache[1], |
| 358 | q, |
| 359 | q_scale, |
| 360 | weights, |
| 361 | TOPK, |
| 362 | HEAD_DIM, |
| 363 | MAX_LEN, |
| 364 | out, |
| 365 | compress_ratio=1, |
| 366 | candidate_blocks=pool, |
| 367 | candidate_block_size=CAND_BLOCK, |
| 368 | candidate_write=True, |
| 369 | ) |
| 370 | allowed = [] |
| 371 | for i, (req, pos) in enumerate(rows): |
| 372 | scores = _scores(q_ref[i], weights[i], case.keys(1, req, pos + 1)) |
| 373 | _assert_topk(out[i], scores, TOPK) |
| 374 | _assert_topk(pool[i], _block_maxima(scores), POOL_BLOCKS) |
| 375 | mask = torch.zeros_like(scores, dtype=torch.bool) |
| 376 | for block in pool[i][pool[i] >= 0].tolist(): |
| 377 | mask[block * CAND_BLOCK : (block + 1) * CAND_BLOCK] = True |
| 378 | allowed.append(torch.where(mask, scores, float("-inf"))) |
| 379 | |
| 380 | pa = ops._aiter() |
| 381 | build = pa.build_candidate_gather |
| 382 | builds = [] |
| 383 | |
| 384 | def counting_build(*args, **kwargs): |
| 385 | builds.append(1) |
| 386 | return build(*args, **kwargs) |
| 387 | |
| 388 | monkeypatch.setattr(pa, "build_candidate_gather", counting_build) |
| 389 | for gather in (True, False): |
| 390 | consumer = metadata(1, gather) |
| 391 | consumer.decode_use_gather = gather |
| 392 | _forward_context(monkeypatch, consumer) |
| 393 | builds.clear() |
| 394 | # Two consumer layers of one step: the second reuses the first's pools. |
| 395 | for _ in range(2): |
| 396 | out = topk_buffer() |
| 397 | ops.rocm_mxfp4_sparse_mqa_indexer( |
| 398 | hidden, |
| 399 | "indexer", |
| 400 | case.cache[1], |
| 401 | q, |
| 402 | q_scale, |
| 403 | weights, |
| 404 | TOPK, |
| 405 | HEAD_DIM, |
| 406 | MAX_LEN, |
| 407 | out, |
| 408 | 1, |
| 409 | pool, |
| 410 | CAND_BLOCK, |
| 411 | POOL_BLOCKS * CAND_BLOCK, |
| 412 | ) |
| 413 | for i in range(num_rows): |
| 414 | _assert_topk(out[i], allowed[i], TOPK) |
| 415 | if gather: |
| 416 | launches = len(consumer.gather_launches) + (consumer.decode is not None) |
| 417 | assert len(builds) == launches |
| 418 | |
| 419 | |
| 420 | @BLOCKS |
| 421 | @pytest.mark.parametrize( |
| 422 | "query_lens", |
| 423 | # Uniform steps launch the dense layers on next_n-row sequences, also with |
| 424 | # cudagraph padding (query length 0) after them; a ragged step keeps a row |
| 425 | # per token. |
| 426 | [[6, 6, 6, 6], [2, 2, 2], [2, 2, 2, 0], [2, 1, 2]], |
| 427 | ids=["native6", "native2", "padded", "ragged"], |
| 428 | ) |
| 429 | def test_decode_layers_match_reference(monkeypatch, block, query_lens): |
| 430 | set_random_seed(0) |
| 431 | next_n = max(query_lens) |
| 432 | case = _Case([q and n for q, n in zip(query_lens, [900, 333, 610, 1200])], block) |
| 433 | # a padding request has next_n rows of its own and no context |
| 434 | rows = [ |
| 435 | (req, n - q + j if q else -1) |
| 436 | for req, (n, q) in enumerate(zip(case.seq_lens, query_lens)) |
| 437 | for j in range(q or next_n) |
| 438 | ] |
| 439 | _run_layers( |
| 440 | monkeypatch, |
| 441 | case, |
| 442 | rows, |
| 443 | lambda r, gather=False: _decode_metadata(case, rows, r, query_lens), |
| 444 | ) |
| 445 | |
| 446 | |
| 447 | @pytest.mark.parametrize( |
| 448 | "seq_lens,new_tokens,chunks", |
| 449 | [ |
| 450 | # one request past a cached prefix, then one sliced into two chunks |
| 451 | ([760, 500], [300, 500], [(0, 1, 0, 300), (1, 2, 300, 560), (1, 2, 560, 800)]), |
| 452 | # three requests in one chunk, the first two with equal rows |
| 453 | ([600, 450, 900], [200, 200, 120], [(0, 3, 0, 520)]), |
| 454 | # four requests with four different query lengths in one chunk |
| 455 | ([900, 600, 450, 700], [150, 90, 200, 60], [(0, 4, 0, 500)]), |
| 456 | ], |
| 457 | ids=["sliced", "batched", "ragged"], |
| 458 | ) |
| 459 | # The consumers rejoin a request's rows across chunks, then cut them at most |
| 460 | # gather_rows a launch. |
| 461 | @pytest.mark.parametrize("gather_rows", [1 << 20, 128], ids=["joined", "cut"]) |
| 462 | @BLOCKS |
| 463 | def test_prefill_layers_match_reference( |
| 464 | monkeypatch, seq_lens, new_tokens, chunks, gather_rows, block |
| 465 | ): |
| 466 | set_random_seed(0) |
| 467 | case = _Case(seq_lens, block) |
| 468 | rows = [ |
| 469 | (req, n - q + i) |
| 470 | for req, (n, q) in enumerate(zip(seq_lens, new_tokens)) |
| 471 | for i in range(q) |
| 472 | ] |
| 473 | query_start_loc = [0, *itertools.accumulate(new_tokens)] |
| 474 | _run_layers( |
| 475 | monkeypatch, |
| 476 | case, |
| 477 | rows, |
| 478 | lambda r, gather=True: _prefill_metadata( |
| 479 | case, rows, r, chunks, query_start_loc, gather_rows, gather |
| 480 | ), |
| 481 | ) |
| 482 | |
| 483 | |
| 484 | def _rope_gptj(x, cos_sin, positions): |
| 485 | """GPT-J RoPE on the last 64 dims, as the indexer rotates Q and K: bf16 in |
| 486 | and out.""" |
| 487 | x = x.bfloat16().float() |
| 488 | cos, sin = cos_sin[positions].chunk(2, dim=-1) |
| 489 | if x.dim() == 3: |
| 490 | cos, sin = cos[:, None], sin[:, None] |
| 491 | rot = x[..., -64:] |
| 492 | even, odd = rot[..., 0::2].clone(), rot[..., 1::2].clone() |
| 493 | rot[..., 0::2] = even * cos - odd * sin |
| 494 | rot[..., 1::2] = odd * cos + even * sin |
| 495 | return x.bfloat16() |
| 496 | |
| 497 | |
| 498 | def _mxfp4_round_trip(x): |
| 499 | """MXFP4 round trip of x, rounded as the writers do: a scale of |
| 500 | 2^ceil(log2(amax / 6)) per 32 values, then the nearest e2m1 value, ties to |
| 501 | even.""" |
| 502 | blocks = x.float().unflatten(-1, (-1, 32)) |
| 503 | amax = blocks.abs().amax(-1, keepdim=True).clamp(min=2**-126) |
| 504 | scale = torch.exp2(torch.ceil(torch.log2(amax / 6))) |
| 505 | y = (blocks / scale).clamp(-6, 6) |
| 506 | # Even codes first: argmin keeps the first of two equal distances. |
| 507 | grid = torch.cat([_E2M1[0::2], _E2M1[1::2]]).to(x.device) |
| 508 | idx = (y.abs()[..., None] - grid).abs().argmin(-1) |
| 509 | return (grid[idx] * y.sign() * scale).flatten(-2) |
| 510 | |
| 511 | |
| 512 | def _torch_keys(case, ratio, req): |
| 513 | """A request's compressed keys in torch: k_norm of each group's last |
| 514 | token, rotated at the group's first position, through MXFP4.""" |
| 515 | k = case.k_pre[req][ratio - 1 :: ratio].float() |
| 516 | k = k * torch.rsqrt(k.pow(2).mean(-1, keepdim=True) + 1e-6) * case.norm |
| 517 | pos = torch.arange(k.shape[0], device=DEVICE) * ratio |
| 518 | return _mxfp4_round_trip(_rope_gptj(k, case.cos_sin, pos)) |
| 519 | |
| 520 | |
| 521 | def _topk_indices(q, weights, keys, ends, k): |
| 522 | """Top-k of sum_h relu(q_h . key) * weights_h over each row's [0, end).""" |
| 523 | scores = (torch.einsum("thd,cd->thc", q, keys).relu() * weights[..., None]).sum(1) |
| 524 | cols = torch.arange(keys.shape[0], device=DEVICE) |
| 525 | return scores.masked_fill(cols >= ends[:, None], float("-inf")).topk(k).indices |
| 526 | |
| 527 | |
| 528 | def _recall(selected, reference): |
| 529 | """Per row, the fraction of the reference top-k that ``selected`` picked.""" |
| 530 | return (selected[:, :, None] == reference[:, None, :]).any(1).float().mean(1) |
| 531 | |
| 532 | |
| 533 | @BLOCKS |
| 534 | @pytest.mark.parametrize("ratio", [2, 1]) |
| 535 | @pytest.mark.parametrize("step", ["decode", "prefill"]) |
| 536 | def test_selection_from_bf16_matches_torch(monkeypatch, step, ratio, block): |
| 537 | """End to end from bf16 K and Q: vLLM's MXFP4 K store and Q quant, the |
| 538 | paged pool, metadata from the planning functions the builder calls, and |
| 539 | the indexer op pick what the same MXFP4 math picks in torch.""" |
| 540 | set_random_seed(0) |
| 541 | if step == "decode": |
| 542 | case = _Case([3000, 2600, 1900, 3500, 1200, 2200, 3900, 1600], block) |
| 543 | rows = [(req, n - 1) for req, n in enumerate(case.seq_lens)] |
| 544 | metadata = _decode_metadata(case, rows, ratio, [1] * len(rows)) |
| 545 | else: |
| 546 | new_tokens = [300, 200, 250] |
| 547 | case = _Case([1500, 900, 2000], block) |
| 548 | rows = [ |
| 549 | (req, n - q + i) |
| 550 | for req, (n, q) in enumerate(zip(case.seq_lens, new_tokens)) |
| 551 | for i in range(q) |
| 552 | ] |
| 553 | metadata = _prefill_metadata( |
| 554 | case, |
| 555 | rows, |
| 556 | ratio, |
| 557 | [(0, len(new_tokens), 0, len(rows))], |
| 558 | [0, *itertools.accumulate(new_tokens)], |
| 559 | 1 << 20, |
| 560 | gather=False, |
| 561 | ) |
| 562 | num_rows = len(rows) |
| 563 | q_bf16 = torch.randn(num_rows, HEADS, HEAD_DIM, device=DEVICE).to(torch.bfloat16) |
| 564 | w = torch.randn(num_rows, HEADS, device=DEVICE) |
| 565 | pos = torch.tensor([p for _, p in rows], device=DEVICE) |
| 566 | scales = (HEAD_DIM**-0.5, HEADS**-0.5) |
| 567 | |
| 568 | (q, q_scale), weights = rocm_mxfp4_indexer_q_quant( |
| 569 | pos, q_bf16, case.cos_sin, w, *scales |
| 570 | ) |
| 571 | _forward_context(monkeypatch, metadata) |
| 572 | out = torch.full((num_rows, TOPK), 7, dtype=torch.int32, device=DEVICE) |
| 573 | ops.rocm_mxfp4_sparse_attn_indexer( |
| 574 | torch.empty(num_rows, 1, device=DEVICE), |
| 575 | "indexer", |
| 576 | case.cache[ratio], |
| 577 | q, |
| 578 | q_scale, |
| 579 | weights, |
| 580 | TOPK, |
| 581 | HEAD_DIM, |
| 582 | MAX_LEN // ratio, |
| 583 | out, |
| 584 | compress_ratio=ratio, |
| 585 | ) |
| 586 | ends = (pos + 1) // ratio |
| 587 | assert bool(((out >= 0) & (out < ends[:, None])).all()) |
| 588 | assert all(row.unique().numel() == TOPK for row in out) |
| 589 | |
| 590 | q_ref = _mxfp4_round_trip(_rope_gptj(q_bf16, case.cos_sin, pos)) |
| 591 | w_ref = w * scales[0] * scales[1] |
| 592 | agree = [] |
| 593 | for req in range(len(case.seq_lens)): |
| 594 | mine = torch.tensor([r == req for r, _ in rows], device=DEVICE) |
| 595 | ref = _topk_indices( |
| 596 | q_ref[mine], w_ref[mine], _torch_keys(case, ratio, req), ends[mine], TOPK |
| 597 | ) |
| 598 | agree.append(_recall(out[mine], ref)) |
| 599 | assert torch.cat(agree).mean().item() >= 0.99 |