metrale_model_layers/layers/qwen3_ssm/
w8a8_decode.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: The declared-W8A8 arm of the GDN decode projections (2026-10-01: and the fixed
4//! `nvfp4` arm, `pinned_qkvz` / `pinned_out`): `in_proj_qkv|z` (two
5//! stacked segments, the sequential `[QKV | Z]` layout) and `out_proj`, on the W8A8 decode
6//! family (`crate::layers::W8a8Mixer`). Every decode path tries it first: the single-token
7//! `ssm_forward`, the batched verify (`decode_batched`) and the multi-sequence batch
8//! (`ssm_batched_proj`).
9//!
10//! Owner: model-layers (qwen3 SSM).
11//! Invariants:
12//! - Installed only on a sequential-QKVZ layer at TP 1 (`set_w8a8_decode_weights`), so the
13//!   stacked output is the layout every reader of the QKVZ buffer expects and no TP
14//!   all-reduce is skipped.
15//! - `Ok(false)` launches nothing: no W8A8 weights, or the row count is outside the family.
16
17use anyhow::{Result, ensure};
18use metrale_gpu_runtime::gpu::DevicePtr;
19
20use super::Qwen3SsmLayer;
21use crate::layer::ForwardContext;
22use crate::layers::ops::W8a8Weight;
23use crate::layers::{FfnComponent, W8a8Ctx, W8a8Ffn, W8a8Mixer};
24use crate::weight_map::WeightQuantFormat;
25
26impl Qwen3SsmLayer {
27    /// 2026-09-28: Install the W8A8 QKV|Z and out_proj. Refuses an interleaved-QKVZ layer or
28    /// shapes that are not `[qkvz, hidden]` and `[hidden, value_dim]`.
29    pub fn set_w8a8_decode_weights(
30        &mut self,
31        w: W8a8Mixer,
32        qkvz_size: u32,
33        hidden: u32,
34        value_dim: u32,
35    ) -> Result<()> {
36        ensure!(
37            self.sequential_qkvz,
38            "W8A8 QKV|Z needs the sequential [QKV | Z] layout"
39        );
40        ensure!(
41            (w.input.n(), w.input.k(), w.output.n(), w.output.k())
42                == (qkvz_size, hidden, hidden, value_dim),
43            "W8A8 GDN shapes [{}, {}] / [{}, {}], want [{qkvz_size}, {hidden}] / [{hidden}, {value_dim}]",
44            w.input.n(),
45            w.input.k(),
46            w.output.n(),
47            w.output.k()
48        );
49        self.w8a8 = Some(w);
50        Ok(())
51    }
52
53    /// 2026-09-28: Install W8A8 gate/up/down on this layer's dense FFN
54    /// (`DenseFfnLayer::set_w8a8_decode_weights`); refuses a MoE or absent FFN.
55    pub fn set_w8a8_ffn_weights(&mut self, w: W8a8Ffn, hidden: u32, inter: u32) -> Result<()> {
56        match self.ffn {
57            FfnComponent::Dense(ref mut d) => d.set_w8a8_decode_weights(w, hidden, inter),
58            _ => anyhow::bail!("W8A8 FFN weights need a dense FFN"),
59        }
60    }
61
62    /// 2026-09-28: Run this layer's block-scaled FP8 `[QKV | Z]` and out_proj (the decode FP8
63    /// weights of a native FP8 checkpoint, `set_fp8_decode_weights`) W8A8 from now on, with
64    /// per-(token, 128) activation scales. `Ok(false)` installs nothing: they are absent or not
65    /// block-scaled.
66    pub fn adopt_fp8_block_w8a8(&mut self, ctx: W8a8Ctx, hidden: u32) -> Result<bool> {
67        let block = |w: Option<crate::weight_map::Fp8Weight>| {
68            w.filter(|f| f.scale_format == WeightQuantFormat::Fp8BlockScaled)
69        };
70        let (Some(qkvz), Some(out)) = (block(self.qkvz_fp8w), block(self.out_proj_fp8w)) else {
71            return Ok(false);
72        };
73        let (input, output) = (W8a8Weight::new(&[qkvz])?, W8a8Weight::new(&[out])?);
74        self.set_w8a8_decode_weights(W8a8Mixer { ctx, input, output }, qkvz.n, hidden, out.k)?;
75        Ok(true)
76    }
77
78    /// 2026-09-28: `out[rows, ldc]` = W8A8 QKV|Z of `normed[rows, ldx]`; `Ok(false)` when
79    /// this layer has no W8A8 weights or `rows` is outside the family. 2026-10-01: Under a fixed
80    /// `nvfp4` GDN format (`--activation-quantization`), the NVFP4 QKV|Z of a sequential layer
81    /// instead, through the row-invariant W4A4 mx path at every row count.
82    #[allow(clippy::too_many_arguments)]
83    pub(super) fn pinned_qkvz(
84        &self,
85        ctx: &ForwardContext,
86        normed: DevicePtr,
87        ldx: u32,
88        rows: usize,
89        out: DevicePtr,
90        ldc: u32,
91        stream: u64,
92    ) -> Result<bool> {
93        if self.sequential_qkvz
94            && let Some(ref nvfp4) = self.qkvz_nvfp4
95            && crate::layers::ops::w4a4_proj::fixed_nvfp4_proj(
96                ctx.gpu,
97                metrale_config::ProjFamily::Gdn,
98                nvfp4,
99                normed,
100                out,
101                rows as u32,
102                ldc,
103                ldx,
104                stream,
105            )?
106        {
107            return Ok(true);
108        }
109        match self.w8a8 {
110            Some(ref w) => w
111                .ctx
112                .proj(ctx.gpu, &w.input, normed, ldx, rows, out, ldc, stream),
113            None => Ok(false),
114        }
115    }
116
117    /// 2026-09-28: `out[rows, ldc]` = W8A8 out_proj of `normed_out[rows, ldx]`, as
118    /// [`Self::pinned_qkvz`]; under a fixed `nvfp4` GDN format, the NVFP4 out_proj on the W4A4 mx
119    /// path.
120    #[allow(clippy::too_many_arguments)]
121    pub(super) fn pinned_out(
122        &self,
123        ctx: &ForwardContext,
124        normed_out: DevicePtr,
125        ldx: u32,
126        rows: usize,
127        out: DevicePtr,
128        ldc: u32,
129        stream: u64,
130    ) -> Result<bool> {
131        let nvfp4 = &self.ssm.out_proj;
132        if self.out_proj_fp8w.is_none()
133            && self.out_proj_dense.is_none()
134            && crate::layers::ops::w4a4_proj::fixed_nvfp4_proj(
135                ctx.gpu,
136                metrale_config::ProjFamily::Gdn,
137                nvfp4,
138                normed_out,
139                out,
140                rows as u32,
141                ldc,
142                ldx,
143                stream,
144            )?
145        {
146            return Ok(true);
147        }
148        match self.w8a8 {
149            Some(ref w) => w
150                .ctx
151                .proj(ctx.gpu, &w.output, normed_out, ldx, rows, out, ldc, stream),
152            None => Ok(false),
153        }
154    }
155}
metrale_model_layers/layers/qwen3_ssm/w8a8_decode.rs

metrale_model_layers/layers/qwen3_ssm/
w8a8_decode.rs
