spark_model/forward/qwen3_5/
linear_attention.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//! GDN (linear-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    LinearAttentionLayer, LinearAttentionScratch, LinearAttentionState, Qwen35ForwardConfig,
10    Qwen35Kernels,
11};
12
13/// Single-token GDN (linear-attention) decoder forward. Returns the
14/// `DevicePtr` containing the layer's output residual stream — that
15/// pointer is `x_buf`, into which `scratch.x_final` was copied so the
16/// caller's residual-stream buffer stays stable across layers.
17#[allow(clippy::too_many_arguments)]
18pub fn forward_linear_attention<Q: QuantWeights>(
19    gpu: &dyn GpuBackend,
20    cfg: &Qwen35ForwardConfig,
21    k: &Qwen35Kernels,
22    layer: &LinearAttentionLayer<'_, Q>,
23    state: &LinearAttentionState,
24    scratch: &LinearAttentionScratch,
25    x_in: DevicePtr,
26    x_buf: DevicePtr,
27    stream: u64,
28    intra_dump: Option<&dyn Fn(&str, DevicePtr, u32) -> Result<()>>,
29) -> Result<DevicePtr> {
30    // 1. norm
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    // 2. projections — fused in_proj_a/in_proj_b share `x_norm`.
46    layer.in_proj_a.gemv_gate_up_with(
47        layer.in_proj_b,
48        gpu,
49        scratch.x_norm,
50        scratch.dt_raw,
51        scratch.b_raw,
52        stream,
53    )?;
54    layer
55        .in_proj_qkv
56        .gemv(gpu, scratch.x_norm, scratch.qkv, stream)?;
57    layer
58        .in_proj_z
59        .gemv(gpu, scratch.x_norm, scratch.z, stream)?;
60
61    // 3. fused causal_conv1d_update_l2norm: conv + SiLU + per-head
62    // L2-norm on Q+K, SiLU only on V.
63    let batch_one: u32 = 1;
64    let block_x: u32 = cfg.k_head_dim_lin;
65    let qkv_total_lin = cfg.qkv_total_lin();
66    let blocks_per_batch = qkv_total_lin.div_ceil(block_x);
67    let qk_channels: u32 = 2 * cfg.num_k_heads_lin * cfg.k_head_dim_lin;
68    let l2_eps: f32 = 1e-6;
69    gpu.launch_typed(
70        k.conv1d,
71        [blocks_per_batch * batch_one, 1, 1],
72        [block_x, 1, 1],
73        0,
74        stream,
75        &[
76            KernelArg::Buffer(state.conv1d_state),
77            KernelArg::Buffer(scratch.qkv),
78            KernelArg::Buffer(layer.conv1d_weight),
79            KernelArg::Buffer(scratch.qkv_smooth),
80            KernelArg::Bytes(&batch_one.to_le_bytes()),
81            KernelArg::Bytes(&qkv_total_lin.to_le_bytes()),
82            KernelArg::Bytes(&cfg.conv_kernel_size.to_le_bytes()),
83            KernelArg::Bytes(&qk_channels.to_le_bytes()),
84            KernelArg::Bytes(&cfg.k_head_dim_lin.to_le_bytes()),
85            KernelArg::Bytes(&l2_eps.to_le_bytes()),
86        ],
87    )?;
88    // Atlas applies the GDN `1/sqrt(d)` factor at the kernel output;
89    // MLX applies its inv_scale at the rms_norm input. The two paths
90    // are mathematically equivalent — see the rationale in
91    // `/Users/.../atlas/memory/feedback_mlx_rms_norm_vs_l2_norm.md`.
92
93    // 4. gate = exp(softplus(dt + dt_bias) * -exp(A_log))
94    let num_state_heads = cfg.num_state_heads();
95    gpu.launch_typed(
96        k.gdn_gate,
97        [num_state_heads.div_ceil(32), 1, 1],
98        [32, 1, 1],
99        0,
100        stream,
101        &[
102            KernelArg::Bytes(&num_state_heads.to_le_bytes()),
103            KernelArg::Buffer(scratch.dt_raw),
104            KernelArg::Buffer(layer.dt_bias),
105            KernelArg::Buffer(layer.a_log),
106            KernelArg::Buffer(scratch.gate),
107        ],
108    )?;
109    // 5. beta = sigmoid(b_raw) → FP32
110    gpu.launch_typed(
111        k.sigmoid,
112        [num_state_heads.div_ceil(32), 1, 1],
113        [32, 1, 1],
114        0,
115        stream,
116        &[
117            KernelArg::Bytes(&num_state_heads.to_le_bytes()),
118            KernelArg::Buffer(scratch.b_raw),
119            KernelArg::Buffer(scratch.beta),
120        ],
121    )?;
122
123    // 6. Split qkv_smooth: Q | K | V (sequential).
124    let k_offset = (cfg.num_k_heads_lin * cfg.k_head_dim_lin) as usize * 2;
125    let v_offset = (2 * cfg.num_k_heads_lin * cfg.k_head_dim_lin) as usize * 2;
126    let q_view = scratch.qkv_smooth;
127    let k_view = scratch.qkv_smooth.offset(k_offset);
128    let v_view = scratch.qkv_smooth.offset(v_offset);
129
130    // 7. gated_delta_rule_decode
131    let batch_size = 1u32;
132    let total_groups = cfg.num_v_heads_lin * batch_size;
133    gpu.launch_typed(
134        k.gdn_dec,
135        [total_groups, 1, 1],
136        [128, 1, 1],
137        0,
138        stream,
139        &[
140            KernelArg::Buffer(state.gdn_state),
141            KernelArg::Buffer(q_view),
142            KernelArg::Buffer(k_view),
143            KernelArg::Buffer(v_view),
144            KernelArg::Buffer(scratch.gate),
145            KernelArg::Buffer(scratch.beta),
146            KernelArg::Buffer(scratch.y),
147            KernelArg::Bytes(&batch_size.to_le_bytes()),
148            KernelArg::Bytes(&cfg.num_k_heads_lin.to_le_bytes()),
149            KernelArg::Bytes(&cfg.num_v_heads_lin.to_le_bytes()),
150            KernelArg::Bytes(&cfg.k_head_dim_lin.to_le_bytes()),
151            KernelArg::Bytes(&cfg.v_head_dim_lin.to_le_bytes()),
152        ],
153    )?;
154
155    // 8. per-head rms_norm at head_dim=v_head_dim_lin over num_v_heads_lin tokens
156    gpu.launch_typed(
157        k.rms,
158        [cfg.num_v_heads_lin, 1, 1],
159        [128, 1, 1],
160        0,
161        stream,
162        &[
163            KernelArg::Bytes(&cfg.v_head_dim_lin.to_le_bytes()),
164            KernelArg::Bytes(&cfg.rms_eps.to_le_bytes()),
165            KernelArg::Buffer(scratch.y),
166            KernelArg::Buffer(layer.norm_weight),
167            KernelArg::Buffer(scratch.y_norm),
168        ],
169    )?;
170
171    // 9+10+11. Fused: out = out_proj @ (silu(z) ⊙ y_norm).
172    layer
173        .out_proj
174        .gemv_silu_gate(gpu, scratch.z, scratch.y_norm, scratch.out, stream)?;
175
176    // 12+13. Fused residual + post-attention RMSNorm.
177    gpu.launch_typed(
178        k.add_rms,
179        [1, 1, 1],
180        [128, 1, 1],
181        0,
182        stream,
183        &[
184            KernelArg::Bytes(&cfg.hidden.to_le_bytes()),
185            KernelArg::Bytes(&cfg.rms_eps.to_le_bytes()),
186            KernelArg::Buffer(x_in),
187            KernelArg::Buffer(scratch.out),
188            KernelArg::Buffer(layer.post_ln),
189            KernelArg::Buffer(scratch.x_resid),
190            KernelArg::Buffer(scratch.x_norm2),
191        ],
192    )?;
193    // Fused dual-output GEMV: shares x_norm2 across gate_proj and up_proj.
194    layer.gate_proj.gemv_gate_up_with(
195        layer.up_proj,
196        gpu,
197        scratch.x_norm2,
198        scratch.gate_act,
199        scratch.up_act,
200        stream,
201    )?;
202    // Fused: x_final = x_resid + down_proj @ (silu(gate_act) ⊙ up_act).
203    layer.down_proj.gemv_silu_gate_resid(
204        gpu,
205        scratch.gate_act,
206        scratch.up_act,
207        scratch.x_resid,
208        scratch.x_final,
209        stream,
210    )?;
211
212    // Intra-layer dumps (debug-only; gated externally via Option).
213    if let Some(dump) = intra_dump {
214        gpu.synchronize(stream)?;
215        let z_dim_lin = cfg.z_dim_lin();
216        dump("gdn_x_norm", scratch.x_norm, cfg.hidden)?;
217        dump("gdn_qkv_pre", scratch.qkv, qkv_total_lin)?;
218        dump("gdn_qkv_smooth", scratch.qkv_smooth, qkv_total_lin)?;
219        dump("gdn_y", scratch.y, z_dim_lin)?;
220        dump("gdn_y_norm", scratch.y_norm, z_dim_lin)?;
221        dump("gdn_out", scratch.out, cfg.hidden)?;
222        dump("gdn_x_resid", scratch.x_resid, cfg.hidden)?;
223        dump("gdn_x_final", scratch.x_final, cfg.hidden)?;
224    }
225
226    // Copy x_final (post-MLP-residual) to caller's stable buffer so the
227    // next layer's input pointer stays the same across layers.
228    gpu.copy_d2d_async(scratch.x_final, x_buf, cfg.hidden as usize * 2, stream)?;
229    Ok(x_buf)
230}
spark_model/forward/qwen3_5/linear_attention.rs

spark_model/forward/qwen3_5/
linear_attention.rs
