spark_model/layers/qwen3_ssm/
mod.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Qwen3-Next SSM (Gated Delta Net) layer implementing TransformerLayer.
4//!
5//! Corrected pipeline matching the HuggingFace reference implementation:
6//!   1. QKVZ projection (interleaved output)
7//!   2. Deinterleave QKVZ → sequential [Q | K | V | Z]
8//!   3. BA projection (interleaved output)
9//!   4. Compute GDN gates: gate = exp(-A * softplus(alpha + dt_bias)), beta = sigmoid(b)
10//!   5. Conv1d update on [Q | K | V] concatenated (d_inner=8192)
11//!   6. Split conv output → Q', K', V'
12//!   7. GDN decode (Q', K', V', gate, beta) — kernel handles GQA internally
13//!   8. Gated RMS norm (GDN output, Z gate)
14//!   9. Output projection [value_dim → hidden_size]
15//!  10. MoE FFN
16
17use anyhow::Result;
18use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
19use spark_runtime::kv_cache::PagedKvCache;
20
21use crate::layer::{ForwardContext, GdnPrefillBuffers, LayerState, SsmLayerState};
22use crate::layers::FfnComponent;
23use crate::layers::ops;
24use crate::weight_map::{DenseWeight, QuantizedWeight, SsmWeights};
25
26// The `Qwen3SsmLayer` declaration lives in `layer_struct.rs` (≤500 LoC split).
27mod layer_struct;
28mod ple_seq;
29pub use layer_struct::Qwen3SsmLayer;
30
31// Kernel-selection helpers moved to `kernel_select.rs` (≤500 LoC split).
32
33// ── Sub-files (split for ≤500 LoC) ────────────────────────────────────────
34mod debug;
35mod decode_w8a8_proj;
36pub mod gdn_flags;
37mod init;
38mod init_fp8;
39mod init_q2;
40mod kernel_select;
41mod lora;
42mod prefill_out_w8a8;
43mod prefill_w8a8;
44mod rowwise_bf16;
45mod ssm_forward;
46pub(crate) mod ssm_h_fp16;
47mod trait_decode;
48mod trait_decode_batched;
49mod trait_decode_batched_conv_gdn;
50mod trait_decode_batched_conv_gdn_exact;
51mod trait_decode_batched_conv_gdn_multi;
52mod trait_decode_batched_conv_gdn_multi_exact;
53mod trait_decode_batched_conv_gdn_wyn;
54mod trait_decode_hc;
55mod trait_decode_multi_seq;
56mod trait_layer;
57mod trait_prefill;
58mod trait_prefill_block;
59mod trait_prefill_gdn;
60mod trait_prefill_hc;
61mod trait_prefill_helper;
62mod trait_prefill_phase1;
63mod trait_prefill_phase3;
64mod trait_prefill_proj;
65mod trait_prefill_recur;
66mod woa;
67
68pub use gdn_flags::{
69    GdnFlags, MAX_F16_TWIN_DFLASH_GAMMA, MAX_F16_TWIN_K, default_dflash_gamma,
70    gdn_fused_norm_enabled, ssm_batched_recurrent_enabled, ssm_h_dtype_bits,
71    ssm_h_f16_pool_enabled, ssm_h_fp16_enabled, verify_exact_enabled,
72};
73
74// ── TransformerLayer impl (delegates to per-file inherent _inner methods) ──
75
76#[cfg(test)]
77#[path = "prefill_alloc_tests.rs"]
78mod prefill_alloc_tests;
79#[cfg(test)]
80#[path = "rowwise_alloc_tests.rs"]
81mod rowwise_alloc_tests;
82#[cfg(test)]
83mod tests;
84
85#[path = "hc.rs"]
86mod hc;
spark_model/layers/qwen3_ssm/mod.rs

spark_model/layers/qwen3_ssm/
mod.rs
