From d3778e56c422a7e17eb097c60e29ecd92142d09f Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 20 Aug 2026 01:35:58 +0900 Subject: [PATCH] perf(speculative): single-slice the accept hook and drop the drafter's materialized causal mask Two structural cleanups from issue #1185 Phase 4, both confined to the qwen3_5_mtp drafter: - accept_verified_tokens built its paired-hidden block with a chain of per-position concatenate calls. The positions are the contiguous range starting at keep, so one slice of verify_hidden replaces n slices and n-1 intermediate buffers. - forward_hidden_stack materialized an [s, s+offset] additive causal mask on every multi-token call. Multi-token forwards now route through the shared causal_attention wrapper, which derives the same bottom-right alignment from the cache K length and selects MLX's native causal mask mode (and the M5 NAX causal kernel where present) without building a mask array. Single-token draft steps keep the maskless dispatch unchanged. The caching suggested in the issue does not apply: the offset advances every round, so a (s, offset) key never repeats. Drafter numerics only influence acceptance (module docs, layer.rs): the target verifies every proposal, so output remains byte-identical to classic decode either way. --- .../src/drafter/qwen3_5_mtp/layer.rs | 26 +++++++---- .../src/drafter/qwen3_5_mtp/model.rs | 46 +++++++------------ 2 files changed, 34 insertions(+), 38 deletions(-) diff --git a/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/layer.rs b/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/layer.rs index c13e5d2be..8546bada0 100644 --- a/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/layer.rs +++ b/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/layer.rs @@ -72,16 +72,16 @@ impl Qwen35MtpAttention { /// Attention forward over the drafter's own KV cache. /// /// - `x`: `[B, L, hidden_size]` input (post `input_layernorm`). - /// - `mask`: `Some(causal mask)` for multi-token forwards (drafter - /// prompt prefill and accepted-token extension), `None` for the - /// single-token draft steps. + /// - `causal`: `true` for multi-token forwards (drafter prompt prefill + /// and accepted-token extension), `false` for the single-token draft + /// steps, which see no future positions. /// - `cache`: this layer's own KV cache; keys/values are appended. /// - `rope_offset`: absolute position of `x`'s first token in the /// target sequence (the drafter's `next_position`). pub fn forward( &self, x: &MlxArray, - mask: Option<&MlxArray>, + causal: bool, cache: &mut KVCache, rope_offset: i32, ) -> UniquePtr { @@ -142,8 +142,18 @@ impl Qwen35MtpAttention { let (cache_k, cache_v) = cache.update_and_fetch(keys, values); - let attn_out = - crate::layers::attention(&queries, &cache_k, &cache_v, self.scale, mask, 0.0, 0); + // The causal wrapper derives the query offset from the cache K + // length (bottom-right alignment), which equals the explicit + // `create_causal_mask(s, cache.offset)` this call site used to + // build, in the degraded empty-cache mode included. It selects + // MLX's native "causal" mask mode instead of materializing an + // `[L, k_len]` additive mask per multi-token call; the kernel + // difference only touches drafter numerics (see module docs). + let attn_out = if causal { + crate::causal_attention(&queries, &cache_k, &cache_v, self.scale, 0.0, 0) + } else { + crate::layers::attention(&queries, &cache_k, &cache_v, self.scale, None, 0.0, 0) + }; let output = ffi::transpose_axes(&attn_out, &[0, 2, 1, 3]); let output = ffi::reshape(&output, &[b, l, -1]); @@ -216,12 +226,12 @@ impl Qwen35MtpDecoderLayer { pub fn forward( &self, x: &MlxArray, - mask: Option<&MlxArray>, + causal: bool, cache: &mut KVCache, rope_offset: i32, ) -> UniquePtr { let normed = self.input_layernorm.forward(x); - let r = self.attention.forward(&normed, mask, cache, rope_offset); + let r = self.attention.forward(&normed, causal, cache, rope_offset); let h = ffi::add(x, &r); let mlp_out = self.mlp.forward(&self.post_attention_layernorm.forward(&h)); ffi::add(&h, &mlp_out) diff --git a/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/model.rs b/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/model.rs index 6cfe04203..ae33f7b31 100644 --- a/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/model.rs +++ b/src/lib/mlxcel-core/src/drafter/qwen3_5_mtp/model.rs @@ -519,21 +519,17 @@ impl Qwen35MtpDraftModel { let fused = crate::ops::concatenate(&a, &b, -1); let mut h = self.fc.forward(&fused); - // Multi-token forwards use a causal mask offset by the drafter + // Multi-token forwards need causal masking offset by the drafter // cache's key length (mirrors upstream // `create_attention_mask(h, layer_cache)`); single-token draft steps - // need none. The RoPE offset is the drafter's logical position, which - // equals the cache offset while state is intact and runs ahead of it - // in the empty-cache degraded mode. - let mask = if s > 1 { - let cache_offset = self.cache.first().map(|c| c.offset).unwrap_or(0); - Some(crate::utils::create_causal_mask(s, cache_offset)) - } else { - None - }; + // need none. The layer derives that offset from its own cache K + // length, which stays correct in the empty-cache degraded mode where + // the RoPE offset (the drafter's logical position) runs ahead of the + // cache. + let causal = s > 1; let rope_offset = self.next_position; for (layer, cache) in self.layers.iter().zip(self.cache.iter_mut()) { - h = layer.forward(&h, mask.as_deref(), cache, rope_offset); + h = layer.forward(&h, causal, cache, rope_offset); } let h = self.norm.forward(&h); maybe_sync(&h); @@ -839,28 +835,18 @@ impl Drafter for Qwen35MtpDraftModel { } // Extend with the accepted tokens not yet in the cache, paired with - // the target's true verify hidden, plus the newly emitted bonus. - let h_dim = hshape[2]; - let mut tokens: Vec = Vec::new(); - let mut hidden_cat: Option> = None; - let push_slice = |pos: usize, hidden_cat: &mut Option>| { - let pos = pos as i32; - let s = ffi::slice(verify_hidden, &[0, pos, 0], &[hshape[0], pos + 1, h_dim]); - *hidden_cat = Some(match hidden_cat.take() { - None => s, - Some(prev) => crate::ops::concatenate(&prev, &s, 1), - }); - }; - for (draft_idx, &draft_tok) in draft_tokens.iter().enumerate().take(accepted).skip(keep) { - tokens.push(draft_tok); - push_slice(draft_idx, &mut hidden_cat); - } + // the target's true verify hidden, plus the newly emitted bonus. The + // hidden positions are the contiguous range starting at `keep` + // (drafts `keep..accepted`, then position `accepted` for the bonus), + // so one slice covers them all. + let mut tokens: Vec = draft_tokens[keep..accepted].to_vec(); if let Some(&last) = new_tokens.last() { tokens.push(last); - push_slice(accepted, &mut hidden_cat); } - - if let (false, Some(hiddens)) = (tokens.is_empty(), hidden_cat) { + if !tokens.is_empty() { + let start = keep as i32; + let end = start + tokens.len() as i32; + let hiddens = ffi::slice(verify_hidden, &[0, start, 0], &[hshape[0], end, hshape[2]]); let h = self.forward_hidden_stack(&tokens, &hiddens)?; self.set_seed_from_hidden(&h, sampler)?; }