metrale_model_layers/layers/ops/w8a8_decode/
parts.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: The parts of a W8A8 projection the circuit executor launches one group at a
4//! time: a stacked weight's segments, the kernel each launch takes, and a scratch view over
5//! one activation buffer. The launches themselves stay [`super::w8a8_act_quant`],
6//! [`super::w8a8_act_quant_silu`] and [`super::w8a8_gemv`].
7//!
8//! Owner: model-layers ops.
9//! Invariants:
10//! - [`W8a8Kernels::gemv_entry`] is the entry [`super::w8a8_gemv`] launches for a launch of
11//!   that many rows; [`W8a8Kernels::quant`] and [`W8a8Kernels::quant_silu`] are the ones the
12//!   quantizers launch.
13
14use anyhow::{Result, ensure};
15use metrale_gpu_runtime::gpu::{DevicePtr, KernelHandle};
16
17use super::{W8A8_LAUNCH_ROWS, W8a8Kernels, W8a8Scale, W8a8Scratch, W8a8Weight, entry_index};
18
19impl W8a8Weight {
20    /// 2026-09-30: Segment `i` as a weight of its own. A segment's output rows read only its
21    /// own weight rows and scales, so its launch writes the bits the stacked launch writes
22    /// for those rows.
23    pub fn segment(&self, i: usize) -> Result<W8a8Weight> {
24        ensure!(i < self.count, "W8a8Weight: segment {i} of {}", self.count);
25        W8a8Weight::new(&[self.segs[i]])
26    }
27
28    /// 2026-09-30: Stacked segments.
29    pub fn segments(&self) -> usize {
30        self.count
31    }
32}
33
34impl W8a8Kernels {
35    /// 2026-09-30: The quantizer of `scale`'s layout.
36    pub fn quant(&self, scale: W8a8Scale) -> KernelHandle {
37        match scale {
38            W8a8Scale::PerRow => self.quant_row,
39            W8a8Scale::Block128 => self.quant_g128,
40        }
41    }
42
43    /// 2026-09-30: The SiLU·mul quantizer of `scale`'s layout.
44    pub fn quant_silu(&self, scale: W8a8Scale) -> KernelHandle {
45        match scale {
46            W8a8Scale::PerRow => self.quant_silu_row,
47            W8a8Scale::Block128 => self.quant_silu_g128,
48        }
49    }
50
51    /// 2026-09-30: The GEMV entry of one launch of `rows` rows; `None` above one launch.
52    pub fn gemv_entry(&self, scale: W8a8Scale, rows: usize) -> Option<KernelHandle> {
53        if !(1..=W8A8_LAUNCH_ROWS).contains(&rows) {
54            return None;
55        }
56        let e = entry_index(rows);
57        Some(match scale {
58            W8a8Scale::PerRow => self.rowscale[e],
59            W8a8Scale::Block128 => self.blk128[e],
60        })
61    }
62}
63
64impl W8a8Scratch {
65    /// 2026-09-30: One activation of `rows` rows packed from `base`: `[rows, k]` E4M3, then
66    /// its FP32 scales (the circuit's `fp8/token` edge, `Format::bytes`).
67    pub fn packed(base: DevicePtr, rows: usize, k: u32, scale: W8a8Scale) -> Self {
68        let q_bytes = rows * k as usize;
69        Self {
70            q: base,
71            q_bytes,
72            scale: base.offset(q_bytes),
73            scale_bytes: rows * scale.act_scales_per_row(k) * 4,
74        }
75    }
76}
metrale_model_layers/layers/ops/w8a8_decode/parts.rs

metrale_model_layers/layers/ops/w8a8_decode/
parts.rs
