spark_model/layers/qwen3_ssm/layer_struct.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! The `Qwen3SsmLayer` struct itself.
4//!
5//! Split out of `mod.rs` under the 500-LoC cap: the struct and its field
6//! documentation are 417 of that file's 506 lines, and the 2026-09-04
7//! batched-verify line pushed it over. Moving the declaration leaves `mod.rs`
8//! as what it should be -- the module's table of contents -- and is an exact
9//! piecewise copy of the declaration -- no field, attribute or doc comment is
10//! changed -- with ONE deliberate adjustment: fields that carried no
11//! visibility modifier are now `pub(super)`. That is not a widening. A private
12//! field in `mod.rs` was visible to every descendant of `qwen3_ssm`, which is
13//! where all the sibling impl files live; private in a CHILD module it would
14//! be visible to nobody but this file. `pub(super)` from here reaches exactly
15//! the set it reached before.
16
17use crate::layers::FfnComponent;
18use crate::layers::ops;
19use crate::layers::w4a16_gemv_tiers::W4a16BatchmTiers;
20use crate::weight_map::{DenseWeight, Fp8Weight, QuantizedWeight, SsmWeights};
21use spark_runtime::gpu::{DevicePtr, KernelHandle};
22
23/// Qwen3-Next SSM/GDN layer (36 of 48 layers).
24///
25/// Supports two QKVZ projection modes:
26/// - **Interleaved** (80B): `w4a16_gemv_qkvz` or GEMV + `deinterleave_qkvz`
27/// - **Sequential** (3.5-35B): plain GEMV → `[Q|K|V|Z]` already in order
28#[allow(dead_code)]
29pub struct Qwen3SsmLayer {
30 /// mHC weights when the model carries a `hc_mult`-wide highway; see `hc`.
31 pub(crate) hc: Option<crate::layers::qwen3_attention::HcWeights>,
32 /// PLE n-gram injection. `Some` on exactly ONE model layer (layer 1 on
33 /// this checkpoint); it runs at the TOP of the mHC forward, before this
34 /// layer's own hyper-connection, matching the reference's
35 /// `hidden_states = hidden_states + self.ple(...)`.
36 pub(crate) ple: Option<crate::layers::ple::PleLayer>,
37 /// mHC kernel handles. Resolved only when `config.hc_mult > 0`, so a
38 /// plain GDN model issues no lookup and leaves no row in the startup
39 /// audit. See `qwen3_attention::init_arch_gates`.
40 pub(super) hc_pre_k: KernelHandle,
41 pub(super) hc_post_k: KernelHandle,
42 /// Seeds the highway on MODEL layer 0 — which on a 3:1 GDN:attention
43 /// interleave is a GDN layer, so this side owns the expand that the
44 /// attention side used to do.
45 pub(super) hc_expand_k: KernelHandle,
46 pub(super) input_norm: DenseWeight,
47 pub(super) ssm: SsmWeights,
48 pub(super) post_attn_norm: DenseWeight,
49 pub(super) ffn: FfnComponent,
50 /// GDN `out_proj` LoRA delta for this layer, with the kernels to apply it.
51 /// `None` on every base serve, which keeps the base path byte-identical.
52 pub(super) lora_out_proj: Option<(
53 crate::layers::ops::lora_delta::LoraPair,
54 crate::layers::ops::lora_delta::LoraKernels,
55 )>,
56 // NVFP4-quantized QKVZ weight (quarters bandwidth vs BF16)
57 pub(super) qkvz_nvfp4: Option<QuantizedWeight>,
58 // Transposed [K/2, N] copy for coalesced w4a16_gemm reads (prefill)
59 pub(super) qkvz_nvfp4_t: Option<QuantizedWeight>,
60 // Transposed out_proj for prefill GEMM
61 pub(super) out_proj_nvfp4_t: Option<QuantizedWeight>,
62 // BF16 out_proj for models where SSM weights are not pre-quantized
63 pub out_proj_dense: Option<DenseWeight>,
64 // FP8 E4M3 checkpoint weights for native FP8 serving (w8a16_gemv LUT kernel)
65 pub(super) qkvz_fp8w: Option<Fp8Weight>,
66 pub(super) out_proj_fp8w: Option<Fp8Weight>,
67 /// PER-ROW FP8 (`Fp8PerRow`) for PREFILL ONLY, from mixed-precision
68 /// compressed-tensors checkpoints (`AVAROK_FP8_ROWWISE=1`).
69 ///
70 /// Separate fields rather than reusing `qkvz_fp8w`/`out_proj_fp8w`, and
71 /// that separation is the safety property: those two are read by
72 /// `w8a16_gemv` in `ssm_forward.rs` and `trait_decode_batched.rs`, which
73 /// index the scale as a `[N/128, K/128]` block grid. A per-row buffer is
74 /// SMALLER than that index space, so it would not fault — it would return
75 /// plausible garbage. Only the row-wise cuBLASLt prefill arm reads these;
76 /// decode keeps the NVFP4 copy.
77 pub(super) qkvz_fp8w_rowwise: Option<Fp8Weight>,
78 pub(super) out_proj_fp8w_rowwise: Option<Fp8Weight>,
79 /// This layer's slice of the LEDGERED row-wise BF16-weight slab
80 /// (`BufferSizes::ssm_rowwise_w_bf16`), or 0 before the first prefill
81 /// through the matching arm carves and fills it. See `rowwise_bf16.rs`
82 /// for the #917 receipt that moved these bytes out of a by-pointer cache
83 /// of `gpu.alloc`s and into the arena.
84 ///
85 /// Per LAYER rather than per weight pointer, because that is what the
86 /// lifetime is: the slab is the arena's, and the arena outlives every
87 /// prefill this layer runs.
88 pub(super) qkvz_rowwise_bf16: std::sync::atomic::AtomicU64,
89 pub(super) out_proj_rowwise_bf16: std::sync::atomic::AtomicU64,
90 /// Tier-1c keep-packed ternary Q2_0 fused in_proj_qkvz (`AVAROK_GGUF_NATIVE_Q2`).
91 /// [Q|K|V|Z] rows byte-concatenated from packed `in_proj_qkv` (V-region
92 /// row-permuted) + `in_proj_z` (row-permuted) at load, so the 2-bit weight is
93 /// HF-correct. `out_proj` stays NVFP4 (column reorder not packed-permutable here).
94 pub(super) qkvz_q2: Option<crate::weight_map::PackedQ2Weight>,
95 /// Q2_0 kernels for the packed qkvz: `gemv` = `q2_0_gemv_vec` decode; `dequant`
96 /// = load-time packed→BF16 for the transient-dequant prefill fallback;
97 /// `mmq_{nc,wc}` = Tier-2 keep-packed tensor-core MMQ prefill (`KernelHandle(0)`
98 /// → fallback); `q4k_quant_act` = shared q8_1 activation quantizer.
99 pub(super) q2_0_gemv_k: KernelHandle,
100 pub(super) dequant_q2_0_gn_k: KernelHandle,
101 pub(super) q2_0_mmq_nc_k: KernelHandle,
102 pub(super) q2_0_mmq_wc_k: KernelHandle,
103 pub(super) q4k_quant_act_k: KernelHandle,
104 /// When true, QKVZ projection output is already sequential [Q|K|V|Z].
105 /// Skips the deinterleave kernel (used by Qwen3.5 where QKV+Z are
106 /// concatenated at load time rather than interleaved per-group).
107 pub(super) sequential_qkvz: bool,
108 /// Streaming multiprocessor count, read from the driver ONCE at
109 /// construction (`GpuBackend::sm_count`). `ms_proj_gemm` needs to know how
110 /// wide the machine is to decide whether halving the CTA rows is a saving
111 /// or an under-fill; a compiled-in constant would be wrong on every part
112 /// that is not the one it was tuned on.
113 pub(super) sm_count: u32,
114 // Kernels — decode path (single-token GEMV)
115 pub(super) rms_norm_residual_k: KernelHandle,
116 pub(super) gated_rms_norm_k: KernelHandle,
117 pub(super) gated_rms_norm_f32_k: KernelHandle,
118 /// `gated_rms_norm_f32_input_strided` — the SAME per-(sequence, head) math
119 /// as `gated_rms_norm_f32_k`, with `blockIdx.y` walking the sequences, so
120 /// a batched decode step spends ONE launch per layer instead of one per
121 /// row. 0 when absent (notably on a `gdn_norm_sigmoid` model, which has no
122 /// strided sigmoid twin), which keeps the per-seq loop. #927: the H100
123 /// batch-16 trace showed the per-seq loop at 768 launches / 1.61 ms.
124 pub(super) gated_rms_norm_f32_strided_k: KernelHandle,
125 pub(super) dense_gemv_k: KernelHandle,
126 /// K=2 verify: batched (M=2) BF16 GDN in_proj_qkvz — one weight pass for
127 /// both verify tokens instead of two M=1 `dense_gemv` reads.
128 pub(super) dense_gemv_batch2_k: KernelHandle,
129 pub(super) w4a16_gemv_k: KernelHandle,
130 /// Single-warp `w4a16_gemv_sw`. `KernelHandle(0)` on miss → base GEMV.
131 pub(super) w4a16_gemv_sw_k: KernelHandle,
132 pub(super) w8a16_gemv_k: KernelHandle,
133 pub(super) w4a16_gemv_qkvz_k: KernelHandle,
134 pub(super) deinterleave_k: KernelHandle,
135 pub(super) conv1d_k: KernelHandle,
136 pub(super) conv1d_l2norm_k: KernelHandle,
137 pub(super) conv1d_l2norm_f32_k: KernelHandle,
138 /// `conv1d_l2norm_f32_k` with explicit input/output row strides, letting
139 /// the concurrent-decode path batch all N sequences into one launch.
140 /// `KernelHandle(0)` on kernel sets that predate it — the multi-seq path
141 /// then falls back to the per-sequence conv loop.
142 pub(super) conv1d_l2norm_f32_strided_k: KernelHandle,
143 pub(super) gdn_k: KernelHandle,
144 pub(super) gdn_f32_k: KernelHandle,
145 pub(super) gdn_f32_norm_k: KernelHandle,
146 pub(super) gdn_f32_conv_norm_k: KernelHandle,
147 pub(super) gdn_f32_strided_k: KernelHandle,
148 pub(super) gdn_f32_strided_norm_k: KernelHandle,
149 /// Half-width register retention (k_dim==v_dim==128): retains the first 64 H
150 /// columns so the update re-reads only the rest (2R+1W -> 1.5R+1W).
151 pub(super) gdn_f32_strided_norm_half_k: KernelHandle,
152 /// SRAM-staged full retention (k_dim==v_dim==128): the columns the register
153 /// file cannot hold are staged in shared memory on the first pass instead of
154 /// being re-read from H (1.5R+1W -> 1.0R+1W). Bit-identical to
155 /// `gdn_f32_strided_norm_half_k` but measured throughput-NEUTRAL, so it is
156 /// OPT-IN via `gdn_smem_stage_enabled()` (`AVAROK_GDN_SMEM_STAGE`).
157 pub(super) gdn_f32_strided_norm_smem_k: KernelHandle,
158 /// FP16 h-state twin of `gdn_f32_strided_norm_half_k` (`AVAROK_SSM_H_FP16`).
159 /// Additive: it never replaces the FP32 kernel, it is selected instead of
160 /// it when the sequence's `SsmLayerState::h_is_f16` is set.
161 pub(super) gdn_f16_strided_norm_half_k: KernelHandle,
162 /// FP16 h-state twin of `gdn_f32_norm_k` — the per-sequence arm the batched
163 /// dispatch falls back to at n == 1 and whenever pool slots fragment out of
164 /// slice order. Without it the FP16 pool would be read as FP32 on exactly
165 /// those steps.
166 pub(super) gdn_f16_norm_k: KernelHandle,
167 pub(super) ba_gates_k: KernelHandle,
168 pub(super) residual_add_k: KernelHandle,
169 pub(super) l2_norm_k: KernelHandle,
170 pub(super) residual_add_rms_norm_k: KernelHandle,
171 /// Dual-output (bf16 + f32) MoE-input norm for AVAROK_FP32_ROUTING. Zero if absent.
172 pub(super) residual_add_rms_norm_gatef32_k: KernelHandle,
173 pub(super) gated_rms_norm_prefill_k: KernelHandle,
174 // Kernels — batched verification path (multi-token GEMM)
175 pub(super) w4a16_gemm_k: KernelHandle,
176 pub(super) w4a16_gemm_t_k: KernelHandle, // Transposed B layout [K/2, N] — K_STEP_T=32
177 pub(super) w4a16_gemm_t_k64_k: KernelHandle, // K64 variant: K_STEP_T=64, halves outer loop
178 /// K64 with a 64-wide N tile: same math, 2x the CTAs. `KernelHandle(0)`
179 /// when absent or killed by `AVAROK_NO_K64_N64`.
180 pub(super) w4a16_gemm_t_k64_n64_k: KernelHandle,
181 pub(super) w4a16_gemm_t_m128_k: KernelHandle, // M128 variant: 2 M-chunks per CTA, halves B re-reads
182 pub(super) w4a16_gemm_t_m128_v2_k: KernelHandle, // M128 8-warp pipelined (fast at small M; the FFN's kernel)
183 pub(super) w4a16_gemv_batch2_k: KernelHandle,
184 pub(super) dense_gemm_k: KernelHandle,
185 pub(super) dense_gemm_pipelined_k: KernelHandle,
186 pub(super) gdn_prefill_k: KernelHandle,
187 pub(super) gdn_prefill_split_k: KernelHandle,
188 pub(super) gdn_prefill_split4_k: KernelHandle,
189 pub(super) gdn_prefill_persistent_k: KernelHandle,
190 pub(super) gdn_prefill_persistent_wy4_k: KernelHandle,
191 /// Register-resident token-sequential warm-replay recurrence (H in regs, >=2
192 /// CTA/SM, no barriers). Token-equal to WY4 (cosine 1.0), ~2.9x faster.
193 /// DEFAULT-ON since 2026-07-25 (serve-validated: full MLPerf-edge e2e, wall
194 /// −7.25%, BFCL identical); kill switch `AVAROK_NO_GDN_REGRESIDENT=1`.
195 pub(super) gdn_prefill_regresident_k: KernelHandle,
196 /// FLA multi-kernel chunked prefill (baked default for 128-dim GDN): recompute_wu →
197 /// chunk_delta_h_ksplit (k-split occupancy) → chunk_fwd_o. 1.75x vs wy4 @16k,
198 /// token-equal (cos=1.0 vs scalar). Three handles; all must be non-null.
199 pub(super) gdn_prefill_fla_recompute_wu_k: KernelHandle,
200 /// Hopper twins of the two SCALAR REMNANTS of the FLA prefill (#928):
201 /// `gdn_recompute_wu_hopper.cu`'s blocked triangular solve on tensor cores
202 /// and `gdn_fwd_o_hopper.cu`'s masked `tril(kq).uc` square. They live only
203 /// under `kernels/hopper`, so `try_kernel` gives 0 everywhere else and the
204 /// launcher then runs the unchanged parents. Selected by the family lever
205 /// `[defaults] gdn_prefill_tc` (`AVAROK_GDN_PREFILL_TC` overriding), which
206 /// Hopper ships ON since round 13; `AVAROK_NO_GDN_PREFILL_TC_REMNANTS=1` pins them
207 /// off while keeping the tensor-core state spine, which is the A/B that
208 /// separates the three kernels. The nsys receipt that motivates them is in
209 /// `GDN-PREFILL-ATTRIBUTION.md`: 5.5% and 4.0% of a 1193-token H100
210 /// prefill, both with their big matmuls already on `mma.sync` and their
211 /// remainder — a triangular product on 128 of 512 threads, and two forward
212 /// substitutions worth 79-85% of their kernel — still scalar.
213 pub(super) gdn_prefill_fla_recompute_wu_hopper_k: KernelHandle,
214 pub(super) gdn_prefill_fla_chunk_fwd_o_hopper_k: KernelHandle,
215 pub(super) gdn_prefill_fla_chunk_delta_h_k: KernelHandle,
216 /// Tensor-core / DV-block-split variant of the FLA chunk_delta_h spine
217 /// (`gated_delta_rule_chunk_delta_h_tc_vblock`). Loaded by default but not
218 /// yet wired into the prefill dispatch — the cos-gate validates it in
219 /// isolation first. `allow(dead_code)` until the launch site reads it.
220 #[allow(dead_code)]
221 pub(super) gdn_prefill_fla_chunk_delta_h_tc_vblock_k: KernelHandle,
222 /// TENSOR-CORE chunked-prefill state spine
223 /// (`gated_delta_rule_chunk_tc::gated_delta_rule_chunk_delta_h_tcfuse`),
224 /// behind `[defaults] gdn_prefill_tc` — ON for `kernels/hopper` since round
225 /// 13, OFF elsewhere, with `AVAROK_GDN_PREFILL_TC` overriding either way.
226 /// Both per-chunk
227 /// products run on `mma.sync.m16n8k16` with bf16 operands and an f32
228 /// accumulator that IS the recurrent state; `h` stays f32 in memory. The
229 /// nsys receipt that motivates it is in `GDN-PREFILL-ATTRIBUTION.md`
230 /// (#928): the shipped scalar spine is 26.6%/32.3% of the 1193/4593-token
231 /// H100 prefill at 3.7 TFLOP/s, i.e. latency-bound at 4.5% warp residency.
232 /// `try_kernel` => 0 on images without it, and the launcher additionally
233 /// refuses any head/chunk that differs from the compile-time tile.
234 pub(super) gdn_prefill_fla_chunk_delta_h_tcfuse_k: KernelHandle,
235 /// Warp-dense fused GDN state spine (`gated_delta_rule_chunk_delta_h_vtile`).
236 /// 512 threads = 16 warps/CTA against ksplit's 8, with the SAME grid (one CTA
237 /// per head) so `W`/`K` global loads are not duplicated — an ncu profile put
238 /// ksplit at L2 60.6% / L1 56.2%, i.e. memory-pipeline bound, which is why the
239 /// DV-split variants (which duplicate those loads) lost. Fusing the two
240 /// per-chunk passes deletes ksplit's `duc[CHUNK]` register array, and that is
241 /// what pays for the extra warps: 118 registers, no spills. Measured 2.15-2.18x
242 /// vs ksplit at 2048/8192/16384 with cos=1.0000 (`gdn_chunk_shapetest`).
243 pub(super) gdn_prefill_fla_chunk_delta_h_fused_k: KernelHandle,
244 /// TMA (`cp.async.bulk.tensor`) build of the state spine, behind
245 /// `AVAROK_GDN_TMA=1`. `try_kernel` => 0 on images that lack it, and the
246 /// launcher additionally refuses varlen and any head narrower than the
247 /// compile-time tile — the descriptors encode that tile, and a mismatched
248 /// shape loads the wrong columns without erroring.
249 pub(super) gdn_prefill_fla_chunk_delta_h_tma_k: KernelHandle,
250 pub(super) gdn_prefill_fla_chunk_fwd_o_k: KernelHandle,
251 /// WY32 chunked prefill: processes 32 tokens per WY iteration with H in
252 /// shared memory. ~30x faster than per-token for 14k+ sequences.
253 pub(super) gdn_prefill_wy32_k: KernelHandle,
254 // ── Q12 Phase 2b: same-chunk-len batched GDN prefill kernels ──
255 // Each takes `float* const* h_state_ptrs` plus stacked QKV/gate/beta/output.
256 // Used by `Qwen3SsmLayer::prefill_batched` when N≥2 streams have matching
257 // chunk_len. Null on targets that don't carry the corresponding kernel.
258 pub(super) gdn_prefill_wy32_batched_k: KernelHandle,
259 pub(super) gdn_prefill_persistent_batched_k: KernelHandle,
260 pub(super) gdn_prefill_persistent_wy4_batched_k: KernelHandle,
261 pub(super) gdn_prefill_split4_batched_k: KernelHandle,
262 pub(super) compute_gdn_gates_k: KernelHandle,
263 pub(super) ba_gates_prefill_k: KernelHandle,
264 /// Hopper twin of `ba_gates_prefill_k`
265 /// (`ssm_ba_gates_hopper::dense_gemm_ba_gates_prefill_hopper`, #928): one
266 /// CTA per token instead of `ceil(N/4)`, bit-identical output. Null on
267 /// every target but hopper, and declined below the token-count floor —
268 /// `ops::ba_gates_pick` owns both rules.
269 pub(super) ba_gates_prefill_hopper_k: KernelHandle,
270 // Kernels — prefill (multi-token sequential)
271 pub(super) conv1d_prefill_k: KernelHandle,
272 /// Token-parallel prefill conv1d (`causal_conv1d_update_prefill_tp`).
273 pub(super) conv1d_prefill_tp_k: KernelHandle,
274 // Kernels — fused chunk2 path (2-token verification)
275 pub(super) gdn_chunk2_k: KernelHandle,
276 pub(super) conv1d_chunk2_k: KernelHandle,
277 // Kernels — fused chunk3 path (3-token verification)
278 pub(super) gdn_chunk3_k: KernelHandle,
279 pub(super) w4a16_gemv_batch3_k: KernelHandle,
280 // NVFP4 batched decode GEMV (multi-seq concurrency + chain verify):
281 // the narrow batch{4,5,6,7,8} family plus batch16 (M<=16) — siblings of
282 // w8a16_gemv_batch4/16 for the FP4 QKVZ + out_proj, so FP4 decode
283 // amortizes the weight read at C=4..16 like FP8.
284 pub(super) w4a16_batchm: W4a16BatchmTiers,
285 pub(super) w4a16_gemv_batch16_k: KernelHandle,
286 // Kernels — WY-chunkwise path (2-pass verification)
287 pub(super) gdn_wy2_k: KernelHandle,
288 /// Register-resident wy2 twin (K=2 verify, the C=32 hot shape): Pass 2
289 /// is served from the Pass 1 H read retained in registers
290 /// (`__launch_bounds__(128,1)`, 128 floats/thread — the regresident
291 /// prefill pattern), cutting the kernel's HBM state traffic from 2R+2W
292 /// to 1R+2W. Byte-identical accumulation order to `gdn_wy2_k`
293 /// (bitwise-asserted by gdn_wy_verify_microtest's parity leg).
294 /// KernelHandle(0) when not linked (e.g. strix module sets). Selection +
295 /// kd/vd==128 guard + width gate (n >= wy_resident_min_width(); the
296 /// 1-block/SM kernel loses at narrow launches) live in `wy2_kernel`
297 /// (trait_decode_batched_conv_gdn);
298 /// kill switch AVAROK_NO_GDN_WY2_RESIDENT (PRESENCE — `=0` is NOT off).
299 pub(super) gdn_wy2_resident_k: KernelHandle,
300 pub(super) gdn_wy3_k: KernelHandle,
301 /// Register-resident wy3 twin (K=3 verify — the 16:2 ladder rung's 3
302 /// rows/seq shape, plus the 24:2/32:2 rungs of the 96-row envelope):
303 /// Pass 2 served from the Pass 1 H read retained in registers, cutting
304 /// HBM state traffic from 2R+3W to 1R+3W. Byte-identical accumulation
305 /// order to `gdn_wy3_k` (bitwise-asserted by gdn_wy_verify_microtest's
306 /// wy3 parity leg). KernelHandle(0) when not linked. Selection +
307 /// kd/vd==128 guard + width gate (n >= wy_resident_min_width()) live in
308 /// `wy3_kernel` (trait_decode_batched_conv_gdn);
309 /// kill switch AVAROK_NO_GDN_WY3_RESIDENT (PRESENCE — `=0` is NOT off).
310 pub(super) gdn_wy3_resident_k: KernelHandle,
311 pub(super) gdn_wy4_k: KernelHandle,
312 /// Write-on-accept K=4 twin + its post-verdict fold (2026-09-03, see
313 /// `woa.rs`). `KernelHandle(0)` when the module is absent; OPT-IN via
314 /// `AVAROK_GDN_WOA=1` (see gdn_flags.rs). Engaged only on a per-call
315 /// request.
316 pub(super) gdn_wy4_woa_k: KernelHandle,
317 pub(super) gdn_wy4_fold_k: KernelHandle,
318 pub(super) gdn_wy4_clear_k: KernelHandle,
319 /// Device addresses bound by the model on the first request (NULL until
320 /// then): the engaged u32 (1 = the twin ran for the last requesting
321 /// verify, written inside the graph) and this layer's stash slab of
322 /// `woa_seqs` sequences.
323 pub(super) woa_flag: std::sync::atomic::AtomicU64,
324 pub(super) woa_stash: std::sync::atomic::AtomicU64,
325 pub(super) woa_seqs: std::sync::atomic::AtomicUsize,
326 pub(super) woa_dims: [usize; 4],
327 /// FP16 h-state twins of the five WY verify kernels above
328 /// (`AVAROK_SSM_H_FP16` stage 2). Same launch contracts, same float
329 /// expressions and accumulation orders as their FP32 parents — the h-state
330 /// and its rollback intermediates are simply `__half` in memory, with the
331 /// state rounded once per token boundary so a rollback checkpoint holds
332 /// exactly the bits the forward chain carried.
333 ///
334 /// Stage 1 narrowed only the NON-speculative decode scan, so `--speculative`
335 /// and the flag were mutually exclusive (preflight refused). These close
336 /// that: with speculation on, the WY kernels are the only GDN h-state
337 /// readers/writers in the step, so the rungs whose best config is spec-ON
338 /// could not use FP16 at all.
339 ///
340 /// KernelHandle(0) when not linked. The selectors (`wy2_kernel`,
341 /// `wy3_kernel`, and the K=4 sites) gate on `.0 != 0` and fall back to the
342 /// FP32 parent — which is why preflight must independently refuse the flag
343 /// when a reachable K has no twin, since that fallback would read an FP16
344 /// pool through an FP32 kernel and produce fluent garbage.
345 pub(super) gdn_wy2_f16_k: KernelHandle,
346 pub(super) gdn_wy2_resident_f16_k: KernelHandle,
347 pub(super) gdn_wy3_f16_k: KernelHandle,
348 pub(super) gdn_wy3_resident_f16_k: KernelHandle,
349 pub(super) gdn_wy4_f16_k: KernelHandle,
350 /// Stage-3 f16-SIZED pool (`--ssm-h-dtype f16-pool`): the two h-state
351 /// width converters (`ssm_h_dtype.cu`). PREFILL uses them as a matched
352 /// pair around its FP32 kernels — widen the narrow slot into the
353 /// sequence's FP32 staging blob, run, narrow back — so unlike the
354 /// decode-side one-shot conversion these launch once per SSM layer per
355 /// prefill pass and are self-cancelling. A 0 handle is a hard error at
356 /// the first prefill, never an FP32 fallback: an FP32 kernel writing a
357 /// 2-byte-sized slot is an OOB write into the neighbouring slot.
358 pub(super) ssm_h_f16_to_f32_k: KernelHandle,
359 pub(super) ssm_h_f32_to_f16_k: KernelHandle,
360 /// STAGE 1 fused K=2 MTP-verify epilogue: conv1d+L2norm ×2 and
361 /// gated-RMS-norm ×2 each folded into a single launch. Dispatched only
362 /// when the `AVAROK_GDN_FUSED_VERIFY` env flag is set (default OFF); the
363 /// per-token path runs unchanged otherwise. Bit-identical (cos == 1.0).
364 pub(super) gdn_verify_fused_conv_k2_k: KernelHandle,
365 pub(super) gdn_verify_fused_norm_k2_k: KernelHandle,
366 /// Fused generic-K verify conv1d+L2norm (one launch for all K positions,
367 /// rollback snapshots written inline). Used by the K=17 DFlash verify arm;
368 /// default ON when present, kill-switch `AVAROK_GDN_FUSED_CONV17=0`.
369 /// NULL handle on targets lacking the .cu → per-token loop unchanged.
370 pub(super) gdn_verify_fused_conv_kn_k: KernelHandle,
371 /// Batched twin (gridDim.y = n_seq) — batched spec decode. 0 when absent.
372 pub(super) gdn_verify_fused_conv_kn_batched_k: KernelHandle,
373 /// Exact-verify `_snap` twins (issue #435 route (a)): the fused-norm
374 /// decode kernels with an inline per-token h-state rollback snapshot, and
375 /// the FP32-output fused verify conv. All OPTIONAL (model-shadow staged,
376 /// currently qwen3.6-27b/nvfp4 only): a 0 handle makes the exact arm fall
377 /// back to the parent kernel + `copy_d2d_async` snapshots — the same
378 /// bits, more launches.
379 pub(super) gdn_f32_norm_snap_k: KernelHandle,
380 pub(super) gdn_f32_strided_norm_snap_k: KernelHandle,
381 pub(super) gdn_verify_fused_conv_kn_f32_k: KernelHandle,
382 /// WY-Chunkwise K=17 GDN verify (DFlash γ+1). Only present in
383 /// qwen3.6-35b-a3b's PTX module set; NULL handle for other targets,
384 /// in which case decode_batched(K=17) falls through to the sequential
385 /// per-token path.
386 pub(super) gdn_wy17_k: KernelHandle,
387 /// WY-Chunkwise K∈{5..16} GDN verify (every chain-verify width between
388 /// the dedicated wy4 and the DFlash wy17; K=9..16 added 2026-08-29 for
389 /// the γ>8 window class, which previously fell to the sequential
390 /// per-token loop — the measured γ10 tax). One K-templated source
391 /// (`gated_delta_rule_wyn.cu`, gb10 common) instantiates wy5..wy16 with
392 /// the same pool-layout intermediates contract as wy17. Index = K-5;
393 /// NULL handles on targets lacking the module → sequential fallback.
394 /// Kill-switch: `AVAROK_GDN_WYN=0` (default ON).
395 pub(super) gdn_wyn_k: [KernelHandle; 12],
396 /// FP16 h-state twins of the wyN family (K=5..16), stage 2 of
397 /// `AVAROK_SSM_H_FP16` — added 2026-08-29 (#812: the FP16 pool is the
398 /// lever that lets MTP serve wide batch; DFlash was refused it for
399 /// want of these twins). Same index contract (K-5); zero handles on
400 /// targets lacking the module. Under the f16 pool a missing twin is a
401 /// HARD ERROR at dispatch (never a silent FP32 fallback over FP16
402 /// state). provenance-id: 526f6e616c6420522e205374657369616b
403 pub(super) gdn_wyn_f16_k: [KernelHandle; 12],
404 /// Pointer-table twins for the cross-sequence batched verify
405 /// (K=5..16, `state_is_table` compiled in); index = K - 5.
406 pub(super) gdn_wyn_table_k: [KernelHandle; 12],
407 /// FP16 twins of the above; same index contract.
408 pub(super) gdn_wyn_f16_table_k: [KernelHandle; 12],
409 // State allocation sizes (pre-computed from config)
410 pub(super) h_state_bytes: usize,
411 pub(super) conv_state_bytes: usize,
412 // Pre-dequanted FP8 weights for zero-overhead prefill GEMMs
413 pub(super) qkvz_fp8: Option<DevicePtr>,
414 pub(super) out_proj_fp8: Option<DevicePtr>,
415 pub(super) fp8_gemm_k: KernelHandle,
416 pub(super) fp8_gemm_t_m128_k: KernelHandle, // M128: halves B re-reads for out_proj at ISL > 128
417 // Block-scaled W8A16 prefill kernels (preferred over single-scale
418 // fp8_gemm_n128 when block-scaled FP8 weights are available — matches
419 // vLLM's per-128-block scale precision instead of single-scale).
420 pub(super) w8a16_gemm_k: KernelHandle,
421 // Pipelined (cp.async) rewrite of w8a16_gemm: bit-identical, ~4.6× faster.
422 // KernelHandle(0) when not linked into the image. Gated ON only when
423 // AVAROK_W8A16_PIPELINED=1 (default OFF — production dispatch unchanged).
424 pub(super) w8a16_gemm_pipelined_k: KernelHandle,
425 // M<=4 weight-streaming block-scaled FP8 GEMV. Replaces the M-padded
426 // w8a16_gemm_pipelined for n<=4 batched decode (qkvz + out_proj): pipelined
427 // pads M=4 to a 128-row MMA tile (32× compute over-provision, issue-bound);
428 // this streams the weight once with 4 FP32 accumulators. Bit-identical per
429 // row to w8a16_gemv. KernelHandle(0) when not linked.
430 pub(super) w8a16_gemv_batch4_k: KernelHandle,
431 // M<=16 sibling of batch4 for high-concurrency decode (n=5..16): same
432 // weight-streaming GEMV, avoids the M-padded MMA at C=8/16.
433 pub(super) w8a16_gemv_batch16_k: KernelHandle,
434 pub(super) w8a16_gemm_t_k: KernelHandle,
435 // W8A8 + FP32 epilogue (vLLM-equivalent) prefill kernels.
436 // `per_token_group_quant_fp8` produces FP8 activations + per-token-per-128
437 // FP32 scale; `fp8_gemm_t_blockscaled` consumes both with FP8 MMA and
438 // applies a_scale × b_scale in the FP32 epilogue. Gated behind
439 // `AVAROK_FP8_W8A8=1` for staged rollout.
440 pub(super) per_token_group_quant_fp8_k: ops::Fp8ActQuant,
441 pub(super) fp8_gemm_t_blockscaled_k: KernelHandle,
442 /// `fp8_act_scale_to_kmajor` — rewrites the quantizer's `[M, K/128]`
443 /// VEC128 activation scales into the `[K/128, ceil16(M)]` layout cuBLASLt
444 /// documents. 0 when the module is absent, which makes the cuBLASLt QKVZ
445 /// arm decline (see `prefill_w8a8.rs`); the in-tree kernel reads the
446 /// quantizer's own order and needs no adapter.
447 pub(super) fp8_act_scale_kmajor_k: KernelHandle,
448}