@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 211e252d0b

[Bugfix][Model] Fix MiMo-V2.5 fused fp8 qkv_proj sharding (pre-shard count is num_key_value_heads; MTP path too) (#57508)

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

vllmellm committed Sep 19, 2026 at 21:05 UTC 211e252d0b4f8429f9b15fc52bdfed07782c7f70
3 files changed +415 -63
tests/models/quantization/test_mimo_v2_qkv_shard.py new
+225
@@ -0,0 +1,225 @@
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 +
23 +pytestmark = pytest.mark.cpu_test
24 +
25 +BLOCK = 128
26 +FP8_DTYPE = torch.float8_e4m3fn
27 +FP8_MAX = torch.finfo(FP8_DTYPE).max
28 +COLS = 128 # one column block is enough to exercise the row mapping
29 +# fp8 e4m3 has 3 mantissa bits, so a single block-quantization round trip costs
30 +# ~2-3% relative L2; a wrong row -> scale pairing costs >= 10%.
31 +TOL = 0.06
32 +
33 +# (num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp): MiMo-V2.5 GA, its
34 +# SWA/MTP layers, and MiMo-V2.5-Pro.
35 +GEOMETRIES = {
36 + "ga": (64, 4, 192, 128, 4),
37 + "swa": (64, 8, 192, 128, 4),
38 + "pro": (128, 8, 192, 128, 8),
39 +}
40 +
41 +
42 +def _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp):
43 + return (
44 + num_heads // ckpt_tp * head_dim
45 + + num_kv_heads // ckpt_tp * head_dim
46 + + num_kv_heads // ckpt_tp * v_head_dim
47 + )
48 +
49 +
50 +def _quantize_chunks(truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp):
51 + """Block-quantize each checkpoint chunk on its own grid, like the export."""
52 + rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp)
53 + weights, scales = [], []
54 + for chunk in range(ckpt_tp):
55 + rows = truth[chunk * rows_per_chunk : (chunk + 1) * rows_per_chunk]
56 + quantized = torch.zeros_like(rows)
57 + scale = torch.zeros(
58 + cdiv(rows_per_chunk, BLOCK), rows.shape[1] // BLOCK, dtype=torch.float32
59 + )
60 + for row in range(0, rows_per_chunk, BLOCK):
61 + for col in range(0, rows.shape[1], BLOCK):
62 + block = rows[row : row + BLOCK, col : col + BLOCK]
63 + amax = block.abs().max().clamp(min=1e-12) / FP8_MAX
64 + scale[row // BLOCK, col // BLOCK] = amax
65 + quantized[row : row + BLOCK, col : col + BLOCK] = block / amax
66 + weights.append(quantized.to(FP8_DTYPE))
67 + scales.append(scale)
68 + return torch.cat(weights, dim=0), torch.cat(scales, dim=0)
69 +
70 +
71 +def _owned_rows(
72 + truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp, tp_rank, tp_size
73 +):
74 + """Rows the rank must own, in the forward's `[Q | K | V]` order."""
75 + rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp)
76 + q_per_chunk = num_heads // ckpt_tp * head_dim
77 + k_per_chunk = num_kv_heads // ckpt_tp * head_dim
78 + kv_heads: list[int] = []
79 + if tp_size <= num_kv_heads:
80 + kv_heads = list(
81 + range(
82 + tp_rank * (num_kv_heads // tp_size),
83 + (tp_rank + 1) * (num_kv_heads // tp_size),
84 + )
85 + )
86 + else:
87 + kv_heads = [tp_rank // (tp_size // num_kv_heads)]
88 + index = []
89 + for head in range(
90 + tp_rank * (num_heads // tp_size), (tp_rank + 1) * (num_heads // tp_size)
91 + ):
92 + chunk, offset = divmod(head, num_heads // ckpt_tp)
93 + start = chunk * rows_per_chunk + offset * head_dim
94 + index += list(range(start, start + head_dim))
95 + for head in kv_heads:
96 + chunk, offset = divmod(head, num_kv_heads // ckpt_tp)
97 + start = chunk * rows_per_chunk + q_per_chunk + offset * head_dim
98 + index += list(range(start, start + head_dim))
99 + for head in kv_heads:
100 + chunk, offset = divmod(head, num_kv_heads // ckpt_tp)
101 + start = chunk * rows_per_chunk + q_per_chunk + k_per_chunk + offset * v_head_dim
102 + index += list(range(start, start + v_head_dim))
103 + return truth[torch.tensor(index)]
104 +
105 +
106 +def _dequantize(weight, scale):
107 + """What the kernel computes: scale[row // block] applied per weight row."""
108 + per_row = scale.repeat_interleave(BLOCK, dim=1)
109 + return weight.to(torch.float32) * per_row[torch.arange(weight.shape[0]) // BLOCK]
110 +
111 +
112 +def _relative_l2(got, expected):
113 + return ((got - expected).norm() / expected.norm()).item()
114 +
115 +
116 +@pytest.mark.parametrize("geometry", list(GEOMETRIES))
117 +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
118 +def test_fused_qkv_proj_sharding(geometry, tp_size):
119 + num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp = GEOMETRIES[geometry]
120 + if num_heads % tp_size or (tp_size > num_kv_heads and tp_size % num_kv_heads):
121 + pytest.skip("TP size does not divide the head counts")
122 +
123 + rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp)
124 + truth = torch.randn(ckpt_tp * rows_per_chunk, COLS)
125 + for chunk in range(ckpt_tp): # per-chunk magnitude spread, as in the export
126 + truth[chunk * rows_per_chunk : (chunk + 1) * rows_per_chunk] *= 1.0 + chunk
127 + weight, scale = _quantize_chunks(
128 + truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp
129 + )
130 +
131 + kv_per_rank = (num_kv_heads // tp_size) if tp_size <= num_kv_heads else 1
132 + rows_rank = (num_heads // tp_size) * head_dim + kv_per_rank * (
133 + head_dim + v_head_dim
134 + )
135 + for tp_rank in range(tp_size):
136 + w_rank, s_rank = _shard_fp8_qkv_proj(
137 + weight,
138 + scale,
139 + num_heads=num_heads,
140 + num_kv_heads=num_kv_heads,
141 + head_dim=head_dim,
142 + v_head_dim=v_head_dim,
143 + tp_rank=tp_rank,
144 + tp_size=tp_size,
145 + ckpt_tp=ckpt_tp,
146 + )
147 + assert w_rank.shape == (rows_rank, COLS)
148 + assert s_rank.shape == (cdiv(rows_rank, BLOCK), scale.shape[1])
149 + expected = _owned_rows(
150 + truth,
151 + num_heads,
152 + num_kv_heads,
153 + head_dim,
154 + v_head_dim,
155 + ckpt_tp,
156 + tp_rank,
157 + tp_size,
158 + )
159 + assert _relative_l2(_dequantize(w_rank, s_rank), expected) <= TOL
160 +
161 +
162 +@pytest.mark.parametrize("geometry", ["ga", "swa"])
163 +@pytest.mark.parametrize("tp_size", [1, 2, 4])
164 +def test_exact_chunk_matches_checkpoint(geometry, tp_size):
165 + """At tp_size <= ckpt_tp the chunk count is the only thing that matters."""
166 + num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp = GEOMETRIES[geometry]
167 + rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp)
168 + truth = torch.randn(ckpt_tp * rows_per_chunk, COLS)
169 + weight, scale = _quantize_chunks(
170 + truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp
171 + )
172 + if tp_size == ckpt_tp:
173 + w_rank, s_rank = _shard_fp8_qkv_proj(
174 + weight,
175 + scale,
176 + num_heads=num_heads,
177 + num_kv_heads=num_kv_heads,
178 + head_dim=head_dim,
179 + v_head_dim=v_head_dim,
180 + tp_rank=0,
181 + tp_size=tp_size,
182 + ckpt_tp=ckpt_tp,
183 + )
184 + # The checkpoint chunk *is* the rank's slice: no re-quantization.
185 + # (torch.equal has no fp8 kernel, so compare the underlying bytes.)
186 + expected_w = weight.chunk(tp_size, dim=0)[0]
187 + assert torch.equal(w_rank.view(torch.uint8), expected_w.view(torch.uint8))
188 + assert torch.equal(s_rank, scale.chunk(tp_size, dim=0)[0])
189 +
190 +
191 +def test_wrong_chunk_count_is_detected():
192 + """Reading SWA scales as one slice per KV head must not look correct.
193 +
194 + This is the regression the fix addresses: 116 scale rows are 4 x 29 (29 =
195 + ceil(3712 / 128)), not 8 x 15, so ``ckpt_tp = num_kv_heads`` pairs every
196 + weight row with the wrong scale.
197 + """
198 + num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp = GEOMETRIES["swa"]
199 + rows_per_chunk = _chunk_rows(num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp)
200 + truth = torch.randn(ckpt_tp * rows_per_chunk, COLS)
201 + for chunk in range(ckpt_tp):
202 + truth[chunk * rows_per_chunk : (chunk + 1) * rows_per_chunk] *= 1.0 + chunk
203 + weight, scale = _quantize_chunks(
204 + truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp
205 + )
206 + assert scale.shape[0] == ckpt_tp * cdiv(rows_per_chunk, BLOCK)
207 + assert scale.shape[0] != num_kv_heads * cdiv(rows_per_chunk, BLOCK)
208 +
209 + kw = dict(
210 + num_heads=num_heads,
211 + num_kv_heads=num_kv_heads,
212 + head_dim=head_dim,
213 + v_head_dim=v_head_dim,
214 + tp_rank=0,
215 + tp_size=1,
216 + )
217 + correct = _shard_fp8_qkv_proj(weight, scale, ckpt_tp=ckpt_tp, **kw)
218 + wrong = _shard_fp8_qkv_proj(weight, scale, ckpt_tp=num_kv_heads, **kw)
219 + assert correct[0].shape == wrong[0].shape
220 +
221 + expected = _owned_rows(
222 + truth, num_heads, num_kv_heads, head_dim, v_head_dim, ckpt_tp, 0, 1
223 + )
224 + assert _relative_l2(_dequantize(*correct), expected) <= TOL
225 + assert _relative_l2(_dequantize(*wrong), expected) > 5 * TOL
vllm/model_executor/models/mimo_v2.py
+141 -55
@@ -50,6 +50,7 @@ from vllm.model_executor.model_loader.weight_utils import (
50 )
51 from vllm.model_executor.models.utils import sequence_parallel_chunk
52 from vllm.sequence import IntermediateTensors
53 +from vllm.utils.math_utils import cdiv
54 from vllm.v1.attention.backend import AttentionType
55 from vllm.v1.attention.backends.registry import AttentionBackendEnum
56
@@ -460,6 +461,28 @@ class MiMoV2FlashDecoderLayer(nn.Module):
461 return self.config.hybrid_layer_pattern[self.layer_id] == 1
462
463
464 +def _requantize_fp8(
465 + grouped: torch.Tensor, rows_rank: int, block: int, dtype: torch.dtype
466 +) -> tuple[torch.Tensor, torch.Tensor]:
467 + """Block-quantize a rank's ``[Q | K | V]`` rows back to fp8.
468 +
469 + A rank's rows need not end on a block boundary (a single 1856-row slice is
470 + 14.5 blocks) while ``scaled_quantize`` requires both dims to be multiples of
471 + the block size, so pad the tail with zeros: zero rows cannot raise a block's
472 + amax, so every scale -- and the number of scale rows -- is unchanged. The
473 + padding is dropped again here.
474 + """
475 + padded = cdiv(rows_rank, block) * block
476 + if padded != rows_rank:
477 + grouped = torch.cat(
478 + [grouped, grouped.new_zeros(padded - rows_rank, grouped.shape[1])], dim=0
479 + )
480 + w_rank, s_rank = scaled_quantize(
481 + grouped, GroupShape(block, block), dtype, compute_dtype=torch.float32
482 + )
483 + return w_rank[:rows_rank], s_rank
484 +
485 +
486 def _shard_fp8_qkv_proj(
487 w_full: torch.Tensor,
488 s_full: torch.Tensor,
@@ -469,75 +492,136 @@ def _shard_fp8_qkv_proj(
492 v_head_dim: int,
493 tp_rank: int,
494 tp_size: int,
495 + ckpt_tp: int,
496 block: int = 128,
497 ) -> tuple[torch.Tensor, torch.Tensor]:
498 """Shard the fp8 qkv_proj weights for ``tp_rank``.
499
476 - The checkpoint stores the fused QKV as ``num_kv_heads`` contiguous groups
477 - (one per KV head; ``n`` below), each ordered ``[Q | K | V]``:
500 + The checkpoint stores the fused QKV pre-sharded for ``ckpt_tp`` ranks (the
501 + model config's ``num_key_value_heads``), each chunk holding that slice's
502 + Q, K and V rows:
503
479 - [Q_1 | K_1 | V_1 | Q_2 | K_2 | V_2 | ... | Q_n | K_n | V_n]
504 + [Q_0 | K_0 | V_0 | Q_1 | K_1 | V_1 | ... | Q_n | K_n | V_n] (n = ckpt_tp)
505
481 - Per group, Q has ``(num_heads / num_kv_heads) * head_dim`` rows, K has
482 - ``head_dim`` rows, and V has ``v_head_dim`` rows.
506 + Per chunk, Q has ``(num_heads / ckpt_tp) * head_dim`` rows, K has
507 + ``(num_kv_heads / ckpt_tp) * head_dim`` rows, and V has
508 + ``(num_kv_heads / ckpt_tp) * v_head_dim`` rows, and the fp8 block scales
509 + are tiled per chunk too (``ceil(rows_per_chunk / block)`` rows each).
510
484 - Each TP rank owns ``g = num_kv_heads / tp_size`` of these groups, and the
485 - forward expects them de-interleaved into a single Q, K, and V block:
511 + ``ckpt_tp`` is not the layer's KV-head count: a MiMo-V2.5 SWA layer has 8
512 + KV heads over 4 chunks, so each chunk carries two KV heads (3712 rows = 29
513 + blocks) while a GA layer has 4 KV heads over 4 chunks (3392 rows each).
514 +
515 + The forward expects each rank's slice de-interleaved:
516
517 [Q_1 | Q_2 | ... | Q_g | K_1 | K_2 | ... | K_g | V_1 | V_2 | ... | V_g]
518
489 - When ``g == 1`` the rank's slice is already ``[Q | K | V]``, so a plain
490 - chunk suffices. When ``g > 1`` we cannot reach the de-interleaved layout by
491 - re-permuting the fp8 block scales: each scale covers a 128-row block, and
492 - since K is 192 rows (1.5 blocks) a block straddles the K/V boundary, so no
493 - whole-block permutation produces it. Instead we dequantize this rank's
494 - groups to float (dropping the block constraint), reorder the rows into the
495 - layout above (Q, K, and V then each span a whole number of blocks), and
496 - re-quantize to fp8.
519 + When ``tp_size == ckpt_tp`` the checkpoint chunk *is* that layout, so a
520 + plain chunk of both weight and scale suffices. Otherwise each rank's Q, K
521 + and V rows are gathered from the chunks that hold them, dequantized with
522 + the chunk's own scales, reordered, and re-quantized to fp8.
523 """
498 - assert tp_size <= num_kv_heads and num_kv_heads % tp_size == 0, (
499 - "TP size must evenly split the number of KV heads."
524 + assert num_heads % tp_size == 0, (
525 + f"num_heads={num_heads} must be divisible by tp_size={tp_size}."
526 )
501 -
502 - kv_heads_per_rank = num_kv_heads // tp_size
503 - if kv_heads_per_rank == 1:
504 - # One KV head per rank. The weights and scale can be trivially sharded
505 - # without re-quantization.
506 - w = w_full.chunk(tp_size, dim=0)[tp_rank]
507 - s = s_full.chunk(tp_size, dim=0)[tp_rank]
508 - return w, s
509 -
510 - q_rows_per_group = (num_heads // num_kv_heads) * head_dim
511 - k_rows_per_group = head_dim
512 - v_rows_per_group = v_head_dim
513 - rows_per_group = q_rows_per_group + k_rows_per_group + v_rows_per_group
514 - scale_rows_per_group = s_full.shape[0] // num_kv_heads
515 - qs, ks, vs = [], [], []
516 - for g_idx in range(tp_rank * kv_heads_per_rank, (tp_rank + 1) * kv_heads_per_rank):
517 - row_start = g_idx * rows_per_group
518 - scale_row_start = g_idx * scale_rows_per_group
519 - # Dequantize this group's weights.
520 - w_g = w_full[row_start : row_start + rows_per_group].to(torch.float32)
521 - s_g = s_full[scale_row_start : scale_row_start + scale_rows_per_group].to(
522 - torch.float32
523 - )
524 - s_g_expanded = s_g.repeat_interleave(block, dim=0).repeat_interleave(
525 - block, dim=1
526 - )[:rows_per_group]
527 - w_g_dequant = w_g * s_g_expanded
528 - # Track the dequantized q, k, and v weights separately.
529 - qs.append(w_g_dequant[:q_rows_per_group])
530 - ks.append(w_g_dequant[q_rows_per_group : q_rows_per_group + k_rows_per_group])
531 - vs.append(w_g_dequant[q_rows_per_group + k_rows_per_group :])
532 -
533 - # Combine the q, k, and v weights into the following layout:
534 - # [Q_1, Q_2, .., Q_g, K_1, K_2, ..., K_g, V_1, V_2, ..., V_g]
535 - grouped = torch.cat([torch.cat(qs), torch.cat(ks), torch.cat(vs)], dim=0)
536 - # Quantize back to fp8.
537 - return scaled_quantize(
538 - grouped, GroupShape(block, block), w_full.dtype, compute_dtype=torch.float32
527 + if ckpt_tp <= 0 or num_heads % ckpt_tp or num_kv_heads % ckpt_tp:
528 + raise ValueError(
529 + f"fused qkv_proj is pre-sharded at {ckpt_tp} chunks, which do not "
530 + f"divide num_heads={num_heads} / num_kv_heads={num_kv_heads}."
531 + )
532 + # When there are fewer KV heads than ranks, vLLM replicates them
533 + # (`num_kv_head_replicas`) and rank r owns KV head r // replicas, which
534 + # keeps every Q head grouped with the KV head it attends to.
535 + if tp_size <= num_kv_heads:
536 + assert num_kv_heads % tp_size == 0, (
537 + f"num_kv_heads={num_kv_heads} must be divisible by tp_size={tp_size}."
538 + )
539 + kv_head_ids = list(
540 + range(
541 + tp_rank * (num_kv_heads // tp_size),
542 + (tp_rank + 1) * (num_kv_heads // tp_size),
543 + )
544 + )
545 + else:
546 + assert tp_size % num_kv_heads == 0, (
547 + f"tp_size={tp_size} must be divisible by num_kv_heads={num_kv_heads}."
548 + )
549 + kv_head_ids = [tp_rank // (tp_size // num_kv_heads)]
550 + q_head_ids = list(
551 + range(tp_rank * (num_heads // tp_size), (tp_rank + 1) * (num_heads // tp_size))
552 )
553
554 + rows_per_chunk = w_full.shape[0] // ckpt_tp
555 + q_per_chunk = (num_heads // ckpt_tp) * head_dim
556 + k_per_chunk = (num_kv_heads // ckpt_tp) * head_dim
557 + v_per_chunk = (num_kv_heads // ckpt_tp) * v_head_dim
558 + if q_per_chunk + k_per_chunk + v_per_chunk != rows_per_chunk:
559 + raise ValueError(
560 + f"fused qkv_proj has {w_full.shape[0]} rows, not {ckpt_tp} chunks of "
561 + f"{q_per_chunk} Q + {k_per_chunk} K + {v_per_chunk} V rows."
562 + )
563 +
564 + # The scales are tiled per chunk; they collapse to one continuous grid when
565 + # a chunk is a whole number of blocks (the SWA chunks are: 3712 = 29 * 128).
566 + chunk_scale_rows = cdiv(rows_per_chunk, block)
567 + rows = torch.arange(w_full.shape[0])
568 + per_chunk_scales = s_full.shape[0] == ckpt_tp * chunk_scale_rows
569 + if per_chunk_scales:
570 + scale_index = (rows // rows_per_chunk) * chunk_scale_rows + (
571 + rows % rows_per_chunk
572 + ) // block
573 + elif s_full.shape[0] == cdiv(w_full.shape[0], block):
574 + scale_index = rows // block
575 + else:
576 + raise ValueError(
577 + f"fused qkv_proj scale has {s_full.shape[0]} rows, expected either "
578 + f"{ckpt_tp * chunk_scale_rows} ({ckpt_tp} chunks of "
579 + f"{chunk_scale_rows} rows) or {cdiv(w_full.shape[0], block)} "
580 + f"(one continuous grid)"
581 + )
582 +
583 + if tp_size == ckpt_tp and per_chunk_scales:
584 + # One checkpoint chunk per rank: already [Q | K | V] for that rank, and
585 + # its scale rows line up with ceil(rows_per_chunk / block).
586 + return (
587 + w_full.chunk(ckpt_tp, dim=0)[tp_rank],
588 + s_full.chunk(ckpt_tp, dim=0)[tp_rank],
589 + )
590 +
591 + # Gather this rank's Q, K and V rows from the chunks that hold them.
592 + q_heads_per_chunk = num_heads // ckpt_tp
593 + kv_heads_per_chunk = num_kv_heads // ckpt_tp
594 + head_rows = torch.arange(head_dim)
595 + v_head_rows = torch.arange(v_head_dim)
596 + row_index: list[torch.Tensor] = []
597 + for head in q_head_ids:
598 + chunk = head // q_heads_per_chunk
599 + row_index.append(
600 + chunk * rows_per_chunk + (head % q_heads_per_chunk) * head_dim + head_rows
601 + )
602 + for head in kv_head_ids:
603 + chunk = head // kv_heads_per_chunk
604 + row_index.append(
605 + chunk * rows_per_chunk
606 + + q_per_chunk
607 + + (head % kv_heads_per_chunk) * head_dim
608 + + head_rows
609 + )
610 + for head in kv_head_ids:
611 + chunk = head // kv_heads_per_chunk
612 + row_index.append(
613 + chunk * rows_per_chunk
614 + + q_per_chunk
615 + + k_per_chunk
616 + + (head % kv_heads_per_chunk) * v_head_dim
617 + + v_head_rows
618 + )
619 + index = torch.cat(row_index)
620 + grouped = w_full[index].to(torch.float32) * s_full[
621 + scale_index[index]
622 + ].repeat_interleave(block, dim=1)
623 + return _requantize_fp8(grouped, index.numel(), block, w_full.dtype)
624 +
625
626 @support_torch_compile
627 class MiMoV2Model(nn.Module, EagleModelMixin):
@@ -819,6 +903,8 @@ class MiMoV2Model(nn.Module, EagleModelMixin):
903 v_head_dim=attn.v_head_dim,
904 tp_rank=tp_rank,
905 tp_size=tp_size,
906 + # The fused qkv_proj is pre-sharded for this many ranks.
907 + ckpt_tp=self.config.num_key_value_heads,
908 )
909 sharded = {"weight": w_rank, "weight_scale_inv": s_rank}
910 for kind, tensor in sharded.items():
vllm/model_executor/models/mimo_v2_mtp.py
+49 -8
@@ -45,7 +45,7 @@ from .interfaces import (
45 SupportsMultiModal,
46 _require_is_multimodal,
47 )
48 -from .mimo_v2 import MiMoV2Attention, MiMoV2MLP
48 +from .mimo_v2 import MiMoV2Attention, MiMoV2MLP, _shard_fp8_qkv_proj
49 from .utils import _merge_multimodal_embeddings, maybe_prefix
50
51 # MiMo-V2 checkpoints contain multiple MTP layers, but vLLM currently supports
@@ -266,6 +266,7 @@ class MiMoV2MTP(nn.Module):
266
267 params_dict = dict(self.named_parameters())
268 loaded_params: set[str] = set()
269 + pending_qkv_proj: dict[str, dict[str, torch.Tensor]] = {}
270
271 for name, loaded_weight in weights:
272 if "rotary_emb.inv_freq" in name:
@@ -279,14 +280,54 @@ class MiMoV2MTP(nn.Module):
280 ):
281 continue
282
282 - # Support fused qkv_proj checkpoint (Pro format).
283 - # The checkpoint is stored pre-sharded for TP=8 as
284 - # [Q_rank0, K_rank0, V_rank0, Q_rank1, ...], so splitting along
285 - # dim 0 with chunk(tp_size) gives each rank its Q+K+V slice for
286 - # both the FP8 weight and the block weight_scale_inv. This matches
287 - # how the main model loads the same layout.
283 + # Fused qkv_proj: the fp8 weight and its block scales arrive as
284 + # separate tensors and are pre-sharded at `config.num_key_value_heads`
285 + # chunks of [Q_c | K_c | V_c] (see _shard_fp8_qkv_proj for the
286 + # layout). Shard them together: a plain chunk(tp_size) only happens
287 + # to be right when tp_size is that chunk count *and* the scales are
288 + # tiled per chunk (they are 4 x 29 = 116 rows for MiMo-V2.5, not
289 + # tp_size slices of an even split).
290 if "qkv_proj" in name:
289 - if name in params_dict:
291 + is_fp8_weight = (
292 + name.endswith("qkv_proj.weight")
293 + and loaded_weight.dtype == torch.float8_e4m3fn
294 + )
295 + is_fp8_scale = name.endswith("qkv_proj.weight_scale_inv")
296 + if is_fp8_weight or is_fp8_scale:
297 + prefix, kind = name.rsplit(".", 1)
298 + if (
299 + f"{prefix}.weight" not in params_dict
300 + or f"{prefix}.weight_scale_inv" not in params_dict
301 + ):
302 + continue
303 + entry = pending_qkv_proj.setdefault(prefix, {})
304 + entry[kind] = loaded_weight
305 + if "weight" not in entry or "weight_scale_inv" not in entry:
306 + # Waiting for the other half of the fused projection.
307 + continue
308 + del pending_qkv_proj[prefix]
309 +
310 + # Geometry comes from the module the weights belong to.
311 + attn = self.get_submodule(prefix.rsplit(".", 1)[0])
312 + w_rank, s_rank = _shard_fp8_qkv_proj(
313 + entry["weight"],
314 + entry["weight_scale_inv"],
315 + num_heads=attn.total_num_heads,
316 + num_kv_heads=attn.total_num_kv_heads,
317 + head_dim=attn.head_dim,
318 + v_head_dim=attn.v_head_dim,
319 + tp_rank=tp_rank,
320 + tp_size=tp_size,
321 + ckpt_tp=self.config.num_key_value_heads,
322 + )
323 + for k, tensor in (("weight", w_rank), ("weight_scale_inv", s_rank)):
324 + param = params_dict[f"{prefix}.{k}"]
325 + if tensor.shape[0] > param.shape[0]:
326 + tensor = tensor[: param.shape[0]]
327 + default_weight_loader(param, tensor)
328 + loaded_params.add(f"{prefix}.{k}")
329 + elif name in params_dict:
330 + # Non-fp8 fused checkpoint: unchanged behaviour.
331 param = params_dict[name]
332 loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]
333 default_weight_loader(param, loaded_weight)