diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py index 143c4d40..247b85ce 100644 --- a/src/exo/worker/engines/mlx/generator/batch_generate.py +++ b/src/exo/worker/engines/mlx/generator/batch_generate.py @@ -1,4 +1,5 @@ import contextlib +import os import time from dataclasses import dataclass, field from typing import Callable, cast @@ -96,11 +97,45 @@ class ExoBatchGenerator: _active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False) def __post_init__(self) -> None: - self._mlx_gen = MlxBatchGenerator( - model=self.model, - stop_tokens=set(eos_ids_from_tokenizer(self.tokenizer)), - prefill_step_size=4096, - ) + use_speculative = os.environ.get("EXO_SPECULATIVE", "0") == "1" + stop_tokens = set(eos_ids_from_tokenizer(self.tokenizer)) + + if use_speculative: + try: + from exo.worker.engines.mlx.speculative.mtp_module import MTPPredictor + from exo.worker.engines.mlx.speculative.mtp_batch_generator import MTPBatchGenerator + + mtp_weights = self._resolve_mtp_weights() + gamma = int(os.environ.get("EXO_SPECULATIVE_GAMMA", "2")) + + if mtp_weights: + mtp = MTPPredictor(self.model, mtp_weights, quantize=False) + temp = float(os.environ.get("EXO_SPECULATIVE_TEMP", "0.7")) + alpha = float(os.environ.get("EXO_SPECULATIVE_ALPHA", "1.0")) + self._mlx_gen = MTPBatchGenerator( + model=self.model, + mtp_predictor=mtp, + gamma=gamma, + temp=temp, + alpha=alpha, + stop_tokens=stop_tokens, + prefill_step_size=4096, + ) + logger.info(f"MTP speculative decoding enabled (γ={gamma}, T={temp})") + else: + logger.warning("EXO_SPECULATIVE=1 but could not find MTP weights. Falling back.") + self._mlx_gen = MlxBatchGenerator( + model=self.model, stop_tokens=stop_tokens, prefill_step_size=4096, + ) + except Exception as e: + logger.warning(f"Failed to init MTP speculative: {e}. Falling back.") + self._mlx_gen = MlxBatchGenerator( + model=self.model, stop_tokens=stop_tokens, prefill_step_size=4096, + ) + else: + self._mlx_gen = MlxBatchGenerator( + model=self.model, stop_tokens=stop_tokens, prefill_step_size=4096, + ) self._mlx_gen._needs_topk = False # pyright: ignore[reportAttributeAccessIssue] @property @@ -111,6 +146,109 @@ class ExoBatchGenerator: or self._mlx_gen.active_batch is not None ) + def _resolve_mtp_weights(self) -> str | None: + """Find MTP weights: explicit path or auto-detect.""" + explicit_path = os.environ.get("EXO_MTP_WEIGHTS", "") + if explicit_path and os.path.exists(explicit_path): + return explicit_path + + mtp_model = os.environ.get("EXO_MTP_MODEL", "") + if not mtp_model: + try: + inner = getattr(self.model, 'model', None) or self.model.language_model.model + args = getattr(inner, 'args', None) + if args and getattr(args, 'mtp_num_hidden_layers', 0) > 0: + model_type = getattr(args, 'model_type', '') + if 'qwen3_5' in model_type: + mtp_model = "Qwen/Qwen3.5-27B" + logger.info(f"Auto-detected MTP model: {mtp_model}") + except Exception: + pass + + if not mtp_model: + return None + + try: + return self._extract_mtp_from_hf(mtp_model) + except Exception as e: + logger.warning(f"Failed to extract MTP weights from {mtp_model}: {e}") + return None + + def _extract_mtp_from_hf(self, repo_id: str) -> str: + """Download MTP tensors from HF repo and cache.""" + import hashlib + from pathlib import Path + from huggingface_hub import snapshot_download + from safetensors.torch import load_file, save_file + + cache_dir = Path.home() / ".cache" / "exo" / "mtp_weights" + cache_dir.mkdir(parents=True, exist_ok=True) + cache_key = hashlib.md5(repo_id.encode()).hexdigest()[:12] + cached_path = cache_dir / f"mtp_{cache_key}.safetensors" + + if cached_path.exists(): + logger.info(f"Using cached MTP weights: {cached_path}") + return str(cached_path) + + logger.info(f"Downloading MTP weights from {repo_id}...") + model_dir = snapshot_download(repo_id, allow_patterns=["*.safetensors", "*.json"]) + + mtp_tensors = {} + model_path = Path(model_dir) + for sf_file in sorted(model_path.glob("*.safetensors")): + tensors = load_file(str(sf_file)) + for k, v in tensors.items(): + if k.startswith("model.mtp."): + mtp_tensors[k[len("model."):]] = v + + if not mtp_tensors: + raise ValueError(f"No MTP tensors found in {repo_id}") + + save_file(mtp_tensors, str(cached_path)) + logger.info(f"Extracted {len(mtp_tensors)} MTP tensors → {cached_path}") + return str(cached_path) + + def warmup_speculative(self, model, tokenizer) -> None: + """Warm up speculative decoding kernels.""" + if not hasattr(self._mlx_gen, 'mtp'): + return + + from exo.worker.engines.mlx.speculative.mtp_module import speculative_forward, draft_tokens + from mlx_lm.models import cache as cache_mod + + logger.info("Warming up speculative decoding kernels...") + mtp = self._mlx_gen.mtp + gamma = self._mlx_gen.gamma + + warmup_prompt = tokenizer.encode("Warm up speculative decoding.") + cache = cache_mod.make_prompt_cache(model) + mtp.reset_cache() + + pre_norm, logits = speculative_forward(model, mx.array([warmup_prompt]), cache) + mx.eval(pre_norm, logits) + nt = mx.argmax(logits[0, -1], axis=-1).item() + + if pre_norm.shape[1] > 1: + _ = mtp.predict(pre_norm[:, :-1, :], mx.array([warmup_prompt[1:]])) + mx.eval(_) + + last_pn = pre_norm[:, -1:, :] + next_arr = mx.array([[nt]]) + for _ in range(3): + d_ids, _ = draft_tokens(mtp, last_pn, next_arr, gamma, 0.0) + dc = mx.concatenate([d.reshape(1, 1) for d in d_ids], axis=1) + vi = mx.concatenate([next_arr, dc], axis=1) + vpn, vl = speculative_forward(model, vi, cache, speculative=True) + an = mx.argmax(vl[0], axis=-1) + mx.eval(vpn, an) + next_arr = an[0].reshape(1, 1) + last_pn = vpn[:, 0:1, :] + for i, c in enumerate(cache): + if hasattr(c, 'base'): + cache[i] = c.base + + logger.info("Speculative warmup complete") + def submit( self, task_params: TextGenerationTaskParams, @@ -172,10 +310,17 @@ class ExoBatchGenerator: seed = task_params.seed if task_params.seed is not None else 42 mx.random.seed(seed) + # EXO_SPECULATIVE_TEMP overrides sampling temperature when set + spec_temp_override = os.environ.get("EXO_SPECULATIVE_TEMP") + if spec_temp_override is not None: + sampling_temp = float(spec_temp_override) + elif task_params.temperature is not None: + sampling_temp = task_params.temperature + else: + sampling_temp = 0.7 + sampler = make_sampler( - temp=task_params.temperature - if task_params.temperature is not None - else 0.7, + temp=sampling_temp, top_p=task_params.top_p if task_params.top_p is not None else 1.0, min_p=task_params.min_p if task_params.min_p is not None else 0.05, top_k=task_params.top_k if task_params.top_k is not None else 0, @@ -200,6 +345,22 @@ class ExoBatchGenerator: distributed_prompt_progress_callback, ) + # MTP prefill: build MTP KV cache from prompt hidden states + if hasattr(self._mlx_gen, 'mtp'): + prompt_pre_norm = self._mlx_gen._captured.get('prompt_pre_norm') + if prompt_pre_norm is not None: + mx.eval(prompt_pre_norm) + self._mlx_gen.mtp.reset_cache() + S_pre = prompt_pre_norm.shape[1] + if S_pre > 0 and len(all_prompt_tokens) > S_pre: + mtp_toks = all_prompt_tokens[1:S_pre + 1].tolist() + _ = self._mlx_gen.mtp.predict( + prompt_pre_norm, + mx.array([mtp_toks]) + ) + mx.eval(_) + logger.info(f"MTP cache prefilled ({S_pre} positions)") + # We need to clamp rotating kv caches to max size so that mlx lm's _merge_caches behaves for c in cache: if ( @@ -254,6 +415,15 @@ class ExoBatchGenerator: uid = uids[0] + # Pass request temperature to speculative cycle + if hasattr(self._mlx_gen, '_request_temp'): + env_temp = os.environ.get("EXO_SPECULATIVE_TEMP") + if env_temp is not None: + self._mlx_gen._request_temp[uid] = float(env_temp) + else: + request_temp = task_params.temperature if task_params.temperature is not None else 0.7 + self._mlx_gen._request_temp[uid] = request_temp + self._active_tasks[uid] = _EngineTask( uid=uid, task_params=task_params, @@ -337,7 +507,7 @@ class ExoBatchGenerator: logprob: float | None = None top_logprobs: list[TopLogprobItem] | None = None - if task_params.logprobs: + if task_params.logprobs and os.environ.get("EXO_DISABLE_LOGPROBS") != "1": with mx.stream(generation_stream): logprob, top_logprobs = extract_top_logprobs( logprobs=response.logprobs, diff --git a/src/exo/worker/engines/mlx/patches/__init__.py b/src/exo/worker/engines/mlx/patches/__init__.py index b05a045e..3deabd52 100644 --- a/src/exo/worker/engines/mlx/patches/__init__.py +++ b/src/exo/worker/engines/mlx/patches/__init__.py @@ -1,3 +1,10 @@ +import json +import os +from pathlib import Path + +import mlx.nn as nn +from loguru import logger + from exo.worker.engines.mlx.patches.opt_batch_gen import apply_batch_gen_patch from exo.worker.engines.mlx.patches.standard_yarn_rope import patch_yarn_rope @@ -12,3 +19,32 @@ def apply_mlx_patches() -> None: patch_yarn_rope() # patch_gdn_softplus() apply_batch_gen_patch() + + +def maybe_apply_patches(model: nn.Module, model_path: Path) -> None: + """Detect model type and apply kernel fusion patches if available.""" + fused_mode = os.environ.get("EXO_FUSED_KERNELS", "1") + if fused_mode == "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_batched_fused_patches + + logger.info("Detected Qwen3.5 MoE model, applying batched fused kernel patches") + apply_qwen35_batched_fused_patches(model) + + elif model_type == "qwen3_5": + from .qwen3_5.lpb_patch import apply_lpb_patches + + logger.info("Detected Qwen3.5 dense model, applying LpB kernel patches") + apply_lpb_patches(model, batch_size=4) diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5/__init__.py b/src/exo/worker/engines/mlx/patches/qwen3_5/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5/custom_qmv_loop_over_b.py b/src/exo/worker/engines/mlx/patches/qwen3_5/custom_qmv_loop_over_b.py new file mode 100644 index 00000000..cfdfd917 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5/custom_qmv_loop_over_b.py @@ -0,0 +1,135 @@ +#!/usr/bin/env python3 +"""Isolated loop-over-B GEMV kernel for quantized matmul. + +Extracts the loop-over-B pattern from batched_fused_gdn_projections_8bit +but without any epilogues — pure Y = X @ dequant(W)^T output. + +For comparing our GEMV approach against MLX's affine_qmv_fast on +an isolated QuantizedLinear operation (e.g., in_proj_qkv: N=8192, K=2048). + +TG: (32, 2, 1) = 64 threads = 2 SGs. +Each SG: 4 output rows. +B loop inside row loop for low register pressure (R = 4B + 5). + +Usage: + from custom_qmv_loop_over_b import custom_qmv_loop_over_b + y = custom_qmv_loop_over_b(x, w, scales, biases, M=8, N=8192, K=2048) +""" + +import mlx.core as mx + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def _gen_custom_qmv_source(M_val, N_val, K_val, group_size=64): + gs = group_size + sc_stride = 256 // gs + slid_div = gs // 8 + K_groups = K_val // gs + B = M_val # batch size = M + + return f""" + const int RESULTS_PER_SG = 4; + const int VALUES_PER_THREAD = 8; + const int BLOCK_SIZE = 256; + const int K = {K_val}; + const int N = {N_val}; + const int M = {M_val}; + const int K_groups = {K_groups}; + const int SC_STRIDE = {sc_stride}; + const int SLID_DIV = {slid_div}; + + uint3 tgid = threadgroup_position_in_grid; + uint sgid = simdgroup_index_in_threadgroup; + uint slid = thread_index_in_simdgroup; + int tg = tgid.y; + + int out_row = tg * 8 + sgid * RESULTS_PER_SG; + if (out_row >= N) return; + + // Weight pointers + const device uint8_t* ws = (const device uint8_t*)w + (long)out_row * K + slid * VALUES_PER_THREAD; + const device bfloat16_t* sc = (const device bfloat16_t*)scales + (long)out_row * K_groups + slid / SLID_DIV; + const device bfloat16_t* bi = (const device bfloat16_t*)biases + (long)out_row * K_groups + slid / SLID_DIV; + + // Result accumulators: 4 rows × B batches + float result[{4 * B}]; + for (int i = 0; i < {4 * B}; i++) result[i] = 0; + + int x_base = slid * VALUES_PER_THREAD; + + // K-loop: loop over B inside row loop + for (int k_off = 0; k_off < K; k_off += BLOCK_SIZE) {{ + + for (int row = 0; row < RESULTS_PER_SG; row++) {{ + const device uint8_t* wl = ws + row * K; + float s_val = float(sc[row * K_groups]); + float b_val = float(bi[row * K_groups]); + + for (int b = 0; b < {B}; b++) {{ + float accum = 0, xsum = 0; + for (int i = 0; i < VALUES_PER_THREAD; i++) {{ + float xi = float(((const device bfloat16_t*)x)[b * K + x_base + i]); + accum += xi * float(wl[i]); + xsum += xi; + }} + result[b * 4 + row] += s_val * accum + xsum * b_val; + }} + }} + + ws += BLOCK_SIZE; sc += SC_STRIDE; bi += SC_STRIDE; x_base += BLOCK_SIZE; + }} + + // Reduction + for (int i = 0; i < {4 * B}; i++) result[i] = simd_sum(result[i]); + + // Write output (bf16) + if (slid < 4u) {{ + for (int b = 0; b < {B}; b++) {{ + int r = out_row + (int)slid; + if (r < N) {{ + y[b * N + r] = static_cast(result[b * 4 + slid]); + }} + }} + }} +""" + + +_custom_qmv_cache = {} + + +def custom_qmv_loop_over_b(x, w, scales, biases, M, N, K, group_size=64): + """Loop-over-B GEMV for quantized matmul. + + Args: + x: (M, K) bfloat16 input + w: (N, K/4) uint32 packed 8-bit weights + scales: (N, K/gs) bfloat16 + biases: (N, K/gs) bfloat16 + M, N, K: dimensions + Returns: + y: (M, N) bfloat16 + """ + key = (M, N, K, group_size) + if key not in _custom_qmv_cache: + _custom_qmv_cache[key] = mx.fast.metal_kernel( + name=f"custom_qmv_loop_b_M{M}_N{N}_K{K}", + input_names=["x", "w", "scales", "biases"], + output_names=["y"], + source=_gen_custom_qmv_source(M, N, K, group_size), + ) + kern = _custom_qmv_cache[key] + + n_tg = ceil_div(N, 8) + + result = kern( + inputs=[x, w, scales, biases], + output_shapes=[(M * N,)], + output_dtypes=[mx.bfloat16], + grid=(32, n_tg * 2, 1), + threadgroup=(32, 2, 1), + ) + + return result[0].reshape(M, N) diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5/lpb_patch.py b/src/exo/worker/engines/mlx/patches/qwen3_5/lpb_patch.py new file mode 100644 index 00000000..1551f1fd --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5/lpb_patch.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Loop-over-B patches for Qwen3.5-27B dense model. + +Replaces vanilla QuantizedLinear calls with custom loop-over-B GEMV +for projections where N > K (expanding projections). Falls back to +vanilla for N <= K (contracting projections like down_proj, o_proj). + +Usage: + from lpb_patch import apply_lpb_patches + apply_lpb_patches(model, batch_size=4) +""" + +import mlx.core as mx +import mlx.nn as nn + +from .custom_qmv_loop_over_b import custom_qmv_loop_over_b + + +def _make_lpb_forward(original_module, N, K, BS, GS=64): + """Create a patched forward that uses loop-over-B.""" + w = original_module.weight + s = original_module.scales + b = original_module.biases + + MAX_M = 16 # Max total tokens (B*S) for custom kernel; above this use vanilla + + def forward(self_unused, x): + # Use LpB for small M=B*S. Large prefill falls back to vanilla. + M_total = 1 + for d in x.shape[:-1]: + M_total *= d + if M_total > MAX_M: + return original_module(x) + orig_shape = x.shape + x_2d = x.reshape(-1, K) + M = x_2d.shape[0] + y = custom_qmv_loop_over_b(x_2d, w, s, b, M, N, K, GS) + return y.reshape(*orig_shape[:-1], N) + + return forward + + +def apply_lpb_patches(model, batch_size=4): + """Patch all expanding QuantizedLinear projections with loop-over-B. + + Only patches projections where N > K (expanding): + - gate_proj, up_proj (17408 > 5120) + - in_proj_qkv (10240 > 5120) + - in_proj_z (6144 > 5120) + - q_proj (12288 > 5120) + + Skips N <= K projections (down_proj, o_proj, k_proj, v_proj) + where vanilla is already efficient. + """ + inner = getattr(model, 'model', None) or model.language_model.model + patched = 0 + + for li, layer in enumerate(inner.layers): + # MLP: gate_proj, up_proj (N=17408, K=5120) + mlp = layer.mlp + for proj_name in ['gate_proj', 'up_proj', 'down_proj']: + proj = getattr(mlp, proj_name) + if isinstance(proj, nn.QuantizedLinear): + N = proj.weight.shape[0] # output dim + K_packed = proj.weight.shape[1] + K = K_packed * 4 # 8-bit: 4 values per uint32 + setattr(mlp, proj_name, type('LpBLinear', (), { + '__call__': _make_lpb_forward(proj, N, K, batch_size), + 'weight': proj.weight, + 'scales': proj.scales, + 'biases': proj.biases, + })()) + patched += 1 + + # Attention projections + if layer.is_linear: + attn = layer.linear_attn + for proj_name in ['in_proj_qkv', 'in_proj_z', 'out_proj']: + if hasattr(attn, proj_name): + proj = getattr(attn, proj_name) + if isinstance(proj, nn.QuantizedLinear): + N = proj.weight.shape[0] + K = proj.weight.shape[1] * 4 + setattr(attn, proj_name, type('LpBLinear', (), { + '__call__': _make_lpb_forward(proj, N, K, batch_size), + 'weight': proj.weight, + 'scales': proj.scales, + 'biases': proj.biases, + })()) + patched += 1 + else: + attn = layer.self_attn + for proj_name in ['q_proj', 'o_proj']: + if hasattr(attn, proj_name): + proj = getattr(attn, proj_name) + if isinstance(proj, nn.QuantizedLinear): + N = proj.weight.shape[0] + K = proj.weight.shape[1] * 4 + setattr(attn, proj_name, type('LpBLinear', (), { + '__call__': _make_lpb_forward(proj, N, K, batch_size), + 'weight': proj.weight, + 'scales': proj.scales, + 'biases': proj.biases, + })()) + patched += 1 + + print(f" Patched {patched} projections with loop-over-B") + return patched diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/__init__.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/apply.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/apply.py new file mode 100644 index 00000000..e1f38bd2 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/apply.py @@ -0,0 +1,74 @@ +"""Apply batched fused kernel 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 ( + _patch_swiglu_weights, + _patch_shared_expert, + _patch_down_proj, + _patch_oproj_gate_rms, + _patch_gdn_proj_weights, + _patch_gqa_proj_weights, +) +from mlx_lm.models.qwen3_5 import DecoderLayer +from mlx_lm.models.qwen3_next import Qwen3NextAttention, Qwen3NextSparseMoeBlock +from mlx_lm.models.qwen3_5 import GatedDeltaNet + + +def apply_qwen35_batched_fused_patches(model: nn.Module) -> None: + """Apply batched fused patches (GDN + GQA attention + oproj MoE) to all layers. + + Fused GDN attention (3/4 layers) + fused GQA projections (1/4 layers) + + batched oproj MoE (4 custom dispatches). Works with BatchGenerator for + any batch size 1..8. Falls back to vanilla for B>8 or S>1. + """ + layers = model.layers # type: ignore[attr-defined] + n_layers = len(layers) + + t0 = time.time() + + n_gdn = 0 + n_gqa = 0 + for li, layer in enumerate(layers): + moe = layer.mlp + if isinstance(moe, Qwen3NextSparseMoeBlock): + # MoE weight prep + _patch_swiglu_weights(moe) + _patch_shared_expert(moe) + _patch_down_proj(moe) + _patch_oproj_gate_rms(layer, gate_bm=8) + + # Attention weight prep + 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 (li + 1) % 10 == 0 or li == 0: + logger.info(f" Patched layer {li+1}/{n_layers}") + + # Import patched __call__ methods + from .fused_gdn_attention import _fused_gdn_call + from .batched_fused_gqa_attention import _batched_fused_gqa_call + from .batched_moe import _batched_oproj_moe_call + from .decoder import _fused_gdn_decoder_call + + # Class-level method replacement + GatedDeltaNet.__call__ = _fused_gdn_call + Qwen3NextAttention.__call__ = _batched_fused_gqa_call + Qwen3NextSparseMoeBlock.__call__ = _batched_oproj_moe_call + DecoderLayer.__call__ = _fused_gdn_decoder_call + + t_patch = time.time() - t0 + logger.info( + f"Qwen3.5 batched fused: {n_gdn} GDN + {n_gqa} GQA layers, " + f"{n_layers} total in {t_patch:.1f}s" + ) diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/batched_fused_gqa_attention.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/batched_fused_gqa_attention.py new file mode 100644 index 00000000..1c1045a8 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/batched_fused_gqa_attention.py @@ -0,0 +1,102 @@ +"""Batched fused GQA attention for Qwen3.5 (projections + norm+rope fused, vanilla SDPA). + +Dispatches: + 1. batched_fused_gqa_projections — merged q+gate+k+v GEMV with register weight sharing + 2. fused_qk_norm_rope — per-head RMSNorm + RoPE (already supports B>1 via grid z) + 3. Vanilla cache update (BatchKVCache) + 4. Vanilla SDPA (MLX built-in, handles batching natively) + 5. Vanilla gate multiply + +Returns pre-out_proj output for the oproj MoE block. +Falls back to vanilla (with o_proj) for B>8 or S>1. +""" + +from typing import Any, Optional + +import mlx.core as mx + +from .kernels.batched_fused_gqa_projections_8bit import batched_fused_gqa_projections + + +def _batched_fused_gqa_call( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, +) -> mx.array: + """Batched fused GQA attention with custom projection + norm/rope kernels. + + For 1<=B<=8, S=1: fused projections + fused norm+rope + vanilla SDPA. + For B>8 or S>1: vanilla fallback. + Returns pre-out_proj output [B, S, H_q*D]. + """ + B, S, _ = x.shape + + if S > 1 or B > 8: + # Vanilla fallback + q_proj_output = self.q_proj(x) + queries, gate = mx.split( + q_proj_output.reshape(B, S, self.num_attention_heads, -1), 2, axis=-1 + ) + gate = gate.reshape(B, S, -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, S, self.num_key_value_heads, -1) + ).transpose(0, 2, 1, 3) + values = values.reshape(B, S, 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, S, -1) + return self.o_proj(output * mx.sigmoid(gate)) + + H_q = self.num_attention_heads + H_kv = self.num_key_value_heads + D = self.head_dim + + # ── Dispatch 1: batched fused projections ── + queries, gate_sigmoid, keys, values = batched_fused_gqa_projections( + x, + self._merged_proj_w, self._merged_proj_s, self._merged_proj_b, + self._merged_proj_dims, + batch_size=B, + total_tg=getattr(self, '_d1_total_tg', None), + ) + + # ── Dispatch 2+: vanilla norm + rope (avoids mx.eval sync on BatchKVCache offset) ── + queries = self.q_norm(queries.reshape(B, 1, H_q, D)).transpose(0, 2, 1, 3) + keys = self.k_norm(keys.reshape(B, 1, H_kv, D)).transpose(0, 2, 1, 3) + values = values.reshape(B, 1, H_kv, D).transpose(0, 2, 1, 3) + + if cache is not None: + queries = self.rope(queries, offset=cache.offset) + keys = self.rope(keys, offset=cache.offset) + else: + queries = self.rope(queries) + keys = self.rope(keys) + + # ── Dispatch 3: KV cache update ── + if cache is not None: + keys, values = cache.update_and_fetch(keys, values) + + # ── Dispatch 4: vanilla SDPA ── + 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, S, -1) + + # ── Gate multiply ── + return output * gate_sigmoid.astype(output.dtype) diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/batched_moe.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/batched_moe.py new file mode 100644 index 00000000..cb5df200 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/batched_moe.py @@ -0,0 +1,94 @@ +"""Batched oproj MoE with 4 custom Metal kernel dispatches. + +Fuses o_proj + RMSNorm + gate GEMV + softmax + topk + SwiGLU + down_proj + epilogue +into 4 dispatches with register-level weight sharing for the shared expert. +Falls back to vanilla MoE when called without _residual (from vanilla decoder path). +""" + +import mlx.core as mx + +from .kernels.batched_merged_down_proj_8bit import batched_merged_down_proj_8bit +from .kernels.batched_oproj_gate_gemv_8bit import batched_oproj_gate_gemv +from .kernels.batched_softmax_topk_swiglu_8bit import batched_softmax_topk_swiglu_8bit +from .kernels.batched_moe_epilogue import batched_moe_epilogue + + +def _batched_oproj_moe_call(self, attn_out_3d, _residual=None): + """Batched MoE with full oproj fusion (4 custom dispatches). + + Receives raw attention output (pre-o_proj) and residual for B tokens. + All 4 dispatches use register-level weight sharing for shared weights. + + When _residual is None, called from vanilla decoder — do vanilla MoE. + """ + if _residual is None: + # Vanilla MoE path (called from vanilla decoder for B>8 or S>1) + x = attn_out_3d + 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 + + B_dim = attn_out_3d.shape[0] + + K = self._oproj_M + K_attn = self._oproj_K_attn + n_active = self.top_k + E = self._oproj_n_experts + + attn_out = attn_out_3d.reshape(B_dim, K_attn).astype(mx.bfloat16) + residual = _residual.reshape(B_dim, K).astype(mx.bfloat16) + + # ── Dispatch 1: batched o_proj + gate GEMVs ── + h_scaled, h_out, x2_partials, gate_part_a, gate_part_b = \ + batched_oproj_gate_gemv( + 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, batch_size=B_dim, + n_experts=E, gate_bm=self._oproj_gate_bm, + K_hidden=self._oproj_K_hidden, + ) + + n_oproj_tg = (K + 31) // 32 + N_INTER = self.switch_mlp._fused_n_inter + SHARED_INTER = self._shared_inter + + # ── Dispatch 2: batched softmax + topk + SwiGLU ── + y_routed, y_shared, out_inds, norm_scores, gate_raw = \ + batched_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, + self._seg_w, self._seg_s, self._seg_b, + h_scaled, gate_part_a, gate_part_b, x2_partials, + n_inter=N_INTER, k_hidden=K, batch_size=B_dim, + n_active=n_active, n_oproj_tg=n_oproj_tg, + n_experts=E, shared_inter=SHARED_INTER, + ) + + # ── Dispatch 3: batched merged down_proj ── + d_routed, d_shared = batched_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.reshape(B_dim * SHARED_INTER), out_inds, + k_out=K, n_in=self._down_N, batch_size=B_dim, + n_active=n_active, shared_n_in=SHARED_INTER, + ) + + # ── Dispatch 4: batched epilogue ── + Y = batched_moe_epilogue( + d_routed, d_shared, norm_scores, + h_out, gate_raw, + k_val=K, batch_size=B_dim, n_active=n_active, + ) + + return Y.reshape(B_dim, 1, K).astype(attn_out_3d.dtype) diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/common.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/common.py new file mode 100644 index 00000000..841c7f19 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/common.py @@ -0,0 +1,500 @@ +"""Common weight preparation functions for Qwen3.5 fused kernel patches. + +Functions: + ceil_div — integer ceiling division + _patch_swiglu_weights — stack gate+up weights for fused SwiGLU kernel + _patch_down_proj — extract down_proj weights for merged kernel dispatch + _patch_shared_expert — prepare shared expert weights (8-bit) + dequantize_shared_expert — convert shared expert from 8-bit to bf16 + _patch_oproj_gate_rms — precompute M1/W_fused for fused o_proj + gate GEMV + _patch_gdn_proj_weights — merge GDN projection weights for fused GEMV + _patch_gqa_proj_weights — merge GQA q/k/v weights with q_proj permutation + make_qwen_random_cache — create pre-filled cache for testing + build_model — build Qwen3.5 MoE layers with 8-bit quantization +""" + +from types import SimpleNamespace + +import mlx.core as mx +import mlx.nn as nn +from mlx_lm.models.qwen3_5 import ( + DecoderLayer, + TextModelArgs, +) +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. + + Creates concatenated (E, 2*N_INTER, K/4) weight, (E, 2*N_INTER, K/gs) scales/biases + from the separate gate_proj and up_proj QuantizedSwitchLinear layers. + """ + 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. + + Stacks shared gate+up quantized weights (weight, scales, biases). + Stores down_proj quantized weights separately. + Shared expert stays in 8-bit — same as vanilla MLX dispatch. + """ + shared = moe.shared_expert + gp = shared.gate_proj + up = shared.up_proj + dp = shared.down_proj + + # Gate+up stacked: (2*SHARED_INTER, K/4) uint32, (2*SHARED_INTER, K/gs) bf16 + 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) + + # Down_proj: (K, SHARED_INTER/4) uint32, (K, SHARED_INTER/gs) bf16 + moe._shared_down_w = dp.weight + moe._shared_down_s = dp.scales + moe._shared_down_b = dp.biases + + # QuantizedLinear: weight is (out_features, in_features/pack_factor) uint32 + # For 8-bit: pack_factor = 4, so in_features = weight.shape[1] * 4 + moe._shared_inter = gp.weight.shape[0] # SHARED_INTER (= out_features) + moe._shared_gs = gp.group_size # gs (64) + + 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 # (E, K_OUT, N_IN/4) uint32 + moe._down_s = dp.scales # (E, K_OUT, N_IN/gs) bf16 + moe._down_b = dp.biases # (E, K_OUT, N_IN/gs) bf16 + moe._down_K = dp.output_dims # K = 4096 + moe._down_N = dp.input_dims # N = 1024 + moe._down_gs = dp.group_size # gs = 64 + mx.eval(moe._down_w, moe._down_s, moe._down_b) + + +def dequantize_shared_expert(moe): + """Convert shared expert from 8-bit QuantizedLinear to bf16 weight wrappers. + + The fused kernels expect bf16 shared expert weights. The real model (and our + random model) has shared expert quantized to 8-bit. This dequantizes in-place. + """ + shared = moe.shared_expert + for proj_name in ["gate_proj", "up_proj", "down_proj"]: + proj = getattr(shared, proj_name) + if hasattr(proj, 'scales') and hasattr(proj, 'biases'): + w_bf16 = mx.dequantize( + proj.weight, proj.scales, proj.biases, + group_size=proj.group_size, bits=proj.bits, + ).astype(mx.bfloat16) + mx.eval(w_bf16) + setattr(shared, proj_name, SimpleNamespace(weight=w_bf16)) + + +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 + + Also stores o_proj quantized weights and shared_expert_gate weights + on the moe block for use by Dispatch 1 and Dispatch 2. + + Args: + layer: DecoderLayer instance + gate_bm: SGs per gate TG in Dispatch 1 (1,2,4,8) + """ + moe = layer.mlp + + # ── Get attention output projection (works for both attention types) ── + 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) ── + # Eval incrementally to limit peak memory (E=512: dequant temps are ~140 MB) + 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) + + # ── RMSNorm weight ── + rms_weight = layer.post_attention_layernorm.weight.astype(mx.bfloat16) + + # ── W_fused = dequant(W_gate) · diag(w_rms) ── + 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 # free ~8 MB (E=512) or ~1 MB (E=64) + + # ── M1 = W_fused @ W_oproj — precomputed in f32, stored bf16 ── + M1 = (W_fused.astype(mx.float32) @ W_oproj_f32).astype(mx.bfloat16) + mx.eval(M1) + del W_oproj_f32 # free ~128 MB + + # Store on moe block + moe._oproj_M1 = M1 # (E, K_attn) bf16 + moe._oproj_W_fused = W_fused # (E, K) bf16 + moe._oproj_rms_weight = rms_weight # (K,) bf16 + + # ── O_proj quantized weights (for 8-bit GEMV in Dispatch 1) ── + moe._oproj_w = oproj.weight # (K, K_attn/4) uint32 + moe._oproj_s = oproj.scales # (K, K_attn/gs) bf16 + moe._oproj_b = oproj.biases # (K, K_attn/gs) bf16 + moe._oproj_K_attn = oproj.weight.shape[1] * 4 # 8192 (8-bit: pack_factor=4) + + # ── Shared expert gate weights (for TG(0,0,0) fusion in Dispatch 2) ── + seg = moe.shared_expert_gate + moe._seg_w = seg.weight # (1, K/4) uint32 + moe._seg_s = seg.scales # (1, K/gs) bf16 + moe._seg_b = seg.biases # (1, K/gs) bf16 + + # ── Dimensions ── + M = oproj.weight.shape[0] # 4096 (hidden_size) + K_hidden = W_fused.shape[1] # 4096 (same as M for Qwen) + n_experts = W_fused.shape[0] # E + moe._oproj_M = M + moe._oproj_K_hidden = K_hidden + moe._oproj_n_experts = n_experts + moe._oproj_n_tg = ceil_div(M, 32) # 128 for M=4096 + moe._oproj_gate_bm = gate_bm + + mx.eval(moe._oproj_rms_weight) + + + +def _patch_gdn_proj_weights(attn): + """Merge all 4 GDN projection weights into contiguous buffers. + + Concatenates in_proj_qkv/z/b/a weights, scales, biases into single + contiguous arrays for better memory locality in the fused GEMV kernel. + Stored on the GatedDeltaNet module as _merged_proj_*. + """ + 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], # N_QKV = 8192 + attn.in_proj_z.weight.shape[0], # N_Z = 4096 + attn.in_proj_b.weight.shape[0], # N_B = 32 + attn.in_proj_a.weight.shape[0], # N_A = 32 + ) + mx.eval(W_merged, S_merged, B_merged) + + +def _patch_gqa_proj_weights(attn): + """Merge GQA q_proj, k_proj, v_proj weights into contiguous buffers. + + q_proj outputs (H_q * 2 * D) = interleaved [queries, gate] per head. + We permute rows so queries (H_q * D) come first, then gate (H_q * D), + then k_proj, then v_proj. This gives clean contiguous regions for + the fused GEMV kernel's TG routing. + + Permutation for q_proj: + Original row layout: [head0_q[0:D], head0_gate[0:D], head1_q[0:D], ...] + After permutation: [head0_q, head1_q, ..., head0_gate, head1_gate, ...] + + Stored on Qwen3NextAttention as _merged_proj_*. + """ + q = attn.q_proj + k = attn.k_proj + v = attn.v_proj + + H_q = attn.num_attention_heads + D = attn.head_dim + + # Permute q_proj weights: separate queries and gate rows + # q_proj.weight shape: (H_q * 2 * D, K / pack_factor) for 8-bit + # Reshape to (H_q, 2*D, ...), split into queries[:, :D, :] and gate[:, D:, :] + 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) + + # Merge: [queries, gate, keys, values] + 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 + attn._merged_proj_dims = ( + H_q * D, # N_Q = 4096 (queries) + H_q * D, # N_GATE = 4096 (gate) + k.weight.shape[0], # N_K = 512 + v.weight.shape[0], # N_V = 512 + ) + mx.eval(W_merged, S_merged, B_merged) + + # Pre-cache constant scalar arrays for kernel dispatch (avoid per-call creation) + N_Q, N_GATE, N_K, N_V = attn._merged_proj_dims + N_TOTAL = N_Q + N_GATE + N_K + N_V + K = q.weight.shape[1] * 4 # 8-bit: pack_factor=4 + attn._kernel_scalars = { + # Dispatch 1: fused_gqa_projections + 'K': mx.array(K, 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), + # Dispatch 4-5: custom SDPA + '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()) + + # Precompute grid/TG dims for Dispatch 1 + 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 for fused norm+rope kernel (Dispatch 2) + # inv_freq[d] = theta^(-d / half_dims) for d in {0, ..., half_dims-1} + rope_dims = attn.rope.dims # 64 (partial_rotary_factor * head_dim) + half_dims = rope_dims // 2 # 32 + theta = attn.rope.base # 10000000 + 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 make_qwen_random_cache(layer, config, prefill_len): + """Create a pre-filled cache for a single Qwen3.5 decoder layer. + + GatedDeltaNet layers get ArraysCache(size=2) with fixed-size state: + cache[0] = conv state: (B, conv_kernel_size-1, conv_dim) bf16 + cache[1] = SSM state: (B, num_v_heads, head_k_dim, head_v_dim) bf16 + + GQA layers get KVCache with prefill_len tokens: + keys: (B, n_kv_heads, alloc_len, head_dim) bf16 + values: (B, n_kv_heads, alloc_len, head_dim) bf16 + """ + if layer.is_linear: + from mlx_lm.models.cache import ArraysCache + cache = ArraysCache(size=2) + attn = layer.linear_attn + cache[0] = mx.random.normal( + (1, attn.conv_kernel_size - 1, attn.conv_dim) + ).astype(mx.bfloat16) + cache[1] = mx.random.normal( + (1, attn.num_v_heads, attn.head_k_dim, attn.head_v_dim) + ).astype(mx.bfloat16) + return cache + else: + from mlx_lm.models.cache import KVCache + cache = KVCache() + n_steps = (prefill_len + KVCache.step - 1) // KVCache.step + alloc_len = n_steps * KVCache.step + n_kv = config.num_key_value_heads + hd = config.head_dim + cache.keys = mx.random.normal((1, n_kv, alloc_len, hd)).astype(mx.bfloat16) + cache.values = mx.random.normal((1, n_kv, alloc_len, hd)).astype(mx.bfloat16) + cache.offset = prefill_len + return cache + + +def build_model(n_experts=16, n_layers=1, top_k=4, + hidden_size=4096, moe_intermediate_size=1024, + shared_expert_intermediate_size=2048, tp=1, + n_attn_heads=32, n_kv_heads=2, + lin_v_heads=64, lin_k_heads=16, + head_dim=256): + """Build Qwen3.5 MoE decoder layers with 8-bit gs=64 quantization. + + Matches real mlx-community Qwen3.5 quantization: + Everything 8-bit gs=64 (gate, experts, shared expert, shared_expert_gate, attention/SSM). + RMSNorm weights: bf16. + + Default dimensions are for Qwen3.5-397B-A17B. For 35B-A3B, pass: + n_attn_heads=16, lin_v_heads=32, hidden_size=2048 + + Uses qwen3_5.DecoderLayer (same class as real model) with hybrid attention: + 3/4 GatedDeltaNet (SSM-like), 1/4 full attention (full_attention_interval=4). + + TP=2 halves all column-parallel dimensions. + + Returns: + layers: list of DecoderLayer instances + config: TextModelArgs + GROUP_SIZE: int (64) + """ + GROUP_SIZE = 64 + BITS = 8 + + # Apply TP sharding + moe_inter = moe_intermediate_size // tp + shared_inter = shared_expert_intermediate_size // tp + n_attn_heads_tp = n_attn_heads // tp + n_kv_heads_tp = max(1, n_kv_heads // tp) + lin_v_heads_tp = lin_v_heads // tp + lin_k_heads_tp = lin_k_heads // tp + + config = TextModelArgs( + model_type="qwen3_5_moe", + hidden_size=hidden_size, + num_hidden_layers=n_layers, + intermediate_size=moe_inter, + num_attention_heads=n_attn_heads_tp, + num_key_value_heads=n_kv_heads_tp, + linear_num_value_heads=lin_v_heads_tp, + linear_num_key_heads=lin_k_heads_tp, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_conv_kernel_dim=4, + num_experts=n_experts, + num_experts_per_tok=top_k, + decoder_sparse_step=1, + shared_expert_intermediate_size=shared_inter, + moe_intermediate_size=moe_inter, + norm_topk_prob=True, + rms_norm_eps=1e-6, + vocab_size=248320, + head_dim=head_dim, + full_attention_interval=4, + max_position_embeddings=262144, + rope_theta=10000000, + partial_rotary_factor=0.25, + rope_parameters={ + "type": "default", + "rope_theta": 10000000, + "partial_rotary_factor": 0.25, + }, + ) + + tp_str = f" (TP={tp})" if tp > 1 else "" + print(f" Config: {n_layers} layer(s), {n_experts} experts, top_k={top_k}, " + f"hidden={hidden_size}, inter={moe_inter}, shared={shared_inter}{tp_str}") + print(f" Quant: {BITS}-bit gs={GROUP_SIZE} (all weights)") + + layers = [DecoderLayer(config, idx) for idx in range(n_layers)] + + for li, layer in enumerate(layers): + # Cast all attention/SSM params to bf16 before quantizing, matching + # real safetensors model where all non-quantized params are bf16. + # nn.quantize only touches nn.Linear; other params (conv1d, dt_bias, + # norm, A_log) must be cast manually. + attn_mod = layer.linear_attn if layer.is_linear else layer.self_attn + for name, mod in attn_mod.named_modules(): + if isinstance(mod, nn.Linear): + mod.weight = mod.weight.astype(mx.bfloat16) + elif isinstance(mod, nn.Conv1d): + mod.weight = mod.weight.astype(mx.bfloat16) + # Cast leaf parameters (dt_bias, norm.weight, q/k_norm) to bf16 + # A_log stays f32 (matches real model) + if layer.is_linear: + gdn = layer.linear_attn + gdn.dt_bias = gdn.dt_bias.astype(mx.bfloat16) + gdn.norm.weight = gdn.norm.weight.astype(mx.bfloat16) + else: + gqa = layer.self_attn + gqa.q_norm.weight = gqa.q_norm.weight.astype(mx.bfloat16) + gqa.k_norm.weight = gqa.k_norm.weight.astype(mx.bfloat16) + nn.quantize(attn_mod, bits=BITS, group_size=GROUP_SIZE) + mx.eval(attn_mod.parameters()) + + # RMSNorm to bf16 (norms are never quantized) + layer.input_layernorm.weight = layer.input_layernorm.weight.astype(mx.bfloat16) + layer.post_attention_layernorm.weight = layer.post_attention_layernorm.weight.astype(mx.bfloat16) + mx.eval(layer.input_layernorm.weight, layer.post_attention_layernorm.weight) + + # MoE block: quantize everything to 8-bit gs=64 + moe = layer.mlp + if isinstance(moe, Qwen3NextSparseMoeBlock): + # Gate: random init (zeros get optimized away), then quantize. + # nn.quantize on a leaf nn.Linear is a no-op (walks children, finds none). + # Use QuantizedLinear.from_linear directly. + moe.gate.weight = ( + mx.random.normal(moe.gate.weight.shape) * 0.01 + ).astype(mx.float32) + moe.gate = nn.QuantizedLinear.from_linear( + moe.gate, group_size=GROUP_SIZE, bits=BITS) + mx.eval(moe.gate.parameters()) + + # Routed experts: quantize per-projection to limit peak memory + nn.quantize(moe.switch_mlp, bits=BITS, group_size=GROUP_SIZE) + mx.eval(moe.switch_mlp.gate_proj.parameters()) + mx.eval(moe.switch_mlp.up_proj.parameters()) + mx.eval(moe.switch_mlp.down_proj.parameters()) + + # Shared expert: quantize to 8-bit (matching real model) + nn.quantize(moe.shared_expert, bits=BITS, group_size=GROUP_SIZE) + mx.eval(moe.shared_expert.parameters()) + + # shared_expert_gate: quantize to 8-bit gs=64 (leaf nn.Linear fix) + moe.shared_expert_gate = nn.QuantizedLinear.from_linear( + moe.shared_expert_gate, group_size=GROUP_SIZE, bits=BITS) + mx.eval(moe.shared_expert_gate.parameters()) + + if (li + 1) % 10 == 0 or li == 0: + print(f" Layer {li+1}/{n_layers} ready") + + return layers, config, GROUP_SIZE diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/decoder.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/decoder.py new file mode 100644 index 00000000..47c286fc --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/decoder.py @@ -0,0 +1,199 @@ +"""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. +""" + +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 _vanilla_decoder_call(self, x, mask=None, cache=None): + """Original vanilla DecoderLayer.__call__ (fallback for B>8 or S>1).""" + 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)) + return h + out + + +def _fused_gdn_decoder_call(self, x, mask=None, cache=None): + """Decoder with batched fused kernels. Falls back to vanilla for B>8 or S>1. + + When fused: attention returns pre-out_proj output, MoE handles oproj + gate + experts. + When vanilla: original DecoderLayer flow (attention + residual + layernorm + MoE). + """ + B = x.shape[0] + S = x.shape[1] + + # Full vanilla fallback for large batch or prefill + if B > 8 or S > 1: + return _vanilla_decoder_call(self, x, mask, cache) + + # Fused path: attention returns pre-oproj, MoE handles the rest + 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=True, + ) + + if cache is not None: + cache[1] = state + + out = self.norm(out, z) + return out.reshape(B, S, -1) # skip out_proj diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/fused_gdn_attention.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/fused_gdn_attention.py new file mode 100644 index 00000000..e081653e --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/fused_gdn_attention.py @@ -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.batched_fused_gdn_projections_8bit import batched_fused_gdn_projections as 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 self.out_proj(out.reshape(B, S, -1)) # include out_proj for vanilla decoder + + +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 + + # Vanilla fallback: fused kernels are decode-only (S=1, B<=8) + if S > 1 or B > 8: + 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) diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/__init__.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_fused_gdn_projections_8bit.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_fused_gdn_projections_8bit.py new file mode 100644 index 00000000..05399257 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_fused_gdn_projections_8bit.py @@ -0,0 +1,265 @@ +"""Batched fused GDN projections for Qwen3.5-35B-A3B, batch_size 1..8. + +Register-level weight sharing: each TG loads weights once, computes B outputs. +Adapts fused_gdn_projections_8bit with the same pattern as +batched_fused_gqa_projections_8bit. + +4 regions with different epilogues: + - QKV: GEMV → conv1d(4-tap) → SiLU → bf16 + cache update + - Z: GEMV → SiLU → f32 + - B: GEMV → sigmoid → f32 (beta for GDN kernel) + - A: GEMV → g=exp(-exp(A_log)*softplus(a+dt_bias)) → f32 + +All constants baked into Metal source. B unrolled at code-generation time. + +Grid: (32, total_tg * 2, 1), TG: (32, 2, 1) +No grid z for batch — batch is handled in registers. +""" + +import mlx.core as mx + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def _gen_batched_fused_gdn_proj_source(K, N_QKV, N_Z, N_B, N_A, B, group_size=64): + gs = int(group_size) + sc_stride = 256 // gs + slid_div = gs // 8 + N_TOTAL = N_QKV + N_Z + N_B + N_A + K_groups = K // gs + N_QKV_TG = ceil_div(N_QKV, 8) + N_Z_TG = ceil_div(N_Z, 8) + N_B_TG = ceil_div(N_B, 8) + + # Per-batch x loading (B unrolled) + x_load = "\n".join(f""" + float x{b}_thread[VALUES_PER_THREAD]; float xsum{b} = 0; + for (int i = 0; i < VALUES_PER_THREAD; i++) {{ + float xi = float(x[{b} * K + x_base + i]); x{b}_thread[i] = xi; xsum{b} += xi; + }}""" for b in range(B)) + + # Per-batch dot product with weights in registers + qdot = "\n".join(f""" + float accum{b} = 0; + for (int i = 0; i < VALUES_PER_THREAD; i++) accum{b} += x{b}_thread[i] * w_vals[i]; + result{b}[row] += s_val * accum{b} + xsum{b} * b_val;""" for b in range(B)) + + result_decls = " ".join(f"float result{b}[4] = {{0,0,0,0}};" for b in range(B)) + simd_reduce = "\n ".join( + f"for (int row = 0; row < 4; row++) result{b}[row] = simd_sum(result{b}[row]);" for b in range(B)) + + # QKV epilogue: conv1d + SiLU + cache update, per batch + qkv_write = "\n".join(f""" + if (slid < 4u && c < N_QKV) {{ + float qkv_val = result{b}[slid]; + long cs_base = (long){b} * 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(s1); + conv_state_out[cs_base + 1 * conv_dim + c] = static_cast(s2); + conv_state_out[cs_base + 2 * conv_dim + c] = static_cast(qkv_val); + qkv_out[{b} * conv_dim + c] = static_cast(silu_out); + }}""" for b in range(B)) + + # Z epilogue: SiLU per batch + z_write = "\n".join(f""" + if (slid < 4u && z_row < N_Z) {{ + float val = result{b}[slid]; + z_silu_out[{b} * N_Z + z_row] = val / (1.0f + metal::exp(-val)); + }}""" for b in range(B)) + + # B epilogue: sigmoid per batch + b_write = "\n".join(f""" + if (slid < 4u && b_row < N_B) {{ + b_out[{b} * N_B + b_row] = 1.0f / (1.0f + metal::exp(-result{b}[slid])); + }}""" for b in range(B)) + + # A epilogue: g computation per batch + a_write = "\n".join(f""" + if (slid < 4u && a_row < N_A_val) {{ + float a_val = result{b}[slid]; + float dt = float(dt_bias_arr[a_row]); + float x_g = a_val + dt; + 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} * N_A_val + a_row] = g_val; + }}""" for b in range(B)) + + return f""" + const int RESULTS_PER_SG = 4; + const int VALUES_PER_THREAD = 8; + const int BLOCK_SIZE = 256; + const int GROUP_SIZE = {gs}; + const int SC_STRIDE = {sc_stride}; + const int SLID_DIV = {slid_div}; + const int K = {K}; + const int K_groups = {K_groups}; + const int N_QKV = {N_QKV}; + const int N_Z = {N_Z}; + const int N_B = {N_B}; + const int N_TOTAL = {N_TOTAL}; + const int N_QKV_TG = {N_QKV_TG}; + const int N_Z_TG = {N_Z_TG}; + const int N_B_TG = {N_B_TG}; + + uint3 tgid = threadgroup_position_in_grid; + uint sgid = simdgroup_index_in_threadgroup; + uint slid = thread_index_in_simdgroup; + int tg = tgid.y; + + int out_row, region; + 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; + + // Weight pointers (shared across all batch elements) + 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; + + {result_decls} + int x_base = slid * VALUES_PER_THREAD; + + // K-loop: load weights into registers once, compute {B} batch elements + for (int k_off = 0; k_off < K; k_off += BLOCK_SIZE) {{ +{x_load} + + for (int row = 0; row < RESULTS_PER_SG; row++) {{ + const device uint8_t* wl = ws + row * K; + float s_val = float(sc[row * K_groups]); + float b_val = float(bi[row * K_groups]); + float w_vals[VALUES_PER_THREAD]; + for (int i = 0; i < VALUES_PER_THREAD; i++) w_vals[i] = float(wl[i]); +{qdot} + }} + + ws += BLOCK_SIZE; sc += SC_STRIDE; bi += SC_STRIDE; x_base += BLOCK_SIZE; + }} + + {simd_reduce} + + // Region-specific epilogues for all {B} batches + if (region == 0) {{ + int c = out_row + (int)slid; + int conv_dim = N_QKV; +{qkv_write} + }} else if (region == 1) {{ + int z_row = out_row - N_QKV + (int)slid; +{z_write} + }} else if (region == 2) {{ + int b_row = out_row - N_QKV - N_Z + (int)slid; +{b_write} + }} else {{ + int a_row = out_row - N_QKV - N_Z - N_B + (int)slid; + int N_A_val = N_TOTAL - N_QKV - N_Z - N_B; +{a_write} + }} +""" + + +_batched_gdn_proj_cache = {} + + +def _get_batched_gdn_proj_kernel(K, N_QKV, N_Z, N_B, N_A, B, group_size=64): + key = (K, N_QKV, N_Z, N_B, N_A, B, group_size) + if key not in _batched_gdn_proj_cache: + _batched_gdn_proj_cache[key] = mx.fast.metal_kernel( + name=f"batched_fused_gdn_proj_K{K}_NQKV{N_QKV}_B{B}", + input_names=[ + "x", + "W_merged", "S_merged", "B_merged", + "conv_state", "conv_w", + "A_log_arr", "dt_bias_arr", + ], + output_names=["qkv_out", "z_silu_out", "b_out", "a_out", "conv_state_out"], + source=_gen_batched_fused_gdn_proj_source(K, N_QKV, N_Z, N_B, N_A, B, group_size), + ) + return _batched_gdn_proj_cache[key] + + +def batched_fused_gdn_projections( + x, + W_merged, S_merged, B_merged, + proj_dims, + conv_state, conv_weights, + A_log, dt_bias, + batch_size=1, +): + """Batched fused GDN projections with register-level weight sharing. + + Same as fused_gdn_projections but loads weights once per TG and computes + B outputs from registers. No grid z for batch. + + Args: + x: [B, 1, K] bf16 — post-RMSNorm hidden state + W_merged, S_merged, B_merged: merged quantized weights + proj_dims: (N_QKV, N_Z, N_B, N_A) + conv_state: [B, 3, conv_dim] bf16 + conv_weights: [conv_dim, 4, 1] or [conv_dim, 4] bf16 + A_log: [Hv] f32, dt_bias: [Hv] f32 + batch_size: int (1..8) + + Returns: + qkv_conv_silu: [B, 1, N_QKV] bf16 + z_silu: [B, 1, N_Z] f32 + beta: [B, 1, N_B] f32 + g: [B, 1, N_A] f32 + conv_state_out: [B, 3, N_QKV] bf16 + """ + B = batch_size + + N_QKV, N_Z, N_B, N_A = proj_dims + K = x.shape[-1] + + kern = _get_batched_gdn_proj_kernel(K, N_QKV, N_Z, N_B, N_A, B) + + 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 + + conv_w_flat = conv_weights.reshape(-1, 4) if conv_weights.ndim == 3 else conv_weights + 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, + ], + output_shapes=[ + (B * N_QKV,), + (B * N_Z,), + (B * N_B,), + (B * N_A,), + (B * 3 * N_QKV,), + ], + output_dtypes=[mx.bfloat16, mx.float32, mx.float32, mx.float32, mx.bfloat16], + grid=(32, total_tg * 2, 1), # No grid z — batch in registers + 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 diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_fused_gqa_projections_8bit.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_fused_gqa_projections_8bit.py new file mode 100644 index 00000000..96f0b1e9 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_fused_gqa_projections_8bit.py @@ -0,0 +1,205 @@ +"""Batched fused GQA projections (Dispatch 1) for batch_size 1..8. + +Adapts fused_gqa_projections_8bit for B>1 with register-level weight sharing. +Each TG loads weights once, computes B outputs from registers. + +4 regions with different epilogues (same as B=1): + - Queries: GEMV → raw bf16 + - Gate: GEMV → sigmoid → f32 + - Keys: GEMV → raw bf16 + - Values: GEMV → raw bf16 + +All constants baked into Metal source. B unrolled at code-generation time. +""" + +import mlx.core as mx + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def _gen_batched_fused_gqa_proj_source(K, N_Q, N_GATE, N_K, N_V, B, group_size=64): + gs = int(group_size) + sc_stride = 256 // gs + slid_div = gs // 8 + N_TOTAL = N_Q + N_GATE + N_K + N_V + K_groups = K // gs + N_Q_TG = ceil_div(N_Q, 8) + N_GATE_TG = ceil_div(N_GATE, 8) + N_K_TG = ceil_div(N_K, 8) + + # Per-batch x loading + x_load = "\n".join(f""" + float x{b}_thread[VALUES_PER_THREAD]; float xsum{b} = 0; + for (int i = 0; i < VALUES_PER_THREAD; i++) {{ + float xi = float(x[{b} * K + x_base + i]); x{b}_thread[i] = xi; xsum{b} += xi; + }}""" for b in range(B)) + + # Per-batch qdot (weights in registers) + qdot = "\n".join(f""" + float accum{b} = 0; + for (int i = 0; i < VALUES_PER_THREAD; i++) accum{b} += x{b}_thread[i] * w_vals[i]; + result{b}[row] += s_val * accum{b} + xsum{b} * b_val;""" for b in range(B)) + + result_decls = " ".join(f"float result{b}[4] = {{0,0,0,0}};" for b in range(B)) + simd_reduce = "\n ".join( + f"for (int row = 0; row < 4; row++) result{b}[row] = simd_sum(result{b}[row]);" for b in range(B)) + + # Queries epilogue (bf16 write per batch) + q_write = "\n".join(f""" + if (slid < 4u && q_row < N_Q) q_out[{b} * N_Q + q_row] = static_cast(result{b}[slid]);""" + for b in range(B)) + + # Gate epilogue (sigmoid → f32 per batch) + gate_write = "\n".join(f""" + if (slid < 4u && g_row < N_GATE) {{ + float sig{b} = 1.0f / (1.0f + metal::exp(-result{b}[slid])); + gate_out[{b} * N_GATE + g_row] = sig{b}; + }}""" for b in range(B)) + + # Keys epilogue + k_write = "\n".join(f""" + if (slid < 4u && k_row < N_K) k_out[{b} * N_K + k_row] = static_cast(result{b}[slid]);""" + for b in range(B)) + + # Values epilogue + v_write = "\n".join(f""" + if (slid < 4u && v_row < N_V) v_out[{b} * N_V + v_row] = static_cast(result{b}[slid]);""" + for b in range(B)) + + N_V_val = N_TOTAL - N_Q - N_GATE - N_K + + return f""" + const int RESULTS_PER_SG = 4; + const int VALUES_PER_THREAD = 8; + const int BLOCK_SIZE = 256; + const int GROUP_SIZE = {gs}; + const int SC_STRIDE = {sc_stride}; + const int SLID_DIV = {slid_div}; + const int K = {K}; + const int K_groups = {K_groups}; + const int N_Q = {N_Q}; + const int N_GATE = {N_GATE}; + const int N_K = {N_K}; + const int N_V = {N_V_val}; + const int N_TOTAL = {N_TOTAL}; + const int N_Q_TG = {N_Q_TG}; + const int N_GATE_TG = {N_GATE_TG}; + const int N_K_TG = {N_K_TG}; + + uint3 tgid = threadgroup_position_in_grid; + uint sgid = simdgroup_index_in_threadgroup; + uint slid = thread_index_in_simdgroup; + int b_idx = tgid.z; + int tg = tgid.y; + + int out_row, region; + 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; + + // Weight pointers (shared across all batch elements) + 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; + + {result_decls} + int x_base = slid * VALUES_PER_THREAD; + + // K-loop: load weights once, compute {B} batch elements + for (int k_off = 0; k_off < K; k_off += BLOCK_SIZE) {{ +{x_load} + + for (int row = 0; row < RESULTS_PER_SG; row++) {{ + const device uint8_t* wl = ws + row * K; + float s_val = float(sc[row * K_groups]); + float b_val = float(bi[row * K_groups]); + float w_vals[VALUES_PER_THREAD]; + for (int i = 0; i < VALUES_PER_THREAD; i++) w_vals[i] = float(wl[i]); +{qdot} + }} + + ws += BLOCK_SIZE; sc += SC_STRIDE; bi += SC_STRIDE; x_base += BLOCK_SIZE; + }} + + {simd_reduce} + + // Region-specific epilogues for all {B} batches + if (region == 0) {{ + int q_row = out_row + (int)slid; +{q_write} + }} else if (region == 1) {{ + int g_row = out_row - N_Q + (int)slid; +{gate_write} + }} else if (region == 2) {{ + int k_row = out_row - N_Q - N_GATE + (int)slid; +{k_write} + }} else {{ + int v_row = out_row - N_Q - N_GATE - N_K + (int)slid; +{v_write} + }} +""" + + +_batched_proj_cache = {} + + +def _get_batched_proj_kernel(K, N_Q, N_GATE, N_K, N_V, B, group_size=64): + key = (K, N_Q, N_GATE, N_K, N_V, B, group_size) + if key not in _batched_proj_cache: + _batched_proj_cache[key] = mx.fast.metal_kernel( + name=f"batched_fused_gqa_proj_K{K}_NQ{N_Q}_B{B}", + input_names=["x", "W_merged", "S_merged", "B_merged"], + output_names=["q_out", "gate_out", "k_out", "v_out"], + source=_gen_batched_fused_gqa_proj_source(K, N_Q, N_GATE, N_K, N_V, B, group_size), + ) + return _batched_proj_cache[key] + + +def batched_fused_gqa_projections(x, W_merged, S_merged, B_merged, proj_dims, + batch_size, total_tg=None): + """Batched fused GQA projections with register weight sharing. + + Args: + x: [B, 1, K] bf16 + W_merged, S_merged, B_merged: merged q+gate+k+v weights + proj_dims: (N_Q, N_GATE, N_K, N_V) + batch_size: B (1..8) + + Returns: + queries (B, 1, N_Q) bf16, gate_sigmoid (B, 1, N_GATE) f32, + keys (B, 1, N_K) bf16, values (B, 1, N_V) bf16 + """ + B = batch_size + N_Q, N_GATE, N_K, N_V = proj_dims + K = x.shape[-1] + + kern = _get_batched_proj_kernel(K, N_Q, N_GATE, N_K, N_V, B) + + 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) + + x_flat = x.reshape(B, K) + + results = kern( + inputs=[x_flat, W_merged, S_merged, B_merged], + output_shapes=[ + (B * N_Q,), (B * N_GATE,), (B * N_K,), (B * N_V,), + ], + output_dtypes=[mx.bfloat16, mx.float32, mx.bfloat16, mx.bfloat16], + grid=(32, total_tg * 2, 1), + threadgroup=(32, 2, 1), + ) + + return (results[0].reshape(B, 1, N_Q), + results[1].reshape(B, 1, N_GATE), + results[2].reshape(B, 1, N_K), + results[3].reshape(B, 1, N_V)) diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_merged_down_proj_8bit.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_merged_down_proj_8bit.py new file mode 100644 index 00000000..4f8d4313 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_merged_down_proj_8bit.py @@ -0,0 +1,252 @@ +"""Batched merged 8-bit down_proj GEMV for Qwen3.5 routed + shared experts. + +Adapts merged_down_proj_8bit for batch_size B (1..8). + +Grid z-dimension: B * n_active + 1 + - tgid.z < B * n_active: routed experts (one TG per batch×expert pair) + Same structure as affine_gather_qmv — each TG independently indexes into + expert weights via inds[flat_idx]. + - tgid.z == B * n_active: shared expert (ONE TG handles ALL B batch elements + with register-level weight sharing — loads weights once, computes B outputs) + +All constants baked into Metal source (no scalar kernel inputs). +""" +import mlx.core as mx + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def _gen_batched_merged_down_8bit_source(K_OUT, N_IN, SHARED_N_IN, n_active, B, group_size=64): + gs = int(group_size) + sc_stride = 256 // gs + slid_divisor = gs // 8 + N_groups = N_IN // gs + SHARED_N_groups = SHARED_N_IN // gs + total_routed = B * n_active + + # Shared expert: generate unrolled batch loops + shared_x_load_lines = [] + for b in range(B): + shared_x_load_lines.append(f""" + float x{b}_thread[VALUES_PER_THREAD]; + float xsum{b} = 0; + for (int i = 0; i < VALUES_PER_THREAD; i++) {{ + float xi = X_shared[{b} * SHARED_N_IN + x_base + i]; + x{b}_thread[i] = xi; + xsum{b} += xi; + }}""") + shared_x_load = "\n".join(shared_x_load_lines) + + shared_qdot_lines = [] + for b in range(B): + shared_qdot_lines.append(f""" + float accum{b} = 0; + for (int i = 0; i < VALUES_PER_THREAD; i++) {{ + accum{b} += x{b}_thread[i] * w_vals[i]; + }} + result{b}[row] += s_val * accum{b} + xsum{b} * b_val;""") + shared_qdot = "\n".join(shared_qdot_lines) + + shared_result_decls = "\n ".join( + f"float result{b}[RESULTS_PER_SG] = {{0, 0, 0, 0}};" + for b in range(B) + ) + + shared_write_lines = [] + for b in range(B): + shared_write_lines.append(f""" + for (int row = 0; row < RESULTS_PER_SG; row++) {{ + float r{b} = simd_sum(result{b}[row]); + if (slid == 0) {{ + Y_shared[{b} * K_OUT + out_row + row] = r{b}; + }} + }}""") + shared_write = "\n".join(shared_write_lines) + + return f""" + const int RESULTS_PER_SG = 4; + const int VALUES_PER_THREAD = 8; + const int BLOCK_SIZE = 256; + const int K_OUT = {K_OUT}; + const int N_IN = {N_IN}; + const int SHARED_N_IN = {SHARED_N_IN}; + const int N_GROUPS = {N_groups}; + const int SHARED_N_GROUPS = {SHARED_N_groups}; + const int N_ACTIVE = {n_active}; + const int TOTAL_ROUTED = {total_routed}; + const int BATCH_SIZE = {B}; + + uint3 tgid = threadgroup_position_in_grid; + uint sgid = simdgroup_index_in_threadgroup; + uint slid = thread_index_in_simdgroup; + + int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG; + if (out_row >= K_OUT) return; + + if (tgid.z < (uint)TOTAL_ROUTED) {{ + // ═══════ ROUTED EXPERT PATH (same as gather_qmv) ═══════ + // tgid.z indexes flat (batch, expert) pairs + int flat_idx = (int)tgid.z; + int expert = inds[flat_idx]; + + 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 bfloat16_t* sc = (const device bfloat16_t*)S + + (long)expert * K_OUT * N_GROUPS + out_row * N_GROUPS + slid / {slid_divisor}; + const device bfloat16_t* bi = (const device bfloat16_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 + + flat_idx * 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 + flat_idx * 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 {{ + // ═══════ SHARED EXPERT PATH (register-level weight sharing) ═══════ + // ONE TG handles ALL {B} batch elements. + // Load shared expert weights once, compute {B} outputs from registers. + + const device uint8_t* ws = (const device uint8_t*)W_shared_down + + (long)out_row * SHARED_N_IN + slid * VALUES_PER_THREAD; + const device bfloat16_t* sc = (const device bfloat16_t*)S_shared_down + + (long)out_row * SHARED_N_GROUPS + slid / {slid_divisor}; + const device bfloat16_t* bi = (const device bfloat16_t*)B_shared_down + + (long)out_row * SHARED_N_GROUPS + slid / {slid_divisor}; + + int x_base = slid * VALUES_PER_THREAD; + {shared_result_decls} + + for (int k = 0; k < SHARED_N_IN; k += BLOCK_SIZE) {{ + // Load x for all {B} batch elements +{shared_x_load} + + // Load weights once into registers, compute all {B} batches + for (int row = 0; row < RESULTS_PER_SG; row++) {{ + const device uint8_t* wl = ws + row * SHARED_N_IN; + float s_val = float(sc[row * SHARED_N_GROUPS]); + float b_val = float(bi[row * SHARED_N_GROUPS]); + + float w_vals[VALUES_PER_THREAD]; + for (int i = 0; i < VALUES_PER_THREAD; i++) {{ + w_vals[i] = float(wl[i]); + }} + +{shared_qdot} + }} + + ws += BLOCK_SIZE; + sc += {sc_stride}; + bi += {sc_stride}; + x_base += BLOCK_SIZE; + }} + + // Write all {B} outputs +{shared_write} + }} +""" + + +_batched_down_cache = {} + + +def _get_batched_down_kernel(K_OUT, N_IN, SHARED_N_IN, n_active, B, group_size=64): + key = (K_OUT, N_IN, SHARED_N_IN, n_active, B, group_size) + if key not in _batched_down_cache: + _batched_down_cache[key] = mx.fast.metal_kernel( + name=f"batched_down_K{K_OUT}_N{N_IN}_SN{SHARED_N_IN}_na{n_active}_B{B}", + input_names=["W", "S", "B_q", + "W_shared_down", "S_shared_down", "B_shared_down", + "X_routed", "X_shared", "inds"], + output_names=["Y_routed", "Y_shared"], + source=_gen_batched_merged_down_8bit_source( + K_OUT, N_IN, SHARED_N_IN, n_active, B, group_size), + ) + return _batched_down_cache[key] + + +def batched_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, batch_size, + n_active, group_size=64, + shared_n_in=None): + """Batched merged down_proj for 8-bit routed + shared experts. + + Args: + w_q: routed weights (E, K_OUT, N_IN/4) uint32 + s: routed scales (E, K_OUT, N_IN/gs) bf16 + b_q: routed biases (E, K_OUT, N_IN/gs) bf16 + w_shared_down: shared weight (K_OUT, SHARED_N_IN/4) uint32 + s_shared_down: shared scales (K_OUT, SHARED_N_IN/gs) bf16 + b_shared_down: shared biases (K_OUT, SHARED_N_IN/gs) bf16 + x_routed: (B * n_active, N_IN) f32 + x_shared: (B * SHARED_N_IN,) f32 + inds: (B * n_active,) uint32 + k_out: output dimension + n_in: routed input dimension + batch_size: B + n_active: experts per token (top_k) + shared_n_in: shared input dim (defaults to n_in) + + Returns: + Y_routed: (B * n_active, k_out) f32 + Y_shared: (B, k_out) f32 + """ + B = batch_size + 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 + + kern = _get_batched_down_kernel(k_out_val, n_in_val, shared_n_in_val, n_active, B) + + y_groups = ceil_div(k_out_val, 8) + total_routed = B * n_active + + Y = kern( + inputs=[w_q, s, b_q, + w_shared_down, s_shared_down, b_shared_down, + x_routed, x_shared, inds], + output_shapes=[(total_routed * k_out_val,), (B * k_out_val,)], + output_dtypes=[mx.float32, mx.float32], + grid=(32, y_groups * 2, total_routed + 1), + threadgroup=(32, 2, 1), + ) + + return Y[0].reshape(total_routed, k_out_val), Y[1].reshape(B, k_out_val) diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_moe_epilogue.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_moe_epilogue.py new file mode 100644 index 00000000..b0ad2f9d --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_moe_epilogue.py @@ -0,0 +1,92 @@ +"""Batched MoE epilogue for Qwen3.5: weighted sum + shared expert gate + residual. + +Computes per batch element: + Y[b, j] = bf16( Σ_a(scores[b,a] * D_routed[b*n_active+a, j]) + + sigmoid(gate_raw[b]) * D_shared[b, j] + + H[b, j] ) + +Grid z = B (one set of threads per batch element). +All constants baked into Metal source. +""" +import mlx.core as mx + + +def _gen_batched_epilogue_source(K, n_active, B): + return f""" + const int K_const = {K}; + const int n_active_const = {n_active}; + const int B_const = {B}; + + uint tid = thread_position_in_grid.x; + uint batch_id = thread_position_in_grid.z; + + if (tid >= K_const || batch_id >= B_const) return; + + // Weighted sum of routed expert outputs for this batch element + float acc = 0.0f; + int routed_base = (int)batch_id * n_active_const * K_const; + int score_base = (int)batch_id * n_active_const; + for (int a = 0; a < n_active_const; a++) {{ + acc += scores[score_base + a] * D_routed[routed_base + a * K_const + tid]; + }} + + // Shared expert: sigmoid(gate_raw) * D_shared + float gate_raw_val = gate_raw[(int)batch_id]; + float gate = 1.0f / (1.0f + metal::exp(-gate_raw_val)); + float shared_val = D_shared[(int)batch_id * K_const + tid] * gate; + + // Add residual and write + Y[(int)batch_id * K_const + tid] = static_cast( + acc + shared_val + float(H[(int)batch_id * K_const + tid]) + ); +""" + + +_batched_epilogue_cache = {} + + +def _get_batched_epilogue_kernel(K, n_active, B): + key = (K, n_active, B) + if key not in _batched_epilogue_cache: + _batched_epilogue_cache[key] = mx.fast.metal_kernel( + name=f"batched_epilogue_K{K}_na{n_active}_B{B}", + input_names=["D_routed", "D_shared", "scores", "H", "gate_raw"], + output_names=["Y"], + source=_gen_batched_epilogue_source(K, n_active, B), + ) + return _batched_epilogue_cache[key] + + +def batched_moe_epilogue(d_routed, d_shared, scores, h, gate_raw, + k_val, batch_size, n_active): + """Batched MoE epilogue with fused sigmoid. + + Args: + d_routed: (B * n_active, K) f32 + d_shared: (B, K) f32 + scores: (B * n_active,) f32 + h: (B, K) bf16 — residual + gate_raw: (B,) f32 — raw shared expert gate (pre-sigmoid) + k_val: hidden dimension + batch_size: B + n_active: experts per token + + Returns: + Y: (B, K) bf16 + """ + K = int(k_val) + B = batch_size + kern = _get_batched_epilogue_kernel(K, n_active, B) + + tg_size = min(K, 1024) + n_tg = (K + tg_size - 1) // tg_size + + Y = kern( + inputs=[d_routed, d_shared.reshape(B * K), scores, h.reshape(B * K), gate_raw], + output_shapes=[(B * K,)], + output_dtypes=[mx.bfloat16], + grid=(n_tg * tg_size, 1, B), + threadgroup=(tg_size, 1, 1), + ) + + return Y[0].reshape(B, K) diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_oproj_gate_gemv_8bit.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_oproj_gate_gemv_8bit.py new file mode 100644 index 00000000..cc6c7433 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_oproj_gate_gemv_8bit.py @@ -0,0 +1,332 @@ +"""Batched Dispatch 1: Fused o_proj (8-bit) + gate GEMV parts + x² partials. + +Adapts custom_oproj_gate_gemv_8bit for batch_size B (1..8). +Register-level weight sharing: each TG loads weights once, computes B outputs. + +Three GEMV regions (same as B=1): + TGs 0..N_OPROJ_TG-1: o_proj GEMV (8-bit) + residual + h_scaled + x² + TGs N_OPROJ_TG..+N_M1_TG-1: M1 × attn_out → gate_part_a (bf16 GEMV) + TGs +N_M1_TG..end: W_fused × residual → gate_part_b (bf16 GEMV) + +All constants baked into Metal source. B is unrolled at code-generation time. +""" +import mlx.core as mx + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def _gen_batched_oproj_source(n_experts, M, K_attn, K_hidden, B, group_size=64, gate_bm=8): + E = int(n_experts) + gs = group_size + oproj_slid_divisor = gs // 8 + oproj_sc_stride = 256 // gs + blockM_gate = gate_bm * 4 + n_m1_tg = ceil_div(E, blockM_gate) + + # Generate unrolled per-batch code for o_proj epilogue + oproj_epilogue = [] + for b in range(B): + oproj_epilogue.append(f""" + float x2_acc{b} = 0.0f; + for (int tm = 0; tm < TM; tm++) {{ + int k = out_row + tm; + float h{b} = result{b}[tm] + float(residual[{b} * M_DIM + k]); + x2_acc{b} += h{b} * h{b}; + h_scaled[{b} * M_DIM + k] = static_cast(h{b} * float(w_rms[k])); + h_out[{b} * M_DIM + k] = static_cast(h{b}); + }}""") + oproj_epilogue_code = "\n".join(oproj_epilogue) + + oproj_x2_write = [] + for b in range(B): + oproj_x2_write.append(f""" + total{b} += tgp_x2[s * {B} + {b}];""") + oproj_x2_sum = "\n".join(oproj_x2_write) + + oproj_x2_final = [] + for b in range(B): + oproj_x2_final.append(f""" + x2_partials[{b} * N_OPROJ_TG_DIM + tg_x] = total{b};""") + oproj_x2_final_code = "\n".join(oproj_x2_final) + + # o_proj K-loop: load weights once, compute B batch elements + oproj_x_load = "\n".join(f""" + float xv{b}[VPT]; float xsum{b} = 0.0f; + for (int i = 0; i < VPT; i++) {{ xv{b}[i] = float(attn_out[{b} * K_ATTN_DIM + xb + i]); xsum{b} += xv{b}[i]; }}""" for b in range(B)) + + oproj_qdot = "\n".join(f""" + acc{b}[row] += s_val * wdot(xv{b}, w_vals) + xsum{b} * b_val;""" for b in range(B)) + + oproj_result_decls = " ".join(f"float acc{b}[TM] = {{0,0,0,0}};" for b in range(B)) + oproj_simd_reduce = "\n".join(f" float result{b}[TM]; for (int tm=0;tm float {{ + float a = 0; + for (int i = 0; i < 8; i++) a += x[i] * w[i]; + return a; + }}; + + const int N_OPROJ_TG = N_OPROJ_TG_DIM; + const int N_M1_TG = {n_m1_tg}; + const int blockM_gate = {blockM_gate}; + + uint tg_x = threadgroup_position_in_grid.x; + uint sgid = simdgroup_index_in_threadgroup; + uint slid = thread_index_in_simdgroup; + + if (tg_x < (uint)N_OPROJ_TG) {{ + // ══════ O_PROJ GEMV (8-bit, register-sharing for {B} batches) ══════ + const int blockM = 32; + const int VPT = 8; + const int BLOCK_SIZE = 256; + + int out_row = int(tg_x) * blockM + int(sgid) * TM; + if (out_row >= M_DIM) return; + out_row = (out_row + TM <= M_DIM) ? out_row : (M_DIM - TM); + + threadgroup float tgp_x2[8 * {B}]; + + {oproj_result_decls} + int K_groups = K_ATTN_DIM / {gs}; + + // Weight pointers (shared across batch) + const device uint8_t* ws = (const device uint8_t*)W_oproj + + (long)out_row * K_ATTN_DIM + slid * VPT; + const device bfloat16_t* sc = (const device bfloat16_t*)S_oproj + + (long)out_row * K_groups + slid / {oproj_slid_divisor}; + const device bfloat16_t* bi = (const device bfloat16_t*)B_oproj + + (long)out_row * K_groups + slid / {oproj_slid_divisor}; + + int xb = slid * VPT; + + for (int k = 0; k < K_ATTN_DIM; k += BLOCK_SIZE) {{ + // Load x for all {B} batches +{oproj_x_load} + + // Load weights once, compute all batches + for (int row = 0; row < TM; row++) {{ + const device uint8_t* wl = ws + row * K_ATTN_DIM; + float s_val = float(sc[row * K_groups]); + float b_val = float(bi[row * K_groups]); + float w_vals[VPT]; + for (int i = 0; i < VPT; i++) w_vals[i] = float(wl[i]); +{oproj_qdot} + }} + + ws += BLOCK_SIZE; sc += {oproj_sc_stride}; bi += {oproj_sc_stride}; + xb += BLOCK_SIZE; + }} + + // simd_sum for all batches +{oproj_simd_reduce} + + // Epilogue: residual add + x² + h_scaled + h_out + if (slid == 0) {{ +{oproj_epilogue_code} +{oproj_tgp_write} + }} + + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sgid == 0 && slid == 0) {{ + {oproj_total_decls} + for (int s = 0; s < 8; s++) {{ +{oproj_x2_sum} + }} +{oproj_x2_final_code} + }} + + }} else if (tg_x < (uint)(N_OPROJ_TG + N_M1_TG)) {{ + // ══════ M1 GEMV (bf16, register-sharing for {B} batches) ══════ + 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); + + {gate_a_decls} + int bn = int(slid) * TN; + int n_iter = K_ATTN_DIM / blockN; + + for (int i = 0; i < n_iter; i++) {{ +{gate_a_x_load} + for (int tm = 0; tm < TM; tm++) {{ + float w_row[TN]; + for (int tn = 0; tn < TN; tn++) w_row[tn] = float(M1[(out_row + tm) * K_ATTN_DIM + bn + tn]); +{gate_a_dot} + }} + bn += blockN; + }} + + for (int tm = 0; tm < TM; tm++) {{ +{gate_a_reduce} + }} + + if (slid == 0) {{ + for (int tm = 0; tm < TM; tm++) {{ + int e = out_row + tm; + if (e < E_CONST) {{ +{gate_a_write} + }} + }} + }} + + }} else {{ + // ══════ W_FUSED GEMV (bf16, register-sharing for {B} batches) ══════ + 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); + + {gate_b_decls} + int bn = int(slid) * TN; + int n_iter = K_HIDDEN_DIM / blockN; + + for (int i = 0; i < n_iter; i++) {{ +{gate_b_x_load} + for (int tm = 0; tm < TM; tm++) {{ + float w_row[TN]; + for (int tn = 0; tn < TN; tn++) w_row[tn] = float(W_fused[(out_row + tm) * K_HIDDEN_DIM + bn + tn]); +{gate_b_dot} + }} + bn += blockN; + }} + + for (int tm = 0; tm < TM; tm++) {{ +{gate_b_reduce} + }} + + if (slid == 0) {{ + for (int tm = 0; tm < TM; tm++) {{ + int e = out_row + tm; + if (e < E_CONST) {{ +{gate_b_write} + }} + }} + }} + }} +""" + + +_batched_oproj_cache = {} + + +def _get_batched_oproj_kernel(n_experts, M, K_attn, K_hidden, B, group_size=64, gate_bm=8): + key = (n_experts, M, K_attn, K_hidden, B, group_size, gate_bm) + if key not in _batched_oproj_cache: + _batched_oproj_cache[key] = mx.fast.metal_kernel( + name=f"batched_oproj_E{n_experts}_M{M}_Ka{K_attn}_Kh{K_hidden}_B{B}", + input_names=[ + "W_oproj", "S_oproj", "B_oproj", + "attn_out", "residual", "w_rms", + "M1", "W_fused", + ], + output_names=["h_scaled", "h_out", "x2_partials", + "gate_part_a", "gate_part_b"], + source=_gen_batched_oproj_source(n_experts, M, K_attn, K_hidden, B, group_size, gate_bm), + ) + return _batched_oproj_cache[key] + + +def batched_oproj_gate_gemv(W_oproj, S_oproj, B_oproj, + attn_out, residual, w_rms, + M1, W_fused, + M, K_attn, batch_size, + n_experts=256, gate_bm=8, + K_hidden=None, group_size=64): + """Batched fused 8-bit o_proj + bf16 gate GEMVs. + + Args: + W_oproj/S_oproj/B_oproj: 8-bit o_proj weights + attn_out: (B, K_attn) bf16 + residual: (B, K) bf16 + w_rms: (K,) bf16 — RMSNorm weight (shared) + M1: (E, K_attn) bf16 (shared) + W_fused: (E, K) bf16 (shared) + M: hidden size + K_attn: attention output dim + batch_size: B + n_experts: E + gate_bm: SGs per gate TG + + Returns: + h_scaled (B, M) bf16, h_out (B, M) bf16, + x2_partials (B, N_TG) f32, gate_part_a (B, E) f32, gate_part_b (B, E) f32 + """ + B = batch_size + M_val = int(M) + K_attn_val = int(K_attn) + K_hidden_val = int(K_hidden) if K_hidden is not None else M_val + + kern = _get_batched_oproj_kernel(n_experts, M_val, K_attn_val, K_hidden_val, B, group_size, gate_bm) + + n_oproj_tg = ceil_div(M_val, 32) + blockM_gate = gate_bm * 4 + 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.reshape(B * K_attn_val), residual.reshape(B * M_val), w_rms, + M1, W_fused], + output_shapes=[ + (B * M_val,), (B * M_val,), + (B * n_oproj_tg,), + (B * n_experts,), (B * n_experts,), + ], + output_dtypes=[mx.bfloat16, mx.bfloat16, mx.float32, mx.float32, mx.float32], + grid=(total_tg * 32, 8, 1), + threadgroup=(32, 8, 1), + ) + + return (results[0].reshape(B, M_val), + results[1].reshape(B, M_val), + results[2].reshape(B, n_oproj_tg), + results[3].reshape(B, n_experts), + results[4].reshape(B, n_experts)) diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_softmax_topk_swiglu_8bit.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_softmax_topk_swiglu_8bit.py new file mode 100644 index 00000000..a75f192d --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/batched_softmax_topk_swiglu_8bit.py @@ -0,0 +1,578 @@ +"""Dispatch 2 (batched): Softmax prologue + 8-bit SwiGLU for Qwen3.5 MoE. + +Combines the prologue from oproj_softmax_topk_swiglu_8bit (B=1) with the +batched body from batched_merged_swiglu_8bit. All constants are baked into +the Metal source at Python code-generation time. + +Grid z-dimension: B * n_active + 1 + - tgid.z < B * n_active: routed expert TGs + batch_id = tgid.z / n_active, local_z = tgid.z % n_active + - tgid.z == B * n_active: shared expert TG (register-level weight sharing + for B batch elements, including shared_expert_gate GEMV) + +Prologue (all TGs): + Phase 1: distributed x2 partial sum -> inv_rms (per batch_id) + Phase 2 (routed TGs only): gate scores -> softmax -> parallel top-k -> + norm_topk_prob -> write out_inds / norm_scores + Phase 3 (shared TG, SG 0): shared_expert_gate 8-bit GEMV with + register-level weight sharing -> gate_raw[B] + +Body (after TG barrier): + Routed: 8-bit gate+up+SwiGLU with h_scaled[batch_id] input + Shared: register-level weight sharing for B batch elements +""" +import mlx.core as mx + + +def ceil_div(a, b): + return (a + b - 1) // b + + +def _gen_batched_softmax_topk_swiglu_source( + N_INTER, SHARED_INTER, K, n_active, B, + n_experts=256, top_k=10, norm_topk=True, group_size=64, + n_oproj_tg=64, +): + """Generate Metal source for batched softmax + top-k + SwiGLU. + + All routed TGs compute the full softmax+topk for their batch_id into TG + memory. Only the TG with local_z==0 writes out_inds/norm_scores to device + memory. Each routed TG reads its own expert from tg_inds[local_z]. + """ + gs = group_size + sc_stride = 256 // gs + slid_divisor = gs // 8 + N_TOTAL = 2 * N_INTER + K_groups = K // gs + SHARED_K_groups = K // gs + total_routed = B * n_active + E = int(n_experts) + K_TOP = int(top_k) + SPT = (E + 63) // 64 + + # ── Shared expert body: unrolled per-batch code ── + shared_x_load = "\n".join(f""" + float x{b}_thread[VALUES_PER_THREAD]; + float xsum{b} = 0; + for (int i = 0; i < VALUES_PER_THREAD; i++) {{ + float xi = float(X[{b} * K_DIM + x_base + i]); + x{b}_thread[i] = xi; + xsum{b} += xi; + }}""" for b in range(B)) + + shared_gate_qdot = "\n".join(f""" + float accum_g{b} = 0; + for (int i = 0; i < VALUES_PER_THREAD; i++) accum_g{b} += x{b}_thread[i] * wg_vals[i]; + gate{b}[row] += sg * accum_g{b} + xsum{b} * bg;""" for b in range(B)) + + shared_up_qdot = "\n".join(f""" + float accum_u{b} = 0; + for (int i = 0; i < VALUES_PER_THREAD; i++) accum_u{b} += x{b}_thread[i] * wu_vals[i]; + up{b}[row] += su * accum_u{b} + xsum{b} * bu;""" for b in range(B)) + + shared_result_decls = "\n ".join( + f"float gate{b}[RESULTS_PER_SG] = {{0,0,0,0}}; float up{b}[RESULTS_PER_SG] = {{0,0,0,0}};" + for b in range(B)) + + shared_write_lines = [] + for b in range(B): + shared_write_lines.append(f""" + for (int row = 0; row < RESULTS_PER_SG; row++) {{ + float g{b} = simd_sum(gate{b}[row]) * inv_rms_{b}; + float u{b} = simd_sum(up{b}[row]) * inv_rms_{b}; + if (slid == 0) {{ + float silu_g{b} = g{b} / (1.0f + metal::exp(-g{b})); + Y_shared[{b} * SHARED_INTER_DIM + out_row + row] = silu_g{b} * u{b}; + }} + }}""") + shared_write = "\n".join(shared_write_lines) + + # inv_rms for all B batches in the shared TG + shared_inv_rms_lines = [] + for b in range(B): + shared_inv_rms_lines.append(f""" + float local_x2_{b} = 0.0f; + for (int i = x2_start; i < x2_end; i++) local_x2_{b} += x2_partials[{b} * N_OPROJ_TG_DIM + i]; + float sg_x2_{b} = simd_sum(local_x2_{b}); + if (slid == 0) tg_x2_sg[sgid] = sg_x2_{b}; + threadgroup_barrier(mem_flags::mem_threadgroup); + float inv_rms_{b} = metal::precise::rsqrt((tg_x2_sg[0] + tg_x2_sg[1]) / (float)K_DIM + 1e-6f);""") + shared_inv_rms_block = "\n".join(shared_inv_rms_lines) + + # Shared expert gate GEMV accumulator declarations + seg_acc_decls = "\n ".join( + f"float seg_gate_acc{b} = 0.0f;" for b in range(B)) + + seg_write = "\n".join( + f" gate_raw[{b}] = seg_gate_acc{b} * inv_rms_{b};" + for b in range(B)) + + return f""" + const int RESULTS_PER_SG = 4; + const int VALUES_PER_THREAD = 8; + const int BLOCK_SIZE = 256; + const int N_INTER_DIM = {N_INTER}; + const int SHARED_INTER_DIM = {SHARED_INTER}; + const int K_DIM = {K}; + const int K_GROUPS = {K_groups}; + const int N_TOTAL = {N_TOTAL}; + const int N_ACTIVE = {n_active}; + const int TOTAL_ROUTED = {total_routed}; + const int BATCH_SIZE = {B}; + const int E_CONST = {E}; + const int K_TOP_CONST = {K_TOP}; + const int SPT = {SPT}; + const int N_OPROJ_TG_DIM = {n_oproj_tg}; + + 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 + + threadgroup float tg_x2_sg[2]; + threadgroup int tg_inds[{K_TOP}]; + threadgroup float tg_selected_scores[{K_TOP}]; + + if (tgid.z < (uint)TOTAL_ROUTED) {{ + // ═══════════════════════════════════════════════════════════════ + // ROUTED TG PROLOGUE + // ═══════════════════════════════════════════════════════════════ + int flat_idx = (int)tgid.z; + int batch_id = flat_idx / N_ACTIVE; + int local_z = flat_idx % N_ACTIVE; + + // Phase 1: distributed x2 sum -> inv_rms for batch_id + int chunk = (N_OPROJ_TG_DIM + 63) / 64; + int x2_start = tid * chunk; + int x2_end = min(x2_start + chunk, N_OPROJ_TG_DIM); + float local_x2 = 0.0f; + for (int i = x2_start; i < x2_end; i++) + local_x2 += x2_partials[batch_id * N_OPROJ_TG_DIM + i]; + float sg_x2_sum = simd_sum(local_x2); + 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_DIM + 1e-6f); + + // Phase 2: ALL routed TGs compute full softmax + top-k for their + // batch_id. Each TG gets its own copy in TG memory (tg_inds, + // tg_selected_scores). This avoids cross-TG communication. + {{ + 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[batch_id * E_CONST + e] + + gate_part_b[batch_id * E_CONST + 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++) {{ + 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; + }} + }} + + 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); + + 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; + + if (tid == winner_tid) {{ + tg_inds[round] = best_e; + tg_selected_scores[round] = best; + 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); + }} + + // Only local_z==0 writes norm_scores + out_inds to device memory + if (local_z == 0 && tid == 0) {{ + float total_score = 0.0f; + for (int a = 0; a < {K_TOP}; a++) total_score += tg_selected_scores[a]; + float inv_total = {"1.0f / total_score" if norm_topk else "1.0f"}; + for (int a = 0; a < {K_TOP}; a++) {{ + norm_scores[batch_id * {K_TOP} + a] = tg_selected_scores[a] * inv_total; + out_inds[batch_id * {K_TOP} + a] = (uint)tg_inds[a]; + }} + }} + }} + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // ═══════════════════════════════════════════════════════════════ + // ROUTED BODY: 8-bit gate+up+SwiGLU + // ═══════════════════════════════════════════════════════════════ + int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG; + if (out_row >= N_INTER_DIM) return; + + int expert = tg_inds[local_z]; + + const device uint8_t* ws_gate = (const device uint8_t*)W + + (long)expert * N_TOTAL * K_DIM + out_row * K_DIM + slid * VALUES_PER_THREAD; + const device bfloat16_t* sc_gate = (const device bfloat16_t*)S + + (long)expert * N_TOTAL * K_GROUPS + out_row * K_GROUPS + slid / {slid_divisor}; + const device bfloat16_t* bi_gate = (const device bfloat16_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_DIM + (out_row + N_INTER_DIM) * K_DIM + slid * VALUES_PER_THREAD; + const device bfloat16_t* sc_up = (const device bfloat16_t*)S + + (long)expert * N_TOTAL * K_GROUPS + (out_row + N_INTER_DIM) * K_GROUPS + slid / {slid_divisor}; + const device bfloat16_t* bi_up = (const device bfloat16_t*)B_q + + (long)expert * N_TOTAL * K_GROUPS + (out_row + N_INTER_DIM) * K_GROUPS + slid / {slid_divisor}; + + int x_base = batch_id * K_DIM + slid * VALUES_PER_THREAD; + + float gate_result[4] = {{0, 0, 0, 0}}; + float up_result[4] = {{0, 0, 0, 0}}; + + for (int k = 0; k < K_DIM; 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 = ws_gate + row * K_DIM; + 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; + + const device uint8_t* wu = ws_up + row * K_DIM; + 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 + flat_idx * N_INTER_DIM + 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 {{ + // ═══════════════════════════════════════════════════════════════ + // SHARED TG (tgid.z == TOTAL_ROUTED) + // ═══════════════════════════════════════════════════════════════ + + // Phase 1: compute inv_rms for ALL B batch elements + int chunk = (N_OPROJ_TG_DIM + 63) / 64; + int x2_start = tid * chunk; + int x2_end = min(x2_start + chunk, N_OPROJ_TG_DIM); +{shared_inv_rms_block} + + // Phase 3: shared_expert_gate 8-bit GEMV (SG 0 only) + // Load W_seg once, compute B dot products via register-level weight sharing + if (sgid == 0) {{ + const int VPT = 8; + const int SEG_BLOCK = 256; // 32 * VPT + + const device uint8_t* seg_w_ptr = (const device uint8_t*)W_seg + + slid * VPT; + const device bfloat16_t* seg_sc = (const device bfloat16_t*)S_seg + + slid / {slid_divisor}; + const device bfloat16_t* seg_bi = (const device bfloat16_t*)B_seg + + slid / {slid_divisor}; + int seg_xb = slid * VPT; + + {seg_acc_decls} + + for (int k = 0; k < K_DIM; k += SEG_BLOCK) {{ + // Load weight block once into registers + float seg_w_regs[VPT]; + for (int i = 0; i < VPT; i++) seg_w_regs[i] = float(seg_w_ptr[i]); + float seg_sc_val = float(*seg_sc); + float seg_bi_val = float(*seg_bi); + + // Compute B dot products from the same weight registers +{chr(10).join(f''' {{ + float xsum{b} = 0.0f, wacc{b} = 0.0f; + for (int i = 0; i < VPT; i++) {{ + float xi = float(X[{b} * K_DIM + seg_xb + i]); + xsum{b} += xi; + wacc{b} += xi * seg_w_regs[i]; + }} + seg_gate_acc{b} += seg_sc_val * wacc{b} + xsum{b} * seg_bi_val; + }}''' for b in range(B))} + + seg_w_ptr += SEG_BLOCK; + seg_sc += {sc_stride}; + seg_bi += {sc_stride}; + seg_xb += SEG_BLOCK; + }} + + // Reduce across SG and write gate_raw[B] +{chr(10).join(f" seg_gate_acc{b} = simd_sum(seg_gate_acc{b});" for b in range(B))} + if (slid == 0) {{ +{seg_write} + }} + }} + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // ═══════════════════════════════════════════════════════════════ + // SHARED BODY: register-level weight sharing for B batch elements + // ═══════════════════════════════════════════════════════════════ + int out_row = tgid.y * 8 + sgid * RESULTS_PER_SG; + if (out_row >= SHARED_INTER_DIM) return; + + const device uint8_t* ws_gate = (const device uint8_t*)W_shared + + (long)out_row * K_DIM + slid * VALUES_PER_THREAD; + const device bfloat16_t* sc_gate = (const device bfloat16_t*)S_shared + + (long)out_row * {SHARED_K_groups} + slid / {slid_divisor}; + const device bfloat16_t* bi_gate = (const device bfloat16_t*)B_shared + + (long)out_row * {SHARED_K_groups} + slid / {slid_divisor}; + + const device uint8_t* ws_up = (const device uint8_t*)W_shared + + (long)(out_row + SHARED_INTER_DIM) * K_DIM + slid * VALUES_PER_THREAD; + const device bfloat16_t* sc_up = (const device bfloat16_t*)S_shared + + (long)(out_row + SHARED_INTER_DIM) * {SHARED_K_groups} + slid / {slid_divisor}; + const device bfloat16_t* bi_up = (const device bfloat16_t*)B_shared + + (long)(out_row + SHARED_INTER_DIM) * {SHARED_K_groups} + slid / {slid_divisor}; + + int x_base = slid * VALUES_PER_THREAD; + {shared_result_decls} + + for (int k = 0; k < K_DIM; k += BLOCK_SIZE) {{ + // Load x for all {B} batch elements +{shared_x_load} + + for (int row = 0; row < RESULTS_PER_SG; row++) {{ + // Load gate weights once into registers + const device uint8_t* wg = ws_gate + row * K_DIM; + float sg = float(sc_gate[row * {SHARED_K_groups}]); + float bg = float(bi_gate[row * {SHARED_K_groups}]); + float wg_vals[VALUES_PER_THREAD]; + for (int i = 0; i < VALUES_PER_THREAD; i++) wg_vals[i] = float(wg[i]); + + // Compute gate for all {B} batches from registers +{shared_gate_qdot} + + // Load up weights once into registers + const device uint8_t* wu = ws_up + row * K_DIM; + float su = float(sc_up[row * {SHARED_K_groups}]); + float bu = float(bi_up[row * {SHARED_K_groups}]); + float wu_vals[VALUES_PER_THREAD]; + for (int i = 0; i < VALUES_PER_THREAD; i++) wu_vals[i] = float(wu[i]); + + // Compute up for all {B} batches from registers +{shared_up_qdot} + }} + + 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; + }} + + // SwiGLU epilogue + write for all {B} batches (with inv_rms) +{shared_write} + }} +""" + +_batched_softmax_topk_swiglu_cache = {} + + +def _get_batched_softmax_topk_swiglu_kernel( + N_INTER, SHARED_INTER, K, n_active, B, + n_experts=256, top_k=10, norm_topk=True, group_size=64, + n_oproj_tg=64, +): + key = (N_INTER, SHARED_INTER, K, n_active, B, n_experts, top_k, norm_topk, group_size, n_oproj_tg) + if key not in _batched_softmax_topk_swiglu_cache: + nt_tag = "_nt" if norm_topk else "" + _batched_softmax_topk_swiglu_cache[key] = mx.fast.metal_kernel( + name=(f"batched_softmax_topk_swiglu_8bit" + f"_NI{N_INTER}_SI{SHARED_INTER}_K{K}" + f"_na{n_active}_B{B}_E{n_experts}_k{top_k}" + f"_gs{group_size}{nt_tag}"), + input_names=[ + "W", "S", "B_q", # routed expert weights + "W_shared", "S_shared", "B_shared", # shared expert weights + "W_seg", "S_seg", "B_seg", # shared_expert_gate weights + "X", # h_scaled (B, K) bf16 + "gate_part_a", "gate_part_b", # (B, E) f32 + "x2_partials", # (B, N_OPROJ_TG) f32 + ], + output_names=["Y_routed", "Y_shared", "out_inds", + "norm_scores", "gate_raw"], + source=_gen_batched_softmax_topk_swiglu_source( + N_INTER, SHARED_INTER, K, n_active, B, + n_experts, top_k, norm_topk, group_size, n_oproj_tg), + ) + return _batched_softmax_topk_swiglu_cache[key] + + +def batched_softmax_topk_swiglu_8bit( + w_gu, s_gu, b_gu, # routed gate+up weights (E, 2*N_INTER, K/4) + w_shared, s_shared, b_shared, # shared gate+up weights (2*SHARED_INTER, K/4) + w_seg, s_seg, b_seg, # shared_expert_gate weights (1, K/4) + h_scaled, # (B, K) bf16 — from Dispatch 1 + gate_part_a, # (B, E) f32 — from Dispatch 1 + gate_part_b, # (B, E) f32 — from Dispatch 1 + x2_partials, # (B, N_OPROJ_TG) f32 — from Dispatch 1 + n_inter, k_hidden, batch_size, n_active, + n_oproj_tg, n_experts=256, + shared_inter=None, group_size=64, +): + """Batched Dispatch 2: softmax + top-k + merged 8-bit SwiGLU with oproj prologue. + + Prologue (per-batch): + Phase 1: distributed x2 -> inv_rms (all TGs, indexed by batch_id) + Phase 2: softmax(gate_scores) -> top-k -> norm_topk_prob (all routed TGs) + Phase 3: shared_expert_gate 8-bit GEMV -> gate_raw[B] (shared TG, SG 0) + + Body: + Routed: 8-bit gate+up+SwiGLU with h_scaled[batch_id] input, inv_rms factored + Shared: register-level weight sharing for B batch elements + + Args: + w_gu: stacked routed weights (E, 2*N_INTER, K/4) uint32 + s_gu: routed scales (E, 2*N_INTER, K/gs) bf16 + b_gu: routed biases (E, 2*N_INTER, K/gs) bf16 + w_shared: shared gate+up stacked (2*SHARED_INTER, K/4) uint32 + s_shared: shared scales (2*SHARED_INTER, K/gs) bf16 + b_shared: shared biases (2*SHARED_INTER, K/gs) bf16 + w_seg/s_seg/b_seg: shared_expert_gate 8-bit weights (1, K/4) uint32 + h_scaled: (B, K) bf16 — h * w_rms from Dispatch 1 + gate_part_a: (B, E) f32 — partial gate scores from Dispatch 1 + gate_part_b: (B, E) f32 — partial gate scores from Dispatch 1 + x2_partials: (B, N_OPROJ_TG) f32 — per-TG x2 sums from Dispatch 1 + n_inter: routed intermediate size + k_hidden: hidden size K + batch_size: B (1..8) + n_active: experts per token (top_k) + n_oproj_tg: number of o_proj TGs (for x2 partial sum reduction) + n_experts: total number of experts E + shared_inter: shared intermediate size (defaults to n_inter) + group_size: quantization group size (default 64) + + Returns: + (Y_routed, Y_shared, out_inds, norm_scores, gate_raw): + Y_routed: (B * n_active, n_inter) f32 + Y_shared: (B, shared_inter) f32 + out_inds: (B * n_active,) uint32 + norm_scores: (B * n_active,) f32 + gate_raw: (B,) f32 — raw shared expert gate values (sigmoid in epilogue) + """ + B = int(batch_size) + n_inter_val = int(n_inter) + shared_inter_val = int(shared_inter) if shared_inter is not None else n_inter_val + k_val = int(k_hidden) + n_active_val = int(n_active) + top_k = n_active_val + E = int(n_experts) + n_oproj_tg_val = int(n_oproj_tg) + + kern = _get_batched_softmax_topk_swiglu_kernel( + n_inter_val, shared_inter_val, k_val, n_active_val, B, + E, top_k, True, int(group_size), n_oproj_tg_val, + ) + + max_inter = max(n_inter_val, shared_inter_val) + total_routed = B * n_active_val + + results = kern( + inputs=[ + w_gu, s_gu, b_gu, + w_shared, s_shared, b_shared, + w_seg, s_seg, b_seg, + h_scaled, + gate_part_a, gate_part_b, + x2_partials, + ], + output_shapes=[ + (total_routed * n_inter_val,), # Y_routed flat + (B * shared_inter_val,), # Y_shared flat + (total_routed,), # out_inds + (total_routed,), # norm_scores + (B,), # 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, total_routed + 1), + threadgroup=(32, 2, 1), + ) + + Y_routed = results[0].reshape(total_routed, n_inter_val) + Y_shared = results[1].reshape(B, shared_inter_val) + out_inds = results[2] + norm_scores = results[3] + gate_raw = results[4] + + return Y_routed, Y_shared, out_inds, norm_scores, gate_raw diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_gdn_projections_8bit.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_gdn_projections_8bit.py new file mode 100644 index 00000000..9ac82695 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_gdn_projections_8bit.py @@ -0,0 +1,288 @@ +"""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(K, N_QKV, N_Z, N_B, N_A, group_size=64): + """Generate Metal source for fused GDN projections with merged weights. + + All constants baked into Metal source (no scalar kernel inputs). + """ + gs = int(group_size) + sc_stride = 256 // gs + slid_div = gs // 8 + N_TOTAL = N_QKV + N_Z + N_B + N_A + K_groups = K // gs + N_QKV_TG = ceil_div(N_QKV, 8) + N_Z_TG = ceil_div(N_Z, 8) + + return f""" + const int RESULTS_PER_SG = 4; + const int VALUES_PER_THREAD = 8; + const int BLOCK_SIZE = 256; + const int GROUP_SIZE = {gs}; + const int SC_STRIDE = {sc_stride}; + const int SLID_DIV = {slid_div}; + const int K = {K}; + const int K_groups = {K_groups}; + const int N_QKV = {N_QKV}; + const int N_Z = {N_Z}; + const int N_B = {N_B}; + const int N_TOTAL = {N_TOTAL}; + const int N_QKV_TG = {N_QKV_TG}; + const int N_Z_TG = {N_Z_TG}; + const int N_B_TG = {ceil_div(N_B, 8)}; + + 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(s1); + conv_state_out[cs_base + 1 * conv_dim + c] = static_cast(s2); + conv_state_out[cs_base + 2 * conv_dim + c] = static_cast(qkv_val); + + qkv_out[b_idx * conv_dim + c] = static_cast(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_cache = {} + + +def _get_fused_gdn_proj_kernel(K, N_QKV, N_Z, N_B, N_A, group_size=64): + key = (K, N_QKV, N_Z, N_B, N_A, group_size) + if key not in _fused_gdn_proj_cache: + _fused_gdn_proj_cache[key] = mx.fast.metal_kernel( + name=f"fused_gdn_proj_K{K}_NQKV{N_QKV}_NZ{N_Z}_NB{N_B}_NA{N_A}", + input_names=[ + "x", + "W_merged", "S_merged", "B_merged", + "conv_state", "conv_w", + "A_log_arr", "dt_bias_arr", + ], + output_names=["qkv_out", "z_silu_out", "b_out", "a_out", "conv_state_out"], + source=_gen_fused_gdn_projections_source(K, N_QKV, N_Z, N_B, N_A, group_size), + ) + return _fused_gdn_proj_cache[key] + + +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 + + N_QKV, N_Z, N_B, N_A = proj_dims + K = x.shape[-1] + + kern = _get_fused_gdn_proj_kernel(K, N_QKV, N_Z, N_B, N_A) + + 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 + + conv_w_flat = conv_weights.reshape(-1, 4) if conv_weights.ndim == 3 else conv_weights + 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, + ], + 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 diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_qk_rmsnorm.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_qk_rmsnorm.py new file mode 100644 index 00000000..11bb73f8 --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_qk_rmsnorm.py @@ -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(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) diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_rms_norm_gated.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_rms_norm_gated.py new file mode 100644 index 00000000..f110f4cc --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/fused_rms_norm_gated.py @@ -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(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) diff --git a/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/gdn_step_precomputed.py b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/gdn_step_precomputed.py new file mode 100644 index 00000000..a8205e4c --- /dev/null +++ b/src/exo/worker/engines/mlx/patches/qwen3_5_moe/kernels/gdn_step_precomputed.py @@ -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(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(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(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], + ) diff --git a/src/exo/worker/engines/mlx/speculative/__init__.py b/src/exo/worker/engines/mlx/speculative/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/exo/worker/engines/mlx/speculative/mtp_batch_generator.py b/src/exo/worker/engines/mlx/speculative/mtp_batch_generator.py new file mode 100644 index 00000000..fa1fe412 --- /dev/null +++ b/src/exo/worker/engines/mlx/speculative/mtp_batch_generator.py @@ -0,0 +1,324 @@ +#!/usr/bin/env python3 +"""MTP Speculative Decoding integrated with mlx_lm's BatchGenerator. + +Subclasses BatchGenerator to add MTP drafting + S>1 verification with +correct GDN state rollback via SpeculativeArraysCache. + +At BS=1: drafts γ tokens with MTP, verifies at S=γ+1, buffers accepted tokens. +At BS>1: falls back to standard BatchGenerator (no speculative). + +Usage: + from mtp_batch_generator import MTPBatchGenerator + gen = MTPBatchGenerator(model, mtp_predictor, gamma=2, ...) + gen.insert([prompt_tokens]) + while True: + responses = gen.next() +""" + +import time + +import mlx.core as mx +from mlx_lm.generate import BatchGenerator, generation_stream + +from .mtp_module import MTPPredictor, speculative_forward, draft_tokens + + +class MTPBatchGenerator(BatchGenerator): + """BatchGenerator with MTP speculative decoding for BS=1.""" + + def __init__( + self, + model, + mtp_predictor: MTPPredictor, + gamma: int = 2, + temp: float = 0.0, + alpha: float = 1.0, + **kwargs, + ): + super().__init__(model, **kwargs) + self.mtp = mtp_predictor + self.gamma = gamma + self.temp = temp + self.alpha = alpha + + self._token_buffer = {} # uid → [(token, logprobs), ...] + self._captured = {} # pre_norm / prompt_pre_norm from norm wrapper + self._mtp_pre_norm = {} # uid → (B, 1, D) pre-norm hidden state + self._mtp_prefilled = set() # uids with MTP cache prefilled + self._request_temp = {} # uid → temperature from request + + self._setup_hidden_capture() + + def _setup_hidden_capture(self): + """Monkey-patch model's final norm to capture pre-norm hidden state. + + Captures: + - pre_norm: hidden states before final RMSNorm (for MTP input) + - prompt_pre_norm: same but only when S>1 (prefill) + """ + inner = getattr(self.model, 'model', None) or self.model.language_model.model + original_norm = inner.norm + captured = self._captured + + class _CapturingNorm: + def __init__(self, orig): + self._orig = orig + self.weight = orig.weight + + def __call__(self, x): + captured['pre_norm'] = x + if x.shape[1] > 1: + captured['prompt_pre_norm'] = x + return self._orig(x) + + def __getattr__(self, name): + return getattr(self._orig, name) + + inner.norm = _CapturingNorm(original_norm) + + def _next(self): + batch = self.active_batch + + # Yield buffered tokens first + if batch is not None and len(batch) == 1: + uid = batch.uids[0] + if uid in self._token_buffer and self._token_buffer[uid]: + return self._yield_buffered(batch, uid) + + # BS=1 speculative path + if (batch is not None + and len(batch) == 1 + and self.gamma > 0 + and len(self.unprocessed_prompts) == 0): + uid = batch.uids[0] + if uid not in self._mtp_prefilled: + return self._first_step_and_prefill(batch, uid) + return self._speculative_next() + + # Standard path (BS>1 or no batch) + responses = super()._next() + if responses and batch is not None and len(batch) == 1: + if 'pre_norm' in self._captured: + uid = batch.uids[0] + self._mtp_pre_norm[uid] = self._captured['pre_norm'][:, -1:, :] + return responses + + def _first_step_and_prefill(self, batch, uid): + """First decode step. MTP cache already prefilled by ExoBatchGenerator.submit().""" + responses = super()._next() + if not responses: + return responses + + # Capture decode pre_norm from this standard step for first speculative cycle + decode_pre_norm = self._captured.get('pre_norm') + if decode_pre_norm is not None: + mx.eval(decode_pre_norm) + self._mtp_pre_norm[uid] = decode_pre_norm[:, -1:, :] + + self._mtp_prefilled.add(uid) + return responses + + def _speculative_next(self): + """Core speculative cycle with correct GDN rollback.""" + tic = time.perf_counter() + batch = self.active_batch + uid = batch.uids[0] + y = batch.y # (1,) — token from previous step, to be yielded + y_val = y[0].item() + y_logprobs = batch.logprobs[0] + + # Append current y to token history + batch.tokens[0] = mx.concatenate((batch.tokens[0], y[0:1])) + + pre_norm = self._mtp_pre_norm.get(uid) + if pre_norm is None: + return super()._next() + + gamma = self.gamma + temp = self._request_temp.get(uid, self.temp) + alpha = self.alpha + + # 1. Draft γ tokens (lazy chain, no eval) + next_token_arr = y.reshape(1, 1) + draft_ids, draft_probs = draft_tokens( + self.mtp, pre_norm, next_token_arr, gamma, temp) + + # 2. Verify via speculative_forward (handles GDN cache wrapping + kernel swap) + draft_concat = mx.concatenate( + [d.reshape(1, 1) for d in draft_ids], axis=1) # (1, γ) + verify_input = mx.concatenate( + [next_token_arr, draft_concat], axis=1) # (1, γ+1) + verify_pre_norm, verify_logits = speculative_forward( + self.model, verify_input, batch.cache, speculative=True) + + # 3. Build acceptance check lazily + target_tokens = mx.argmax(verify_logits[:, :gamma, :], axis=-1) + + if temp == 0: + matches = mx.equal(target_tokens, draft_concat).squeeze(0) + all_next = mx.argmax(verify_logits[0], axis=-1) + logprobs_all = verify_logits[0] - mx.logsumexp( + verify_logits[0], axis=-1, keepdims=True) + mx.async_eval(matches, all_next, logprobs_all, verify_pre_norm) + else: + accept_ratios = [] + for i in range(gamma): + p = mx.softmax(verify_logits[0, i] / temp, axis=-1) + q = draft_probs[i] + p_di = p[draft_ids[i].squeeze()] + q_di = q[0, draft_ids[i].squeeze()] + ratio = p_di / mx.maximum(q_di, 1e-10) + accept_ratios.append(mx.minimum(ratio ** alpha, 1.0)) + uniforms = mx.random.uniform(shape=(gamma,)) + corrections = [] + for i in range(gamma): + p = mx.softmax(verify_logits[0, i] / temp, axis=-1) + q = draft_probs[i][0] + residual = mx.maximum(p - q, 0.0) + corrections.append(mx.random.categorical(mx.log(residual + 1e-10))) + bonus_token = mx.random.categorical(verify_logits[0, gamma] * (1.0 / temp)) + logprobs_all = verify_logits[0] - mx.logsumexp( + verify_logits[0], axis=-1, keepdims=True) + mx.async_eval(accept_ratios, uniforms, corrections, bonus_token, + logprobs_all, verify_pre_norm, draft_concat) + + # 4. Determine acceptance + n_accepted = 0 + for i in range(gamma): + if temp == 0: + if matches[i].item(): + n_accepted += 1 + else: + break + else: + if uniforms[i].item() < accept_ratios[i].item(): + n_accepted += 1 + else: + break + + # 5. Rollback cache + rollback = gamma - n_accepted + if rollback > 0: + for c in batch.cache: + if hasattr(c, 'offset'): + c.offset -= rollback + elif hasattr(c, 'rollback'): + c.rollback(n_accepted) + + # Unwrap SpeculativeArraysCache + for i, c in enumerate(batch.cache): + if hasattr(c, 'base'): + batch.cache[i] = c.base + + # 6. Bonus/correction token + logprobs + if n_accepted == gamma: + if temp == 0: + bonus_val = all_next[gamma].item() + else: + bonus_val = bonus_token.item() + bonus_lp = logprobs_all[gamma] + else: + if temp == 0: + bonus_val = all_next[n_accepted].item() + else: + bonus_val = corrections[n_accepted].item() + bonus_lp = logprobs_all[n_accepted] + + # 7. Update MTP pre_norm for next cycle + self._mtp_pre_norm[uid] = verify_pre_norm[ + :, (gamma if n_accepted == gamma else n_accepted): + (gamma if n_accepted == gamma else n_accepted) + 1, :] + + # 8. Build token list: current y + accepted drafts + draft_int_values = draft_concat[0].tolist() + all_tokens = [(y_val, y_logprobs)] + for i in range(n_accepted): + all_tokens.append((draft_int_values[i], logprobs_all[i])) + + # 9. Set batch.y = bonus for next cycle + batch.y = mx.array([bonus_val]) + batch.logprobs = [bonus_lp] + + # Append accepted drafts to token history + if n_accepted > 0: + batch.tokens[0] = mx.concatenate( + (batch.tokens[0], mx.array([t for t, _ in all_tokens[1:]]))) + batch.num_tokens[0] += len(all_tokens) + + # 10. Check stop conditions — truncate at stop token + toc = time.perf_counter() + self._stats.generation_time += toc - tic + self._stats.generation_tokens += len(all_tokens) + + # Find first stop token or length limit in all_tokens + stop_idx = None + for idx, (tok, _) in enumerate(all_tokens): + if tok in self.stop_tokens: + stop_idx = idx + break + if batch.num_tokens[0] >= batch.max_tokens[0]: + stop_idx = idx + break + + first_tok, first_lp = all_tokens[0] + + if stop_idx is not None: + # Tokens before the stop are valid output — buffer them + # The stop token itself triggers finish_reason + valid_tokens = all_tokens[:stop_idx] + if valid_tokens: + # Yield first, buffer rest + a final stop entry + if len(valid_tokens) > 1: + self._token_buffer[uid] = valid_tokens[1:] + # Append stop marker as last buffered token + stop_tok, stop_lp = all_tokens[stop_idx] + if uid not in self._token_buffer: + self._token_buffer[uid] = [] + self._token_buffer[uid].append((stop_tok, stop_lp)) + mx.async_eval(batch.y) + return [self.Response(uid, first_tok, first_lp, None, lambda: None)] + else: + # Stop token is the first token — finish immediately + cache = batch.extract_cache(0) + self.active_batch = None + self._cleanup_uid(uid) + return [self.Response(uid, first_tok, first_lp, "stop", cache)] + + # Buffer remaining tokens + if len(all_tokens) > 1: + self._token_buffer[uid] = all_tokens[1:] + + mx.async_eval(batch.y) + return [self.Response(uid, first_tok, first_lp, None, lambda: None)] + + def _yield_buffered(self, batch, uid): + """Yield one buffered token from a previous speculative cycle.""" + tic = time.perf_counter() + buf = self._token_buffer[uid] + tok, lp = buf.pop(0) + + if not buf: + del self._token_buffer[uid] + + finish_reason = None + if tok in self.stop_tokens: + finish_reason = "stop" + elif batch.num_tokens[0] >= batch.max_tokens[0]: + finish_reason = "length" + + cache = None + if finish_reason: + cache = batch.extract_cache(0) + self.active_batch = None + self._cleanup_uid(uid) + + toc = time.perf_counter() + self._stats.generation_time += toc - tic + return [self.Response(uid, tok, lp, finish_reason, cache or (lambda: None))] + + def _cleanup_uid(self, uid): + """Clean up MTP state for a finished request.""" + self._mtp_pre_norm.pop(uid, None) + self._mtp_prefilled.discard(uid) + self._token_buffer.pop(uid, None) + self._request_temp.pop(uid, None) diff --git a/src/exo/worker/engines/mlx/speculative/mtp_module.py b/src/exo/worker/engines/mlx/speculative/mtp_module.py new file mode 100644 index 00000000..071a1a87 --- /dev/null +++ b/src/exo/worker/engines/mlx/speculative/mtp_module.py @@ -0,0 +1,515 @@ +#!/usr/bin/env python3 +"""MTP (Multi-Token Prediction) module for Qwen3.5-27B. + +Architecture (from llama.cpp build_mtp_head + HuggingFace config): + 1. Normalize: pre_fc_norm_hidden(hidden_state) || pre_fc_norm_embedding(embed(token)) + 2. Combine: fc(concat([e_norm, h_norm])) → 5120 + 3. 1 GQA decoder layer (same config as main model's full-attention layers) + - Attention with Q/K RMSNorm + partial RoPE + output gate + 4. Final norm → shared lm_head → vocab logits + +Predicts token t+2 given the main model's hidden state at position t +and the token sampled at position t+1. + +Usage: + from .mtp_module import MTPPredictor + mtp = MTPPredictor(model, "mtp_weights.safetensors") + # During decode: + pre_norm, normed = mtp.get_hidden_state(input_tokens, cache) + logits_t1 = mtp.apply_lm_head(normed) # token t+1 + logits_t2 = mtp.predict(pre_norm, token_t1) # token t+2 +""" + +import mlx.core as mx +import mlx.nn as nn + + +def speculative_forward(model, inputs, cache, speculative=False): + """Run model forward pass, optionally capturing GDN per-step states for rollback. + + This is the shared core for both MTP and draft-model speculative decoding. + It manually iterates model layers to: + 1. Wrap GDN caches in SpeculativeArraysCache when speculative=True + 2. Patch gated_delta_update to use the speculative kernel + 3. Capture per-step recurrent states and reconstruct conv_input + + Args: + model: the loaded model (e.g. from mlx_lm.load) + inputs: (B, S) int token ids + cache: cache list from make_prompt_cache() + speculative: if True, saves per-step GDN states for rollback + + Returns: + (pre_norm, logits) — pre-RMSNorm hidden states and vocab logits + """ + inner = getattr(model, 'model', None) or model.language_model.model + text_model = getattr(model, 'model', None) or model.language_model + S = inputs.shape[1] + do_spec = speculative and S > 1 + + if hasattr(inner, 'embed_tokens'): + hidden_states = inner.embed_tokens(inputs) + else: + hidden_states = inputs + + cache_list = cache if cache is not None else [None] * len(inner.layers) + + gdn_spec_data = [] + if do_spec: + from .speculative_cache import SpeculativeArraysCache + for i, c in enumerate(cache_list): + if c is not None and hasattr(c, 'cache') and not hasattr(c, 'offset'): + cache_list[i] = SpeculativeArraysCache(c, S=S) + if cache is not None: + for i in range(len(cache)): + cache[i] = cache_list[i] + + spec_all_states = [] + if do_spec: + import mlx_lm.models.qwen3_5 as _qwen3_5_mod + _orig_gdu = _qwen3_5_mod.gated_delta_update + _qwen3_5_mod.gated_delta_update = _make_speculative_gdu(spec_all_states) + + from mlx_lm.models.qwen3_5 import create_attention_mask, create_ssm_mask + fa_mask = create_attention_mask(hidden_states, cache_list[inner.fa_idx]) + ssm_mask = create_ssm_mask(hidden_states, cache_list[inner.ssm_idx]) + + for layer, c in zip(inner.layers, cache_list): + mask = ssm_mask if layer.is_linear else fa_mask + + if do_spec and layer.is_linear: + from .speculative_cache import SpeculativeArraysCache as _SAC + if isinstance(c, _SAC): + pre_conv = c[0] + if pre_conv is None: + gdn = layer.linear_attn + pre_conv = mx.zeros( + (hidden_states.shape[0], gdn.conv_kernel_size - 1, + gdn.conv_dim), dtype=hidden_states.dtype) + gdn_spec_data.append((hidden_states, pre_conv, c, layer)) + + hidden_states = layer(hidden_states, mask=mask, cache=c) + + if do_spec: + _qwen3_5_mod.gated_delta_update = _orig_gdu + + gdn_idx = 0 + for layer_input, pre_conv, spec_cache, parent_layer in gdn_spec_data: + if gdn_idx < len(spec_all_states): + spec_cache.all_states = spec_all_states[gdn_idx] + gdn_idx += 1 + + gdn = parent_layer.linear_attn + normed = parent_layer.input_layernorm(layer_input) + if hasattr(gdn, 'in_proj_qkv'): + qkv = gdn.in_proj_qkv(normed) + else: + q, k, v, z, b, a = gdn.fix_query_key_value_ordering( + gdn.in_proj_qkvz(normed), gdn.in_proj_ba(normed)) + B_dim = normed.shape[0] + qkv = mx.concatenate( + [q.reshape(B_dim, S, -1), k.reshape(B_dim, S, -1), + v.reshape(B_dim, S, -1)], axis=-1) + spec_cache.conv_input = mx.concatenate([pre_conv, qkv], axis=1) + + pre_norm = hidden_states + normed = inner.norm(hidden_states) + + if hasattr(text_model, 'lm_head'): + logits = text_model.lm_head(normed) + else: + logits = inner.embed_tokens.as_linear(normed) + + return pre_norm, logits + + +def _make_speculative_gdu(all_states_list): + """Create a gated_delta_update replacement that uses the speculative kernel. + + The speculative kernel is identical to the original but also outputs + per-step recurrent states (all_states). These are appended to + all_states_list for later assignment to SpeculativeArraysCache wrappers. + + Returns (y, state_out) — same interface as original gated_delta_update. + """ + from .speculative_gdn_kernel import speculative_gated_delta_kernel + from mlx_lm.models.gated_delta import compute_g + + def speculative_gated_delta_update(q, k, v, a, b, A_log, dt_bias, + state=None, mask=None, use_kernel=True): + beta = mx.sigmoid(b) + g = compute_g(A_log, a, dt_bias) + if state is None: + B, _, Hk, Dk = q.shape + Hv, Dv = v.shape[-2:] + state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype) + y, state_out, all_states = speculative_gated_delta_kernel( + q, k, v, g, beta, state, mask) + all_states_list.append(all_states) + return y, state_out + + return speculative_gated_delta_update + + +class MTPPredictor: + """MTP draft predictor for speculative decoding. + + Wraps the MTP module weights and provides: + - get_hidden_state(): extract pre-lm_head hidden state from main model + - predict(): run MTP to get next-next-token logits + """ + + def __init__(self, model, mtp_weights_path, quantize=True): + """Load MTP weights and attach to the main model. + + Args: + model: loaded Qwen3.5-27B model + mtp_weights_path: path to mtp_weights.safetensors + quantize: quantize MTP linears to 8-bit gs=64 + """ + self.model = model + self._inner = getattr(model, 'model', None) or model.language_model.model + self._text_model = getattr(model, 'model', None) or model.language_model + + # Shared components + self.embed_tokens = self._inner.embed_tokens + if hasattr(self._text_model, 'lm_head'): + self.lm_head = self._text_model.lm_head + else: + # tie_word_embeddings case + self.lm_head = None + + # Load MTP weights + weights = mx.load(mtp_weights_path) + + # ---- Sanitize norm weights ---- + # CRITICAL: Qwen3.5 HuggingFace format stores ALL norm weights as (actual - 1.0). + # mlx-lm's TextModel.sanitize() adds +1.0 back for the main model norms, but + # MTP weights are stripped before sanitize runs. We must apply the same shift + # to ALL 1-D norm weights in the MTP. + # + # Evidence: pre_fc_norm_hidden has mean=-0.17 raw → 0.83 after shift (plausible). + # Linear projection weights (2-D) are NOT shifted. + shifted = [] + for k in list(weights.keys()): + if weights[k].ndim == 1: + weights[k] = weights[k] + 1.0 + shifted.append(k) + if shifted: + print(f" Sanitized {len(shifted)} norm weights (+1.0 shift)") + + # Infer all dimensions from weight shapes (works for any Qwen3.5 size) + fc_w = weights['mtp.fc.weight'] + hidden_size = fc_w.shape[0] # 4096 (9B) or 5120 (27B) + fc_in = fc_w.shape[1] # 2 * hidden_size + + q_w = weights['mtp.layers.0.self_attn.q_proj.weight'] + q_out = q_w.shape[0] # num_heads * head_dim * 2 (gate) + k_w = weights['mtp.layers.0.self_attn.k_proj.weight'] + kv_out = k_w.shape[0] # num_kv_heads * head_dim + o_w = weights['mtp.layers.0.self_attn.o_proj.weight'] + o_in = o_w.shape[1] # num_heads * head_dim + + # Detect MoE vs dense MLP + self.is_moe = 'mtp.layers.0.mlp.gate.weight' in weights + + if not self.is_moe: + gate_w = weights['mtp.layers.0.mlp.gate_proj.weight'] + intermediate = gate_w.shape[0] + else: + intermediate = 0 # MoE experts handle this + + # head_dim from q_norm weight (always per-head) + head_dim = weights.get('mtp.layers.0.self_attn.q_norm.weight', + mx.ones(256)).shape[0] + num_heads = o_in // head_dim + num_kv_heads = kv_out // head_dim + + print(f" Dims: hidden={hidden_size}, heads={num_heads}, kv_heads={num_kv_heads}, " + f"head_dim={head_dim}, MLP={'MoE' if self.is_moe else f'dense({intermediate})'}") + + # Build layers from weights — all dimension-agnostic + def make_linear(w): + out_dim, in_dim = w.shape + l = nn.Linear(in_dim, out_dim, bias=False) + l.weight = w + return l + + self.pre_fc_norm_hidden = nn.RMSNorm(hidden_size) + self.pre_fc_norm_hidden.weight = weights['mtp.pre_fc_norm_hidden.weight'] + + self.pre_fc_norm_embedding = nn.RMSNorm(hidden_size) + self.pre_fc_norm_embedding.weight = weights['mtp.pre_fc_norm_embedding.weight'] + + self.fc = make_linear(fc_w) + self.q_proj = make_linear(q_w) + self.k_proj = make_linear(k_w) + self.v_proj = make_linear(weights['mtp.layers.0.self_attn.v_proj.weight']) + self.o_proj = make_linear(o_w) + + self.q_norm = nn.RMSNorm(head_dim) + self.k_norm = nn.RMSNorm(head_dim) + q_norm_key = 'mtp.layers.0.self_attn.q_norm.weight' + k_norm_key = 'mtp.layers.0.self_attn.k_norm.weight' + if q_norm_key in weights: + self.q_norm.weight = weights[q_norm_key] + self.k_norm.weight = weights[k_norm_key] + + self.input_layernorm = nn.RMSNorm(hidden_size) + self.input_layernorm.weight = weights['mtp.layers.0.input_layernorm.weight'] + + self.post_attention_layernorm = nn.RMSNorm(hidden_size) + self.post_attention_layernorm.weight = weights['mtp.layers.0.post_attention_layernorm.weight'] + + if self.is_moe: + # Reuse mlx-lm's SparseMoeBlock from the target model + moe_layer = None + for layer in self._inner.layers: + if hasattr(layer, 'mlp') and hasattr(layer.mlp, 'gate'): + moe_layer = layer.mlp + break + if moe_layer is None: + raise RuntimeError("MTP has MoE weights but target model has no MoE layer") + + # Create a new MoE block with same class/config as target + moe_class = type(moe_layer) + args = getattr(self._text_model, 'args', None) + if args is None and hasattr(self._text_model, 'model'): + args = getattr(self._text_model.model, 'args', None) + self.mlp = moe_class(args) + + # Load MTP MoE weights — remap HF expert names to mlx-lm SwitchLinear + prefix = 'mtp.layers.0.mlp.' + # Direct weights: gate, shared_expert, shared_expert_gate + direct_keys = {} + expert_weights = {} # {proj_name: {expert_idx: weight}} + for k, v in weights.items(): + if not k.startswith(prefix): + continue + name = k[len(prefix):] + # Check if it's an individual expert weight + if name.startswith('experts.'): + # experts.N.{gate,up,down}_proj.weight → stack into switch_mlp + parts = name.split('.') + idx = int(parts[1]) + proj = parts[2] # gate_proj, up_proj, down_proj + key = f'{proj}.{parts[3]}' # gate_proj.weight + if key not in expert_weights: + expert_weights[key] = {} + expert_weights[key][idx] = v + else: + direct_keys[name] = v + + # Stack individual expert weights into SwitchLinear format + moe_weights = [] + for proj_key, idx_map in expert_weights.items(): + n_experts = max(idx_map.keys()) + 1 + stacked = mx.stack([idx_map[i] for i in range(n_experts)]) + moe_weights.append((f'switch_mlp.{proj_key}', stacked)) + + # Add direct weights + for name, v in direct_keys.items(): + moe_weights.append((name, v)) + + self.mlp.load_weights(moe_weights) + print(f" MoE MLP: {len(moe_weights)} weight groups loaded " + f"({len(expert_weights)} stacked expert projections)") + else: + self.gate_proj = make_linear(gate_w) + self.up_proj = make_linear(weights['mtp.layers.0.mlp.up_proj.weight']) + self.down_proj = make_linear(weights['mtp.layers.0.mlp.down_proj.weight']) + + self.norm = nn.RMSNorm(hidden_size) + self.norm.weight = weights['mtp.norm.weight'] + + # RoPE from main model's GQA layers + for layer in self._inner.layers: + if not layer.is_linear: + self.rope = layer.self_attn.rope + break + + # GQA config + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_dim = head_dim + self.scale = head_dim ** -0.5 + + # MTP KV cache (separate from main model) + self.kv_cache = None + + mx.eval(self.pre_fc_norm_hidden.weight, self.pre_fc_norm_embedding.weight, + self.fc.weight, self.input_layernorm.weight, + self.post_attention_layernorm.weight, self.norm.weight, + self.q_norm.weight, self.k_norm.weight) + + if quantize: + self._quantize_linears() + + total_params = sum(w.size for w in weights.values()) + print(f" MTP loaded: {len(weights)} tensors, {total_params / 1e6:.1f}M params" + f"{' (quantized 8-bit gs=64)' if quantize else ' (bf16)'}") + + def _quantize_linears(self): + """Quantize all MTP linear layers to 8-bit gs=64.""" + for name in ['fc', 'q_proj', 'k_proj', 'v_proj', 'o_proj', + 'gate_proj', 'up_proj', 'down_proj']: + linear = getattr(self, name) + linear.weight = linear.weight.astype(mx.bfloat16) + q = nn.QuantizedLinear.from_linear(linear, group_size=64, bits=8) + mx.eval(q.parameters()) + setattr(self, name, q) + + def reset_cache(self): + """Reset the MTP KV cache (call at start of generation).""" + from mlx_lm.models.cache import KVCache + self.kv_cache = KVCache() + + def get_hidden_state(self, inputs, cache, speculative=False): + """Run main model and return pre-norm hidden states + logits. + + Delegates to the shared speculative_forward() function. + """ + return speculative_forward(self.model, inputs, cache, speculative) + + def _attn_mlp(self, h): + """Run GQA attention + MLP. Shared by predict, predict_hidden, predict_from_hidden.""" + B, S = h.shape[0], h.shape[1] + + residual = h + h = self.input_layernorm(h) + + q_out = self.q_proj(h) + q_out, gate = mx.split( + q_out.reshape(B, S, self.num_heads, -1), 2, axis=-1 + ) + gate = gate.reshape(B, S, -1) + + queries = self.q_norm(q_out).transpose(0, 2, 1, 3) + keys = self.k_norm( + self.k_proj(h).reshape(B, S, self.num_kv_heads, self.head_dim) + ).transpose(0, 2, 1, 3) + values = self.v_proj(h).reshape( + B, S, self.num_kv_heads, self.head_dim + ).transpose(0, 2, 1, 3) + + if self.kv_cache is not None: + offset = self.kv_cache.offset + queries = self.rope(queries, offset=offset) + keys = self.rope(keys, offset=offset) + keys, values = self.kv_cache.update_and_fetch(keys, values) + else: + queries = self.rope(queries) + keys = self.rope(keys) + + mask = None + if S > 1: + total_kv = keys.shape[2] + q_pos = mx.arange(S) + (total_kv - S) + k_pos = mx.arange(total_kv) + mask = mx.where(k_pos[None, :] <= q_pos[:, None], + mx.array(0, dtype=queries.dtype), + mx.array(-1e9, dtype=queries.dtype)) + + output = mx.fast.scaled_dot_product_attention( + queries, keys, values, scale=self.scale, mask=mask + ) + output = output.transpose(0, 2, 1, 3).reshape(B, S, -1) + + h = residual + self.o_proj(output * mx.sigmoid(gate)) + + residual = h + h = self.post_attention_layernorm(h) + if self.is_moe: + h = residual + self.mlp(h) + else: + h = residual + self.down_proj(nn.silu(self.gate_proj(h)) * self.up_proj(h)) + + return h # post-FFN, pre-norm + + def _combine(self, hidden_state, token_ids): + """Combine hidden state + token embedding → fc input.""" + B, S = hidden_state.shape[0], hidden_state.shape[1] + embed = self.embed_tokens(token_ids.reshape(B, S)) + h_norm = self.pre_fc_norm_hidden(hidden_state) + e_norm = self.pre_fc_norm_embedding(embed) + return self.fc(mx.concatenate([e_norm, h_norm], axis=-1)) + + def predict(self, hidden_state, token_ids, return_hidden=False, draft_mode=False): + """Predict next-next-token logits using MTP. + + Args: + hidden_state: (B, S, D) bf16 — PRE-NORM hidden states + token_ids: (B, S) or (S,) int — tokens at each position + return_hidden: if True, also return pre-norm hidden for chaining + draft_mode: if True, use truncated lm_head (32K vocab) for speed + Returns: + logits: (B, S, vocab_size) if S>1, (B, vocab_size) if S=1 + If return_hidden: (logits, hidden) + """ + S = hidden_state.shape[1] + h = self._combine(hidden_state, token_ids) + pre_norm_out = self._attn_mlp(h) + + normed = self.norm(pre_norm_out) + if draft_mode: + logits = normed @ self.draft_lm_head_weight.T + elif self.lm_head is not None: + logits = self.lm_head(normed) + else: + logits = self.embed_tokens.as_linear(normed) + + if S == 1: + logits = logits.squeeze(1) + + if return_hidden: + return logits, pre_norm_out + return logits + + def predict_hidden(self, hidden_state, token_ids): + """Like predict() but returns only post-FFN hidden state (no lm_head).""" + h = self._combine(hidden_state, token_ids) + return self._attn_mlp(h) + + def predict_from_hidden(self, prev_hidden): + """MTP step using post_norm of prev_hidden instead of token embedding. + + Replaces embed_tokens + pre_fc_norm_embedding with just norm(prev_hidden). + This skips the lm_head → argmax → embed_tokens roundtrip. + """ + post_norm = self.norm(prev_hidden) + h_norm = self.pre_fc_norm_hidden(prev_hidden) + h = self.fc(mx.concatenate([post_norm, h_norm], axis=-1)) + return self._attn_mlp(h) + + +def draft_tokens(mtp_pred, hidden, first_token_arr, gamma, temp, fast_lm_head=False): + """Draft γ tokens by chaining MTP predictions — fully lazy, no mx.eval. + + The entire chain stays in the MLX computation graph. Draft token ids + are lazy mx.arrays (argmax/categorical results), not Python ints. + + Args: + first_token_arr: mx.array of shape (1,1) — the token to start from + Returns: (draft_ids, draft_probs) where draft_ids[i] is a lazy mx.array + scalar, draft_probs[i] is the full draft distribution (or None if greedy) + """ + draft_ids = [] + draft_probs = [] + h = hidden + tok_arr = first_token_arr + + for i in range(gamma): + logits, h = mtp_pred.predict(h, tok_arr, return_hidden=True, + draft_mode=fast_lm_head) + + if temp == 0: + tok_arr = mx.argmax(logits, axis=-1).reshape(1, 1) + draft_ids.append(tok_arr.reshape(-1)) + draft_probs.append(None) + else: + q = mx.softmax(logits / temp, axis=-1) + tok_arr = mx.random.categorical(logits * (1.0 / temp)).reshape(1, 1) + draft_ids.append(tok_arr.reshape(-1)) + draft_probs.append(q) + + return draft_ids, draft_probs diff --git a/src/exo/worker/engines/mlx/speculative/speculative_cache.py b/src/exo/worker/engines/mlx/speculative/speculative_cache.py new file mode 100644 index 00000000..57e63717 --- /dev/null +++ b/src/exo/worker/engines/mlx/speculative/speculative_cache.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +"""SpeculativeArraysCache — wraps ArraysCache for correct GDN rollback. + +During speculative verification (S>1), captures: +- all_states: per-step recurrent states from the speculative kernel +- conv_input: full conv_input tensor for conv state rollback + +On rejection, rollback(n_accepted) restores both recurrent and conv +state to the correct intermediate position. +""" + +import mlx.core as mx + + +class SpeculativeArraysCache: + """Wrapper around ArraysCache that supports rollback for speculative decode. + + Delegates all normal cache operations to the underlying ArraysCache. + Adds all_states/conv_input storage and a rollback() method. + """ + + def __init__(self, base_cache, S, conv_kernel_size=4): + self.base = base_cache + self._S = S + self.n_keep = conv_kernel_size - 1 # typically 3 + self.all_states = None # [B, T, Hv, Dv, Dk] from speculative kernel + self.conv_input = None # [B, n_keep+S, conv_dim] for conv rollback + + # Delegate cache operations + def __getitem__(self, idx): + return self.base[idx] + + def __setitem__(self, idx, val): + self.base[idx] = val + + @property + def cache(self): + return self.base.cache + + @cache.setter + def cache(self, v): + self.base.cache = v + + @property + def state(self): + return self.base.state + + @state.setter + def state(self, v): + self.base.state = v + + @property + def lengths(self): + return self.base.lengths + + @lengths.setter + def lengths(self, v): + self.base.lengths = v + + @property + def left_padding(self): + return self.base.left_padding + + @left_padding.setter + def left_padding(self, v): + self.base.left_padding = v + + def advance(self, N): + self.base.advance(N) + + def make_mask(self, N): + return self.base.make_mask(N) + + def empty(self): + return self.base.empty() + + @property + def nbytes(self): + return self.base.nbytes + + def rollback(self, n_accepted): + """Roll back to state after processing n_accepted+1 tokens. + + Args: + n_accepted: number of accepted draft tokens (0 = all rejected, + only the first token 'y' was processed correctly) + """ + # Recurrent state: restore intermediate state at accepted position + if self.all_states is not None: + self.base.cache[1] = self.all_states[0, n_accepted] + + # Conv state: slice conv_input to the correct window + if self.conv_input is not None: + self.base.cache[0] = self.conv_input[:, n_accepted + 1: n_accepted + 1 + self.n_keep, :] + + # Clear stored states + self.all_states = None + self.conv_input = None diff --git a/src/exo/worker/engines/mlx/speculative/speculative_gdn_kernel.py b/src/exo/worker/engines/mlx/speculative/speculative_gdn_kernel.py new file mode 100644 index 00000000..8efbebe1 --- /dev/null +++ b/src/exo/worker/engines/mlx/speculative/speculative_gdn_kernel.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python3 +"""Speculative variant of the GatedDeltaNet kernel. + +Identical to the original gated_delta_kernel but also outputs per-step +recurrent states for rollback during speculative decoding. + +The original kernel only writes the FINAL state. This variant writes +the state at EVERY timestep to an extra output buffer `all_states`. +""" + +from typing import Optional, Tuple +import mlx.core as mx + + +def _make_speculative_gated_delta_kernel(has_mask=False, vectorized=False): + mask_source = "mask[b_idx * T + t]" if has_mask else "true" + + if vectorized: + g_comment = "// g: [B, T, Hv, Dk]" + g_setup = "auto g_ = g + (b_idx * T * Hv + hv_idx) * Dk;" + g_access = "g_[s_idx]" + g_advance = "g_ += Hv * Dk;" + else: + g_comment = "// g: [B, T, Hv]" + g_setup = "auto g_ = g + b_idx * T * Hv;" + g_access = "g_[hv_idx]" + g_advance = "g_ += Hv;" + + 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; + + // all_states: [B, T, Hv, Dv, Dk] — per-step state output + auto a_state = all_states + (b_idx * T * Hv * Dv + hv_idx * Dv + dv_idx) * Dk; + auto a_stride = Hv * Dv * 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(i_state[s_idx]); + }} + + {g_comment} + {g_setup} + auto beta_ = beta + b_idx * T * Hv; + + for (int t = 0; t < T; ++t) {{ + if ({mask_source}) {{ + 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_access}; + kv_mem += state[i] * k_[s_idx]; + }} + kv_mem = simd_sum(kv_mem); + + auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx]; + + 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(out); + }} + }} + + // Save per-step state for speculative rollback + for (int i = 0; i < n_per_t; ++i) {{ + auto s_idx = n_per_t * dk_idx + i; + a_state[s_idx] = static_cast(state[i]); + }} + a_state += a_stride; + + q_ += Hk * Dk; + k_ += Hk * Dk; + v_ += Hv * Dv; + y += Hv * Dv; + {g_advance} + beta_ += Hv; + }} + // Write final state (same as original kernel) + for (int i = 0; i < n_per_t; ++i) {{ + auto s_idx = n_per_t * dk_idx + i; + o_state[s_idx] = static_cast(state[i]); + }} + """ + inputs = ["q", "k", "v", "g", "beta", "state_in", "T"] + if has_mask: + inputs.append("mask") + + suffix = "_spec" + if vectorized: + suffix += "_vec" + if has_mask: + suffix += "_mask" + + return mx.fast.metal_kernel( + name=f"gated_delta_step{suffix}", + input_names=inputs, + output_names=["y", "state_out", "all_states"], + source=source, + ) + + +# Pre-build kernel variants +_spec_kernel = _make_speculative_gated_delta_kernel(has_mask=False, vectorized=False) +_spec_kernel_masked = _make_speculative_gated_delta_kernel(has_mask=True, vectorized=False) +_spec_kernel_vec = _make_speculative_gated_delta_kernel(has_mask=False, vectorized=True) +_spec_kernel_vec_masked = _make_speculative_gated_delta_kernel(has_mask=True, vectorized=True) + + +def speculative_gated_delta_kernel( + 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, mx.array]: + """Like gated_delta_kernel but also returns per-step states. + + Returns: + y: [B, T, Hv, Dv] — output (same as original) + state_out: [B, Hv, Dv, Dk] — final state (same as original) + all_states: [B, T, Hv, Dv, Dk] — state after each timestep + """ + B, T, Hk, Dk = k.shape + Hv, Dv = v.shape[2:] + input_type = q.dtype + + if g.ndim == 4: + kernel = _spec_kernel_vec + inputs = [q, k, v, g, beta, state, T] + if mask is not None: + kernel = _spec_kernel_vec_masked + inputs.append(mask) + else: + kernel = _spec_kernel + inputs = [q, k, v, g, beta, state, T] + if mask is not None: + kernel = _spec_kernel_masked + 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, (B, T, Hv, Dv, Dk)], + output_dtypes=[input_type, input_type, input_type], + ) diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 790dcd8d..ad60124b 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -224,6 +224,10 @@ def load_mlx_items( else: vision_processor = None + # Apply model-specific kernel fusion patches (Qwen3.5 MoE, Qwen3.5 dense LpB, etc.) + from exo.worker.engines.mlx.patches import maybe_apply_patches + maybe_apply_patches(cast(Model, model), model_path if group is None else build_model_path(bound_instance.bound_shard.model_card.model_id)) + return cast(Model, model), tokenizer, vision_processor diff --git a/src/exo/worker/runner/llm_inference/batch_generator.py b/src/exo/worker/runner/llm_inference/batch_generator.py index 5284d53a..74259773 100644 --- a/src/exo/worker/runner/llm_inference/batch_generator.py +++ b/src/exo/worker/runner/llm_inference/batch_generator.py @@ -358,6 +358,7 @@ class BatchGenerator(InferenceGenerator): group=self.group, model_id=self.model_id, ) + self._mlx_gen.warmup_speculative(self.model, self.tokenizer) def submit( self,