test adding new kernels in exo
This commit is contained in:
davidmcc73
@@ -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)
|
||||
@@ -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")
|
||||
@@ -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)")
|
||||
@@ -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
|
||||
+324
@@ -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<bfloat16_t>(h * float(w_rms[k]));
|
||||
h_out[k] = static_cast<bfloat16_t>(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]
|
||||
@@ -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<bfloat16_t>(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]
|
||||
@@ -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]
|
||||
+274
@@ -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]
|
||||
+493
@@ -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]
|
||||
@@ -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)
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user