metrale_model_layers/layers/ops/w4a4_proj/
steps.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: The two launches of a W4A4 projection, quantize and MX GEMV, over an NVFP4
4//! activation the caller places: `proj` runs both on the shared scratch, and the circuit
5//! executor runs them as its `act_quant` and projection groups over the activation edge.
6//!
7//! Owner: model-layers ops.
8//! Invariants:
9//! - [`W4a4Proj::gemv`] launches the entry `mx_plan` picks at the process's levers, and
10//!   [`W4a4Proj::mx_kernel`] names that same entry, so a caller that checks a planned kernel
11//!   against it checks the launch.
12//! - A [`W4a4Proj`] exists only for a launch `w4a4_state_for` admits, or (from
13//!   [`W4a4Proj::prepared`]) for the quantize alone once `prepare` found the kernels.
14
15use anyhow::{Result, ensure};
16use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
17use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
18
19use super::{MxLaunch, W4a4State, mx_nt, mx_plan, mx_ps, state, w4a4_state_for};
20use crate::weight_map::QuantizedWeight;
21
22/// 2026-09-30: One NVFP4 activation of `m` rows: `[m, k / 2]` E2M1 and `[m, k / 16]` E4M3
23/// group scales, both in the GEMV's fragment order, and `[m]` FP32 per-row global scales.
24#[derive(Clone, Copy, Debug, PartialEq, Eq)]
25pub struct Nvfp4ActBuf {
26    pub aq: DevicePtr,
27    pub a_scale: DevicePtr,
28    pub a_gs: DevicePtr,
29}
30
31impl Nvfp4ActBuf {
32    /// 2026-09-30: The three parts packed in that order from `base` (the circuit's `nvfp4/g16`
33    /// edge, `Format::bytes`). Every part stays 8-byte aligned when `k % 128 == 0`.
34    pub fn packed(base: DevicePtr, m: u32, k: u32) -> Self {
35        let (m, k) = (m as usize, k as usize);
36        let a_scale = base.offset(m * k / 2);
37        Self {
38            aq: base,
39            a_scale,
40            a_gs: a_scale.offset(m * k / 16),
41        }
42    }
43}
44
45/// 2026-09-30: The W4A4 kernels of one backend (see the module header).
46#[derive(Clone, Copy)]
47pub struct W4a4Proj(pub(super) W4a4State);
48
49impl W4a4Proj {
50    /// 2026-09-30: The kernels for an `m`-row `[n, k]` projection over `weight`, when the tier
51    /// admits it there (`nvfp4_proj_small_m` then runs W4A4); `None` when it runs W4A16.
52    pub fn for_launch(
53        gpu: &dyn GpuBackend,
54        weight: &QuantizedWeight,
55        m: u32,
56        n: u32,
57        k: u32,
58    ) -> Option<Self> {
59        w4a4_state_for(gpu, weight, m, n, k).map(Self)
60    }
61
62    /// 2026-09-30: The kernels once `prepare` found them, for the quantize launch alone.
63    pub fn prepared(gpu: &dyn GpuBackend) -> Option<Self> {
64        state(gpu).map(Self)
65    }
66
67    /// 2026-09-30: The shared scratch `proj` quantizes into.
68    pub(super) fn scratch(&self) -> Nvfp4ActBuf {
69        Nvfp4ActBuf {
70            aq: self.0.aq,
71            a_scale: self.0.a_scale,
72            a_gs: self.0.a_gs,
73        }
74    }
75
76    /// 2026-09-30: `w4a4_quant_rows`.
77    pub fn quant_kernel(&self) -> KernelHandle {
78        self.0.quant
79    }
80
81    /// 2026-09-30: The MX entry an `m`-row `[n, k]` launch takes.
82    pub fn mx_kernel(&self, m: u32, n: u32, k: u32) -> KernelHandle {
83        match mx_plan(&self.0, m, n, k, mx_nt(), mx_ps()) {
84            MxLaunch::Tiles { kernel, .. } | MxLaunch::Persistent { kernel, .. } => kernel,
85        }
86    }
87
88    /// 2026-09-30: Quantize `m` rows of the BF16 `input` (`[m, k]`) into `act`.
89    #[allow(clippy::too_many_arguments)]
90    pub fn quantize(
91        &self,
92        gpu: &dyn GpuBackend,
93        input: DevicePtr,
94        act: Nvfp4ActBuf,
95        m: u32,
96        k: u32,
97        stream: u64,
98    ) -> Result<()> {
99        KernelLaunch::new(gpu, self.0.quant)
100            .grid([m, 1, 1])
101            .block([256, 1, 1])
102            .arg_ptr(input)
103            .arg_ptr(act.aq)
104            .arg_ptr(act.a_scale)
105            .arg_ptr(act.a_gs)
106            .arg_u32(k)
107            .launch(stream)
108    }
109
110    /// 2026-09-30: `output[m, n]` (BF16) = the W4A4 product of `act` and `weight`.
111    #[allow(clippy::too_many_arguments)]
112    pub fn gemv(
113        &self,
114        gpu: &dyn GpuBackend,
115        act: Nvfp4ActBuf,
116        weight: &QuantizedWeight,
117        output: DevicePtr,
118        m: u32,
119        n: u32,
120        k: u32,
121        stream: u64,
122    ) -> Result<()> {
123        let s = &self.0;
124        let (mx, grid, smem, sst) = match mx_plan(s, m, n, k, mx_nt(), mx_ps()) {
125            MxLaunch::Tiles {
126                kernel,
127                rows_per_cta,
128            } => (kernel, div_ceil(n, rows_per_cta), 0, None),
129            MxLaunch::Persistent { kernel, sst, smem } => (kernel, s.sms, smem, Some(sst)),
130        };
131        ensure!(mx.0 != 0, "w4a4: no kernel for {m} rows");
132        let launch = KernelLaunch::new(gpu, mx)
133            .grid([grid, 1, 1])
134            .block([256, 1, 1])
135            .shared_mem(smem)
136            .arg_ptr(act.aq)
137            .arg_ptr(act.a_scale)
138            .arg_ptr(act.a_gs)
139            .arg_ptr(weight.weight)
140            .arg_ptr(weight.weight_scale)
141            .arg_f32(weight.weight_scale_2)
142            .arg_ptr(output)
143            .arg_u32(m)
144            .arg_u32(n)
145            .arg_u32(k);
146        match sst {
147            Some(sst) => launch.arg_u32(sst).launch(stream),
148            None => launch.launch(stream),
149        }
150    }
151}
metrale_model_layers/layers/ops/w4a4_proj/steps.rs

metrale_model_layers/layers/ops/w4a4_proj/
steps.rs
