spark_model/layers/qwen3_attention/
helpers.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! `Qwen3AttentionLayer` setters and small per-layer compute helpers
4//! (`apply_layer_scalar`, `effective_attn_scale`).
5
6use anyhow::Result;
7use spark_runtime::gpu::{DevicePtr, GpuBackend};
8
9use super::types::{HcWeights, HeadGateActivation, MlaWeights, Qwen3AttentionLayer};
10use crate::layers::FfnComponent;
11use crate::layers::ops;
12use crate::weight_map::{DenseWeight, QuantizedWeight};
13
14/// YaRN attention-temperature factor for a single `mscale` value.
15/// Matches HF `yarn_get_mscale`: `0.1 * mscale * ln(scale) + 1.0` for
16/// `scale > 1`, else `1.0`.
17fn yarn_get_mscale(scale: f32, mscale: f32) -> f32 {
18    if scale <= 1.0 {
19        1.0
20    } else {
21        0.1 * mscale * scale.ln() + 1.0
22    }
23}
24
25/// Compute the YaRN `_mscale` ratio that DeepSeek folds into the rope
26/// cos/sin: `get_mscale(factor, mscale) / get_mscale(factor, mscale_all_dim)`.
27/// Returns 1.0 when YaRN is disabled (`yarn_factor <= 1`).
28pub(crate) fn yarn_rope_mscale(config: &atlas_core::config::ModelConfig) -> f32 {
29    let factor = config.yarn_factor;
30    if factor <= 1.0 {
31        return 1.0;
32    }
33    let num = yarn_get_mscale(factor, config.yarn_mscale);
34    let den = yarn_get_mscale(factor, config.yarn_mscale_all_dim);
35    num / den
36}
37
38impl Qwen3AttentionLayer {
39    /// Set MLA weights for 2-step latent decode. When set, decode uses
40    /// latent→norm→expand instead of single-step GEMV.
41    pub fn set_mla_weights(&mut self, mla: MlaWeights) {
42        self.mla = Some(mla);
43    }
44
45    /// Set per-block Manifold-Constrained Hyper-Connection weights
46    /// (DeepSeek-V4). When set, the attn/ffn residual sites route through
47    /// `hc_pre`/`hc_post` against the model-level `hc_streams` buffer.
48    pub fn set_hc_weights(&mut self, hc: HcWeights) {
49        self.hc = Some(hc);
50    }
51
52    /// Attach the QSA indexer (Qwen3.8-Flash-Next full-attention layers).
53    pub fn set_qsa(&mut self, qsa: crate::layers::qsa::QsaIndexer) {
54        self.qsa = Some(qsa);
55    }
56
57    /// Set per-layer dimension overrides for heterogeneous models (Gemma-4).
58    /// Full-attention layers have different Q/KV head counts and head_dim
59    /// than sliding layers.
60    pub fn set_dimension_overrides(
61        &mut self,
62        head_dim: usize,
63        num_q_heads: usize,
64        num_kv_heads: usize,
65    ) {
66        self.head_dim_override = Some(head_dim);
67        self.num_q_heads_override = Some(num_q_heads);
68        self.num_kv_heads_override = Some(num_kv_heads);
69    }
70
71    /// Set per-layer sliding-window size (Gemma-4 hybrid attention).
72    /// Call with `Some(window_size)` on sliding layers, `None` on
73    /// full-attention layers. Non-Gemma-4 models never call this.
74    pub fn set_sliding_window(&mut self, window: Option<u32>) {
75        self.sliding_window = window;
76    }
77
78    /// Set per-head attention gate weight (Step 3.7 g_proj).
79    /// The weight is BF16 shape [num_q_heads, hidden_size].
80    pub(crate) fn set_head_gate_weight(&mut self, w: DenseWeight, activation: HeadGateActivation) {
81        self.head_gate_weight = Some(w);
82        self.head_gate_activation = activation;
83    }
84
85    pub(crate) fn set_yarn_rope(&mut self, inv_freq: DevicePtr, attention_factor: f32) {
86        self.yarn_inv_freq = inv_freq;
87        self.yarn_attention_factor = attention_factor;
88    }
89
90    /// Set per-layer RoPE overrides (theta, rotary_dim) for dual-RoPE
91    /// models (Gemma-4).
92    pub fn set_rope_overrides(&mut self, theta: f32, rotary_dim: u32) {
93        self.rope_theta_override = Some(theta);
94        self.rotary_dim_override = Some(rotary_dim);
95    }
96
97    /// Enable proportional RoPE (Gemma-4 full-attention layers). Must be
98    /// called AFTER `set_rope_overrides`; the `rotary_dim` set there is
99    /// reinterpreted as the number of non-zero rotation pairs.
100    pub fn set_rope_proportional(&mut self, enable: bool) {
101        self.rope_proportional = enable;
102    }
103
104    /// Set per-layer attention scale override. Gemma-4 uses QK-norm, so
105    /// attention scale should be 1.0 (not 1/sqrt(head_dim)).
106    pub fn set_attn_scale_override(&mut self, scale: f32) {
107        self.attn_scale_override = Some(scale);
108    }
109
110    /// NVFP4 M=1 decode GEMV. Single-warp when the lever is on and the kernel
111    /// resolved; otherwise the 64-thread kernel.
112    #[allow(clippy::too_many_arguments)]
113    pub(super) fn nvfp4_decode_gemv(
114        &self,
115        gpu: &dyn GpuBackend,
116        use_sw: bool,
117        input: DevicePtr,
118        weight: &QuantizedWeight,
119        output: DevicePtr,
120        n: u32,
121        k: u32,
122        stream: u64,
123    ) -> Result<()> {
124        ops::w4a16_decode_gemv(
125            gpu,
126            self.w4a16_gemv_k,
127            self.w4a16_gemv_sw_k,
128            use_sw,
129            input,
130            weight,
131            output,
132            n,
133            k,
134            stream,
135        )
136    }
137
138    /// Set K=V mode (Gemma-4 full-attention layers).
139    ///
140    /// `v_norm_weight` is a BF16 weight buffer of size `[head_dim]`. For
141    /// Gemma-4 it's ones-filled because Gemma-4's rms_norm kernel uses
142    /// the absolute convention `out = x * rms * weight`, and `weight =
143    /// 1.0` gives pure RMSNorm (matching HF
144    /// `Gemma4RMSNorm(with_scale=False)`).
145    pub fn set_k_eq_v(&mut self, v_norm_weight: DenseWeight) {
146        self.k_eq_v = true;
147        self.v_norm_weight = Some(v_norm_weight);
148    }
149
150    /// Install a pure-RMSNorm v_norm WITHOUT enabling K=V aliasing. Used
151    /// for Gemma-4 sliding-attention layers where V_proj exists on disk
152    /// but HF `Gemma4TextAttention.forward()` still applies
153    /// `value_states = self.v_norm(value_states)` with
154    /// `Gemma4RMSNorm(with_scale=False)` — pure `x * rms`.
155    pub fn set_v_norm(&mut self, v_norm_weight: DenseWeight) {
156        self.v_norm_weight = Some(v_norm_weight);
157    }
158
159    /// Install a BF16 dense fallback for the output projection. When
160    /// set, decode + prefill skip the NVFP4 `attn.o_proj` path and use
161    /// this BF16 dense_gemv / dense_gemm instead. Required for Gemma-4
162    /// dense (Nvidia ModelOpt's official ignore list keeps ALL
163    /// self_attn projections at BF16).
164    pub fn set_o_dense_bf16(&mut self, o_dense: DenseWeight) {
165        self.o_dense_bf16 = Some(o_dense);
166    }
167
168    /// Set post-sublayer norms (Gemma-4: 4-norm residual structure).
169    pub fn set_post_sublayer_norms(
170        &mut self,
171        post_attn_out: DenseWeight,
172        post_ffn_out: DenseWeight,
173    ) {
174        self.post_attn_out_norm = Some(post_attn_out);
175        self.post_ffn_out_norm = Some(post_ffn_out);
176    }
177
178    /// Set per-layer scalar (Gemma-4: hidden_states *= scalar at end of
179    /// layer).
180    pub fn set_layer_scalar(&mut self, scalar: f32) {
181        self.layer_scalar = Some(scalar);
182    }
183
184    /// Set secondary MoE FFN (Gemma-4 26B dual-FFN: dense + MoE per
185    /// layer).
186    pub fn set_moe_ffn(
187        &mut self,
188        ffn: FfnComponent,
189        pre_norm: DenseWeight,
190        post_norm: DenseWeight,
191        post_dense_norm: DenseWeight,
192    ) {
193        self.moe_ffn = Some(ffn);
194        self.pre_moe_norm = Some(pre_norm);
195        self.post_moe_out_norm = Some(post_norm);
196        self.post_dense_ffn_norm = Some(post_dense_norm);
197    }
198
199    /// LongCat: install the shortcut MoE on the FIRST sublayer of a
200    /// dual-sublayer block. The MoE runs on this sublayer's post-attention
201    /// normed input; its output is stashed into `carry` (capacity
202    /// `carry_tokens` tokens) and added by the SECOND sublayer via
203    /// [`Self::set_shortcut_carry_in`].
204    pub fn set_shortcut_moe(
205        &mut self,
206        moe: FfnComponent,
207        carry: spark_runtime::gpu::DevicePtr,
208        carry_tokens: usize,
209    ) {
210        self.moe_ffn = Some(moe);
211        self.shortcut_carry_out = Some((carry, carry_tokens));
212    }
213
214    /// LongCat: the SECOND sublayer of a dual-sublayer block adds the paired
215    /// first sublayer's stashed shortcut-MoE output at its end.
216    pub fn set_shortcut_carry_in(
217        &mut self,
218        carry: spark_runtime::gpu::DevicePtr,
219        carry_tokens: usize,
220    ) {
221        self.shortcut_carry_in = Some((carry, carry_tokens));
222    }
223
224    /// Apply layer_scalar in-place: `hidden *= scalar`. Uses
225    /// `bf16_scale_inplace` for the (always BF16) residual stream.
226    pub(crate) fn apply_layer_scalar(
227        &self,
228        gpu: &dyn spark_runtime::gpu::GpuBackend,
229        hidden: spark_runtime::gpu::DevicePtr,
230        hidden_size: usize,
231        scalar: f32,
232        stream: u64,
233    ) -> anyhow::Result<()> {
234        use spark_runtime::kernel_args::KernelLaunch;
235        let scale_k = gpu.kernel("embed_scale", "bf16_scale_inplace")?;
236        let n = hidden_size as u32;
237        KernelLaunch::new(gpu, scale_k)
238            .grid([n.div_ceil(256), 1, 1])
239            .block([256, 1, 1])
240            .arg_ptr(hidden)
241            .arg_u32(n)
242            .arg_f32(scalar)
243            .launch(stream)
244    }
245
246    /// Compute effective attention scale: override if set, else
247    /// `1/sqrt(head_dim)`.
248    pub(crate) fn effective_attn_scale(&self, head_dim: u32) -> f32 {
249        self.attn_scale_override
250            .unwrap_or_else(|| 1.0f32 / (head_dim as f32).sqrt())
251    }
252}
253
254/// The QSA per-seq carry from a sequence's [`crate::layer::AttnLayerState`],
255/// lazily created on first use (Avarok #753 item B).
256pub(in crate::layers::qwen3_attention) fn qsa_seq_state<'a>(
257    qsa: &crate::layers::qsa::QsaIndexer,
258    state: &'a mut dyn crate::layer::LayerState,
259    gpu: &dyn spark_runtime::gpu::GpuBackend,
260) -> anyhow::Result<&'a mut crate::layers::qsa::QsaSeqState> {
261    let attn = state
262        .as_any_mut()
263        .downcast_mut::<crate::layer::AttnLayerState>()
264        .ok_or_else(|| anyhow::anyhow!("QSA host layer state is not AttnLayerState"))?;
265    if attn.qsa.is_none() {
266        attn.qsa = Some(qsa.new_seq_state(gpu)?);
267    }
268    Ok(attn.qsa.as_mut().expect("just created"))
269}
270
271#[cfg(test)]
272mod yarn_mscale_tests {
273    use super::yarn_rope_mscale;
274    use atlas_core::config::ModelConfig;
275
276    // Test 1 + Test 4: with the DS4F-forced config (yarn_mscale ==
277    // yarn_mscale_all_dim == 0.0, factor 16), yarn_rope_mscale returns EXACTLY
278    // 1.0 — the single value fed to all nine DS4F rope call sites, removing the
279    // erroneous 1.2772589 amplitude on CSA/HCA layers.
280    #[test]
281    fn ds4f_forced_config_yields_mscale_one() {
282        let mut c = ModelConfig::qwen3_next_80b_nvfp4();
283        c.yarn_factor = 16.0;
284        c.yarn_mscale = 0.0;
285        c.yarn_mscale_all_dim = 0.0;
286        assert_eq!(yarn_rope_mscale(&c), 1.0);
287    }
288
289    // Test 5 (helper side): the helper itself is UNCHANGED. Under the generic
290    // HF-DeepseekV3 default (mscale 1.0, mscale_all_dim 0.0) it still returns the
291    // 1.2772589 ratio, so any legitimate YaRN-mscale caller (a different model
292    // whose config sets these fields) is unaffected. Only the DS4F *config* flips
293    // the result, not this function.
294    #[test]
295    fn generic_yarn_default_unchanged_1277() {
296        let mut c = ModelConfig::qwen3_next_80b_nvfp4();
297        c.yarn_factor = 16.0;
298        c.yarn_mscale = 1.0;
299        c.yarn_mscale_all_dim = 0.0;
300        let m = yarn_rope_mscale(&c);
301        assert!((m - 1.2772589).abs() < 1e-5, "expected ~1.2772589, got {m}");
302    }
303
304    // YaRN disabled (factor <= 1) short-circuits to 1.0 (unchanged behavior).
305    #[test]
306    fn yarn_disabled_factor_one_is_mscale_one() {
307        let mut c = ModelConfig::qwen3_next_80b_nvfp4();
308        c.yarn_factor = 1.0;
309        c.yarn_mscale = 1.0;
310        c.yarn_mscale_all_dim = 0.0;
311        assert_eq!(yarn_rope_mscale(&c), 1.0);
312    }
313}
spark_model/layers/qwen3_attention/helpers.rs

spark_model/layers/qwen3_attention/
helpers.rs
