spark_model/forward/qwen3_5/
full_attention.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! Full-attention layer forward (single-token decode).
3
4use anyhow::Result;
5use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelArg};
6
7use super::super::quant_weights::QuantWeights;
8use super::{
9    FullAttentionLayer, FullAttentionScratch, LayerKvCache, Qwen35ForwardConfig, Qwen35Kernels,
10};
11
12/// Single-token full-attention decoder forward. Returns the
13/// `DevicePtr` containing the layer's output residual stream
14/// (caller-owned `scratch.x_out`).
15#[allow(clippy::too_many_arguments)]
16pub fn forward_full_attention<Q: QuantWeights>(
17    gpu: &dyn GpuBackend,
18    cfg: &Qwen35ForwardConfig,
19    k: &Qwen35Kernels,
20    layer: &FullAttentionLayer<'_, Q>,
21    scratch: &FullAttentionScratch,
22    kv: &LayerKvCache,
23    inv_freq_ptr: DevicePtr,
24    positions_ptr: DevicePtr,
25    x_in: DevicePtr,
26    cache_pos: u32,
27    seq_len_attn: u32,
28    stream: u64,
29) -> Result<DevicePtr> {
30    // norm1
31    gpu.launch_typed(
32        k.rms,
33        [1, 1, 1],
34        [128, 1, 1],
35        0,
36        stream,
37        &[
38            KernelArg::Bytes(&cfg.hidden.to_le_bytes()),
39            KernelArg::Bytes(&cfg.rms_eps.to_le_bytes()),
40            KernelArg::Buffer(x_in),
41            KernelArg::Buffer(layer.input_ln),
42            KernelArg::Buffer(scratch.x_norm),
43        ],
44    )?;
45    layer
46        .q_proj
47        .gemv(gpu, scratch.x_norm, scratch.q_full, stream)?;
48    // Fused k_proj + v_proj — both share x_norm and have identical
49    // (N=KV_DIM, K=HIDDEN, group_size) shapes for Qwen3.5.
50    layer.k_proj.gemv_gate_up_with(
51        layer.v_proj,
52        gpu,
53        scratch.x_norm,
54        scratch.k,
55        scratch.v,
56        stream,
57    )?;
58
59    // Qwen3.5 q_proj output is [num_heads, head_dim * 2] interleaved
60    // per head as [Q_h | gate_h]. Deinterleave into separate buffers
61    // before normalisation / RoPE / attention.
62    gpu.launch_typed(
63        k.qkv_split,
64        [cfg.head_dim, cfg.num_heads, 1],
65        [1, 1, 1],
66        0,
67        stream,
68        &[
69            KernelArg::Bytes(&cfg.num_heads.to_le_bytes()),
70            KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
71            KernelArg::Buffer(scratch.q_full),
72            KernelArg::Buffer(scratch.q_split),
73            KernelArg::Buffer(scratch.gate_split),
74        ],
75    )?;
76    let gate_view = scratch.gate_split;
77
78    // per-head q/k norm (treat each head as a token)
79    gpu.launch_typed(
80        k.rms,
81        [cfg.num_heads, 1, 1],
82        [128, 1, 1],
83        0,
84        stream,
85        &[
86            KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
87            KernelArg::Bytes(&cfg.rms_eps.to_le_bytes()),
88            KernelArg::Buffer(scratch.q_split),
89            KernelArg::Buffer(layer.q_norm),
90            KernelArg::Buffer(scratch.q_norm_out),
91        ],
92    )?;
93    gpu.launch_typed(
94        k.rms,
95        [cfg.num_kv_heads, 1, 1],
96        [128, 1, 1],
97        0,
98        stream,
99        &[
100            KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
101            KernelArg::Bytes(&cfg.rms_eps.to_le_bytes()),
102            KernelArg::Buffer(scratch.k),
103            KernelArg::Buffer(layer.k_norm),
104            KernelArg::Buffer(scratch.k_norm_out),
105        ],
106    )?;
107
108    // RoPE on the q_norm_out / k_norm_out buffers directly. Saves the
109    // d2d copy that an in-place norm would have cost.
110    let half_dim = cfg.rotary_dim / 2;
111    let n_tokens = 1u32;
112    gpu.launch_typed(
113        k.rope,
114        [half_dim, cfg.num_heads, 1],
115        [1, 1, 1],
116        0,
117        stream,
118        &[
119            KernelArg::Bytes(&n_tokens.to_le_bytes()),
120            KernelArg::Bytes(&cfg.num_heads.to_le_bytes()),
121            KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
122            KernelArg::Bytes(&cfg.rotary_dim.to_le_bytes()),
123            KernelArg::Buffer(positions_ptr),
124            KernelArg::Buffer(inv_freq_ptr),
125            KernelArg::Buffer(scratch.q_norm_out),
126        ],
127    )?;
128    gpu.launch_typed(
129        k.rope,
130        [half_dim, cfg.num_kv_heads, 1],
131        [1, 1, 1],
132        0,
133        stream,
134        &[
135            KernelArg::Bytes(&n_tokens.to_le_bytes()),
136            KernelArg::Bytes(&cfg.num_kv_heads.to_le_bytes()),
137            KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
138            KernelArg::Bytes(&cfg.rotary_dim.to_le_bytes()),
139            KernelArg::Buffer(positions_ptr),
140            KernelArg::Buffer(inv_freq_ptr),
141            KernelArg::Buffer(scratch.k_norm_out),
142        ],
143    )?;
144
145    // KV-cache append uses the post-RoPE k_norm_out.
146    let scale: f32 = 1.0 / (cfg.head_dim as f32).sqrt();
147    if kv.dtype != super::MetalKvDtype::Bf16 {
148        // ── Turbo path (symmetric Turbo8/4/3/2 + safer-asym Bf16K+TurboNV) ──
149        // Quantized sides are stored in the WHT-rotated basis. Per-side
150        // gating mirrors the CUDA bookends: rotate K at append + WHT(Q)
151        // before attention only when the K side is rotated; rotate V at
152        // append + iWHT(out) after attention only when the V side is.
153        // For the safer-asym family 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        // Sparse-V gate threshold (0.0 disables). ATLAS_SPARSE_V_THRESHOLD
182        // overrides the default 1e-3 from the attention-gated dequant work.
183        let sparse_v: f32 = std::env::var("ATLAS_SPARSE_V_THRESHOLD")
184            .ok()
185            .and_then(|v| v.parse().ok())
186            .unwrap_or(1e-3);
187        use super::MetalKvDtype as D;
188        match dt {
189            D::Turbo8 | D::Turbo4 | D::Turbo3 | D::Turbo2 => {
190                let (kvap_turbo, attn_turbo) = match dt {
191                    D::Turbo8 => (k.kvap_turbo8, k.attn_turbo8),
192                    D::Turbo4 => (k.kvap_turbo4, k.attn_turbo4),
193                    D::Turbo3 => (k.kvap_turbo3, k.attn_turbo3),
194                    _ => (k.kvap_turbo2, k.attn_turbo2),
195                };
196                let (k_scales, v_scales) = (
197                    kv.k_scales.expect("sym turbo cache has k_scales"),
198                    kv.v_scales.expect("sym turbo cache has v_scales"),
199                );
200                gpu.launch_typed(
201                    kvap_turbo,
202                    append_grid,
203                    [64, 1, 1],
204                    0,
205                    stream,
206                    &[
207                        KernelArg::Bytes(&cfg.num_kv_heads.to_le_bytes()),
208                        KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
209                        KernelArg::Bytes(&cache_pos.to_le_bytes()),
210                        KernelArg::Buffer(scratch.k_norm_out),
211                        KernelArg::Buffer(scratch.v),
212                        KernelArg::Buffer(kv.k),
213                        KernelArg::Buffer(kv.v),
214                        KernelArg::Buffer(k_scales),
215                        KernelArg::Buffer(v_scales),
216                    ],
217                )?;
218                gpu.launch_typed(
219                    k.wht,
220                    [cfg.num_heads, 1, 1],
221                    [32, 1, 1],
222                    0,
223                    stream,
224                    &[
225                        KernelArg::Bytes(&hd_bytes),
226                        KernelArg::Buffer(scratch.q_norm_out),
227                    ],
228                )?;
229                gpu.launch_typed(
230                    attn_turbo,
231                    [cfg.num_heads, 1, 1],
232                    [32, 1, 1],
233                    0,
234                    stream,
235                    &[
236                        KernelArg::Bytes(&seq_len_attn.to_le_bytes()),
237                        KernelArg::Bytes(&cfg.num_heads.to_le_bytes()),
238                        KernelArg::Bytes(&cfg.num_kv_heads.to_le_bytes()),
239                        KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
240                        KernelArg::Bytes(&scale.to_le_bytes()),
241                        KernelArg::Bytes(&sparse_v.to_le_bytes()),
242                        KernelArg::Buffer(scratch.q_norm_out),
243                        KernelArg::Buffer(kv.k),
244                        KernelArg::Buffer(kv.v),
245                        KernelArg::Buffer(k_scales),
246                        KernelArg::Buffer(v_scales),
247                        KernelArg::Buffer(scratch.attn_out),
248                    ],
249                )?;
250            }
251            D::Bf16KTurbo4V | D::Bf16KTurbo3V | D::Bf16KTurbo2V => {
252                let (kvap_asym, attn_asym) = match dt {
253                    D::Bf16KTurbo4V => (k.kvap_bf16k_turbo4v, k.attn_bf16k_turbo4v),
254                    D::Bf16KTurbo3V => (k.kvap_bf16k_turbo3v, k.attn_bf16k_turbo3v),
255                    _ => (k.kvap_bf16k_turbo2v, k.attn_bf16k_turbo2v),
256                };
257                let v_scales = kv.v_scales.expect("asym cache has v_scales");
258                gpu.launch_typed(
259                    kvap_asym,
260                    append_grid,
261                    [64, 1, 1],
262                    0,
263                    stream,
264                    &[
265                        KernelArg::Bytes(&cfg.num_kv_heads.to_le_bytes()),
266                        KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
267                        KernelArg::Bytes(&cache_pos.to_le_bytes()),
268                        KernelArg::Buffer(scratch.k_norm_out),
269                        KernelArg::Buffer(scratch.v),
270                        KernelArg::Buffer(kv.k),
271                        KernelArg::Buffer(kv.v),
272                        KernelArg::Buffer(v_scales),
273                    ],
274                )?;
275                // K is un-rotated, so Q stays un-rotated: no WHT(Q).
276                gpu.launch_typed(
277                    attn_asym,
278                    [cfg.num_heads, 1, 1],
279                    [32, 1, 1],
280                    0,
281                    stream,
282                    &[
283                        KernelArg::Bytes(&seq_len_attn.to_le_bytes()),
284                        KernelArg::Bytes(&cfg.num_heads.to_le_bytes()),
285                        KernelArg::Bytes(&cfg.num_kv_heads.to_le_bytes()),
286                        KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
287                        KernelArg::Bytes(&scale.to_le_bytes()),
288                        KernelArg::Bytes(&sparse_v.to_le_bytes()),
289                        KernelArg::Buffer(scratch.q_norm_out),
290                        KernelArg::Buffer(kv.k),
291                        KernelArg::Buffer(kv.v),
292                        KernelArg::Buffer(v_scales),
293                        KernelArg::Buffer(scratch.attn_out),
294                    ],
295                )?;
296            }
297            D::Bf16 => unreachable!("outer branch excludes Bf16"),
298        }
299        if dt.v_is_rotated() {
300            gpu.launch_typed(
301                k.wht_inv,
302                [cfg.num_heads, 1, 1],
303                [32, 1, 1],
304                0,
305                stream,
306                &[
307                    KernelArg::Bytes(&hd_bytes),
308                    KernelArg::Buffer(scratch.attn_out),
309                ],
310            )?;
311        }
312    } else {
313        gpu.launch_typed(
314            k.kvap,
315            [cfg.head_dim, cfg.num_kv_heads, 1],
316            [1, 1, 1],
317            0,
318            stream,
319            &[
320                KernelArg::Bytes(&cfg.num_kv_heads.to_le_bytes()),
321                KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
322                KernelArg::Bytes(&cache_pos.to_le_bytes()),
323                KernelArg::Buffer(scratch.k_norm_out),
324                KernelArg::Buffer(scratch.v),
325                KernelArg::Buffer(kv.k),
326                KernelArg::Buffer(kv.v),
327            ],
328        )?;
329
330        // attention_decode with seq_len = seq_len_attn.
331        gpu.launch_typed(
332            k.attn,
333            [cfg.num_heads, 1, 1],
334            [32, 1, 1],
335            0,
336            stream,
337            &[
338                KernelArg::Bytes(&seq_len_attn.to_le_bytes()),
339                KernelArg::Bytes(&cfg.num_heads.to_le_bytes()),
340                KernelArg::Bytes(&cfg.num_kv_heads.to_le_bytes()),
341                KernelArg::Bytes(&cfg.head_dim.to_le_bytes()),
342                KernelArg::Bytes(&scale.to_le_bytes()),
343                KernelArg::Buffer(scratch.q_norm_out),
344                KernelArg::Buffer(kv.k),
345                KernelArg::Buffer(kv.v),
346                KernelArg::Buffer(scratch.attn_out),
347            ],
348        )?;
349    }
350
351    // sigmoid_gate(attn_gate, attn_out)
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    // o_proj
368    layer
369        .o_proj
370        .gemv(gpu, scratch.gated_attn, scratch.o, stream)?;
371
372    // Fused residual + post-attention RMSNorm.
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    // Fused dual-output GEMV: shares x_norm2 across gate_proj and up_proj.
390    layer.gate_proj.gemv_gate_up_with(
391        layer.up_proj,
392        gpu,
393        scratch.x_norm2,
394        scratch.gate_act,
395        scratch.up_act,
396        stream,
397    )?;
398    // Fused: x_out = x_resid + down_proj @ (silu(gate_act) ⊙ up_act).
399    layer.down_proj.gemv_silu_gate_resid(
400        gpu,
401        scratch.gate_act,
402        scratch.up_act,
403        scratch.x_resid,
404        scratch.x_out,
405        stream,
406    )?;
407    Ok(scratch.x_out)
408}
spark_model/forward/qwen3_5/full_attention.rs

spark_model/forward/qwen3_5/
full_attention.rs
