metrale_model_layers/layers/
w8a8_layer.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: The W8A8 decode weights a layer carries when its checkpoint
4//! declares FP8 weights and FP8 activations for a projection (compressed-tensors
5//! `float-quantized` groups, HF `fp8`), and the shared per-model resources they
6//! run with (`ops::w8a8_proj`).
7//!
8//! A layer holding one of these runs that projection W8A8 at every decode row
9//! count in `1..=ops::W8A8_MAX_ROWS`, ahead of its other decode arms; wider
10//! launches and prefill keep the layer's other copies. Whether a checkpoint's
11//! projection gets one is the loader's decision (the declared-precision
12//! policy); these types only say that it can run.
13//!
14//! Owner: model-layers.
15//! Invariants:
16//! - Every `W8a8Ctx` of one model shares one scratch; the W8A8 projections run
17//!   on the decode stream only, one after another, so the quantize -> GEMV
18//!   pairs never interleave.
19//! - A row's output bits do not depend on how many rows share the launch
20//!   (`ops::w8a8_decode`), so a W8A8 projection is row-invariant under every
21//!   row-tier policy.
22
23use anyhow::Result;
24use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
25
26use crate::layers::ops::{self, W8a8Kernels, W8a8Scratch, W8a8Weight};
27
28/// 2026-09-28: The kernels and the quantized-activation scratch one model's
29/// W8A8 projections share.
30#[derive(Clone, Copy, Debug)]
31pub struct W8a8Ctx {
32    pub kernels: W8a8Kernels,
33    pub scratch: W8a8Scratch,
34}
35
36impl W8a8Ctx {
37    /// 2026-09-28: Look the kernels up and allocate a scratch for a `max_k`-wide
38    /// activation of `ops::W8A8_MAX_ROWS` rows.
39    pub fn new(gpu: &dyn GpuBackend, max_k: u32) -> Result<Self> {
40        Ok(Self {
41            kernels: W8a8Kernels::load(gpu),
42            scratch: W8a8Scratch::alloc(gpu, max_k)?,
43        })
44    }
45
46    /// 2026-09-28: Whether `w` can run W8A8 at `rows` rows.
47    pub fn available(&self, w: &W8a8Weight, rows: usize) -> bool {
48        ops::w8a8_decode_available(&self.kernels, w, rows, &self.scratch)
49    }
50
51    /// 2026-09-28: `out[rows, ldc]` = W8A8 of `x[rows, ldx]` (BF16) and `w`;
52    /// `Ok(false)`, launching nothing, when not [`Self::available`].
53    #[allow(clippy::too_many_arguments)]
54    pub fn proj(
55        &self,
56        gpu: &dyn GpuBackend,
57        w: &W8a8Weight,
58        x: DevicePtr,
59        ldx: u32,
60        rows: usize,
61        out: DevicePtr,
62        ldc: u32,
63        stream: u64,
64    ) -> Result<bool> {
65        if !self.available(w, rows) {
66            return Ok(false);
67        }
68        ops::w8a8_proj(
69            gpu,
70            &self.kernels,
71            w,
72            x,
73            ldx,
74            rows,
75            out,
76            ldc,
77            &self.scratch,
78            stream,
79        )?;
80        Ok(true)
81    }
82
83    /// 2026-09-28: The down projection of a SiLU FFN: `out` = W8A8 of
84    /// `bf16(silu(gate) * up)` (gate and up `[rows, ld]` BF16) and `w`, with the
85    /// SiLU product quantized in the same launch; `Ok(false)` when not available.
86    #[allow(clippy::too_many_arguments)]
87    pub fn silu_proj(
88        &self,
89        gpu: &dyn GpuBackend,
90        w: &W8a8Weight,
91        gate: DevicePtr,
92        up: DevicePtr,
93        ld: u32,
94        rows: usize,
95        out: DevicePtr,
96        ldc: u32,
97        stream: u64,
98    ) -> Result<bool> {
99        if !self.available(w, rows) {
100            return Ok(false);
101        }
102        ops::w8a8_act_quant_silu(
103            gpu,
104            &self.kernels,
105            w.scale(),
106            gate,
107            up,
108            ld,
109            rows,
110            w.k(),
111            &self.scratch,
112            stream,
113        )?;
114        ops::w8a8_gemv(gpu, &self.kernels, w, &self.scratch, rows, out, ldc, stream)?;
115        Ok(true)
116    }
117}
118
119/// 2026-09-28: The input and output projections of an attention layer
120/// (Q|K|V stacked, O) or a GDN layer (QKV|Z stacked, out_proj).
121#[derive(Clone, Copy, Debug)]
122pub struct W8a8Mixer {
123    pub ctx: W8a8Ctx,
124    pub input: W8a8Weight,
125    pub output: W8a8Weight,
126}
127
128/// 2026-09-28: A SiLU dense FFN's gate, up and down projections.
129#[derive(Clone, Copy, Debug)]
130pub struct W8a8Ffn {
131    pub ctx: W8a8Ctx,
132    pub gate: W8a8Weight,
133    pub up: W8a8Weight,
134    pub down: W8a8Weight,
135}
metrale_model_layers/layers/w8a8_layer.rs

metrale_model_layers/layers/
w8a8_layer.rs
