metrale_model_layers/layers/
dense_ffn_init.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `DenseFfnLayer::new_with_activation`: looks up every kernel handle the layer
4//! holds, with no FP8, BF16, packed-Q2 or LoRA overlay installed.
5//!
6//! Owner: model-layers (dense FFN).
7//! Invariants:
8//! - A kernel looked up with `gpu.kernel` must resolve or construction fails; one looked up
9//!   with `try_kernel` or `try_target_kernel` leaves a zero handle when absent.
10
11use anyhow::Result;
12use metrale_gpu_runtime::gpu::{GpuBackend, KernelHandle};
13
14use super::{DenseFfnLayer, DenseFfnWeights, FfnActivation, batch16_decode, gateup_fused, m16_tc};
15use crate::layers::ops;
16use crate::layers::w4a16_gemv_tiers::W4a16BatchmTiers;
17
18impl DenseFfnLayer {
19    pub fn new_with_activation(
20        weights: DenseFfnWeights,
21        activation: FfnActivation,
22        gpu: &dyn GpuBackend,
23    ) -> Result<Self> {
24        let act_mul = match activation {
25            FfnActivation::SiLU => gpu.kernel("moe_silu_mul", "moe_silu_mul")?,
26            FfnActivation::GeLU => gpu.kernel("gelu", "gelu_mul")?,
27        };
28        // 2026-09-25: One value gates both the `silu_mul_strided` lookup below and
29        // `gateup_fused_plan`.
30        let gateup_fused = gateup_fused::ffn_gateup_fused();
31        // 2026-09-25: Optional BF16 kernels: `try_kernel` gives a zero handle when the target lacks
32        // them.
33        let dense_gemv_bf16_k = super::try_kernel(gpu, "gemv", "dense_gemv_bf16");
34        let dense_gemm_bf16_k = super::try_kernel(gpu, "gemm", "dense_gemm_bf16");
35        let dense_gemm_tc_k = super::try_kernel(gpu, "gemm_tc", "dense_gemm_tc");
36
37        let layer = Self {
38            weights,
39            activation,
40            w4a16_gemv: gpu.kernel("w4a16_gemv", "w4a16_gemv")?,
41            w4a16_gemv_sw: super::try_kernel(gpu, "w4a16_gemv", "w4a16_gemv_sw"),
42            w4a16_gemv_dual: gpu.kernel("w4a16_gemv_fused", "w4a16_gemv_dual")?,
43            w4a16_gemv_silu_input: gpu.kernel("w4a16_gemv_fused", "w4a16_gemv_silu_input")?,
44            w4a16_gemv_dual_sw: super::try_kernel(gpu, "w4a16_gemv_fused", "w4a16_gemv_dual_sw"),
45            w4a16_gemv_silu_input_sw: super::try_kernel(
46                gpu,
47                "w4a16_gemv_fused",
48                "w4a16_gemv_silu_input_sw",
49            ),
50            w4a16_gemv_dual_batch2: gpu.kernel("w4a16_gemv", "w4a16_gemv_dual_batch2")?,
51            w4a16_gemv_dual_batch3: gpu.kernel("w4a16_gemv", "w4a16_gemv_dual_batch3")?,
52            w4a16_gemv_batch2: gpu.kernel("w4a16_gemv", "w4a16_gemv_batch2")?,
53            w4a16_gemv_batch3: gpu.kernel("w4a16_gemv", "w4a16_gemv_batch3")?,
54            w4a16_batchm: W4a16BatchmTiers::resolve(gpu),
55            w4a16_gemm: gpu.kernel("w4a16", "w4a16_gemm")?,
56            w4a16_gemm_t_m128_k: super::try_kernel(gpu, "w4a16", "w4a16_gemm_t_m128"),
57            w4a16_gemm_t_m128_v2_k: super::w4a16_v2_kernel(gpu),
58            w4a16_gemm_t_m128_bf16_k: super::try_kernel(gpu, "w4a16", "w4a16_gemm_t_m128_bf16"),
59            w4a16_gemm_t_m128_bf16_v2_k: super::try_kernel(
60                gpu,
61                "w4a16",
62                "w4a16_gemm_t_m128_bf16_v2",
63            ),
64            w4a16_gemm_t_k: super::tgemm_kernel(gpu),
65            int8_faith2_k: super::try_kernel(gpu, "w4a16", "int8_gemm_faith2"),
66            int8_faith5_k: super::try_kernel(gpu, "w4a16", "int8_gemm_i32acc"),
67            requant_w_int8_k: super::try_kernel(gpu, "w4a16", "requant_w_nvfp4_int8"),
68            requant_a_int8_k: super::try_kernel(gpu, "w4a16", "requant_a_bf16_int8"),
69            int8_gate: std::sync::OnceLock::new(),
70            int8_up: std::sync::OnceLock::new(),
71            int8_down: std::sync::OnceLock::new(),
72            w4a4_gemm_k: super::try_kernel(gpu, "w4a4", "w4a4_gemm"),
73            quantize_nvfp4_k: super::try_kernel(gpu, "quantize_nvfp4", "quantize_bf16_to_nvfp4"),
74            q4k_mmq_nc_k: super::try_kernel(gpu, "q4k_mmq", "metrale_q4k_mmq128_nc"),
75            q4k_mmq_wc_k: super::try_kernel(gpu, "q4k_mmq", "metrale_q4k_mmq128_wc"),
76            q4k_quant_act_k: super::try_kernel(gpu, "q4k_mmq", "metrale_q8_1_quantize_ds4_bf16"),
77            q4k_quant_w_k: super::try_kernel(gpu, "q4k_quantize", "q4k_quantize"),
78            dequant_nvfp4_bf16_k: super::try_kernel(
79                gpu,
80                "dequant_nvfp4_bf16",
81                "dequant_nvfp4_to_bf16",
82            ),
83            q4k_gate: std::sync::OnceLock::new(),
84            q4k_up: std::sync::OnceLock::new(),
85            q4k_down: std::sync::OnceLock::new(),
86            nvfp4_mmq_nc_k: super::try_kernel(gpu, "nvfp4_mmq", "metrale_nvfp4_mmq128_nc"),
87            nvfp4_mmq_wc_k: super::try_kernel(gpu, "nvfp4_mmq", "metrale_nvfp4_mmq128_wc"),
88            nvfp4_mmq16_nc_k: super::try_kernel(gpu, "nvfp4_mmq", "metrale_nvfp4_mmq16_nc"),
89            nvfp4_mmq16_wc_k: super::try_kernel(gpu, "nvfp4_mmq", "metrale_nvfp4_mmq16_wc"),
90            nvfp4_mmq32_nc_k: super::try_kernel(gpu, "nvfp4_mmq", "metrale_nvfp4_mmq32_nc"),
91            nvfp4_mmq32_wc_k: super::try_kernel(gpu, "nvfp4_mmq", "metrale_nvfp4_mmq32_wc"),
92            nvfp4_mmq64_nc_k: super::try_kernel(gpu, "nvfp4_mmq", "metrale_nvfp4_mmq64_nc"),
93            nvfp4_mmq64_wc_k: super::try_kernel(gpu, "nvfp4_mmq", "metrale_nvfp4_mmq64_wc"),
94            nvfp4_quant_act_k: super::try_kernel(gpu, "nvfp4_mmq", "metrale_nvfp4_quantize_bf16"),
95            nvfp4_repack_k: super::try_kernel(gpu, "nvfp4_mmq", "metrale_nvfp4_repack"),
96            nvfp4_silu_scaled_k: super::try_kernel(
97                gpu,
98                "nvfp4_mmq",
99                "metrale_nvfp4_silu_mul_scaled",
100            ),
101            nvfp4_silu_quant_k: super::try_kernel(gpu, "nvfp4_mmq", "metrale_nvfp4_silu_mul_quant"),
102            nvfp4_scale_k: super::try_kernel(gpu, "nvfp4_mmq", "metrale_nvfp4_scale_bf16"),
103            fp4mmq_gate: std::sync::OnceLock::new(),
104            fp4mmq_up: std::sync::OnceLock::new(),
105            fp4mmq_down: std::sync::OnceLock::new(),
106            w4a16_gemm_t_k64_k: super::k64_kernel(gpu).unwrap_or(KernelHandle(0)),
107            act_mul,
108            bf16_weights: None,
109            w8a8: None,
110            dense_gemv_bf16_k,
111            dense_gemm_bf16_k,
112            dense_gemm_tc_k,
113            fp8_weights: None,
114            w8a16_gemv_k: super::try_kernel(gpu, "w8a16_gemv", "w8a16_gemv"),
115            w8a16_gemm_k: super::try_kernel(gpu, "w8a16_gemm", "w8a16_gemm"),
116            w8a16_gemv_batch4_k: super::try_kernel(gpu, "w8a16_gemv_batch4", "w8a16_gemv_batch4"),
117            w8a16_gemv_batch16_k: super::try_kernel(gpu, "w8a16_gemv_batch4", "w8a16_gemv_batch16"),
118            batch16_enabled: batch16_decode::ffn_batch16_enabled(),
119            w8a16_gemm_m16_k: super::try_target_kernel(gpu, "w8a16_gemm_m16", "w8a16_gemm_m16"),
120            w8a16_gemm_m16_n64_k: super::try_target_kernel(
121                gpu,
122                "w8a16_gemm_m16",
123                "w8a16_gemm_m16_n64",
124            ),
125            m16_tc: m16_tc::m16_tc_levers().ffn,
126            m16_tc_n_tile: m16_tc::m16_tc_levers().ffn_n_tile,
127            w8a16_gemm_pipelined_k: super::try_kernel(
128                gpu,
129                "w8a16_gemm_pipelined",
130                "w8a16_gemm_pipelined",
131            ),
132            w8a16_gemv_dual_k: super::try_kernel(gpu, "w8a16_gemv_fused", "w8a16_gemv_dual"),
133            w8a16_gemv_silu_input_k: super::try_kernel(
134                gpu,
135                "w8a16_gemv_fused",
136                "w8a16_gemv_silu_input",
137            ),
138            w8a16_gemm_t_m128_k: super::try_kernel(gpu, "w8a16_gemm_t_m128", "w8a16_gemm_t_m128"),
139            per_token_group_quant_fp8_k: ops::Fp8ActQuant::resolve(gpu),
140            fp8_gemm_t_blockscaled_k: super::try_kernel(
141                gpu,
142                "fp8_gemm_t_blockscaled",
143                "fp8_gemm_t_blockscaled",
144            ),
145            fp8_act_scale_kmajor_k: super::try_kernel(
146                gpu,
147                "fp8_scale_transpose",
148                "fp8_act_scale_to_kmajor",
149            ),
150            fp8_gate_up_fused: None,
151            gateup_fused,
152            silu_mul_strided_k: if gateup_fused {
153                super::try_target_kernel(gpu, "silu_mul_strided", "silu_mul_strided")
154            } else {
155                KernelHandle(0)
156            },
157            lora: None,
158            q2_weights: None,
159            q2_0_gemv_k: super::try_kernel(gpu, "q2_0_gemv_vec", "q2_0_gemv_vec"),
160            q2_0_gemv_batchm_k: super::try_kernel(gpu, "q2_0_gemv_vec", "q2_0_gemv_vec_batchm"),
161            dequant_q2_0_gn_k: super::try_kernel(
162                gpu,
163                "dequant_gguf_bf16",
164                "dequant_q2_0_gn_to_bf16",
165            ),
166            // 2026-09-25: Looked up by `set_q2_weights`, so a layer without packed-Q2 weights never
167            // probes `q2_0_mmq`.
168            q2_0_mmq_nc_k: KernelHandle(0),
169            q2_0_mmq_wc_k: KernelHandle(0),
170        };
171        Ok(layer)
172    }
173}
metrale_model_layers/layers/dense_ffn_init.rs

metrale_model_layers/layers/
dense_ffn_init.rs
