metrale_model_layers/layers/
dense_ffn_weights.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: The dense-MLP weight sets a loader hands to `DenseFfnLayer` (NVFP4, BF16,
4//! native FP8, packed Q2) and the gated activation. Re-exported from `dense_ffn`.
5//!
6//! Owner: model-layers (dense FFN).
7//! Invariants:
8//! - Every field is the loader's device allocation; nothing here frees it.
9
10use metrale_gpu_runtime::gpu::DevicePtr;
11
12use crate::weight_map::{DenseWeight, Fp8Weight, PackedQ2Weight, QuantizedWeight};
13
14pub struct DenseFfnWeights {
15    pub gate_proj: QuantizedWeight,
16    pub up_proj: QuantizedWeight,
17    pub down_proj: QuantizedWeight,
18    /// 2026-09-25: Transposed (`[K/2, N]`) copies read by the tile prefill GEMMs (`w4a16_gemm_t*`).
19    /// `None` (never built, or freed by `finalize_q4k_load` / `finalize_nvfp4_mmq_load`) leaves
20    /// that projection's prefill to the MMQ, W4A4 or int8 arms or the base `w4a16_gemm`.
21    /// Decode reads the non-transposed weights above.
22    pub gate_proj_t: Option<QuantizedWeight>,
23    pub up_proj_t: Option<QuantizedWeight>,
24    pub down_proj_t: Option<QuantizedWeight>,
25}
26
27/// 2026-09-25: BF16 dense MLP weights. Once installed with `set_bf16_weights`, decode runs
28/// `dense_gemv_bf16` and prefill runs cuBLASLt, `dense_gemm_tc` or `dense_gemm_bf16`.
29pub struct DenseFfnWeightsBf16 {
30    pub gate_proj: DenseWeight,
31    pub up_proj: DenseWeight,
32    pub down_proj: DenseWeight,
33}
34
35/// 2026-09-26: Native block-scaled FP8 E4M3 dense MLP weights. Once installed with
36/// `set_fp8_weights`, decode runs the arm `fp8_down::fp8_down_arm` picks and prefill runs
37/// the `w8_gemm!` ladder in `prefill_fp8`.
38pub struct DenseFfnWeightsFp8 {
39    pub gate_proj: Fp8Weight,
40    pub up_proj: Fp8Weight,
41    pub down_proj: Fp8Weight,
42}
43
44/// 2026-09-25: Packed Q2_0 dense MLP weights (`block_q2_0`), kept 2-bit on the device. Once
45/// installed with `set_q2_weights`, decode runs `q2_0_gemv_vec` per projection and prefill
46/// runs the Q2_0 MMQ GEMM or a BF16 dequant (`forward_prefill_inner`).
47pub struct DenseFfnWeightsQ2 {
48    pub gate_proj: PackedQ2Weight,
49    pub up_proj: PackedQ2Weight,
50    pub down_proj: PackedQ2Weight,
51}
52
53/// 2026-09-25: Activation of the gated FFN: `SiLU` runs `moe_silu_mul`, `GeLU` runs `gelu_mul`.
54#[derive(Debug, Clone, Copy, PartialEq, Eq)]
55pub enum FfnActivation {
56    SiLU,
57    GeLU,
58}
59
60/// 2026-09-25: Int8 copy of one NVFP4 projection for the `int8_gemm_faith2` prefill arms: `w_i8`
61/// is `[N, K]` int8 and `w_scale` is `[N, K/32]` F32. Built once by `ensure_int8_weight`.
62#[derive(Debug, Clone, Copy)]
63pub(super) struct Int8Weight {
64    pub(super) w_i8: DevicePtr,
65    pub(super) w_scale: DevicePtr,
66}
67
68/// 2026-09-25: GGML `block_q4_K` copy of one NVFP4 projection for the `METRALE_FFN_MMQ` prefill
69/// arm, built once by `ensure_q4k_weight`.
70#[derive(Debug, Clone, Copy)]
71pub(super) struct Q4kWeight {
72    pub(super) w_q4k: DevicePtr,
73}
74
75/// 2026-09-25: `block_nvfp4` repack of one NVFP4 projection for the NVFP4 MMQ prefill arm, built
76/// once by `ensure_nvfp4_mmq_weight`.
77#[derive(Debug, Clone, Copy)]
78pub(super) struct Fp4MmqWeight {
79    pub(super) w: DevicePtr,
80}
metrale_model_layers/layers/dense_ffn_weights.rs

metrale_model_layers/layers/
dense_ffn_weights.rs
