metrale_model_layers/layers/ops/dispatch_proj/
cutlass.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: CUTLASS projection routes: BF16 from an FP8 weight, and NVFP4
4//! from an NVFP4 or an FP8 weight.
5//!
6//! Owner: model-layers ops.
7//! Invariants:
8//! - Each derived weight is memoized in the model's `DerivedWeights` under the
9//!   source weight's device pointer, and only after it was built without
10//!   error.
11
12use super::*;
13
14/// 2026-09-25: [`dequant_fp8_bf16_into`] into a new allocation, memoized by
15/// the FP8 weight's pointer.
16///
17/// The allocation is a plain `gpu.alloc` with no `BufferSizes` entry. A new
18/// caller should pass its own budgeted destination to
19/// [`dequant_fp8_bf16_into`] instead.
20fn dequant_fp8_bf16_cached(
21    gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
22    derived: &super::DerivedWeights,
23    fp8w: &crate::weight_map::Fp8Weight,
24    stream: u64,
25) -> anyhow::Result<u64> {
26    let cache_key = fp8w.weight.0;
27    if let Some(hit) = derived.get_ptr(super::Derivation::Bf16, cache_key) {
28        return Ok(hit);
29    }
30    let out = gpu.alloc(dequant_fp8_bf16_bytes(fp8w))?;
31    dequant_fp8_bf16_into(gpu, fp8w, out, stream)?;
32    derived.insert_ptr(super::Derivation::Bf16, cache_key, out.0);
33    Ok(out.0)
34}
35
36/// 2026-09-25: [`dequant_fp8_bf16_into`] into a new allocation that the caller
37/// must free.
38fn dequant_fp8_bf16_uncached(
39    gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
40    fp8w: &crate::weight_map::Fp8Weight,
41    stream: u64,
42) -> anyhow::Result<metrale_gpu_runtime::gpu::DevicePtr> {
43    let out = gpu.alloc(dequant_fp8_bf16_bytes(fp8w))?;
44    dequant_fp8_bf16_into(gpu, fp8w, out, stream)?;
45    Ok(out)
46}
47
48/// 2026-09-25: Route a projection `out[M,N] = act[M,K] @ weightᵀ` through
49/// CUTLASS BF16, against the FP8 weight's memoized BF16 copy.
50#[allow(clippy::too_many_arguments)]
51pub fn cutlass_bf16_proj(
52    gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
53    derived: &super::DerivedWeights,
54    act: metrale_gpu_runtime::gpu::DevicePtr,
55    fp8w: &crate::weight_map::Fp8Weight,
56    out: metrale_gpu_runtime::gpu::DevicePtr,
57    m: u32,
58    n: u32,
59    k: u32,
60    stream: u64,
61) -> anyhow::Result<()> {
62    let w_bf16 = dequant_fp8_bf16_cached(gpu, derived, fp8w, stream)?;
63    metrale_gpu_runtime::cutlass::bf16_gemm_act_weight_t(act.0, w_bf16, out.0, m, n, k, stream)
64}
65
66#[allow(clippy::too_many_arguments)]
67/// 2026-09-25: Transpose an NVFP4 weight's packed bytes from `[K/2,N]` into the
68/// `[N,K/2]` layout the CUTLASS GEMM reads, memoized by the source weight's
69/// pointer.
70fn cutlass_nvfp4_weight_transposed_cached(
71    gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
72    derived: &super::DerivedWeights,
73    weight_t: &crate::weight_map::QuantizedWeight,
74    n: u32,
75    k: u32,
76    stream: u64,
77) -> anyhow::Result<u64> {
78    let cache_key = weight_t.weight.0;
79    if let Some(hit) = derived.get_ptr(super::Derivation::CutlassNvfp4Transposed, cache_key) {
80        return Ok(hit);
81    }
82    let dst = gpu.alloc((n as usize) * (k as usize) / 2)?;
83    metrale_gpu_runtime::cutlass::transpose_nvfp4_packed_kton(
84        weight_t.weight.0,
85        dst.0,
86        n,
87        k,
88        stream,
89    )?;
90    gpu.synchronize(stream)?;
91    derived.insert_ptr(super::Derivation::CutlassNvfp4Transposed, cache_key, dst.0);
92    Ok(dst.0)
93}
94
95#[allow(clippy::too_many_arguments)]
96/// 2026-09-25: Route a projection `out[M,N] = act[M,K] @ weightᵀ` through
97/// CUTLASS NVFP4. `weight_t` is an NVFP4 weight with `[K/2,N]` packed bytes;
98/// the bytes are transposed once to `[N,K/2]` and memoized.
99pub fn cutlass_nvfp4_proj(
100    ctx: &crate::layer::ForwardContext<'_>,
101    act: metrale_gpu_runtime::gpu::DevicePtr,
102    weight_t: &crate::weight_map::QuantizedWeight,
103    out: metrale_gpu_runtime::gpu::DevicePtr,
104    m: u32,
105    n: u32,
106    k: u32,
107    stream: u64,
108) -> anyhow::Result<()> {
109    let (gpu, derived) = (ctx.gpu, ctx.derived);
110    let packed = cutlass_nvfp4_weight_transposed_cached(gpu, derived, weight_t, n, k, stream)?;
111    metrale_gpu_runtime::cutlass::nvfp4_gemm_bf16_act_weight_t(
112        act.0,
113        packed,
114        weight_t.weight_scale.0,
115        weight_t.weight_scale_2,
116        out.0,
117        m,
118        n,
119        k,
120        stream,
121    )
122}
123
124fn cutlass_nvfp4_weight_from_fp8_cached(
125    gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
126    derived: &super::DerivedWeights,
127    fp8w: &crate::weight_map::Fp8Weight,
128    stream: u64,
129) -> anyhow::Result<(u64, u64)> {
130    let cache_key = fp8w.weight.0;
131    if let Some(hit) = derived.get_pair(super::Derivation::CutlassNvfp4FromFp8, cache_key) {
132        return Ok(hit);
133    }
134
135    let n = fp8w.n as usize;
136    let k = fp8w.k as usize;
137    let w_bf16 = dequant_fp8_bf16_uncached(gpu, fp8w, stream)?;
138    let packed_t = gpu.alloc(n * k / 2)?;
139    let scale_t = gpu.alloc(n * k / 16)?;
140    metrale_gpu_runtime::cutlass::pack_bf16_weight_to_nvfp4_t(
141        w_bf16.0, packed_t.0, scale_t.0, fp8w.n, fp8w.k, stream,
142    )?;
143    gpu.synchronize(stream)?;
144    gpu.free(w_bf16)?;
145    derived.insert_pair(
146        super::Derivation::CutlassNvfp4FromFp8,
147        cache_key,
148        (packed_t.0, scale_t.0),
149    );
150    Ok((packed_t.0, scale_t.0))
151}
152
153/// 2026-09-25: CUTLASS NVFP4 projection for an FP8 weight. On first use the
154/// weight is dequantized into a transient BF16 buffer, packed into NVFP4 data
155/// and scales, and the pair is memoized for later calls.
156#[allow(clippy::too_many_arguments)]
157pub fn cutlass_nvfp4_proj_from_fp8(
158    ctx: &crate::layer::ForwardContext<'_>,
159    act: metrale_gpu_runtime::gpu::DevicePtr,
160    fp8w: &crate::weight_map::Fp8Weight,
161    out: metrale_gpu_runtime::gpu::DevicePtr,
162    m: u32,
163    n: u32,
164    k: u32,
165    stream: u64,
166) -> anyhow::Result<()> {
167    let (gpu, derived) = (ctx.gpu, ctx.derived);
168    let (packed_t, scale_t) = cutlass_nvfp4_weight_from_fp8_cached(gpu, derived, fp8w, stream)?;
169    metrale_gpu_runtime::cutlass::nvfp4_gemm_bf16_act_weight_t(
170        act.0, packed_t, scale_t, 1.0, out.0, m, n, k, stream,
171    )
172}
metrale_model_layers/layers/ops/dispatch_proj/cutlass.rs

metrale_model_layers/layers/ops/dispatch_proj/
cutlass.rs
