cubecl_std/quant/
round.rs

1use cubecl::prelude::*;
2use cubecl_common::quant::scheme::{F32Grid, ScaleDtype};
3use cubecl_core as cubecl;
4
5/// The smallest value representable in `dtype` that is not below `scale`, in a kernel.
6///
7/// Device-side counterpart of [`ScaleDtype::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/// `dtype`, 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_dtype<F: Float>(scale: F, #[comptime] dtype: ScaleDtype) -> F {
19    #[comptime]
20    match dtype {
21        ScaleDtype::F32 => scale,
22        ScaleDtype::F16 | ScaleDtype::BF16 | ScaleDtype::UE4M3 => {
23            F::cast_from(step_up(f32::cast_from(scale), dtype))
24        }
25        ScaleDtype::UE8M0 => F::cast_from(step_up_to_power_of_two(f32::cast_from(scale))),
26    }
27}
28
29/// The `ue8m0` arm of [`round_up_to_dtype`]: the smallest power of two not below `scale`.
30///
31/// Kept apart from [`step_up`] because both of `ue8m0`'s ends need clamping before the shared
32/// stepping means anything — its bottom, 2^-127, is subnormal in f32, and it has no zero for a
33/// fully-zero block to calibrate to. Mirrors the host's `round_up_to_power_of_two`, and the two
34/// have to keep agreeing or a tensor quantized on one backend reconstructs differently on another.
35#[cube]
36fn step_up_to_power_of_two(scale: f32) -> f32 {
37    let min = comptime!(ScaleDtype::UE8M0_MIN);
38    let max = comptime!(ScaleDtype::UE8M0_MAX);
39
40    if scale <= min {
41        min
42    } else if scale >= max {
43        max
44    } else {
45        round_up_on_grid(scale, comptime!(ScaleDtype::UE8M0.f32_grid()))
46    }
47}
48
49#[cube]
50fn step_up(scale: f32, #[comptime] dtype: ScaleDtype) -> f32 {
51    // Mirrors ScaleDtype::round_up, saturating at the top rather than converting past it: above the
52    // maximum a conversion gives an infinity, and every value scaled by it then reconstructs wrong.
53    // Both paths below work on the f32 bit pattern rather than the storage type, because the
54    // narrowing conversion that would replace them is one the WGSL path leaves unrounded.
55    let grid = comptime!(dtype.f32_grid());
56    let max = comptime!(dtype.max_representable());
57
58    if scale >= max {
59        max
60    } else if comptime!(grid.subnormals.is_some()) {
61        let subnormals = comptime!(grid.subnormals.unwrap());
62        let spacing = comptime!(subnormals.spacing);
63
64        // Below the minimum normal the spacing stops halving, so the answer is a count of steps.
65        if scale < comptime!(subnormals.min_normal) {
66            f32::ceil(scale / spacing) * spacing
67        } else {
68            round_up_on_grid(scale, grid)
69        }
70    } else {
71        round_up_on_grid(scale, grid)
72    }
73}
74
75/// Rounds `scale` up onto `grid`, for a value in the dtype's normal range and below its maximum.
76///
77/// Truncating the low f32 mantissa bits lands on the grid, and biasing first turns that truncation
78/// into a round up.
79#[cube]
80fn round_up_on_grid(scale: f32, #[comptime] grid: F32Grid) -> f32 {
81    let bits = u32::reinterpret(scale);
82    let up_bits = (bits + comptime!(grid.round_up_bias())) & comptime!(grid.truncate_mask());
83    f32::reinterpret(up_bits)
84}
cubecl_std/quant/round.rs

cubecl_std/quant/
round.rs
