metrale_model_layers/layers/qwen3_attention/
helpers.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `Qwen3AttentionLayer` setters that model loaders call after
4//! construction, the YaRN mscale helper, and small per-layer helpers
5//! (`apply_layer_scalar`, `effective_attn_scale`, `nvfp4_decode_gemv`,
6//! `qsa_seq_state`).
7//!
8//! Owner: model-layers (attention).
9//! Invariants: none beyond the types.
10
11use anyhow::Result;
12use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
13
14use super::types::{HcWeights, HeadGateActivation, MlaWeights, Qwen3AttentionLayer};
15use crate::layers::FfnComponent;
16use crate::layers::ops;
17use crate::weight_map::{DenseWeight, QuantizedWeight};
18
19/// 2026-09-25: YaRN attention-temperature factor for one `mscale`:
20/// `0.1 * mscale * ln(scale) + 1.0` for `scale > 1`, else `1.0`.
21fn yarn_get_mscale(scale: f32, mscale: f32) -> f32 {
22    if scale <= 1.0 {
23        1.0
24    } else {
25        0.1 * mscale * scale.ln() + 1.0
26    }
27}
28
29/// 2026-09-25: The YaRN mscale ratio
30/// `get_mscale(factor, yarn_mscale) / get_mscale(factor, yarn_mscale_all_dim)`,
31/// or 1.0 when `yarn_factor <= 1`. The DeepSeek-V4 RoPE call sites pass it to
32/// their kernels.
33pub fn yarn_rope_mscale(config: &metrale_config::ModelConfig) -> f32 {
34    let factor = config.yarn_factor;
35    if factor <= 1.0 {
36        return 1.0;
37    }
38    let num = yarn_get_mscale(factor, config.yarn_mscale);
39    let den = yarn_get_mscale(factor, config.yarn_mscale_all_dim);
40    num / den
41}
42
43impl Qwen3AttentionLayer {
44    /// 2026-09-25: Installs the MLA weights; the decode and prefill paths take
45    /// their MLA arms when `mla` is set.
46    pub fn set_mla_weights(&mut self, mla: MlaWeights) {
47        self.mla = Some(mla);
48    }
49
50    /// 2026-09-25: Installs the hyper-connection weights. When they are set,
51    /// decode and prefill take the hyper-connection paths, which work on the
52    /// model-level `hc_streams` buffer.
53    pub fn set_hc_weights(&mut self, hc: HcWeights) {
54        self.hc = Some(hc);
55    }
56
57    pub fn set_qsa(&mut self, qsa: crate::layers::qsa::QsaIndexer) {
58        self.qsa = Some(qsa);
59    }
60
61    /// 2026-09-25: Per-layer head_dim and Q/KV head counts, for models whose
62    /// attention layers differ in shape.
63    pub fn set_dimension_overrides(
64        &mut self,
65        head_dim: usize,
66        num_q_heads: usize,
67        num_kv_heads: usize,
68    ) {
69        self.head_dim_override = Some(head_dim);
70        self.num_q_heads_override = Some(num_q_heads);
71        self.num_kv_heads_override = Some(num_kv_heads);
72    }
73
74    /// 2026-09-25: Per-layer sliding-window size: `Some(window)` on
75    /// sliding-window layers, `None` on full-attention layers.
76    pub fn set_sliding_window(&mut self, window: Option<u32>) {
77        self.sliding_window = window;
78    }
79
80    /// 2026-09-25: Per-head attention-output gate: a BF16
81    /// `[num_q_heads, hidden_size]` weight and its activation.
82    pub fn set_head_gate_weight(&mut self, w: DenseWeight, activation: HeadGateActivation) {
83        self.head_gate_weight = Some(w);
84        self.head_gate_activation = activation;
85    }
86
87    /// 2026-09-29: Whether this layer rotates Q and K before attention; false when the
88    /// config declares no positional encoding (Nemotron-H).
89    pub(crate) fn rotates_qk(&self) -> bool {
90        self.position_encoding == metrale_config::AttnPositionEncoding::Rope
91    }
92
93    pub fn set_yarn_rope(&mut self, inv_freq: DevicePtr, attention_factor: f32) {
94        assert!(
95            self.rotates_qk(),
96            "YaRN RoPE on a layer without positional encoding"
97        );
98        self.yarn_inv_freq = inv_freq;
99        self.yarn_attention_factor = attention_factor;
100    }
101
102    pub fn set_rope_overrides(&mut self, theta: f32, rotary_dim: u32) {
103        assert!(
104            self.rotates_qk(),
105            "RoPE overrides on a layer without positional encoding"
106        );
107        self.rope_theta_override = Some(theta);
108        self.rotary_dim_override = Some(rotary_dim);
109    }
110
111    /// 2026-09-25: Proportional RoPE: rotation pairs `(i, i + head_dim/2)` for
112    /// `i` below the `rotary_dim` given to `set_rope_overrides`, which this
113    /// mode reads as the number of rotated pairs (`ops::rope_proportional`).
114    /// Both setters only store fields, so their order does not matter.
115    pub fn set_rope_proportional(&mut self, enable: bool) {
116        assert!(
117            !enable || self.rotates_qk(),
118            "proportional RoPE on a layer without positional encoding"
119        );
120        self.rope_proportional = enable;
121    }
122
123    /// 2026-09-25: Replaces the default `1/sqrt(head_dim)` attention scale
124    /// (`effective_attn_scale`).
125    pub fn set_attn_scale_override(&mut self, scale: f32) {
126        self.attn_scale_override = Some(scale);
127    }
128
129    /// 2026-09-25: NVFP4 M=1 decode GEMV: the single-warp-per-output
130    /// `w4a16_gemv_sw` when `use_sw` is set and that kernel resolved, otherwise
131    /// `w4a16_gemv` (`ops::w4a16_decode_gemv`).
132    #[allow(clippy::too_many_arguments)]
133    pub(super) fn nvfp4_decode_gemv(
134        &self,
135        gpu: &dyn GpuBackend,
136        use_sw: bool,
137        input: DevicePtr,
138        weight: &QuantizedWeight,
139        output: DevicePtr,
140        n: u32,
141        k: u32,
142        stream: u64,
143    ) -> Result<()> {
144        // 2026-10-01: Under a fixed `nvfp4` attention format, the row-invariant W4A4 mx path.
145        if ops::w4a4_proj::fixed_nvfp4_proj(
146            gpu,
147            metrale_config::ProjFamily::Attn,
148            weight,
149            input,
150            output,
151            1,
152            n,
153            k,
154            stream,
155        )? {
156            return Ok(());
157        }
158        ops::w4a16_decode_gemv(
159            gpu,
160            self.w4a16_gemv_k,
161            self.w4a16_gemv_sw_k,
162            use_sw,
163            input,
164            weight,
165            output,
166            n,
167            k,
168            stream,
169        )
170    }
171
172    /// 2026-10-01: Whether Q, K and V take the fixed `nvfp4` path at `rows` rows: a fixed
173    /// `nvfp4` attention format (`--activation-quantization`), NVFP4 q/k/v weights, no MLA and
174    /// no LoRA overlay (a q delta folds before the gated deinterleave the path runs itself).
175    pub(super) fn fixed_nvfp4_qkv_serves(&self, rows: usize) -> bool {
176        let nvfp4 = |w: &Option<crate::weight_map::QuantWeight>| {
177            w.as_ref().and_then(|w| w.as_nvfp4()).is_some()
178        };
179        self.mla.is_none()
180            && self.lora.is_none()
181            && nvfp4(&self.q_weight)
182            && nvfp4(&self.k_weight)
183            && nvfp4(&self.v_weight)
184            && crate::layers::fixed_act(metrale_config::ProjFamily::Attn, rows)
185                == Some(metrale_config::ActQuantFormat::Nvfp4)
186    }
187
188    /// 2026-09-25: Sets `k_eq_v` and installs `v_norm_weight`, a BF16
189    /// `[head_dim]` buffer. The Gemma-4 loader calls this, with a ones-filled
190    /// weight, when the checkpoint has no V projection.
191    pub fn set_k_eq_v(&mut self, v_norm_weight: DenseWeight) {
192        self.k_eq_v = true;
193        self.v_norm_weight = Some(v_norm_weight);
194    }
195
196    /// 2026-09-25: Installs `v_norm_weight` without setting `k_eq_v`. The
197    /// Gemma-4 loader calls this, with a ones-filled weight, when the checkpoint
198    /// has a V projection.
199    pub fn set_v_norm(&mut self, v_norm_weight: DenseWeight) {
200        self.v_norm_weight = Some(v_norm_weight);
201    }
202
203    /// 2026-09-25: Installs a BF16 output-projection weight. The decode and
204    /// prefill o_proj dispatch (`decode/attention_forward_oproj.rs`,
205    /// `prefill/paged_oproj.rs`, `trait_impl/multi_seq/attn/o_proj.rs`) uses it
206    /// when it is set and no earlier arm (MLA, a transposed quantized copy)
207    /// matched.
208    pub fn set_o_dense_bf16(&mut self, o_dense: DenseWeight) {
209        self.o_dense_bf16 = Some(o_dense);
210    }
211
212    pub fn set_post_sublayer_norms(
213        &mut self,
214        post_attn_out: DenseWeight,
215        post_ffn_out: DenseWeight,
216    ) {
217        self.post_attn_out_norm = Some(post_attn_out);
218        self.post_ffn_out_norm = Some(post_ffn_out);
219    }
220
221    pub fn set_layer_scalar(&mut self, scalar: f32) {
222        self.layer_scalar = Some(scalar);
223    }
224
225    /// 2026-09-25: Installs a second FFN (an MoE) with the pre-MoE,
226    /// post-MoE-output and post-dense-FFN norms; the dual-FFN path runs only
227    /// when all four are set.
228    pub fn set_moe_ffn(
229        &mut self,
230        ffn: FfnComponent,
231        pre_norm: DenseWeight,
232        post_norm: DenseWeight,
233        post_dense_norm: DenseWeight,
234    ) {
235        self.moe_ffn = Some(ffn);
236        self.pre_moe_norm = Some(pre_norm);
237        self.post_moe_out_norm = Some(post_norm);
238        self.post_dense_ffn_norm = Some(post_dense_norm);
239    }
240
241    /// 2026-09-25: Installs the shortcut MoE on the first sublayer of a
242    /// dual-sublayer block. The MoE runs on this sublayer's post-attention
243    /// normed input, before its dense FFN, and its output is stored in `carry`
244    /// (room for `carry_tokens` tokens; a larger chunk is an error). The second
245    /// sublayer adds it back, see [`Self::set_shortcut_carry_in`].
246    pub fn set_shortcut_moe(
247        &mut self,
248        moe: FfnComponent,
249        carry: metrale_gpu_runtime::gpu::DevicePtr,
250        carry_tokens: usize,
251    ) {
252        self.moe_ffn = Some(moe);
253        self.shortcut_carry_out = Some((carry, carry_tokens));
254    }
255
256    /// 2026-09-25: On the second sublayer of a dual-sublayer block: adds the
257    /// paired first sublayer's stored shortcut-MoE output at the end of this
258    /// sublayer.
259    pub fn set_shortcut_carry_in(
260        &mut self,
261        carry: metrale_gpu_runtime::gpu::DevicePtr,
262        carry_tokens: usize,
263    ) {
264        self.shortcut_carry_in = Some((carry, carry_tokens));
265    }
266
267    /// 2026-09-25: Multiplies the `hidden_size` BF16 values at `hidden` by
268    /// `scalar` in place (`embed_scale::bf16_scale_inplace`).
269    pub(crate) fn apply_layer_scalar(
270        &self,
271        gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
272        hidden: metrale_gpu_runtime::gpu::DevicePtr,
273        hidden_size: usize,
274        scalar: f32,
275        stream: u64,
276    ) -> anyhow::Result<()> {
277        use metrale_gpu_runtime::kernel_args::KernelLaunch;
278        let scale_k = gpu.kernel("embed_scale", "bf16_scale_inplace")?;
279        let n = hidden_size as u32;
280        KernelLaunch::new(gpu, scale_k)
281            .grid([n.div_ceil(256), 1, 1])
282            .block([256, 1, 1])
283            .arg_ptr(hidden)
284            .arg_u32(n)
285            .arg_f32(scalar)
286            .launch(stream)
287    }
288
289    pub(crate) fn effective_attn_scale(&self, head_dim: u32) -> f32 {
290        self.attn_scale_override
291            .unwrap_or_else(|| 1.0f32 / (head_dim as f32).sqrt())
292    }
293}
294
295/// 2026-09-25: The sequence's QSA state inside its
296/// [`crate::layer::AttnLayerState`], created on first use. Errors when the
297/// layer state is not an `AttnLayerState`.
298pub(in crate::layers::qwen3_attention) fn qsa_seq_state<'a>(
299    qsa: &crate::layers::qsa::QsaIndexer,
300    state: &'a mut dyn crate::layer::LayerState,
301    gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
302) -> anyhow::Result<&'a mut crate::layers::qsa::QsaSeqState> {
303    let attn = state
304        .as_any_mut()
305        .downcast_mut::<crate::layer::AttnLayerState>()
306        .ok_or_else(|| anyhow::anyhow!("QSA host layer state is not AttnLayerState"))?;
307    if attn.qsa.is_none() {
308        attn.qsa = Some(qsa.new_seq_state(gpu)?);
309    }
310    Ok(attn.qsa.as_mut().expect("just created"))
311}
312
313#[cfg(test)]
314mod yarn_mscale_tests {
315    use super::yarn_rope_mscale;
316    use metrale_config::ModelConfig;
317
318    // 2026-09-25: `yarn_mscale == yarn_mscale_all_dim == 0.0` at factor 16
319    // gives exactly 1.0.
320    #[test]
321    fn ds4f_forced_config_yields_mscale_one() {
322        let mut c = ModelConfig::qwen3_next_80b_nvfp4();
323        c.yarn_factor = 16.0;
324        c.yarn_mscale = 0.0;
325        c.yarn_mscale_all_dim = 0.0;
326        assert_eq!(yarn_rope_mscale(&c), 1.0);
327    }
328
329    // 2026-09-25: mscale 1.0 and mscale_all_dim 0.0 at factor 16 give
330    // `0.1 * ln(16) + 1`, about 1.2772589.
331    #[test]
332    fn generic_yarn_default_unchanged_1277() {
333        let mut c = ModelConfig::qwen3_next_80b_nvfp4();
334        c.yarn_factor = 16.0;
335        c.yarn_mscale = 1.0;
336        c.yarn_mscale_all_dim = 0.0;
337        let m = yarn_rope_mscale(&c);
338        assert!((m - 1.2772589).abs() < 1e-5, "expected ~1.2772589, got {m}");
339    }
340
341    // 2026-09-25: `yarn_factor <= 1` returns 1.0 without reading the mscales.
342    #[test]
343    fn yarn_disabled_factor_one_is_mscale_one() {
344        let mut c = ModelConfig::qwen3_next_80b_nvfp4();
345        c.yarn_factor = 1.0;
346        c.yarn_mscale = 1.0;
347        c.yarn_mscale_all_dim = 0.0;
348        assert_eq!(yarn_rope_mscale(&c), 1.0);
349    }
350}
metrale_model_layers/layers/qwen3_attention/helpers.rs

metrale_model_layers/layers/qwen3_attention/
helpers.rs
