metrale_model_layers/layers/qwen3_ssm/layer_struct.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The `Qwen3SsmLayer` struct: the GDN layer's weights, kernel
4//! handles and per-layer bindings.
5//!
6//! Owner: model-layers (qwen3_ssm).
7//! Invariants: none beyond the types.
8
9use crate::layers::FfnComponent;
10use crate::layers::ops;
11use crate::layers::w4a16_gemv_tiers::W4a16BatchmTiers;
12use crate::weight_map::{DenseWeight, Fp8Weight, QuantizedWeight, SsmWeights};
13use metrale_gpu_runtime::gpu::{DevicePtr, KernelHandle};
14
15/// 2026-09-25: The Gated DeltaNet layer. Its QKVZ projection comes in two
16/// layouts: interleaved per key-head group, which `deinterleave_qkvz` (or the
17/// fused `w4a16_gemv_qkvz`) turns into `[Q | K | V | Z]`, and sequential
18/// (`sequential_qkvz`), already in that order.
19#[allow(dead_code)]
20pub struct Qwen3SsmLayer {
21 /// 2026-09-25: mHC highway weights (`set_hc_weights`); `Some` routes the
22 /// entry points to the hc paths (see `hc.rs`).
23 pub(crate) hc: Option<crate::layers::qwen3_attention::HcWeights>,
24 /// 2026-09-25: PLE n-gram injection (`set_ple`), run by the hc paths before
25 /// `hc_pre_site`.
26 pub(crate) ple: Option<crate::layers::ple::PleLayer>,
27 /// 2026-09-25: `hyper_connection::hc_pre`. It and the two `hc_*` handles
28 /// below are 0 unless `config.hc_mult > 0` (`hc_kernel`).
29 pub(super) hc_pre_k: KernelHandle,
30 pub(super) hc_post_k: KernelHandle,
31 /// 2026-09-25: `hc_expand`, launched only on the first model layer
32 /// (`HcWeights::is_first_model_layer`).
33 pub(super) hc_expand_k: KernelHandle,
34 pub(super) input_norm: DenseWeight,
35 pub(super) ssm: SsmWeights,
36 pub(super) post_attn_norm: DenseWeight,
37 pub(super) ffn: FfnComponent,
38 /// 2026-09-25: GDN `out_proj` LoRA delta and its kernels
39 /// (`set_out_proj_lora`); `None` makes `apply_lora_out_proj` a no-op.
40 pub(super) lora_out_proj: Option<(
41 crate::layers::ops::lora_delta::LoraPair,
42 crate::layers::ops::lora_delta::LoraKernels,
43 )>,
44 pub(super) qkvz_nvfp4: Option<QuantizedWeight>,
45 pub(super) qkvz_nvfp4_t: Option<QuantizedWeight>,
46 pub(super) out_proj_nvfp4_t: Option<QuantizedWeight>,
47 pub out_proj_dense: Option<DenseWeight>,
48 // 2026-09-25: Block-scaled FP8 weights (`set_fp8_decode_weights`).
49 pub(super) qkvz_fp8w: Option<Fp8Weight>,
50 pub(super) out_proj_fp8w: Option<Fp8Weight>,
51 /// 2026-09-25: Per-row FP8 weights (`set_fp8_rowwise_prefill_weights`),
52 /// read only by the row-wise prefill arms. They are kept apart from
53 /// `qkvz_fp8w` / `out_proj_fp8w` because those fields' readers index the
54 /// scale as a `[N/128, K/128]` grid, and a per-row scale read that way gives
55 /// wrong values without faulting.
56 pub(super) qkvz_fp8w_rowwise: Option<Fp8Weight>,
57 pub(super) out_proj_fp8w_rowwise: Option<Fp8Weight>,
58 /// 2026-09-28: W8A8 QKV|Z and out_proj (`set_w8a8_decode_weights`), for a
59 /// checkpoint that declares them FP8 W8A8; run ahead of every other decode
60 /// arm at 1..=64 rows (`w8a8_decode.rs`).
61 pub(crate) w8a8: Option<crate::layers::W8a8Mixer>,
62 /// 2026-09-25: Addresses of this layer's two slices of the arena slab
63 /// `ssm_rowwise_w_bf16`, 0 until the first row-wise prefill carves them
64 /// (`rowwise_bf16.rs`).
65 pub(super) qkvz_rowwise_bf16: std::sync::atomic::AtomicU64,
66 pub(super) out_proj_rowwise_bf16: std::sync::atomic::AtomicU64,
67 /// 2026-09-25: Packed Q2_0 fused `in_proj_qkvz` (`set_packed_q2_qkvz`).
68 /// Its only installer, the `qwen35_dense` loader, keeps `out_proj` NVFP4.
69 pub(super) qkvz_q2: Option<crate::weight_map::PackedQ2Weight>,
70 /// 2026-09-25: Packed-Q2 kernels: `q2_0_gemv_vec` for decode, and for
71 /// `qkvz_q2_prefill_gemm` the per-call BF16 dequant, the two MMQ GEMMs and
72 /// the q8_1 activation quantizer. The last three are 0 until
73 /// `set_packed_q2_qkvz` looks them up.
74 pub(super) q2_0_gemv_k: KernelHandle,
75 pub(super) dequant_q2_0_gn_k: KernelHandle,
76 pub(super) q2_0_mmq_nc_k: KernelHandle,
77 pub(super) q2_0_mmq_wc_k: KernelHandle,
78 pub(super) q4k_quant_act_k: KernelHandle,
79 /// 2026-09-25: The QKVZ weight is already in `[Q | K | V | Z]` row order
80 /// (`new_sequential`); no path launches `deinterleave_qkvz` then.
81 pub(super) sequential_qkvz: bool,
82 /// 2026-09-25: SM count from `GpuBackend::sm_count` at construction;
83 /// `ms_proj_gemm` compares grid widths with it.
84 pub(super) sm_count: u32,
85 pub(super) rms_norm_residual_k: KernelHandle,
86 pub(super) gated_rms_norm_k: KernelHandle,
87 pub(super) gated_rms_norm_f32_k: KernelHandle,
88 /// 2026-09-25: `gated_rms_norm_f32_input_strided`, one launch for all
89 /// sequences in `ssm_batched_recurrent.rs`. 0 when absent and on a
90 /// `gdn_norm_sigmoid` model; that file then loops per sequence.
91 pub(super) gated_rms_norm_f32_strided_k: KernelHandle,
92 pub(super) dense_gemv_k: KernelHandle,
93 /// 2026-09-25: `dense_gemv_bf16_batch2`, the two-row GEMV over the BF16
94 /// `in_proj_qkvz` in the batched decode projection.
95 pub(super) dense_gemv_batch2_k: KernelHandle,
96 pub(super) w4a16_gemv_k: KernelHandle,
97 /// 2026-09-25: `w4a16_gemv_sw`; `w4a16_decode_gemv` picks it over
98 /// `w4a16_gemv_k` only when `use_gemv_sw` accepts the lever and the handle.
99 pub(super) w4a16_gemv_sw_k: KernelHandle,
100 pub(super) w8a16_gemv_k: KernelHandle,
101 pub(super) w4a16_gemv_qkvz_k: KernelHandle,
102 pub(super) deinterleave_k: KernelHandle,
103 pub(super) conv1d_k: KernelHandle,
104 pub(super) conv1d_l2norm_k: KernelHandle,
105 pub(super) conv1d_l2norm_f32_k: KernelHandle,
106 /// 2026-09-25: `causal_conv1d_update_l2norm_f32_strided`, one launch for all
107 /// sequences in `ssm_batched_recurrent.rs`; with 0 that file loops per
108 /// sequence.
109 pub(super) conv1d_l2norm_f32_strided_k: KernelHandle,
110 pub(super) gdn_k: KernelHandle,
111 pub(super) gdn_f32_k: KernelHandle,
112 pub(super) gdn_f32_norm_k: KernelHandle,
113 pub(super) gdn_f32_conv_norm_k: KernelHandle,
114 pub(super) gdn_f32_strided_k: KernelHandle,
115 pub(super) gdn_f32_strided_norm_k: KernelHandle,
116 pub(super) gdn_f32_strided_norm_half_k: KernelHandle,
117 /// 2026-09-25: Shared-memory-staged strided decode kernel, selected in
118 /// `ssm_batched_recurrent.rs` only when kd == vd == 128 and
119 /// `METRALE_GDN_SMEM_STAGE` is set (any value).
120 pub(super) gdn_f32_strided_norm_smem_k: KernelHandle,
121 /// 2026-09-25: FP16 h-state variant of `gdn_f32_strided_norm_half_k`; the
122 /// strided recurrence in `ssm_batched_recurrent.rs` launches it, and no
123 /// FP32 variant, under the FP16 h-state.
124 pub(super) gdn_f16_strided_norm_half_k: KernelHandle,
125 /// 2026-09-25: FP16 h-state variant of `gdn_f32_norm_k`, used by
126 /// `ssm_forward` and the per-sequence arm of `ssm_batched_recurrent.rs`.
127 pub(super) gdn_f16_norm_k: KernelHandle,
128 pub(super) ba_gates_k: KernelHandle,
129 pub(super) residual_add_k: KernelHandle,
130 pub(super) l2_norm_k: KernelHandle,
131 pub(super) residual_add_rms_norm_k: KernelHandle,
132 /// 2026-09-25: `residual_add_rms_norm_gatef32`; single-token decode uses it
133 /// when FP32 routing is active for the FFN and the handle is non-zero.
134 pub(super) residual_add_rms_norm_gatef32_k: KernelHandle,
135 pub(super) gated_rms_norm_prefill_k: KernelHandle,
136 pub(super) w4a16_gemm_k: KernelHandle,
137 pub(super) w4a16_gemm_t_k: KernelHandle,
138 pub(super) w4a16_gemm_t_k64_k: KernelHandle,
139 /// 2026-09-25: `w4a16_gemm_t_k64_n64_p3`, the narrow-N deep-K GEMM; 0 when
140 /// absent or when `METRALE_NO_K64_N64` is set (any value).
141 pub(super) w4a16_gemm_t_k64_n64_k: KernelHandle,
142 pub(super) w4a16_gemm_t_m128_k: KernelHandle,
143 pub(super) w4a16_gemm_t_m128_v2_k: KernelHandle,
144 pub(super) w4a16_gemv_batch2_k: KernelHandle,
145 pub(super) dense_gemm_k: KernelHandle,
146 pub(super) dense_gemm_pipelined_k: KernelHandle,
147 pub(super) gdn_prefill_k: KernelHandle,
148 pub(super) gdn_prefill_split_k: KernelHandle,
149 pub(super) gdn_prefill_split4_k: KernelHandle,
150 pub(super) gdn_prefill_persistent_k: KernelHandle,
151 pub(super) gdn_prefill_persistent_wy4_k: KernelHandle,
152 /// 2026-09-25: `gated_delta_rule_prefill_regresident`, used by
153 /// `trait_prefill_recur.rs` when the `gdn_regresident` lever is on (the
154 /// default; `METRALE_NO_GDN_REGRESIDENT=1` turns it off), kd == vd == 128
155 /// and the handle is non-zero.
156 pub(super) gdn_prefill_regresident_k: KernelHandle,
157 /// 2026-09-25: `gated_delta_rule_recompute_wu`, the first kernel of the FLA
158 /// chunked prefill (`ops::gdn_prefill_fla`), which then runs a state spine
159 /// and `chunk_fwd_o`.
160 pub(super) gdn_prefill_fla_recompute_wu_k: KernelHandle,
161 /// 2026-09-25: Hopper twins of `recompute_wu` and `chunk_fwd_o`
162 /// (`init_kernels::prefill_wu_hopper_k`, `prefill_fwd_o_hopper_k`); 0 on
163 /// targets without the Hopper sources, where the parents run.
164 pub(super) gdn_prefill_fla_recompute_wu_hopper_k: KernelHandle,
165 pub(super) gdn_prefill_fla_chunk_fwd_o_hopper_k: KernelHandle,
166 pub(super) gdn_prefill_fla_chunk_delta_h_k: KernelHandle,
167 /// 2026-09-25: `gated_delta_rule_chunk_delta_h_tc_vblock`;
168 /// `ops::gdn_prefill_fla` uses it only when the fused spine is not used and
169 /// `METRALE_GDN_TC_VBLOCK=1`.
170 #[allow(dead_code)]
171 pub(super) gdn_prefill_fla_chunk_delta_h_tc_vblock_k: KernelHandle,
172 /// 2026-09-25: The tensor-core state spine, `ops::GDN_TC_SPINE_ENTRY`
173 /// (`init_kernels::gdn_prefill_tc_kernel`); 0 unless the target default
174 /// `gdn_prefill_tc` is on and the target compiles the module.
175 pub(super) gdn_prefill_fla_chunk_delta_h_tcfuse_k: KernelHandle,
176 /// 2026-09-25: The scalar fused state spine that
177 /// `init_kernels::fused_spine_kernel` picks.
178 pub(super) gdn_prefill_fla_chunk_delta_h_fused_k: KernelHandle,
179 /// 2026-09-25: `gated_delta_rule_chunk_delta_h_tma`, requested by
180 /// `METRALE_GDN_TMA=1` (`ops::ssm_gdn_a3`); 0 when absent.
181 pub(super) gdn_prefill_fla_chunk_delta_h_tma_k: KernelHandle,
182 pub(super) gdn_prefill_fla_chunk_fwd_o_k: KernelHandle,
183 /// 2026-09-25: `gated_delta_rule_prefill_wy64` (module
184 /// `gated_delta_rule_wy64_prefill`), which takes WY chunks of 32 tokens
185 /// with H in shared memory.
186 pub(super) gdn_prefill_wy32_k: KernelHandle,
187 // 2026-09-25: Multi-stream GDN prefill kernels for
188 // `trait_prefill_gdn/batched.rs`; 0 when the target lacks them.
189 pub(super) gdn_prefill_wy32_batched_k: KernelHandle,
190 pub(super) gdn_prefill_persistent_batched_k: KernelHandle,
191 pub(super) gdn_prefill_persistent_wy4_batched_k: KernelHandle,
192 pub(super) gdn_prefill_split4_batched_k: KernelHandle,
193 pub(super) compute_gdn_gates_k: KernelHandle,
194 pub(super) ba_gates_prefill_k: KernelHandle,
195 /// 2026-09-25: Hopper twin of `ba_gates_prefill_k`
196 /// (`init_kernels::ba_gates_hopper_k`); `ops::ba_gates_pick` chooses
197 /// between the two per launch.
198 pub(super) ba_gates_prefill_hopper_k: KernelHandle,
199 pub(super) conv1d_prefill_k: KernelHandle,
200 /// 2026-09-25: `causal_conv1d_update_prefill_tp`; 0 when absent.
201 pub(super) conv1d_prefill_tp_k: KernelHandle,
202 pub(super) gdn_chunk2_k: KernelHandle,
203 pub(super) conv1d_chunk2_k: KernelHandle,
204 pub(super) gdn_chunk3_k: KernelHandle,
205 pub(super) w4a16_gemv_batch3_k: KernelHandle,
206 // 2026-09-25: The `w4a16_gemv_batch{4..8}` tiers, and below them
207 // `w4a16_gemv_batch16`.
208 pub(super) w4a16_batchm: W4a16BatchmTiers,
209 pub(super) w4a16_gemv_batch16_k: KernelHandle,
210 pub(super) gdn_wy2_k: KernelHandle,
211 /// 2026-09-25: Register-resident variant of `gdn_wy2_k`. `wy2_kernel` picks
212 /// it when kd == vd == 128, `n >= wy_resident_min_width()`, the handle is
213 /// non-zero and `METRALE_NO_GDN_WY2_RESIDENT` is unset, and logs its choice
214 /// at the first dispatch.
215 pub(super) gdn_wy2_resident_k: KernelHandle,
216 pub(super) gdn_wy3_k: KernelHandle,
217 /// 2026-09-25: Register-resident variant of `gdn_wy3_k`, picked by
218 /// `wy3_kernel` under the same conditions with
219 /// `METRALE_NO_GDN_WY3_RESIDENT`.
220 pub(super) gdn_wy3_resident_k: KernelHandle,
221 pub(super) gdn_wy4_k: KernelHandle,
222 /// 2026-09-25: Write-on-accept K=4 verify kernel, with its fold and
223 /// flag-clear kernels below (`woa::woa_kernels`, each 0 when absent).
224 /// `woa_decision` runs it only when all three resolved, the call requests
225 /// it, `METRALE_GDN_WOA=1`, the h-state is FP32, kd == vd == 128 and a stash
226 /// is bound for at least `n` sequences.
227 pub(super) gdn_wy4_woa_k: KernelHandle,
228 pub(super) gdn_wy4_fold_k: KernelHandle,
229 pub(super) gdn_wy4_clear_k: KernelHandle,
230 /// 2026-09-25: Set by `gdn_woa_bind`, 0 until then: the device address of
231 /// the engaged word, this layer's stash slab, and the number of sequences
232 /// the slab holds.
233 pub(super) woa_flag: std::sync::atomic::AtomicU64,
234 pub(super) woa_stash: std::sync::atomic::AtomicU64,
235 pub(super) woa_seqs: std::sync::atomic::AtomicUsize,
236 pub(super) woa_dims: [usize; 4],
237 /// 2026-09-26: Carried-state verify handles and binding (`carry.rs`).
238 pub(super) carry: super::carry::CarryState,
239 /// 2026-09-25: FP16 h-state variants of the WY verify kernels, 0 when
240 /// absent. Under `ssm_h_fp16_enabled()`, `wy2_kernel` and `wy3_kernel`
241 /// return these, even when 0, and never the FP32 kernels.
242 pub(super) gdn_wy2_f16_k: KernelHandle,
243 pub(super) gdn_wy2_resident_f16_k: KernelHandle,
244 pub(super) gdn_wy3_f16_k: KernelHandle,
245 pub(super) gdn_wy3_resident_f16_k: KernelHandle,
246 pub(super) gdn_wy4_f16_k: KernelHandle,
247 /// 2026-09-25: The h-state width converters (`ssm_h_dtype`). Under the
248 /// 2-byte-sized pool, `prefill_h_begin` / `prefill_h_end` (`ssm_h_fp16.rs`)
249 /// widen a sequence's slot into its FP32 staging buffer before the FP32
250 /// prefill kernels and narrow it back after; a 0 handle there is an error.
251 pub(super) ssm_h_f16_to_f32_k: KernelHandle,
252 pub(super) ssm_h_f32_to_f16_k: KernelHandle,
253 /// 2026-09-25: Fused K=2 verify conv and norm kernels, used only when both
254 /// are non-zero and `METRALE_GDN_FUSED_VERIFY=1` (`fused_verify_k2_enabled`).
255 pub(super) gdn_verify_fused_conv_k2_k: KernelHandle,
256 pub(super) gdn_verify_fused_norm_k2_k: KernelHandle,
257 /// 2026-09-25: Fused verify conv of the wyN / wy17 arm
258 /// (`decode_batched_conv_gdn_wyn`), used when non-zero, the conv
259 /// intermediates are contiguous and `METRALE_GDN_FUSED_CONV17` is not `0`;
260 /// otherwise that arm runs the conv per token.
261 pub(super) gdn_verify_fused_conv_kn_k: KernelHandle,
262 /// 2026-09-25: Multi-sequence variant of the above, for
263 /// `trait_decode_batched_conv_gdn_multi.rs`; 0 when absent.
264 pub(super) gdn_verify_fused_conv_kn_batched_k: KernelHandle,
265 /// 2026-09-25: Exact-verify kernels that write the per-token h-state
266 /// snapshot inline, and the FP32-output fused verify conv. With a 0 handle,
267 /// `trait_decode_batched_conv_gdn_exact.rs` copies snapshots with
268 /// `copy_d2d_async` and `_multi_exact.rs` declines to the per-sequence loop.
269 pub(super) gdn_f32_norm_snap_k: KernelHandle,
270 pub(super) gdn_f32_strided_norm_snap_k: KernelHandle,
271 pub(super) gdn_verify_fused_conv_kn_f32_k: KernelHandle,
272 /// 2026-09-25: `gated_delta_rule_wy17`, used for K = 17 when non-zero, the
273 /// `gdn_wy17` lever is on and the h-state is FP32; otherwise K = 17 takes
274 /// the sequential path.
275 pub(super) gdn_wy17_k: KernelHandle,
276 /// 2026-09-25: The K = 5..16 verify kernels (`init_kernels::wyn_kernels`),
277 /// index = K - 5, read by `wyn_kernel` when the `gdn_wyn` lever is on
278 /// (the default; `METRALE_GDN_WYN=0` turns it off).
279 pub(super) gdn_wyn_k: [KernelHandle; 12],
280 /// 2026-09-25: FP16 h-state variants of `gdn_wyn_k`, same index. Under the
281 /// FP16 h-state a 0 entry makes `decode_batched_conv_gdn` return an error
282 /// rather than run an FP32 kernel on the FP16
283 /// state). provenance-id: 526f6e616c6420522e205374657369616b
284 pub(super) gdn_wyn_f16_k: [KernelHandle; 12],
285 /// 2026-09-25: Pointer-table variants for the cross-sequence batched verify
286 /// (`init_kernels::wyn_table_kernels`, read by `wyn_table_kernel`), index =
287 /// K - 5.
288 pub(super) gdn_wyn_table_k: [KernelHandle; 12],
289 /// 2026-09-25: FP16 variants of `gdn_wyn_table_k`, same index.
290 pub(super) gdn_wyn_f16_table_k: [KernelHandle; 12],
291 // 2026-09-25: Per-sequence FP32 state sizes in bytes: `nv * kd * vd` values
292 // of h, `conv_dim * d_conv` of conv state.
293 pub(super) h_state_bytes: usize,
294 pub(super) conv_state_bytes: usize,
295 // 2026-09-25: Unscaled FP8 weights for the `fp8_gemm_n128` arms
296 // (`set_fp8_prefill_only_weights`, `predequant_for_prefill`).
297 pub(super) qkvz_fp8: Option<DevicePtr>,
298 pub(super) out_proj_fp8: Option<DevicePtr>,
299 pub(super) fp8_gemm_k: KernelHandle,
300 pub(super) fp8_gemm_t_m128_k: KernelHandle,
301 // 2026-09-25: Block-scaled W8A16 GEMMs, which the prefill tries before the
302 // unscaled `fp8_gemm_n128` arm.
303 pub(super) w8a16_gemm_k: KernelHandle,
304 // 2026-09-25: `w8a16_gemm_pipelined`; no environment variable gates it.
305 pub(super) w8a16_gemm_pipelined_k: KernelHandle,
306 // 2026-09-25: 32-row M-tile variant; `ops::w8a16_gemm_pipelined_by_m`
307 // chooses between it and `w8a16_gemm_pipelined` by row count.
308 pub(super) w8a16_gemm_pipelined_m32_k: KernelHandle,
309 // 2026-09-26: 64-row M-tile variant for 33..=64 rows, same module.
310 pub(super) w8a16_gemm_pipelined_m64_k: KernelHandle,
311 // 2026-09-25: `w8a16_gemv_batch4` (M <= 4) and, below it,
312 // `w8a16_gemv_batch16`, both from module `w8a16_gemv_batch4`.
313 pub(super) w8a16_gemv_batch4_k: KernelHandle,
314 pub(super) w8a16_gemv_batch16_k: KernelHandle,
315 pub(super) w8a16_gemm_t_k: KernelHandle,
316 // 2026-09-25: W8A8 prefill: `per_token_group_quant_fp8` writes FP8
317 // activations with one FP32 scale per token per 128 columns, and
318 // `fp8_gemm_t_blockscaled` applies both scale sets in an FP32 epilogue.
319 // The QKVZ arm runs unless `METRALE_FP8_SINGLE_SCALE=1`; the `out_proj`
320 // arm only under `METRALE_FP8_W8A8=1`.
321 pub(super) per_token_group_quant_fp8_k: ops::Fp8ActQuant,
322 pub(super) fp8_gemm_t_blockscaled_k: KernelHandle,
323 /// 2026-09-25: `fp8_act_scale_to_kmajor`: copies the quantizer's
324 /// `[M, K/128]` activation scales into the `[K/128, ceil16(M)]` layout the
325 /// cuBLASLt arms read. With 0, both cuBLASLt W8A8 prefill arms decline under
326 /// the default k-major layout.
327 pub(super) fp8_act_scale_kmajor_k: KernelHandle,
328}