spark_model/layers/qwen3_attention/
trait_impl.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3use anyhow::Result;
4use spark_runtime::gpu::{DevicePtr, GpuBackend};
5use spark_runtime::kv_cache::PagedKvCache;
6
7use super::Qwen3AttentionLayer;
8use crate::layer::{
9    BatchedAttnMetadata, EmptyLayerState, ForwardContext, LayerState, TransformerLayer,
10};
11use crate::layers::FfnComponent;
12
13mod decode_inner;
14mod multi_seq;
15mod prefill_inner;
16
17/// Debug: read back BF16 GPU tensor and compute L2 norm + first 4 values.
18pub(super) fn diag_norm(
19    gpu: &dyn GpuBackend,
20    ptr: DevicePtr,
21    n_elements: usize,
22    stream: u64,
23    label: &str,
24) {
25    let _ = gpu.synchronize(stream);
26    let mut buf = vec![0u16; n_elements];
27    // SAFETY: `buf` is `vec![0u16; n_elements]` on the line above, so
28    // `buf.len() == n_elements` and `n_elements * 2 == buf.len() *
29    // size_of::<u16>()` — the span is exactly the Vec's buffer, all of it
30    // zero-initialised. `bytes` is the sole reference derived from `buf` while it
31    // is live: it is dead after the `copy_d2h` below, before `buf.iter()` runs.
32    let bytes =
33        unsafe { std::slice::from_raw_parts_mut(buf.as_mut_ptr() as *mut u8, n_elements * 2) };
34    if gpu.copy_d2h(ptr, bytes).is_err() {
35        return;
36    }
37    let vals: Vec<f32> = buf
38        .iter()
39        .map(|&b| f32::from_bits((b as u32) << 16))
40        .collect();
41    let norm: f32 = vals.iter().map(|v| v * v).sum::<f32>().sqrt();
42    let max_abs: f32 = vals.iter().map(|v| v.abs()).fold(0.0f32, f32::max);
43    let f4 = if vals.len() >= 4 {
44        format!(
45            "[{:.4},{:.4},{:.4},{:.4}]",
46            vals[0], vals[1], vals[2], vals[3]
47        )
48    } else {
49        format!("{:?}", &vals[..vals.len().min(4)])
50    };
51    tracing::info!("DIAG {label}: norm={norm:.4} max={max_abs:.4} first4={f4} n={n_elements}");
52}
53
54/// Debug: read back FP32 GPU tensor and compute L2 norm + first 4 values.
55/// Used by the DeepSeek-V4 multi-seq decode diagnostic path (post/comb-attn
56/// holographic tensors are FP32). V4-only — no non-V4 caller.
57pub fn diag_norm_f32(
58    gpu: &dyn GpuBackend,
59    ptr: DevicePtr,
60    n_elements: usize,
61    stream: u64,
62    label: &str,
63) {
64    let _ = gpu.synchronize(stream);
65    let mut buf = vec![0f32; n_elements];
66    // SAFETY: `buf` is `vec![0f32; n_elements]` on the line above, so
67    // `buf.len() == n_elements` and `n_elements * 4 == buf.len() *
68    // size_of::<f32>()` — the span is exactly the Vec's buffer, all of it
69    // zero-initialised. `bytes` is the sole reference derived from `buf` while it
70    // is live: it is dead after the `copy_d2h` below, before `buf.iter()` runs.
71    let bytes =
72        unsafe { std::slice::from_raw_parts_mut(buf.as_mut_ptr() as *mut u8, n_elements * 4) };
73    if gpu.copy_d2h(ptr, bytes).is_err() {
74        return;
75    }
76    let norm: f32 = buf.iter().map(|v| v * v).sum::<f32>().sqrt();
77    let max_abs: f32 = buf.iter().map(|v| v.abs()).fold(0.0f32, f32::max);
78    let f4 = if buf.len() >= 4 {
79        format!("[{:.4},{:.4},{:.4},{:.4}]", buf[0], buf[1], buf[2], buf[3])
80    } else {
81        format!("{:?}", &buf[..buf.len().min(4)])
82    };
83    tracing::info!(
84        "DIAG {label}: norm={norm:.4} max={max_abs:.4} first4={f4} n={n_elements} (FP32)"
85    );
86}
87
88// The `OnceLock<bool>` static that lived here is now
89// `layers::ops::ModelLevers::gemma4_diag`, resolved when the model is built
90// and carried on `ForwardContext`.
91
92impl TransformerLayer for Qwen3AttentionLayer {
93    fn uses_local_mla_prefill(&self) -> bool {
94        self.mla.is_some()
95    }
96
97    fn as_any_mut(&mut self) -> Option<&mut dyn std::any::Any> {
98        Some(self)
99    }
100
101    fn fp8_calibration_frozen(&self) -> Option<bool> {
102        self.fp8_calibration
103            .as_ref()
104            .map(|cal| !cal.is_calibrating())
105    }
106
107    /// QSA selection does a host top-k per step — never capturable, and a
108    /// graph captured on the dense path would replay wrong attention once
109    /// selection activates.
110    fn decode_graph_unsupported(&self) -> bool {
111        self.qsa.is_some()
112    }
113
114    fn has_aux_state(&self) -> bool {
115        self.qsa.is_some()
116    }
117
118    fn snapshot_aux(
119        &self,
120        state: &dyn LayerState,
121        gpu: &dyn GpuBackend,
122        stream: u64,
123    ) -> Result<Option<Vec<u8>>> {
124        let Some(qsa) = self.qsa.as_ref() else {
125            return Ok(None);
126        };
127        let attn = state
128            .as_any()
129            .downcast_ref::<crate::layer::AttnLayerState>()
130            .ok_or_else(|| anyhow::anyhow!("QSA host layer state is not AttnLayerState"))?;
131        match attn.qsa.as_ref() {
132            Some(st) => Ok(Some(qsa.snapshot_aux(st, gpu, stream)?)),
133            // Sequence never reached this layer's ingest: nothing to carry.
134            None => Ok(None),
135        }
136    }
137
138    fn restore_aux(
139        &self,
140        state: &mut dyn LayerState,
141        blob: &[u8],
142        gpu: &dyn GpuBackend,
143        stream: u64,
144    ) -> Result<()> {
145        let qsa = self
146            .qsa
147            .as_ref()
148            .ok_or_else(|| anyhow::anyhow!("restore_aux: no QSA on this layer"))?;
149        let attn = state
150            .as_any_mut()
151            .downcast_mut::<crate::layer::AttnLayerState>()
152            .ok_or_else(|| anyhow::anyhow!("QSA host layer state is not AttnLayerState"))?;
153        if attn.qsa.is_none() {
154            attn.qsa = Some(qsa.new_seq_state(gpu)?);
155        }
156        qsa.restore_aux(attn.qsa.as_mut().expect("just created"), blob, gpu, stream)
157    }
158
159    fn decode(
160        &self,
161        hidden: DevicePtr,
162        residual: DevicePtr,
163        state: &mut dyn LayerState,
164        kv_cache: &mut PagedKvCache,
165        seq_len: usize,
166        block_table: &mut Vec<u32>,
167        disk_block_ids: &mut Vec<u32>,
168        disk_last_offloaded_per_layer: &mut Vec<u32>,
169        ctx: &ForwardContext,
170        stream: u64,
171    ) -> Result<()> {
172        self.decode_inner(
173            hidden,
174            residual,
175            state,
176            kv_cache,
177            seq_len,
178            block_table,
179            disk_block_ids,
180            disk_last_offloaded_per_layer,
181            ctx,
182            stream,
183        )
184    }
185
186    #[allow(clippy::too_many_arguments)]
187    fn prefill(
188        &self,
189        hidden: DevicePtr,
190        residual: DevicePtr,
191        num_tokens: usize,
192        state: &mut dyn LayerState,
193        kv_cache: &mut PagedKvCache,
194        seq_len_start: usize,
195        block_table: &mut Vec<u32>,
196        disk_block_ids: &mut Vec<u32>,
197        disk_last_offloaded_per_layer: &mut Vec<u32>,
198        kv_write_start: usize,
199        ctx: &ForwardContext,
200        stream: u64,
201    ) -> Result<()> {
202        self.prefill_inner(
203            hidden,
204            residual,
205            num_tokens,
206            state,
207            kv_cache,
208            seq_len_start,
209            block_table,
210            disk_block_ids,
211            disk_last_offloaded_per_layer,
212            kv_write_start,
213            None, // batched_meta — single-stream
214            ctx,
215            stream,
216        )
217    }
218
219    /// Q12 Path B: batched-mode attention prefill via `prefill_inner` with
220    /// `batched_meta = Some`. The model-level `prefill_attn_batched_layer`
221    /// calls this method. Per-stream block_table is unused under batched
222    /// mode (block_table_ptrs from batched_meta carries them); we still
223    /// pass an empty Vec to satisfy the signature.
224    fn prefill_inner_batched_q12(
225        &self,
226        hidden_stacked: DevicePtr,
227        residual_stacked: DevicePtr,
228        num_tokens: usize,
229        kv_cache: &mut PagedKvCache,
230        seq_len_start: usize,
231        batched_meta: &BatchedAttnMetadata,
232        ctx: &ForwardContext,
233        stream: u64,
234    ) -> Result<()> {
235        let mut empty_state = EmptyLayerState;
236        let mut empty_block_table: Vec<u32> = Vec::new();
237        let mut empty_disk_block_ids: Vec<u32> = Vec::new();
238        let mut empty_disk_last: Vec<u32> = Vec::new();
239        self.prefill_inner(
240            hidden_stacked,
241            residual_stacked,
242            num_tokens,
243            &mut empty_state,
244            kv_cache,
245            seq_len_start,
246            &mut empty_block_table,
247            &mut empty_disk_block_ids,
248            &mut empty_disk_last,
249            0,
250            Some(batched_meta),
251            ctx,
252            stream,
253        )
254    }
255
256    #[allow(clippy::too_many_arguments)]
257    fn decode_multi_seq<'a, 'b: 'a>(
258        &self,
259        hidden: DevicePtr,
260        residual: DevicePtr,
261        num_seqs: usize,
262        states: &'a mut [&'b mut (dyn LayerState + 'static)],
263        kv_cache: &mut PagedKvCache,
264        seq_lens: &[usize],
265        block_tables: &[Vec<u32>],
266        ctx: &ForwardContext,
267        stream: u64,
268    ) -> Result<()> {
269        self.decode_multi_seq_inner(
270            hidden,
271            residual,
272            num_seqs,
273            states,
274            kv_cache,
275            seq_lens,
276            block_tables,
277            ctx,
278            stream,
279        )
280    }
281
282    fn alloc_state(&self, _gpu: &dyn GpuBackend) -> Result<Box<dyn LayerState>> {
283        Ok(Box::new(crate::layer::AttnLayerState::default()))
284    }
285
286    /// Free the QSA indexer carry this sequence lazily attached.
287    ///
288    /// `alloc_state` hands back an EMPTY `AttnLayerState`; the buffers appear
289    /// later, on first use, via `qsa_seq_state`. So the thing to release is
290    /// not what `alloc_state` returned — it is whatever the sequence grew.
291    /// `take()` makes this idempotent and leaves the state in the same shape
292    /// `alloc_state` produced.
293    ///
294    /// A layer with no QSA indexer (plain attention) never populates the
295    /// field, so the `take()` yields `None` and this costs nothing.
296    fn release_state(&self, state: &mut dyn LayerState, gpu: &dyn GpuBackend) -> Result<()> {
297        let Some(attn) = state
298            .as_any_mut()
299            .downcast_mut::<crate::layer::AttnLayerState>()
300        else {
301            return Ok(());
302        };
303        let Some(mut st) = attn.qsa.take() else {
304            return Ok(());
305        };
306        let Some(qsa) = self.qsa.as_ref() else {
307            // Buffers exist but the indexer is gone: nothing can size or
308            // free them correctly, and guessing would be worse than saying so.
309            anyhow::bail!("release_state: QSA seq state present but layer has no QSA indexer");
310        };
311        qsa.release_seq_state(&mut st, gpu)
312    }
313
314    fn transpose_moe_for_prefill(
315        &mut self,
316        gpu: &dyn GpuBackend,
317        config: &atlas_core::config::ModelConfig,
318    ) -> Result<()> {
319        if let FfnComponent::Moe(moe) = &mut self.ffn {
320            moe.transpose_for_prefill(gpu, config)?;
321        }
322        if let Some(FfnComponent::Moe(moe)) = self.moe_ffn.as_mut() {
323            moe.transpose_for_prefill(gpu, config)?;
324        }
325        Ok(())
326    }
327
328    fn transpose_moe_gate_up_for_prefill(
329        &mut self,
330        gpu: &dyn GpuBackend,
331        config: &atlas_core::config::ModelConfig,
332    ) -> Result<()> {
333        if let FfnComponent::Moe(moe) = &mut self.ffn {
334            moe.transpose_gate_up_for_prefill(gpu, config)?;
335        }
336        if let Some(FfnComponent::Moe(moe)) = self.moe_ffn.as_mut() {
337            moe.transpose_gate_up_for_prefill(gpu, config)?;
338        }
339        Ok(())
340    }
341
342    fn set_moe_down_transpose_scratch(
343        &mut self,
344        scratch_packed: DevicePtr,
345        scratch_scale: DevicePtr,
346        packed_ptrs_t: DevicePtr,
347        scale_ptrs_t: DevicePtr,
348    ) {
349        if let FfnComponent::Moe(moe) = &mut self.ffn {
350            moe.set_down_transpose_scratch(
351                scratch_packed,
352                scratch_scale,
353                packed_ptrs_t,
354                scale_ptrs_t,
355            );
356        }
357        if let Some(FfnComponent::Moe(moe)) = self.moe_ffn.as_mut() {
358            moe.set_down_transpose_scratch(
359                scratch_packed,
360                scratch_scale,
361                packed_ptrs_t,
362                scale_ptrs_t,
363            );
364        }
365    }
366
367    fn transpose_moe_for_prefill_unified(
368        &mut self,
369        gpu: &dyn GpuBackend,
370        config: &atlas_core::config::ModelConfig,
371    ) -> Result<()> {
372        if let FfnComponent::Moe(moe) = &mut self.ffn {
373            moe.transpose_for_prefill_unified(gpu, config)?;
374        }
375        if let Some(FfnComponent::Moe(moe)) = self.moe_ffn.as_mut() {
376            moe.transpose_for_prefill_unified(gpu, config)?;
377        }
378        Ok(())
379    }
380
381    fn transpose_moe_for_prefill_hybrid(
382        &mut self,
383        gpu: &dyn GpuBackend,
384        config: &atlas_core::config::ModelConfig,
385    ) -> Result<()> {
386        if let FfnComponent::Moe(moe) = &mut self.ffn {
387            moe.transpose_for_prefill_hybrid(gpu, config)?;
388        }
389        if let Some(FfnComponent::Moe(moe)) = self.moe_ffn.as_mut() {
390            moe.transpose_for_prefill_hybrid(gpu, config)?;
391        }
392        Ok(())
393    }
394}
spark_model/layers/qwen3_attention/trait_impl.rs

spark_model/layers/qwen3_attention/
trait_impl.rs
