metrale_model_layers/layers/ops/
gemm_dense_int8.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the int8 prefill GEMM: the NVFP4-to-int8 weight
4//! requant, and the two-launch prefill (activation requant, then the int8 x
5//! int8 block-scaled GEMM `int8_gemm_faith2`).
6//!
7//! Owner: model-layers ops.
8//! Invariants: none beyond the types.
9
10use anyhow::Result;
11use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
12use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
13
14/// 2026-09-25: Requantize an NVFP4 weight (packed E2M1, per-16 E4M3 block
15/// scales, per-tensor `scale2`) to int8 with one F32 scale per 32 elements,
16/// for `int8_gemm_faith2`. Reads `W_packed[N, K/2]` and `W_e4m3[N, K/16]`;
17/// writes `W_i8[N, K]` (signed int8) and `W_scale[N, K/32]` (F32). The dense
18/// FFN's `ensure_int8_weight` runs it on a weight's first int8 prefill and
19/// keeps the result.
20#[allow(clippy::too_many_arguments)]
21pub fn requant_w_nvfp4_int8(
22    gpu: &dyn GpuBackend,
23    kernel: KernelHandle,
24    w_packed: DevicePtr,
25    w_e4m3: DevicePtr,
26    scale2: f32,
27    w_i8: DevicePtr,
28    w_scale: DevicePtr,
29    n: u32,
30    k: u32,
31    stream: u64,
32) -> Result<()> {
33    let blocks = n * (k / 32);
34    KernelLaunch::new(gpu, kernel)
35        .grid([div_ceil(blocks, 128), 1, 1])
36        .block([128, 1, 1])
37        .arg_ptr(w_packed)
38        .arg_ptr(w_e4m3)
39        .arg_f32(scale2)
40        .arg_ptr(w_i8)
41        .arg_ptr(w_scale)
42        .arg_u32(n)
43        .arg_u32(k)
44        .launch(stream)
45}
46
47/// 2026-09-25: int8 prefill GEMM: requantize the BF16 activations to int8
48/// with one F32 scale per 32 elements (`requant_a_kernel`), then multiply them
49/// with the int8 weight from [`requant_w_nvfp4_int8`] (`faith2_kernel`). Both
50/// launches go on `stream`.
51///
52/// `a_bf16` `[M, K]` BF16, `w_i8` `[N, K]` int8, `w_scale` `[N, K/32]` F32,
53/// `out` `[M, N]` BF16. `a_i8_scratch` and `a_scale_scratch` are caller-owned,
54/// of at least `M*K` and `M*(K/32)*4` bytes.
55#[allow(clippy::too_many_arguments)]
56pub fn int8_gemm_faith2_prefill(
57    gpu: &dyn GpuBackend,
58    faith2_kernel: KernelHandle,
59    requant_a_kernel: KernelHandle,
60    a_bf16: DevicePtr,
61    w_i8: DevicePtr,
62    w_scale: DevicePtr,
63    a_i8_scratch: DevicePtr,
64    a_scale_scratch: DevicePtr,
65    out: DevicePtr,
66    m: u32,
67    n: u32,
68    k: u32,
69    stream: u64,
70) -> Result<()> {
71    let a_blocks = m * (k / 32);
72    KernelLaunch::new(gpu, requant_a_kernel)
73        .grid([div_ceil(a_blocks, 128), 1, 1])
74        .block([128, 1, 1])
75        .arg_ptr(a_bf16)
76        .arg_ptr(a_i8_scratch)
77        .arg_ptr(a_scale_scratch)
78        .arg_u32(m)
79        .arg_u32(k)
80        .launch(stream)?;
81    KernelLaunch::new(gpu, faith2_kernel)
82        .grid([div_ceil(n, 128), div_ceil(m, 128), 1])
83        .block([256, 1, 1])
84        .arg_ptr(a_i8_scratch)
85        .arg_ptr(w_i8)
86        .arg_ptr(a_scale_scratch)
87        .arg_ptr(w_scale)
88        .arg_ptr(out)
89        .arg_u32(m)
90        .arg_u32(n)
91        .arg_u32(k)
92        .launch(stream)
93}
metrale_model_layers/layers/ops/gemm_dense_int8.rs

metrale_model_layers/layers/ops/
gemm_dense_int8.rs
