MTP speculative decoding + kernel patches for Qwen3.5 (rebased on main)
Speculative decoding: - speculative/: MTPPredictor, MTPBatchGenerator, SpeculativeArraysCache, speculative GDN kernel with per-step state output - Correct GDN state + conv state rollback on draft rejection - Proper rejection sampling with residual distribution correction - Lazy drafting + single async_eval per cycle - MTP prefill in submit() with correct token alignment - Speculative warmup at startup - EXO_SPECULATIVE=1, EXO_MTP_WEIGHTS, EXO_SPECULATIVE_GAMMA, EXO_SPECULATIVE_TEMP, EXO_SPECULATIVE_ALPHA env vars - EXO_DISABLE_LOGPROBS=1 to skip per-token logprobs extraction Kernel patches: - qwen3_5_moe/: batched fused kernels for 35B-A3B MoE decode - qwen3_5/: LpB GEMV kernels for 27B/9B dense models - Auto-detected by model_type, EXO_FUSED_KERNELS=0 to disable - maybe_apply_patches() called after model load in utils_mlx.py Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
dmcc73
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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<bfloat16_t>(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)
|
||||
@@ -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
|
||||
@@ -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"
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
+265
@@ -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<bfloat16_t>(s1);
|
||||
conv_state_out[cs_base + 1 * conv_dim + c] = static_cast<bfloat16_t>(s2);
|
||||
conv_state_out[cs_base + 2 * conv_dim + c] = static_cast<bfloat16_t>(qkv_val);
|
||||
qkv_out[{b} * conv_dim + c] = static_cast<bfloat16_t>(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
|
||||
+205
@@ -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<bfloat16_t>(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<bfloat16_t>(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<bfloat16_t>(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))
|
||||
+252
@@ -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)
|
||||
@@ -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<bfloat16_t>(
|
||||
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)
|
||||
+332
@@ -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<bfloat16_t>(h{b} * float(w_rms[k]));
|
||||
h_out[{b} * M_DIM + k] = static_cast<bfloat16_t>(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<TM;tm++) result{b}[tm] = simd_sum(acc{b}[tm]);" for b in range(B))
|
||||
|
||||
oproj_tgp_write = "\n".join(f" tgp_x2[sgid * {B} + {b}] = x2_acc{b};" for b in range(B))
|
||||
oproj_total_decls = " ".join(f"float total{b} = 0.0f;" for b in range(B))
|
||||
|
||||
# Gate M1 GEMV: load M1 weights once, compute B dot products with B attn_outs
|
||||
gate_a_x_load = "\n".join(f"""
|
||||
float v{b}[TN];
|
||||
for (int tn = 0; tn < TN; tn++) v{b}[tn] = float(attn_out[{b} * K_ATTN_DIM + bn + tn]);""" for b in range(B))
|
||||
|
||||
gate_a_dot = "\n".join(f"""
|
||||
float gacc{b} = 0.0f;
|
||||
for (int tn = 0; tn < TN; tn++) gacc{b} += w_row[tn] * v{b}[tn];
|
||||
gresult{b}[tm] += gacc{b};""" for b in range(B))
|
||||
|
||||
gate_a_decls = " ".join(f"float gresult{b}[TM] = {{0,0,0,0}};" for b in range(B))
|
||||
gate_a_reduce = "\n".join(f" gresult{b}[tm] = simd_sum(gresult{b}[tm]);" for b in range(B))
|
||||
gate_a_write = "\n".join(f"""
|
||||
gate_part_a[{b} * E_CONST + e] = gresult{b}[tm];""" for b in range(B))
|
||||
|
||||
# Gate W_fused GEMV: same pattern but with residual input
|
||||
gate_b_x_load = "\n".join(f"""
|
||||
float rv{b}[TN];
|
||||
for (int tn = 0; tn < TN; tn++) rv{b}[tn] = float(residual[{b} * K_HIDDEN_DIM + bn + tn]);""" for b in range(B))
|
||||
|
||||
gate_b_dot = "\n".join(f"""
|
||||
float wdot{b} = 0.0f;
|
||||
for (int tn = 0; tn < TN; tn++) wdot{b} += w_row[tn] * rv{b}[tn];
|
||||
bresult{b}[tm] += wdot{b};""" for b in range(B))
|
||||
|
||||
gate_b_decls = " ".join(f"float bresult{b}[TM] = {{0,0,0,0}};" for b in range(B))
|
||||
gate_b_reduce = "\n".join(f" bresult{b}[tm] = simd_sum(bresult{b}[tm]);" for b in range(B))
|
||||
gate_b_write = "\n".join(f"""
|
||||
gate_part_b[{b} * E_CONST + e] = bresult{b}[tm];""" for b in range(B))
|
||||
|
||||
return f"""
|
||||
const int TM = 4;
|
||||
const int TN = 4;
|
||||
const int blockN = 128;
|
||||
const int E_CONST = {E};
|
||||
const int M_DIM = {M};
|
||||
const int K_ATTN_DIM = {K_attn};
|
||||
const int K_HIDDEN_DIM = {K_hidden};
|
||||
const int N_OPROJ_TG_DIM = {ceil_div(M, 32)};
|
||||
const int BATCH_SIZE = {B};
|
||||
|
||||
// Helper: dot product of x_thread and w_vals (8 elements)
|
||||
auto wdot = [](thread float* x, thread float* w) -> 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))
|
||||
+578
@@ -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
|
||||
@@ -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<bfloat16_t>(s1);
|
||||
conv_state_out[cs_base + 1 * conv_dim + c] = static_cast<bfloat16_t>(s2);
|
||||
conv_state_out[cs_base + 2 * conv_dim + c] = static_cast<bfloat16_t>(qkv_val);
|
||||
|
||||
qkv_out[b_idx * conv_dim + c] = static_cast<bfloat16_t>(silu_out);
|
||||
}}
|
||||
|
||||
}} else if (region == 1) {{
|
||||
// ═══ Z: SiLU → write f32 ═══
|
||||
int z_row = out_row - N_QKV + (int)slid;
|
||||
if (slid < (uint)RESULTS_PER_SG && z_row < N_Z) {{
|
||||
float val = result[slid];
|
||||
float silu_val = val / (1.0f + metal::exp(-val));
|
||||
z_silu_out[b_idx * N_Z + z_row] = silu_val;
|
||||
}}
|
||||
|
||||
}} else if (region == 2) {{
|
||||
// ═══ B: sigmoid(result) → beta (f32) ═══
|
||||
int b_row = out_row - N_QKV - N_Z + (int)slid;
|
||||
if (slid < (uint)RESULTS_PER_SG && b_row < N_B) {{
|
||||
float val = result[slid];
|
||||
float beta = 1.0f / (1.0f + metal::exp(-val));
|
||||
b_out[b_idx * N_B + b_row] = beta;
|
||||
}}
|
||||
|
||||
}} else {{
|
||||
// ═══ A: g = exp(-exp(A_log) * softplus(a + dt_bias)) → f32 ═══
|
||||
int a_row = out_row - N_QKV - N_Z - N_B + (int)slid;
|
||||
int N_A = N_TOTAL - N_QKV - N_Z - N_B;
|
||||
if (slid < (uint)RESULTS_PER_SG && a_row < N_A) {{
|
||||
float a_val = result[slid];
|
||||
float dt = float(dt_bias_arr[a_row]);
|
||||
float x_g = a_val + dt;
|
||||
// softplus(x) = log(1 + exp(x)), with x>20 shortcut for numerical stability
|
||||
float sp = (x_g > 20.0f) ? x_g : metal::log(1.0f + metal::exp(x_g));
|
||||
float g_val = metal::exp(-metal::exp(float(A_log_arr[a_row])) * sp);
|
||||
a_out[b_idx * N_A + a_row] = g_val;
|
||||
}}
|
||||
}}
|
||||
"""
|
||||
|
||||
|
||||
_fused_gdn_proj_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
|
||||
@@ -0,0 +1,128 @@
|
||||
"""Fused Q/K per-head L2-norm for GDN attention (Dispatch 3).
|
||||
|
||||
Performs per-head L2 normalization on q and k vectors with different scaling.
|
||||
Matches vLLM and latest mlx-lm (qwen3_5.py) which use rsqrt(sum(x²) + eps),
|
||||
NOT rms_norm which uses rsqrt(mean(x²) + eps).
|
||||
|
||||
From qwen3_5.py (updated to match vLLM):
|
||||
inv_scale = Dk^(-0.5) = 128^(-0.5)
|
||||
q = inv_scale * q * rsqrt(sum(q²) + 1e-6) → L2-normalize then scale by 1/√Dk
|
||||
k = k * rsqrt(sum(k²) + 1e-6) → L2-normalize only (no extra scale)
|
||||
|
||||
Grid: (32 heads × 32 threads, 1, B).
|
||||
Each TG = 32 threads = 1 SG, handles one 128-dim head.
|
||||
Dk=128 = 32 threads × 4 elements → exactly 1 SG, no cross-SG reduction.
|
||||
"""
|
||||
|
||||
import mlx.core as mx
|
||||
|
||||
|
||||
def _gen_fused_qk_rmsnorm_source():
|
||||
"""Generate Metal source for fused Q/K per-head L2-norm.
|
||||
|
||||
Input: qkv [B, 8192] bf16 (flattened from [B, 1, 8192])
|
||||
- [0, 2048): q = 16 heads × 128
|
||||
- [2048, 4096): k = 16 heads × 128
|
||||
- [4096, 8192): v (untouched)
|
||||
|
||||
Output: qk_out [B, 4096] bf16
|
||||
- [0, 2048): q L2-normalized then scaled by 1/√Dk
|
||||
- [2048, 4096): k L2-normalized (no extra scale)
|
||||
|
||||
Grid: (32 * 32, 1, B), TG: (32, 1, 1)
|
||||
tgid.x 0..15: q heads → scale = 1/√128
|
||||
tgid.x 16..31: k heads → scale = 1.0
|
||||
tgid.z: batch index
|
||||
"""
|
||||
return """
|
||||
const int N_READS = 4;
|
||||
const int DK = 128;
|
||||
const int HK = 16;
|
||||
const float EPS = 1e-6f;
|
||||
const float Q_SCALE = rsqrt(128.0f); // inv_scale = Dk^(-0.5)
|
||||
const float K_SCALE = 1.0f; // no extra scale for k
|
||||
|
||||
uint head_idx = threadgroup_position_in_grid.x;
|
||||
uint slid = thread_index_in_simdgroup;
|
||||
uint b_idx = thread_position_in_grid.z;
|
||||
|
||||
bool is_q = (head_idx < (uint)HK);
|
||||
|
||||
// Input offset: q heads at [0, 2048), k heads at [2048, 4096)
|
||||
int in_base = is_q
|
||||
? (b_idx * 8192 + head_idx * DK)
|
||||
: (b_idx * 8192 + 2048 + (head_idx - HK) * DK);
|
||||
|
||||
// Output offset: q at [0, 2048), k at [2048, 4096)
|
||||
int out_base = b_idx * 4096 + head_idx * DK;
|
||||
|
||||
// ── Phase 1: Load 4 elements + sum of squares ──
|
||||
float vals[4];
|
||||
float partial_sq = 0.0f;
|
||||
int elem_base = slid * N_READS;
|
||||
|
||||
for (int i = 0; i < N_READS; i++) {
|
||||
float xi = float(qkv[in_base + elem_base + i]);
|
||||
vals[i] = xi;
|
||||
partial_sq += xi * xi;
|
||||
}
|
||||
|
||||
// ── Phase 2: simd reduction (32 threads → full sum of 128 elements) ──
|
||||
float sum_sq = simd_sum(partial_sq);
|
||||
|
||||
// ── Phase 3: compute L2 inv-norm (NOT rms_norm — no /Dk) ──
|
||||
float inv_rms = metal::precise::rsqrt(sum_sq + EPS);
|
||||
|
||||
// ── Phase 4: scale and write ──
|
||||
float scale = is_q ? Q_SCALE : K_SCALE;
|
||||
float combined = inv_rms * scale;
|
||||
|
||||
for (int i = 0; i < N_READS; i++) {
|
||||
qk_out[out_base + elem_base + i] = static_cast<bfloat16_t>(vals[i] * combined);
|
||||
}
|
||||
"""
|
||||
|
||||
|
||||
_fused_qk_rmsnorm_kernel = None
|
||||
|
||||
|
||||
def _get_fused_qk_rmsnorm_kernel():
|
||||
"""Get or compile the fused Q/K RMSNorm kernel."""
|
||||
global _fused_qk_rmsnorm_kernel
|
||||
if _fused_qk_rmsnorm_kernel is None:
|
||||
_fused_qk_rmsnorm_kernel = mx.fast.metal_kernel(
|
||||
name="fused_qk_rmsnorm",
|
||||
input_names=["qkv"],
|
||||
output_names=["qk_out"],
|
||||
source=_gen_fused_qk_rmsnorm_source(),
|
||||
)
|
||||
return _fused_qk_rmsnorm_kernel
|
||||
|
||||
|
||||
def fused_qk_rmsnorm(qkv_conv_silu, batch_size=1):
|
||||
"""Fused Q/K per-head RMSNorm for GDN attention.
|
||||
|
||||
Args:
|
||||
qkv_conv_silu: [B, 1, 8192] bf16 — post-conv, post-SiLU output from Dispatch 2.
|
||||
First 2048 = q (16 heads × 128), next 2048 = k, last 4096 = v.
|
||||
batch_size: int — batch dimension.
|
||||
|
||||
Returns:
|
||||
qk_normed: [B, 1, 4096] bf16 — normalized q (first 2048) and k (next 2048).
|
||||
v is NOT copied; Dispatch 4 reads v directly from qkv_conv_silu[:, :, 4096:].
|
||||
"""
|
||||
B = batch_size
|
||||
kern = _get_fused_qk_rmsnorm_kernel()
|
||||
|
||||
# Flatten to [B, 8192] for kernel
|
||||
qkv_flat = qkv_conv_silu.reshape(B, 8192)
|
||||
|
||||
n_heads = 32 # 16 q + 16 k
|
||||
results = kern(
|
||||
inputs=[qkv_flat],
|
||||
output_shapes=[(B * 4096,)],
|
||||
output_dtypes=[mx.bfloat16],
|
||||
grid=(n_heads * 32, 1, B),
|
||||
threadgroup=(32, 1, 1),
|
||||
)
|
||||
return results[0].reshape(B, 1, 4096)
|
||||
@@ -0,0 +1,117 @@
|
||||
"""Fused RMSNormGated for GDN attention (Dispatch 5).
|
||||
|
||||
Fuses RMSNorm(out, weight) × z_silu into one kernel.
|
||||
SiLU on z was already applied in Dispatch 2, so z_silu arrives as f32.
|
||||
|
||||
From qwen3_next.py (Qwen3NextRMSNormGated):
|
||||
x = rms_norm(hidden_states, weight, eps) # weight: [Dv=128]
|
||||
gate = silu(z.float()) # already done in Dispatch 2
|
||||
return (gate * x).to(hidden_states.dtype)
|
||||
|
||||
Grid: (32 heads × 32 threads, 1, B).
|
||||
Each TG = 32 threads = 1 SG, handles one 128-dim head.
|
||||
Dv=128 = 32 threads × 4 elements → exactly 1 SG.
|
||||
"""
|
||||
|
||||
import mlx.core as mx
|
||||
|
||||
|
||||
def _gen_fused_rms_norm_gated_source():
|
||||
"""Generate Metal source for fused RMSNormGated.
|
||||
|
||||
Inputs:
|
||||
gdn_out: [B, Hv*Dv] bf16 — GDN output, flattened (Hv=32, Dv=128)
|
||||
z_silu: [B, Hv*Dv] f32 — post-SiLU z from Dispatch 2
|
||||
weight: [Dv] f32 — RMSNormGated learned weight (128 elements)
|
||||
|
||||
Output:
|
||||
out: [B, Hv*Dv] bf16 — result = z_silu * rms_norm(gdn_out, weight)
|
||||
|
||||
Grid: (32 * 32, 1, B), TG: (32, 1, 1)
|
||||
tgid.x: head index (0..31)
|
||||
tgid.z: batch index
|
||||
"""
|
||||
return """
|
||||
const int N_READS = 4;
|
||||
const int DV = 128;
|
||||
const int HV = 32;
|
||||
const float EPS = 1e-6f;
|
||||
|
||||
uint head_idx = threadgroup_position_in_grid.x;
|
||||
uint slid = thread_index_in_simdgroup;
|
||||
uint b_idx = thread_position_in_grid.z;
|
||||
|
||||
int base = b_idx * HV * DV + head_idx * DV;
|
||||
int elem_base = slid * N_READS;
|
||||
|
||||
// ── Phase 1: Load gdn_out elements + sum of squares ──
|
||||
float gdn_vals[4];
|
||||
float partial_sq = 0.0f;
|
||||
|
||||
for (int i = 0; i < N_READS; i++) {
|
||||
float xi = float(gdn_out[base + elem_base + i]);
|
||||
gdn_vals[i] = xi;
|
||||
partial_sq += xi * xi;
|
||||
}
|
||||
|
||||
// ── Phase 2: simd reduction (32 threads → full sum of 128 elements) ──
|
||||
float sum_sq = simd_sum(partial_sq);
|
||||
|
||||
// ── Phase 3: compute inv_rms ──
|
||||
float inv_rms = metal::precise::rsqrt(sum_sq / float(DV) + EPS);
|
||||
|
||||
// ── Phase 4: RMSNorm × z_silu, write bf16 ──
|
||||
for (int i = 0; i < N_READS; i++) {
|
||||
int idx = elem_base + i;
|
||||
float w = float(weight[idx]); // learned weight[Dv]
|
||||
float normed = gdn_vals[i] * inv_rms * w; // RMSNorm
|
||||
float z_val = z_silu[base + idx]; // already f32, post-SiLU
|
||||
out[base + idx] = static_cast<bfloat16_t>(z_val * normed);
|
||||
}
|
||||
"""
|
||||
|
||||
|
||||
_fused_rms_norm_gated_kernel = None
|
||||
|
||||
|
||||
def _get_fused_rms_norm_gated_kernel():
|
||||
"""Get or compile the fused RMSNormGated kernel."""
|
||||
global _fused_rms_norm_gated_kernel
|
||||
if _fused_rms_norm_gated_kernel is None:
|
||||
_fused_rms_norm_gated_kernel = mx.fast.metal_kernel(
|
||||
name="fused_rms_norm_gated",
|
||||
input_names=["gdn_out", "z_silu", "weight"],
|
||||
output_names=["out"],
|
||||
source=_gen_fused_rms_norm_gated_source(),
|
||||
)
|
||||
return _fused_rms_norm_gated_kernel
|
||||
|
||||
|
||||
def fused_rms_norm_gated(gdn_out, z_silu, weight, batch_size=1):
|
||||
"""Fused RMSNormGated: RMSNorm(out, weight) × z_silu.
|
||||
|
||||
Args:
|
||||
gdn_out: [B, 1, Hv, Dv] bf16 — GDN recurrence output (Hv=32, Dv=128).
|
||||
z_silu: [B, 1, 4096] f32 — post-SiLU z from Dispatch 2.
|
||||
weight: [128] f32 — RMSNormGated learned weight (Dv elements).
|
||||
batch_size: int.
|
||||
|
||||
Returns:
|
||||
out: [B, 1, 4096] bf16 — ready for out_proj in Dispatch 6.
|
||||
"""
|
||||
B = batch_size
|
||||
kern = _get_fused_rms_norm_gated_kernel()
|
||||
|
||||
# Flatten to [B, 4096]
|
||||
gdn_flat = gdn_out.reshape(B, 4096)
|
||||
z_flat = z_silu.reshape(B, 4096)
|
||||
|
||||
n_heads = 32 # Hv
|
||||
results = kern(
|
||||
inputs=[gdn_flat, z_flat, weight],
|
||||
output_shapes=[(B * 4096,)],
|
||||
output_dtypes=[mx.bfloat16],
|
||||
grid=(n_heads * 32, 1, B),
|
||||
threadgroup=(32, 1, 1),
|
||||
)
|
||||
return results[0].reshape(B, 1, 4096)
|
||||
@@ -0,0 +1,177 @@
|
||||
"""GDN recurrence with pre-computed g and beta (Dispatch 4).
|
||||
|
||||
Modified version of gated_delta_step from mlx-lm-fork/mlx_lm/models/gated_delta.py.
|
||||
Instead of computing g = exp(-exp(A_log) * softplus(a + dt_bias)) and beta = sigmoid(b)
|
||||
inside the kernel, accepts them as pre-computed f32 inputs from Dispatch 2.
|
||||
|
||||
Non-vectorized only (Qwen3.5-35B-A3B uses scalar gating per head).
|
||||
|
||||
Grid: (32, Dv, B*Hv) = (32, 128, B*32), TG: (32, 4, 1)
|
||||
"""
|
||||
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import mlx.core as mx
|
||||
|
||||
|
||||
def _make_gdn_precomputed_kernel(has_mask=False):
|
||||
"""Build the GDN kernel with pre-computed g and beta."""
|
||||
if not mx.metal.is_available():
|
||||
return None
|
||||
|
||||
mask_source = "mask[b_idx * T + t]" if has_mask else "true"
|
||||
|
||||
source = f"""
|
||||
auto n = thread_position_in_grid.z;
|
||||
auto b_idx = n / Hv;
|
||||
auto hv_idx = n % Hv;
|
||||
auto hk_idx = hv_idx / (Hv / Hk);
|
||||
constexpr int n_per_t = Dk / 32;
|
||||
|
||||
// q, k: [B, T, Hk, Dk]
|
||||
auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk;
|
||||
auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
|
||||
|
||||
// v, y: [B, T, Hv, Dv]
|
||||
auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
|
||||
y += b_idx * T * Hv * Dv + hv_idx * Dv;
|
||||
|
||||
auto dk_idx = thread_position_in_threadgroup.x;
|
||||
auto dv_idx = thread_position_in_grid.y;
|
||||
|
||||
// state_in, state_out: [B, Hv, Dv, Dk]
|
||||
auto i_state = state_in + (n * Dv + dv_idx) * Dk;
|
||||
auto o_state = state_out + (n * Dv + dv_idx) * Dk;
|
||||
|
||||
float state[n_per_t];
|
||||
for (int i = 0; i < n_per_t; ++i) {{
|
||||
auto s_idx = n_per_t * dk_idx + i;
|
||||
state[i] = static_cast<float>(i_state[s_idx]);
|
||||
}}
|
||||
|
||||
// g: [B, T, Hv] f32 — pre-computed decay gate
|
||||
auto g_ = g + b_idx * T * Hv;
|
||||
// beta: [B, T, Hv] f32 — pre-computed sigmoid(b)
|
||||
auto beta_ = beta + b_idx * T * Hv;
|
||||
|
||||
for (int t = 0; t < T; ++t) {{
|
||||
if ({mask_source}) {{
|
||||
// Pre-computed g and beta (no softplus/exp/sigmoid needed)
|
||||
float g_val = g_[hv_idx];
|
||||
float beta_val = beta_[hv_idx];
|
||||
|
||||
float kv_mem = 0.0f;
|
||||
for (int i = 0; i < n_per_t; ++i) {{
|
||||
auto s_idx = n_per_t * dk_idx + i;
|
||||
state[i] = state[i] * g_val;
|
||||
kv_mem += state[i] * k_[s_idx];
|
||||
}}
|
||||
kv_mem = simd_sum(kv_mem);
|
||||
|
||||
auto delta = (v_[dv_idx] - kv_mem) * beta_val;
|
||||
|
||||
float out = 0.0f;
|
||||
for (int i = 0; i < n_per_t; ++i) {{
|
||||
auto s_idx = n_per_t * dk_idx + i;
|
||||
state[i] = state[i] + k_[s_idx] * delta;
|
||||
out += state[i] * q_[s_idx];
|
||||
}}
|
||||
out = simd_sum(out);
|
||||
if (thread_index_in_simdgroup == 0) {{
|
||||
y[dv_idx] = static_cast<InT>(out);
|
||||
}}
|
||||
}}
|
||||
// Increment data pointers to next time step
|
||||
q_ += Hk * Dk;
|
||||
k_ += Hk * Dk;
|
||||
v_ += Hv * Dv;
|
||||
y += Hv * Dv;
|
||||
g_ += Hv;
|
||||
beta_ += Hv;
|
||||
}}
|
||||
for (int i = 0; i < n_per_t; ++i) {{
|
||||
auto s_idx = n_per_t * dk_idx + i;
|
||||
o_state[s_idx] = static_cast<InT>(state[i]);
|
||||
}}
|
||||
"""
|
||||
|
||||
inputs = ["q", "k", "v", "g", "beta", "state_in", "T"]
|
||||
if has_mask:
|
||||
inputs.append("mask")
|
||||
|
||||
suffix = "_precomputed"
|
||||
if has_mask:
|
||||
suffix += "_mask"
|
||||
|
||||
return mx.fast.metal_kernel(
|
||||
name=f"gated_delta_step{suffix}",
|
||||
input_names=inputs,
|
||||
output_names=["y", "state_out"],
|
||||
source=source,
|
||||
)
|
||||
|
||||
|
||||
_gdn_precomputed_kernel = None
|
||||
_gdn_precomputed_kernel_masked = None
|
||||
|
||||
|
||||
def _get_gdn_precomputed_kernel(has_mask=False):
|
||||
"""Get or compile the pre-computed GDN kernel."""
|
||||
global _gdn_precomputed_kernel, _gdn_precomputed_kernel_masked
|
||||
if has_mask:
|
||||
if _gdn_precomputed_kernel_masked is None:
|
||||
_gdn_precomputed_kernel_masked = _make_gdn_precomputed_kernel(has_mask=True)
|
||||
return _gdn_precomputed_kernel_masked
|
||||
else:
|
||||
if _gdn_precomputed_kernel is None:
|
||||
_gdn_precomputed_kernel = _make_gdn_precomputed_kernel(has_mask=False)
|
||||
return _gdn_precomputed_kernel
|
||||
|
||||
|
||||
def gated_delta_update_precomputed(
|
||||
q: mx.array,
|
||||
k: mx.array,
|
||||
v: mx.array,
|
||||
g: mx.array,
|
||||
beta: mx.array,
|
||||
state: mx.array,
|
||||
mask: Optional[mx.array] = None,
|
||||
) -> Tuple[mx.array, mx.array]:
|
||||
"""GDN recurrence with pre-computed g and beta.
|
||||
|
||||
Args:
|
||||
q: [B, T, Hk, Dk] bf16 — normalized q from Dispatch 3
|
||||
k: [B, T, Hk, Dk] bf16 — normalized k from Dispatch 3
|
||||
v: [B, T, Hv, Dv] bf16 — v from Dispatch 2 (qkv_conv_silu[:, :, 4096:])
|
||||
g: [B, T, Hv] f32 — pre-computed decay gate from Dispatch 2
|
||||
beta: [B, T, Hv] f32 — pre-computed sigmoid(b) from Dispatch 2
|
||||
state: [B, Hv, Dv, Dk] bf16 — recurrent state from cache
|
||||
mask: [B, T] optional
|
||||
|
||||
Returns:
|
||||
y: [B, T, Hv, Dv] bf16
|
||||
new_state: [B, Hv, Dv, Dk] bf16
|
||||
"""
|
||||
B, T, Hk, Dk = k.shape
|
||||
Hv, Dv = v.shape[2:]
|
||||
input_type = q.dtype
|
||||
|
||||
kernel = _get_gdn_precomputed_kernel(has_mask=mask is not None)
|
||||
inputs = [q, k, v, g, beta, state, T]
|
||||
if mask is not None:
|
||||
inputs.append(mask)
|
||||
|
||||
return kernel(
|
||||
inputs=inputs,
|
||||
template=[
|
||||
("InT", input_type),
|
||||
("Dk", Dk),
|
||||
("Dv", Dv),
|
||||
("Hk", Hk),
|
||||
("Hv", Hv),
|
||||
],
|
||||
grid=(32, Dv, B * Hv),
|
||||
threadgroup=(32, 4, 1),
|
||||
output_shapes=[(B, T, Hv, Dv), state.shape],
|
||||
output_dtypes=[input_type, input_type],
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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<float>(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<InT>(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<InT>(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<InT>(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],
|
||||
)
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user