metrale_model_layers/forward/qwen3_5/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Single-token decode forward for one Qwen3.5 decoder layer, full
4//! attention or GDN, through `&dyn GpuBackend` and a `QuantWeights` impl.
5//!
6//! Owner: model-layers.
7//! Invariants: none beyond the types.
8//!
9//! [`forward_full_attention`] and [`forward_linear_attention`] each run norm,
10//! projections, attention or the GDN recurrence, the residual add with the
11//! post-attention norm, and the MLP with its residual. One token per call: the
12//! KV append writes position `cache_pos`, and attention is passed `seq_len_attn`
13//! as the cache length. Tokenizer, sampler and weight loading belong to the caller.
14
15use anyhow::Result;
16use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
17
18use super::quant_weights::QuantWeights;
19
20mod full_attention;
21mod linear_attention;
22
23pub use full_attention::forward_full_attention;
24pub use linear_attention::forward_linear_attention;
25
26/// 2026-09-25: Dimensions of a Qwen3.5 checkpoint, read by both layer forwards.
27#[derive(Debug, Clone, Copy)]
28pub struct Qwen35ForwardConfig {
29    pub hidden: u32,
30    pub intermediate: u32,
31    pub num_layers: u32,
32    pub vocab: u32,
33    pub group_size: u32,
34    pub rms_eps: f32,
35
36    pub num_heads: u32,
37    pub num_kv_heads: u32,
38    pub head_dim: u32,
39    pub rope_theta: f32,
40    /// 2026-09-25: RoPE rotates only the first `rotary_dim` elements of each head;
41    /// the rest pass through unchanged (`rope_apply.metal`).
42    pub rotary_dim: u32,
43
44    pub num_k_heads_lin: u32,
45    pub num_v_heads_lin: u32,
46    pub k_head_dim_lin: u32,
47    pub v_head_dim_lin: u32,
48    pub conv_kernel_size: u32,
49}
50
51impl Qwen35ForwardConfig {
52    /// 2026-09-25: The dimensions the `metal_qwen35_inference` example runs with; its
53    /// default model directory is `Qwen3.5-4B-MLX-8bit`.
54    pub const fn qwen3_5_4b_mlx_int8() -> Self {
55        Self {
56            hidden: 2560,
57            intermediate: 9216,
58            num_layers: 32,
59            vocab: 248_320,
60            group_size: 64,
61            rms_eps: 1e-6,
62            num_heads: 16,
63            num_kv_heads: 4,
64            head_dim: 256,
65            rope_theta: 10_000_000.0,
66            rotary_dim: 64,
67            num_k_heads_lin: 16,
68            num_v_heads_lin: 32,
69            k_head_dim_lin: 128,
70            v_head_dim_lin: 128,
71            conv_kernel_size: 4,
72        }
73    }
74
75    /// 2026-09-25: q_proj output width. Qwen3.5 packs the attention output gate
76    /// into the Q projection: each head's row is `[Q_h | gate_h]`, which
77    /// `qwen35_qkv_split` separates before the norm.
78    #[inline]
79    pub const fn q_total(&self) -> u32 {
80        self.num_heads * self.head_dim * 2
81    }
82    /// 2026-09-25: Q width after the gate is split off: half of [`Self::q_total`].
83    #[inline]
84    pub const fn q_only(&self) -> u32 {
85        self.num_heads * self.head_dim
86    }
87    #[inline]
88    pub const fn kv_dim(&self) -> u32 {
89        self.num_kv_heads * self.head_dim
90    }
91    #[inline]
92    pub const fn z_dim_lin(&self) -> u32 {
93        self.num_v_heads_lin * self.v_head_dim_lin
94    }
95    #[inline]
96    pub const fn qkv_total_lin(&self) -> u32 {
97        2 * self.num_k_heads_lin * self.k_head_dim_lin + self.num_v_heads_lin * self.v_head_dim_lin
98    }
99    /// 2026-09-25: The number of GDN heads that gate, beta, dt_bias and A_log are
100    /// indexed by.
101    #[inline]
102    pub const fn num_state_heads(&self) -> u32 {
103        self.num_v_heads_lin
104    }
105}
106
107/// 2026-09-25: Kernel handles for both layer forwards, looked up once by
108/// [`Self::resolve`].
109pub struct Qwen35Kernels {
110    pub rms: KernelHandle,
111    pub rope: KernelHandle,
112    pub kvap: KernelHandle,
113    pub attn: KernelHandle,
114    pub sg: KernelHandle,
115    pub add_rms: KernelHandle,
116    pub qkv_split: KernelHandle,
117    pub conv1d: KernelHandle,
118    pub gdn_gate: KernelHandle,
119    pub sigmoid: KernelHandle,
120    pub gdn_dec: KernelHandle,
121    /// 2026-09-25: Turbo KV cache kernels: quantizing appends, dequantizing decode
122    /// attentions, and the WHT rotations around them. `resolve` looks them up
123    /// even for a bf16 cache, so a missing one fails at startup.
124    pub kvap_turbo8: KernelHandle,
125    pub attn_turbo8: KernelHandle,
126    pub kvap_turbo4: KernelHandle,
127    pub attn_turbo4: KernelHandle,
128    pub kvap_turbo3: KernelHandle,
129    pub attn_turbo3: KernelHandle,
130    pub kvap_turbo2: KernelHandle,
131    pub attn_turbo2: KernelHandle,
132    pub kvap_bf16k_turbo4v: KernelHandle,
133    pub attn_bf16k_turbo4v: KernelHandle,
134    pub kvap_bf16k_turbo3v: KernelHandle,
135    pub attn_bf16k_turbo3v: KernelHandle,
136    pub kvap_bf16k_turbo2v: KernelHandle,
137    pub attn_bf16k_turbo2v: KernelHandle,
138    pub wht: KernelHandle,
139    pub wht_inv: KernelHandle,
140}
141
142impl Qwen35Kernels {
143    /// 2026-09-25: Looks up every kernel both layer forwards use; returns an error
144    /// if any is missing.
145    pub fn resolve(gpu: &dyn GpuBackend) -> Result<Self> {
146        Ok(Self {
147            rms: gpu.kernel("rms_norm", "rms_norm")?,
148            rope: gpu.kernel("rope_apply", "rope_apply")?,
149            kvap: gpu.kernel("kv_cache_append", "kv_cache_append")?,
150            attn: gpu.kernel("attention_decode", "attention_decode")?,
151            sg: gpu.kernel("sigmoid_gate", "sigmoid_gate")?,
152            add_rms: gpu.kernel("add_rms_norm", "add_rms_norm")?,
153            qkv_split: gpu.kernel("qwen35_qkv_split", "qwen35_qkv_split")?,
154            conv1d: gpu.kernel("causal_conv1d_update_l2norm", "causal_conv1d_update_l2norm")?,
155            gdn_gate: gpu.kernel("gdn_helpers", "gdn_compute_gate")?,
156            sigmoid: gpu.kernel("gdn_helpers", "sigmoid_bf16_to_f32")?,
157            gdn_dec: gpu.kernel("gated_delta_rule_decode", "gated_delta_rule_decode")?,
158            kvap_turbo8: gpu.kernel("kv_cache_append_turbo8", "kv_cache_append_turbo8")?,
159            attn_turbo8: gpu.kernel("attention_decode_turbo8", "attention_decode_turbo8")?,
160            kvap_turbo4: gpu.kernel("kv_cache_append_turbo4", "kv_cache_append_turbo4")?,
161            attn_turbo4: gpu.kernel("attention_decode_turbo4", "attention_decode_turbo4")?,
162            kvap_turbo3: gpu.kernel("kv_cache_append_turbo3", "kv_cache_append_turbo3")?,
163            attn_turbo3: gpu.kernel("attention_decode_turbo3", "attention_decode_turbo3")?,
164            kvap_turbo2: gpu.kernel("kv_cache_append_turbo2", "kv_cache_append_turbo2")?,
165            attn_turbo2: gpu.kernel("attention_decode_turbo2", "attention_decode_turbo2")?,
166            kvap_bf16k_turbo4v: gpu.kernel(
167                "kv_cache_append_bf16k_turbov",
168                "kv_cache_append_bf16k_turbo4v",
169            )?,
170            attn_bf16k_turbo4v: gpu.kernel(
171                "attention_decode_bf16k_turbov",
172                "attention_decode_bf16k_turbo4v",
173            )?,
174            kvap_bf16k_turbo3v: gpu.kernel(
175                "kv_cache_append_bf16k_turbov",
176                "kv_cache_append_bf16k_turbo3v",
177            )?,
178            attn_bf16k_turbo3v: gpu.kernel(
179                "attention_decode_bf16k_turbov",
180                "attention_decode_bf16k_turbo3v",
181            )?,
182            kvap_bf16k_turbo2v: gpu.kernel(
183                "kv_cache_append_bf16k_turbov",
184                "kv_cache_append_bf16k_turbo2v",
185            )?,
186            attn_bf16k_turbo2v: gpu.kernel(
187                "attention_decode_bf16k_turbov",
188                "attention_decode_bf16k_turbo2v",
189            )?,
190            wht: gpu.kernel("wht_bf16", "wht_bf16_inplace")?,
191            wht_inv: gpu.kernel("wht_bf16", "wht_bf16_inplace_inv")?,
192        })
193    }
194}
195
196/// 2026-09-25: Storage format of the contiguous KV cache ([`LayerKvCache`]).
197/// Per-element sizes are what [`LayerKvCache::alloc`] allocates, scales
198/// included.
199#[derive(Debug, Clone, Copy, PartialEq, Eq)]
200pub enum MetalKvDtype {
201    /// 2026-09-25: BF16, 2 bytes per element.
202    Bf16,
203    /// 2026-09-25: FP8 E4M3 data and a BF16 scale per group of 16, in the
204    /// WHT-rotated basis: 1.125 bytes per element.
205    Turbo8,
206    /// 2026-09-25: 4-bit Lloyd-Max codebook indices and an FP8 E4M3 scale per
207    /// group of 16 that keeps the group's L2 norm, in the WHT-rotated basis:
208    /// 0.5625 bytes per element.
209    Turbo4,
210    /// 2026-09-25: 3-bit Lloyd-Max indices (8 values in 3 bytes) and FP8 E4M3
211    /// group scales: 0.4375 bytes per element.
212    Turbo3,
213    /// 2026-09-25: 2-bit Lloyd-Max indices (4 per byte) and FP8 E4M3 group
214    /// scales: 0.3125 bytes per element.
215    Turbo2,
216    /// 2026-09-25: K in BF16, not rotated; V as Turbo4.
217    Bf16KTurbo4V,
218    /// 2026-09-25: K in BF16, not rotated; V as Turbo3.
219    Bf16KTurbo3V,
220    /// 2026-09-25: K in BF16, not rotated; V as Turbo2.
221    Bf16KTurbo2V,
222}
223
224impl MetalKvDtype {
225    /// 2026-09-25: K is stored in the WHT-rotated basis. When true, the forward
226    /// rotates K before the append and Q before attention.
227    pub fn k_is_rotated(self) -> bool {
228        matches!(
229            self,
230            Self::Turbo8 | Self::Turbo4 | Self::Turbo3 | Self::Turbo2
231        )
232    }
233    /// 2026-09-25: V is stored in the WHT-rotated basis. When true, the forward
234    /// rotates V before the append and applies the inverse WHT to the
235    /// attention output.
236    pub fn v_is_rotated(self) -> bool {
237        self != Self::Bf16
238    }
239}
240
241impl std::str::FromStr for MetalKvDtype {
242    type Err = anyhow::Error;
243    fn from_str(s: &str) -> Result<Self> {
244        match s {
245            "bf16" => Ok(Self::Bf16),
246            "turbo8" => Ok(Self::Turbo8),
247            "turbo4" => Ok(Self::Turbo4),
248            "turbo3" => Ok(Self::Turbo3),
249            "turbo2" => Ok(Self::Turbo2),
250            "bf16k_turbo4v" => Ok(Self::Bf16KTurbo4V),
251            "bf16k_turbo3v" => Ok(Self::Bf16KTurbo3V),
252            "bf16k_turbo2v" => Ok(Self::Bf16KTurbo2V),
253            other => {
254                anyhow::bail!(
255                    "kv dtype {other:?} not supported on metal (bf16 | turbo8 | turbo4 | turbo3 | turbo2 | bf16k_turbo4v/3v/2v)"
256                )
257            }
258        }
259    }
260}
261
262/// 2026-09-25: KV cache of one full-attention layer, for one sequence.
263///
264/// `dtype` selects the storage format. A quantized side holds packed data in
265/// `k`/`v` and its per-16-element group scales in `k_scales`/`v_scales`
266/// (BF16 for Turbo8, FP8 E4M3 otherwise).
267pub struct LayerKvCache {
268    pub k: DevicePtr,
269    pub v: DevicePtr,
270    /// 2026-09-25: Capacity in tokens: the `max_seq` given to [`Self::alloc`].
271    #[allow(dead_code)]
272    pub capacity: u32,
273    pub dtype: MetalKvDtype,
274    /// 2026-09-25: Group-scale buffers, `Some` only for a quantized side: both
275    /// for Turbo8/4/3/2, V only for the `Bf16KTurbo*V` formats, neither for
276    /// `Bf16`.
277    pub k_scales: Option<DevicePtr>,
278    pub v_scales: Option<DevicePtr>,
279}
280
281impl LayerKvCache {
282    /// 2026-09-25: Allocates a cache of `max_seq` tokens by `kv_dim` elements in
283    /// the given format. Panics for a turbo format when `kv_dim` is not a
284    /// multiple of 16.
285    pub fn alloc(
286        gpu: &dyn GpuBackend,
287        dtype: MetalKvDtype,
288        max_seq: u32,
289        kv_dim: u32,
290    ) -> Result<Self> {
291        assert!(
292            dtype == MetalKvDtype::Bf16 || kv_dim.is_multiple_of(16),
293            "turbo dtypes need KV_DIM divisible by 16"
294        );
295        let n = (max_seq * kv_dim) as usize;
296        let scale_bytes_e4m3 = (max_seq * kv_dim / 16) as usize;
297        // 2026-09-25: (k_bytes, v_bytes, k_scale_bytes, v_scale_bytes)
298        let (kb, vb, ksb, vsb) = match dtype {
299            MetalKvDtype::Bf16 => (n * 2, n * 2, 0, 0),
300            // 2026-09-25: Turbo8 scales are BF16: 2 bytes per group of 16.
301            MetalKvDtype::Turbo8 => (n, n, scale_bytes_e4m3 * 2, scale_bytes_e4m3 * 2),
302            MetalKvDtype::Turbo4 => (n / 2, n / 2, scale_bytes_e4m3, scale_bytes_e4m3),
303            MetalKvDtype::Turbo3 => (n * 3 / 8, n * 3 / 8, scale_bytes_e4m3, scale_bytes_e4m3),
304            MetalKvDtype::Turbo2 => (n / 4, n / 4, scale_bytes_e4m3, scale_bytes_e4m3),
305            MetalKvDtype::Bf16KTurbo4V => (n * 2, n / 2, 0, scale_bytes_e4m3),
306            MetalKvDtype::Bf16KTurbo3V => (n * 2, n * 3 / 8, 0, scale_bytes_e4m3),
307            MetalKvDtype::Bf16KTurbo2V => (n * 2, n / 4, 0, scale_bytes_e4m3),
308        };
309        let alloc_opt = |bytes: usize| -> Result<Option<DevicePtr>> {
310            Ok(if bytes > 0 {
311                Some(gpu.alloc(bytes)?)
312            } else {
313                None
314            })
315        };
316        Ok(Self {
317            k: gpu.alloc(kb)?,
318            v: gpu.alloc(vb)?,
319            capacity: max_seq,
320            dtype,
321            k_scales: alloc_opt(ksb)?,
322            v_scales: alloc_opt(vsb)?,
323        })
324    }
325}
326
327pub struct FullAttentionLayer<'a, Q: QuantWeights> {
328    pub input_ln: DevicePtr,
329    pub q_norm: DevicePtr,
330    pub k_norm: DevicePtr,
331    pub post_ln: DevicePtr,
332    pub q_proj: &'a Q,
333    pub k_proj: &'a Q,
334    pub v_proj: &'a Q,
335    pub o_proj: &'a Q,
336    pub gate_proj: &'a Q,
337    pub up_proj: &'a Q,
338    pub down_proj: &'a Q,
339}
340
341/// 2026-09-25: Scratch buffers for [`forward_full_attention`].
342pub struct FullAttentionScratch {
343    pub x_norm: DevicePtr,
344    pub q_full: DevicePtr,
345    pub q_split: DevicePtr,
346    pub gate_split: DevicePtr,
347    pub k: DevicePtr,
348    pub v: DevicePtr,
349    pub q_norm_out: DevicePtr,
350    pub k_norm_out: DevicePtr,
351    pub attn_out: DevicePtr,
352    pub gated_attn: DevicePtr,
353    pub o: DevicePtr,
354    pub x_resid: DevicePtr,
355    pub x_norm2: DevicePtr,
356    pub gate_act: DevicePtr,
357    pub up_act: DevicePtr,
358    pub x_out: DevicePtr,
359}
360
361pub struct LinearAttentionLayer<'a, Q: QuantWeights> {
362    pub input_ln: DevicePtr,
363    /// 2026-09-25: FP32 `[num_state_heads]`.
364    pub a_log: DevicePtr,
365    /// 2026-09-25: BF16 `[num_state_heads]`.
366    pub dt_bias: DevicePtr,
367    /// 2026-09-25: BF16 `[qkv_total_lin, conv_kernel_size]`.
368    pub conv1d_weight: DevicePtr,
369    pub in_proj_a: &'a Q,
370    pub in_proj_b: &'a Q,
371    pub in_proj_qkv: &'a Q,
372    pub in_proj_z: &'a Q,
373    /// 2026-09-25: BF16 `[v_head_dim_lin]`.
374    pub norm_weight: DevicePtr,
375    pub out_proj: &'a Q,
376    pub post_ln: DevicePtr,
377    pub gate_proj: &'a Q,
378    pub up_proj: &'a Q,
379    pub down_proj: &'a Q,
380}
381
382/// 2026-09-25: Conv and GDN state of one linear-attention layer, carried from
383/// token to token. The caller allocates and zeroes it.
384pub struct LinearAttentionState {
385    /// 2026-09-25: FP32 `[qkv_total_lin, conv_kernel_size]`.
386    pub conv1d_state: DevicePtr,
387    /// 2026-09-25: FP32 `[num_v_heads_lin, k_head_dim_lin, v_head_dim_lin]`.
388    pub gdn_state: DevicePtr,
389}
390
391/// 2026-09-25: Scratch buffers for [`forward_linear_attention`].
392pub struct LinearAttentionScratch {
393    pub x_norm: DevicePtr,
394    pub dt_raw: DevicePtr,
395    pub b_raw: DevicePtr,
396    pub qkv: DevicePtr,
397    pub qkv_smooth: DevicePtr,
398    pub z: DevicePtr,
399    /// 2026-09-25: FP32 `[num_state_heads]`.
400    pub gate: DevicePtr,
401    /// 2026-09-25: FP32 `[num_state_heads]`.
402    pub beta: DevicePtr,
403    pub y: DevicePtr,
404    pub y_norm: DevicePtr,
405    pub out: DevicePtr,
406    pub x_resid: DevicePtr,
407    pub x_norm2: DevicePtr,
408    pub gate_act: DevicePtr,
409    pub up_act: DevicePtr,
410    pub x_final: DevicePtr,
411}
metrale_model_layers/forward/qwen3_5/mod.rs

metrale_model_layers/forward/qwen3_5/
mod.rs
