spark_model/layers/ops/
gemm_dense_int8.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! int8 W4A8 prefill GEMM wrappers, extracted piecewise from `gemm_dense.rs`
4//! (500-LoC cap): one-time NVFP4→int8 weight requant + the faith2 two-launch
5//! prefill (activation requant → int8×int8 block-scaled GEMM).
6
7use anyhow::Result;
8use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
9use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
10
11/// Requant an NVFP4 weight (packed E2M1 + per-16 E4M3 block scales + per-tensor
12/// `scale2`) into an int8 weight + per-32 F32 block scale, for the int8 W4A8
13/// prefill GEMM (`int8_gemm_faith2`). One-time conversion per weight at load (or
14/// lazily on first int8 prefill).
15///
16/// Reads `W_packed[N, K/2]`, `W_e4m3[N, K/16]`, `scale2` → `W_i8[N, K]` (signed
17/// int8) + `W_scale[N, K/32]` (F32). The per-16 NVFP4 scales are re-blocked to
18/// per-32 int8 scales by the kernel.
19///
20/// Grid: (ceil(N*(K/32) / 128), 1, 1)  Block: (128, 1, 1)
21#[allow(clippy::too_many_arguments)]
22pub fn requant_w_nvfp4_int8(
23    gpu: &dyn GpuBackend,
24    kernel: KernelHandle,
25    w_packed: DevicePtr,
26    w_e4m3: DevicePtr,
27    scale2: f32,
28    w_i8: DevicePtr,
29    w_scale: DevicePtr,
30    n: u32,
31    k: u32,
32    stream: u64,
33) -> Result<()> {
34    let blocks = n * (k / 32);
35    KernelLaunch::new(gpu, kernel)
36        .grid([div_ceil(blocks, 128), 1, 1])
37        .block([128, 1, 1])
38        .arg_ptr(w_packed)
39        .arg_ptr(w_e4m3)
40        .arg_f32(scale2)
41        .arg_ptr(w_i8)
42        .arg_ptr(w_scale)
43        .arg_u32(n)
44        .arg_u32(k)
45        .launch(stream)
46}
47
48/// int8 W4A8 prefill GEMM: requant BF16 activations to int8 (per-32 F32 scale)
49/// then `C = (A_i8 * W_i8)` folded with per-32 A/W block scales via
50/// `int8_gemm_faith2`. The weight is already int8 (see `requant_w_nvfp4_int8`).
51///
52/// A_bf16: [M, K] BF16 activations. W_i8: [N, K] int8 weights. W_scale: [N, K/32]
53/// F32. `a_i8_scratch` / `a_scale_scratch` are caller-owned scratch buffers of at
54/// least `M*K` bytes and `M*(K/32)*4` bytes respectively. Out: [M, N] BF16.
55///
56/// Two launches on `stream` (stream-ordered): requant_a → faith2.
57///   requant_a grid: (ceil(M*(K/32) / 128), 1, 1)  block: (128, 1, 1)
58///   faith2    grid: (ceil(N/128), ceil(M/128), 1)  block: (256, 1, 1)
59#[allow(clippy::too_many_arguments)]
60pub fn int8_gemm_faith2_prefill(
61    gpu: &dyn GpuBackend,
62    faith2_kernel: KernelHandle,
63    requant_a_kernel: KernelHandle,
64    a_bf16: DevicePtr,
65    w_i8: DevicePtr,
66    w_scale: DevicePtr,
67    a_i8_scratch: DevicePtr,
68    a_scale_scratch: DevicePtr,
69    out: DevicePtr,
70    m: u32,
71    n: u32,
72    k: u32,
73    stream: u64,
74) -> Result<()> {
75    // (a) BF16 acts → int8 + per-32 F32 scale.
76    let a_blocks = m * (k / 32);
77    KernelLaunch::new(gpu, requant_a_kernel)
78        .grid([div_ceil(a_blocks, 128), 1, 1])
79        .block([128, 1, 1])
80        .arg_ptr(a_bf16)
81        .arg_ptr(a_i8_scratch)
82        .arg_ptr(a_scale_scratch)
83        .arg_u32(m)
84        .arg_u32(k)
85        .launch(stream)?;
86    // (b) int8 × int8 GEMM with per-32 block scales → BF16 out.
87    KernelLaunch::new(gpu, faith2_kernel)
88        .grid([div_ceil(n, 128), div_ceil(m, 128), 1])
89        .block([256, 1, 1])
90        .arg_ptr(a_i8_scratch)
91        .arg_ptr(w_i8)
92        .arg_ptr(a_scale_scratch)
93        .arg_ptr(w_scale)
94        .arg_ptr(out)
95        .arg_u32(m)
96        .arg_u32(n)
97        .arg_u32(k)
98        .launch(stream)
99}
spark_model/layers/ops/gemm_dense_int8.rs

spark_model/layers/ops/
gemm_dense_int8.rs
