spark_model/layers/ops/
dispatch_proj.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! cuBLAS / CUTLASS projection routers + their cached weight-prep helpers.
4//! Extracted from `dispatch_helpers.rs` during the ≤500-line split. Re-exported
5//! at `crate::layers::ops::*` via `ops.rs`.
6
7#![allow(unused_imports)]
8
9use super::*;
10
11/// Route a projection through native-FP8 cuBLASLt block-scaled matmul: quantize
12/// the activation to FP8 + per-[token,128-of-K] VEC128 scales (the existing
13/// `per_token_group_quant_fp8` kernel), feed the FP8 weight + its per-128×128
14/// block scales directly (zero dequant, zero extra weight memory). Both operands
15/// 128-block-scaled (cuBLASLt requires it). ~1.8× the bf16 path (152 vs 85 TF).
16///
17/// `act_fp8_scratch`/`act_scale_scratch` must hold the padded extents (the
18/// `buffers.fp8_act`/`fp8_act_scale` arena buffers, sized for max_batch_tokens).
19#[allow(clippy::too_many_arguments)]
20pub fn cublas_fp8_proj(
21    gpu: &dyn spark_runtime::gpu::GpuBackend,
22    ptg_quant_k: spark_runtime::gpu::KernelHandle,
23    act_bf16: spark_runtime::gpu::DevicePtr,
24    act_fp8_scratch: spark_runtime::gpu::DevicePtr,
25    act_scale_scratch: spark_runtime::gpu::DevicePtr,
26    fp8w: &crate::weight_map::Fp8Weight,
27    out: spark_runtime::gpu::DevicePtr,
28    m: u32,
29    n: u32,
30    k: u32,
31    stream: u64,
32) -> anyhow::Result<()> {
33    // Quantize the real M tokens → fp8 bytes + VEC128 scales [M, K/128].
34    per_token_group_quant_fp8(
35        gpu,
36        ptg_quant_k,
37        act_bf16,
38        act_fp8_scratch,
39        act_scale_scratch,
40        m,
41        k,
42        stream,
43    )?;
44    // cuBLASLt requires the scale-tensor M extent to be a multiple of 4; pad to
45    // 16 (TC-friendly) and zero the padding scale rows so the phantom output
46    // columns (ignored by the caller) are well-defined.
47    let m_pad = m.div_ceil(16) * 16;
48    if m_pad > m {
49        let kg = (k / 128) as usize;
50        let pad_off = m as usize * kg * 4;
51        let pad_bytes = (m_pad - m) as usize * kg * 4;
52        gpu.memset_async(act_scale_scratch.offset(pad_off), 0, pad_bytes, stream)?;
53    }
54    spark_runtime::cublaslt::fp8_gemm_act_weight_t_blkscaled(
55        act_fp8_scratch.0,
56        act_scale_scratch.0,
57        fp8w.weight.0,
58        fp8w.row_scale.0,
59        out.0,
60        m_pad,
61        n,
62        k,
63        stream,
64    )
65}
66
67/// Dequantize a block-scaled FP8 weight `[N,K]` → BF16 on-GPU once, cached by the
68/// FP8 weight pointer (weights are immutable after load). 128×128 blocks + FP32
69/// scales (the holo layout). Backs [`cublas_bf16_proj`].
70fn dequant_fp8_bf16_cached(
71    gpu: &dyn spark_runtime::gpu::GpuBackend,
72    derived: &super::DerivedWeights,
73    fp8w: &crate::weight_map::Fp8Weight,
74    stream: u64,
75) -> anyhow::Result<u64> {
76    use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
77    let cache_key = fp8w.weight.0;
78    if let Some(hit) = derived.get_ptr(super::Derivation::Bf16, cache_key) {
79        return Ok(hit);
80    }
81    let (n, kk) = (fp8w.n, fp8w.k);
82    let out = gpu.alloc(n as usize * kk as usize * 2)?; // BF16 [N,K]
83    // The kernel reads `scale[(n / block_n) * sk + (k / block_k)]`, so the
84    // SAME kernel serves both layouts — the block geometry is what selects
85    // between them, not a second kernel:
86    //
87    //   block-scaled   block_n = block_k = 128, sk = K/128
88    //   PER-ROW        block_n = 1, block_k = K, sk = 1
89    //                  -> offset = n * 1 + 0 = n, one multiplier per row
90    //
91    // That per-row case is what a mixed-precision compressed-tensors
92    // checkpoint ships, and dequantising it here is lossless: every FP8 E4M3
93    // value is exactly representable in BF16, so this is the fold's
94    // no-double-quant path even though the GEMM downstream is BF16.
95    let per_row = fp8w.scale_format == crate::weight_map::WeightQuantFormat::Fp8PerRow;
96    let (block_n, block_k, sk) = if per_row {
97        (1u32, kk, 1u32)
98    } else {
99        (128u32, 128u32, kk / 128)
100    };
101    let kernel = gpu.kernel(
102        "dequant_fp8_blockscaled_bf16",
103        "dequant_fp8_blockscaled_bf16",
104    )?;
105    KernelLaunch::new(gpu, kernel)
106        .grid([div_ceil(kk, 64), div_ceil(n, 4), 1])
107        .block([64, 4, 1])
108        .arg_ptr(fp8w.weight)
109        .arg_ptr(fp8w.row_scale)
110        .arg_ptr(out)
111        .arg_u32(n)
112        .arg_u32(kk)
113        .arg_u32(block_n)
114        .arg_u32(block_k)
115        .arg_u32(sk)
116        .arg_u32(1) // scale_is_fp32
117        .launch(stream)?;
118    derived.insert_ptr(super::Derivation::Bf16, cache_key, out.0);
119    Ok(out.0)
120}
121
122fn dequant_fp8_bf16_uncached(
123    gpu: &dyn spark_runtime::gpu::GpuBackend,
124    fp8w: &crate::weight_map::Fp8Weight,
125    stream: u64,
126) -> anyhow::Result<spark_runtime::gpu::DevicePtr> {
127    use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
128    let (n, kk) = (fp8w.n, fp8w.k);
129    let out = gpu.alloc(n as usize * kk as usize * 2)?;
130    let block = 128u32;
131    let sk = kk / block;
132    let kernel = gpu.kernel(
133        "dequant_fp8_blockscaled_bf16",
134        "dequant_fp8_blockscaled_bf16",
135    )?;
136    KernelLaunch::new(gpu, kernel)
137        .grid([div_ceil(kk, 64), div_ceil(n, 4), 1])
138        .block([64, 4, 1])
139        .arg_ptr(fp8w.weight)
140        .arg_ptr(fp8w.row_scale)
141        .arg_ptr(out)
142        .arg_u32(n)
143        .arg_u32(kk)
144        .arg_u32(block)
145        .arg_u32(block)
146        .arg_u32(sk)
147        .arg_u32(1)
148        .launch(stream)?;
149    Ok(out)
150}
151
152/// Route a projection `out[M,N] = act[M,K] @ weightᵀ` through cuBLASLt BF16.
153/// The FP8 weight is dequantized to BF16 once (cached); W16A16 here is strictly
154/// more accurate than the blockscaled W8A8 path it replaces.
155#[allow(clippy::too_many_arguments)]
156pub fn cublas_bf16_proj(
157    gpu: &dyn spark_runtime::gpu::GpuBackend,
158    derived: &super::DerivedWeights,
159    act: spark_runtime::gpu::DevicePtr,
160    fp8w: &crate::weight_map::Fp8Weight,
161    out: spark_runtime::gpu::DevicePtr,
162    m: u32,
163    n: u32,
164    k: u32,
165    stream: u64,
166) -> anyhow::Result<()> {
167    let w_bf16 = dequant_fp8_bf16_cached(gpu, derived, fp8w, stream)?;
168    spark_runtime::cublaslt::bf16_gemm_act_weight_t(act.0, w_bf16, out.0, m, n, k, stream)
169}
170
171/// Route a projection `out[M,N] = act[M,K] @ weightᵀ` through cuBLASLt BF16 for
172/// a weight that is already native BF16 `[N,K]` (no dequant step). Used by
173/// models whose attention/shared-expert weights ship unquantized (e.g. Laguna),
174/// which can never satisfy the `as_fp8()` gate of [`cublas_bf16_proj`].
175pub fn cublas_bf16_proj_dense(
176    act: spark_runtime::gpu::DevicePtr,
177    weight_bf16: spark_runtime::gpu::DevicePtr,
178    out: spark_runtime::gpu::DevicePtr,
179    m: u32,
180    n: u32,
181    k: u32,
182    stream: u64,
183) -> anyhow::Result<()> {
184    spark_runtime::cublaslt::bf16_gemm_act_weight_t(act.0, weight_bf16.0, out.0, m, n, k, stream)
185}
186
187/// Route a projection `out[M,N] = act[M,K] @ weightᵀ` through CUTLASS BF16.
188///
189/// ★ A REFERENCE PATH FOR BENCHMARKING, NOT A SHIPPING ONE. Opt-in behind
190/// `ATLAS_CUTLASS_GEMM=1` and OFF by default; a build without `CUTLASS_HOME`
191/// cannot reach it at all. It exists so a shape can be A/B'd against the
192/// industry reference on the same box — if CUTLASS wins a shape, the fix is
193/// a faster Atlas kernel, not a promotion. See the module docs on
194/// `spark_runtime::cutlass` for the full rationale (SSOT).
195#[allow(clippy::too_many_arguments)]
196pub fn cutlass_bf16_proj(
197    gpu: &dyn spark_runtime::gpu::GpuBackend,
198    derived: &super::DerivedWeights,
199    act: spark_runtime::gpu::DevicePtr,
200    fp8w: &crate::weight_map::Fp8Weight,
201    out: spark_runtime::gpu::DevicePtr,
202    m: u32,
203    n: u32,
204    k: u32,
205    stream: u64,
206) -> anyhow::Result<()> {
207    let w_bf16 = dequant_fp8_bf16_cached(gpu, derived, fp8w, stream)?;
208    spark_runtime::cutlass::bf16_gemm_act_weight_t(act.0, w_bf16, out.0, m, n, k, stream)
209}
210
211/// Route a projection `out[M,N] = act[M,K] @ weightᵀ` through native CUTLASS
212/// NVFP4. The activation is packed to CUTLASS NVFP4 inside the runtime wrapper.
213/// `weight_t` must be Atlas's transposed NVFP4 layout `[K/2,N]` plus
214/// `[K/16,N]` scales, as produced by `QuantizedWeight::transpose_for_gemm`.
215#[allow(clippy::too_many_arguments)]
216/// Transpose a native NVFP4 checkpoint weight from Atlas `[K/2,N]` into the
217/// CUTLASS `[N,K/2]` byte layout the GEMM consumes, caching the result by
218/// source weight ptr. Without this the ColumnMajor B operand is read
219/// transposed and the GEMM produces garbage (cos≈0 vs reference).
220fn cutlass_nvfp4_weight_transposed_cached(
221    gpu: &dyn spark_runtime::gpu::GpuBackend,
222    derived: &super::DerivedWeights,
223    weight_t: &crate::weight_map::QuantizedWeight,
224    n: u32,
225    k: u32,
226    stream: u64,
227) -> anyhow::Result<u64> {
228    let cache_key = weight_t.weight.0;
229    if let Some(hit) = derived.get_ptr(super::Derivation::CutlassNvfp4Transposed, cache_key) {
230        return Ok(hit);
231    }
232    let dst = gpu.alloc((n as usize) * (k as usize) / 2)?;
233    spark_runtime::cutlass::transpose_nvfp4_packed_kton(weight_t.weight.0, dst.0, n, k, stream)?;
234    gpu.synchronize(stream)?;
235    derived.insert_ptr(super::Derivation::CutlassNvfp4Transposed, cache_key, dst.0);
236    Ok(dst.0)
237}
238
239#[allow(clippy::too_many_arguments)]
240pub fn cutlass_nvfp4_proj(
241    // The backend and this model's derived-weight cache travel together
242    // everywhere they are used; taking the context instead of the pair keeps
243    // the call sites one line each.
244    ctx: &crate::layer::ForwardContext<'_>,
245    act: spark_runtime::gpu::DevicePtr,
246    weight_t: &crate::weight_map::QuantizedWeight,
247    out: spark_runtime::gpu::DevicePtr,
248    m: u32,
249    n: u32,
250    k: u32,
251    stream: u64,
252) -> anyhow::Result<()> {
253    let (gpu, derived) = (ctx.gpu, ctx.derived);
254    let packed = cutlass_nvfp4_weight_transposed_cached(gpu, derived, weight_t, n, k, stream)?;
255    spark_runtime::cutlass::nvfp4_gemm_bf16_act_weight_t(
256        act.0,
257        packed,
258        weight_t.weight_scale.0,
259        weight_t.weight_scale_2,
260        out.0,
261        m,
262        n,
263        k,
264        stream,
265    )
266}
267
268fn cutlass_nvfp4_weight_from_fp8_cached(
269    gpu: &dyn spark_runtime::gpu::GpuBackend,
270    derived: &super::DerivedWeights,
271    fp8w: &crate::weight_map::Fp8Weight,
272    stream: u64,
273) -> anyhow::Result<(u64, u64)> {
274    let cache_key = fp8w.weight.0;
275    if let Some(hit) = derived.get_pair(super::Derivation::CutlassNvfp4FromFp8, cache_key) {
276        return Ok(hit);
277    }
278
279    let n = fp8w.n as usize;
280    let k = fp8w.k as usize;
281    let w_bf16 = dequant_fp8_bf16_uncached(gpu, fp8w, stream)?;
282    let packed_t = gpu.alloc(n * k / 2)?;
283    let scale_t = gpu.alloc(n * k / 16)?;
284    spark_runtime::cutlass::pack_bf16_weight_to_nvfp4_t(
285        w_bf16.0, packed_t.0, scale_t.0, fp8w.n, fp8w.k, stream,
286    )?;
287    gpu.synchronize(stream)?;
288    gpu.free(w_bf16)?;
289    derived.insert_pair(
290        super::Derivation::CutlassNvfp4FromFp8,
291        cache_key,
292        (packed_t.0, scale_t.0),
293    );
294    Ok((packed_t.0, scale_t.0))
295}
296
297/// Native CUTLASS NVFP4 projection for FP8 checkpoint weights. The FP8 weight
298/// is dequantized to BF16 using the existing cache, then packed once into
299/// Atlas-transposed NVFP4 data/scales and reused for future calls.
300#[allow(clippy::too_many_arguments)]
301pub fn cutlass_nvfp4_proj_from_fp8(
302    // The backend and this model's derived-weight cache travel together
303    // everywhere they are used; taking the context instead of the pair keeps
304    // the call sites one line each.
305    ctx: &crate::layer::ForwardContext<'_>,
306    act: spark_runtime::gpu::DevicePtr,
307    fp8w: &crate::weight_map::Fp8Weight,
308    out: spark_runtime::gpu::DevicePtr,
309    m: u32,
310    n: u32,
311    k: u32,
312    stream: u64,
313) -> anyhow::Result<()> {
314    let (gpu, derived) = (ctx.gpu, ctx.derived);
315    let (packed_t, scale_t) = cutlass_nvfp4_weight_from_fp8_cached(gpu, derived, fp8w, stream)?;
316    spark_runtime::cutlass::nvfp4_gemm_bf16_act_weight_t(
317        act.0, packed_t, scale_t, 1.0, out.0, m, n, k, stream,
318    )
319}
spark_model/layers/ops/dispatch_proj.rs

spark_model/layers/ops/
dispatch_proj.rs
