metrale_model_layers/layers/qwen3_ssm/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `Qwen3SsmLayer`, the Gated DeltaNet (linear-attention) layer,
4//! and the submodules its constructor, forward paths and tests live in.
5//!
6//! Owner: model-layers (qwen3_ssm).
7//! Invariants: none beyond the types.
8
9use anyhow::Result;
10use metrale_cache::kv_cache::PagedKvCache;
11use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
12
13use crate::layer::{ForwardContext, GdnPrefillBuffers, LayerState, SsmLayerState};
14use crate::layers::FfnComponent;
15use crate::layers::ops;
16use crate::weight_map::{DenseWeight, QuantizedWeight, SsmWeights};
17
18mod layer_struct;
19mod ple_seq;
20pub use layer_struct::Qwen3SsmLayer;
21
22mod carry;
23pub use carry::carry_flush_kernel;
24mod circuit;
25mod debug;
26mod decode_w8a8_proj;
27pub mod gdn_flags;
28mod init;
29mod init_fp8;
30mod init_q2;
31mod kernel_select;
32mod lora;
33mod prefill_out_w8a8;
34mod prefill_w8a8;
35mod row_tier_proj;
36mod rowwise_bf16;
37mod ssm_forward;
38pub(crate) mod ssm_h_fp16;
39mod trait_decode;
40mod trait_decode_batched;
41mod trait_decode_batched_conv_gdn;
42mod trait_decode_batched_conv_gdn_exact;
43mod trait_decode_batched_conv_gdn_exact_chain;
44mod trait_decode_batched_conv_gdn_exact_f16;
45mod trait_decode_batched_conv_gdn_multi;
46mod trait_decode_batched_conv_gdn_multi_exact;
47mod trait_decode_batched_conv_gdn_wyn;
48mod trait_decode_hc;
49mod trait_decode_multi_seq;
50mod trait_layer;
51mod trait_prefill;
52mod trait_prefill_block;
53mod trait_prefill_gdn;
54mod trait_prefill_hc;
55mod trait_prefill_helper;
56mod trait_prefill_phase1;
57mod trait_prefill_phase3;
58mod trait_prefill_proj;
59mod trait_prefill_recur;
60mod w8a8_decode;
61mod woa;
62
63pub use gdn_flags::{
64    GdnFlags, MAX_F16_TWIN_DFLASH_GAMMA, MAX_F16_TWIN_K, UNRESOLVED_DFLASH_GAMMA,
65    default_dflash_gamma, gdn_fused_norm_enabled, resolve_dflash_gamma,
66    ssm_batched_recurrent_enabled, ssm_h_dtype_bits, ssm_h_f16_pool_enabled, ssm_h_fp16_enabled,
67    verify_exact_enabled,
68};
69
70#[cfg(test)]
71#[path = "prefill_alloc_tests.rs"]
72mod prefill_alloc_tests;
73#[cfg(test)]
74#[path = "rowwise_alloc_tests.rs"]
75mod rowwise_alloc_tests;
76#[cfg(test)]
77mod tests;
78#[cfg(test)]
79#[path = "tests_inpass_capture.rs"]
80mod tests_inpass_capture;
81
82#[path = "hc.rs"]
83mod hc;