@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 233 lines 9.29 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Unit test for MiMo-V2 fused fp8 ``qkv_proj`` sharding.
4
5 The checkpoint stores the fused QKV pre-sharded at ``num_key_value_heads``
6 chunks (``ckpt_tp``), each holding that slice's ``[Q_c | K_c | V_c]`` rows with
7 its own block scale grid. ``ckpt_tp`` is not the layer's KV-head count: a
8 MiMo-V2.5 SWA layer has 8 KV heads but 4 chunks of 3712 rows (= 29 blocks), so
9 each chunk carries two KV heads. Reading the scales as one slice per KV head
10 expands 14 x 128 = 1792 rows for an 1856-row slice and fails.
11
12 These tests build such checkpoints for the GA, SWA and Pro geometries and check
13 that every rank gets exactly the weights it should own, for TP sizes below,
14 equal to and above ``ckpt_tp`` (the last case replicating KV heads).
15 """
16
17 import pytest
18 import torch
19
20 from vllm.model_executor.models.mimo_v2 import _shard_fp8_qkv_proj
21 from vllm.utils.math_utils import cdiv
22 from vllm.utils.torch_utils import set_default_torch_num_threads
23
24 pytestmark = pytest.mark.cpu_test
25
26 BLOCK = 128
27 FP8_DTYPE = torch.float8_e4m3fn
28 FP8_MAX = torch.finfo(FP8_DTYPE).max
29 COLS = 128 # one column block is enough to exercise the row mapping
30 # fp8 e4m3 has 3 mantissa bits, so a single block-quantization round trip costs
31 # ~2-3% relative L2; a wrong row -> scale pairing costs >= 10%.
32 TOL = 0.06
33
34 # (num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp): MiMo-V2.5 GA, its
35 # SWA/MTP layers, and MiMo-V2.5-Pro.
36 GEOMETRIES = {
37 "ga": (64, 4, 192, 128, 4),
38 "swa": (64, 8, 192, 128, 4),
39 "pro": (128, 8, 192, 128, 8),
40 }
41
42
43 @pytest.fixture(scope="module", autouse=True)
44 def single_threaded_cpu_ops():
45 # Do not prime the parent's thread pool before later engine forks.
46 with set_default_torch_num_threads(1):
47 yield
48
49
50 def _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp):
51 return (
52 num_heads // ckpt_tp * head_dim
53 + num_kv_heads // ckpt_tp * head_dim
54 + num_kv_heads // ckpt_tp * v_head_dim
55 )
56
57
58 def _quantize_chunks(truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp):
59 """Block-quantize each checkpoint chunk on its own grid, like the export."""
60 rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp)
61 weights, scales = [], []
62 for chunk in range(ckpt_tp):
63 rows = truth[chunk * rows_per_chunk : (chunk + 1) * rows_per_chunk]
64 quantized = torch.zeros_like(rows)
65 scale = torch.zeros(
66 cdiv(rows_per_chunk, BLOCK), rows.shape[1] // BLOCK, dtype=torch.float32
67 )
68 for row in range(0, rows_per_chunk, BLOCK):
69 for col in range(0, rows.shape[1], BLOCK):
70 block = rows[row : row + BLOCK, col : col + BLOCK]
71 amax = block.abs().max().clamp(min=1e-12) / FP8_MAX
72 scale[row // BLOCK, col // BLOCK] = amax
73 quantized[row : row + BLOCK, col : col + BLOCK] = block / amax
74 weights.append(quantized.to(FP8_DTYPE))
75 scales.append(scale)
76 return torch.cat(weights, dim=0), torch.cat(scales, dim=0)
77
78
79 def _owned_rows(
80 truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp, tp_rank, tp_size
81 ):
82 """Rows the rank must own, in the forward's `[Q | K | V]` order."""
83 rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp)
84 q_per_chunk = num_heads // ckpt_tp * head_dim
85 k_per_chunk = num_kv_heads // ckpt_tp * head_dim
86 kv_heads: list[int] = []
87 if tp_size <= num_kv_heads:
88 kv_heads = list(
89 range(
90 tp_rank * (num_kv_heads // tp_size),
91 (tp_rank + 1) * (num_kv_heads // tp_size),
92 )
93 )
94 else:
95 kv_heads = [tp_rank // (tp_size // num_kv_heads)]
96 index = []
97 for head in range(
98 tp_rank * (num_heads // tp_size), (tp_rank + 1) * (num_heads // tp_size)
99 ):
100 chunk, offset = divmod(head, num_heads // ckpt_tp)
101 start = chunk * rows_per_chunk + offset * head_dim
102 index += list(range(start, start + head_dim))
103 for head in kv_heads:
104 chunk, offset = divmod(head, num_kv_heads // ckpt_tp)
105 start = chunk * rows_per_chunk + q_per_chunk + offset * head_dim
106 index += list(range(start, start + head_dim))
107 for head in kv_heads:
108 chunk, offset = divmod(head, num_kv_heads // ckpt_tp)
109 start = chunk * rows_per_chunk + q_per_chunk + k_per_chunk + offset * v_head_dim
110 index += list(range(start, start + v_head_dim))
111 return truth[torch.tensor(index)]
112
113
114 def _dequantize(weight, scale):
115 """What the kernel computes: scale[row // block] applied per weight row."""
116 per_row = scale.repeat_interleave(BLOCK, dim=1)
117 return weight.to(torch.float32) * per_row[torch.arange(weight.shape[0]) // BLOCK]
118
119
120 def _relative_l2(got, expected):
121 return ((got - expected).norm() / expected.norm()).item()
122
123
124 @pytest.mark.parametrize("geometry", list(GEOMETRIES))
125 @pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
126 def test_fused_qkv_proj_sharding(geometry, tp_size):
127 num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp = GEOMETRIES[geometry]
128 if num_heads % tp_size or (tp_size > num_kv_heads and tp_size % num_kv_heads):
129 pytest.skip("TP size does not divide the head counts")
130
131 rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp)
132 truth = torch.randn(ckpt_tp * rows_per_chunk, COLS)
133 for chunk in range(ckpt_tp): # per-chunk magnitude spread, as in the export
134 truth[chunk * rows_per_chunk : (chunk + 1) * rows_per_chunk] *= 1.0 + chunk
135 weight, scale = _quantize_chunks(
136 truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp
137 )
138
139 kv_per_rank = (num_kv_heads // tp_size) if tp_size <= num_kv_heads else 1
140 rows_rank = (num_heads // tp_size) * head_dim + kv_per_rank * (
141 head_dim + v_head_dim
142 )
143 for tp_rank in range(tp_size):
144 w_rank, s_rank = _shard_fp8_qkv_proj(
145 weight,
146 scale,
147 num_heads=num_heads,
148 num_kv_heads=num_kv_heads,
149 head_dim=head_dim,
150 v_head_dim=v_head_dim,
151 tp_rank=tp_rank,
152 tp_size=tp_size,
153 ckpt_tp=ckpt_tp,
154 )
155 assert w_rank.shape == (rows_rank, COLS)
156 assert s_rank.shape == (cdiv(rows_rank, BLOCK), scale.shape[1])
157 expected = _owned_rows(
158 truth,
159 num_heads,
160 num_kv_heads,
161 head_dim,
162 v_head_dim,
163 ckpt_tp,
164 tp_rank,
165 tp_size,
166 )
167 assert _relative_l2(_dequantize(w_rank, s_rank), expected) <= TOL
168
169
170 @pytest.mark.parametrize("geometry", ["ga", "swa"])
171 @pytest.mark.parametrize("tp_size", [1, 2, 4])
172 def test_exact_chunk_matches_checkpoint(geometry, tp_size):
173 """At tp_size <= ckpt_tp the chunk count is the only thing that matters."""
174 num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp = GEOMETRIES[geometry]
175 rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp)
176 truth = torch.randn(ckpt_tp * rows_per_chunk, COLS)
177 weight, scale = _quantize_chunks(
178 truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp
179 )
180 if tp_size == ckpt_tp:
181 w_rank, s_rank = _shard_fp8_qkv_proj(
182 weight,
183 scale,
184 num_heads=num_heads,
185 num_kv_heads=num_kv_heads,
186 head_dim=head_dim,
187 v_head_dim=v_head_dim,
188 tp_rank=0,
189 tp_size=tp_size,
190 ckpt_tp=ckpt_tp,
191 )
192 # The checkpoint chunk *is* the rank's slice: no re-quantization.
193 # (torch.equal has no fp8 kernel, so compare the underlying bytes.)
194 expected_w = weight.chunk(tp_size, dim=0)[0]
195 assert torch.equal(w_rank.view(torch.uint8), expected_w.view(torch.uint8))
196 assert torch.equal(s_rank, scale.chunk(tp_size, dim=0)[0])
197
198
199 def test_wrong_chunk_count_is_detected():
200 """Reading SWA scales as one slice per KV head must not look correct.
201
202 This is the regression the fix addresses: 116 scale rows are 4 x 29 (29 =
203 ceil(3712 / 128)), not 8 x 15, so ``ckpt_tp = num_kv_heads`` pairs every
204 weight row with the wrong scale.
205 """
206 num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp = GEOMETRIES["swa"]
207 rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp)
208 truth = torch.randn(ckpt_tp * rows_per_chunk, COLS)
209 for chunk in range(ckpt_tp):
210 truth[chunk * rows_per_chunk : (chunk + 1) * rows_per_chunk] *= 1.0 + chunk
211 weight, scale = _quantize_chunks(
212 truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp
213 )
214 assert scale.shape[0] == ckpt_tp * cdiv(rows_per_chunk, BLOCK)
215 assert scale.shape[0] != num_kv_heads * cdiv(rows_per_chunk, BLOCK)
216
217 kw = dict(
218 num_heads=num_heads,
219 num_kv_heads=num_kv_heads,
220 head_dim=head_dim,
221 v_head_dim=v_head_dim,
222 tp_rank=0,
223 tp_size=1,
224 )
225 correct = _shard_fp8_qkv_proj(weight, scale, ckpt_tp=ckpt_tp, **kw)
226 wrong = _shard_fp8_qkv_proj(weight, scale, ckpt_tp=num_kv_heads, **kw)
227 assert correct[0].shape == wrong[0].shape
228
229 expected = _owned_rows(
230 truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp, 0, 1
231 )
232 assert _relative_l2(_dequantize(*correct), expected) <= TOL
233 assert _relative_l2(_dequantize(*wrong), expected) > 5 * TOL