diff --git a/src/exo/worker/engines/mlx/patches/__init__.py b/src/exo/worker/engines/mlx/patches/__init__.py new file mode 100644 index 00000000..38d96dc7 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/__init__.py @@ -0,0 +1,38 @@ +"""Model-specific kernel fusion patches for MLX inference. + +Detects model type after loading and applies optimized kernel patches +where available. Currently supports: +- Qwen3.5 MoE (model_type: qwen3_5_moe): oproj fusion, 4 custom Metal dispatches + +Set EXO_FUSED_KERNELS=0 to disable patches (baseline mode). +Default: enabled (EXO_FUSED_KERNELS=1). +""" + +import json +import os +from pathlib import Path + +import mlx.nn as nn +from loguru import logger + + +def maybe_apply_patches(model: nn.Module, model_path: Path) -> None: + """Detect model type and apply kernel fusion patches if available.""" + if os.environ.get("EXO_FUSED_KERNELS", "1") == "0": + logger.info("Kernel fusion patches disabled (EXO_FUSED_KERNELS=0)") + return + + config_path = model_path / "config.json" + if not config_path.exists(): + return + + with open(config_path) as f: + config = json.load(f) + + model_type = config.get("model_type", "") + + if model_type == "qwen3_5_moe": + from .qwen3_5_moe.apply import apply_qwen35_oproj_patches + + logger.info("Detected Qwen3.5 MoE model, applying oproj fusion patches") + apply_qwen35_oproj_patches(model) diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/__init__.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/apply.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/apply.py new file mode 100644 index 00000000..e230f392 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/apply.py @@ -0,0 +1,27 @@ +"""Apply oproj fusion patches to Qwen3.5 MoE models. + +Entry point called from patches/__init__.py after model type detection. +""" + +import time + +import mlx.nn as nn +from loguru import logger + +from .common import apply_oproj_fused_patches + + +def apply_qwen35_oproj_patches(model: nn.Module) -> None: + """Apply oproj 4-dispatch fusion to all layers of a Qwen3.5 MoE model. + + Patches decoder, attention, and MoE __call__ methods to use custom Metal + kernels for decode (seq_len=1). Prefill (seq_len>1) falls back to vanilla. + """ + layers = model.layers # type: ignore[attr-defined] + n_layers = len(layers) + + t0 = time.time() + apply_oproj_fused_patches(layers, gate_bm=8, free_originals=False) + t_patch = time.time() - t0 + + logger.info(f"Qwen3.5 oproj fusion: patched {n_layers} layers in {t_patch:.1f}s") diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/common.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/common.py new file mode 100644 index 00000000..1c47d4a0 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/common.py @@ -0,0 +1,187 @@ +"""Weight preparation and patch orchestration for Qwen3.5 oproj fusion. + +Adapted from mlx_bench/model_patches/qwen/common.py. +""" + +import mlx.core as mx +import mlx.nn as nn +from loguru import logger +from mlx_lm.models.qwen3_5 import DecoderLayer +from mlx_lm.models.qwen3_next import Qwen3NextSparseMoeBlock + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def _patch_swiglu_weights(moe): + """Stack gate+up weights for fused 8-bit SwiGLU kernel.""" + gate_proj = moe.switch_mlp.gate_proj + up_proj = moe.switch_mlp.up_proj + + moe.switch_mlp._fused_w_gu = mx.concatenate( + [gate_proj.weight, up_proj.weight], axis=1) + moe.switch_mlp._fused_s_gu = mx.concatenate( + [gate_proj.scales, up_proj.scales], axis=1) + moe.switch_mlp._fused_b_gu = mx.concatenate( + [gate_proj.biases, up_proj.biases], axis=1) + moe.switch_mlp._fused_n_inter = gate_proj.output_dims + moe.switch_mlp._fused_k_hidden = gate_proj.input_dims + moe.switch_mlp._fused_group_size = gate_proj.group_size + + mx.eval(moe.switch_mlp._fused_w_gu, + moe.switch_mlp._fused_s_gu, + moe.switch_mlp._fused_b_gu) + + +def _patch_shared_expert(moe): + """Prepare shared expert quantized weights for fused 8-bit path.""" + shared = moe.shared_expert + gp = shared.gate_proj + up = shared.up_proj + dp = shared.down_proj + + moe._shared_w_gu = mx.concatenate([gp.weight, up.weight], axis=0) + moe._shared_s_gu = mx.concatenate([gp.scales, up.scales], axis=0) + moe._shared_b_gu = mx.concatenate([gp.biases, up.biases], axis=0) + + moe._shared_down_w = dp.weight + moe._shared_down_s = dp.scales + moe._shared_down_b = dp.biases + + moe._shared_inter = gp.weight.shape[0] + moe._shared_gs = gp.group_size + + mx.eval(moe._shared_w_gu, moe._shared_s_gu, moe._shared_b_gu, + moe._shared_down_w, moe._shared_down_s, moe._shared_down_b) + + +def _patch_down_proj(moe): + """Extract down_proj weights for merged 8-bit kernel dispatch.""" + dp = moe.switch_mlp.down_proj + moe._down_w = dp.weight + moe._down_s = dp.scales + moe._down_b = dp.biases + moe._down_K = dp.output_dims + moe._down_N = dp.input_dims + moe._down_gs = dp.group_size + mx.eval(moe._down_w, moe._down_s, moe._down_b) + + +def _patch_oproj_gate_rms(layer, gate_bm=8): + """Precompute M1/W_fused for fused o_proj + gate GEMV (oproj 4-dispatch mode). + + Gate decomposition: + gate_score[e] = W_gate[e,:] @ rms_norm(h) + where h = residual + W_oproj @ attn_out + rms_norm(h) = h * w_rms * inv_rms + + Expanding: + gate_score[e] = (W_fused @ residual + M1 @ attn_out) * inv_rms + + Precomputed offline (per layer, stored on moe block): + W_fused = dequant(W_gate) · diag(w_rms) — (E, K) bf16 + M1 = W_fused @ dequant(W_oproj) — (E, K_attn) bf16 + """ + moe = layer.mlp + + if layer.is_linear: + oproj = layer.linear_attn.out_proj + else: + oproj = layer.self_attn.o_proj + + # Dequantize gate and o_proj (temporary, for M1 computation) + gate = moe.gate + W_gate_f32 = mx.dequantize( + gate.weight, gate.scales, gate.biases, + group_size=gate.group_size, bits=gate.bits, + ).astype(mx.float32) + + W_oproj_f32 = mx.dequantize( + oproj.weight, oproj.scales, oproj.biases, + group_size=oproj.group_size, bits=oproj.bits, + ).astype(mx.float32) + mx.eval(W_gate_f32, W_oproj_f32) + + rms_weight = layer.post_attention_layernorm.weight.astype(mx.bfloat16) + + w_rms_f32 = rms_weight.astype(mx.float32) + W_fused = (W_gate_f32 * w_rms_f32).astype(mx.bfloat16) + mx.eval(W_fused) + del W_gate_f32 + + M1 = (W_fused.astype(mx.float32) @ W_oproj_f32).astype(mx.bfloat16) + mx.eval(M1) + del W_oproj_f32 + + moe._oproj_M1 = M1 + moe._oproj_W_fused = W_fused + moe._oproj_rms_weight = rms_weight + + moe._oproj_w = oproj.weight + moe._oproj_s = oproj.scales + moe._oproj_b = oproj.biases + moe._oproj_K_attn = oproj.weight.shape[1] * 4 + + seg = moe.shared_expert_gate + moe._seg_w = seg.weight + moe._seg_s = seg.scales + moe._seg_b = seg.biases + + M = oproj.weight.shape[0] + K_hidden = W_fused.shape[1] + n_experts = W_fused.shape[0] + moe._oproj_M = M + moe._oproj_K_hidden = K_hidden + moe._oproj_n_experts = n_experts + moe._oproj_n_tg = ceil_div(M, 32) + moe._oproj_gate_bm = gate_bm + + mx.eval(moe._oproj_rms_weight) + + +def apply_oproj_fused_patches(layers, gate_bm=8, free_originals=False): + """Apply oproj-mode fused MoE patches (4-dispatch) to all layers. + + 1. Prepare SwiGLU weights (_patch_swiglu_weights) + 2. Prepare shared expert 8-bit weights (_patch_shared_expert) + 3. Prepare down_proj weights (_patch_down_proj) + 4. Precompute M1/W_fused + store o_proj/gate weights (_patch_oproj_gate_rms) + 5. Replace __call__ methods for decoder + MoE + attention + """ + from .moe import _oproj_moe_call + from .decoder import ( + _oproj_decoder_call, + _pre_oproj_attention_call, + _pre_oproj_qwen35_linear_attn_call, + ) + from mlx_lm.models.qwen3_next import Qwen3NextAttention + from mlx_lm.models.qwen3_5 import GatedDeltaNet + + n_patched = 0 + for li, layer in enumerate(layers): + moe = layer.mlp + if isinstance(moe, Qwen3NextSparseMoeBlock): + _patch_swiglu_weights(moe) + _patch_shared_expert(moe) + _patch_down_proj(moe) + _patch_oproj_gate_rms(layer, gate_bm=gate_bm) + + if free_originals: + for attr in ('weight', 'scales', 'biases'): + for proj in (moe.switch_mlp.gate_proj, + moe.switch_mlp.up_proj): + try: + delattr(proj, attr) + except AttributeError: + pass + + n_patched += 1 + if (li + 1) % 10 == 0 or li == 0: + logger.info(f" Patched layer {li+1}/{len(layers)} (oproj mode)") + + Qwen3NextAttention.__call__ = _pre_oproj_attention_call + GatedDeltaNet.__call__ = _pre_oproj_qwen35_linear_attn_call + Qwen3NextSparseMoeBlock.__call__ = _oproj_moe_call + DecoderLayer.__call__ = _oproj_decoder_call + logger.info(f" Patched {n_patched} MoE blocks (oproj mode, 4 dispatches)") diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/decoder.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/decoder.py new file mode 100644 index 00000000..d00c543c --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/decoder.py @@ -0,0 +1,168 @@ +"""Decoder layer __call__ variants for Qwen3.5. + +Two modes: + _fused_decoder_call: passes residual to fused MoE epilogue (~15 dispatches) + _oproj_decoder_call: fuses o_proj + RMSNorm + gate GEMV (4 dispatches) + +Attention patches for oproj mode: + _pre_oproj_attention_call: Qwen3NextAttention.__call__ that skips o_proj + _pre_oproj_qwen35_linear_attn_call: qwen3_5.GatedDeltaNet.__call__ that skips out_proj + +Note: qwen3_5.GatedDeltaNet (used by DecoderLayer) is a DIFFERENT class from +qwen3_next.Qwen3NextGatedDeltaNet. They have different projection layouts: + - qwen3_5.GatedDeltaNet: separate in_proj_qkv, in_proj_z, in_proj_b, in_proj_a + - qwen3_next.Qwen3NextGatedDeltaNet: merged in_proj_qkvz, in_proj_ba +The patch must match qwen3_5.GatedDeltaNet's __call__ structure. + +Adapted from mlx_bench/model_patches/qwen/decoder.py. +""" + +from typing import Any, Optional + +import mlx.core as mx +import mlx.nn as nn +from mlx.nn.layers.activations import silu as nn_silu + +# Map moe block id → parent decoder layer (avoids circular refs in model tree) +_parent_layer_map = {} + + +def _fused_decoder_call(self, x, mask=None, cache=None): + """Decoder layer with residual passed to fused MoE epilogue. + + Replaces: + h = x + attn(norm(x)) + out = h + mlp(norm(h)) # mlp returns MoE output, then adds h + + With: + h = x + attn(norm(x)) + out = mlp(norm(h), _residual=h) # epilogue fuses: moe_out + h + """ + if self.is_linear: + r = self.linear_attn(self.input_layernorm(x), mask, cache) + else: + r = self.self_attn(self.input_layernorm(x), mask, cache) + h = x + r + out = self.mlp(self.post_attention_layernorm(h), _residual=h) + return out # already includes residual add from epilogue + + +def _oproj_decoder_call(self, x, mask=None, cache=None): + """Decoder with fused o_proj + RMSNorm + gate GEMV (oproj 4-dispatch mode). + + Skips o_proj, addmm, and post_attention_layernorm — all fused into Dispatch 1. + Attention __call__ is patched to return pre-o_proj output. + + Flow: + pre_oproj = attn(input_layernorm(x)) # returns BEFORE o_proj + MoE receives (pre_oproj, residual=x) and handles o_proj + RMSNorm + gate internally + """ + if self.is_linear: + pre_oproj = self.linear_attn(self.input_layernorm(x), mask, cache) + else: + pre_oproj = self.self_attn(self.input_layernorm(x), mask, cache) + _parent_layer_map[id(self.mlp)] = self + return self.mlp(pre_oproj, _residual=x) + + +def _pre_oproj_attention_call(self, x, mask=None, cache=None): + """Qwen3NextAttention.__call__ that returns pre-o_proj output. + + Identical to original except final line returns output*sigmoid(gate) + instead of self.o_proj(output*sigmoid(gate)). + """ + B, L, D = x.shape + q_proj_output = self.q_proj(x) + queries, gate = mx.split( + q_proj_output.reshape(B, L, self.num_attention_heads, -1), 2, axis=-1 + ) + gate = gate.reshape(B, L, -1) + keys, values = self.k_proj(x), self.v_proj(x) + queries = self.q_norm(queries).transpose(0, 2, 1, 3) + keys = self.k_norm( + keys.reshape(B, L, self.num_key_value_heads, -1) + ).transpose(0, 2, 1, 3) + values = values.reshape(B, L, self.num_key_value_heads, -1).transpose( + 0, 2, 1, 3 + ) + 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) + + from mlx_lm.models.qwen3_next import scaled_dot_product_attention + output = scaled_dot_product_attention( + queries, keys, values, cache=cache, scale=self.scale, mask=mask + ) + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + return output * mx.sigmoid(gate) # skip o_proj + + +def _pre_oproj_qwen35_linear_attn_call( + self, + inputs: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, +) -> mx.array: + """qwen3_5.GatedDeltaNet.__call__ that returns pre-out_proj output. + + Identical to qwen3_5.GatedDeltaNet.__call__ except final line returns + out.reshape(B,S,-1) instead of self.out_proj(out.reshape(B,S,-1)). + + Note: this targets qwen3_5.GatedDeltaNet (separate projections), NOT + qwen3_next.Qwen3NextGatedDeltaNet (merged projections). They are + different classes with different __call__ bodies. + """ + from mlx_lm.models.gated_delta import gated_delta_update + + B, S, _ = inputs.shape + + qkv = self.in_proj_qkv(inputs) + z = self.in_proj_z(inputs).reshape(B, S, self.num_v_heads, self.head_v_dim) + b = self.in_proj_b(inputs) + a = self.in_proj_a(inputs) + + if cache is not None and cache[0] is not None: + conv_state = cache[0] + else: + conv_state = mx.zeros( + (B, self.conv_kernel_size - 1, self.conv_dim), + dtype=inputs.dtype, + ) + + if mask is not None: + qkv = mx.where(mask[..., None], qkv, 0) + conv_input = mx.concatenate([conv_state, qkv], axis=1) + if cache is not None: + cache[0] = conv_input[:, -(self.conv_kernel_size - 1) :] + conv_out = nn.silu(self.conv1d(conv_input)) + + q, k, v = [ + t.reshape(B, S, h, d) + for t, h, d in zip( + mx.split(conv_out, [self.key_dim, 2 * self.key_dim], -1), + [self.num_k_heads, self.num_k_heads, self.num_v_heads], + [self.head_k_dim, self.head_k_dim, self.head_v_dim], + ) + ] + + state = cache[1] if cache else None + inv_scale = k.shape[-1] ** -0.5 + q = (inv_scale**2) * mx.fast.rms_norm(q, None, 1e-6) + k = inv_scale * mx.fast.rms_norm(k, None, 1e-6) + + out, state = gated_delta_update( + q, k, v, a, b, + self.A_log, self.dt_bias, + state, mask, + use_kernel=not self.training, + ) + + if cache is not None: + cache[1] = state + + out = self.norm(out, z) + return out.reshape(B, S, -1) # skip out_proj diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/__init__.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/custom_oproj_gate_gemv_8bit.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/custom_oproj_gate_gemv_8bit.py new file mode 100644 index 00000000..ddd51bd9 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/custom_oproj_gate_gemv_8bit.py @@ -0,0 +1,324 @@ +"""Dispatch 1: Fused o_proj (8-bit) + gate GEMV parts + x² partials for Qwen3.5. + +Port of Kimi's custom_oproj_gate_gemv.py adapted for 8-bit quantized o_proj. + +Single dispatch with 3 GEMV types sharing 256-thread TGs (8 SGs of 32): + +TGs 0 to N_OPROJ_TG-1: o_proj GEMV (8-bit, M=4096, K=8192) + 8-bit affine dequant: result = scale * Σ(x*w_uint8) + bias * Σ(x) + Epilogue: h = oproj_result + residual, h_scaled = h*w_rms, h_out = h, x²_acc + TG reduction: 8 SG x² → 1 float per TG. + +TGs N_OPROJ_TG to +N_M1_TG-1: M1 GEMV (bf16, E × K_attn) + M1 = W_fused @ W_oproj (pre-computed). Input = attn_out. + Output: gate_part_a (E,) f32. + +TGs +N_M1_TG to end: W_fused GEMV (bf16, E × K) + W_fused × residual → gate_part_b (E,) f32. +""" +import mlx.core as mx + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def _gen_custom_oproj_8bit_source(n_experts=64, group_size=64, scale_bf16=True): + """Generate Metal source for fused 8-bit o_proj + bf16 gate GEMVs.""" + E = int(n_experts) + gs = int(group_size) + sc_t = "bfloat16_t" if scale_bf16 else "float" + # 8-bit dequant params for o_proj + oproj_slid_divisor = gs // 8 # 64/8 = 8 + oproj_sc_stride = 256 // gs # 256/64 = 4 + + return f""" + // ── Constants ── + const int TM = 4; // rows per SG for bf16 GEMVs + const int TN = 4; // elements per thread per iter for bf16 GEMVs + const int blockN = 128; // 32 threads × TN=4 + const int E_CONST = {E}; + + int M = M_val; // 4096 (hidden_size) + int K_attn = K_attn_val; // 8192 (o_proj input dim) + int K_hidden = K_hidden_val; // 4096 (hidden_size, same as M for Qwen) + int N_OPROJ_TG = N_OPROJ_TG_val; + int N_M1_TG = N_M1_TG_val; + int blockM_gate = BM_GATE_val * TM; // rows per gate GEMV TG + + uint tg_x = threadgroup_position_in_grid.x; + uint sgid = simdgroup_index_in_threadgroup; // 0..7 + uint slid = thread_index_in_simdgroup; // 0..31 + + if (tg_x < (uint)N_OPROJ_TG) {{ + // ════════════════════════════════════════════════════════════════ + // O_PROJ GEMV: 8-bit quantized, M=4096, K_attn=8192 + // Each TG handles 32 rows (8 SGs × 4 rows each) + // 8-bit affine dequant: result = scale * Σ(x*w) + bias * Σ(x) + // ════════════════════════════════════════════════════════════════ + const int blockM = 32; // 8 SGs × TM=4 + const int VPT = 8; // values per thread per iteration + const int BLOCK_SIZE = 256; // 32 threads × 8 values + + int out_row = int(tg_x) * blockM + int(sgid) * TM; + if (out_row >= M) return; + out_row = (out_row + TM <= M) ? out_row : (M - TM); + + threadgroup float tgp_x2[8]; + + // 8-bit GEMV K-loop: K-outer, tm-inner (input loaded once, reused across TM rows) + float acc[TM] = {{0.0f, 0.0f, 0.0f, 0.0f}}; + float result[TM]; + int K_groups = K_attn / {gs}; + + // Per-row weight/scale/bias pointers + const device uint8_t* ws0 = (const device uint8_t*)W_oproj + (long)(out_row + 0) * K_attn + slid * VPT; + const device uint8_t* ws1 = (const device uint8_t*)W_oproj + (long)(out_row + 1) * K_attn + slid * VPT; + const device uint8_t* ws2 = (const device uint8_t*)W_oproj + (long)(out_row + 2) * K_attn + slid * VPT; + const device uint8_t* ws3 = (const device uint8_t*)W_oproj + (long)(out_row + 3) * K_attn + slid * VPT; + + const device {sc_t}* sc0 = (const device {sc_t}*)S_oproj + (long)(out_row + 0) * K_groups + slid / {oproj_slid_divisor}; + const device {sc_t}* sc1 = (const device {sc_t}*)S_oproj + (long)(out_row + 1) * K_groups + slid / {oproj_slid_divisor}; + const device {sc_t}* sc2 = (const device {sc_t}*)S_oproj + (long)(out_row + 2) * K_groups + slid / {oproj_slid_divisor}; + const device {sc_t}* sc3 = (const device {sc_t}*)S_oproj + (long)(out_row + 3) * K_groups + slid / {oproj_slid_divisor}; + + const device {sc_t}* bi0 = (const device {sc_t}*)B_oproj + (long)(out_row + 0) * K_groups + slid / {oproj_slid_divisor}; + const device {sc_t}* bi1 = (const device {sc_t}*)B_oproj + (long)(out_row + 1) * K_groups + slid / {oproj_slid_divisor}; + const device {sc_t}* bi2 = (const device {sc_t}*)B_oproj + (long)(out_row + 2) * K_groups + slid / {oproj_slid_divisor}; + const device {sc_t}* bi3 = (const device {sc_t}*)B_oproj + (long)(out_row + 3) * K_groups + slid / {oproj_slid_divisor}; + + int xb = slid * VPT; + + for (int k = 0; k < K_attn; k += BLOCK_SIZE) {{ + // Load input ONCE per K-block + float xv[VPT]; + float xsum = 0.0f; + for (int i = 0; i < VPT; i++) {{ + xv[i] = float(attn_out[xb + i]); + xsum += xv[i]; + }} + + // Row 0 + {{ float wacc = 0.0f; + for (int i = 0; i < VPT; i++) wacc += xv[i] * float(ws0[i]); + acc[0] += float(*sc0) * wacc + xsum * float(*bi0); + ws0 += BLOCK_SIZE; sc0 += {oproj_sc_stride}; bi0 += {oproj_sc_stride}; }} + // Row 1 + {{ float wacc = 0.0f; + for (int i = 0; i < VPT; i++) wacc += xv[i] * float(ws1[i]); + acc[1] += float(*sc1) * wacc + xsum * float(*bi1); + ws1 += BLOCK_SIZE; sc1 += {oproj_sc_stride}; bi1 += {oproj_sc_stride}; }} + // Row 2 + {{ float wacc = 0.0f; + for (int i = 0; i < VPT; i++) wacc += xv[i] * float(ws2[i]); + acc[2] += float(*sc2) * wacc + xsum * float(*bi2); + ws2 += BLOCK_SIZE; sc2 += {oproj_sc_stride}; bi2 += {oproj_sc_stride}; }} + // Row 3 + {{ float wacc = 0.0f; + for (int i = 0; i < VPT; i++) wacc += xv[i] * float(ws3[i]); + acc[3] += float(*sc3) * wacc + xsum * float(*bi3); + ws3 += BLOCK_SIZE; sc3 += {oproj_sc_stride}; bi3 += {oproj_sc_stride}; }} + + xb += BLOCK_SIZE; + }} + + for (int tm = 0; tm < TM; tm++) + result[tm] = simd_sum(acc[tm]); + + // Epilogue: addmm + x² + h_scaled + h_out + float x2_acc = 0.0f; + if (slid == 0) {{ + for (int tm = 0; tm < TM; tm++) {{ + int k = out_row + tm; + float h = result[tm] + float(residual[k]); + x2_acc += h * h; + h_scaled[k] = static_cast(h * float(w_rms[k])); + h_out[k] = static_cast(h); + }} + }} + + // TG x² reduction: 8 SGs → 1 value + if (slid == 0) tgp_x2[sgid] = x2_acc; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sgid == 0 && slid == 0) {{ + float total = 0.0f; + for (int s = 0; s < 8; s++) total += tgp_x2[s]; + x2_partials[tg_x] = total; + }} + + }} else if (tg_x < (uint)(N_OPROJ_TG + N_M1_TG)) {{ + // ════════════════════════════════════════════════════════════════ + // M1 GEMV: bf16, M=E, K=K_attn — M1 × attn_out → gate_part_a + // ════════════════════════════════════════════════════════════════ + int local_tg = int(tg_x) - N_OPROJ_TG; + int out_row = local_tg * blockM_gate + int(sgid) * TM; + if (out_row >= E_CONST) return; + out_row = (out_row + TM <= E_CONST) ? out_row : (E_CONST - TM); + + float result[TM] = {{0.0f, 0.0f, 0.0f, 0.0f}}; + int bn = int(slid) * TN; + int n_iter = K_attn / blockN; + + for (int i = 0; i < n_iter; i++) {{ + float v[TN]; + for (int tn = 0; tn < TN; tn++) + v[tn] = float(attn_out[bn + tn]); + for (int tm = 0; tm < TM; tm++) {{ + float acc = 0.0f; + for (int tn = 0; tn < TN; tn++) + acc += float(M1[(out_row + tm) * K_attn + bn + tn]) * v[tn]; + result[tm] += acc; + }} + bn += blockN; + }} + + // Handle remainder if K_attn not divisible by blockN + if (K_attn > n_iter * blockN) {{ + for (int tm = 0; tm < TM; tm++) {{ + for (int tn = 0; tn < TN; tn++) {{ + if (bn + tn < K_attn) + result[tm] += float(M1[(out_row + tm) * K_attn + bn + tn]) + * float(attn_out[bn + tn]); + }} + }} + }} + + for (int tm = 0; tm < TM; tm++) + result[tm] = simd_sum(result[tm]); + + if (slid == 0) {{ + for (int tm = 0; tm < TM; tm++) {{ + int e = out_row + tm; + if (e < E_CONST) + gate_part_a[e] = result[tm]; + }} + }} + + }} else {{ + // ════════════════════════════════════════════════════════════════ + // W_FUSED GEMV: bf16, M=E, K=K_hidden — W_fused × residual → gate_part_b + // ════════════════════════════════════════════════════════════════ + int local_tg = int(tg_x) - N_OPROJ_TG - N_M1_TG; + int out_row = local_tg * blockM_gate + int(sgid) * TM; + if (out_row >= E_CONST) return; + out_row = (out_row + TM <= E_CONST) ? out_row : (E_CONST - TM); + + float result[TM] = {{0.0f, 0.0f, 0.0f, 0.0f}}; + int bn = int(slid) * TN; + int n_iter = K_hidden / blockN; + + for (int i = 0; i < n_iter; i++) {{ + float v[TN]; + for (int tn = 0; tn < TN; tn++) + v[tn] = float(residual[bn + tn]); + for (int tm = 0; tm < TM; tm++) {{ + float acc = 0.0f; + for (int tn = 0; tn < TN; tn++) + acc += float(W_fused[(out_row + tm) * K_hidden + bn + tn]) * v[tn]; + result[tm] += acc; + }} + bn += blockN; + }} + + if (K_hidden > n_iter * blockN) {{ + for (int tm = 0; tm < TM; tm++) {{ + for (int tn = 0; tn < TN; tn++) {{ + if (bn + tn < K_hidden) + result[tm] += float(W_fused[(out_row + tm) * K_hidden + bn + tn]) + * float(residual[bn + tn]); + }} + }} + }} + + for (int tm = 0; tm < TM; tm++) + result[tm] = simd_sum(result[tm]); + + if (slid == 0) {{ + for (int tm = 0; tm < TM; tm++) {{ + int e = out_row + tm; + if (e < E_CONST) + gate_part_b[e] = result[tm]; + }} + }} + }} +""" + + +_custom_oproj_8bit_kernels = {} + + +def _get_custom_oproj_8bit_kernel(n_experts=64, group_size=64, scale_bf16=True): + key = (n_experts, group_size, scale_bf16) + if key not in _custom_oproj_8bit_kernels: + sc_tag = "_bf16sc" if scale_bf16 else "" + _custom_oproj_8bit_kernels[key] = mx.fast.metal_kernel( + name=f"custom_oproj_gate_gemv_8bit_e{n_experts}_gs{group_size}{sc_tag}", + input_names=[ + "W_oproj", "S_oproj", "B_oproj", # o_proj 8-bit weights + "attn_out", # (K_attn,) bf16 + "residual", # (K,) bf16 + "w_rms", # (K,) bf16 — RMSNorm weight + "M1", # (E, K_attn) bf16 + "W_fused", # (E, K) bf16 + "M_val", "K_attn_val", "K_hidden_val", + "N_OPROJ_TG_val", "N_M1_TG_val", "BM_GATE_val", + ], + output_names=["h_scaled", "h_out", "x2_partials", + "gate_part_a", "gate_part_b"], + source=_gen_custom_oproj_8bit_source(n_experts, group_size, scale_bf16), + ) + return _custom_oproj_8bit_kernels[key] + + +def fused_custom_oproj_8bit(W_oproj, S_oproj, B_oproj, + attn_out, residual, w_rms, + M1, W_fused, + M, K_attn, K_hidden=None, + n_experts=64, gate_bm=8, + group_size=64): + """Dispatch the fused 8-bit o_proj + bf16 gate GEMVs kernel. + + Three GEMV types in one dispatch: + - TGs 0..N_OPROJ_TG-1: 8-bit o_proj GEMV → h_scaled, h_out, x2_partials + - TGs N_OPROJ_TG..+N_M1_TG: bf16 M1 GEMV → gate_part_a + - TGs +N_M1_TG..end: bf16 W_fused GEMV → gate_part_b + + Args: + W_oproj/S_oproj/B_oproj: 8-bit quantized o_proj weights + attn_out: (K_attn,) bf16 — pre-o_proj attention output + residual: (K,) bf16 — input residual + w_rms: (K,) bf16 — post_attention_layernorm weight + M1: (E, K_attn) bf16 — precomputed W_fused @ W_oproj + W_fused: (E, K) bf16 — precomputed dequant(W_gate) * w_rms + M: hidden size (4096) + K_attn: attention output dim (8192) + K_hidden: hidden size for W_fused (defaults to M) + gate_bm: SGs per gate TG (1,2,4,8). Controls TG count. + + Returns: + (h_scaled, h_out, x2_partials, gate_part_a, gate_part_b) + """ + M = int(M) + K_attn = int(K_attn) + K_hidden = int(K_hidden) if K_hidden is not None else M + scale_bf16 = (S_oproj.dtype == mx.bfloat16) + + kern = _get_custom_oproj_8bit_kernel(n_experts, group_size, scale_bf16) + + n_oproj_tg = ceil_div(M, 32) # 128 for M=4096 + blockM_gate = gate_bm * 4 # rows per gate GEMV TG + n_m1_tg = ceil_div(n_experts, blockM_gate) + n_wf_tg = ceil_div(n_experts, blockM_gate) + total_tg = n_oproj_tg + n_m1_tg + n_wf_tg + + results = kern( + inputs=[W_oproj, S_oproj, B_oproj, + attn_out, residual, w_rms, M1, W_fused, + M, K_attn, K_hidden, n_oproj_tg, n_m1_tg, gate_bm], + output_shapes=[(M,), (M,), (n_oproj_tg,), (n_experts,), (n_experts,)], + output_dtypes=[mx.bfloat16, mx.bfloat16, mx.float32, + mx.float32, mx.float32], + grid=(total_tg * 32, 8, 1), # total threads: total_tg * 256 + threadgroup=(32, 8, 1), # 256 threads per TG (8 SGs of 32) + ) + return results[0], results[1], results[2], results[3], results[4] diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_moe_epilogue_qwen.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_moe_epilogue_qwen.py new file mode 100644 index 00000000..f6319acb --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_moe_epilogue_qwen.py @@ -0,0 +1,97 @@ +"""Fused MoE epilogue for Qwen3.5: weighted sum + shared expert gate + residual. + +Computes: + Y[j] = bf16( Σ_a(scores[a] * D_routed[a,j]) + gate_shared * D_shared[j] + H[j] ) + +Two modes: + fuse_sigmoid=False: gate_shared is pre-computed sigmoid output (scalar f32) + fuse_sigmoid=True: gate_raw is raw dot product; sigmoid computed internally +""" +import mlx.core as mx + + +def _gen_epilogue_qwen_source(K, n_active, fuse_sigmoid=False): + """Metal source for fused MoE epilogue with shared expert gate.""" + if fuse_sigmoid: + gate_line = """ + // Compute sigmoid from raw gate value + float gate = 1.0f / (1.0f + metal::exp(-shared_gate_val));""" + else: + gate_line = """ + float gate = shared_gate_val;""" + + return f""" + const int K_const = {K}; + const int n_active_const = {n_active}; + + uint tid = thread_position_in_grid.x; + if (tid >= K_const) return; + + // Weighted sum of routed expert outputs + float acc = 0.0f; + for (int a = 0; a < n_active_const; a++) {{ + acc += scores[a] * D_routed[a * K_const + tid]; + }} +{gate_line} + + // Shared expert: multiply by gate, add to accumulator + float shared_val = float(D_shared[tid]) * gate; + + // Add residual and write + Y[tid] = static_cast(acc + shared_val + float(H[tid])); +""" + + +_epilogue_qwen_kernels = {} + + +def _get_epilogue_qwen_kernel(K, n_active, fuse_sigmoid=False): + key = (K, n_active, fuse_sigmoid) + if key not in _epilogue_qwen_kernels: + sig_tag = "_fsig" if fuse_sigmoid else "" + _epilogue_qwen_kernels[key] = mx.fast.metal_kernel( + name=f"fused_moe_epilogue_qwen_K{K}_n{n_active}{sig_tag}", + input_names=["D_routed", "D_shared", "scores", "H", + "shared_gate_val"], + output_names=["Y"], + source=_gen_epilogue_qwen_source(K, n_active, fuse_sigmoid), + ) + return _epilogue_qwen_kernels[key] + + +def fused_moe_epilogue_qwen(d_routed, d_shared, scores, h, + shared_gate, k_val, fuse_sigmoid=False): + """Fused MoE epilogue with shared expert gate. + + Args: + d_routed: routed expert outputs (n_active, K) float32 + d_shared: shared expert output (K,) float32 + scores: normalized routing scores (n_active,) float32 + h: residual hidden state (K,) bfloat16 + shared_gate: sigmoid output (fuse_sigmoid=False) or raw dot product + (fuse_sigmoid=True), scalar float32 + k_val: hidden dimension + fuse_sigmoid: if True, compute sigmoid(shared_gate) internally + + Returns: + Y: (K,) bfloat16 — final layer output + """ + K = int(k_val) + n_active = scores.shape[0] + kern = _get_epilogue_qwen_kernel(K, n_active, fuse_sigmoid) + + # shared_gate is a scalar — pass as 0-d array + if shared_gate.ndim > 0: + shared_gate = shared_gate.reshape(()) + + tg_size = min(K, 1024) + n_tg = (K + tg_size - 1) // tg_size + + Y = kern( + inputs=[d_routed, d_shared, scores, h, shared_gate], + output_shapes=[(K,)], + output_dtypes=[mx.bfloat16], + grid=(n_tg * tg_size, 1, 1), + threadgroup=(tg_size, 1, 1), + ) + return Y[0] diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/merged_down_proj_8bit.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/merged_down_proj_8bit.py new file mode 100644 index 00000000..0ab2dabf --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/merged_down_proj_8bit.py @@ -0,0 +1,218 @@ +"""Merged 8-bit down_proj GEMV for Qwen3.5 routed + shared experts. + +Port of the 4-bit down_proj kernel, adapted for 8-bit quantization (gs=64). +Maps intermediate → hidden: (n_active, N_IN) → (n_active, K_OUT). + +Two paths via tgid.z: +- 8-bit routed experts (z < n_active): uint8 dequant with expert index lookup, f32 output +- 8-bit shared expert (z == n_active): uint8 dequant (no expert index), f32 output +""" +import mlx.core as mx + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def _gen_merged_down_8bit_source(group_size=64, scale_bf16=True): + """Metal source for merged 8-bit down_proj GEMV. + + Same structure as gate GEMV: 2 SGs × 4 rows/SG = 8 output rows per TG. + 8-bit dequant: result = scale * Σ(x[i]*w[i]) + bias * Σ(x[i]) + Both routed and shared paths use 8-bit dequantization. + """ + gs = int(group_size) + sc_stride = 256 // gs # groups per K-block (256/64 = 4) + slid_divisor = gs // 8 # threads per group (64/8 = 8) + sc_t = "bfloat16_t" if scale_bf16 else "float" + + return f""" + const int RESULTS_PER_SG = 4; + const int VALUES_PER_THREAD = 8; + const int BLOCK_SIZE = 256; + + int K_OUT = K_OUT_val; + int N_IN = N_IN_val; + int SHARED_N_IN = SHARED_N_IN_val; + int n_active = n_active_val; + + uint3 tgid = threadgroup_position_in_grid; + uint sgid = simdgroup_index_in_threadgroup; // 0 or 1 + uint slid = thread_index_in_simdgroup; // 0..31 + + int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG; + if (out_row >= K_OUT) return; + + if (tgid.z < (uint)n_active) {{ + // ═══════ 8-BIT ROUTED EXPERT PATH ═══════ + int N_groups = N_IN / {gs}; + + int expert = inds[tgid.z]; + + const device uint8_t* ws = (const device uint8_t*)W + + (long)expert * K_OUT * N_IN + out_row * N_IN + slid * VALUES_PER_THREAD; + const device {sc_t}* sc = (const device {sc_t}*)S + + (long)expert * K_OUT * N_groups + out_row * N_groups + slid / {slid_divisor}; + const device {sc_t}* bi = (const device {sc_t}*)B_q + + (long)expert * K_OUT * N_groups + out_row * N_groups + slid / {slid_divisor}; + + const device float* x_ptr = (const device float*)X_routed + + tgid.z * N_IN; + + int x_base = slid * VALUES_PER_THREAD; + float result[4] = {{0, 0, 0, 0}}; + + for (int k = 0; k < N_IN; k += BLOCK_SIZE) {{ + float x_thread[8]; + float xsum = 0; + for (int i = 0; i < 8; i++) {{ + float xi = x_ptr[x_base + i]; + x_thread[i] = xi; + xsum += xi; + }} + + for (int row = 0; row < RESULTS_PER_SG; row++) {{ + const device uint8_t* wl = ws + row * N_IN; + float s = float(sc[row * N_groups]); + float b = float(bi[row * N_groups]); + float accum = 0; + for (int i = 0; i < 8; i++) {{ + accum += x_thread[i] * float(wl[i]); + }} + result[row] += s * accum + xsum * b; + }} + + ws += BLOCK_SIZE; + sc += {sc_stride}; + bi += {sc_stride}; + x_base += BLOCK_SIZE; + }} + + device float* yp = Y_routed + tgid.z * K_OUT + out_row; + for (int row = 0; row < RESULTS_PER_SG; row++) {{ + float r = simd_sum(result[row]); + if (slid == 0) {{ + yp[row] = r; + }} + }} + + }} else {{ + // ═══════ 8-BIT SHARED EXPERT PATH ═══════ + // Same dequant as routed path, but no expert index lookup. + // W_shared_down is (K_OUT, SHARED_N_IN/4) uint32 → (K_OUT, SHARED_N_IN) uint8 + int N_groups = SHARED_N_IN / {gs}; + + const device uint8_t* ws = (const device uint8_t*)W_shared_down + + (long)out_row * SHARED_N_IN + slid * VALUES_PER_THREAD; + const device {sc_t}* sc = (const device {sc_t}*)S_shared_down + + (long)out_row * N_groups + slid / {slid_divisor}; + const device {sc_t}* bi = (const device {sc_t}*)B_shared_down + + (long)out_row * N_groups + slid / {slid_divisor}; + + // X_shared is float32 (output of Kernel 1 shared path) + const device float* x_ptr = (const device float*)X_shared; + + int x_base = slid * VALUES_PER_THREAD; + float result[4] = {{0, 0, 0, 0}}; + + for (int k = 0; k < SHARED_N_IN; k += BLOCK_SIZE) {{ + float x_thread[8]; + float xsum = 0; + for (int i = 0; i < 8; i++) {{ + float xi = x_ptr[x_base + i]; + x_thread[i] = xi; + xsum += xi; + }} + + for (int row = 0; row < RESULTS_PER_SG; row++) {{ + const device uint8_t* wl = ws + row * SHARED_N_IN; + float s = float(sc[row * N_groups]); + float b = float(bi[row * N_groups]); + float accum = 0; + for (int i = 0; i < 8; i++) {{ + accum += x_thread[i] * float(wl[i]); + }} + result[row] += s * accum + xsum * b; + }} + + ws += BLOCK_SIZE; + sc += {sc_stride}; + bi += {sc_stride}; + x_base += BLOCK_SIZE; + }} + + device float* yp = Y_shared + out_row; + for (int tm = 0; tm < RESULTS_PER_SG; tm++) {{ + float r = simd_sum(result[tm]); + if (slid == 0) {{ + yp[tm] = r; + }} + }} + }} +""" + + +_merged_down_8bit_kernels = {} + + +def _get_merged_down_8bit_kernel(group_size=64, scale_bf16=True): + key = (group_size, scale_bf16) + if key not in _merged_down_8bit_kernels: + sc_tag = "_bf16sc" if scale_bf16 else "" + _merged_down_8bit_kernels[key] = mx.fast.metal_kernel( + name=f"merged_down_proj_8bit_gs{group_size}{sc_tag}", + input_names=["W", "S", "B_q", + "W_shared_down", "S_shared_down", "B_shared_down", + "X_routed", "X_shared", "inds", + "K_OUT_val", "N_IN_val", "SHARED_N_IN_val", "n_active_val"], + output_names=["Y_routed", "Y_shared"], + source=_gen_merged_down_8bit_source(group_size, scale_bf16), + ) + return _merged_down_8bit_kernels[key] + + +def fused_merged_down_proj_8bit(w_q, s, b_q, + w_shared_down, s_shared_down, b_shared_down, + x_routed, x_shared, inds, + k_out, n_in, group_size=64, + shared_n_in=None): + """Single-dispatch merged down_proj for 8-bit routed + 8-bit shared experts. + + Args: + w_q: routed quantized down weights (E, K_OUT, N_IN/4) uint32 + s: routed scales (E, K_OUT, N_IN/gs) bfloat16 + b_q: routed biases (E, K_OUT, N_IN/gs) bfloat16 + w_shared_down: shared expert down weight (K_OUT, SHARED_N_IN/4) uint32 + s_shared_down: shared expert down scales (K_OUT, SHARED_N_IN/gs) bfloat16 + b_shared_down: shared expert down biases (K_OUT, SHARED_N_IN/gs) bfloat16 + x_routed: routed SwiGLU output (n_active, N_IN) float32 + x_shared: shared SwiGLU output (SHARED_N_IN,) float32 + inds: selected expert indices (n_active,) uint32 + k_out: output dimension (4096 for Qwen3.5) + n_in: routed expert input dimension (1024 for Qwen3.5) + group_size: quantization group size (64) + shared_n_in: shared expert input dimension (defaults to n_in) + + Returns: + (Y_routed, Y_shared): + Y_routed: (n_active, k_out) float32 + Y_shared: (k_out,) float32 + """ + scale_bf16 = (s.dtype == mx.bfloat16) + kern = _get_merged_down_8bit_kernel(group_size, scale_bf16) + n_active = inds.shape[0] + k_out_val = int(k_out) + n_in_val = int(n_in) + shared_n_in_val = int(shared_n_in) if shared_n_in is not None else n_in_val + y_groups = ceil_div(k_out_val, 8) + Y = kern( + inputs=[w_q, s, b_q, + w_shared_down, s_shared_down, b_shared_down, + x_routed, x_shared, inds, + k_out_val, n_in_val, shared_n_in_val, n_active], + output_shapes=[(n_active, k_out_val), (k_out_val,)], + output_dtypes=[mx.float32, mx.float32], + grid=(32, y_groups * 2, n_active + 1), + threadgroup=(32, 2, 1), + ) + return Y[0], Y[1] diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/merged_routed_shared_swiglu_8bit.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/merged_routed_shared_swiglu_8bit.py new file mode 100644 index 00000000..79aa1336 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/merged_routed_shared_swiglu_8bit.py @@ -0,0 +1,274 @@ +"""Merged 8-bit fused gate+up+SwiGLU for Qwen3.5 routed + shared experts. + +Port of the 4-bit kernel for Kimi K2.5, adapted for: +- 8-bit quantization (direct uint8 byte reads, no nibble extraction) +- group_size=64 (one scale+bias per 64 elements) +- Qwen3.5 dimensions (K=4096, N_INTER=1024, E=512, top_k=10) +- No score normalization in kernel (handled by gate dispatch) + +Two paths via tgid.z: +- 8-bit routed experts (z < n_active): uint8 dequant with expert index lookup, f32 output +- 8-bit shared expert (z == n_active): uint8 dequant (no expert index), f32 output +""" +import mlx.core as mx + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def _gen_merged_8bit_source(group_size=64, scale_bf16=True): + """Generate Metal source for merged 8-bit fused gate+up+SwiGLU. + + Both routed and shared expert paths use 8-bit dequant: + result = scale * Σ(x[i]*w[i]) + bias * Σ(x[i]) + + Each thread processes VALUES_PER_THREAD=8 elements of K per iteration. + BLOCK_SIZE = 32 * 8 = 256 elements per K-block. + """ + gs = int(group_size) + sc_stride = 256 // gs # groups consumed per K-block (256/64 = 4) + slid_divisor = gs // 8 # threads per group (64/8 = 8) + sc_t = "bfloat16_t" if scale_bf16 else "float" + + return f""" + const int RESULTS_PER_SG = 4; + + int N_INTER = N_INTER_val; + int SHARED_INTER = SHARED_INTER_val; + int K = K_val; + int n_active = n_active_val; + + uint3 tgid = threadgroup_position_in_grid; + uint sgid = simdgroup_index_in_threadgroup; // 0 or 1 + uint slid = thread_index_in_simdgroup; // 0..31 + + int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG; + + // Routed path uses N_INTER rows, shared path uses SHARED_INTER rows + int row_limit = (tgid.z < (uint)n_active) ? N_INTER : SHARED_INTER; + if (out_row >= row_limit) return; + + float gate_result[4] = {{0, 0, 0, 0}}; + float up_result[4] = {{0, 0, 0, 0}}; + + if (tgid.z < (uint)n_active) {{ + // ═══════ 8-BIT ROUTED EXPERT PATH ═══════ + const int VALUES_PER_THREAD = 8; + const int BLOCK_SIZE = 256; // 32 * 8 + + int N_TOTAL = 2 * N_INTER; // gate + up stacked + int K_groups = K / {gs}; + + int expert = inds[tgid.z]; + + // W is stored as uint32 (4 bytes per uint32), cast to uint8_t + // Layout: (E, N_TOTAL, K/4) uint32 → (E, N_TOTAL, K) uint8 + const device uint8_t* ws_gate = (const device uint8_t*)W + + (long)expert * N_TOTAL * K + out_row * K + slid * VALUES_PER_THREAD; + const device {sc_t}* sc_gate = (const device {sc_t}*)S + + (long)expert * N_TOTAL * K_groups + out_row * K_groups + slid / {slid_divisor}; + const device {sc_t}* bi_gate = (const device {sc_t}*)B_q + + (long)expert * N_TOTAL * K_groups + out_row * K_groups + slid / {slid_divisor}; + + const device uint8_t* ws_up = (const device uint8_t*)W + + (long)expert * N_TOTAL * K + (out_row + N_INTER) * K + slid * VALUES_PER_THREAD; + const device {sc_t}* sc_up = (const device {sc_t}*)S + + (long)expert * N_TOTAL * K_groups + (out_row + N_INTER) * K_groups + slid / {slid_divisor}; + const device {sc_t}* bi_up = (const device {sc_t}*)B_q + + (long)expert * N_TOTAL * K_groups + (out_row + N_INTER) * K_groups + slid / {slid_divisor}; + + int x_base = slid * VALUES_PER_THREAD; + + for (int k = 0; k < K; k += BLOCK_SIZE) {{ + // Load 8 x values and compute sum + float x_thread[8]; + float xsum = 0; + for (int i = 0; i < 8; i++) {{ + float xi = float(X[x_base + i]); + x_thread[i] = xi; + xsum += xi; + }} + + for (int row = 0; row < RESULTS_PER_SG; row++) {{ + // Gate projection + const device uint8_t* wg = ws_gate + row * K; + float sg = float(sc_gate[row * K_groups]); + float bg = float(bi_gate[row * K_groups]); + float accum_g = 0; + for (int i = 0; i < 8; i++) {{ + accum_g += x_thread[i] * float(wg[i]); + }} + gate_result[row] += sg * accum_g + xsum * bg; + + // Up projection + const device uint8_t* wu = ws_up + row * K; + float su = float(sc_up[row * K_groups]); + float bu = float(bi_up[row * K_groups]); + float accum_u = 0; + for (int i = 0; i < 8; i++) {{ + accum_u += x_thread[i] * float(wu[i]); + }} + up_result[row] += su * accum_u + xsum * bu; + }} + + ws_gate += BLOCK_SIZE; + ws_up += BLOCK_SIZE; + sc_gate += {sc_stride}; + sc_up += {sc_stride}; + bi_gate += {sc_stride}; + bi_up += {sc_stride}; + x_base += BLOCK_SIZE; + }} + + // Epilogue: SwiGLU + write f32 to Y_routed + device float* yp = Y_routed + tgid.z * N_INTER + out_row; + for (int row = 0; row < RESULTS_PER_SG; row++) {{ + float g = simd_sum(gate_result[row]); + float u = simd_sum(up_result[row]); + if (slid == 0) {{ + float silu_g = g / (1.0f + metal::exp(-g)); + yp[row] = silu_g * u; + }} + }} + + }} else {{ + // ═══════ 8-BIT SHARED EXPERT PATH ═══════ + // Same dequant as routed path, but no expert index lookup. + // W_shared is (2*SHARED_INTER, K/4) uint32 → (2*SHARED_INTER, K) uint8 + const int VALUES_PER_THREAD = 8; + const int BLOCK_SIZE = 256; // 32 * 8 + + int K_groups = K / {gs}; + + const device uint8_t* ws_gate = (const device uint8_t*)W_shared + + (long)out_row * K + slid * VALUES_PER_THREAD; + const device {sc_t}* sc_gate = (const device {sc_t}*)S_shared + + (long)out_row * K_groups + slid / {slid_divisor}; + const device {sc_t}* bi_gate = (const device {sc_t}*)B_shared + + (long)out_row * K_groups + slid / {slid_divisor}; + + const device uint8_t* ws_up = (const device uint8_t*)W_shared + + (long)(out_row + SHARED_INTER) * K + slid * VALUES_PER_THREAD; + const device {sc_t}* sc_up = (const device {sc_t}*)S_shared + + (long)(out_row + SHARED_INTER) * K_groups + slid / {slid_divisor}; + const device {sc_t}* bi_up = (const device {sc_t}*)B_shared + + (long)(out_row + SHARED_INTER) * K_groups + slid / {slid_divisor}; + + int x_base = slid * VALUES_PER_THREAD; + + for (int k = 0; k < K; k += BLOCK_SIZE) {{ + float x_thread[8]; + float xsum = 0; + for (int i = 0; i < 8; i++) {{ + float xi = float(X[x_base + i]); + x_thread[i] = xi; + xsum += xi; + }} + + for (int row = 0; row < RESULTS_PER_SG; row++) {{ + // Gate projection + const device uint8_t* wg = ws_gate + row * K; + float sg = float(sc_gate[row * K_groups]); + float bg = float(bi_gate[row * K_groups]); + float accum_g = 0; + for (int i = 0; i < 8; i++) {{ + accum_g += x_thread[i] * float(wg[i]); + }} + gate_result[row] += sg * accum_g + xsum * bg; + + // Up projection + const device uint8_t* wu = ws_up + row * K; + float su = float(sc_up[row * K_groups]); + float bu = float(bi_up[row * K_groups]); + float accum_u = 0; + for (int i = 0; i < 8; i++) {{ + accum_u += x_thread[i] * float(wu[i]); + }} + up_result[row] += su * accum_u + xsum * bu; + }} + + ws_gate += BLOCK_SIZE; + ws_up += BLOCK_SIZE; + sc_gate += {sc_stride}; + sc_up += {sc_stride}; + bi_gate += {sc_stride}; + bi_up += {sc_stride}; + x_base += BLOCK_SIZE; + }} + + // Epilogue: SwiGLU + write f32 to Y_shared + device float* yp = Y_shared + out_row; + for (int tm = 0; tm < RESULTS_PER_SG; tm++) {{ + float g = simd_sum(gate_result[tm]); + float u = simd_sum(up_result[tm]); + if (slid == 0) {{ + float silu_g = g / (1.0f + metal::exp(-g)); + yp[tm] = silu_g * u; + }} + }} + }} +""" + + +_merged_8bit_kernels = {} + + +def _get_merged_8bit_kernel(group_size=64, scale_bf16=True): + key = (group_size, scale_bf16) + if key not in _merged_8bit_kernels: + sc_tag = "_bf16sc" if scale_bf16 else "" + _merged_8bit_kernels[key] = mx.fast.metal_kernel( + name=f"merged_routed_shared_swiglu_8bit_gs{group_size}{sc_tag}", + input_names=["W", "S", "B_q", "W_shared", "S_shared", "B_shared", + "X", "inds", + "N_INTER_val", "SHARED_INTER_val", "K_val", "n_active_val"], + output_names=["Y_routed", "Y_shared"], + source=_gen_merged_8bit_source(group_size, scale_bf16), + ) + return _merged_8bit_kernels[key] + + +def fused_merged_gate_up_swiglu_8bit(w_q, s, b_q, + w_shared, s_shared, b_shared, + x, inds, + n_inter, k_hidden, group_size=64, + shared_inter=None): + """Single-dispatch merged gate+up+SwiGLU for 8-bit routed + 8-bit shared experts. + + Args: + w_q: stacked routed quantized weights (E, 2*N_INTER, K/4) uint32 + s: routed scales (E, 2*N_INTER, K/gs) bfloat16 + b_q: routed biases (E, 2*N_INTER, K/gs) bfloat16 + w_shared: shared expert gate+up stacked (2*SHARED_INTER, K/4) uint32 + s_shared: shared expert scales (2*SHARED_INTER, K/gs) bfloat16 + b_shared: shared expert biases (2*SHARED_INTER, K/gs) bfloat16 + x: input vector (K,) bfloat16 + inds: selected expert indices (n_active,) uint32 + n_inter: routed expert intermediate size (1024 for Qwen3.5) + k_hidden: hidden size (4096 for Qwen3.5) + group_size: quantization group size (64) + shared_inter: shared expert intermediate size (defaults to n_inter) + + Returns: + (Y_routed, Y_shared): + Y_routed: (n_active, n_inter) float32 + Y_shared: (shared_inter,) float32 + """ + scale_bf16 = (s.dtype == mx.bfloat16) + kern = _get_merged_8bit_kernel(group_size, scale_bf16) + n_active = inds.shape[0] + n_inter_val = int(n_inter) + shared_inter_val = int(shared_inter) if shared_inter is not None else n_inter_val + # Grid y must cover max(n_inter, shared_inter) + max_inter = max(n_inter_val, shared_inter_val) + Y = kern( + inputs=[w_q, s, b_q, w_shared, s_shared, b_shared, + x, inds, + n_inter_val, shared_inter_val, int(k_hidden), n_active], + output_shapes=[(n_active, n_inter_val), (shared_inter_val,)], + output_dtypes=[mx.float32, mx.float32], + grid=(32, ceil_div(max_inter, 8) * 2, n_active + 1), + threadgroup=(32, 2, 1), + ) + return Y[0], Y[1] diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/oproj_softmax_topk_swiglu_8bit.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/oproj_softmax_topk_swiglu_8bit.py new file mode 100644 index 00000000..e5c3a3c2 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/oproj_softmax_topk_swiglu_8bit.py @@ -0,0 +1,493 @@ +"""Dispatch 2: SwiGLU with softmax prologue for Qwen3.5 oproj fusion. + +Port of Kimi's oproj_topk_v2_swiglu.py adapted for: +- 8-bit quantized weights with gs=64 (Kimi uses 4-bit gs=32) +- Softmax routing (Kimi uses sigmoid) +- E up to 512 with multiple scores per thread (SPT = ceil(E/64)) +- Shared expert gate GEMV hidden in TG(0,0,0) SG 0 +- Both routed and shared paths use 8-bit quantized weights + +Prologue (all TGs): + Phase 1: distributed x² sum → inv_rms + Phase 2 (routed TGs only): gate scores → softmax → parallel top-k → norm_topk_prob + Phase 3 (TG(0,0,0) SG 0): shared_expert_gate 8-bit GEMV → gate_raw + +Main SwiGLU (after prologue): + Routed (z < n_active): 8-bit gate+up+SwiGLU with h_scaled input (inv_rms factored out) + Shared (z == n_active): 8-bit SwiGLU for shared expert +""" +import mlx.core as mx + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def _gen_oproj_softmax_topk_swiglu_8bit_source(group_size=64, scale_bf16=True, + n_experts=64, top_k=10, + norm_topk=True): + """Generate Metal source for softmax + top-k + SwiGLU with oproj prologue.""" + gs = int(group_size) + sc_stride = 256 // gs # groups per K-block (256/64 = 4) + slid_divisor = gs // 8 # threads per group (64/8 = 8) + sc_t = "bfloat16_t" if scale_bf16 else "float" + E = int(n_experts) + K_TOP = int(top_k) + SPT = (E + 63) // 64 # scores per thread + + # Score normalization block (TG(0,0,0) thread 0 only) + if norm_topk: + score_norm_block = f""" + // ── norm_topk_prob + write indices (TG(0,0,0) thread 0) ── + if (tgid.y == 0 && tgid.z == 0 && tid == 0) {{ + float total = 0.0f; + for (int a = 0; a < {K_TOP}; a++) total += tg_selected_scores[a]; + float inv_total = 1.0f / total; + for (int a = 0; a < {K_TOP}; a++) {{ + norm_scores[a] = tg_selected_scores[a] * inv_total; + out_inds[a] = (uint)tg_inds[a]; + }} + }}""" + else: + score_norm_block = f""" + // ── Write indices and raw scores (TG(0,0,0) thread 0) ── + if (tgid.y == 0 && tgid.z == 0 && tid == 0) {{ + for (int a = 0; a < {K_TOP}; a++) {{ + norm_scores[a] = tg_selected_scores[a]; + out_inds[a] = (uint)tg_inds[a]; + }} + }}""" + + return f""" + const int RESULTS_PER_SG = 4; + const int E_CONST = {E}; + const int K_TOP_CONST = {K_TOP}; + const int SPT = {SPT}; + + int N_INTER = N_INTER_val; + int SHARED_INTER = SHARED_INTER_val; + int K = K_val; + int n_active = n_active_val; + int N_OPROJ_TG = N_OPROJ_TG_val; + + uint3 tgid = threadgroup_position_in_grid; + uint sgid = simdgroup_index_in_threadgroup; // 0 or 1 + uint slid = thread_index_in_simdgroup; // 0..31 + int tid = int(sgid) * 32 + int(slid); // 0..63 + + // ═══════════════════════════════════════════════════════════════ + // PROLOGUE PHASE 1: distributed x² sum → inv_rms (ALL TGs) + // ═══════════════════════════════════════════════════════════════ + int chunk = (N_OPROJ_TG + 63) / 64; + int x2_start = tid * chunk; + int x2_end = min(x2_start + chunk, N_OPROJ_TG); + float local_x2 = 0.0f; + for (int i = x2_start; i < x2_end; i++) local_x2 += x2_partials[i]; + float sg_x2_sum = simd_sum(local_x2); + + threadgroup float tg_x2_sg[2]; + if (slid == 0) tg_x2_sg[sgid] = sg_x2_sum; + threadgroup_barrier(mem_flags::mem_threadgroup); + float total_x2 = tg_x2_sg[0] + tg_x2_sg[1]; + float inv_rms = metal::precise::rsqrt(total_x2 / (float)K + 1e-6f); + + // ═══════════════════════════════════════════════════════════════ + // PROLOGUE PHASE 2: Softmax + Top-k (routed TGs only) + // ═══════════════════════════════════════════════════════════════ + threadgroup int tg_inds[{K_TOP}]; + threadgroup float tg_selected_scores[{K_TOP}]; + + if (tgid.z < (uint)n_active) {{ + // Load gate scores and apply inv_rms + float my_scores[SPT]; + for (int j = 0; j < SPT; j++) {{ + int e = tid * SPT + j; + if (e < E_CONST) + my_scores[j] = (gate_part_a[e] + gate_part_b[e]) * inv_rms; + else + my_scores[j] = -1e30f; + }} + + // ── Softmax: distributed max ── + float local_max = -1e30f; + for (int j = 0; j < SPT; j++) + local_max = max(local_max, my_scores[j]); + float sg_max_val = simd_max(local_max); + threadgroup float tg_softmax_sg[2]; + if (slid == 0) tg_softmax_sg[sgid] = sg_max_val; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tg_max = max(tg_softmax_sg[0], tg_softmax_sg[1]); + + // ── Softmax: exp + distributed sum ── + float local_sum = 0.0f; + for (int j = 0; j < SPT; j++) {{ + float e_val = metal::exp(my_scores[j] - tg_max); + my_scores[j] = e_val; + local_sum += e_val; + }} + float sg_sum_val = simd_sum(local_sum); + if (slid == 0) tg_softmax_sg[sgid] = sg_sum_val; + threadgroup_barrier(mem_flags::mem_threadgroup); + float tg_sum = tg_softmax_sg[0] + tg_softmax_sg[1]; + + // ── Softmax: normalize ── + float inv_sum = 1.0f / tg_sum; + for (int j = 0; j < SPT; j++) + my_scores[j] *= inv_sum; + + // ── Parallel top-k: K_TOP rounds ── + threadgroup float tg_tk_val[2]; + threadgroup int tg_tk_info[2]; + + for (int round = 0; round < K_TOP_CONST; round++) {{ + // Find local best among SPT scores + float best = -1.0f; + int best_e = -1; + for (int j = 0; j < SPT; j++) {{ + int e = tid * SPT + j; + if (e < E_CONST && my_scores[j] > best) {{ + best = my_scores[j]; + best_e = e; + }} + }} + + // SG-level max + winner identification + float sg_best = simd_max(best); + int candidate = (best == sg_best && best > 0.0f) ? int(slid) : 999; + int sg_winner = simd_min(candidate); + + if (slid == 0) {{ + tg_tk_val[sgid] = sg_best; + tg_tk_info[sgid] = sg_winner; + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Determine global winner (SG with higher max; tie-break: SG 0) + int winner_sg = (tg_tk_val[0] >= tg_tk_val[1]) ? 0 : 1; + int winner_lane = tg_tk_info[winner_sg]; + int winner_tid = winner_sg * 32 + winner_lane; + + // Winner writes expert index and score to TG memory + if (tid == winner_tid) {{ + tg_inds[round] = best_e; + tg_selected_scores[round] = best; + // Disable the winning score in register + for (int j = 0; j < SPT; j++) {{ + if (tid * SPT + j == best_e) {{ + my_scores[j] = -1.0f; + break; + }} + }} + }} + threadgroup_barrier(mem_flags::mem_threadgroup); + }} + }} +{score_norm_block} + + // ═══════════════════════════════════════════════════════════════ + // PROLOGUE PHASE 3: Shared expert gate (TG(0,0,0) SG 0 only) + // 8-bit GEMV: W_seg (1, K) × X (K,) → gate_raw scalar + // Input is h_scaled; multiply result by inv_rms for true gate value + // ═══════════════════════════════════════════════════════════════ + if (tgid.y == 0 && tgid.z == 0 && sgid == 0) {{ + const int VPT = 8; + const int BLOCK = 256; // 32 * VPT + int K_groups_seg = K / {gs}; + + const device uint8_t* wg_seg = (const device uint8_t*)W_seg + + slid * VPT; + const device {sc_t}* sc_seg = (const device {sc_t}*)S_seg + + slid / {slid_divisor}; + const device {sc_t}* bi_seg = (const device {sc_t}*)B_seg + + slid / {slid_divisor}; + int xb = slid * VPT; + + float gate_acc = 0.0f; + for (int k = 0; k < K; k += BLOCK) {{ + float xsum = 0.0f, wacc = 0.0f; + for (int i = 0; i < VPT; i++) {{ + float xi = float(X[xb + i]); + xsum += xi; + wacc += xi * float(wg_seg[i]); + }} + gate_acc += float(*sc_seg) * wacc + xsum * float(*bi_seg); + wg_seg += BLOCK; + sc_seg += {sc_stride}; + bi_seg += {sc_stride}; + xb += BLOCK; + }} + gate_acc = simd_sum(gate_acc); + if (slid == 0) gate_raw[0] = gate_acc * inv_rms; + }} + + // ═══════════════════════════════════════════════════════════════ + // MAIN SwiGLU BODY + // ═══════════════════════════════════════════════════════════════ + int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG; + int row_limit = (tgid.z < (uint)n_active) ? N_INTER : SHARED_INTER; + if (out_row >= row_limit) return; + + float gate_result[4] = {{0, 0, 0, 0}}; + float up_result[4] = {{0, 0, 0, 0}}; + + if (tgid.z < (uint)n_active) {{ + // ═══════ 8-BIT ROUTED EXPERT PATH ═══════ + const int VALUES_PER_THREAD = 8; + const int BLOCK_SIZE = 256; + + int N_TOTAL = 2 * N_INTER; + int K_groups = K / {gs}; + + int expert = tg_inds[tgid.z]; + + const device uint8_t* ws_gate = (const device uint8_t*)W + + (long)expert * N_TOTAL * K + out_row * K + slid * VALUES_PER_THREAD; + const device {sc_t}* sc_gate = (const device {sc_t}*)S + + (long)expert * N_TOTAL * K_groups + out_row * K_groups + + slid / {slid_divisor}; + const device {sc_t}* bi_gate = (const device {sc_t}*)B_q + + (long)expert * N_TOTAL * K_groups + out_row * K_groups + + slid / {slid_divisor}; + + const device uint8_t* ws_up = (const device uint8_t*)W + + (long)expert * N_TOTAL * K + (out_row + N_INTER) * K + + slid * VALUES_PER_THREAD; + const device {sc_t}* sc_up = (const device {sc_t}*)S + + (long)expert * N_TOTAL * K_groups + (out_row + N_INTER) * K_groups + + slid / {slid_divisor}; + const device {sc_t}* bi_up = (const device {sc_t}*)B_q + + (long)expert * N_TOTAL * K_groups + (out_row + N_INTER) * K_groups + + slid / {slid_divisor}; + + int x_base = slid * VALUES_PER_THREAD; + + for (int k = 0; k < K; k += BLOCK_SIZE) {{ + float x_thread[8]; + float xsum = 0; + for (int i = 0; i < 8; i++) {{ + float xi = float(X[x_base + i]); + x_thread[i] = xi; + xsum += xi; + }} + + for (int row = 0; row < RESULTS_PER_SG; row++) {{ + // Gate projection + const device uint8_t* wg = ws_gate + row * K; + float sg = float(sc_gate[row * K_groups]); + float bg = float(bi_gate[row * K_groups]); + float accum_g = 0; + for (int i = 0; i < 8; i++) + accum_g += x_thread[i] * float(wg[i]); + gate_result[row] += sg * accum_g + xsum * bg; + + // Up projection + const device uint8_t* wu = ws_up + row * K; + float su = float(sc_up[row * K_groups]); + float bu = float(bi_up[row * K_groups]); + float accum_u = 0; + for (int i = 0; i < 8; i++) + accum_u += x_thread[i] * float(wu[i]); + up_result[row] += su * accum_u + xsum * bu; + }} + + ws_gate += BLOCK_SIZE; + ws_up += BLOCK_SIZE; + sc_gate += {sc_stride}; + sc_up += {sc_stride}; + bi_gate += {sc_stride}; + bi_up += {sc_stride}; + x_base += BLOCK_SIZE; + }} + + // Epilogue: apply inv_rms (factored), SwiGLU, write f32 + device float* yp = Y_routed + tgid.z * N_INTER + out_row; + for (int row = 0; row < RESULTS_PER_SG; row++) {{ + float g = simd_sum(gate_result[row]) * inv_rms; + float u = simd_sum(up_result[row]) * inv_rms; + if (slid == 0) {{ + float silu_g = g / (1.0f + metal::exp(-g)); + yp[row] = silu_g * u; + }} + }} + + }} else {{ + // ═══════ 8-BIT SHARED EXPERT PATH ═══════ + const int VALUES_PER_THREAD = 8; + const int BLOCK_SIZE = 256; + + int K_groups = K / {gs}; + + const device uint8_t* ws_gate = (const device uint8_t*)W_shared + + (long)out_row * K + slid * VALUES_PER_THREAD; + const device {sc_t}* sc_gate = (const device {sc_t}*)S_shared + + (long)out_row * K_groups + slid / {slid_divisor}; + const device {sc_t}* bi_gate = (const device {sc_t}*)B_shared + + (long)out_row * K_groups + slid / {slid_divisor}; + + const device uint8_t* ws_up = (const device uint8_t*)W_shared + + (long)(out_row + SHARED_INTER) * K + slid * VALUES_PER_THREAD; + const device {sc_t}* sc_up = (const device {sc_t}*)S_shared + + (long)(out_row + SHARED_INTER) * K_groups + slid / {slid_divisor}; + const device {sc_t}* bi_up = (const device {sc_t}*)B_shared + + (long)(out_row + SHARED_INTER) * K_groups + slid / {slid_divisor}; + + int x_base = slid * VALUES_PER_THREAD; + + for (int k = 0; k < K; k += BLOCK_SIZE) {{ + float x_thread[8]; + float xsum = 0; + for (int i = 0; i < 8; i++) {{ + float xi = float(X[x_base + i]); + x_thread[i] = xi; + xsum += xi; + }} + + for (int row = 0; row < RESULTS_PER_SG; row++) {{ + const device uint8_t* wg_s = ws_gate + row * K; + float sg_s = float(sc_gate[row * K_groups]); + float bg_s = float(bi_gate[row * K_groups]); + float accum_g = 0; + for (int i = 0; i < 8; i++) + accum_g += x_thread[i] * float(wg_s[i]); + gate_result[row] += sg_s * accum_g + xsum * bg_s; + + const device uint8_t* wu_s = ws_up + row * K; + float su_s = float(sc_up[row * K_groups]); + float bu_s = float(bi_up[row * K_groups]); + float accum_u = 0; + for (int i = 0; i < 8; i++) + accum_u += x_thread[i] * float(wu_s[i]); + up_result[row] += su_s * accum_u + xsum * bu_s; + }} + + ws_gate += BLOCK_SIZE; + ws_up += BLOCK_SIZE; + sc_gate += {sc_stride}; + sc_up += {sc_stride}; + bi_gate += {sc_stride}; + bi_up += {sc_stride}; + x_base += BLOCK_SIZE; + }} + + // Epilogue: apply inv_rms (factored), SwiGLU, write f32 + device float* yp = Y_shared + out_row; + for (int tm = 0; tm < RESULTS_PER_SG; tm++) {{ + float g = simd_sum(gate_result[tm]) * inv_rms; + float u = simd_sum(up_result[tm]) * inv_rms; + if (slid == 0) {{ + float silu_g = g / (1.0f + metal::exp(-g)); + yp[tm] = silu_g * u; + }} + }} + }} +""" + + +_oproj_softmax_swiglu_8bit_kernels = {} + + +def _get_oproj_softmax_swiglu_8bit_kernel(group_size=64, scale_bf16=True, + n_experts=64, top_k=10, + norm_topk=True): + key = (group_size, scale_bf16, n_experts, top_k, norm_topk) + if key not in _oproj_softmax_swiglu_8bit_kernels: + sc_tag = "_bf16sc" if scale_bf16 else "" + nt_tag = "_nt" if norm_topk else "" + _oproj_softmax_swiglu_8bit_kernels[key] = mx.fast.metal_kernel( + name=(f"oproj_softmax_topk_swiglu_8bit_gs{group_size}" + f"_e{n_experts}_k{top_k}{sc_tag}{nt_tag}"), + input_names=[ + "W", "S", "B_q", # routed expert weights + "W_shared", "S_shared", "B_shared", # shared expert weights + "X", # h_scaled (K,) bf16 + "gate_part_a", "gate_part_b", # (E,) f32 + "x2_partials", # (N_OPROJ_TG,) f32 + "W_seg", "S_seg", "B_seg", # shared_expert_gate weights + "N_INTER_val", "SHARED_INTER_val", + "K_val", "n_active_val", "N_OPROJ_TG_val", + ], + output_names=["Y_routed", "Y_shared", "out_inds", + "norm_scores", "gate_raw"], + source=_gen_oproj_softmax_topk_swiglu_8bit_source( + group_size, scale_bf16, n_experts, top_k, norm_topk), + ) + return _oproj_softmax_swiglu_8bit_kernels[key] + + +def fused_oproj_softmax_topk_swiglu_8bit( + w_q, s, b_q, # routed expert weights + w_shared, s_shared, b_shared, # shared expert weights + h_scaled, # (K,) bf16 — h * w_rms + gate_part_a, gate_part_b, # (E,) f32 — gate decomposition + x2_partials, # (N_OPROJ_TG,) f32 — per-TG x² + w_seg, s_seg, b_seg, # shared_expert_gate 8-bit weights + n_inter, k_hidden, # MoE dimensions + n_experts, top_k, # routing params + n_oproj_tg, # number of o_proj TGs (for x²) + group_size=64, + shared_inter=None, + norm_topk=True, +): + """Single-dispatch softmax + top-k + merged 8-bit SwiGLU with oproj prologue. + + Prologue: + Phase 1: distributed x² → inv_rms (all TGs) + Phase 2: softmax(gate_part_a + gate_part_b) → top-k → norm (routed TGs) + Phase 3: shared_expert_gate 8-bit GEMV → gate_raw (TG(0,0,0) SG 0) + + Main: 8-bit gate+up+SwiGLU for routed + shared experts. + inv_rms is factored out of the inner loop and applied once in epilogue. + + Args: + h_scaled: (K,) bf16 — h * w_rms from Dispatch 1 + gate_part_a: (E,) f32 — M1 @ attn_out from Dispatch 1 + gate_part_b: (E,) f32 — W_fused @ residual from Dispatch 1 + x2_partials: (N_OPROJ_TG,) f32 — per-TG Σh² from Dispatch 1 + w_seg/s_seg/b_seg: shared_expert_gate 8-bit weights (1, K/4) uint32 + + Returns: + (Y_routed, Y_shared, out_inds, norm_scores, gate_raw): + Y_routed: (top_k, n_inter) f32 + Y_shared: (shared_inter,) f32 + out_inds: (top_k,) uint32 + norm_scores: (top_k,) f32 + gate_raw: (1,) f32 — raw shared expert gate value (sigmoid in epilogue) + """ + scale_bf16 = (s.dtype == mx.bfloat16) + kern = _get_oproj_softmax_swiglu_8bit_kernel( + group_size, scale_bf16, n_experts, top_k, norm_topk) + + n_inter_val = int(n_inter) + shared_inter_val = int(shared_inter) if shared_inter is not None else n_inter_val + n_active = int(top_k) + max_inter = max(n_inter_val, shared_inter_val) + + results = kern( + inputs=[ + w_q, s, b_q, + w_shared, s_shared, b_shared, + h_scaled, + gate_part_a, gate_part_b, + x2_partials, + w_seg, s_seg, b_seg, + n_inter_val, shared_inter_val, + int(k_hidden), n_active, int(n_oproj_tg), + ], + output_shapes=[ + (n_active, n_inter_val), # Y_routed + (shared_inter_val,), # Y_shared + (n_active,), # out_inds + (n_active,), # norm_scores + (1,), # gate_raw + ], + output_dtypes=[ + mx.float32, # Y_routed + mx.float32, # Y_shared + mx.uint32, # out_inds + mx.float32, # norm_scores + mx.float32, # gate_raw + ], + grid=(32, ceil_div(max_inter, 8) * 2, n_active + 1), + threadgroup=(32, 2, 1), + ) + return results[0], results[1], results[2], results[3], results[4] diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/moe.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/moe.py new file mode 100644 index 00000000..5a733f15 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/moe.py @@ -0,0 +1,245 @@ +"""MoE __call__ variants for Qwen3.5. + +Two modes: + _fused_moe_call: gate→SwiGLU→down_proj→epilogue (~15 dispatches, fused to ~4+MLX) + _oproj_moe_call: o_proj+gate→SwiGLU(w/softmax prologue)→down_proj→epilogue (4 dispatches) + +Kernels used: + merged_routed_shared_swiglu_8bit.py — 8-bit gate+up+SwiGLU (_fused_moe_call) + merged_down_proj_8bit.py — 8-bit down_proj (both modes) + fused_moe_epilogue_qwen.py — weighted sum + sigmoid gate + residual (both modes) + custom_oproj_gate_gemv_8bit.py — 8-bit o_proj + bf16 gate GEMVs (_oproj_moe_call) + oproj_softmax_topk_swiglu_8bit.py — SwiGLU with softmax prologue (_oproj_moe_call) + +Adapted from mlx_bench/model_patches/qwen/moe.py. +""" + +import mlx.core as mx + +from .kernels.merged_routed_shared_swiglu_8bit import fused_merged_gate_up_swiglu_8bit +from .kernels.merged_down_proj_8bit import fused_merged_down_proj_8bit +from .kernels.fused_moe_epilogue_qwen import fused_moe_epilogue_qwen +from .kernels.custom_oproj_gate_gemv_8bit import fused_custom_oproj_8bit +from .kernels.oproj_softmax_topk_swiglu_8bit import fused_oproj_softmax_topk_swiglu_8bit + + +def _vanilla_moe_call(self, x): + """Original Qwen3NextSparseMoeBlock.__call__ (for prefill fallback).""" + gates = self.gate(x) + gates = mx.softmax(gates, axis=-1, precise=True) + + k = self.top_k + inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:] + scores = mx.take_along_axis(gates, inds, axis=-1) + if self.norm_topk_prob: + scores = scores / scores.sum(axis=-1, keepdims=True) + + y = self.switch_mlp(x, inds) + y = (y * scores[..., None]).sum(axis=-2) + + shared_y = self.shared_expert(x) + shared_y = mx.sigmoid(self.shared_expert_gate(x)) * shared_y + + return y + shared_y + + +def _fused_moe_call(self, x, _residual=None): + """Qwen3.5 MoE with fused kernels (4 custom dispatches). + + Falls back to vanilla for prefill (seq_len > 1). + + Args: + x: (B, S, K) bf16 — post-layernorm hidden state + _residual: (B, S, K) bf16 — pre-layernorm hidden state for residual add. + If None, epilogue skips residual (returns MoE output only). + """ + # Fused kernels are decode-only (seq_len=1). Fall back for prefill. + if x.shape[-2] > 1: + # Try vanilla switch_mlp path if weights still exist + has_switch_weights = hasattr(self.switch_mlp, 'gate_proj') and \ + hasattr(self.switch_mlp.gate_proj, 'weight') + if has_switch_weights: + out = _vanilla_moe_call(self, x) + else: + # Weights were freed — process tokens one by one through fused path + outs = [] + for t in range(x.shape[-2]): + xt = x[:, t:t+1, :] + res_t = _residual[:, t:t+1, :] if _residual is not None else None + outs.append(_fused_moe_call(self, xt, _residual=res_t)) + return mx.concatenate(outs, axis=1) + if _residual is not None: + out = out + _residual + return out + + # ── Gate routing (vanilla MLX ops) ── + gates = self.gate(x) + gates = mx.softmax(gates, axis=-1, precise=True) + + k = self.top_k + inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:] + scores = mx.take_along_axis(gates, inds, axis=-1) + if self.norm_topk_prob: + scores = scores / scores.sum(axis=-1, keepdims=True) + + x_flat = x.reshape(-1).astype(mx.bfloat16) # (K,) + inds_flat = inds.reshape(-1).astype(mx.uint32) + scores_flat = scores.reshape(-1).astype(mx.float32) + + # ── Dispatch 1: Merged gate+up+SwiGLU (8-bit routed + 8-bit shared) ── + y_routed, y_shared = fused_merged_gate_up_swiglu_8bit( + self.switch_mlp._fused_w_gu, + self.switch_mlp._fused_s_gu, + self.switch_mlp._fused_b_gu, + self._shared_w_gu, + self._shared_s_gu, + self._shared_b_gu, + x_flat, + inds_flat, + self.switch_mlp._fused_n_inter, + self.switch_mlp._fused_k_hidden, + group_size=self.switch_mlp._fused_group_size, + shared_inter=self._shared_inter, + ) + + # ── Dispatch 2: Merged down_proj (8-bit routed + 8-bit shared) ── + y_down_routed, y_down_shared = fused_merged_down_proj_8bit( + self._down_w, + self._down_s, + self._down_b, + self._shared_down_w, + self._shared_down_s, + self._shared_down_b, + y_routed, + y_shared, + inds_flat, + self._down_K, + self._down_N, + group_size=self._down_gs, + shared_n_in=self._shared_inter, + ) + + # ── Dispatch 3: Shared expert gate (small GEMV) ── + shared_gate_out = self.shared_expert_gate(x.reshape(1, 1, -1)) + shared_gate_val = mx.sigmoid(shared_gate_out.reshape(())) + + # ── Dispatch 4: Fused epilogue ── + if _residual is not None: + h_flat = _residual.reshape(-1).astype(mx.bfloat16) + else: + h_flat = mx.zeros((self.switch_mlp._fused_k_hidden,), dtype=mx.bfloat16) + + y = fused_moe_epilogue_qwen( + y_down_routed, # (n_active, K) f32 + y_down_shared, # (K,) f32 + scores_flat, # (n_active,) f32 + h_flat, # (K,) bf16 + shared_gate_val, # scalar f32 + self._down_K, # K + ) + return y.reshape(1, 1, -1) + + +def _oproj_moe_call(self, attn_out_3d, _residual=None): + """Qwen3.5 MoE with fused o_proj + gate GEMVs (4 custom dispatches). + + Receives raw attention output (pre-o_proj) and residual. + Fuses o_proj, RMSNorm, gate softmax, top-k, score norm, shared_expert_gate, + SwiGLU, down_proj, and epilogue into 4 dispatches. + + Dispatch 1: 8-bit o_proj + bf16 M1/W_fused GEMVs → h_scaled, h_out, x2_partials, + gate_part_a, gate_part_b + Dispatch 2: SwiGLU with softmax prologue (inv_rms + softmax + top-k + score norm + + shared_expert_gate in TG(0,0,0)) → y_routed, y_shared, inds, scores, + gate_raw + Dispatch 3: Merged down_proj (unchanged) + Dispatch 4: Fused epilogue with sigmoid(gate_raw) + + Args: + attn_out_3d: (1, 1, K_attn) bf16 — raw attention output (pre-o_proj) + _residual: (1, 1, K) bf16 — input residual (before attention) + """ + # Prefill fallback (S > 1): restore o_proj + vanilla MoE + if attn_out_3d.shape[-2] > 1: + from .decoder import _parent_layer_map + parent = _parent_layer_map[id(self)] + if parent.is_linear: + oproj_mod = parent.linear_attn.out_proj + else: + oproj_mod = parent.self_attn.o_proj + projected = oproj_mod(attn_out_3d) + h = _residual + projected + out = _vanilla_moe_call(self, parent.post_attention_layernorm(h)) + return out + h + + attn_out = attn_out_3d.reshape(-1).astype(mx.bfloat16) # (K_attn,) + residual = _residual.reshape(-1).astype(mx.bfloat16) # (K,) + + K = self._oproj_M + K_attn = self._oproj_K_attn + + # ── Dispatch 1: fused o_proj + M1 GEMV + W_fused GEMV ── + h_scaled, h_out, x2_partials, gate_part_a, gate_part_b = \ + fused_custom_oproj_8bit( + self._oproj_w, self._oproj_s, self._oproj_b, + attn_out, residual, self._oproj_rms_weight, + self._oproj_M1, self._oproj_W_fused, + M=K, K_attn=K_attn, K_hidden=self._oproj_K_hidden, + n_experts=self._oproj_n_experts, + gate_bm=self._oproj_gate_bm, + ) + + # ── Dispatch 2: SwiGLU with softmax prologue ── + y_routed, y_shared, inds, scores, gate_raw = \ + fused_oproj_softmax_topk_swiglu_8bit( + self.switch_mlp._fused_w_gu, + self.switch_mlp._fused_s_gu, + self.switch_mlp._fused_b_gu, + self._shared_w_gu, + self._shared_s_gu, + self._shared_b_gu, + h_scaled, + gate_part_a, + gate_part_b, + x2_partials, + self._seg_w, + self._seg_s, + self._seg_b, + n_inter=self.switch_mlp._fused_n_inter, + k_hidden=K, + n_experts=self._oproj_n_experts, + top_k=self.top_k, + n_oproj_tg=self._oproj_n_tg, + group_size=self.switch_mlp._fused_group_size, + shared_inter=self._shared_inter, + norm_topk=self.norm_topk_prob, + ) + + # ── Dispatch 3: Merged down_proj (8-bit routed + 8-bit shared) ── + y_down_routed, y_down_shared = fused_merged_down_proj_8bit( + self._down_w, + self._down_s, + self._down_b, + self._shared_down_w, + self._shared_down_s, + self._shared_down_b, + y_routed, + y_shared, + inds, + self._down_K, + self._down_N, + group_size=self._down_gs, + shared_n_in=self._shared_inter, + ) + + # ── Dispatch 4: Fused epilogue with sigmoid(gate_raw) ── + y = fused_moe_epilogue_qwen( + y_down_routed, # (n_active, K) f32 + y_down_shared, # (K,) f32 + scores, # (n_active,) f32 — norm_topk_prob from prologue + h_out, # (K,) bf16 — post-o_proj hidden for residual add + gate_raw, # scalar f32 — raw dot product, sigmoid fused + K, + fuse_sigmoid=True, + ) + return y.reshape(1, 1, -1) diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 8b22e6b3..3798e7d1 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -190,6 +190,10 @@ def load_mlx_items( mx.eval(model) end_time = time.perf_counter() logger.info(f"Time taken to load model: {(end_time - start_time):.2f}s") + + from exo.worker.engines.mlx.patches import maybe_apply_patches + maybe_apply_patches(model, model_path) + tokenizer = get_tokenizer(model_path, bound_instance.bound_shard) else: