Add a tiled dequantize kernel - #502
Draft
ThierryCantin-Demers wants to merge 15 commits into
Draft
Conversation
QuantTileArg and the source builder gain the per-tensor scale as an optional 1-element f32 binding, read once per kernel and carried on QuantInfo. A staged side-channel folds it into its scale grid (stage_scales writes global * local), so everything below a stage serves the one-level form of the scheme; only the unstaged path hands the global to the dequantizing view. The binding contract is cubecl's check_global_bindings, at the builder and again at tile().
A window whose extent fits inside a block has a single scale, already folded into window_start at descent, so its per-axis divide and multiply can only answer zero. QuantInfo carries the window's comptime extent and ScaleLayout drops those terms; a tile cut to its block size reads scales[window_start] with no arithmetic left.
A window that sits inside a block reconstructs against a single scale, so the descent reads it there, folds the per-tensor factor in, and serves the values through a view that never consults its scales at all. What was a load per read becomes a register multiply.
cubecl replaced QuantLevel with ScaleLevels, an innermost-first list where each level owns its block size and scale precision, and split the quantized view's one constructor into a named one per scale-in-register mode. Per-tensor is now the full block rather than its own variant, so a level test reads block_size().is_full() and a two-level scheme reads levels().len() > 1. Three rewrites cubek was spelling out come from the scheme now: staged_scheme is scale_levels().inner(), the swap_dims block rewrite is swap_block_dims, and block_dims is resolved_dims. The levels list being private also turns the elementwise kernels' level: patterns into guards on the level count, since a private field cannot be matched. Restores the tile path's uniform-scale read, which had been left calling a constructor that only existed on an unmerged cubecl branch.
`quantized` and `quantized_two_level` said the same thing twice, and their `global` named the per-tensor scale as an exception rather than the outer level it is. One constructor now takes one binding per scheme level, innermost first, and splits to (inner, outer) once, mirroring cubecl's own `ScaleBindings`. The count is checked against the scheme where it already was, so a caller that forgets the outer binding is refused exactly as before. Only the innermost level is addressed per position, so the list stops at the builder: below it the kernel argument keeps a scales tensor plus an optional whole-tensor scale, which is what the levels bind as. `global` becomes `outer` everywhere, cubecl's word for the same thing (`ScaleInRegister::OuterLevels` is the product of every outer level), so the name survives a level cap above two. Also puts `window`'s doc comment back on `window`, which `uniform_scale` had been inserted in front of.
cubecl, pinned here at `7eead27c`, split whole-tensor granularity out of `BlockSize`: a level now states `ScaleGranularity::Tensor` or `Block(BlockSize)`, `block_size()` answers `Option<BlockSize>`, and `full()`/`is_full()` are gone. Every per-tensor test becomes the `None` arm. The `Option` fits the control flow that was already there: the two guards in the naive setup fold into their let-chains and stop asking for the block twice, and the scales and global scale layouts become a `match`, which also retires a doubled brace block in each. No scheme changes meaning. `BlockSize` is rank-relative now, so an all-`FULL` block no longer collapses to whole-tensor, but nothing here builds a block with a `FULL` dimension.
cubecl replaced the scale-level list with two optional fields on the scheme, set additively by `per_tensor` and `per_block` in any order, and renamed `QuantParam` to `ScaleDtype` with `param()` becoming `scale_dtype()`. `ScaleLevels`, `ScaleGranularity` and the level-list accessors are gone, so a level count reads `num_levels()` and the two test helpers that took a level list now take none: their callers chain the level on instead. `staged_scheme` rebuilds the scheme rather than stripping its outer level. The fields are private and set additively, so there is no way to clear one.
…emes block_scheme()/scheme_block() chained .per_block() onto a scheme that already had .per_tensor() set. QuantScheme's new builder is additive, so this silently built a two-level scheme instead of a one-level block one, and every extended-tier block test and benchmark panicked instead of running. Also: - add the missing num_levels() > 1 guards in InputBinding::swap_dims and check_block_size_compat, matching every other rewritten call site - give per-tensor schemes a real block edge (usize::MAX, not 1) so uniform()'s read-once fast path actually fires for the common case - cover uniform_scale()'s outer-scale fold with a test whose window fits inside one block - dedupe the three-way DequantView dispatch into QuantInfo::dequant_view - drop ScaleLayout::kept() in favor of Sequence::len() - parameterize run_quantized_block with an outer scale instead of forking it as run_quantized_two_level - fix a misleading panic-message comment and drop leftover whitespace
The tile path reads a two-level scheme through QuantizedView, which narrowed the effective scale to the output type before multiplying it in, so a scale below f16's smallest subnormal came back as zero for the whole tensor. cubecl bdf99182 forms that product in f32.
The scale-levels migration rewrote these two against an intermediate cubecl API, and nothing compiles the extended tier by default, so the rename that followed left them behind: `as_type_native_unchecked` no longer exists and `cargo test --features extended` failed to build, taking every matmul quantization test with it.
…obal cubecl renamed it, on the grounds that `outer` names the level by its position in the list while `global` names what it holds: one scale for the whole tensor, the factor every block scale is normalized against.
ThierryCantin-Demers
changed the title
A per-tensor scale covers every window, but its block edges were stated as 1, so validate_scheme rejected any vectorized per-tensor operand (1 does not divide a wider line) and uniform_window never held, costing a scale address computation per line where one uniform read suffices. State the edges as usize::MAX and skip the straddling rules when a scheme has no block: there is one scale, nothing can straddle it.
ThierryCantin-Demers
force-pushed
the
feat/tile-copy-op
branch
from
August 17, 2026 19:22
b939fee to
ea6aa1f
Compare
ThierryCantin-Demers
force-pushed
the
feat/quant-dequantize-tiled
branch
from
August 17, 2026 19:22
9704989 to
9a785e3
Compare
Reading one tile into another was a hand-written walk at every call site: build the witnessed space, iterate its regions, and transport each one. The walk is a property of the tiles, not of the caller, so it belongs on the tile. `dst.copy(&src)` walks while a level spreads regions across cubes, and transports once it does not. The transport fills a whole tile with every unit of the cube, so a level that spreads inside a cube is already its own doing and ends the walk. Dropping that stopping rule leaves three quarters of a multi-plane tile unwritten, which the new spread test pins.
The tile-path dequantize was a single-cube per-tensor 2-D demo. Give it the shape of a real elementwise kernel: rank-agnostic axes, a served width derived from the device's IO widths against both bindings' layout and the scheme's blocks, a cube grid over every axis with an optional plane level raising the cube's unit count, and a region walk in the kernel so each cube copies only its own tiles. Block, packed-u32 and two-level schemes ride the same body; a shape nothing divides degrades to a scalar bounds-checked launch instead of being refused. The native i8 gate moves to utils so both dequantize paths share it.
Block, packed-u32 block, two-level (with the zero-global mutation check), an awkward shape that degrades to the scalar plan, and a rank-3 tensor, each compared against host expectations built from the raw q values and scales rather than a round-trip.
ThierryCantin-Demers
force-pushed
the
feat/tile-copy-op
branch
from
August 17, 2026 20:09
ea6aa1f to
8733a67
Compare
ThierryCantin-Demers
force-pushed
the
feat/quant-dequantize-tiled
branch
from
August 17, 2026 20:09
9a785e3 to
4f1f9cd
Compare
Base automatically changed from
feat/tile-copy-op
to
feat/tile-two-level-quant
August 18, 2026 14:08
ThierryCantin-Demers
force-pushed
the
feat/tile-two-level-quant
branch
from
August 18, 2026 14:17
2641777 to
f29d201
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
A dequantize kernel built on the tile engine: it plans its geometry host-side and its body is one
Tile::copy, since the input tile serves the output element and decodes on read. Stacked on #501.Not wired into
dequantize::launch_ref. Nothing routes to it yet, on purpose: it becomes a real path once it is fully implemented and tested.What changed
u32words and sits inside one scale block. It then picks a cube tile edge, and a plane cut when one raises the cube's unit count without straddling a scale block.Testing
Unit tests for the width and geometry choices, and device tests over the schemes it serves plus the shapes and strides that fall back. 8 tiled tests and 10 unit tests pass; tile quant tests 38 passed.