spark_model/layers/ops/dispatch_proj/
cutlass.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! CUTLASS-backed projection paths: the benchmark-only BF16 reference route
4//! and the native NVFP4 routes (checkpoint-native and dequant-from-FP8).
5//! Split out of `dispatch_proj.rs` during the ≤500-line split (pure move).
6
7use super::*;
8
9/// [`dequant_fp8_bf16_into`] into a FRESH allocation, memoised by FP8 weight
10/// pointer (weights are immutable after load).
11///
12/// ⚠ OFF-LEDGER. The allocation has no `spark_runtime::buffers::sizes::BufferSizes`
13/// entry, so `--gpu-memory-utilization` cannot see it — the #917 defect named
14/// on [`dequant_fp8_bf16_into`]. The last caller is [`cutlass_bf16_proj`], a
15/// benchmark-only reference path that a shipping recipe cannot reach; the GDN
16/// row-wise arms moved to the ledgered slab in
17/// `qwen3_ssm/rowwise_bf16.rs`. Do NOT add callers — take a ledgered
18/// destination and call [`dequant_fp8_bf16_into`] instead.
19fn dequant_fp8_bf16_cached(
20    gpu: &dyn spark_runtime::gpu::GpuBackend,
21    derived: &super::DerivedWeights,
22    fp8w: &crate::weight_map::Fp8Weight,
23    stream: u64,
24) -> anyhow::Result<u64> {
25    let cache_key = fp8w.weight.0;
26    if let Some(hit) = derived.get_ptr(super::Derivation::Bf16, cache_key) {
27        return Ok(hit);
28    }
29    let out = gpu.alloc(dequant_fp8_bf16_bytes(fp8w))?; // BF16 [N,K]
30    dequant_fp8_bf16_into(gpu, fp8w, out, stream)?;
31    derived.insert_ptr(super::Derivation::Bf16, cache_key, out.0);
32    Ok(out.0)
33}
34
35/// [`dequant_fp8_bf16_into`] into a FRESH allocation the caller FREES — the
36/// NVFP4 packer's transient, which never outlives the pack.
37fn dequant_fp8_bf16_uncached(
38    gpu: &dyn spark_runtime::gpu::GpuBackend,
39    fp8w: &crate::weight_map::Fp8Weight,
40    stream: u64,
41) -> anyhow::Result<spark_runtime::gpu::DevicePtr> {
42    let out = gpu.alloc(dequant_fp8_bf16_bytes(fp8w))?;
43    dequant_fp8_bf16_into(gpu, fp8w, out, stream)?;
44    Ok(out)
45}
46
47/// Route a projection `out[M,N] = act[M,K] @ weightᵀ` through CUTLASS BF16.
48///
49/// ★ A REFERENCE PATH FOR BENCHMARKING, NOT A SHIPPING ONE. Opt-in behind
50/// `AVAROK_CUTLASS_GEMM=1` and OFF by default; a build without `CUTLASS_HOME`
51/// cannot reach it at all. It exists so a shape can be A/B'd against the
52/// industry reference on the same box — if CUTLASS wins a shape, the fix is
53/// a faster Atlas kernel, not a promotion. See the module docs on
54/// `spark_runtime::cutlass` for the full rationale (SSOT).
55#[allow(clippy::too_many_arguments)]
56pub fn cutlass_bf16_proj(
57    gpu: &dyn spark_runtime::gpu::GpuBackend,
58    derived: &super::DerivedWeights,
59    act: spark_runtime::gpu::DevicePtr,
60    fp8w: &crate::weight_map::Fp8Weight,
61    out: spark_runtime::gpu::DevicePtr,
62    m: u32,
63    n: u32,
64    k: u32,
65    stream: u64,
66) -> anyhow::Result<()> {
67    let w_bf16 = dequant_fp8_bf16_cached(gpu, derived, fp8w, stream)?;
68    spark_runtime::cutlass::bf16_gemm_act_weight_t(act.0, w_bf16, out.0, m, n, k, stream)
69}
70
71/// Route a projection `out[M,N] = act[M,K] @ weightᵀ` through native CUTLASS
72/// NVFP4. The activation is packed to CUTLASS NVFP4 inside the runtime wrapper.
73/// `weight_t` must be Atlas's transposed NVFP4 layout `[K/2,N]` plus
74/// `[K/16,N]` scales, as produced by `QuantizedWeight::transpose_for_gemm`.
75#[allow(clippy::too_many_arguments)]
76/// Transpose a native NVFP4 checkpoint weight from Atlas `[K/2,N]` into the
77/// CUTLASS `[N,K/2]` byte layout the GEMM consumes, caching the result by
78/// source weight ptr. Without this the ColumnMajor B operand is read
79/// transposed and the GEMM produces garbage (cos≈0 vs reference).
80fn cutlass_nvfp4_weight_transposed_cached(
81    gpu: &dyn spark_runtime::gpu::GpuBackend,
82    derived: &super::DerivedWeights,
83    weight_t: &crate::weight_map::QuantizedWeight,
84    n: u32,
85    k: u32,
86    stream: u64,
87) -> anyhow::Result<u64> {
88    let cache_key = weight_t.weight.0;
89    if let Some(hit) = derived.get_ptr(super::Derivation::CutlassNvfp4Transposed, cache_key) {
90        return Ok(hit);
91    }
92    let dst = gpu.alloc((n as usize) * (k as usize) / 2)?;
93    spark_runtime::cutlass::transpose_nvfp4_packed_kton(weight_t.weight.0, dst.0, n, k, stream)?;
94    gpu.synchronize(stream)?;
95    derived.insert_ptr(super::Derivation::CutlassNvfp4Transposed, cache_key, dst.0);
96    Ok(dst.0)
97}
98
99#[allow(clippy::too_many_arguments)]
100pub fn cutlass_nvfp4_proj(
101    // The backend and this model's derived-weight cache travel together
102    // everywhere they are used; taking the context instead of the pair keeps
103    // the call sites one line each.
104    ctx: &crate::layer::ForwardContext<'_>,
105    act: spark_runtime::gpu::DevicePtr,
106    weight_t: &crate::weight_map::QuantizedWeight,
107    out: spark_runtime::gpu::DevicePtr,
108    m: u32,
109    n: u32,
110    k: u32,
111    stream: u64,
112) -> anyhow::Result<()> {
113    let (gpu, derived) = (ctx.gpu, ctx.derived);
114    let packed = cutlass_nvfp4_weight_transposed_cached(gpu, derived, weight_t, n, k, stream)?;
115    spark_runtime::cutlass::nvfp4_gemm_bf16_act_weight_t(
116        act.0,
117        packed,
118        weight_t.weight_scale.0,
119        weight_t.weight_scale_2,
120        out.0,
121        m,
122        n,
123        k,
124        stream,
125    )
126}
127
128fn cutlass_nvfp4_weight_from_fp8_cached(
129    gpu: &dyn spark_runtime::gpu::GpuBackend,
130    derived: &super::DerivedWeights,
131    fp8w: &crate::weight_map::Fp8Weight,
132    stream: u64,
133) -> anyhow::Result<(u64, u64)> {
134    let cache_key = fp8w.weight.0;
135    if let Some(hit) = derived.get_pair(super::Derivation::CutlassNvfp4FromFp8, cache_key) {
136        return Ok(hit);
137    }
138
139    let n = fp8w.n as usize;
140    let k = fp8w.k as usize;
141    let w_bf16 = dequant_fp8_bf16_uncached(gpu, fp8w, stream)?;
142    let packed_t = gpu.alloc(n * k / 2)?;
143    let scale_t = gpu.alloc(n * k / 16)?;
144    spark_runtime::cutlass::pack_bf16_weight_to_nvfp4_t(
145        w_bf16.0, packed_t.0, scale_t.0, fp8w.n, fp8w.k, stream,
146    )?;
147    gpu.synchronize(stream)?;
148    gpu.free(w_bf16)?;
149    derived.insert_pair(
150        super::Derivation::CutlassNvfp4FromFp8,
151        cache_key,
152        (packed_t.0, scale_t.0),
153    );
154    Ok((packed_t.0, scale_t.0))
155}
156
157/// Native CUTLASS NVFP4 projection for FP8 checkpoint weights. The FP8 weight
158/// is dequantized to BF16 using the existing cache, then packed once into
159/// Atlas-transposed NVFP4 data/scales and reused for future calls.
160#[allow(clippy::too_many_arguments)]
161pub fn cutlass_nvfp4_proj_from_fp8(
162    // The backend and this model's derived-weight cache travel together
163    // everywhere they are used; taking the context instead of the pair keeps
164    // the call sites one line each.
165    ctx: &crate::layer::ForwardContext<'_>,
166    act: spark_runtime::gpu::DevicePtr,
167    fp8w: &crate::weight_map::Fp8Weight,
168    out: spark_runtime::gpu::DevicePtr,
169    m: u32,
170    n: u32,
171    k: u32,
172    stream: u64,
173) -> anyhow::Result<()> {
174    let (gpu, derived) = (ctx.gpu, ctx.derived);
175    let (packed_t, scale_t) = cutlass_nvfp4_weight_from_fp8_cached(gpu, derived, fp8w, stream)?;
176    spark_runtime::cutlass::nvfp4_gemm_bf16_act_weight_t(
177        act.0, packed_t, scale_t, 1.0, out.0, m, n, k, stream,
178    )
179}
spark_model/layers/ops/dispatch_proj/cutlass.rs

spark_model/layers/ops/dispatch_proj/
cutlass.rs
