Add MTP speculative decoding for Qwen3.5 models
Integrates MTP-based speculative decoding into exo's BatchGenerator. When enabled via EXO_SPECULATIVE=1 and EXO_MTP_WEIGHTS=<path>, MTPBatchGenerator replaces the standard MlxBatchGenerator for BS=1 inference, drafting γ tokens with the model's built-in MTP head and verifying at S=γ+1. New files in speculative/: - mtp_module.py: MTPPredictor + speculative_forward (kernel swap for GDN rollback) + draft_tokens (lazy MTP chaining) - mtp_batch_generator.py: MTPBatchGenerator subclassing mlx_lm's BatchGenerator with token buffering and BS>1 fallback - speculative_cache.py: SpeculativeArraysCache for GDN state rollback - speculative_gdn_kernel.py: Metal kernel with per-step state output Environment variables: EXO_SPECULATIVE=1 Enable speculative decoding EXO_MTP_WEIGHTS=/path/to/file Path to MTP weights safetensors EXO_SPECULATIVE_GAMMA=2 Draft tokens per cycle (default: 2) MTP weights must be extracted from the original HF model (e.g. Qwen/Qwen3.5-27B) as they are stripped during MLX quantization. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
dmcc73
@@ -1,3 +1,4 @@
|
||||
import os
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Callable, cast
|
||||
@@ -79,11 +80,47 @@ class ExoBatchGenerator:
|
||||
_active_tasks: dict[int, _EngineTask] = field(default_factory=dict, init=False)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self._exo_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 = os.environ.get("EXO_MTP_WEIGHTS", "")
|
||||
gamma = int(os.environ.get("EXO_SPECULATIVE_GAMMA", "2"))
|
||||
|
||||
if mtp_weights and os.path.exists(mtp_weights):
|
||||
mtp = MTPPredictor(self.model, mtp_weights, quantize=False)
|
||||
self._exo_gen = MTPBatchGenerator(
|
||||
model=self.model,
|
||||
mtp_predictor=mtp,
|
||||
gamma=gamma,
|
||||
stop_tokens=stop_tokens,
|
||||
prefill_step_size=4096,
|
||||
)
|
||||
logger.info(f"MTP speculative decoding enabled (γ={gamma})")
|
||||
else:
|
||||
logger.warning(f"EXO_SPECULATIVE=1 but MTP weights not found at '{mtp_weights}'. Falling back to standard generation.")
|
||||
self._exo_gen = MlxBatchGenerator(
|
||||
model=self.model,
|
||||
stop_tokens=stop_tokens,
|
||||
prefill_step_size=4096,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to initialize MTP speculative decoding: {e}. Falling back to standard generation.")
|
||||
self._exo_gen = MlxBatchGenerator(
|
||||
model=self.model,
|
||||
stop_tokens=stop_tokens,
|
||||
prefill_step_size=4096,
|
||||
)
|
||||
else:
|
||||
self._exo_gen = MlxBatchGenerator(
|
||||
model=self.model,
|
||||
stop_tokens=stop_tokens,
|
||||
prefill_step_size=4096,
|
||||
)
|
||||
|
||||
@property
|
||||
def has_work(self) -> bool:
|
||||
|
||||
@@ -0,0 +1,327 @@
|
||||
#!/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._setup_hidden_capture()
|
||||
|
||||
def _setup_hidden_capture(self):
|
||||
"""Monkey-patch model's final norm to capture pre-norm hidden state.
|
||||
|
||||
Needed for the first decode step and BS>1 fallback path.
|
||||
During speculative verify, we use speculative_forward()'s return value instead.
|
||||
"""
|
||||
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: run standard step, then prefill MTP cache."""
|
||||
responses = super()._next()
|
||||
if not responses:
|
||||
return responses
|
||||
|
||||
self.mtp.reset_cache()
|
||||
|
||||
prompt_pre_norm = self._captured.get('prompt_pre_norm')
|
||||
decode_pre_norm = self._captured.get('pre_norm')
|
||||
|
||||
# Batched MTP prefill
|
||||
# prompt_pre_norm has S positions from the prefill chunk (prompt[0:S])
|
||||
# Pair position i with token i+1 for MTP cache building
|
||||
if prompt_pre_norm is not None:
|
||||
mx.eval(prompt_pre_norm)
|
||||
toks = batch.tokens[0]
|
||||
mx.eval(toks)
|
||||
toks_list = toks.tolist()
|
||||
S_pre = prompt_pre_norm.shape[1]
|
||||
if S_pre > 1:
|
||||
mtp_tokens = toks_list[1:S_pre] # tokens 1..S_pre-1
|
||||
_ = self.mtp.predict(
|
||||
prompt_pre_norm[:, :-1, :],
|
||||
mx.array([mtp_tokens])
|
||||
)
|
||||
mx.eval(_)
|
||||
|
||||
# Use decode pre_norm (from the S=1 step that just ran)
|
||||
if decode_pre_norm is not None:
|
||||
mx.eval(decode_pre_norm)
|
||||
self._mtp_pre_norm[uid] = decode_pre_norm[:, -1:, :]
|
||||
elif prompt_pre_norm is not None:
|
||||
self._mtp_pre_norm[uid] = prompt_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.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
|
||||
toc = time.perf_counter()
|
||||
self._stats.generation_time += toc - tic
|
||||
self._stats.generation_tokens += len(all_tokens)
|
||||
|
||||
finish_reason = None
|
||||
for tok, _ in all_tokens:
|
||||
if tok in self.stop_tokens:
|
||||
finish_reason = "stop"
|
||||
break
|
||||
if batch.num_tokens[0] >= batch.max_tokens[0]:
|
||||
finish_reason = "length"
|
||||
break
|
||||
|
||||
first_tok, first_lp = all_tokens[0]
|
||||
|
||||
if finish_reason:
|
||||
cache = batch.extract_cache(0)
|
||||
self.active_batch = None
|
||||
self._cleanup_uid(uid)
|
||||
return [self.Response(uid, first_tok, first_lp, finish_reason, 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)
|
||||
@@ -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],
|
||||
)
|
||||
Reference in New Issue
Block a user