""" Block-FP8 variant of MiMo-V2. Subclasses upstream mimo_v2 with: 1. FusedQKVAttention: one qkv_proj matmul instead of three q/k/v. Matches the on-disk Xiaomi weight layout; saves ~12% decode bandwidth. 2. sanitize_block_fp8(): strips MTP/vision/audio only. Assumes the weight dict is already pre-stacked block_fp8 (see convert_mimo.py). 3. apply_block_fp8(): walks the model, replaces every Linear / SwitchLinear that has a `_scale_inv` companion in the weights with the QuantizedLinear / QuantizedSwitchLinear in mode='block_fp8'. Layers without a scale companion remain plain bf16. Triggered by load_model(model_config={"model_type": "mimo_v2_block_fp8"}). """ from typing import Optional, Any import mlx.core as mx import mlx.nn as nn from mlx.utils import tree_map_with_path from .mimo_v2 import ( Attention, Model as MiMoModel, ModelArgs, MoE, sum_gradients, scaled_dot_product_attention, initialize_rope, ) from .switch_layers import SwitchLinear, QuantizedSwitchLinear try: import mlx_block_fp8 as _bfp8 except ImportError: _bfp8 = None # extension not installed; falls back to mx.quantized_matmul # -------------------------------------------------------------------------- # Fused MoE-gate kernel. # # Replaces the upstream mx.compile'd group_expert_select with a single # custom Metal kernel that does sigmoid + bias + top-k + normalize in one # threadgroup dispatch per token. Eliminates the ~25% decode time spent # in compiled GatherAxis + BroadcastDivideMultiply chains. # # Specialized for MiMo-V2.5: n_routed_experts=256, top_k=8, n_group=1, # norm_topk_prob=True, routed_scaling_factor=1.0. Other configs fall back # to the original group_expert_select. # -------------------------------------------------------------------------- _MOE_GATE_KERNEL = None def _get_moe_gate_kernel(): global _MOE_GATE_KERNEL if _MOE_GATE_KERNEL is None: _MOE_GATE_KERNEL = mx.fast.metal_kernel( name="mimo_moe_gate_fused", input_names=["gates", "bias"], output_names=["out_inds", "out_scores"], source=_MOE_GATE_KERNEL_SOURCE, ensure_row_contiguous=True, ) return _MOE_GATE_KERNEL _MOE_GATE_KERNEL_SOURCE = """ uint tg = threadgroup_position_in_grid.x; uint tid = thread_position_in_threadgroup.x; threadgroup float shared_scored[256]; threadgroup float shared_sig[256]; threadgroup uint chosen_inds[8]; threadgroup float chosen_sigs[8]; threadgroup float scratch_val[256]; threadgroup uint scratch_idx[256]; threadgroup float sum_buf; float g = float(gates[tg * 256 + tid]); float b = float(bias[tid]); float sig = 1.0f / (1.0f + metal::exp(-g)); float scored = sig + b; shared_scored[tid] = scored; shared_sig[tid] = sig; threadgroup_barrier(metal::mem_flags::mem_threadgroup); for (int k = 0; k < 8; k++) { scratch_val[tid] = shared_scored[tid]; scratch_idx[tid] = tid; threadgroup_barrier(metal::mem_flags::mem_threadgroup); for (uint stride = 128; stride > 0; stride >>= 1) { if (tid < stride) { float a = scratch_val[tid]; float c = scratch_val[tid + stride]; if (c > a) { scratch_val[tid] = c; scratch_idx[tid] = scratch_idx[tid + stride]; } } threadgroup_barrier(metal::mem_flags::mem_threadgroup); } if (tid == 0) { uint winner = scratch_idx[0]; chosen_inds[k] = winner; chosen_sigs[k] = shared_sig[winner]; shared_scored[winner] = -1.0e30f; } threadgroup_barrier(metal::mem_flags::mem_threadgroup); } if (tid == 0) { float s = 0.0f; for (int i = 0; i < 8; i++) s += chosen_sigs[i]; sum_buf = s + 1.0e-20f; } threadgroup_barrier(metal::mem_flags::mem_threadgroup); if (tid < 8) { out_inds[tg * 8 + tid] = chosen_inds[tid]; out_scores[tg * 8 + tid] = chosen_sigs[tid] / sum_buf; } """ def _fused_group_expert_select( gates, e_score_correction_bias, top_k, n_group, topk_group, routed_scaling_factor, norm_topk_prob, ): # Fast path only for the MiMo-V2.5 config we're targeting. if (gates.shape[-1] != 256 or top_k != 8 or n_group != 1 or not norm_topk_prob or routed_scaling_factor != 1.0): return _orig_group_expert_select( gates, e_score_correction_bias, top_k, n_group, topk_group, routed_scaling_factor, norm_topk_prob, ) B = gates.size // 256 kernel = _get_moe_gate_kernel() # Kernel reads gates as flat (B*256,) and writes outputs as flat (B*8,). # The output shape we declare to MLX determines its array shape; we pass # the original leading dims directly so no reshape is needed at any point. out_shape_inds = (*gates.shape[:-1], 8) inds, scores = kernel( inputs=[gates, e_score_correction_bias], output_shapes=[out_shape_inds, out_shape_inds], output_dtypes=[mx.uint32, mx.float32], grid=(B * 256, 1, 1), threadgroup=(256, 1, 1), ) return inds, scores # Install the monkey-patch at import time. from . import mimo_v2 as _mimo_v2_mod _orig_group_expert_select = _mimo_v2_mod.group_expert_select _mimo_v2_mod.group_expert_select = _fused_group_expert_select # -------------------------------------------------------------------------- # MoE output combination: fp32 accumulator (Xiaomi reference numerics). # # Upstream mimo_v2.MoE.__call__ does: # y = (y * scores[..., None]).sum(axis=-2).astype(x.dtype) # which accumulates in bf16 and silently drifts vs Xiaomi's intended math. # # Xiaomi's MiMo-V2-Flash technical report specifies FP32 precision for the # MoE router. Reference implementations (vLLM, ROCm/aiter, Cursor warp-decode) # all accumulate the weighted expert sum in fp32 and cast once at the end. # # We monkey-patch MoE.__call__ here so the block_fp8 path runs faithful # numerics. Both correctness AND speed improve (one einsum kernel vs a # chain of expand_dims + broadcast_multiply + reduce_sum + astype). # -------------------------------------------------------------------------- def _moe_call_fp32(self, x): if self.sharding_group is not None: x = sum_gradients(self.sharding_group)(x) inds, scores = self.gate(x) y = self.switch_mlp(x, inds) # fp32-accumulator weighted sum: equivalent to # (y.float() * scores.float()[..., None]).sum(axis=-2).astype(x.dtype) # but einsum does this in one fused kernel with fp32 accum internally. y = mx.einsum("...kd,...k->...d", y, scores).astype(x.dtype) if self.sharding_group is not None: y = mx.distributed.all_sum(y, group=self.sharding_group) return y MoE.__call__ = _moe_call_fp32 # -------------------------------------------------------------------------- # Extension-backed quantized layers. These are STANDALONE nn.Modules — they do # NOT subclass nn.QuantizedLinear / QuantizedSwitchLinear and never call # mx.quantized_matmul / mx.gather_qmm with mode="block_fp8". That mode is not in # upstream MLX (it was the rejected PR #3600), so going through it would require # a forked MLX. Instead every path uses the mlx_block_fp8 extension kernels, # which build against stock upstream MLX. Result: the model runs on stock MLX + # this extension, no fork. # # Weights are pre-quantized E4M3 (loaded from the safetensors); we just hold the # uint8 codes + fp32 scales and dispatch to the kernels. # -------------------------------------------------------------------------- import mlx.nn as _nn class ExtBlockFp8Linear(_nn.Module): """Standalone dense block_fp8 linear. Decode (M==1) -> block_fp8_qmv_fast; prefill (M>1) -> block_fp8_qmm_t. No native block_fp8 mode.""" def __init__(self, codes, scales, bias=None): super().__init__() self.weight = codes # uint8 E4M3 [N, K] self.scales = scales # fp32 [N/128, K/128] if bias is not None: self.bias = bias def __call__(self, x): codes = self["weight"] scales = self["scales"] *batch, K = x.shape x2 = x.reshape(-1, K) M = x2.shape[0] if M == 1: y = _bfp8.block_fp8_qmv_fast(x2, codes, scales) else: y = _bfp8.block_fp8_qmm_t(x2, codes, scales) N = codes.shape[0] y = y.reshape(*batch, N) if "bias" in self: y = y + self["bias"] return y class ExtBlockFp8SwitchLinear(_nn.Module): """Standalone MoE block_fp8 switch linear. Both prefill (sorted) and decode (unsorted) route to block_fp8_gather_qmm_rhs — the prefill kernel is M-agnostic and serves decode once indices are sorted. No native block_fp8 mode, no fork dependency.""" def __init__(self, codes, scales, bias=None): super().__init__() self.weight = codes # uint8 E4M3 [E, N, K] self.scales = scales # fp32 [E, N/128, K/128] if bias is not None: self.bias = bias def __call__(self, x, indices, sorted_indices=False): codes = self["weight"] scales = self["scales"] if sorted_indices: # Prefill: indices already globally sorted upstream by _gather_sort. y = _bfp8.block_fp8_gather_qmm_rhs( x, codes, scales, indices, transpose=True) if "bias" in self: y = y + self["bias"] return y # Decode (unsorted): the purpose-built block_fp8_gather_qmv_fast kernel — # one kernel call, no sort, no dequant, at native gather_qmm speed # (~0.3 ms, bit-exact). SwitchGLU passes indices [...lead, top_k] and x # [...lead, Xtk, 1, K] where Xtk is 1 (up/gate: x broadcast across # experts) or top_k (down: x already materialized per expert). Build the # flat lhs (x-row) / rhs (expert) index arrays accordingly. N = codes.shape[1] idx_shape = indices.shape top_k = idx_shape[-1] lead = idx_shape[:-1] K_ = x.shape[-1] Xtk = x.shape[len(lead)] n_lead = 1 for d in lead: n_lead *= d rhs = indices.reshape(-1).astype(mx.uint32) if Xtk == 1: # up/gate: one x-row per lead entry, routed to top_k experts x_rows = x.reshape(n_lead, 1, K_) lhs = mx.repeat(mx.arange(n_lead, dtype=mx.uint32), top_k) else: # down: x already has one row per (lead, expert) x_rows = x.reshape(n_lead * top_k, 1, K_) lhs = mx.arange(n_lead * top_k, dtype=mx.uint32) y = _bfp8.block_fp8_gather_qmv_fast(x_rows, codes, scales, lhs, rhs) y = y.reshape(*idx_shape, 1, N) if "bias" in self: y = y + self["bias"] return y class FusedQKVAttention(Attention): """Same logic as upstream Attention but with a single fused qkv_proj Linear instead of three independent q/k/v projections.""" def __init__(self, args: ModelArgs, is_sliding_window: bool): # Skip super().__init__() — we replicate its body without # allocating three unused q/k/v Linears that would just be replaced. nn.Module.__init__(self) dim = args.hidden_size self.is_sliding_window = is_sliding_window if is_sliding_window: self.n_heads = args.swa_num_attention_heads self.n_kv_heads = args.swa_num_key_value_heads head_dim = args.swa_head_dim v_head_dim = args.swa_v_head_dim rope_theta = args.swa_rope_theta has_sinks = args.add_swa_attention_sink_bias else: self.n_heads = args.num_attention_heads self.n_kv_heads = args.num_key_value_heads head_dim = args.head_dim v_head_dim = args.v_head_dim rope_theta = args.rope_theta has_sinks = args.add_full_attention_sink_bias self.head_dim = head_dim self.v_head_dim = v_head_dim self.scale = head_dim ** -0.5 self.v_scale = args.attention_value_scale self.q_dim = self.n_heads * head_dim self.k_dim = self.n_kv_heads * head_dim self.v_dim = self.n_kv_heads * v_head_dim # Weights are stored in Xiaomi's TP rank-major layout: # per-rank [Q|K|V] blocks concatenated across TP=4 ranks. # We reshape (TP, per_rank, K) and slice at forward time. self.tp = 4 self.qkv_proj = nn.Linear( dim, self.q_dim + self.k_dim + self.v_dim, bias=args.attention_bias, ) self.o_proj = nn.Linear(self.n_heads * v_head_dim, dim, bias=False) self.attention_sink_bias = ( mx.zeros((self.n_heads,)) if has_sinks else None ) self.rope = initialize_rope( int(args.partial_rotary_factor * head_dim), base=rope_theta, traditional=False, scaling_config=args.rope_scaling, max_position_embeddings=args.max_position_embeddings, ) def __call__( self, x: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None, ) -> mx.array: B, L, _ = x.shape qkv = self.qkv_proj(x) # Rank-major layout: per-rank [Q|K|V] blocks across TP=4 ranks. per_rank = (self.q_dim + self.k_dim + self.v_dim) // self.tp q_pr = self.q_dim // self.tp k_pr = self.k_dim // self.tp qkv_r = qkv.reshape(B, L, self.tp, per_rank) q_flat = qkv_r[..., :q_pr].reshape(B, L, self.q_dim) k_flat = qkv_r[..., q_pr:q_pr+k_pr].reshape(B, L, self.k_dim) v_flat = qkv_r[..., q_pr+k_pr:].reshape(B, L, self.v_dim) queries = q_flat.reshape(B, L, self.n_heads, self.head_dim).swapaxes(1, 2) keys = k_flat.reshape(B, L, self.n_kv_heads, self.head_dim).swapaxes(1, 2) values = v_flat.reshape(B, L, self.n_kv_heads, self.v_head_dim).swapaxes(1, 2) if self.v_scale is not None: values = values * self.v_scale if cache is not None: queries = self.rope(queries, offset=cache.offset) keys = self.rope(keys, offset=cache.offset) keys, values = cache.update_and_fetch(keys, values) else: queries = self.rope(queries) keys = self.rope(keys) output = scaled_dot_product_attention( queries, keys, values, cache=cache, scale=self.scale, mask=mask, sinks=self.attention_sink_bias, ) return self.o_proj(output.swapaxes(1, 2).reshape(B, L, -1)) class Model(MiMoModel): """MiMo-V2 with fused QKV attention + native block_fp8 weight loading.""" def __init__(self, args: ModelArgs): super().__init__(args) # Replace each layer's Attention with FusedQKVAttention. # Done after super() because the parent's __init__ instantiates # layers with the base Attention class; we swap in place. for layer_idx, layer in enumerate(self.model.layers): if layer is None: continue is_sliding_window = bool(args.hybrid_layer_pattern[layer_idx]) layer.self_attn = FusedQKVAttention(args, is_sliding_window) # ----- Load-side helpers ----- SKIP_PREFIXES = ( "model.mtp.", "visual.", "audio_encoder.", "speech_embeddings.", ) def sanitize_block_fp8(self, weights): """Pre-stacked file means sanitize is trivial: just drop auxiliaries.""" return { k: v for k, v in weights.items() if not k.startswith(self.SKIP_PREFIXES) } def apply_block_fp8(self, weights): """Install pre-quantized weights into block_fp8 QuantizedLinear / QuantizedSwitchLinear modules. Bf16 weights flow through normally.""" # First pass: collect every (path, module) that needs swapping. # We can't safely mutate during tree_map_with_path traversal, so # gather then apply. swaps = [] # list of (path, new_module) def visit(path, module): wkey = f"{path}.weight" skey = f"{path}.weight_scale_inv" if skey not in weights: return module # leave alone — bf16 passthrough w = weights[wkey] s = weights[skey] new = None if isinstance(module, nn.Linear): out_dims, in_dims = w.shape has_bias = hasattr(module, "bias") bias = module.bias if has_bias else None new = ExtBlockFp8Linear(w, s, bias=bias) elif isinstance(module, SwitchLinear): n_experts, out_dims, in_dims = w.shape has_bias = hasattr(module, "bias") bias = module.bias if has_bias else None new = ExtBlockFp8SwitchLinear(w, s, bias=bias) return new if new is not None else module leaves = self.leaf_modules() leaves = tree_map_with_path( visit, leaves, is_leaf=nn.Module.is_module, ) self.update_modules(leaves) # Second pass: load non-quantized weights. Exclude: # - any *_scale_inv (consumed above) # - any .weight whose _scale_inv was consumed (already installed) consumed = { k[:-len("_scale_inv")] for k in weights if k.endswith("weight_scale_inv") } # consumed entries are like "model.layers.5.mlp.switch_mlp.gate_proj.weight" remaining = { k: v for k, v in weights.items() if not k.endswith("weight_scale_inv") and k not in consumed } self.load_weights(list(remaining.items()), strict=False)