metrale_model_layers/layers/ops/
dispatch_proj_rowwise.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Row-wise FP8 projection through cuBLASLt, with the weight's
4//! per-row FP8 pair passed through or re-quantized from block-scaled FP8.
5//!
6//! Owner: model-layers ops.
7//! Invariants:
8//! - A weight tagged `Fp8PerRow` is used through its own pointers; nothing is
9//!   allocated or memoized for it.
10//! - The requant panics on a format other than `Fp8PerRow` or
11//!   `Fp8BlockScaled`, before it reads the scales.
12
13/// 2026-09-25: `(weight, scale)` as they are when `fp8w` is tagged
14/// `Fp8PerRow`, else `None`. Pure, so it is tested without a GPU.
15pub(super) fn rowwise_pair_passthrough(fp8w: &crate::weight_map::Fp8Weight) -> Option<(u64, u64)> {
16    use crate::weight_map::WeightQuantFormat;
17    (fp8w.scale_format == WeightQuantFormat::Fp8PerRow).then_some((fp8w.weight.0, fp8w.row_scale.0))
18}
19
20/// 2026-09-25: Re-quantize a block-scaled FP8 weight `[N,K]` to row-wise FP8
21/// (E4M3 plus one FP32 scale per row) on the GPU, through a transient BF16
22/// copy, and memoize it by the FP8 weight's pointer. Returns
23/// `(fp8_weight_ptr, per_row_scale_ptr)`. A weight that is already row-wise
24/// returns its own pointers.
25fn requant_weight_rowwise_fp8_cached(
26    gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
27    derived: &super::DerivedWeights,
28    fp8w: &crate::weight_map::Fp8Weight,
29    stream: u64,
30) -> anyhow::Result<(u64, u64)> {
31    use crate::weight_map::WeightQuantFormat;
32    use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
33
34    // 2026-09-25: A mixed-precision compressed-tensors checkpoint (for example
35    // unsloth/Qwen3.8-27B-NVFP4) ships its attention and GDN projections as FP8
36    // with a per-channel scale, which is already the pair this function
37    // produces. Re-quantizing it would round it through BF16 for nothing.
38    if let Some(pair) = rowwise_pair_passthrough(fp8w) {
39        return Ok(pair);
40    }
41    // 2026-09-25: The conversion reads `row_scale` as a `[N/128, K/128]` FP32
42    // grid, so any other scale format is refused here.
43    fp8w.scale_format
44        .expect(WeightQuantFormat::Fp8BlockScaled, "rowwise-fp8 requant");
45    let cache_key = fp8w.weight.0;
46    if let Some(hit) = derived.get_pair(super::Derivation::RowwiseFp8, cache_key) {
47        return Ok(hit);
48    }
49    let (n, k) = (fp8w.n, fp8w.k);
50    // 2026-09-25: Block-scaled FP8 to a transient BF16 `[N,K]`.
51    let bf16 = gpu.alloc(n as usize * k as usize * 2)?;
52    let block = 128u32;
53    let sk = k / block;
54    let dq = gpu.kernel(
55        "dequant_fp8_blockscaled_bf16",
56        "dequant_fp8_blockscaled_bf16",
57    )?;
58    KernelLaunch::new(gpu, dq)
59        .grid([div_ceil(k, 64), div_ceil(n, 4), 1])
60        .block([64, 4, 1])
61        .arg_ptr(fp8w.weight)
62        .arg_ptr(fp8w.row_scale)
63        .arg_ptr(bf16)
64        .arg_u32(n)
65        .arg_u32(k)
66        .arg_u32(block)
67        .arg_u32(block)
68        .arg_u32(sk)
69        .arg_u32(1)
70        .launch(stream)?;
71    // 2026-09-25: BF16 to row-wise FP8 `[N,K]` plus a per-row scale `[N]`.
72    let w_fp8 = gpu.alloc(n as usize * k as usize)?;
73    let w_scale = gpu.alloc(n as usize * 4)?;
74    let qk = gpu.kernel("quant_rowwise_fp8", "quant_rowwise_fp8")?;
75    KernelLaunch::new(gpu, qk)
76        .grid([n, 1, 1])
77        .block([256, 1, 1])
78        .arg_ptr(bf16)
79        .arg_ptr(w_fp8)
80        .arg_ptr(w_scale)
81        .arg_u32(n)
82        .arg_u32(k)
83        .launch(stream)?;
84    gpu.synchronize(stream)?; // 2026-09-25: finish the re-quant before freeing the BF16 copy
85    gpu.free(bf16)?;
86    derived.insert_pair(
87        super::Derivation::RowwiseFp8,
88        cache_key,
89        (w_fp8.0, w_scale.0),
90    );
91    Ok((w_fp8.0, w_scale.0))
92}
93
94/// 2026-09-25: Route a projection through row-wise FP8 cuBLASLt. The weight
95/// pair comes from `requant_weight_rowwise_fp8_cached`; the activation is
96/// quantized per token on every call.
97///
98/// cuBLASLt is handed `ceil16(m)` rows, so `act_fp8_scratch` must hold
99/// `ceil16(m) * k` bytes, `act_scale_scratch` `ceil16(m)` f32 and `out`
100/// `ceil16(m) * n` BF16 elements. The pad rows of both scratch buffers are
101/// zeroed here, and the pad output rows are written.
102#[allow(clippy::too_many_arguments)]
103pub fn cublas_fp8_rowwise_proj(
104    gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
105    derived: &super::DerivedWeights,
106    act_bf16: metrale_gpu_runtime::gpu::DevicePtr,
107    act_fp8_scratch: metrale_gpu_runtime::gpu::DevicePtr,
108    act_scale_scratch: metrale_gpu_runtime::gpu::DevicePtr,
109    fp8w: &crate::weight_map::Fp8Weight,
110    out: metrale_gpu_runtime::gpu::DevicePtr,
111    m: u32,
112    n: u32,
113    k: u32,
114    stream: u64,
115) -> anyhow::Result<()> {
116    use metrale_gpu_runtime::kernel_args::KernelLaunch;
117    let (w_fp8, w_scale) = requant_weight_rowwise_fp8_cached(gpu, derived, fp8w, stream)?;
118    // 2026-09-25: Per-token quant of the activation to FP8 `[M,K]` plus scale `[M]`.
119    let qk = gpu.kernel("quant_rowwise_fp8", "quant_rowwise_fp8")?;
120    KernelLaunch::new(gpu, qk)
121        .grid([m, 1, 1])
122        .block([256, 1, 1])
123        .arg_ptr(act_bf16)
124        .arg_ptr(act_fp8_scratch)
125        .arg_ptr(act_scale_scratch)
126        .arg_u32(m)
127        .arg_u32(k)
128        .launch(stream)?;
129    // 2026-09-25: Pad M to 16, as `cublas_fp8_m_pad` does, and zero both the
130    // pad scales and the pad activation rows: a zero scale drops their
131    // contribution, and zeroed bytes cannot carry a NaN into the accumulator.
132    let m_pad = m.div_ceil(16) * 16;
133    if m_pad > m {
134        let pad_rows = (m_pad - m) as usize;
135        gpu.memset_async(
136            act_scale_scratch.offset(m as usize * 4),
137            0,
138            pad_rows * 4,
139            stream,
140        )?;
141        gpu.memset_async(
142            act_fp8_scratch.offset(m as usize * k as usize),
143            0,
144            pad_rows * k as usize,
145            stream,
146        )?;
147    }
148    metrale_gpu_runtime::cublaslt::fp8_gemm_act_weight_t_rowwise(
149        act_fp8_scratch.0,
150        act_scale_scratch.0,
151        w_fp8,
152        w_scale,
153        out.0,
154        m_pad,
155        n,
156        k,
157        stream,
158    )
159}
160
161#[cfg(test)]
162mod rowwise_passthrough_tests {
163    use super::{requant_weight_rowwise_fp8_cached, rowwise_pair_passthrough};
164    use crate::layers::ops::DerivedWeights;
165    use crate::weight_map::{Fp8Weight, WeightQuantFormat};
166    use metrale_gpu_runtime::gpu::DevicePtr;
167    use metrale_gpu_runtime::gpu::mock::MockGpuBackend;
168
169    fn weight(scale_format: WeightQuantFormat) -> Fp8Weight {
170        Fp8Weight {
171            weight: DevicePtr(0xBEEF),
172            row_scale: DevicePtr(0x5CA1E),
173            n: 4096,
174            k: 5120,
175            scale_format,
176        }
177    }
178
179    /// 2026-09-25: A weight that already has per-row scales is passed through
180    /// with its own pointers.
181    #[test]
182    fn an_already_rowwise_weight_passes_through_verbatim() {
183        let w = weight(WeightQuantFormat::Fp8PerRow);
184        assert_eq!(
185            rowwise_pair_passthrough(&w),
186            Some((w.weight.0, w.row_scale.0)),
187            "the checkpoint's own pointers, not a converted copy"
188        );
189    }
190
191    /// 2026-09-25: Every format other than `Fp8PerRow` returns `None`.
192    #[test]
193    fn other_formats_still_requantize() {
194        for f in [
195            WeightQuantFormat::Fp8BlockScaled,
196            WeightQuantFormat::Fp8SingleScale,
197            WeightQuantFormat::Bf16,
198            WeightQuantFormat::Nvfp4,
199        ] {
200            assert_eq!(
201                rowwise_pair_passthrough(&weight(f)),
202                None,
203                "{f:?} is not a row-wise pair and must not be passed through"
204            );
205        }
206    }
207
208    #[test]
209    fn cached_requant_returns_rowwise_checkpoint_pointers_without_gpu_work() {
210        let gpu = MockGpuBackend::new();
211        let w = weight(WeightQuantFormat::Fp8PerRow);
212
213        assert_eq!(
214            requant_weight_rowwise_fp8_cached(&gpu, &DerivedWeights::new(), &w, 0).unwrap(),
215            (w.weight.0, w.row_scale.0)
216        );
217        assert_eq!(gpu.alloc_count(), 0, "passthrough must not allocate a copy");
218    }
219}
metrale_model_layers/layers/ops/dispatch_proj_rowwise.rs

metrale_model_layers/layers/ops/
dispatch_proj_rowwise.rs
