cubek_quant/
dequantize_tiled.rs

1use cubecl::{
2    ir::ElemType,
3    prelude::*,
4    quant::scheme::{QuantScheme, QuantStore, ScaleDtype},
5};
6use cubek_tile::{
7    Axis, Launcher, Partitioning, Projection, Space, TileArg, TileArgLaunch, TileSpec, kind::Field,
8    launch::Grid, layout::PhysicalAxisMap,
9};
10
11// Input axes
12const M: Axis = Axis(0);
13const N: Axis = Axis(1);
14
15/// Convert the tensor back to a higher precision data type.
16/// Uses the tile-based implementation for dequantization.
17/// Very WIP and naive implementation for now.
18///
19/// The decode is the kernel's own statement: the values are fields of the scheme's element read
20/// out of words, and `output.copy_from(&values.mul(&scale))` multiplies each by the scale where it
21/// is copied. Nothing decodes behind a read.
22pub fn launch_ref(
23    client: &Client,
24    input: TensorBinding,
25    output: TensorBinding,
26    scales: TensorBinding,
27    scheme: &QuantScheme,
28    output_dtype: ElemType,
29) -> Result<(), LaunchError> {
30    assert!(
31        scheme.store == QuantStore::Native,
32        "only native quantization is supported for now."
33    );
34    assert!(
35        scheme.num_levels() == 1 && scheme.block_size().is_none(),
36        "only per tensor quantization is supported for now."
37    );
38    assert!(
39        scheme.scale_dtype() == ScaleDtype::F32,
40        "only f32 scales are supported for now."
41    );
42    assert!(
43        input.shape.len() == 2 && output.shape.len() == 2,
44        "dequantize_tiled: both operands must be plain 2-D tensors, got {:?} and {:?}",
45        input.shape,
46        output.shape
47    );
48
49    // A native store keeps one value a byte, so the buffer is read as words of fields: each
50    // `u32` holds `per_word` of them, decoded by the scheme's own field. The binding still counts
51    // values, as a packed operand's does.
52    let field = Field::Quant(scheme.value);
53    let per_word = field.per_word();
54    let (rows, cols) = (input.shape[0], input.shape[1]);
55    assert!(
56        field.size_bits() == 8,
57        "dequantize_tiled: a native store holds one value a byte, and {:?} is {} bits",
58        scheme.value,
59        field.size_bits()
60    );
61    assert!(
62        input.strides[1] == 1 && input.strides[0] == cols && cols.is_multiple_of(per_word),
63        "dequantize_tiled: the values are read as whole words, so each row must be contiguous \
64         and a whole number of words ({per_word} values), got shape {:?} and strides {:?}",
65        input.shape,
66        input.strides
67    );
68    let max = client.properties().hardware.max_vector_size;
69    assert!(
70        per_word <= max,
71        "dequantize_tiled: a word of {per_word} values is read as one line, and this device's \
72         lines hold {max}"
73    );
74
75    // One tile covering every axis, walked by a single cube: no level cuts it, so there is
76    // nothing to list and the grid is one cube.
77    let space = Space::new(&[(M, rows), (N, cols)]);
78    let plane_size = client.properties().hardware.plane_size_max;
79    let (cube_count, cube_dim) = (CubeCount::Static(1, 1, 1), CubeDim::new_2d(plane_size, 1));
80    let launch = Launcher::new(
81        client,
82        // Static: the decoding copy reads through the values' own axes, which it needs the
83        // extents of at expansion, so the kernel is compiled per shape.
84        Partitioning::new(space.clone(), vec![]),
85        &space,
86        Grid::Stated {
87            cube_count: cube_count.clone(),
88            cube_dim,
89        },
90    )
91    .unwrap_or_else(|refusal| panic!("dequantize_tiled: {refusal}"));
92    // One scale for the whole tensor: a scale tile over no axis of the values.
93    let per_tensor = Projection::new(&[M, N], &[PhysicalAxisMap::broadcast()]);
94    dequantize::launch(
95        client,
96        cube_count,
97        cube_dim,
98        per_word,
99        TileArgLaunch::new(
100            input.into_tensor_arg(),
101            TileSpec::direct(&[M, N]).packed(field),
102        ),
103        TileArgLaunch::new(scales.into_tensor_arg(), TileSpec::new(per_tensor)),
104        TileArgLaunch::new(output.into_tensor_arg(), TileSpec::direct(&[M, N])),
105        launch.partitioning_arg(),
106        output_dtype,
107    );
108
109    Ok(())
110}
111
112#[cube(launch)]
113/// The values are fields of the scheme's element read out of `u32` words, served as `O`; the
114/// scale is a second operand, multiplied in by the copy that says so.
115pub fn dequantize<O: Numeric, V: Size>(
116    input: &TileArg<'_, u32, Const<1>>,
117    scale: &TileArg<'_, f32, Const<1>>,
118    output: &TileArg<'_, O, V>,
119    space: Partitioning,
120    #[define(O)] _output_dtype: ElemType,
121) {
122    let input = input.tile_as::<O>(comptime!(space.clone()));
123    let scale = scale.tile(comptime!(space.clone()));
124    let mut output = output.tile(comptime!(space.clone()));
125    output.copy_from(&input.mul(&scale));
126}
cubek_quant/dequantize_tiled.rs

cubek_quant/
dequantize_tiled.rs
