metrale_model_layers/forward/qwen3_5/
linear_attention.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Single-token decode forward of one Qwen3.5 GDN (linear-attention)
4//! layer.
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    LinearAttentionLayer, LinearAttentionScratch, LinearAttentionState, Qwen35ForwardConfig,
15    Qwen35Kernels,
16};
17
18/// 2026-09-25: Runs one GDN decoder layer for one token and returns `x_buf`,
19/// into which the layer's output residual stream (`scratch.x_final`) is copied.
20/// `intra_dump`, when given, receives named intermediate buffers after a
21/// stream synchronize.
22#[allow(clippy::too_many_arguments)]
23pub fn forward_linear_attention<Q: QuantWeights>(
24    gpu: &dyn GpuBackend,
25    cfg: &Qwen35ForwardConfig,
26    k: &Qwen35Kernels,
27    layer: &LinearAttentionLayer<'_, Q>,
28    state: &LinearAttentionState,
29    scratch: &LinearAttentionScratch,
30    x_in: DevicePtr,
31    x_buf: DevicePtr,
32    stream: u64,
33    intra_dump: Option<&dyn Fn(&str, DevicePtr, u32) -> Result<()>>,
34) -> Result<DevicePtr> {
35    gpu.launch_typed(
36        k.rms,
37        [1, 1, 1],
38        [128, 1, 1],
39        0,
40        stream,
41        &[
42            KernelArg::Bytes(&cfg.hidden.to_le_bytes()),
43            KernelArg::Bytes(&cfg.rms_eps.to_le_bytes()),
44            KernelArg::Buffer(x_in),
45            KernelArg::Buffer(layer.input_ln),
46            KernelArg::Buffer(scratch.x_norm),
47        ],
48    )?;
49    // 2026-09-25: in_proj_a and in_proj_b read the same x_norm: one dual-output
50    // GEMV computes both.
51    layer.in_proj_a.gemv_gate_up_with(
52        layer.in_proj_b,
53        gpu,
54        scratch.x_norm,
55        scratch.dt_raw,
56        scratch.b_raw,
57        stream,
58    )?;
59    layer
60        .in_proj_qkv
61        .gemv(gpu, scratch.x_norm, scratch.qkv, stream)?;
62    layer
63        .in_proj_z
64        .gemv(gpu, scratch.x_norm, scratch.z, stream)?;
65
66    // 2026-09-25: One kernel runs the causal conv update and SiLU on every
67    // channel, then a per-head L2 norm on the Q and K channels only. A block
68    // is one head wide (block_x = k_head_dim_lin), so no block straddles the
69    // Q/K and V ranges.
70    let batch_one: u32 = 1;
71    let block_x: u32 = cfg.k_head_dim_lin;
72    let qkv_total_lin = cfg.qkv_total_lin();
73    let blocks_per_batch = qkv_total_lin.div_ceil(block_x);
74    let qk_channels: u32 = 2 * cfg.num_k_heads_lin * cfg.k_head_dim_lin;
75    let l2_eps: f32 = 1e-6;
76    gpu.launch_typed(
77        k.conv1d,
78        [blocks_per_batch * batch_one, 1, 1],
79        [block_x, 1, 1],
80        0,
81        stream,
82        &[
83            KernelArg::Buffer(state.conv1d_state),
84            KernelArg::Buffer(scratch.qkv),
85            KernelArg::Buffer(layer.conv1d_weight),
86            KernelArg::Buffer(scratch.qkv_smooth),
87            KernelArg::Bytes(&batch_one.to_le_bytes()),
88            KernelArg::Bytes(&qkv_total_lin.to_le_bytes()),
89            KernelArg::Bytes(&cfg.conv_kernel_size.to_le_bytes()),
90            KernelArg::Bytes(&qk_channels.to_le_bytes()),
91            KernelArg::Bytes(&cfg.k_head_dim_lin.to_le_bytes()),
92            KernelArg::Bytes(&l2_eps.to_le_bytes()),
93        ],
94    )?;
95    // 2026-09-25: The GDN query scale 1/sqrt(k_head_dim_lin) is applied by
96    // gated_delta_rule_decode to its output, not here.
97
98    // 2026-09-25: gate = exp(softplus(dt + dt_bias) * -exp(A_log)), FP32.
99    let num_state_heads = cfg.num_state_heads();
100    gpu.launch_typed(
101        k.gdn_gate,
102        [num_state_heads.div_ceil(32), 1, 1],
103        [32, 1, 1],
104        0,
105        stream,
106        &[
107            KernelArg::Bytes(&num_state_heads.to_le_bytes()),
108            KernelArg::Buffer(scratch.dt_raw),
109            KernelArg::Buffer(layer.dt_bias),
110            KernelArg::Buffer(layer.a_log),
111            KernelArg::Buffer(scratch.gate),
112        ],
113    )?;
114    // 2026-09-25: beta = sigmoid(b_raw), FP32.
115    gpu.launch_typed(
116        k.sigmoid,
117        [num_state_heads.div_ceil(32), 1, 1],
118        [32, 1, 1],
119        0,
120        stream,
121        &[
122            KernelArg::Bytes(&num_state_heads.to_le_bytes()),
123            KernelArg::Buffer(scratch.b_raw),
124            KernelArg::Buffer(scratch.beta),
125        ],
126    )?;
127
128    // 2026-09-25: qkv_smooth is BF16 `[Q | K | V]`; the views are byte offsets.
129    let k_offset = (cfg.num_k_heads_lin * cfg.k_head_dim_lin) as usize * 2;
130    let v_offset = (2 * cfg.num_k_heads_lin * cfg.k_head_dim_lin) as usize * 2;
131    let q_view = scratch.qkv_smooth;
132    let k_view = scratch.qkv_smooth.offset(k_offset);
133    let v_view = scratch.qkv_smooth.offset(v_offset);
134
135    let batch_size = 1u32;
136    let total_groups = cfg.num_v_heads_lin * batch_size;
137    gpu.launch_typed(
138        k.gdn_dec,
139        [total_groups, 1, 1],
140        [128, 1, 1],
141        0,
142        stream,
143        &[
144            KernelArg::Buffer(state.gdn_state),
145            KernelArg::Buffer(q_view),
146            KernelArg::Buffer(k_view),
147            KernelArg::Buffer(v_view),
148            KernelArg::Buffer(scratch.gate),
149            KernelArg::Buffer(scratch.beta),
150            KernelArg::Buffer(scratch.y),
151            KernelArg::Bytes(&batch_size.to_le_bytes()),
152            KernelArg::Bytes(&cfg.num_k_heads_lin.to_le_bytes()),
153            KernelArg::Bytes(&cfg.num_v_heads_lin.to_le_bytes()),
154            KernelArg::Bytes(&cfg.k_head_dim_lin.to_le_bytes()),
155            KernelArg::Bytes(&cfg.v_head_dim_lin.to_le_bytes()),
156        ],
157    )?;
158
159    // 2026-09-25: Per-head RMSNorm: each of the num_v_heads_lin heads is one row
160    // of v_head_dim_lin.
161    gpu.launch_typed(
162        k.rms,
163        [cfg.num_v_heads_lin, 1, 1],
164        [128, 1, 1],
165        0,
166        stream,
167        &[
168            KernelArg::Bytes(&cfg.v_head_dim_lin.to_le_bytes()),
169            KernelArg::Bytes(&cfg.rms_eps.to_le_bytes()),
170            KernelArg::Buffer(scratch.y),
171            KernelArg::Buffer(layer.norm_weight),
172            KernelArg::Buffer(scratch.y_norm),
173        ],
174    )?;
175
176    // 2026-09-25: out = out_proj @ (silu(z) ⊙ y_norm).
177    layer
178        .out_proj
179        .gemv_silu_gate(gpu, scratch.z, scratch.y_norm, scratch.out, stream)?;
180
181    // 2026-09-25: One kernel writes x_resid = x_in + out and x_norm2 =
182    // RMSNorm(x_resid) with post_ln.
183    gpu.launch_typed(
184        k.add_rms,
185        [1, 1, 1],
186        [128, 1, 1],
187        0,
188        stream,
189        &[
190            KernelArg::Bytes(&cfg.hidden.to_le_bytes()),
191            KernelArg::Bytes(&cfg.rms_eps.to_le_bytes()),
192            KernelArg::Buffer(x_in),
193            KernelArg::Buffer(scratch.out),
194            KernelArg::Buffer(layer.post_ln),
195            KernelArg::Buffer(scratch.x_resid),
196            KernelArg::Buffer(scratch.x_norm2),
197        ],
198    )?;
199    // 2026-09-25: gate_proj and up_proj read the same x_norm2: one dual-output
200    // GEMV.
201    layer.gate_proj.gemv_gate_up_with(
202        layer.up_proj,
203        gpu,
204        scratch.x_norm2,
205        scratch.gate_act,
206        scratch.up_act,
207        stream,
208    )?;
209    // 2026-09-25: x_final = x_resid + down_proj @ (silu(gate_act) ⊙ up_act).
210    layer.down_proj.gemv_silu_gate_resid(
211        gpu,
212        scratch.gate_act,
213        scratch.up_act,
214        scratch.x_resid,
215        scratch.x_final,
216        stream,
217    )?;
218
219    if let Some(dump) = intra_dump {
220        gpu.synchronize(stream)?;
221        let z_dim_lin = cfg.z_dim_lin();
222        dump("gdn_x_norm", scratch.x_norm, cfg.hidden)?;
223        dump("gdn_qkv_pre", scratch.qkv, qkv_total_lin)?;
224        dump("gdn_qkv_smooth", scratch.qkv_smooth, qkv_total_lin)?;
225        dump("gdn_y", scratch.y, z_dim_lin)?;
226        dump("gdn_y_norm", scratch.y_norm, z_dim_lin)?;
227        dump("gdn_out", scratch.out, cfg.hidden)?;
228        dump("gdn_x_resid", scratch.x_resid, cfg.hidden)?;
229        dump("gdn_x_final", scratch.x_final, cfg.hidden)?;
230    }
231
232    gpu.copy_d2d_async(scratch.x_final, x_buf, cfg.hidden as usize * 2, stream)?;
233    Ok(x_buf)
234}
metrale_model_layers/forward/qwen3_5/linear_attention.rs

metrale_model_layers/forward/qwen3_5/
linear_attention.rs
