electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Adding more kernels to qwen3.5 8bit models

Browse Source
This commit is contained in:
dmcc73
2026-03-15 01:32:48 +00:00
parent 613d47aa1f
commit d3507d2cd6
14 changed files with 2136 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/exo/worker/engines/mlx/patches/__init__.py
+16 -7
View File
@@ -2,10 +2,12 @@
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
- Qwen3.5 MoE (model_type: qwen3_5_moe):
EXO_FUSED_KERNELS=1: oproj mode (4-dispatch MoE fusion)
EXO_FUSED_KERNELS=2: fused_gqa_gdn mode (full GDN + GQA + MoE fusion, default)
Set EXO_FUSED_KERNELS=0 to disable patches (baseline mode).
Default: enabled (EXO_FUSED_KERNELS=1).
Set EXO_FUSED_KERNELS=0 to disable all patches (vanilla mode).
Default: EXO_FUSED_KERNELS=2 (fused_gqa_gdn, best performance).
"""
import json
@@ -18,7 +20,8 @@ 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":
fused_mode = os.environ.get("EXO_FUSED_KERNELS", "2")
if fused_mode == "0":
logger.info("Kernel fusion patches disabled (EXO_FUSED_KERNELS=0)")
return
@@ -32,7 +35,13 @@ def maybe_apply_patches(model: nn.Module, model_path: Path) -> None:
model_type = config.get("model_type", "")
if model_type == "qwen3_5_moe":
from .qwen3_5_moe.apply import apply_qwen35_oproj_patches
if fused_mode == "1":
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)
logger.info("Detected Qwen3.5 MoE model, applying oproj fusion patches")
apply_qwen35_oproj_patches(model)
else:
from .qwen3_5_moe.apply import apply_qwen35_fused_gqa_gdn_patches
logger.info("Detected Qwen3.5 MoE model, applying fused GQA+GDN patches")
apply_qwen35_fused_gqa_gdn_patches(model)
src/exo/worker/engines/mlx/patches/qwen3_5_moe/apply.py
+19 -2
View File
@@ -1,6 +1,7 @@
"""Apply oproj fusion patches to Qwen3.5 MoE models.
"""Apply kernel fusion patches to Qwen3.5 MoE models.
Entry point called from patches/__init__.py after model type detection.
Supports oproj mode (MoE only) and fused_gqa_gdn mode (full attention + MoE).
"""
import time
@@ -8,7 +9,7 @@ import time
import mlx.nn as nn
from loguru import logger
from .common import apply_oproj_fused_patches
from .common import apply_oproj_fused_patches, apply_fused_gqa_gdn_patches
def apply_qwen35_oproj_patches(model: nn.Module) -> None:
@@ -25,3 +26,19 @@ def apply_qwen35_oproj_patches(model: nn.Module) -> None:
t_patch = time.time() - t0
logger.info(f"Qwen3.5 oproj fusion: patched {n_layers} layers in {t_patch:.1f}s")
def apply_qwen35_fused_gqa_gdn_patches(model: nn.Module) -> None:
"""Apply full fusion (GDN + GQA attention + oproj MoE) to all layers.
Fused GDN attention (3/4 layers) + fused GQA attention (1/4 layers)
+ oproj MoE (all layers). 44% faster than vanilla on Qwen3.5-35B-A3B.
"""
layers = model.layers # type: ignore[attr-defined]
n_layers = len(layers)
t0 = time.time()
apply_fused_gqa_gdn_patches(layers, gate_bm=8, free_originals=False)
t_patch = time.time() - t0
logger.info(f"Qwen3.5 fused GQA+GDN: patched {n_layers} layers in {t_patch:.1f}s")
src/exo/worker/engines/mlx/patches/qwen3_5_moe/common.py
+160 -1
View File
@@ -1,4 +1,8 @@
"""Weight preparation and patch orchestration for Qwen3.5 oproj fusion.
"""Weight preparation and patch orchestration for Qwen3.5 kernel fusion.
Supports:
apply_oproj_fused_patches — oproj mode (4-dispatch MoE only)
apply_fused_gqa_gdn_patches — full fusion (fused GDN + GQA attention + oproj MoE)
Adapted from mlx_bench/model_patches/qwen/common.py.
"""
@@ -185,3 +189,158 @@ def apply_oproj_fused_patches(layers, gate_bm=8, free_originals=False):
Qwen3NextSparseMoeBlock.__call__ = _oproj_moe_call
DecoderLayer.__call__ = _oproj_decoder_call
logger.info(f" Patched {n_patched} MoE blocks (oproj mode, 4 dispatches)")
def _patch_gdn_proj_weights(attn):
"""Merge all 4 GDN projection weights into contiguous buffers."""
W_merged = mx.concatenate([
attn.in_proj_qkv.weight,
attn.in_proj_z.weight,
attn.in_proj_b.weight,
attn.in_proj_a.weight,
], axis=0)
S_merged = mx.concatenate([
attn.in_proj_qkv.scales,
attn.in_proj_z.scales,
attn.in_proj_b.scales,
attn.in_proj_a.scales,
], axis=0)
B_merged = mx.concatenate([
attn.in_proj_qkv.biases,
attn.in_proj_z.biases,
attn.in_proj_b.biases,
attn.in_proj_a.biases,
], axis=0)
attn._merged_proj_w = W_merged
attn._merged_proj_s = S_merged
attn._merged_proj_b = B_merged
attn._merged_proj_dims = (
attn.in_proj_qkv.weight.shape[0],
attn.in_proj_z.weight.shape[0],
attn.in_proj_b.weight.shape[0],
attn.in_proj_a.weight.shape[0],
)
mx.eval(W_merged, S_merged, B_merged)
def _patch_gqa_proj_weights(attn):
"""Merge GQA q_proj, k_proj, v_proj weights with q_proj row permutation.
q_proj rows are interleaved [head0_q, head0_gate, head1_q, ...].
Permute so queries come first, then gate, then k, then v.
Pre-cache constant scalar arrays for kernel dispatch.
"""
q = attn.q_proj
k = attn.k_proj
v = attn.v_proj
H_q = attn.num_attention_heads
D = attn.head_dim
W_q = q.weight.reshape(H_q, 2 * D, -1)
S_q = q.scales.reshape(H_q, 2 * D, -1)
B_q = q.biases.reshape(H_q, 2 * D, -1)
W_queries = W_q[:, :D, :].reshape(H_q * D, -1)
W_gate = W_q[:, D:, :].reshape(H_q * D, -1)
S_queries = S_q[:, :D, :].reshape(H_q * D, -1)
S_gate = S_q[:, D:, :].reshape(H_q * D, -1)
B_queries = B_q[:, :D, :].reshape(H_q * D, -1)
B_gate = B_q[:, D:, :].reshape(H_q * D, -1)
W_merged = mx.contiguous(mx.concatenate([W_queries, W_gate, k.weight, v.weight], axis=0))
S_merged = mx.contiguous(mx.concatenate([S_queries, S_gate, k.scales, v.scales], axis=0))
B_merged = mx.contiguous(mx.concatenate([B_queries, B_gate, k.biases, v.biases], axis=0))
attn._merged_proj_w = W_merged
attn._merged_proj_s = S_merged
attn._merged_proj_b = B_merged
N_Q = H_q * D
N_GATE = H_q * D
N_K = k.weight.shape[0]
N_V = v.weight.shape[0]
attn._merged_proj_dims = (N_Q, N_GATE, N_K, N_V)
mx.eval(W_merged, S_merged, B_merged)
# Pre-cache constant scalar arrays for kernel dispatch
N_TOTAL = N_Q + N_GATE + N_K + N_V
K_dim = q.weight.shape[1] * 4 # 8-bit: pack_factor=4
attn._kernel_scalars = {
'K': mx.array(K_dim, dtype=mx.int32),
'N_Q': mx.array(N_Q, dtype=mx.int32),
'N_GATE': mx.array(N_GATE, dtype=mx.int32),
'N_K': mx.array(N_K, dtype=mx.int32),
'N_TOTAL': mx.array(N_TOTAL, dtype=mx.int32),
'N_Q_TG': mx.array(ceil_div(N_Q, 8), dtype=mx.int32),
'N_GATE_TG': mx.array(ceil_div(N_GATE, 8), dtype=mx.int32),
'N_K_TG': mx.array(ceil_div(N_K, 8), dtype=mx.int32),
'scale': mx.array(attn.head_dim ** -0.5, dtype=mx.float32),
'H_Q': mx.array(attn.num_attention_heads, dtype=mx.int32),
'H_KV': mx.array(attn.num_key_value_heads, dtype=mx.int32),
'N_blocks': mx.array(128, dtype=mx.int32),
}
mx.eval(*attn._kernel_scalars.values())
N_V_TG = ceil_div(N_V, 8)
attn._d1_total_tg = ceil_div(N_Q, 8) + ceil_div(N_GATE, 8) + ceil_div(N_K, 8) + N_V_TG
# Precompute RoPE inv_freq
rope_dims = attn.rope.dims
half_dims = rope_dims // 2
theta = attn.rope.base
d_indices = mx.arange(half_dims, dtype=mx.float32)
attn._rope_inv_freq = theta ** (-d_indices / half_dims)
mx.eval(attn._rope_inv_freq)
def apply_fused_gqa_gdn_patches(layers, gate_bm=8, free_originals=False):
"""Apply all fusions: fused GDN + fused GQA + oproj MoE.
Combines:
- Fused GDN attention (3/4 layers: GatedDeltaNet)
- Fused GQA attention (1/4 layers: Qwen3NextAttention)
- Oproj MoE (all layers: oproj_gate_gemv + fused MoE dispatches)
"""
from .moe import _oproj_moe_call
from .decoder import _fused_gdn_decoder_call
from .fused_gdn_attention import _fused_gdn_call
from .fused_gqa_attention import _fused_gqa_call
from mlx_lm.models.qwen3_next import Qwen3NextAttention
from mlx_lm.models.qwen3_5 import GatedDeltaNet
n_patched = 0
n_gdn = 0
n_gqa = 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 layer.is_linear:
_patch_gdn_proj_weights(layer.linear_attn)
n_gdn += 1
else:
_patch_gqa_proj_weights(layer.self_attn)
n_gqa += 1
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)} (fused GQA+GDN mode)")
GatedDeltaNet.__call__ = _fused_gdn_call
Qwen3NextAttention.__call__ = _fused_gqa_call
Qwen3NextSparseMoeBlock.__call__ = _oproj_moe_call
DecoderLayer.__call__ = _fused_gdn_decoder_call
logger.info(f" Patched {n_patched} MoE blocks ({n_gdn} GDN + {n_gqa} GQA, fused GQA+GDN mode)")
src/exo/worker/engines/mlx/patches/qwen3_5_moe/decoder.py
+19 -1
View File
@@ -1,8 +1,9 @@
"""Decoder layer __call__ variants for Qwen3.5.
Two modes:
Three modes:
_fused_decoder_call: passes residual to fused MoE epilogue (~15 dispatches)
_oproj_decoder_call: fuses o_proj + RMSNorm + gate GEMV (4 dispatches)
_fused_gdn_decoder_call: fused GDN/GQA attention + oproj MoE (6 dispatches)
Attention patches for oproj mode:
_pre_oproj_attention_call: Qwen3NextAttention.__call__ that skips o_proj
@@ -65,6 +66,23 @@ def _oproj_decoder_call(self, x, mask=None, cache=None):
return self.mlp(pre_oproj, _residual=x)
def _fused_gdn_decoder_call(self, x, mask=None, cache=None):
"""Decoder with fused GDN/GQA attention + oproj MoE (6-dispatch mode).
GDN layers use fused kernels, GQA layers use fused or vanilla attention.
Both return pre-out_proj output. MoE handles oproj_gate_gemv + MoE dispatches.
Flow is identical to oproj mode — the difference is that GatedDeltaNet.__call__
and/or Qwen3NextAttention.__call__ are patched with fused kernel implementations.
"""
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.
src/exo/worker/engines/mlx/patches/qwen3_5_moe/fused_gdn_attention.py
+161
View File
@@ -0,0 +1,161 @@
"""Fused GDN attention __call__ for qwen3_5.GatedDeltaNet (Dispatches 2-5).
Replaces the vanilla GatedDeltaNet.__call__ with fused kernel dispatches:
Dispatch 2: fused_gdn_projections — merged 8-bit GEMV + conv1d + SiLU(qkv) + SiLU(z)
+ sigmoid(b)→beta + g=exp(-exp(A_log)*softplus(a+dt_bias))
Dispatch 3: fused_qk_rmsnorm — per-head L2-norm on q (×Dk^(-½)) and k
Dispatch 4: gated_delta_kernel — GDN recurrence (receives pre-computed g, beta)
Dispatch 5: fused_rms_norm_gated — RMSNorm(out, weight) × z_silu
All 4 projection weights are pre-merged into contiguous buffers at patch time
(_patch_gdn_proj_weights) for better memory locality.
g/beta computation is fused into Dispatch 2 epilogues, eliminating ~8 micro-
dispatches that gated_delta_update would otherwise generate.
Fused path is decode-only (S=1). For prefill (S>1), falls back to vanilla ops.
Returns pre-out_proj output (same interface as _pre_oproj_qwen35_linear_attn_call).
Dispatch 1 (input_layernorm) is handled by the decoder.
Dispatch 6 (oproj_gate_gemv) is handled by the MoE __call__.
"""
from typing import Any, Optional
import mlx.core as mx
import mlx.nn as nn
from .kernels.fused_gdn_projections_8bit import fused_gdn_projections
from .kernels.fused_qk_rmsnorm import fused_qk_rmsnorm
from .kernels.fused_rms_norm_gated import fused_rms_norm_gated
def _vanilla_gdn_call(self, inputs, mask, cache):
"""Vanilla GDN path for prefill (S>1). Returns pre-out_proj output."""
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 * q * mx.rsqrt(
(q * q).sum(axis=-1, keepdims=True) + 1e-6
)
k = k * mx.rsqrt(
(k * k).sum(axis=-1, keepdims=True) + 1e-6
)
out, state = gated_delta_update(
q, k, v, a, b,
self.A_log, self.dt_bias,
state, mask,
use_kernel=True,
)
if cache is not None:
cache[1] = state
out = self.norm(out, z)
return out.reshape(B, S, -1) # skip out_proj
def _fused_gdn_call(
self,
inputs: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array:
"""Fused GDN attention: merged projections + existing GDN kernel.
Decode (S=1): uses fused kernels with merged weight buffers.
Prefill (S>1): falls back to vanilla ops.
Returns pre-out_proj output [B, S, value_dim] for Dispatch 6.
"""
B, S, _ = inputs.shape
# Prefill fallback: fused kernels are decode-only (S=1)
if S > 1:
return _vanilla_gdn_call(self, inputs, mask, cache)
from mlx_lm.models.gated_delta import gated_delta_kernel
# ── Cache: conv state ──
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,
)
# ── Dispatch 2: fused projections (merged GEMV + conv + SiLU + g/beta) ──
qkv_conv_silu, z_silu, beta, g, conv_state_out = fused_gdn_projections(
inputs,
self._merged_proj_w, self._merged_proj_s, self._merged_proj_b,
self._merged_proj_dims,
conv_state, self.conv1d.weight,
self.A_log, self.dt_bias,
batch_size=B,
)
if cache is not None:
cache[0] = conv_state_out
# ── Dispatch 3: fused Q/K L2-norm ──
qk_normed = fused_qk_rmsnorm(qkv_conv_silu, batch_size=B)
# ── Split q, k from normed output; v from conv output ──
q = qk_normed[:, :, :self.key_dim].reshape(B, S, self.num_k_heads, self.head_k_dim)
k = qk_normed[:, :, self.key_dim:].reshape(B, S, self.num_k_heads, self.head_k_dim)
v = qkv_conv_silu[:, :, 2 * self.key_dim:].reshape(B, S, self.num_v_heads, self.head_v_dim)
# ── Dispatch 4: GDN recurrence with pre-computed g/beta ──
state = cache[1] if cache else None
if state is None:
state = mx.zeros(
(B, self.num_v_heads, self.head_v_dim, self.head_k_dim),
dtype=inputs.dtype,
)
out, state_new = gated_delta_kernel(
q, k, v, g, beta, state, mask,
)
if cache is not None:
cache[1] = state_new
# ── Dispatch 5: fused RMSNorm × z_silu ──
norm_weight = self.norm.weight
result = fused_rms_norm_gated(out, z_silu, norm_weight, batch_size=B)
return result # [B, S, value_dim] — skip out_proj (handled by Dispatch 6)
src/exo/worker/engines/mlx/patches/qwen3_5_moe/fused_gqa_attention.py
+137
View File
@@ -0,0 +1,137 @@
"""Fused GQA attention __call__ for qwen3_next.Qwen3NextAttention.
Replaces vanilla Qwen3NextAttention.__call__ with fused kernel dispatches:
Dispatch 1: fused_gqa_projections — merged 8-bit GEMV (q+gate+k+v) + sigmoid(gate)
Dispatch 2: fused_qk_norm_rope — RMSNorm + RoPE (TODO: custom kernel)
Dispatch 3: KV cache update (MLX built-in)
Dispatch 4: custom_sdpa_pass1 (TODO: custom kernel)
Dispatch 5: custom_sdpa_pass2_gate (TODO: custom kernel, includes gate multiply)
Dispatch 6: oproj_gate_gemv (existing, handled by MoE __call__)
Dispatches 1-2, 4-5 are custom kernels, Dispatch 3 is MLX built-in.
Returns pre-out_proj output (output * sigmoid(gate)) for Dispatch 6.
Fused path is decode-only (S=1). For prefill (S>1), falls back to vanilla ops.
"""
from typing import Any, Optional
import mlx.core as mx
import mlx.nn as nn
from .kernels.fused_gqa_projections_8bit import fused_gqa_projections
from .kernels.fused_qk_norm_rope import fused_qk_norm_rope
from .kernels.custom_sdpa_pass1 import custom_sdpa_pass1
from .kernels.custom_sdpa_pass2_gate import custom_sdpa_pass2_gate
def _vanilla_gqa_call(self, x, mask, cache):
"""Vanilla GQA path for prefill (S>1). Returns pre-out_proj output."""
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 _fused_gqa_call(
self,
x: mx.array,
mask: Optional[mx.array] = None,
cache: Optional[Any] = None,
) -> mx.array:
"""Fused GQA attention: custom projection + norm/rope kernels.
Decode (S=1): Dispatch 1 (fused GEMV) + Dispatch 2 (fused norm+rope)
+ vanilla cache + SDPA (Dispatches 3-5).
Prefill (S>1): falls back to fully vanilla ops.
Returns pre-out_proj output [B, S, H_q*D] for Dispatch 6.
"""
B, S, _ = x.shape
# Prefill fallback: fused kernels are decode-only (S=1)
if S > 1:
return _vanilla_gqa_call(self, x, mask, cache)
H_q = self.num_attention_heads
H_kv = self.num_key_value_heads
D = self.head_dim
# ── Dispatch 1: fused projections (merged GEMV + sigmoid(gate)) ──
sc = getattr(self, '_kernel_scalars', None)
queries, gate_sigmoid, keys, values = fused_gqa_projections(
x,
self._merged_proj_w, self._merged_proj_s, self._merged_proj_b,
self._merged_proj_dims,
batch_size=B,
scalars=sc, total_tg=getattr(self, '_d1_total_tg', None),
)
# ── Dispatch 2: fused Q/K RMSNorm + RoPE ──
queries, keys = fused_qk_norm_rope(
queries, keys,
self.q_norm.weight, self.k_norm.weight,
self._rope_inv_freq, cache.offset,
H_q, H_kv, D, batch_size=B,
)
# queries: [B, H_q, 1, D], keys: [B, H_kv, 1, D]
# Reshape directly to (B, H_kv, 1, D) — no transpose needed since S=1.
# Avoids a 4 MiB copy dispatch that transpose would trigger.
values = values.reshape(B, H_kv, 1, D)
# ── Dispatch 3: KV cache update ──
cache.update_and_fetch(keys, values)
N = cache.offset # actual sequence length after update
alloc_len = cache.keys.shape[2] # allocated buffer length
# ── Dispatch 4: SDPA Pass 1 (online softmax + partial V accumulation) ──
blocks = 128 # M3 Ultra default for N >= 1024
if N < 1024:
# Short sequence: fall back to vanilla SDPA + gate multiply
# Use sliced views for built-in SDPA (handles strides natively)
k_sliced = cache.keys[:, :, :N, :]
v_sliced = cache.values[:, :, :N, :]
from mlx_lm.models.qwen3_next import scaled_dot_product_attention
output = scaled_dot_product_attention(
queries, k_sliced, v_sliced, cache=cache, scale=self.scale, mask=mask
)
output = output.transpose(0, 2, 1, 3).reshape(B, S, -1)
return output * gate_sigmoid.astype(output.dtype)
# Pass full (contiguous) cache buffers + alloc_len to avoid copy dispatch
o_partials, sums, maxs = custom_sdpa_pass1(
queries, cache.keys, cache.values, self.scale,
H_q, H_kv, D, blocks=blocks, batch_size=B,
N=N, alloc_len=alloc_len, scalars=sc,
)
# ── Dispatch 5: SDPA Pass 2 + gate multiply ──
return custom_sdpa_pass2_gate(
o_partials, sums, maxs, gate_sigmoid,
H_q, D, blocks=blocks, V_SPLIT=4, batch_size=B,
scalars=sc,
)
src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/custom_sdpa_pass1.py
+216
View File
@@ -0,0 +1,216 @@
"""Custom SDPA Pass 1 for GQA attention (Dispatch 4).
Replicates MLX's sdpa_vector_2pass_1 kernel for decode (query_len=1).
Online softmax over N key positions, strided by `blocks`.
Each TG = (32, gqa_factor, 1):
- 32 threads in x = 1 SIMD group, each handles D/32 = 8 Q/K/V elements
- gqa_factor threads in y = one per Q head sharing the same KV head
- Grid z = blocks (each block handles a stride of the key sequence)
Algorithm per block:
1. Load query (once), scale by 1/sqrt(D)
2. For each key position i in [block_idx, N) with stride `blocks`:
a. Dot product: score = sum(q[j] * k[i,j]) via simd_sum
b. Online softmax: update max, sum_exp, rescale accumulators
c. Accumulate: o[j] += exp_score * v[i,j]
3. Write partials: o_partials[block, head, D], sums[block, head], maxs[block, head]
Matches MLX's sdpa_vector_2pass_1<bfloat16_t, 256, 256> exactly.
"""
import mlx.core as mx
def _gen_sdpa_pass1_source(D=256):
"""Generate Metal source for SDPA Pass 1.
Replicates MLX's sdpa_vector_2pass_1 for decode (S_q=1, no mask).
Inputs:
queries: [B, H_q, 1, D] bf16 (pre-scaled by 1/sqrt(D) NOT done here)
keys: [B, H_kv, alloc_len, D] bf16 (full cache buffer)
values: [B, H_kv, alloc_len, D] bf16 (full cache buffer)
scale: scalar f32
KV_alloc_len: int (allocated cache length, for head stride computation)
Outputs:
o_partials: [B, H_q, blocks, D] bf16
sums: [B, H_q, blocks] f32
maxs: [B, H_q, blocks] f32
Grid: (H_kv * 32, gqa_factor, blocks * B)
TG: (32, gqa_factor, 1)
We use grid.z = blocks * B to encode both block_idx and batch_idx:
block_idx = tgid.z % blocks
batch_idx = tgid.z / blocks
"""
EPT = D // 32 # elements per thread = 8
return f"""
const int D_DIM = {D};
const int EPT = {EPT};
uint simd_lid = thread_index_in_simdgroup;
uint kv_head_idx = threadgroup_position_in_grid.x;
uint gqa_lane = thread_index_in_threadgroup / 32; // which Q head within GQA group
// Decode block_idx and batch_idx from grid z
uint block_idx = threadgroup_position_in_grid.z % (uint)N_blocks;
uint batch_idx = threadgroup_position_in_grid.z / (uint)N_blocks;
int gqa_factor_val = (int)H_Q / (int)H_KV;
int q_head_idx = (int)kv_head_idx * gqa_factor_val + (int)gqa_lane;
int N_val = (int)N_keys;
// ── Pointer offsets ──
// queries: [B, H_q, 1, D] — one query per head
int q_offset = ((int)batch_idx * (int)H_Q + q_head_idx) * D_DIM;
// keys/values: [B, H_kv, alloc_len, D] — full cache buffer, read only first N_val positions
// Use alloc_len for head stride (not N_val) to handle non-contiguous cache slices
int alloc = (int)KV_alloc_len;
int kv_head_offset = ((int)batch_idx * (int)H_KV + (int)kv_head_idx) * alloc * D_DIM;
int k_offset = kv_head_offset + (int)block_idx * D_DIM;
int v_offset = kv_head_offset + (int)block_idx * D_DIM;
// Output offsets: indexed by [batch, q_head, block]
int out_head_offset = ((int)batch_idx * (int)H_Q + q_head_idx);
int o_offset = (out_head_offset * (int)N_blocks + (int)block_idx) * D_DIM;
int s_offset = out_head_offset * (int)N_blocks + (int)block_idx;
// ── Load query (once, scaled) ──
float q[{EPT}];
for (int i = 0; i < EPT; i++) {{
q[i] = (float)scale * (float)queries[q_offset + (int)simd_lid * EPT + i];
}}
// ── Online softmax + V accumulation ──
float max_score = -__FLT_MAX__;
float sum_exp = 0.0f;
float o[{EPT}] = {{0}};
int k_stride = (int)N_blocks * D_DIM; // stride between consecutive keys for this block
for (int pos = (int)block_idx; pos < N_val; pos += (int)N_blocks) {{
// Dot product: q @ k[pos]
float score = 0.0f;
for (int i = 0; i < EPT; i++) {{
score += q[i] * (float)keys[k_offset + (int)simd_lid * EPT + i];
}}
score = simd_sum(score);
// Online softmax update
float new_max = metal::max(max_score, score);
float factor = metal::fast::exp(max_score - new_max);
float exp_score = metal::fast::exp(score - new_max);
max_score = new_max;
sum_exp = sum_exp * factor + exp_score;
// Accumulate weighted value
for (int i = 0; i < EPT; i++) {{
o[i] = o[i] * factor + exp_score * (float)values[v_offset + (int)simd_lid * EPT + i];
}}
// Advance to next position for this block
k_offset += k_stride;
v_offset += k_stride;
}}
// ── Write partials ──
if (simd_lid == 0) {{
sums[s_offset] = sum_exp;
maxs[s_offset] = max_score;
}}
for (int i = 0; i < EPT; i++) {{
o_partials[o_offset + (int)simd_lid * EPT + i] = static_cast<bfloat16_t>(o[i]);
}}
"""
_pass1_cache = {}
def _get_pass1_kernel(D, gqa_factor):
"""Get or compile the SDPA Pass 1 kernel."""
key = (D, gqa_factor)
if key not in _pass1_cache:
_pass1_cache[key] = mx.fast.metal_kernel(
name="custom_sdpa_pass1",
input_names=["queries", "keys", "values",
"scale", "H_Q", "H_KV", "N_keys", "N_blocks",
"KV_alloc_len"],
output_names=["o_partials", "sums", "maxs"],
source=_gen_sdpa_pass1_source(D),
)
return _pass1_cache[key]
def custom_sdpa_pass1(queries, keys, values, scale, H_q, H_kv, D,
blocks=128, batch_size=1, N=None, alloc_len=None,
scalars=None):
"""SDPA Pass 1: online softmax + partial V accumulation.
Args:
queries: [B, H_q, 1, D] bf16
keys: [B, H_kv, N_or_alloc, D] bf16 — can be full cache buffer
values: [B, H_kv, N_or_alloc, D] bf16 — can be full cache buffer
scale: float (1/sqrt(D))
H_q, H_kv, D: int
blocks: int (number of blocks, default 128 for M3 Ultra)
batch_size: int
N: actual sequence length (if None, inferred from keys.shape[2])
alloc_len: allocated cache length (if None, same as N — contiguous)
scalars: dict of pre-cached mx.array scalars (optional)
Returns:
o_partials: [B, H_q, blocks, D] bf16
sums: [B, H_q, blocks] f32
maxs: [B, H_q, blocks] f32
"""
B = batch_size
gqa_factor = H_q // H_kv
kern = _get_pass1_kernel(D, gqa_factor)
if N is None:
N = keys.shape[2]
if alloc_len is None:
alloc_len = N
if scalars is not None:
s = scalars
# N and alloc_len change per call, must create fresh
n_keys_arr = mx.array(N, dtype=mx.int32)
alloc_arr = mx.array(alloc_len, dtype=mx.int32)
inputs = [queries, keys, values,
s['scale'], s['H_Q'], s['H_KV'], n_keys_arr, s['N_blocks'],
alloc_arr]
else:
inputs = [queries, keys, values,
mx.array(scale, dtype=mx.float32),
mx.array(H_q, dtype=mx.int32),
mx.array(H_kv, dtype=mx.int32),
mx.array(N, dtype=mx.int32),
mx.array(blocks, dtype=mx.int32),
mx.array(alloc_len, dtype=mx.int32)]
results = kern(
inputs=inputs,
output_shapes=[
(B * H_q * blocks * D,), # o_partials
(B * H_q * blocks,), # sums
(B * H_q * blocks,), # maxs
],
output_dtypes=[mx.bfloat16, mx.float32, mx.float32],
grid=(H_kv * 32, gqa_factor, blocks * B),
threadgroup=(32, gqa_factor, 1),
)
o_partials = results[0].reshape(B, H_q, blocks, D)
sums = results[1].reshape(B, H_q, blocks)
maxs = results[2].reshape(B, H_q, blocks)
return o_partials, sums, maxs
src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/custom_sdpa_pass2_gate.py
+200
View File
@@ -0,0 +1,200 @@
"""Custom SDPA Pass 2 + gate multiply for GQA attention (Dispatch 5).
Combines the block-partial results from Pass 1 into final SDPA output,
with gate multiply epilogue: output *= gate_sigmoid.
Modeled on MLX's sdpa_vector_2pass_2 kernel with two improvements:
1. V-dimension splitting (V_SPLIT=4): Instead of 1 TG per head (16 TGs,
20% M3 Ultra utilization), split the 256-dim output across 4 TGs per
head -> 64 TGs (80% utilization).
2. Gate multiply epilogue: Apply output *= gate_sigmoid in the same kernel.
Each TG = 1024 threads (32 SGs × 32 threads), matching MLX's structure:
- Handles D/V_SPLIT = 64 V-dimensions per TG
- 32 SGs process 32 blocks in parallel per iteration
- blocks/32 = 4 iterations to cover all 128 blocks
- Shared memory transpose + simd_sum to reduce across SGs
- Each thread handles EPT = 2 V-elements
Grid: (H_q * V_SPLIT * 32, 32, B)
TG: (32, 32, 1)
"""
import mlx.core as mx
def _gen_sdpa_pass2_gate_source(D=256, V_SPLIT=4):
"""Generate Metal source for SDPA Pass 2 + gate multiply.
Modeled on MLX's sdpa_vector_2pass_2 but with V_SPLIT for higher
TG count and fused gate multiply epilogue.
Inputs:
o_partials: [B, H_q, blocks, D] bf16
sums: [B, H_q, blocks] f32
maxs: [B, H_q, blocks] f32
gate_sigmoid: [B, 1, H_q * D] f32
Output:
attn_output: [B, 1, H_q * D] bf16
TG assignment:
tgid.x = head_idx * V_SPLIT + split_idx
Each TG handles D/V_SPLIT V-elements for one head.
32 SGs within TG process blocks in parallel.
"""
BN = 32 # number of SGs per TG = blocks processed in parallel
V_PER_SPLIT = D // V_SPLIT # 64
EPT = V_PER_SPLIT // BN # 2 elements per thread
return f"""
const int D_DIM = {D};
const int V_SPLIT = {V_SPLIT};
const int V_PER_SPLIT = {V_PER_SPLIT};
const int EPT = {EPT};
const int BN = {BN};
uint tg_idx = threadgroup_position_in_grid.x;
uint simd_gid = simdgroup_index_in_threadgroup; // 0..31 (which SG = block offset)
uint simd_lid = thread_index_in_simdgroup; // 0..31 (thread within SG)
uint b_idx = threadgroup_position_in_grid.z;
int head_idx = (int)tg_idx / V_SPLIT;
int split_idx = (int)tg_idx % V_SPLIT;
int v_offset = split_idx * V_PER_SPLIT;
int blocks_val = (int)N_blocks;
int h_q_val = (int)H_Q;
// Base offset into partials/maxs/sums for this head: [B, H_q, blocks, ...]
int head_base = ((int)b_idx * h_q_val + head_idx);
// Pointer setup for partials: [B, H_q, blocks, D]
// Each SG starts at a different block, reads EPT V-elements at v_offset
int p_base = head_base * blocks_val * D_DIM
+ (int)simd_gid * D_DIM
+ v_offset + (int)simd_lid * EPT;
int ms_base = head_base * blocks_val;
// ── Phase 1: Find global max across all blocks ──
// Each of 32 simd_lids processes blocks/32 = 4 blocks (strided)
float local_max = -__FLT_MAX__;
for (int b = 0; b < blocks_val / BN; ++b) {{
local_max = metal::max(local_max, maxs[ms_base + (int)simd_lid + BN * b]);
}}
float global_max = simd_max(local_max);
// ── Phase 2: Compute global sum_exp ──
float local_sum = 0.0f;
for (int b = 0; b < blocks_val / BN; ++b) {{
float factor = metal::fast::exp(maxs[ms_base + (int)simd_lid + BN * b] - global_max);
local_sum += factor * sums[ms_base + (int)simd_lid + BN * b];
}}
float global_sum = simd_sum(local_sum);
// ── Phase 3: Accumulate V-partials (block-parallel via SGs) ──
// 32 SGs process 32 blocks simultaneously, 4 iterations for 128 blocks.
// Each SG reads EPT V-elements from its assigned block.
float o[EPT] = {{0}};
for (int b = 0; b < blocks_val / BN; ++b) {{
float factor = metal::fast::exp(maxs[ms_base + (int)simd_gid] - global_max);
for (int i = 0; i < EPT; i++) {{
o[i] += factor * (float)o_partials[p_base + i];
}}
// Advance maxs by BN and partials by BN * D
ms_base += BN;
p_base += BN * D_DIM;
}}
// ── Phase 4: Shared memory transpose + reduce across SGs ──
// Each SG has accumulated partials for different blocks but same V-elements.
// Transpose so each SG's threads hold all block contributions for the same V position,
// then simd_sum reduces them.
threadgroup float tg_mem[BN * BN]; // 32×32 = 1024 floats = 4 KB
for (int i = 0; i < EPT; i++) {{
tg_mem[(int)simd_lid * BN + (int)simd_gid] = o[i];
threadgroup_barrier(mem_flags::mem_threadgroup);
o[i] = simd_sum(tg_mem[(int)simd_gid * BN + (int)simd_lid]);
threadgroup_barrier(mem_flags::mem_threadgroup);
}}
// ── Phase 5: Normalize + gate multiply + write (simd_lid == 0 only) ──
if (simd_lid == 0) {{
float inv_sum = (global_sum > 0.0f) ? (1.0f / global_sum) : 0.0f;
int out_base = (int)b_idx * h_q_val * D_DIM + head_idx * D_DIM
+ v_offset + (int)simd_gid * EPT;
for (int i = 0; i < EPT; i++) {{
float val = o[i] * inv_sum;
float gate = gate_sigmoid[out_base + i];
attn_output[out_base + i] = static_cast<bfloat16_t>(val * gate);
}}
}}
"""
_pass2_cache = {}
def _get_pass2_kernel(D, V_SPLIT):
"""Get or compile the SDPA Pass 2 + gate kernel."""
key = (D, V_SPLIT)
if key not in _pass2_cache:
_pass2_cache[key] = mx.fast.metal_kernel(
name="custom_sdpa_pass2_gate",
input_names=["o_partials", "sums", "maxs", "gate_sigmoid",
"H_Q", "N_blocks"],
output_names=["attn_output"],
source=_gen_sdpa_pass2_gate_source(D, V_SPLIT),
)
return _pass2_cache[key]
def custom_sdpa_pass2_gate(o_partials, sums, maxs, gate_sigmoid,
H_q, D, blocks=128, V_SPLIT=4, batch_size=1,
scalars=None):
"""SDPA Pass 2: reduce block partials + gate multiply.
Args:
o_partials: [B, H_q, blocks, D] bf16 — from Pass 1.
sums: [B, H_q, blocks] f32 — from Pass 1.
maxs: [B, H_q, blocks] f32 — from Pass 1.
gate_sigmoid: [B, 1, H_q*D] f32 — from Dispatch 1.
H_q: int — number of query heads.
D: int — head dimension.
blocks: int — number of blocks from Pass 1.
V_SPLIT: int — V-dimension split factor (4 for 80% M3 Ultra util).
batch_size: int
scalars: dict of pre-cached mx.array scalars (optional)
Returns:
attn_output: [B, 1, H_q*D] bf16 — final gated attention output.
"""
B = batch_size
kern = _get_pass2_kernel(D, V_SPLIT)
if scalars is not None:
h_q_arr = scalars['H_Q']
n_blocks_arr = scalars['N_blocks']
else:
h_q_arr = mx.array(H_q, dtype=mx.int32)
n_blocks_arr = mx.array(blocks, dtype=mx.int32)
# Flatten inputs for kernel
o_flat = o_partials.reshape(B * H_q * blocks * D)
s_flat = sums.reshape(B * H_q * blocks)
m_flat = maxs.reshape(B * H_q * blocks)
g_flat = gate_sigmoid.reshape(B * H_q * D)
n_tgs = H_q * V_SPLIT
BN = 32 # SGs per TG
results = kern(
inputs=[o_flat, s_flat, m_flat, g_flat, h_q_arr, n_blocks_arr],
output_shapes=[(B * H_q * D,)],
output_dtypes=[mx.bfloat16],
grid=(n_tgs * BN, BN, B),
threadgroup=(BN, BN, 1),
)
return results[0].reshape(B, 1, H_q * D)
src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_gdn_projections_8bit.py
+313
View File
@@ -0,0 +1,313 @@
"""Fused GDN projections for Qwen3.5-35B-A3B (Dispatch 2).
Single dispatch fuses 4 quantized 8-bit GEMVs + depthwise conv1d + activations:
- in_proj_qkv (8192×2048): GEMV → conv1d(4-tap) → SiLU → write bf16 + cache update
- in_proj_z (4096×2048): GEMV → SiLU → write f32
- in_proj_b (32×2048): GEMV → sigmoid → write f32 (beta for GDN kernel)
- in_proj_a (32×2048): GEMV → g=exp(-exp(A_log)*softplus(a+dt_bias)) → write f32
All 4 projection weight matrices are pre-merged into one contiguous buffer
(W_merged, S_merged, B_merged) for better memory locality and cache behavior.
Merging is done offline at patch time by _patch_gdn_proj_weights().
B/A epilogues compute g and beta in-kernel, eliminating ~8 micro-dispatches
that gated_delta_update would otherwise generate (sigmoid, exp, log, etc.).
The caller can pass g/beta directly to gated_delta_kernel.
TG-level multiplexing: tgid.y routes to different epilogues.
Each TG: 64 threads = 2 SGs of 32, produces 8 output rows (4 per SG).
Standard 8-bit affine GEMV: result = scale * Σ(x[i]*w[i]) + bias * Σ(x[i])
Grid: (32, total_tg * 2, B), TG: (32, 2, 1)
"""
import mlx.core as mx
def ceil_div(a, b):
return (a + b - 1) // b
def _gen_fused_gdn_projections_source(group_size=64):
"""Generate Metal source for fused GDN projections with merged weights.
8-bit dequantization (group_size=64):
result = scale * Σ(x[i]*w[i]) + bias * Σ(x[i])
Single merged weight buffer indexed by absolute out_row.
TG routing via tgid.y determines region (epilogue):
[0, N_QKV_TG): QKV GEMV + conv1d + SiLU + cache
[N_QKV_TG, +N_Z_TG): Z GEMV + SiLU
[+N_Z_TG, +N_B_TG): B GEMV → sigmoid → beta (f32)
[+N_B_TG, +N_A_TG): A GEMV → g=exp(-exp(A_log)*softplus(a+dt_bias)) (f32)
"""
gs = int(group_size)
sc_stride = 256 // gs # groups consumed per K-block = 4
slid_div = gs // 8 # threads per group = 8
return f"""
const int RESULTS_PER_SG = 4;
const int VALUES_PER_THREAD = 8;
const int BLOCK_SIZE = 256; // 32 * 8
const int GROUP_SIZE = {gs};
const int SC_STRIDE = {sc_stride};
const int SLID_DIV = {slid_div};
int K = K_val;
int K_groups = K / GROUP_SIZE;
// Dimension boundaries
int N_QKV = N_QKV_val;
int N_Z = N_Z_val;
int N_B = N_B_val;
int N_TOTAL = N_TOTAL_val;
// TG boundaries
int N_QKV_TG = N_QKV_TG_val;
int N_Z_TG = N_Z_TG_val;
int N_B_TG = N_B_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 b_idx = tgid.z;
int tg = tgid.y;
// ─── Determine region and absolute out_row in merged matrix ───
int out_row;
int region; // 0=QKV, 1=Z, 2=B, 3=A
if (tg < N_QKV_TG) {{
region = 0;
out_row = tg * 8 + sgid * RESULTS_PER_SG;
}} else if (tg < N_QKV_TG + N_Z_TG) {{
region = 1;
out_row = N_QKV + (tg - N_QKV_TG) * 8 + sgid * RESULTS_PER_SG;
}} else if (tg < N_QKV_TG + N_Z_TG + N_B_TG) {{
region = 2;
out_row = N_QKV + N_Z + (tg - N_QKV_TG - N_Z_TG) * 8 + sgid * RESULTS_PER_SG;
}} else {{
region = 3;
out_row = N_QKV + N_Z + N_B + (tg - N_QKV_TG - N_Z_TG - N_B_TG) * 8 + sgid * RESULTS_PER_SG;
}}
if (out_row >= N_TOTAL) return;
// ─── Single pointer into merged weight buffer ───
const device uint8_t* ws = (const device uint8_t*)W_merged + (long)out_row * K + slid * VALUES_PER_THREAD;
const device bfloat16_t* sc = (const device bfloat16_t*)S_merged + (long)out_row * K_groups + slid / SLID_DIV;
const device bfloat16_t* bi = (const device bfloat16_t*)B_merged + (long)out_row * K_groups + slid / SLID_DIV;
// ─── 8-bit GEMV K-loop (unified for all regions) ───
float result[4] = {{0, 0, 0, 0}};
int x_base = b_idx * K + slid * VALUES_PER_THREAD;
for (int k_off = 0; k_off < K; k_off += 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* w = ws + row * K;
float s_val = float(sc[row * K_groups]);
float b_val = float(bi[row * K_groups]);
float accum = 0;
for (int i = 0; i < 8; i++) {{
accum += x_thread[i] * float(w[i]);
}}
result[row] += s_val * accum + xsum * b_val;
}}
ws += BLOCK_SIZE;
sc += SC_STRIDE;
bi += SC_STRIDE;
x_base += BLOCK_SIZE;
}}
// ─── Reduction ───
for (int row = 0; row < RESULTS_PER_SG; row++) {{
result[row] = simd_sum(result[row]);
}}
// ─── Region-specific epilogues ───
// After simd_sum, all 32 threads have result[0..3].
// Threads 0-3 each handle one output row.
if (region == 0) {{
// ═══ QKV: conv1d(4-tap) + SiLU + cache update ═══
int c = out_row + (int)slid; // channel index (= absolute row for QKV)
if (slid < (uint)RESULTS_PER_SG && c < N_QKV) {{
float qkv_val = result[slid];
int conv_dim = N_QKV;
long cs_base = (long)b_idx * 3 * conv_dim;
float s0 = float(conv_state[cs_base + 0 * conv_dim + c]);
float s1 = float(conv_state[cs_base + 1 * conv_dim + c]);
float s2 = float(conv_state[cs_base + 2 * conv_dim + c]);
float conv_out = float(conv_w[c * 4 + 0]) * s0
+ float(conv_w[c * 4 + 1]) * s1
+ float(conv_w[c * 4 + 2]) * s2
+ float(conv_w[c * 4 + 3]) * qkv_val;
float silu_out = conv_out / (1.0f + metal::exp(-conv_out));
conv_state_out[cs_base + 0 * conv_dim + c] = static_cast<bfloat16_t>(s1);
conv_state_out[cs_base + 1 * conv_dim + c] = static_cast<bfloat16_t>(s2);
conv_state_out[cs_base + 2 * conv_dim + c] = static_cast<bfloat16_t>(qkv_val);
qkv_out[b_idx * conv_dim + c] = static_cast<bfloat16_t>(silu_out);
}}
}} else if (region == 1) {{
// ═══ Z: SiLU → write f32 ═══
int z_row = out_row - N_QKV + (int)slid;
if (slid < (uint)RESULTS_PER_SG && z_row < N_Z) {{
float val = result[slid];
float silu_val = val / (1.0f + metal::exp(-val));
z_silu_out[b_idx * N_Z + z_row] = silu_val;
}}
}} else if (region == 2) {{
// ═══ B: sigmoid(result) → beta (f32) ═══
int b_row = out_row - N_QKV - N_Z + (int)slid;
if (slid < (uint)RESULTS_PER_SG && b_row < N_B) {{
float val = result[slid];
float beta = 1.0f / (1.0f + metal::exp(-val));
b_out[b_idx * N_B + b_row] = beta;
}}
}} else {{
// ═══ A: g = exp(-exp(A_log) * softplus(a + dt_bias)) → f32 ═══
int a_row = out_row - N_QKV - N_Z - N_B + (int)slid;
int N_A = N_TOTAL - N_QKV - N_Z - N_B;
if (slid < (uint)RESULTS_PER_SG && a_row < N_A) {{
float a_val = result[slid];
float dt = float(dt_bias_arr[a_row]);
float x_g = a_val + dt;
// softplus(x) = log(1 + exp(x)), with x>20 shortcut for numerical stability
float sp = (x_g > 20.0f) ? x_g : metal::log(1.0f + metal::exp(x_g));
float g_val = metal::exp(-metal::exp(float(A_log_arr[a_row])) * sp);
a_out[b_idx * N_A + a_row] = g_val;
}}
}}
"""
_fused_gdn_proj_kernel = None
def _get_fused_gdn_proj_kernel():
"""Get or compile the fused GDN projections kernel."""
global _fused_gdn_proj_kernel
if _fused_gdn_proj_kernel is None:
_fused_gdn_proj_kernel = mx.fast.metal_kernel(
name="fused_gdn_projections_8bit_merged",
input_names=[
"x",
"W_merged", "S_merged", "B_merged",
"conv_state", "conv_w",
"A_log_arr", "dt_bias_arr",
"K_val",
"N_QKV_val", "N_Z_val", "N_B_val", "N_TOTAL_val",
"N_QKV_TG_val", "N_Z_TG_val", "N_B_TG_val",
],
output_names=["qkv_out", "z_silu_out", "b_out", "a_out", "conv_state_out"],
source=_gen_fused_gdn_projections_source(),
)
return _fused_gdn_proj_kernel
def fused_gdn_projections(
x,
W_merged, S_merged, B_merged,
proj_dims,
conv_state, conv_weights,
A_log, dt_bias,
batch_size=1,
):
"""Fused GDN projections: 4 GEMVs + conv1d + activations + g/beta.
Uses pre-merged contiguous weight buffers for all 4 projections.
B epilogue computes beta = sigmoid(b) in f32.
A epilogue computes g = exp(-exp(A_log) * softplus(a + dt_bias)) in f32.
Caller passes g/beta directly to gated_delta_kernel (no micro-dispatches).
Args:
x: [B, 1, K] bf16 — post-RMSNorm hidden state
W_merged: [N_TOTAL, K/4] uint32 — merged quantized weights
S_merged: [N_TOTAL, K/gs] bf16 — merged scales
B_merged: [N_TOTAL, K/gs] bf16 — merged biases
proj_dims: (N_QKV, N_Z, N_B, N_A) — per-projection output dims
conv_state: [B, 3, conv_dim] bf16 — previous 3 timesteps
conv_weights: [conv_dim, 4, 1] or [conv_dim, 4] bf16 — depthwise conv filters
A_log: [Hv] f32 — GDN decay log-parameter
dt_bias: [Hv] f32 — GDN time constant bias
batch_size: int
Returns:
qkv_conv_silu: [B, 1, N_QKV] bf16 — post-conv, post-SiLU
z_silu: [B, 1, N_Z] f32 — post-SiLU
beta: [B, 1, N_B] f32 — sigmoid(b), ready for GDN kernel
g: [B, 1, N_A] f32 — gating, ready for GDN kernel
conv_state_out: [B, 3, N_QKV] bf16
"""
B = batch_size
kern = _get_fused_gdn_proj_kernel()
N_QKV, N_Z, N_B, N_A = proj_dims
N_TOTAL = N_QKV + N_Z + N_B + N_A
K = x.shape[-1]
# TG counts (8 rows per TG)
N_QKV_TG = ceil_div(N_QKV, 8)
N_Z_TG = ceil_div(N_Z, 8)
N_B_TG = ceil_div(N_B, 8)
N_A_TG = ceil_div(N_A, 8)
total_tg = N_QKV_TG + N_Z_TG + N_B_TG + N_A_TG
# Flatten conv weights to [conv_dim, 4] if needed
conv_w_flat = conv_weights.reshape(-1, 4) if conv_weights.ndim == 3 else conv_weights
# Flatten x to [B, K]
x_flat = x.reshape(B, K)
results = kern(
inputs=[
x_flat,
W_merged, S_merged, B_merged,
conv_state, conv_w_flat,
A_log, dt_bias,
mx.array(K, dtype=mx.int32),
mx.array(N_QKV, dtype=mx.int32),
mx.array(N_Z, dtype=mx.int32),
mx.array(N_B, dtype=mx.int32),
mx.array(N_TOTAL, dtype=mx.int32),
mx.array(N_QKV_TG, dtype=mx.int32),
mx.array(N_Z_TG, dtype=mx.int32),
mx.array(N_B_TG, dtype=mx.int32),
],
output_shapes=[
(B * N_QKV,), # qkv_out
(B * N_Z,), # z_silu_out
(B * N_B,), # beta_out (f32)
(B * N_A,), # g_out (f32)
(B * 3 * N_QKV,), # conv_state_out
],
output_dtypes=[mx.bfloat16, mx.float32, mx.float32, mx.float32, mx.bfloat16],
grid=(32, total_tg * 2, B),
threadgroup=(32, 2, 1),
)
qkv_out = results[0].reshape(B, 1, N_QKV)
z_silu = results[1].reshape(B, 1, N_Z)
beta = results[2].reshape(B, 1, N_B)
g = results[3].reshape(B, 1, N_A)
conv_state_out = results[4].reshape(B, 3, N_QKV)
return qkv_out, z_silu, beta, g, conv_state_out
src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_gqa_projections_8bit.py
+272
View File
@@ -0,0 +1,272 @@
"""Fused GQA projections for Qwen3.5-35B-A3B (Dispatch 1).
Single dispatch fuses 4 quantized 8-bit GEMVs with region-specific epilogues:
- q_proj queries (4096×2048): GEMV → raw bf16 write
- q_proj gate (4096×2048): GEMV → sigmoid → f32 write
- k_proj (512×2048): GEMV → raw bf16 write
- v_proj (512×2048): GEMV → raw bf16 write
All 4 projection weight matrices are pre-merged into one contiguous buffer
(W_merged, S_merged, B_merged) for better memory locality.
Merging is done offline at patch time by _patch_gqa_proj_weights().
Gate sigmoid is computed in f32 directly from the GEMV accumulator,
avoiding bf16 round-trip and eliminating a separate sigmoid dispatch.
TG-level multiplexing: tgid.y routes to different epilogues.
Each TG: 64 threads = 2 SGs of 32, produces 8 output rows (4 per SG).
Standard 8-bit affine GEMV: result = scale * Σ(x[i]*w[i]) + bias * Σ(x[i])
Grid: (32, total_tg * 2, B), TG: (32, 2, 1)
"""
import mlx.core as mx
def ceil_div(a, b):
return (a + b - 1) // b
def _gen_fused_gqa_projections_source(group_size=64):
"""Generate Metal source for fused GQA projections with merged weights.
8-bit dequantization (group_size=64):
result = scale * Σ(x[i]*w[i]) + bias * Σ(x[i])
Single merged weight buffer indexed by absolute out_row.
TG routing via tgid.y determines region (epilogue):
[0, N_Q_TG): Queries GEMV → raw bf16
[N_Q_TG, +N_GATE_TG): Gate GEMV → sigmoid → f32
[+N_GATE_TG, +N_K_TG): Keys GEMV → raw bf16
[+N_K_TG, +N_V_TG): Values GEMV → raw bf16
"""
gs = int(group_size)
sc_stride = 256 // gs # groups consumed per K-block = 4
slid_div = gs // 8 # threads per group = 8
return f"""
const int RESULTS_PER_SG = 4;
const int VALUES_PER_THREAD = 8;
const int BLOCK_SIZE = 256; // 32 * 8
const int GROUP_SIZE = {gs};
const int SC_STRIDE = {sc_stride};
const int SLID_DIV = {slid_div};
int K = K_val;
int K_groups = K / GROUP_SIZE;
// Dimension boundaries
int N_Q = N_Q_val;
int N_GATE = N_GATE_val;
int N_K = N_K_val;
int N_TOTAL = N_TOTAL_val;
// TG boundaries
int N_Q_TG = N_Q_TG_val;
int N_GATE_TG = N_GATE_TG_val;
int N_K_TG = N_K_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 b_idx = tgid.z;
int tg = tgid.y;
// ─── Determine region and absolute out_row in merged matrix ───
int out_row;
int region; // 0=Q, 1=Gate, 2=K, 3=V
if (tg < N_Q_TG) {{
region = 0;
out_row = tg * 8 + sgid * RESULTS_PER_SG;
}} else if (tg < N_Q_TG + N_GATE_TG) {{
region = 1;
out_row = N_Q + (tg - N_Q_TG) * 8 + sgid * RESULTS_PER_SG;
}} else if (tg < N_Q_TG + N_GATE_TG + N_K_TG) {{
region = 2;
out_row = N_Q + N_GATE + (tg - N_Q_TG - N_GATE_TG) * 8 + sgid * RESULTS_PER_SG;
}} else {{
region = 3;
out_row = N_Q + N_GATE + N_K + (tg - N_Q_TG - N_GATE_TG - N_K_TG) * 8 + sgid * RESULTS_PER_SG;
}}
if (out_row >= N_TOTAL) return;
// ─── Single pointer into merged weight buffer ───
const device uint8_t* ws = (const device uint8_t*)W_merged + (long)out_row * K + slid * VALUES_PER_THREAD;
const device bfloat16_t* sc = (const device bfloat16_t*)S_merged + (long)out_row * K_groups + slid / SLID_DIV;
const device bfloat16_t* bi = (const device bfloat16_t*)B_merged + (long)out_row * K_groups + slid / SLID_DIV;
// ─── 8-bit GEMV K-loop (unified for all regions) ───
float result[4] = {{0, 0, 0, 0}};
int x_base = b_idx * K + slid * VALUES_PER_THREAD;
for (int k_off = 0; k_off < K; k_off += 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* w = ws + row * K;
float s_val = float(sc[row * K_groups]);
float b_val = float(bi[row * K_groups]);
float accum = 0;
for (int i = 0; i < 8; i++) {{
accum += x_thread[i] * float(w[i]);
}}
result[row] += s_val * accum + xsum * b_val;
}}
ws += BLOCK_SIZE;
sc += SC_STRIDE;
bi += SC_STRIDE;
x_base += BLOCK_SIZE;
}}
// ─── Reduction ───
for (int row = 0; row < RESULTS_PER_SG; row++) {{
result[row] = simd_sum(result[row]);
}}
// ─── Region-specific epilogues ───
// After simd_sum, all 32 threads have result[0..3].
// Threads 0-3 each handle one output row.
if (region == 0) {{
// ═══ Queries: raw bf16 write ═══
int q_row = out_row + (int)slid;
if (slid < (uint)RESULTS_PER_SG && q_row < N_Q) {{
q_out[b_idx * N_Q + q_row] = static_cast<bfloat16_t>(result[slid]);
}}
}} else if (region == 1) {{
// ═══ Gate: sigmoid → f32 ═══
int g_row = out_row - N_Q + (int)slid;
if (slid < (uint)RESULTS_PER_SG && g_row < N_GATE) {{
float val = result[slid];
float sig = 1.0f / (1.0f + metal::exp(-val));
gate_out[b_idx * N_GATE + g_row] = sig;
}}
}} else if (region == 2) {{
// ═══ Keys: raw bf16 write ═══
int k_row = out_row - N_Q - N_GATE + (int)slid;
if (slid < (uint)RESULTS_PER_SG && k_row < N_K) {{
k_out[b_idx * N_K + k_row] = static_cast<bfloat16_t>(result[slid]);
}}
}} else {{
// ═══ Values: raw bf16 write ═══
int v_row = out_row - N_Q - N_GATE - N_K + (int)slid;
int N_V = N_TOTAL - N_Q - N_GATE - N_K;
if (slid < (uint)RESULTS_PER_SG && v_row < N_V) {{
v_out[b_idx * N_V + v_row] = static_cast<bfloat16_t>(result[slid]);
}}
}}
"""
_fused_gqa_proj_kernel = None
def _get_fused_gqa_proj_kernel():
"""Get or compile the fused GQA projections kernel."""
global _fused_gqa_proj_kernel
if _fused_gqa_proj_kernel is None:
_fused_gqa_proj_kernel = mx.fast.metal_kernel(
name="fused_gqa_projections_8bit_merged",
input_names=[
"x",
"W_merged", "S_merged", "B_merged",
"K_val",
"N_Q_val", "N_GATE_val", "N_K_val", "N_TOTAL_val",
"N_Q_TG_val", "N_GATE_TG_val", "N_K_TG_val",
],
output_names=["q_out", "gate_out", "k_out", "v_out"],
source=_gen_fused_gqa_projections_source(),
)
return _fused_gqa_proj_kernel
def fused_gqa_projections(
x,
W_merged, S_merged, B_merged,
proj_dims,
batch_size=1,
scalars=None, total_tg=None,
):
"""Fused GQA projections: 4 GEMVs with region-specific epilogues.
Uses pre-merged contiguous weight buffers for all 4 projections.
Gate epilogue computes sigmoid in f32 directly from GEMV accumulator.
Args:
x: [B, 1, K] bf16 — post-RMSNorm hidden state
W_merged: [N_TOTAL, K/4] uint32 — merged quantized weights
S_merged: [N_TOTAL, K/gs] bf16 — merged scales
B_merged: [N_TOTAL, K/gs] bf16 — merged biases
proj_dims: (N_Q, N_GATE, N_K, N_V) — per-projection output dims
batch_size: int
scalars: dict of pre-cached mx.array scalars (optional, avoids per-call creation)
total_tg: pre-computed total TG count (optional)
Returns:
queries: [B, 1, N_Q] bf16
gate_sigmoid: [B, 1, N_GATE] f32 — sigmoid(gate), ready for post-SDPA multiply
keys: [B, 1, N_K] bf16
values: [B, 1, N_V] bf16
"""
B = batch_size
kern = _get_fused_gqa_proj_kernel()
N_Q, N_GATE, N_K, N_V = proj_dims
K = x.shape[-1]
# Flatten x to [B, K]
x_flat = x.reshape(B, K)
if scalars is not None:
s = scalars
inputs = [x_flat, W_merged, S_merged, B_merged,
s['K'], s['N_Q'], s['N_GATE'], s['N_K'], s['N_TOTAL'],
s['N_Q_TG'], s['N_GATE_TG'], s['N_K_TG']]
else:
N_TOTAL = N_Q + N_GATE + N_K + N_V
total_tg = ceil_div(N_Q, 8) + ceil_div(N_GATE, 8) + ceil_div(N_K, 8) + ceil_div(N_V, 8)
inputs = [x_flat, W_merged, S_merged, B_merged,
mx.array(K, dtype=mx.int32),
mx.array(N_Q, dtype=mx.int32),
mx.array(N_GATE, dtype=mx.int32),
mx.array(N_K, dtype=mx.int32),
mx.array(N_TOTAL, dtype=mx.int32),
mx.array(ceil_div(N_Q, 8), dtype=mx.int32),
mx.array(ceil_div(N_GATE, 8), dtype=mx.int32),
mx.array(ceil_div(N_K, 8), dtype=mx.int32)]
if total_tg is None:
total_tg = ceil_div(N_Q, 8) + ceil_div(N_GATE, 8) + ceil_div(N_K, 8) + ceil_div(N_V, 8)
results = kern(
inputs=inputs,
output_shapes=[
(B * N_Q,), # q_out
(B * N_GATE,), # gate_out (f32)
(B * N_K,), # k_out
(B * N_V,), # v_out
],
output_dtypes=[mx.bfloat16, mx.float32, mx.bfloat16, mx.bfloat16],
grid=(32, total_tg * 2, B),
threadgroup=(32, 2, 1),
)
queries = results[0].reshape(B, 1, N_Q)
gate_sigmoid = results[1].reshape(B, 1, N_GATE)
keys = results[2].reshape(B, 1, N_K)
values = results[3].reshape(B, 1, N_V)
return queries, gate_sigmoid, keys, values
src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_qk_norm_rope.py
+201
View File
@@ -0,0 +1,201 @@
"""Fused Q/K RMSNorm + RoPE for GQA attention (Dispatch 2).
Performs per-head RMSNorm with learned weight (head_dim=256) then applies
RoPE on the first 64 dims (partial_rotary_factor=0.25) using non-traditional
pairing: element p pairs with p+32 (not p+1).
RMSNorm with weight:
inv_rms = rsqrt(mean(x^2) + eps) = rsqrt(sum(x^2)/D + eps)
out[i] = x[i] * inv_rms * weight[i]
RoPE (non-traditional, partial):
Pairs (p, p+32) for p in {0, ..., 31}:
x'[p] = x[p]*cos(m*f_p) - x[p+32]*sin(m*f_p)
x'[p+32] = x[p]*sin(m*f_p) + x[p+32]*cos(m*f_p)
where f_p = theta^(-p/32), m = cache position
Only first 64 of 256 dims rotated; remaining 192 unchanged.
Grid: ((H_q + H_kv) * 32, 1, B)
Each TG = 32 threads = 1 SG, handles one 256-dim head.
D=256 = 32 threads x 8 elements -> exactly 1 SG, no cross-SG reduction.
"""
import mlx.core as mx
def _gen_fused_qk_norm_rope_source(H_q=16, H_kv=2, D=256, rope_dims=64):
"""Generate Metal source for fused Q/K RMSNorm + RoPE.
Input: queries [B, H_q*D] bf16, keys [B, H_kv*D] bf16
Output: q_out [B, H_q*D] bf16, k_out [B, H_kv*D] bf16
TG assignment:
tgid.x 0..H_q-1: Q heads
tgid.x H_q..H_q+H_kv-1: K heads
Thread layout (32 threads, N_READS=8):
Thread t handles elements [8t, 8t+7]
RoPE thread assignment (non-traditional pairing):
Threads 0-3: first half of pair (elements 0-31)
Threads 4-7: second half of pair (elements 32-63)
Thread t paired with t^4 via simd_shuffle
Threads 8-31: elements 64-255 (unrotated, skip RoPE)
"""
N_READS = D // 32
ROPE_HALF = rope_dims // 2 # 32 = number of rotation pairs
ROPE_THREADS = rope_dims // N_READS # 8 = threads touching rotated dims
FIRST_HALF = ROPE_THREADS // 2 # 4 = threads in first half of pairs
PARTNER_XOR = FIRST_HALF # 4 = XOR to find partner thread
return f"""
// ── Constants ──
const int H_Q = {H_q};
const int H_KV = {H_kv};
const int D_DIM = {D};
const int N_READS = {N_READS};
const float EPS = 1e-6f;
uint head_idx = threadgroup_position_in_grid.x;
uint slid = thread_index_in_simdgroup;
uint b_idx = thread_position_in_grid.z;
bool is_q = (head_idx < (uint)H_Q);
int head_local = is_q ? (int)head_idx : ((int)head_idx - H_Q);
// ── Phase 1: Load N_READS elements + partial sum of squares ──
int in_base = is_q
? (int)(b_idx * H_Q * D_DIM + (int)head_idx * D_DIM)
: (int)(b_idx * H_KV * D_DIM + head_local * D_DIM);
int elem_base = (int)slid * N_READS;
float vals[{N_READS}];
float partial_sq = 0.0f;
for (int i = 0; i < N_READS; i++) {{
float xi;
if (is_q)
xi = (float)queries[in_base + elem_base + i];
else
xi = (float)keys[in_base + elem_base + i];
vals[i] = xi;
partial_sq += xi * xi;
}}
// ── Phase 2: RMSNorm reduction (32 threads -> full sum of {D} elements) ──
float sum_sq = simd_sum(partial_sq);
float inv_rms = metal::precise::rsqrt(sum_sq / (float)D_DIM + EPS);
// ── Phase 3: Normalize with learned weight ──
for (int i = 0; i < N_READS; i++) {{
float w;
if (is_q)
w = (float)q_norm_w[elem_base + i];
else
w = (float)k_norm_w[elem_base + i];
vals[i] = vals[i] * inv_rms * w;
}}
// ── Phase 4: RoPE on first {rope_dims} dims (threads 0..{ROPE_THREADS - 1}) ──
if (slid < {ROPE_THREADS}u) {{
ushort partner = (ushort)(slid ^ {PARTNER_XOR}u);
int cos_base = (int)(slid & {FIRST_HALF - 1}u) * N_READS;
// Compute cos/sin on-device from precomputed inv_freq
// angle = position * inv_freq[d], where inv_freq[d] = theta^(-d/{ROPE_HALF})
float cos_arr[{N_READS}], sin_arr[{N_READS}];
for (int i = 0; i < N_READS; i++) {{
float angle = (float)position[0] * inv_freq[cos_base + i];
cos_arr[i] = metal::fast::cos(angle);
sin_arr[i] = metal::fast::sin(angle);
}}
// Exchange normalized values with partner thread via simd_shuffle
float partner_vals[{N_READS}];
for (int i = 0; i < N_READS; i++)
partner_vals[i] = simd_shuffle(vals[i], partner);
if (slid < {FIRST_HALF}u) {{
// First half of pair: x'[p] = x[p]*cos - x[p+{ROPE_HALF}]*sin
for (int i = 0; i < N_READS; i++)
vals[i] = vals[i] * cos_arr[i] - partner_vals[i] * sin_arr[i];
}} else {{
// Second half: x'[p+{ROPE_HALF}] = x_first*sin + x[p+{ROPE_HALF}]*cos
for (int i = 0; i < N_READS; i++)
vals[i] = partner_vals[i] * sin_arr[i] + vals[i] * cos_arr[i];
}}
}}
// ── Phase 5: Write output ──
int out_base = is_q
? (int)(b_idx * H_Q * D_DIM + (int)head_idx * D_DIM)
: (int)(b_idx * H_KV * D_DIM + head_local * D_DIM);
for (int i = 0; i < N_READS; i++) {{
if (is_q)
q_out[out_base + elem_base + i] = static_cast<bfloat16_t>(vals[i]);
else
k_out[out_base + elem_base + i] = static_cast<bfloat16_t>(vals[i]);
}}
"""
_kernel_cache = {}
def _get_kernel(H_q, H_kv, D, rope_dims):
"""Get or compile the fused Q/K RMSNorm + RoPE kernel."""
key = (H_q, H_kv, D, rope_dims)
if key not in _kernel_cache:
_kernel_cache[key] = mx.fast.metal_kernel(
name="fused_qk_norm_rope",
input_names=["queries", "keys", "q_norm_w", "k_norm_w",
"inv_freq", "position"],
output_names=["q_out", "k_out"],
source=_gen_fused_qk_norm_rope_source(H_q, H_kv, D, rope_dims),
)
return _kernel_cache[key]
def fused_qk_norm_rope(queries, keys, q_norm_weight, k_norm_weight,
inv_freq, cache_offset, H_q, H_kv, D,
batch_size=1):
"""Fused Q/K per-head RMSNorm + RoPE for GQA attention.
Args:
queries: [B, 1, H_q*D] bf16 — raw queries from Dispatch 1.
keys: [B, 1, H_kv*D] bf16 — raw keys from Dispatch 1.
q_norm_weight: [D] bf16 — RMSNorm learned weight for queries.
k_norm_weight: [D] bf16 — RMSNorm learned weight for keys.
inv_freq: [rope_dims/2] f32 — precomputed theta^(-d/half_dims).
cache_offset: int — sequence position for RoPE angles.
H_q: int — number of query heads.
H_kv: int — number of key/value heads.
D: int — head dimension.
batch_size: int — batch size.
Returns:
q_normed_roped: [B, H_q, 1, D] bf16 — ready for SDPA.
k_normed_roped: [B, H_kv, 1, D] bf16 — ready for cache update.
"""
B = batch_size
rope_dims = inv_freq.shape[0] * 2
kern = _get_kernel(H_q, H_kv, D, rope_dims)
q_flat = queries.reshape(B, H_q * D)
k_flat = keys.reshape(B, H_kv * D)
pos = mx.array(cache_offset, dtype=mx.int32)
n_heads = H_q + H_kv
results = kern(
inputs=[q_flat, k_flat, q_norm_weight, k_norm_weight, inv_freq, pos],
output_shapes=[(B * H_q * D,), (B * H_kv * D,)],
output_dtypes=[mx.bfloat16, mx.bfloat16],
grid=(n_heads * 32, 1, B),
threadgroup=(32, 1, 1),
)
q_out = results[0].reshape(B, H_q, 1, D)
k_out = results[1].reshape(B, H_kv, 1, D)
return q_out, k_out
src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_qk_rmsnorm.py
+128
View File
@@ -0,0 +1,128 @@
"""Fused Q/K per-head L2-norm for GDN attention (Dispatch 3).
Performs per-head L2 normalization on q and k vectors with different scaling.
Matches vLLM and latest mlx-lm (qwen3_5.py) which use rsqrt(sum(x²) + eps),
NOT rms_norm which uses rsqrt(mean(x²) + eps).
From qwen3_5.py (updated to match vLLM):
inv_scale = Dk^(-0.5) = 128^(-0.5)
q = inv_scale * q * rsqrt(sum(q²) + 1e-6) → L2-normalize then scale by 1/√Dk
k = k * rsqrt(sum(k²) + 1e-6) → L2-normalize only (no extra scale)
Grid: (32 heads × 32 threads, 1, B).
Each TG = 32 threads = 1 SG, handles one 128-dim head.
Dk=128 = 32 threads × 4 elements → exactly 1 SG, no cross-SG reduction.
"""
import mlx.core as mx
def _gen_fused_qk_rmsnorm_source():
"""Generate Metal source for fused Q/K per-head L2-norm.
Input: qkv [B, 8192] bf16 (flattened from [B, 1, 8192])
- [0, 2048): q = 16 heads × 128
- [2048, 4096): k = 16 heads × 128
- [4096, 8192): v (untouched)
Output: qk_out [B, 4096] bf16
- [0, 2048): q L2-normalized then scaled by 1/√Dk
- [2048, 4096): k L2-normalized (no extra scale)
Grid: (32 * 32, 1, B), TG: (32, 1, 1)
tgid.x 0..15: q heads → scale = 1/√128
tgid.x 16..31: k heads → scale = 1.0
tgid.z: batch index
"""
return """
const int N_READS = 4;
const int DK = 128;
const int HK = 16;
const float EPS = 1e-6f;
const float Q_SCALE = rsqrt(128.0f); // inv_scale = Dk^(-0.5)
const float K_SCALE = 1.0f; // no extra scale for k
uint head_idx = threadgroup_position_in_grid.x;
uint slid = thread_index_in_simdgroup;
uint b_idx = thread_position_in_grid.z;
bool is_q = (head_idx < (uint)HK);
// Input offset: q heads at [0, 2048), k heads at [2048, 4096)
int in_base = is_q
? (b_idx * 8192 + head_idx * DK)
: (b_idx * 8192 + 2048 + (head_idx - HK) * DK);
// Output offset: q at [0, 2048), k at [2048, 4096)
int out_base = b_idx * 4096 + head_idx * DK;
// ── Phase 1: Load 4 elements + sum of squares ──
float vals[4];
float partial_sq = 0.0f;
int elem_base = slid * N_READS;
for (int i = 0; i < N_READS; i++) {
float xi = float(qkv[in_base + elem_base + i]);
vals[i] = xi;
partial_sq += xi * xi;
}
// ── Phase 2: simd reduction (32 threads → full sum of 128 elements) ──
float sum_sq = simd_sum(partial_sq);
// ── Phase 3: compute L2 inv-norm (NOT rms_norm — no /Dk) ──
float inv_rms = metal::precise::rsqrt(sum_sq + EPS);
// ── Phase 4: scale and write ──
float scale = is_q ? Q_SCALE : K_SCALE;
float combined = inv_rms * scale;
for (int i = 0; i < N_READS; i++) {
qk_out[out_base + elem_base + i] = static_cast<bfloat16_t>(vals[i] * combined);
}
"""
_fused_qk_rmsnorm_kernel = None
def _get_fused_qk_rmsnorm_kernel():
"""Get or compile the fused Q/K RMSNorm kernel."""
global _fused_qk_rmsnorm_kernel
if _fused_qk_rmsnorm_kernel is None:
_fused_qk_rmsnorm_kernel = mx.fast.metal_kernel(
name="fused_qk_rmsnorm",
input_names=["qkv"],
output_names=["qk_out"],
source=_gen_fused_qk_rmsnorm_source(),
)
return _fused_qk_rmsnorm_kernel
def fused_qk_rmsnorm(qkv_conv_silu, batch_size=1):
"""Fused Q/K per-head RMSNorm for GDN attention.
Args:
qkv_conv_silu: [B, 1, 8192] bf16 — post-conv, post-SiLU output from Dispatch 2.
First 2048 = q (16 heads × 128), next 2048 = k, last 4096 = v.
batch_size: int — batch dimension.
Returns:
qk_normed: [B, 1, 4096] bf16 — normalized q (first 2048) and k (next 2048).
v is NOT copied; Dispatch 4 reads v directly from qkv_conv_silu[:, :, 4096:].
"""
B = batch_size
kern = _get_fused_qk_rmsnorm_kernel()
# Flatten to [B, 8192] for kernel
qkv_flat = qkv_conv_silu.reshape(B, 8192)
n_heads = 32 # 16 q + 16 k
results = kern(
inputs=[qkv_flat],
output_shapes=[(B * 4096,)],
output_dtypes=[mx.bfloat16],
grid=(n_heads * 32, 1, B),
threadgroup=(32, 1, 1),
)
return results[0].reshape(B, 1, 4096)
src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_rms_norm_gated.py
+117
View File
@@ -0,0 +1,117 @@
"""Fused RMSNormGated for GDN attention (Dispatch 5).
Fuses RMSNorm(out, weight) × z_silu into one kernel.
SiLU on z was already applied in Dispatch 2, so z_silu arrives as f32.
From qwen3_next.py (Qwen3NextRMSNormGated):
x = rms_norm(hidden_states, weight, eps) # weight: [Dv=128]
gate = silu(z.float()) # already done in Dispatch 2
return (gate * x).to(hidden_states.dtype)
Grid: (32 heads × 32 threads, 1, B).
Each TG = 32 threads = 1 SG, handles one 128-dim head.
Dv=128 = 32 threads × 4 elements → exactly 1 SG.
"""
import mlx.core as mx
def _gen_fused_rms_norm_gated_source():
"""Generate Metal source for fused RMSNormGated.
Inputs:
gdn_out: [B, Hv*Dv] bf16 — GDN output, flattened (Hv=32, Dv=128)
z_silu: [B, Hv*Dv] f32 — post-SiLU z from Dispatch 2
weight: [Dv] f32 — RMSNormGated learned weight (128 elements)
Output:
out: [B, Hv*Dv] bf16 — result = z_silu * rms_norm(gdn_out, weight)
Grid: (32 * 32, 1, B), TG: (32, 1, 1)
tgid.x: head index (0..31)
tgid.z: batch index
"""
return """
const int N_READS = 4;
const int DV = 128;
const int HV = 32;
const float EPS = 1e-6f;
uint head_idx = threadgroup_position_in_grid.x;
uint slid = thread_index_in_simdgroup;
uint b_idx = thread_position_in_grid.z;
int base = b_idx * HV * DV + head_idx * DV;
int elem_base = slid * N_READS;
// ── Phase 1: Load gdn_out elements + sum of squares ──
float gdn_vals[4];
float partial_sq = 0.0f;
for (int i = 0; i < N_READS; i++) {
float xi = float(gdn_out[base + elem_base + i]);
gdn_vals[i] = xi;
partial_sq += xi * xi;
}
// ── Phase 2: simd reduction (32 threads → full sum of 128 elements) ──
float sum_sq = simd_sum(partial_sq);
// ── Phase 3: compute inv_rms ──
float inv_rms = metal::precise::rsqrt(sum_sq / float(DV) + EPS);
// ── Phase 4: RMSNorm × z_silu, write bf16 ──
for (int i = 0; i < N_READS; i++) {
int idx = elem_base + i;
float w = float(weight[idx]); // learned weight[Dv]
float normed = gdn_vals[i] * inv_rms * w; // RMSNorm
float z_val = z_silu[base + idx]; // already f32, post-SiLU
out[base + idx] = static_cast<bfloat16_t>(z_val * normed);
}
"""
_fused_rms_norm_gated_kernel = None
def _get_fused_rms_norm_gated_kernel():
"""Get or compile the fused RMSNormGated kernel."""
global _fused_rms_norm_gated_kernel
if _fused_rms_norm_gated_kernel is None:
_fused_rms_norm_gated_kernel = mx.fast.metal_kernel(
name="fused_rms_norm_gated",
input_names=["gdn_out", "z_silu", "weight"],
output_names=["out"],
source=_gen_fused_rms_norm_gated_source(),
)
return _fused_rms_norm_gated_kernel
def fused_rms_norm_gated(gdn_out, z_silu, weight, batch_size=1):
"""Fused RMSNormGated: RMSNorm(out, weight) × z_silu.
Args:
gdn_out: [B, 1, Hv, Dv] bf16 — GDN recurrence output (Hv=32, Dv=128).
z_silu: [B, 1, 4096] f32 — post-SiLU z from Dispatch 2.
weight: [128] f32 — RMSNormGated learned weight (Dv elements).
batch_size: int.
Returns:
out: [B, 1, 4096] bf16 — ready for out_proj in Dispatch 6.
"""
B = batch_size
kern = _get_fused_rms_norm_gated_kernel()
# Flatten to [B, 4096]
gdn_flat = gdn_out.reshape(B, 4096)
z_flat = z_silu.reshape(B, 4096)
n_heads = 32 # Hv
results = kern(
inputs=[gdn_flat, z_flat, weight],
output_shapes=[(B * 4096,)],
output_dtypes=[mx.bfloat16],
grid=(n_heads * 32, 1, B),
threadgroup=(32, 1, 1),
)
return results[0].reshape(B, 1, 4096)
src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/gdn_step_precomputed.py
+177
View File
@@ -0,0 +1,177 @@
"""GDN recurrence with pre-computed g and beta (Dispatch 4).
Modified version of gated_delta_step from mlx-lm-fork/mlx_lm/models/gated_delta.py.
Instead of computing g = exp(-exp(A_log) * softplus(a + dt_bias)) and beta = sigmoid(b)
inside the kernel, accepts them as pre-computed f32 inputs from Dispatch 2.
Non-vectorized only (Qwen3.5-35B-A3B uses scalar gating per head).
Grid: (32, Dv, B*Hv) = (32, 128, B*32), TG: (32, 4, 1)
"""
from typing import Optional, Tuple
import mlx.core as mx
def _make_gdn_precomputed_kernel(has_mask=False):
"""Build the GDN kernel with pre-computed g and beta."""
if not mx.metal.is_available():
return None
mask_source = "mask[b_idx * T + t]" if has_mask else "true"
source = f"""
auto n = thread_position_in_grid.z;
auto b_idx = n / Hv;
auto hv_idx = n % Hv;
auto hk_idx = hv_idx / (Hv / Hk);
constexpr int n_per_t = Dk / 32;
// q, k: [B, T, Hk, Dk]
auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk;
auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
// v, y: [B, T, Hv, Dv]
auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
y += b_idx * T * Hv * Dv + hv_idx * Dv;
auto dk_idx = thread_position_in_threadgroup.x;
auto dv_idx = thread_position_in_grid.y;
// state_in, state_out: [B, Hv, Dv, Dk]
auto i_state = state_in + (n * Dv + dv_idx) * Dk;
auto o_state = state_out + (n * Dv + dv_idx) * Dk;
float state[n_per_t];
for (int i = 0; i < n_per_t; ++i) {{
auto s_idx = n_per_t * dk_idx + i;
state[i] = static_cast<float>(i_state[s_idx]);
}}
// g: [B, T, Hv] f32 — pre-computed decay gate
auto g_ = g + b_idx * T * Hv;
// beta: [B, T, Hv] f32 — pre-computed sigmoid(b)
auto beta_ = beta + b_idx * T * Hv;
for (int t = 0; t < T; ++t) {{
if ({mask_source}) {{
// Pre-computed g and beta (no softplus/exp/sigmoid needed)
float g_val = g_[hv_idx];
float beta_val = beta_[hv_idx];
float kv_mem = 0.0f;
for (int i = 0; i < n_per_t; ++i) {{
auto s_idx = n_per_t * dk_idx + i;
state[i] = state[i] * g_val;
kv_mem += state[i] * k_[s_idx];
}}
kv_mem = simd_sum(kv_mem);
auto delta = (v_[dv_idx] - kv_mem) * beta_val;
float out = 0.0f;
for (int i = 0; i < n_per_t; ++i) {{
auto s_idx = n_per_t * dk_idx + i;
state[i] = state[i] + k_[s_idx] * delta;
out += state[i] * q_[s_idx];
}}
out = simd_sum(out);
if (thread_index_in_simdgroup == 0) {{
y[dv_idx] = static_cast<InT>(out);
}}
}}
// Increment data pointers to next time step
q_ += Hk * Dk;
k_ += Hk * Dk;
v_ += Hv * Dv;
y += Hv * Dv;
g_ += Hv;
beta_ += Hv;
}}
for (int i = 0; i < n_per_t; ++i) {{
auto s_idx = n_per_t * dk_idx + i;
o_state[s_idx] = static_cast<InT>(state[i]);
}}
"""
inputs = ["q", "k", "v", "g", "beta", "state_in", "T"]
if has_mask:
inputs.append("mask")
suffix = "_precomputed"
if has_mask:
suffix += "_mask"
return mx.fast.metal_kernel(
name=f"gated_delta_step{suffix}",
input_names=inputs,
output_names=["y", "state_out"],
source=source,
)
_gdn_precomputed_kernel = None
_gdn_precomputed_kernel_masked = None
def _get_gdn_precomputed_kernel(has_mask=False):
"""Get or compile the pre-computed GDN kernel."""
global _gdn_precomputed_kernel, _gdn_precomputed_kernel_masked
if has_mask:
if _gdn_precomputed_kernel_masked is None:
_gdn_precomputed_kernel_masked = _make_gdn_precomputed_kernel(has_mask=True)
return _gdn_precomputed_kernel_masked
else:
if _gdn_precomputed_kernel is None:
_gdn_precomputed_kernel = _make_gdn_precomputed_kernel(has_mask=False)
return _gdn_precomputed_kernel
def gated_delta_update_precomputed(
q: mx.array,
k: mx.array,
v: mx.array,
g: mx.array,
beta: mx.array,
state: mx.array,
mask: Optional[mx.array] = None,
) -> Tuple[mx.array, mx.array]:
"""GDN recurrence with pre-computed g and beta.
Args:
q: [B, T, Hk, Dk] bf16 — normalized q from Dispatch 3
k: [B, T, Hk, Dk] bf16 — normalized k from Dispatch 3
v: [B, T, Hv, Dv] bf16 — v from Dispatch 2 (qkv_conv_silu[:, :, 4096:])
g: [B, T, Hv] f32 — pre-computed decay gate from Dispatch 2
beta: [B, T, Hv] f32 — pre-computed sigmoid(b) from Dispatch 2
state: [B, Hv, Dv, Dk] bf16 — recurrent state from cache
mask: [B, T] optional
Returns:
y: [B, T, Hv, Dv] bf16
new_state: [B, Hv, Dv, Dk] bf16
"""
B, T, Hk, Dk = k.shape
Hv, Dv = v.shape[2:]
input_type = q.dtype
kernel = _get_gdn_precomputed_kernel(has_mask=mask is not None)
inputs = [q, k, v, g, beta, state, T]
if mask is not None:
inputs.append(mask)
return kernel(
inputs=inputs,
template=[
("InT", input_type),
("Dk", Dk),
("Dv", Dv),
("Hk", Hk),
("Hv", Hv),
],
grid=(32, Dv, B * Hv),
threadgroup=(32, 4, 1),
output_shapes=[(B, T, Hv, Dv), state.shape],
output_dtypes=[input_type, input_type],
)