@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b5adb027ad

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Models] Fix MiMo v2.x QKV TP sharding + FP4 support (#45200)

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Giancarlo Delfin committed Jun 15, 2026 at 02:13 UTC b5adb027ad03c29b46181752ba3b1cb84eff1dd4

2 files changed +170 -5

vllm/model_executor/layers/quantization/fp8.py

+10

@@ -106,6 +106,7 @@ class Fp8Config(QuantizationConfig):

106 activation_scheme: str = "dynamic",

107 ignored_layers: list[str] | None = None,

108 weight_block_size: list[int] | None = None,

109 + store_dtype: str | None = None,

110 ) -> None:

111 super().__init__()

112

@@ -115,6 +116,7 @@ class Fp8Config(QuantizationConfig):

116 raise ValueError(f"Unsupported activation scheme {activation_scheme}")

117 self.activation_scheme = activation_scheme

118 self.ignored_layers = ignored_layers or []

119 + self.store_dtype = store_dtype

120 if weight_block_size is not None:

121 if not is_checkpoint_fp8_serialized:

122 raise ValueError(

@@ -162,6 +164,7 @@ class Fp8Config(QuantizationConfig):

164 activation_scheme = cls.get_from_keys(config, ["activation_scheme"])

165 ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None)

166 weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None)

167 + store_dtype = cls.get_from_keys_or(config, ["store_dtype"], None)

168 if not ignored_layers:

169 ignored_layers = cls.get_from_keys_or(

170 config, ["modules_to_not_convert"], None

@@ -171,6 +174,7 @@ class Fp8Config(QuantizationConfig):

174 activation_scheme=activation_scheme,

175 ignored_layers=ignored_layers,

176 weight_block_size=weight_block_size,

177 + store_dtype=store_dtype,

178 )

179

180 def get_quant_method(

@@ -198,6 +202,12 @@ class Fp8Config(QuantizationConfig):

202 fused_mapping=self.packed_modules_mapping,

203 ):

204 return UnquantizedFusedMoEMethod(layer.moe_config)

205 + if self.store_dtype == "mxfp4":

206 + from vllm.model_executor.layers.quantization.mxfp4 import (

207 + Mxfp4MoEMethod,

208 + )

209 +

210 + return Mxfp4MoEMethod(layer.moe_config)

211 if self.is_checkpoint_fp8_serialized:

212 moe_quant_method = Fp8MoEMethod(self, layer)

213 else:

vllm/model_executor/models/mimo_v2.py

+160 -5

@@ -35,6 +35,10 @@ from vllm.model_executor.layers.linear import (

35 )

36 from vllm.model_executor.layers.logits_processor import LogitsProcessor

37 from vllm.model_executor.layers.quantization import QuantizationConfig

38 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

39 + GroupShape,

40 + scaled_quantize,

41 +)

42 from vllm.model_executor.layers.rotary_embedding import get_rope

43 from vllm.model_executor.layers.vocab_parallel_embedding import (

44 ParallelLMHead,

@@ -455,6 +459,85 @@ class MiMoV2FlashDecoderLayer(nn.Module):

459 return self.config.hybrid_layer_pattern[self.layer_id] == 1

460

461

462 +def _shard_fp8_qkv_proj(

463 + w_full: torch.Tensor,

464 + s_full: torch.Tensor,

465 + num_heads: int,

466 + num_kv_heads: int,

467 + head_dim: int,

468 + v_head_dim: int,

469 + tp_rank: int,

470 + tp_size: int,

471 + block: int = 128,

472 +) -> tuple[torch.Tensor, torch.Tensor]:

473 + """Shard the fp8 qkv_proj weights for ``tp_rank``.

474 +

475 + The checkpoint stores the fused QKV as ``num_kv_heads`` contiguous groups

476 + (one per KV head; ``n`` below), each ordered ``[Q | K | V]``:

477 +

478 + [Q_1 | K_1 | V_1 | Q_2 | K_2 | V_2 | ... | Q_n | K_n | V_n]

479 +

480 + Per group, Q has ``(num_heads / num_kv_heads) * head_dim`` rows, K has

481 + ``head_dim`` rows, and V has ``v_head_dim`` rows.

482 +

483 + Each TP rank owns ``g = num_kv_heads / tp_size`` of these groups, and the

484 + forward expects them de-interleaved into a single Q, K, and V block:

485 +

486 + [Q_1 | Q_2 | ... | Q_g | K_1 | K_2 | ... | K_g | V_1 | V_2 | ... | V_g]

487 +

488 + When ``g == 1`` the rank's slice is already ``[Q | K | V]``, so a plain

489 + chunk suffices. When ``g > 1`` we cannot reach the de-interleaved layout by

490 + re-permuting the fp8 block scales: each scale covers a 128-row block, and

491 + since K is 192 rows (1.5 blocks) a block straddles the K/V boundary, so no

492 + whole-block permutation produces it. Instead we dequantize this rank's

493 + groups to float (dropping the block constraint), reorder the rows into the

494 + layout above (Q, K, and V then each span a whole number of blocks), and

495 + re-quantize to fp8.

496 + """

497 + assert tp_size <= num_kv_heads and num_kv_heads % tp_size == 0, (

498 + "TP size must evenly split the number of KV heads."

499 + )

500 +

501 + kv_heads_per_rank = num_kv_heads // tp_size

502 + if kv_heads_per_rank == 1:

503 + # One KV head per rank. The weights and scale can be trivially sharded

504 + # without re-quantization.

505 + w = w_full.chunk(tp_size, dim=0)[tp_rank]

506 + s = s_full.chunk(tp_size, dim=0)[tp_rank]

507 + return w, s

508 +

509 + q_rows_per_group = (num_heads // num_kv_heads) * head_dim

510 + k_rows_per_group = head_dim

511 + v_rows_per_group = v_head_dim

512 + rows_per_group = q_rows_per_group + k_rows_per_group + v_rows_per_group

513 + scale_rows_per_group = s_full.shape[0] // num_kv_heads

514 + qs, ks, vs = [], [], []

515 + for g_idx in range(tp_rank * kv_heads_per_rank, (tp_rank + 1) * kv_heads_per_rank):

516 + row_start = g_idx * rows_per_group

517 + scale_row_start = g_idx * scale_rows_per_group

518 + # Dequantize this group's weights.

519 + w_g = w_full[row_start : row_start + rows_per_group].to(torch.float32)

520 + s_g = s_full[scale_row_start : scale_row_start + scale_rows_per_group].to(

521 + torch.float32

522 + )

523 + s_g_expanded = s_g.repeat_interleave(block, dim=0).repeat_interleave(

524 + block, dim=1

525 + )[:rows_per_group]

526 + w_g_dequant = w_g * s_g_expanded

527 + # Track the dequantized q, k, and v weights separately.

528 + qs.append(w_g_dequant[:q_rows_per_group])

529 + ks.append(w_g_dequant[q_rows_per_group : q_rows_per_group + k_rows_per_group])

530 + vs.append(w_g_dequant[q_rows_per_group + k_rows_per_group :])

531 +

532 + # Combine the q, k, and v weights into the following layout:

533 + # [Q_1, Q_2, .., Q_g, K_1, K_2, ..., K_g, V_1, V_2, ..., V_g]

534 + grouped = torch.cat([torch.cat(qs), torch.cat(ks), torch.cat(vs)], dim=0)

535 + # Quantize back to fp8.

536 + return scaled_quantize(

537 + grouped, GroupShape(block, block), w_full.dtype, compute_dtype=torch.float32

538 + )

539 +

540 +

541 @support_torch_compile

542 class MiMoV2Model(nn.Module):

543 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

@@ -561,6 +644,10 @@ class MiMoV2Model(nn.Module):

644 params_dict = dict(self.named_parameters(remove_duplicate=False))

645 loaded_params: set[str] = set()

646 expert_params_mapping = self.get_expert_mapping()

647 + # Pro-format fused qkv_proj arrives as two tensors (weight and

648 + # weight_scale_inv). Store them per-layer so that they can be

649 + # sharded together.

650 + pending_fp8_qkv_proj: dict[str, dict[str, torch.Tensor]] = {}

651 for name, loaded_weight in weights:

652 if "rotary_emb.inv_freq" in name:

653 continue

@@ -604,11 +691,15 @@ class MiMoV2Model(nn.Module):

691 if expert_matched:

692 continue

693 # Support fused qkv_proj checkpoint (Pro format)

607 - if "qkv_proj" in name:

608 - if name in params_dict:

609 - param = params_dict[name]

610 - loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]

611 - default_weight_loader(param, loaded_weight)

694 + if self._try_load_fp8_qkv_proj(

695 + name,

696 + loaded_weight,

697 + pending_fp8_qkv_proj,

698 + params_dict,

699 + loaded_params,

700 + tp_rank,

701 + tp_size,

702 + ):

703 continue

704 stacked_matched = False

705 for param_name, weight_name, shard_id in stacked_params_mapping:

@@ -666,6 +757,70 @@ class MiMoV2Model(nn.Module):

757

758 return loaded_params

759

760 + def _try_load_fp8_qkv_proj(

761 + self,

762 + name: str,

763 + tensor: torch.Tensor,

764 + fp8_qkv_proj_dict: dict[str, dict[str, torch.Tensor]],

765 + params_dict: dict[str, torch.nn.Parameter],

766 + loaded_params: set[str],

767 + tp_rank: int,

768 + tp_size: int,

769 + ) -> bool:

770 + """

771 + The fused fp8 QKV projection weights and scale are stored separately.

772 + Special care must be taken while sharding these tensors across TP ranks.

773 + See _shard_fp8_qkv_proj for more details.

774 +

775 + Returns:

776 + True if ``tensor`` was an fp8 qkv_proj weight/scale and was consumed

777 + (caller should skip it); False otherwise, so the caller falls

778 + through to its normal loading path.

779 + """

780 + is_weight = (

781 + name.endswith("qkv_proj.weight") and tensor.dtype == torch.float8_e4m3fn

782 + )

783 + is_scale = name.endswith("qkv_proj.weight_scale_inv")

784 + if not is_weight and not is_scale:

785 + # Weight is not in FP8 format. Ignore.

786 + return False

787 +

788 + if is_pp_missing_parameter(name, self):

789 + # This qkv_proj is for a layer not on this PP rank.

790 + return True

791 +

792 + prefix, qkv_kind = name.rsplit(".", 1)

793 + entry = fp8_qkv_proj_dict.setdefault(prefix, {})

794 + entry[qkv_kind] = tensor

795 + if "weight" not in entry or "weight_scale_inv" not in entry:

796 + # Still waiting for the other param.

797 + return True

798 + del fp8_qkv_proj_dict[prefix]

799 +

800 + # Get self_attn module, which is a parent of qkv_proj.

801 + attn = self.get_submodule(prefix.rsplit(".", 1)[0])

802 +

803 + # Shard the qkv_proj per-rank.

804 + w_rank, s_rank = _shard_fp8_qkv_proj(

805 + entry["weight"],

806 + entry["weight_scale_inv"],

807 + num_heads=attn.total_num_heads,

808 + num_kv_heads=attn.total_num_kv_heads,

809 + head_dim=attn.head_dim,

810 + v_head_dim=attn.v_head_dim,

811 + tp_rank=tp_rank,

812 + tp_size=tp_size,

813 + )

814 + sharded = {"weight": w_rank, "weight_scale_inv": s_rank}

815 + for kind, tensor in sharded.items():

816 + param_name = f"{prefix}.{kind}"

817 + param = params_dict[param_name]

818 + if tensor.shape[0] > param.shape[0]:

819 + tensor = tensor[: param.shape[0]]

820 + default_weight_loader(param, tensor)

821 + loaded_params.add(param_name)

822 + return True

823 +

824

825 class MiMoV2FlashForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):

826 packed_modules_mapping = {