metrale_model_layers/layers/ops/w4a4_proj/
fixed.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: The NVFP4-activation projection of a fixed `--activation-quantization`
4//! ([`nvfp4_proj_mx`]): the W4A4 row quantizer and mx GEMV at any row count, in chunks of the
5//! widest prepared entry, with no fallback to another format.
6//!
7//! Owner: model-layers ops (W4A4).
8//! Invariants: launches only the prepared kernels, or returns an error before launching.
9
10use anyhow::Result;
11use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
12
13use super::{W4a4Proj, last_quant, w4a4_route};
14use crate::weight_map::QuantizedWeight;
15
16/// 2026-09-30: `[m, k] x [n, k]^T` with NVFP4 activations for any `m`: per chunk of the widest
17/// prepared entry, the row quantizer and then the mx GEMV. Every mx entry sums a row the same
18/// way and the quantizer is per row, so a row's output bits do not depend on `m` or on the
19/// other rows. Errors when the kernels were not prepared or the shape does not fit; it never
20/// falls back to another format. For a fixed `--activation-quantization` only.
21#[allow(clippy::too_many_arguments)]
22pub fn nvfp4_proj_mx(
23    gpu: &dyn GpuBackend,
24    input: DevicePtr,
25    weight: &QuantizedWeight,
26    output: DevicePtr,
27    m: u32,
28    n: u32,
29    k: u32,
30    stream: u64,
31) -> Result<()> {
32    let p = W4a4Proj::prepared(gpu).ok_or_else(|| anyhow::anyhow!("w4a4: kernels not prepared"))?;
33    let max_m = p.0.max_m;
34    anyhow::ensure!(
35        w4a4_route(1, n, k, max_m),
36        "w4a4: shape n={n} k={k} outside the mx kernels' contract"
37    );
38    // 2026-09-30: The scratch may hold another input's quantization now.
39    *last_quant().lock().unwrap_or_else(|p| p.into_inner()) = None;
40    let act = p.scratch();
41    let mut done = 0u32;
42    while done < m {
43        let rows = (m - done).min(max_m);
44        let x = input.offset(done as usize * k as usize * 2);
45        let y = output.offset(done as usize * n as usize * 2);
46        p.quantize(gpu, x, act, rows, k, stream)?;
47        p.gemv(gpu, act, weight, y, rows, n, k, stream)?;
48        done += rows;
49    }
50    Ok(())
51}
52
53/// 2026-10-01: `[m, k] x [n, k]^T` through [`nvfp4_proj_mx`] when `family` takes a fixed `nvfp4`
54/// activation format at `m` rows (`--activation-quantization`) and `weight` is present; inputs
55/// and outputs are contiguous `[m, k]` and `[m, n]` rows. `Ok(false)` launches nothing.
56#[allow(clippy::too_many_arguments)]
57pub fn fixed_nvfp4_proj(
58    gpu: &dyn GpuBackend,
59    family: metrale_config::ProjFamily,
60    weight: &QuantizedWeight,
61    input: DevicePtr,
62    output: DevicePtr,
63    m: u32,
64    n: u32,
65    k: u32,
66    stream: u64,
67) -> Result<bool> {
68    if weight.weight.is_null()
69        || crate::layers::fixed_act(family, m as usize)
70            != Some(metrale_config::ActQuantFormat::Nvfp4)
71    {
72        return Ok(false);
73    }
74    nvfp4_proj_mx(gpu, input, weight, output, m, n, k, stream)?;
75    Ok(true)
76}
metrale_model_layers/layers/ops/w4a4_proj/fixed.rs

metrale_model_layers/layers/ops/w4a4_proj/
fixed.rs
