@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 109 lines 3.91 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """fp8 query and index-K cache on the Triton indexer must agree with bf16.
4
5 The MiniMax M3 lightning-indexer scores only feed a top-k block ranking (the
6 attention itself reads the main KV cache), so an fp8 (e4m3) side cache trades
7 ~2% relative score error for half the per-step index-K read bandwidth — the
8 dominant linear-in-context decode cost at long context. #45892 enabled fp8
9 index caches on the SM100 MSA path; this covers the Triton impl used
10 everywhere else (e.g. SM120), for both the decode and prefill scorers.
11 """
12
13 import pytest
14 import torch
15
16 from vllm.models.minimax_m3.common.ops.index_topk import (
17 SPARSE_BLOCK_SIZE,
18 minimax_m3_index_decode,
19 minimax_m3_index_score,
20 minimax_m3_index_topk,
21 )
22 from vllm.platforms import current_platform
23
24 FP8_CUDA_SUPPORTED = current_platform.is_cuda() and current_platform.supports_fp8()
25 PAGE = SPARSE_BLOCK_SIZE
26 HEAD_DIM = 128
27 HEADS = 1
28 TOPK = 16
29 NUM_PAGES = 100
30 SEQ_LEN = 96 * PAGE + 37
31 PLANTED = (7, 40, 88)
32
33
34 def _make_inputs(device: str) -> tuple[torch.Tensor, torch.Tensor]:
35 torch.manual_seed(3)
36 keys = torch.randn(NUM_PAGES * PAGE, HEAD_DIM, dtype=torch.bfloat16, device=device)
37 # RMS-normalize like the real (post index_k_norm) keys.
38 keys = keys * keys.float().pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt().to(
39 torch.bfloat16
40 )
41 q = torch.randn(1, HEADS, HEAD_DIM, dtype=torch.bfloat16, device=device)
42 for b in PLANTED:
43 keys[b * PAGE + 5] = (q[0, 0] * 3).to(torch.bfloat16)
44 return q, keys
45
46
47 def _check_selection(top_bf16: set[int], top_fp8: set[int]) -> None:
48 # Planted (clearly relevant) blocks must never be lost.
49 assert all(b in top_fp8 for b in PLANTED)
50 # Selection may differ only at the noise floor (near-tied filler blocks).
51 assert len(top_bf16 & top_fp8) >= TOPK - 2
52
53
54 @pytest.mark.skipif(not FP8_CUDA_SUPPORTED, reason="CUDA FP8 support required")
55 @torch.inference_mode()
56 def test_fp8_index_query_and_cache_decode_topk_matches_bf16() -> None:
57 dev = "cuda"
58 q, keys = _make_inputs(dev)
59 block_table = torch.arange(NUM_PAGES, dtype=torch.int32, device=dev).unsqueeze(0)
60 seq_lens = torch.tensor([SEQ_LEN], dtype=torch.int32, device=dev)
61
62 def select(cache: torch.Tensor) -> set[int]:
63 query = q.to(cache.dtype)
64 idx = minimax_m3_index_decode(
65 query,
66 cache.view(NUM_PAGES, PAGE, HEAD_DIM).contiguous(),
67 block_table,
68 seq_lens,
69 SEQ_LEN,
70 TOPK,
71 0,
72 0,
73 HEADS,
74 1,
75 1,
76 )
77 return {int(b) for b in idx[0, 0].tolist() if b >= 0}
78
79 _check_selection(select(keys), select(keys.to(torch.float8_e4m3fn)))
80
81
82 @pytest.mark.skipif(not FP8_CUDA_SUPPORTED, reason="CUDA FP8 support required")
83 @torch.inference_mode()
84 def test_fp8_index_query_and_cache_prefill_topk_matches_bf16() -> None:
85 dev = "cuda"
86 q, keys = _make_inputs(dev)
87 block_table = torch.arange(NUM_PAGES, dtype=torch.int32, device=dev).unsqueeze(0)
88 seq_lens = torch.tensor([SEQ_LEN], dtype=torch.int32, device=dev)
89 # single decode-like query token at the end of the sequence
90 cu_seqlens_q = torch.tensor([0, 1], dtype=torch.int32, device=dev)
91 prefix_lens = torch.tensor([SEQ_LEN - 1], dtype=torch.int32, device=dev)
92
93 def select(cache: torch.Tensor) -> set[int]:
94 query = q.to(cache.dtype)
95 score = minimax_m3_index_score(
96 query,
97 cache.view(NUM_PAGES, PAGE, HEAD_DIM).contiguous(),
98 block_table,
99 cu_seqlens_q,
100 seq_lens,
101 prefix_lens,
102 1,
103 SEQ_LEN,
104 HEADS,
105 )
106 idx = minimax_m3_index_topk(score, cu_seqlens_q, prefix_lens, 1, TOPK, 0, 0)
107 return {int(b) for b in idx[0, 0].tolist() if b >= 0}
108
109 _check_selection(select(keys), select(keys.to(torch.float8_e4m3fn)))