spark_model/layers/qwen3_attention/
types.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Qwen3 attention struct definitions: `MlaWeights` (latent attention
4//! 2-step decode) and `Qwen3AttentionLayer` (full attention layer).
5
6use spark_runtime::gpu::{DevicePtr, KernelHandle};
7use spark_runtime::kv_cache::KvCacheDtype;
8
9use crate::layers::FfnComponent;
10use crate::layers::fp8_calibration::Fp8KvCalibration;
11use crate::layers::w4a16_gemv_tiers::W4a16BatchmTiers;
12use crate::weight_map::{AttentionWeights, DenseWeight, QuantWeight, QuantizedWeight};
13
14pub use super::types_weights::{HcWeights, MlaWeights};
15
16#[derive(Clone, Copy, Debug, Eq, PartialEq)]
17pub(crate) enum HeadGateActivation {
18    Sigmoid,
19    Softplus,
20}
21
22/// Qwen3-Next full attention layer (12 of 48 layers).
23#[allow(dead_code)]
24pub struct Qwen3AttentionLayer {
25    pub(super) input_norm: DenseWeight,
26    pub(crate) attn: AttentionWeights,
27    pub(super) post_attn_norm: DenseWeight,
28    pub(super) ffn: FfnComponent,
29    pub(super) attn_layer_idx: usize,
30    /// Startup-static LoRA adapter overlay for the K/V/O projections (v0;
31    /// q_proj excluded — gated Q+gate interleave). Installed
32    /// post-construction via `set_lora_weights`; `None` = base-only.
33    /// M0: stored only — the compute-path reads land in M1.
34    pub(super) lora: Option<crate::layers::ops::lora_delta::LoraAttnWeights>,
35    /// Whether Q projection includes an output gate (Q+Gate interleaved).
36    /// When true, q_proj output is 2× q_dim; attn output is gated by sigmoid.
37    /// When false (e.g. Qwen3-VL), q_proj output is q_dim; no gating applied.
38    pub(super) gated: bool,
39    /// Whether this layer should apply MRoPE-interleaved instead of scalar
40    /// RoPE. Set when `config.mrope_interleaved = true` (Qwen3.6).
41    pub(crate) mrope_interleaved: bool,
42    /// Per-layer dimension overrides for heterogeneous models (Gemma-4).
43    pub(crate) head_dim_override: Option<usize>,
44    pub(crate) num_q_heads_override: Option<usize>,
45    pub(crate) num_kv_heads_override: Option<usize>,
46    /// Per-layer sliding-window size for Gemma-4 hybrid attention.
47    pub(crate) sliding_window: Option<u32>,
48    /// Per-layer RoPE overrides for heterogeneous models (Gemma-4).
49    pub(crate) rope_theta_override: Option<f32>,
50    pub(crate) rotary_dim_override: Option<u32>,
51    /// Proportional RoPE (Gemma-4 full-attention).
52    pub(crate) rope_proportional: bool,
53    /// Per-layer attention scale override (Gemma-4: 1.0 because QK-norm
54    /// handles scaling). When None, uses the standard 1/sqrt(head_dim).
55    pub(crate) attn_scale_override: Option<f32>,
56    /// K=V mode: V comes from raw K projection output (no separate v_proj).
57    pub(crate) k_eq_v: bool,
58    /// Ones-filled BF16 weight buffer for the pure-RMSNorm v_norm path.
59    pub(crate) v_norm_weight: Option<DenseWeight>,
60    /// Per-head attention gate weight (Step 3.7 g_proj).
61    /// Shape: [num_q_heads, hidden_size] BF16. Applied as:
62    /// attn_out = attn_out * sigmoid(g_proj @ hidden_states)
63    /// with broadcast over head_dim.
64    pub(crate) head_gate_weight: Option<DenseWeight>,
65    pub(crate) head_gate_activation: HeadGateActivation,
66    /// Kernel handle for per-head sigmoid gate broadcast multiply.
67    pub(super) sigmoid_gate_head_broadcast_k: KernelHandle,
68    pub(super) softplus_gate_head_broadcast_k: KernelHandle,
69    /// Optional YaRN frequencies for standard (non-MLA) attention.
70    pub(crate) yarn_inv_freq: DevicePtr,
71    pub(crate) yarn_attention_factor: f32,
72    /// Post-attention output norm (Gemma-4).  
73    pub(crate) post_attn_out_norm: Option<DenseWeight>,
74    /// Post-FFN output norm (Gemma-4).
75    pub(crate) post_ffn_out_norm: Option<DenseWeight>,
76    /// Per-layer scalar (Gemma-4): hidden_states *= layer_scalar at end of forward.
77    pub(crate) layer_scalar: Option<f32>,
78    /// Secondary FFN (Gemma-4 26B MoE): runs in parallel with primary FFN (dense).
79    pub(crate) moe_ffn: Option<FfnComponent>,
80    /// LongCat shortcut-MoE PRODUCER: this sublayer computes `moe_ffn` on its
81    /// post-attention normed input and STASHES the result into the carry
82    /// buffer `(ptr, token_capacity)` instead of adding it — the paired NEXT
83    /// sublayer adds it at its end. Gated separately from the Gemma-4 dual-FFN
84    /// arm (which requires the three Gemma norms, absent here).
85    pub(crate) shortcut_carry_out: Option<(spark_runtime::gpu::DevicePtr, usize)>,
86    /// LongCat shortcut-MoE CONSUMER: after this sublayer's FFN residual add,
87    /// `hidden += carry` (the shortcut MoE output stashed by the previous
88    /// sublayer).
89    pub(crate) shortcut_carry_in: Option<(spark_runtime::gpu::DevicePtr, usize)>,
90    /// Pre-norm for MoE input (pre_feedforward_layernorm_2).
91    pub(crate) pre_moe_norm: Option<DenseWeight>,
92    /// Post-norm for MoE output (post_feedforward_layernorm_2).
93    pub(crate) post_moe_out_norm: Option<DenseWeight>,
94    /// Post-norm for dense FFN output only (post_feedforward_layernorm_1).
95    pub(crate) post_dense_ffn_norm: Option<DenseWeight>,
96    pub(super) kv_dtype: KvCacheDtype,
97    /// Turbo4 sparse-V pruning threshold (0.0 = disabled).
98    pub(super) sparse_v_threshold: f32,
99    // ── Decode weights (QuantWeight enum: Nvfp4 | Fp8 | Dense) ──
100    pub(super) q_weight: Option<QuantWeight>,
101    pub(super) k_weight: Option<QuantWeight>,
102    pub(super) v_weight: Option<QuantWeight>,
103    pub(super) o_weight: Option<QuantWeight>,
104    /// BF16 dense fallback for the output projection. When `Some`, the
105    /// decode/prefill o_proj GEMV uses this BF16 weight instead of the
106    /// NVFP4 path (`attn.o_proj`). Used by Gemma-4 dense which honors
107    /// Nvidia ModelOpt's official ignore list.
108    pub(super) o_dense_bf16: Option<DenseWeight>,
109    // ── MLA (Multi-head Latent Attention) — 2-step decode ──
110    pub(crate) mla: Option<MlaWeights>,
111    // ── Manifold-Constrained Hyper-Connections (mHC) — DeepSeek-V4 ──
112    /// Per-block HC parameters. `Some` only for DeepSeek-V4 (`hc_mult > 0`),
113    /// in which case the attn/ffn residual sites use `hc_pre`/`hc_post`
114    /// against the `hc_streams` buffer instead of the standard residual add.
115    pub(crate) hc: Option<HcWeights>,
116    // ── QSA indexer (Qwen3.8-Flash-Next) ──
117    /// Decode-side sparse-attention selection. `Some` only on the 12
118    /// full-attention layers of qwen4_exp. Presence vetoes decode-graph
119    /// capture (the selection top-k is a host round trip).
120    pub(crate) qsa: Option<crate::layers::qsa::QsaIndexer>,
121    /// HC `hc_pre` kernel handle (NULL when HC disabled).
122    pub(super) hc_pre_k: KernelHandle,
123    /// HC `hc_post` kernel handle (NULL when HC disabled).
124    pub(super) hc_post_k: KernelHandle,
125    /// HC `hc_expand` kernel handle (NULL when HC disabled).
126    pub(super) hc_expand_k: KernelHandle,
127    /// HC `hc_head` kernel handle (NULL when HC disabled).
128    pub(super) hc_head_k: KernelHandle,
129    // ── Transposed weights for prefill GEMM ──
130    /// Fused [q|k|v] transposed twin (N = q_proj_dim + 2*kv_dim). Present only
131    /// when the three projections share one `weight_scale_2` — the GEMM applies
132    /// a single scale2 per launch. `None` => the three separate GEMMs run.
133    pub(super) qkv_nvfp4_t: Option<QuantizedWeight>,
134    pub(super) q_nvfp4_t: Option<QuantizedWeight>,
135    pub(super) k_nvfp4_t: Option<QuantizedWeight>,
136    pub(super) v_nvfp4_t: Option<QuantizedWeight>,
137    pub(super) o_nvfp4_t: Option<QuantizedWeight>,
138    pub(super) q_fp8w_t: Option<crate::weight_map::Fp8WeightTransposed>,
139    pub(super) k_fp8w_t: Option<crate::weight_map::Fp8WeightTransposed>,
140    pub(super) v_fp8w_t: Option<crate::weight_map::Fp8WeightTransposed>,
141    pub(super) o_fp8w_t: Option<crate::weight_map::Fp8WeightTransposed>,
142    pub(super) w8a16_gemm_t_k: KernelHandle,
143    pub(super) w8a16_gemm_t_pipelined_k: KernelHandle,
144    // Fast transposed FP8 prefill GEMM (128x128 / 8-warp / two-level FP32 fold).
145    // Consumes the SAME B_t[K,N] + block_scale_t[K/128,N/128] that
146    // transpose_fp8 / transpose_block_scale already produce. KernelHandle(0) on
147    // miss → fall back to w8a16_gemm_t.
148    pub(super) w8a16_gemm_t_m128_k: KernelHandle,
149    // W8A8 + FP32 epilogue (vLLM-equivalent) — gated by ATLAS_FP8_W8A8=1.
150    pub(super) per_token_group_quant_fp8_k: KernelHandle,
151    pub(super) fp8_gemm_t_blockscaled_k: KernelHandle,
152    // Kernels — decode (GEMV M=1)
153    /// Offset-from-1 `rms_norm` (`out = x * (1 + w) / rms`). Used ONLY for the
154    /// unweighted normalize (`norm_unit_w()` is zero-filled, so `1 + 0 = 1`).
155    pub(super) rms_norm_k: KernelHandle,
156    /// The norm kernel for every weight that comes from the CHECKPOINT.
157    /// Same handle as `rms_norm_k` for offset-from-1 models; `rms_norm_vanilla`
158    /// (`out = x * w / rms`) for models that ship HF-vanilla norm weights.
159    pub(super) rms_norm_w_k: KernelHandle,
160    /// Warp-per-row sibling of `rms_norm_w_k` for short per-head rows; 0 if absent.
161    pub(super) rms_norm_w_warp_row_k: KernelHandle,
162    /// True when `rms_norm_w_k` is the vanilla kernel — i.e. the checkpoint's
163    /// norm weights are loaded exactly, with no `-1` pre-subtraction.
164    pub(super) norm_vanilla: bool,
165    pub(super) rms_norm_residual_k: KernelHandle,
166    /// Gemma-4 FP32-input rms_norm (absolute formula).
167    pub(super) rms_norm_f32_in_k: KernelHandle,
168    pub(super) dense_gemv_k: KernelHandle,
169    /// Load-time packed-Q2 → BF16 dequant (`dequant_gguf_bf16` module). Used by
170    /// the Tier-1c keep-packed attention PREFILL path: dequant q/k/v/o into a
171    /// transient BF16 scratch, run the normal `dense_gemm`, free. Decode uses
172    /// the native `q2_0_gemv_vec` (no dequant). `KernelHandle(0)` when absent.
173    pub(super) dequant_q2_0_gn_k: KernelHandle,
174    /// Native Q2_0 MMQ prefill (Tier-2, `ATLAS_GGUF_NATIVE_Q2_MMQ`): keep-packed
175    /// tensor-core int8 MMA vs a shared q8_1 activation. `KernelHandle(0)` when
176    /// absent → the transient-dequant prefill path is used instead. The q8_1
177    /// activation quantizer is shared with Q4_K (`q4k_quant_act_k`).
178    pub(super) q2_0_mmq_nc_k: KernelHandle,
179    pub(super) q2_0_mmq_wc_k: KernelHandle,
180    pub(super) q4k_quant_act_k: KernelHandle,
181    /// Native `q2_0_gemv_vec` decode kernel for keep-packed q/k/v/o.
182    pub(super) q2_0_gemv_k: KernelHandle,
183    /// Batched BF16 GEMV (M rows, one weight pass). Multi-seq decode q/k/v for
184    /// models whose attention weights are plain BF16 -- the quantized paths have
185    /// w4a16/w8a16 batch tiers, BF16 had none.
186    pub(super) dense_gemv_batchm_k: KernelHandle,
187    pub(super) w4a16_gemv_k: KernelHandle,
188    /// Single-warp `w4a16_gemv_sw`. `KernelHandle(0)` on miss → base GEMV.
189    pub(super) w4a16_gemv_sw_k: KernelHandle,
190    pub(super) w8a16_gemv_k: KernelHandle,
191    pub(super) w8a16_gemm_k: KernelHandle,
192    pub(super) w8a16_gemm_pipelined_k: KernelHandle,
193    pub(super) w4a16_gemv_dual_k: KernelHandle,
194    pub(super) rope_k: KernelHandle,
195    /// Strided sibling: rotates all n sequences in ONE launch. 0 when absent.
196    pub(super) rope_strided_k: KernelHandle,
197    /// Strided sibling of `rms_norm_w_k`: all n sequences in ONE launch. 0 when absent.
198    pub(super) rms_norm_strided_k: KernelHandle,
199    /// MRoPE-interleaved kernel.
200    pub(super) rope_mrope_interleaved_k: KernelHandle,
201    /// K-only MRoPE kernel used when Q RoPE is fused into Q deinterleave/norm.
202    pub(super) rope_mrope_interleaved_k_only_k: KernelHandle,
203    /// YaRN RoPE kernel using pre-computed inv_freq table (Mistral, etc.)
204    pub(super) rope_yarn_k: KernelHandle,
205    pub(super) rope_yarn_scaled_k: KernelHandle,
206    /// Interleaved (GPT-J / is_neox_style=False) YaRN RoPE kernel — DeepSeek MLA.
207    pub(super) rope_yarn_interleaved_k: KernelHandle,
208    /// Conjugate (negated-sin) interleaved YaRN RoPE — DeepSeek-V4 attention
209    /// output de-rotation (eq.26).
210    pub(super) rope_yarn_interleaved_inv_k: KernelHandle,
211    /// Proportional RoPE kernel (Gemma-4 full-attention layers).
212    pub(super) rope_proportional_k: KernelHandle,
213    pub(super) reshape_cache_k: KernelHandle,
214    /// Fused k_norm + RoPE + paged BF16 cache write — eliminates two
215    /// intermediate BF16 rounding steps that cause the documented L35-L39
216    /// cliff in chunked-prefill BF16 KV mode (memory:
217    /// `project_qwen36_phase2b_softmax_expf.md`).
218    pub(super) fused_k_norm_rope_cache_write_bf16_k: KernelHandle,
219    /// MRoPE-interleaved variant of the above. Same precision regime.
220    /// Dispatched when `mrope_interleaved` is true.
221    pub(super) fused_k_norm_rope_mrope_cache_write_bf16_k: KernelHandle,
222    /// V-only paged cache write. Used alongside the fused K-path so the
223    /// K side of the cache stays single-rounded.
224    pub(super) reshape_and_cache_flash_v_only_k: KernelHandle,
225    /// WHT kernel for turbo KV cache.
226    pub(super) wht_bf16_k: KernelHandle,
227    /// Inverse WHT. With TQ_PLUS_SIGNS off this aliases the forward kernel
228    /// (plain WHT is self-inverse); with TQ+ signs the inverse reverses the
229    /// signs1/signs2 order, which is required because (S2·H·S1)·(S2·H·S1) ≠ I.
230    pub(super) wht_bf16_k_inv: KernelHandle,
231    /// InnerQ application kernels (Q pre-WHT scale_inv, K post-WHT scale).
232    /// Returns 0 handle when InnerQ kernel module isn't loaded — caller should
233    /// guard launches with `.0 != 0`.
234    pub(super) innerq_apply_q_k: KernelHandle,
235    pub(super) innerq_apply_k_k: KernelHandle,
236    pub(super) paged_decode_k: KernelHandle,
237    /// HDIM=512 paged decode kernel for Gemma-4 full-attention layers
238    pub(super) paged_decode_512_k: KernelHandle,
239    /// MLA absorbed paged decode kernel (HDIM=320).
240    pub(super) paged_decode_mla_k: KernelHandle,
241    /// MLA paged decode kernel for DeepSeek-V4-Flash (compressed KV cache: 576 dims)
242    pub(super) mla_paged_decode_k: KernelHandle,
243    /// MLA paged decode kernel for DeepSeek-V4-Flash with FP8 KV cache
244    pub(super) mla_paged_decode_fp8_k: KernelHandle,
245    /// MLA batched GEMV for Q absorption and V extraction.
246    pub(super) mla_batched_gemv_k: KernelHandle,
247    /// MLA fused kernels — decode.
248    pub(super) mla_q_rope_scatter_k: KernelHandle,
249    pub(super) mla_q_rope_writeback_k: KernelHandle,
250    pub(super) mla_cache_assemble_k: KernelHandle,
251    /// MLA fused kernels — prefill.
252    pub(super) mla_q_rope_extract_batched_k: KernelHandle,
253    pub(super) mla_q_rope_writeback_batched_k: KernelHandle,
254    pub(super) mla_kv_assemble_batched_k: KernelHandle,
255    pub(super) mla_cache_assemble_batched_k: KernelHandle,
256    /// MLA absorbed prefill flash attention (HDIM=320, GQA 32:1)
257    pub(super) prefill_attn_mla320_k: KernelHandle,
258    /// Grouped GEMM for MLA Q absorption + V extraction.
259    pub(super) grouped_gemm_mla_k: KernelHandle,
260    /// Q_final assembly: [absorbed|rope] per head.
261    pub(super) mla_q_final_assemble_k: KernelHandle,
262    /// Fused MLA prefill: Q_absorb + attention + V_extract in one kernel.
263    pub(super) mla_fused_prefill_k: KernelHandle,
264    /// Split-K GEMM for skinny prefill matrices (M < 64).
265    pub(super) gemm_splitk_partial_k: KernelHandle,
266    pub(super) gemm_splitk_reduce_k: KernelHandle,
267    /// Tensor-core BF16 GEMM (m16n8k16 MMA).
268    pub(super) dense_gemm_tc_k: KernelHandle,
269    pub(super) paged_decode_splitk_k: Option<KernelHandle>,
270    pub(super) paged_decode_reduce_k: Option<KernelHandle>,
271    pub(super) residual_add_k: KernelHandle,
272    pub(super) sigmoid_gate_mul_k: KernelHandle,
273    pub(super) deinterleave_qg_k: KernelHandle,
274    pub(super) w4a16_gemv_qg_k: KernelHandle,
275    pub(super) residual_add_rms_norm_k: KernelHandle,
276    /// Dual-output (bf16 + f32) MoE-input norm for ATLAS_FP32_ROUTING. Zero if absent.
277    pub(super) residual_add_rms_norm_gatef32_k: KernelHandle,
278    // Kernels — batch2 (K=2 verify)
279    pub(super) w4a16_gemv_qg_batch2_k: KernelHandle,
280    pub(super) w4a16_gemv_dual_batch2_k: KernelHandle,
281    pub(super) w4a16_gemv_batch2_k: KernelHandle,
282    // Kernels — batch3 (K=3 verify)
283    pub(super) w4a16_gemv_qg_batch3_k: KernelHandle,
284    pub(super) w4a16_gemv_dual_batch3_k: KernelHandle,
285    pub(super) w4a16_gemv_batch3_k: KernelHandle,
286    /// Narrow `w4a16_gemv_batch{M}` family (M=4..8) for the K=4 verify and the
287    /// K=5..8 chain verify q/k/v/o projections. SSOT for the M -> tier
288    /// decision; individual tiers are 0-handles when the target lacks them.
289    pub(super) w4a16_batchm: W4a16BatchmTiers,
290    // Kernels — prefill (GEMM M=N + Flash Attention)
291    pub(super) w4a16_gemm_k: KernelHandle,
292    pub(super) w4a16_gemm_t_k: KernelHandle,
293    pub(super) w4a16_gemm_t_k64_k: KernelHandle,
294    /// K64 with a 64-wide N tile: same math, 2x the CTAs. `KernelHandle(0)`
295    /// when absent or killed by `ATLAS_NO_K64_N64`.
296    pub(super) w4a16_gemm_t_k64_n64_k: KernelHandle,
297    pub(super) w4a16_gemm_t_m128_k: KernelHandle,
298    /// LOSSLESS BF16-TC variant of t_m128 for QKV/o projection prefill (FP4→BF16
299    /// dequant + BF16 MMA, no FP8 activation crush). Opt-in via ATLAS_BF16_TC_PROJ
300    /// (default off → t_m128 path unchanged). KernelHandle(0) on miss.
301    pub(super) w4a16_gemm_t_m128_bf16_k: KernelHandle,
302    /// MiniMax-only shadow kernel.
303    pub(super) w4a16_gemm_t_m128_v2_k: KernelHandle,
304    /// v3 variant: K_STEP=64.
305    pub(super) w4a16_gemm_t_m128_v3_k: KernelHandle,
306    pub(super) dense_gemm_k: KernelHandle,
307    /// Tensor-core pipelined BF16 GEMM (mma.sync + cp.async, 128×128 tile) —
308    /// ~40× the scalar `dense_gemm_k` on large-M prefill projections, same math
309    /// (cosine 1.0). Used for the BF16-fallback Q/K/V/O projections (Holo's
310    /// native-FP8-dequant-to-BF16 attention path).
311    pub(super) dense_gemm_pipelined_k: KernelHandle,
312    pub(super) prefill_attn_k: KernelHandle,
313    /// HDIM=512 contiguous prefill for Gemma-4 full-attention layers
314    pub(super) prefill_attn_512_k: KernelHandle,
315    /// Did `prefill_attn_512_k` resolve to the TENSOR-CORE instantiation, or the
316    /// scalar reference? Only affects the profile label — but that label has now
317    /// been wrong twice, each time sending an investigation at the wrong kernel,
318    /// so which one ran is recorded rather than assumed.
319    pub(super) prefill_attn_512_is_tc: bool,
320    /// DeepSeek-V4 CSA compressor: window softmax-gated KV compression.
321    pub(super) csa_compress_k: KernelHandle,
322    /// DeepSeek-V4 CSA prefill attention over [raw | compressed] KV + sink.
323    pub(super) prefill_attn_compressed_k: KernelHandle,
324    /// 4b: # compressed blocks prefill wrote to `mla.compressor.pool` for the
325    /// active sequence (= prefill_len / ratio). Decode's compressed arm attends
326    /// blocks `[0, this)`. AtomicU32 for interior mutability under prefill's
327    /// `&self`; V4 serves max_batch=1 so one counter suffices (inc-3: per-seq
328    /// tracking + decode-time append will grow this each boundary crossing).
329    pub(super) v4_comp_pool_filled: std::sync::atomic::AtomicU32,
330    /// 4b inc-3 decode-append state (V4 serves max_batch=1 → scalar per layer).
331    /// `prev_valid`: the CSA `prev_win` ring holds a real previous decode window
332    /// (false until the first decode append, and reset each prefill) — when false
333    /// the CSA append masks Ca (window-0 semantics). `decode_started`/`first_pos`:
334    /// the absolute position of the first decode token this sequence, used to skip
335    /// any prefill/decode straddle window whose ring slots aren't all decode-written
336    /// (that one block is left as prefill/zero — a documented seam, not corruption).
337    pub(super) v4_comp_prev_valid: std::sync::atomic::AtomicBool,
338    pub(super) v4_decode_started: std::sync::atomic::AtomicBool,
339    pub(super) v4_decode_first_pos: std::sync::atomic::AtomicU32,
340    /// HDIM=512 paged prefill (BF16 KV) for Gemma-4 chunked long-context prefill
341    pub(super) prefill_attn_paged_512_k: KernelHandle,
342    pub(super) prefill_attn_64_k: KernelHandle,
343    pub(super) prefill_attn_paged_k: KernelHandle,
344    pub(super) prefill_attn_paged_fp8_k: KernelHandle,
345    pub(super) prefill_attn_paged_nvfp4_k: KernelHandle,
346    pub(super) prefill_attn_paged_turbo4_k: KernelHandle,
347    // BR=64 variants for long-context prefill (q_len >= 256)
348    pub(super) prefill_attn_paged_64_k: KernelHandle,
349    pub(super) prefill_attn_paged_fp8_64_k: KernelHandle,
350    pub(super) prefill_attn_paged_nvfp4_64_k: KernelHandle,
351    pub(super) prefill_attn_paged_turbo2_64_k: KernelHandle,
352    pub(super) prefill_attn_paged_turbo3_64_k: KernelHandle,
353    pub(super) prefill_attn_paged_turbo4_64_k: KernelHandle,
354    pub(super) prefill_attn_paged_turbo8_64_k: KernelHandle,
355    // ── TurboQuant+ asymmetric BR=64 prefill kernels ──
356    // Combined-dtype kernels that read K and V with different on-disk layouts.
357    // Currently: Bf16K + Turbo3V (safer-asym variant — K kept at bf16 precision,
358    // V aggressively compressed to 3-bit Lloyd-Max + FP8 group scale).
359    pub(super) prefill_attn_paged_bf16k_turbo3v_64_k: KernelHandle,
360    pub(super) prefill_attn_paged_bf16k_turbo4v_64_k: KernelHandle,
361    pub(super) prefill_attn_paged_bf16k_turbo2v_64_k: KernelHandle,
362    // Fp8K + TurboNV variants — same shape as bf16k_turbo*v_64 but threads
363    // the FP8 K-side per-tensor `k_scale` through to the dequant in
364    // LOAD_K_TILE. Targets FP8-attention models (Qwen3.6-35B-FP8 etc.).
365    pub(super) prefill_attn_paged_fp8k_turbo3v_64_k: KernelHandle,
366    pub(super) prefill_attn_paged_fp8k_turbo4v_64_k: KernelHandle,
367    pub(super) prefill_attn_paged_fp8k_turbo2v_64_k: KernelHandle,
368    // Both-sides-quantized TurboQuant+ asym (K and V both turbo, separate
369    // pool strides). K-side WHT bookend + Q WHT both fire because K is turbo.
370    pub(super) prefill_attn_paged_turbo4k_turbo3v_64_k: KernelHandle,
371    pub(super) prefill_attn_paged_turbo4k_turbo8v_64_k: KernelHandle,
372    pub(super) prefill_attn_paged_turbo3k_turbo8v_64_k: KernelHandle,
373    // ── Q12 Phase 3: same-chunk-len batched paged-prefill kernels ──
374    // Each takes `const int* const* block_table_ptrs` + per-batch Q/O
375    // offsets. Used by `Qwen3AttentionLayer::prefill_batched` when N≥2
376    // streams share the same chunk_len. Null on targets that don't
377    // carry the corresponding kernel (e.g. CPU backend).
378    pub(super) prefill_attn_paged_batched_k: KernelHandle,
379    pub(super) prefill_attn_paged_fp8_batched_k: KernelHandle,
380    pub(super) prefill_attn_paged_nvfp4_batched_k: KernelHandle,
381    pub(super) prefill_attn_paged_batched_64_k: KernelHandle,
382    pub(super) prefill_attn_paged_fp8_batched_64_k: KernelHandle,
383    pub(super) prefill_attn_paged_nvfp4_batched_64_k: KernelHandle,
384    // Batched prefill kernels
385    pub(super) deinterleave_qg_split_k: KernelHandle,
386    pub(super) deinterleave_qg_split_qnorm_k: KernelHandle,
387    pub(super) deinterleave_qg_split_qnorm_mrope_k: KernelHandle,
388    pub(super) sigmoid_gate_mul_batched_k: KernelHandle,
389    // Pre-dequanted FP8 weights for zero-overhead prefill GEMMs
390    pub(super) q_fp8: Option<DevicePtr>,
391    pub(super) k_fp8: Option<DevicePtr>,
392    pub(super) v_fp8: Option<DevicePtr>,
393    pub(super) o_fp8: Option<DevicePtr>,
394    pub(super) fp8_gemm_k: KernelHandle,
395    // FP8×FP8 GEMM
396    pub(super) bf16_to_fp8_k: KernelHandle,
397    pub(super) fp8_fp8_gemm_k: KernelHandle,
398    // M128 variants
399    pub(super) fp8_gemm_t_m128_k: KernelHandle,
400    pub(super) fp8_fp8_gemm_t_m128_k: KernelHandle,
401    // Native FP4 prefill (mxf4nvf4): present only for models whose kernel dir
402    // ships w4a4_gemm_mfast (try_kernel returns 0 elsewhere).
403    pub(super) w4a4_gemm_k: KernelHandle,
404    pub(super) quantize_nvfp4_k: KernelHandle,
405    /// Online FP8 KV scale calibration.
406    pub(super) fp8_calibration: Option<Fp8KvCalibration>,
407}
spark_model/layers/qwen3_attention/types.rs

spark_model/layers/qwen3_attention/
types.rs
