cubecl_std/quant/
round.rs

1use cubecl::prelude::*;
2use cubecl_common::quant::scheme::{F32Grid, QuantParam};
3use cubecl_core as cubecl;
4
5/// The smallest value representable in `param` that is not below `scale`, in a kernel.
6///
7/// Device-side counterpart of [`QuantParam::round_up`], and the two have to agree: a tensor
8/// quantized on one backend has to reconstruct the same on another.
9///
10/// Returned as `F` rather than the storage type because the result is exactly representable in
11/// `param`, so the caller's cast to it is lossless.
12///
13/// `F` only carries the value in and out. The rule runs in f32, so a narrow `F` cannot turn the
14/// saturation bound into an infinity or the subnormal spacing into a flushed zero.
15///
16/// `scale` must not be negative, as with the host rule.
17#[cube]
18pub fn round_up_to_param<F: Float>(scale: F, #[comptime] param: QuantParam) -> F {
19    #[comptime]
20    match param {
21        QuantParam::F32 => scale,
22        QuantParam::F16 | QuantParam::BF16 | QuantParam::UE4M3 => {
23            F::cast_from(step_up(f32::cast_from(scale), param))
24        }
25        // Returning `scale` would diverge from the host rule, which has no answer here either.
26        QuantParam::UE8M0 => comptime!(unimplemented!("UE8M0 scales are not yet supported")),
27    }
28}
29
30#[cube]
31fn step_up(scale: f32, #[comptime] param: QuantParam) -> f32 {
32    // Mirrors QuantParam::round_up, saturating at the top rather than converting past it: above the
33    // maximum a conversion gives an infinity, and every value scaled by it then reconstructs wrong.
34    // Both paths below work on the f32 bit pattern rather than the storage type, because the
35    // narrowing conversion that would replace them is one the WGSL path leaves unrounded.
36    let grid = comptime!(param.f32_grid());
37    let max = comptime!(param.max_representable());
38
39    if scale >= max {
40        max
41    } else if comptime!(grid.subnormals.is_some()) {
42        let subnormals = comptime!(grid.subnormals.unwrap());
43        let spacing = comptime!(subnormals.spacing);
44
45        // Below the minimum normal the spacing stops halving, so the answer is a count of steps.
46        if scale < comptime!(subnormals.min_normal) {
47            f32::ceil(scale / spacing) * spacing
48        } else {
49            round_up_on_grid(scale, grid)
50        }
51    } else {
52        round_up_on_grid(scale, grid)
53    }
54}
55
56/// Rounds `scale` up onto `grid`, for a value in the param's normal range and below its maximum.
57///
58/// Truncating the low f32 mantissa bits lands on the grid, and biasing first turns that truncation
59/// into a round up.
60#[cube]
61fn round_up_on_grid(scale: f32, #[comptime] grid: F32Grid) -> f32 {
62    let bits = u32::reinterpret(scale);
63    let up_bits = (bits + comptime!(grid.round_up_bias())) & comptime!(grid.truncate_mask());
64    f32::reinterpret(up_bits)
65}
cubecl_std/quant/round.rs

cubecl_std/quant/
round.rs
