metrale_model_layers/forward/qwen3_5/
full_attention.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Single-token decode forward of one Qwen3.5 full-attention layer,
4//! with a bf16 or turbo KV cache.
5//!
6//! Owner: model-layers.
7//! Invariants: none beyond the types.
8
9use anyhow::Result;
10use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelArg};
11
12use super::super::quant_weights::QuantWeights;
13use super::{
14    FullAttentionLayer, FullAttentionScratch, LayerKvCache, Qwen35ForwardConfig, Qwen35Kernels,
15};
16
17/// 2026-09-25: Runs one full-attention decoder layer for one token and returns
18/// `scratch.x_out`, which holds the layer's output residual stream.
19#[allow(clippy::too_many_arguments)]
20pub fn forward_full_attention<Q: QuantWeights>(
21    gpu: &dyn GpuBackend,
22    cfg: &Qwen35ForwardConfig,
23    k: &Qwen35Kernels,
24    layer: &FullAttentionLayer<'_, Q>,
25    scratch: &FullAttentionScratch,
26    kv: &LayerKvCache,
27    inv_freq_ptr: DevicePtr,
28    positions_ptr: DevicePtr,
29    x_in: DevicePtr,
30    cache_pos: u32,
31    seq_len_attn: u32,
32    stream: u64,
33) -> Result<DevicePtr> {
34    gpu.launch_typed(
35        k.rms,
36        [1, 1, 1],
37        [128, 1, 1],
38        0,
39        stream,
40        &[
41            KernelArg::Bytes(&cfg.hidden.to_le_bytes()),
42            KernelArg::Bytes(&cfg.rms_eps.to_le_bytes()),
43            KernelArg::Buffer(x_in),
44            KernelArg::Buffer(layer.input_ln),
45            KernelArg::Buffer(scratch.x_norm),
46        ],
47    )?;
48    layer
49        .q_proj
50        .gemv(gpu, scratch.x_norm, scratch.q_full, stream)?;
51    // 2026-09-25: k_proj and v_proj read the same x_norm, so one dual-output
52    // GEMV computes both.
53    layer.k_proj.gemv_gate_up_with(
54        layer.v_proj,
55        gpu,
56        scratch.x_norm,
57        scratch.k,
58        scratch.v,
59        stream,
60    )?;
61
62    // 2026-09-25: q_proj writes `[num_heads, head_dim * 2]`, each head's row
63    // `[Q_h | gate_h]`. Split Q and the gate into separate buffers.
64    gpu.launch_typed(
65        k.qkv_split,
66        [cfg.head_dim, cfg.num_heads, 1],
67        [1, 1, 1],
68        0,
69        stream,
70        &[
71            KernelArg::Bytes(&cfg.num_heads.to_le_bytes()),
72            KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
73            KernelArg::Buffer(scratch.q_full),
74            KernelArg::Buffer(scratch.q_split),
75            KernelArg::Buffer(scratch.gate_split),
76        ],
77    )?;
78    let gate_view = scratch.gate_split;
79
80    // 2026-09-25: Per-head Q and K RMSNorm: each head is one row of head_dim.
81    gpu.launch_typed(
82        k.rms,
83        [cfg.num_heads, 1, 1],
84        [128, 1, 1],
85        0,
86        stream,
87        &[
88            KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
89            KernelArg::Bytes(&cfg.rms_eps.to_le_bytes()),
90            KernelArg::Buffer(scratch.q_split),
91            KernelArg::Buffer(layer.q_norm),
92            KernelArg::Buffer(scratch.q_norm_out),
93        ],
94    )?;
95    gpu.launch_typed(
96        k.rms,
97        [cfg.num_kv_heads, 1, 1],
98        [128, 1, 1],
99        0,
100        stream,
101        &[
102            KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
103            KernelArg::Bytes(&cfg.rms_eps.to_le_bytes()),
104            KernelArg::Buffer(scratch.k),
105            KernelArg::Buffer(layer.k_norm),
106            KernelArg::Buffer(scratch.k_norm_out),
107        ],
108    )?;
109
110    // 2026-09-25: RoPE rotates q_norm_out and k_norm_out in place.
111    let half_dim = cfg.rotary_dim / 2;
112    let n_tokens = 1u32;
113    gpu.launch_typed(
114        k.rope,
115        [half_dim, cfg.num_heads, 1],
116        [1, 1, 1],
117        0,
118        stream,
119        &[
120            KernelArg::Bytes(&n_tokens.to_le_bytes()),
121            KernelArg::Bytes(&cfg.num_heads.to_le_bytes()),
122            KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
123            KernelArg::Bytes(&cfg.rotary_dim.to_le_bytes()),
124            KernelArg::Buffer(positions_ptr),
125            KernelArg::Buffer(inv_freq_ptr),
126            KernelArg::Buffer(scratch.q_norm_out),
127        ],
128    )?;
129    gpu.launch_typed(
130        k.rope,
131        [half_dim, cfg.num_kv_heads, 1],
132        [1, 1, 1],
133        0,
134        stream,
135        &[
136            KernelArg::Bytes(&n_tokens.to_le_bytes()),
137            KernelArg::Bytes(&cfg.num_kv_heads.to_le_bytes()),
138            KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
139            KernelArg::Bytes(&cfg.rotary_dim.to_le_bytes()),
140            KernelArg::Buffer(positions_ptr),
141            KernelArg::Buffer(inv_freq_ptr),
142            KernelArg::Buffer(scratch.k_norm_out),
143        ],
144    )?;
145
146    // 2026-09-25: The KV append takes the post-RoPE k_norm_out.
147    let scale: f32 = 1.0 / (cfg.head_dim as f32).sqrt();
148    if kv.dtype != super::MetalKvDtype::Bf16 {
149        // 2026-09-25: Turbo formats. A quantized side is stored in the
150        // WHT-rotated basis: K is rotated before the append and Q before
151        // attention only when K is rotated; V is rotated before the append and
152        // the attention output gets the inverse WHT only when V is rotated. In
153        // the `Bf16KTurbo*V` formats K stays raw BF16, so Q stays raw too.
154        let dt = kv.dtype;
155        let hd_bytes = cfg.head_dim.to_le_bytes();
156        if dt.k_is_rotated() {
157            gpu.launch_typed(
158                k.wht,
159                [cfg.num_kv_heads, 1, 1],
160                [32, 1, 1],
161                0,
162                stream,
163                &[
164                    KernelArg::Bytes(&hd_bytes),
165                    KernelArg::Buffer(scratch.k_norm_out),
166                ],
167            )?;
168        }
169        if dt.v_is_rotated() {
170            gpu.launch_typed(
171                k.wht,
172                [cfg.num_kv_heads, 1, 1],
173                [32, 1, 1],
174                0,
175                stream,
176                &[KernelArg::Bytes(&hd_bytes), KernelArg::Buffer(scratch.v)],
177            )?;
178        }
179        let num_groups = cfg.kv_dim() / 16;
180        let append_grid = [num_groups.div_ceil(64), 1, 1];
181        // 2026-09-25: Sparse-V threshold: the turbo attention kernels skip V rows
182        // whose unnormalized weight exp(score - max) is at or below it, so 0.0 skips
183        // only zero-weight rows. METRALE_SPARSE_V_THRESHOLD sets it; unset or
184        // unparsable gives 1e-3.
185        let sparse_v: f32 = std::env::var("METRALE_SPARSE_V_THRESHOLD")
186            .ok()
187            .and_then(|v| v.parse().ok())
188            .unwrap_or(1e-3);
189        use super::MetalKvDtype as D;
190        match dt {
191            D::Turbo8 | D::Turbo4 | D::Turbo3 | D::Turbo2 => {
192                let (kvap_turbo, attn_turbo) = match dt {
193                    D::Turbo8 => (k.kvap_turbo8, k.attn_turbo8),
194                    D::Turbo4 => (k.kvap_turbo4, k.attn_turbo4),
195                    D::Turbo3 => (k.kvap_turbo3, k.attn_turbo3),
196                    _ => (k.kvap_turbo2, k.attn_turbo2),
197                };
198                let (k_scales, v_scales) = (
199                    kv.k_scales.expect("sym turbo cache has k_scales"),
200                    kv.v_scales.expect("sym turbo cache has v_scales"),
201                );
202                gpu.launch_typed(
203                    kvap_turbo,
204                    append_grid,
205                    [64, 1, 1],
206                    0,
207                    stream,
208                    &[
209                        KernelArg::Bytes(&cfg.num_kv_heads.to_le_bytes()),
210                        KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
211                        KernelArg::Bytes(&cache_pos.to_le_bytes()),
212                        KernelArg::Buffer(scratch.k_norm_out),
213                        KernelArg::Buffer(scratch.v),
214                        KernelArg::Buffer(kv.k),
215                        KernelArg::Buffer(kv.v),
216                        KernelArg::Buffer(k_scales),
217                        KernelArg::Buffer(v_scales),
218                    ],
219                )?;
220                gpu.launch_typed(
221                    k.wht,
222                    [cfg.num_heads, 1, 1],
223                    [32, 1, 1],
224                    0,
225                    stream,
226                    &[
227                        KernelArg::Bytes(&hd_bytes),
228                        KernelArg::Buffer(scratch.q_norm_out),
229                    ],
230                )?;
231                gpu.launch_typed(
232                    attn_turbo,
233                    [cfg.num_heads, 1, 1],
234                    [32, 1, 1],
235                    0,
236                    stream,
237                    &[
238                        KernelArg::Bytes(&seq_len_attn.to_le_bytes()),
239                        KernelArg::Bytes(&cfg.num_heads.to_le_bytes()),
240                        KernelArg::Bytes(&cfg.num_kv_heads.to_le_bytes()),
241                        KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
242                        KernelArg::Bytes(&scale.to_le_bytes()),
243                        KernelArg::Bytes(&sparse_v.to_le_bytes()),
244                        KernelArg::Buffer(scratch.q_norm_out),
245                        KernelArg::Buffer(kv.k),
246                        KernelArg::Buffer(kv.v),
247                        KernelArg::Buffer(k_scales),
248                        KernelArg::Buffer(v_scales),
249                        KernelArg::Buffer(scratch.attn_out),
250                    ],
251                )?;
252            }
253            D::Bf16KTurbo4V | D::Bf16KTurbo3V | D::Bf16KTurbo2V => {
254                let (kvap_asym, attn_asym) = match dt {
255                    D::Bf16KTurbo4V => (k.kvap_bf16k_turbo4v, k.attn_bf16k_turbo4v),
256                    D::Bf16KTurbo3V => (k.kvap_bf16k_turbo3v, k.attn_bf16k_turbo3v),
257                    _ => (k.kvap_bf16k_turbo2v, k.attn_bf16k_turbo2v),
258                };
259                let v_scales = kv.v_scales.expect("asym cache has v_scales");
260                gpu.launch_typed(
261                    kvap_asym,
262                    append_grid,
263                    [64, 1, 1],
264                    0,
265                    stream,
266                    &[
267                        KernelArg::Bytes(&cfg.num_kv_heads.to_le_bytes()),
268                        KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
269                        KernelArg::Bytes(&cache_pos.to_le_bytes()),
270                        KernelArg::Buffer(scratch.k_norm_out),
271                        KernelArg::Buffer(scratch.v),
272                        KernelArg::Buffer(kv.k),
273                        KernelArg::Buffer(kv.v),
274                        KernelArg::Buffer(v_scales),
275                    ],
276                )?;
277                // 2026-09-25: K is not rotated, so Q is not rotated either.
278                gpu.launch_typed(
279                    attn_asym,
280                    [cfg.num_heads, 1, 1],
281                    [32, 1, 1],
282                    0,
283                    stream,
284                    &[
285                        KernelArg::Bytes(&seq_len_attn.to_le_bytes()),
286                        KernelArg::Bytes(&cfg.num_heads.to_le_bytes()),
287                        KernelArg::Bytes(&cfg.num_kv_heads.to_le_bytes()),
288                        KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
289                        KernelArg::Bytes(&scale.to_le_bytes()),
290                        KernelArg::Bytes(&sparse_v.to_le_bytes()),
291                        KernelArg::Buffer(scratch.q_norm_out),
292                        KernelArg::Buffer(kv.k),
293                        KernelArg::Buffer(kv.v),
294                        KernelArg::Buffer(v_scales),
295                        KernelArg::Buffer(scratch.attn_out),
296                    ],
297                )?;
298            }
299            D::Bf16 => unreachable!("outer branch excludes Bf16"),
300        }
301        if dt.v_is_rotated() {
302            gpu.launch_typed(
303                k.wht_inv,
304                [cfg.num_heads, 1, 1],
305                [32, 1, 1],
306                0,
307                stream,
308                &[
309                    KernelArg::Bytes(&hd_bytes),
310                    KernelArg::Buffer(scratch.attn_out),
311                ],
312            )?;
313        }
314    } else {
315        gpu.launch_typed(
316            k.kvap,
317            [cfg.head_dim, cfg.num_kv_heads, 1],
318            [1, 1, 1],
319            0,
320            stream,
321            &[
322                KernelArg::Bytes(&cfg.num_kv_heads.to_le_bytes()),
323                KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
324                KernelArg::Bytes(&cache_pos.to_le_bytes()),
325                KernelArg::Buffer(scratch.k_norm_out),
326                KernelArg::Buffer(scratch.v),
327                KernelArg::Buffer(kv.k),
328                KernelArg::Buffer(kv.v),
329            ],
330        )?;
331
332        gpu.launch_typed(
333            k.attn,
334            [cfg.num_heads, 1, 1],
335            [32, 1, 1],
336            0,
337            stream,
338            &[
339                KernelArg::Bytes(&seq_len_attn.to_le_bytes()),
340                KernelArg::Bytes(&cfg.num_heads.to_le_bytes()),
341                KernelArg::Bytes(&cfg.num_kv_heads.to_le_bytes()),
342                KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
343                KernelArg::Bytes(&scale.to_le_bytes()),
344                KernelArg::Buffer(scratch.q_norm_out),
345                KernelArg::Buffer(kv.k),
346                KernelArg::Buffer(kv.v),
347                KernelArg::Buffer(scratch.attn_out),
348            ],
349        )?;
350    }
351
352    let q_only = cfg.q_only();
353    gpu.launch_typed(
354        k.sg,
355        [q_only.div_ceil(64), 1, 1],
356        [64, 1, 1],
357        0,
358        stream,
359        &[
360            KernelArg::Bytes(&q_only.to_le_bytes()),
361            KernelArg::Buffer(gate_view),
362            KernelArg::Buffer(scratch.attn_out),
363            KernelArg::Buffer(scratch.gated_attn),
364        ],
365    )?;
366
367    layer
368        .o_proj
369        .gemv(gpu, scratch.gated_attn, scratch.o, stream)?;
370
371    // 2026-09-25: One kernel writes x_resid = x_in + o and x_norm2 =
372    // RMSNorm(x_resid) with post_ln.
373    gpu.launch_typed(
374        k.add_rms,
375        [1, 1, 1],
376        [128, 1, 1],
377        0,
378        stream,
379        &[
380            KernelArg::Bytes(&cfg.hidden.to_le_bytes()),
381            KernelArg::Bytes(&cfg.rms_eps.to_le_bytes()),
382            KernelArg::Buffer(x_in),
383            KernelArg::Buffer(scratch.o),
384            KernelArg::Buffer(layer.post_ln),
385            KernelArg::Buffer(scratch.x_resid),
386            KernelArg::Buffer(scratch.x_norm2),
387        ],
388    )?;
389    // 2026-09-25: gate_proj and up_proj read the same x_norm2: one dual-output
390    // GEMV.
391    layer.gate_proj.gemv_gate_up_with(
392        layer.up_proj,
393        gpu,
394        scratch.x_norm2,
395        scratch.gate_act,
396        scratch.up_act,
397        stream,
398    )?;
399    // 2026-09-25: x_out = x_resid + down_proj @ (silu(gate_act) ⊙ up_act).
400    layer.down_proj.gemv_silu_gate_resid(
401        gpu,
402        scratch.gate_act,
403        scratch.up_act,
404        scratch.x_resid,
405        scratch.x_out,
406        stream,
407    )?;
408    Ok(scratch.x_out)
409}
metrale_model_layers/forward/qwen3_5/full_attention.rs

metrale_model_layers/forward/qwen3_5/
full_attention.rs
