@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 189 lines 5.35 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5 import torch
6
7 from vllm.platforms import current_platform
8
9 if not torch.cuda.is_available() or not current_platform.supports_fp8():
10 pytest.skip(
11 "MiniMax-M3 FP8 sparse attention scale tests require an FP8 GPU.",
12 allow_module_level=True,
13 )
14
15 if current_platform.is_rocm():
16 from vllm.models.minimax_m3.amd.ops.sparse_attn import (
17 SPARSE_BLOCK_SIZE,
18 minimax_m3_sparse_attn,
19 minimax_m3_sparse_attn_decode,
20 )
21 else:
22 from vllm.models.minimax_m3.common.ops.sparse_attn import (
23 SPARSE_BLOCK_SIZE,
24 minimax_m3_sparse_attn,
25 minimax_m3_sparse_attn_decode,
26 )
27
28
29 DEVICE = "cuda"
30 DTYPE = torch.bfloat16
31 HEAD_DIM = 128
32 NUM_KV_HEADS = 1
33 NUM_HEADS = 2
34 K_SCALE = 0.25
35 V_SCALE = 0.5
36
37
38 def _scale_tensors(mode: str, num_blocks: int):
39 if mode == "scalar":
40 k_scale = torch.tensor(K_SCALE, dtype=torch.float32, device=DEVICE)
41 v_scale = torch.tensor(V_SCALE, dtype=torch.float32, device=DEVICE)
42 else:
43 shape = (NUM_KV_HEADS, num_blocks * SPARSE_BLOCK_SIZE)
44 k_scale = torch.full(shape, K_SCALE, dtype=torch.float32, device=DEVICE)
45 v_scale = torch.full(shape, V_SCALE, dtype=torch.float32, device=DEVICE)
46 return k_scale, v_scale
47
48
49 def _make_kv_cache(num_blocks: int, seed: int):
50 torch.manual_seed(seed)
51 fp8_dtype = current_platform.fp8_dtype()
52 kv_ref = torch.randn(
53 num_blocks,
54 NUM_KV_HEADS,
55 SPARSE_BLOCK_SIZE,
56 2 * HEAD_DIM,
57 dtype=DTYPE,
58 device=DEVICE,
59 )
60 kv_fp8 = torch.empty_like(kv_ref, dtype=fp8_dtype)
61 kv_fp8[..., :HEAD_DIM] = (kv_ref[..., :HEAD_DIM].float() / K_SCALE).to(fp8_dtype)
62 kv_fp8[..., HEAD_DIM:] = (kv_ref[..., HEAD_DIM:].float() / V_SCALE).to(fp8_dtype)
63
64 kv_dequant = torch.empty_like(kv_ref)
65 kv_dequant[..., :HEAD_DIM] = (kv_fp8[..., :HEAD_DIM].float() * K_SCALE).to(DTYPE)
66 kv_dequant[..., HEAD_DIM:] = (kv_fp8[..., HEAD_DIM:].float() * V_SCALE).to(DTYPE)
67 return kv_fp8, kv_dequant
68
69
70 @pytest.mark.parametrize("scale_mode", ["scalar", "per_token_head"])
71 @torch.inference_mode()
72 def test_minimax_m3_sparse_prefill_fp8_kv_scales(scale_mode: str):
73 total_q = 17
74 num_blocks = 1
75 torch.manual_seed(0)
76 q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, dtype=DTYPE, device=DEVICE) * 0.1
77 kv_fp8, kv_dequant = _make_kv_cache(num_blocks, seed=1)
78 k_scale, v_scale = _scale_tensors(scale_mode, num_blocks)
79
80 topk = torch.zeros(NUM_KV_HEADS, total_q, 1, dtype=torch.int32, device=DEVICE)
81 block_table = torch.zeros(1, 1, dtype=torch.int32, device=DEVICE)
82 cu_seqlens = torch.tensor([0, total_q], dtype=torch.int32, device=DEVICE)
83 seq_lens = torch.tensor([total_q], dtype=torch.int32, device=DEVICE)
84 prefix_lens = torch.zeros(1, dtype=torch.int32, device=DEVICE)
85 got = torch.empty_like(q)
86 ref = torch.empty_like(q)
87 unscaled = torch.empty_like(q)
88
89 minimax_m3_sparse_attn(
90 q,
91 kv_fp8,
92 topk,
93 block_table,
94 cu_seqlens,
95 seq_lens,
96 prefix_lens,
97 total_q,
98 NUM_KV_HEADS,
99 HEAD_DIM**-0.5,
100 got,
101 k_scale=k_scale,
102 v_scale=v_scale,
103 )
104 minimax_m3_sparse_attn(
105 q,
106 kv_dequant,
107 topk,
108 block_table,
109 cu_seqlens,
110 seq_lens,
111 prefix_lens,
112 total_q,
113 NUM_KV_HEADS,
114 HEAD_DIM**-0.5,
115 ref,
116 )
117 minimax_m3_sparse_attn(
118 q,
119 kv_fp8,
120 topk,
121 block_table,
122 cu_seqlens,
123 seq_lens,
124 prefix_lens,
125 total_q,
126 NUM_KV_HEADS,
127 HEAD_DIM**-0.5,
128 unscaled,
129 )
130
131 torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2)
132 assert not torch.allclose(unscaled, ref, rtol=1e-1, atol=1e-1)
133
134
135 @pytest.mark.parametrize("scale_mode", ["scalar", "per_token_head"])
136 @torch.inference_mode()
137 def test_minimax_m3_sparse_decode_fp8_kv_scales(scale_mode: str):
138 total_q = 2
139 num_blocks = 1
140 torch.manual_seed(2)
141 q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, dtype=DTYPE, device=DEVICE) * 0.1
142 kv_fp8, kv_dequant = _make_kv_cache(num_blocks, seed=3)
143 k_scale, v_scale = _scale_tensors(scale_mode, num_blocks)
144
145 topk = torch.zeros(NUM_KV_HEADS, total_q, 1, dtype=torch.int32, device=DEVICE)
146 block_table = torch.zeros(total_q, 1, dtype=torch.int32, device=DEVICE)
147 seq_lens = torch.tensor([64, 128], dtype=torch.int32, device=DEVICE)
148 got = torch.empty_like(q)
149 ref = torch.empty_like(q)
150 unscaled = torch.empty_like(q)
151
152 minimax_m3_sparse_attn_decode(
153 q,
154 kv_fp8,
155 topk,
156 block_table,
157 seq_lens,
158 NUM_KV_HEADS,
159 HEAD_DIM**-0.5,
160 got,
161 decode_query_len=1,
162 k_scale=k_scale,
163 v_scale=v_scale,
164 )
165 minimax_m3_sparse_attn_decode(
166 q,
167 kv_dequant,
168 topk,
169 block_table,
170 seq_lens,
171 NUM_KV_HEADS,
172 HEAD_DIM**-0.5,
173 ref,
174 decode_query_len=1,
175 )
176 minimax_m3_sparse_attn_decode(
177 q,
178 kv_fp8,
179 topk,
180 block_table,
181 seq_lens,
182 NUM_KV_HEADS,
183 HEAD_DIM**-0.5,
184 unscaled,
185 decode_query_len=1,
186 )
187
188 torch.testing.assert_close(got, ref, rtol=2e-2, atol=2e-2)
189 assert not torch.allclose(unscaled, ref, rtol=1e-1, atol=1e-1)