Skip to content

Tensor-carried dtype propagation, with quantized signals - #6220

Draft
nathanielsimard wants to merge 11 commits into
mainfrom
feat/quant-propagation
Draft

Tensor-carried dtype propagation, with quantized signals#6220
nathanielsimard wants to merge 11 commits into
mainfrom
feat/quant-propagation

Conversation

@nathanielsimard

@nathanielsimard nathanielsimard commented Oct 25, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

Checklist

  • I have reviewed and understand the changes in this PR.
  • Confirmed that cargo run-checks has been executed.
  • Made sure the book is up to date with changes in this PR.
  • Added or updated changelog entries for user-visible changes, if applicable, following the
    changelog policy in CONTRIBUTING.md.

Related Issues/PRs

Used by Metabolic's signal quantization (activations stored as MXFP8 between kernels, NVFP4 weights).

Changes

A tensor's dtype now propagates to the outputs it reaches, so a quantized activation stays quantized through a model without each op being written for it.

  • Propagation. Beside its DType, each tensor carries a flag: Propagate or Inhibit. Every float op resolves its output type from its inputs' through TensorType::combine. Only parameters inhibit, so a quantized weight doesn't turn its product quantized. Operands are decoded to f32 for compute and the output is quantized back to the resolved format. DTypeResolution::Strict panics when two propagating inputs disagree.
  • Custom ops. ResolvedOutput gives a backend extension's kernel the same resolution as a built-in op (of, written_into, read, write).
  • Autodiff. Quantized tensors are tracked, and the gradient passes straight through the rounding.
  • CubeCL.
    • A single-level, innermost-axis symmetric block scheme quantizes in one launch, laid out contiguously so a reshape of its leading axes stays metadata.
    • A slice that doesn't cut a block is a view of the words and scales.
    • A reshape the blocks can't follow is quantized again.
    • quantize_dynamic is forwarded through dispatch.
  • Fusion. Element-wise fusion decodes ue8m0 block scales in the fused kernel (MXFP8). The fix is that the scale argument is typed from the scheme instead of the handle's u8 storage; otherwise it multiplied by the raw code. ue4m3 scales still decode wrong when fused, so they stay unfused.
  • Book. Covers dtype propagation and gradients through quantized tensors.

The q_* methods are kept.

Testing

Branched from main at dd31726.

  • Flex, all targets: tensor 1913, tensor_f16 1387, autodiff 728. All pass.
  • ROCm with fusion, all targets (gfx1151). The only failures also fail on main at dd31726:
    • tensor: collective::emulated::should_gather_an_empty_dimension (an overflow at burn-cubecl/src/tensor/base.rs:404).
    • tensor_f16 (6) and autodiff_f16 (12): *_with_indices, topk and maxpool3d. The reduce-with-indices autotune plan can't execute at f16. main also fails ctc there.
  • New tests:
    • The one-launch quantize matches the calibrated path byte for byte, at f32. At f16 the calibrated path stores a block of zeros as the format's largest code.
    • A fused MXFP8 decode and add match the float input.
    • A packed slice is a view.
  • Updated tests: quantized ops keep their output quantized now, so the op tests compare against the expected values quantized again. nd_block_split_reshape quantizes again instead of panicking.
  • Metabolic on this branch: qwen3-4b with NVFP4 weights and an MXFP8 signal generates coherently on ROCm.
  • Not run: cargo run-checks, CUDA, wgpu.

AI Usage

Claude Code (Claude Opus 5.5) wrote the code, the tests and this description, and opened the PR, mostly on its own following a design agreed with the contributor. The contributor hasn't reviewed the final diff yet.

🤖 Generated with Claude Code

nathanielsimard and others added 11 commits October 25, 2026 08:02
@nathanielsimard @claude
feat(tensor): resolve every operation's output dtype from its inputs'…
aa9ee18 
… propagation

A tensor carries whether its dtype decides the dtype of the operations it
enters (DTypePropagation, beside the dtype in the bridge, never seen by a
backend). One rule, TensorType::combine, resolves the output of every
combining operation: the dtype of the propagating inputs, the widest on a
tie, f32 for two different dtypes as wide. Writes into a tensor keep the
target's type. Module parameters inhibit, so a weight never lends its
format; inputs, created tensors and casts propagate, so an activation
quantized once keeps every output it reaches quantized.

QuantPropagation leaves QuantConfig and the device settings, and backends
return float from quantized operations: the bridge requantizes when the
resolved type is quantized. DTypeResolution::Strict turns a widening tie
into a panic for test suites.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@nathanielsimard @claude
feat(autodiff): track quantized tensors, gradients passing straight t…
22b7e3a 
…hrough

Autodiff's quantized primitive is now AutodiffQTensor: the packed values with
the graph node gradients reach them through. Quantize and dequantize pass the
gradient through unchanged (a straight-through estimator, cast back to the
input's precision), layout operations undo themselves on the gradient, and
every other quantized operation differentiates through its float composition.

Dispatch carries a tracked quantized tensor as BackendTensor::AutodiffQuantized,
and the routing macro sends an operation with a quantized input in an enabled
context through the autodiff backend, lifting an untracked one to a leaf.
Module parameters hold their float values inhibiting, so a weight never
decides the dtype of the activations it meets.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@nathanielsimard @claude
feat(tensor): resolve a custom operation's output as a built-in's
5a91f28 
ResolvedOutput, the bridge's resolver, is public: a backend extension's kernel
resolves its output type from its inputs' with ResolvedOutput::of (or
written_into for an operation writing into a tensor), reads its float operands
in the precision it computes in and writes its output back to the resolved
type, so a custom operation keeps a quantized signal quantized as a built-in
one does. Tensor::tensor_type exposes the dtype with its propagation.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@nathanielsimard @claude
fix(cubecl): quantize a reshape again where its blocks cannot follow
5109f43 
A block-quantized reshape whose blocks or packing cannot follow as metadata
(a transposed view laid back in rows, a block split across rows) used to be
unimplemented!: it now decodes, reshapes and quantizes again, the block along
the innermost axis as an operation's output is. innermost_scheme is public for
it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@nathanielsimard @claude
test(quantization): reshape transposed block-quantized heads back int…
ad3a610 
…o rows

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@nathanielsimard @claude
docs(book): dtype propagation, and gradients through quantized tensors
fc1e64b 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@nathanielsimard @claude
perf(cubecl): one-pass block scales for quantize_dynamic
ba393c5 
A single-level block scheme along the innermost axis (the MX formats, q8/q4
blocks) computes its symmetric block scales in one launch, a unit per block,
instead of min-max calibration's chain of reshapes, reductions and arithmetic
(about eight launches); the existing quantize kernel then packs. Fusion runs
quantize_dynamic as one backend operation, so the backend's calibration is the
one that runs. Every other scheme keeps the composed calibration.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@nathanielsimard @claude
fix(dispatch): forward quantize_dynamic to the backend
902c7ea 
The dispatch layer took QTensorOps' default, so every quantize_dynamic was
decomposed into the composed calibration before fusion or a backend saw it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@nathanielsimard @claude
perf: quantize a single-level block scheme in one launch, laid out co…
27a82e7 
…ntiguously

The block scale and the packed values come from one unit per block, in place of
a scales launch then a quantize launch. The output is allocated contiguously, so
merging an activation's heads back into rows stays a change of metadata instead
of decoding and quantizing it again. Byte-identical to quantizing under the
min-max calibration's scales.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@nathanielsimard @claude
feat(fusion): decode ue8m0 block scales inside a fused kernel
c5c8592 
An MXFP8 or MXFP4 tensor read by an element-wise expression is decoded in the
fused kernel, as an integer or float-scaled one already was, rather than by a
dequantize launch of its own. A scale argument takes its type from the scheme,
not from its handle, which carries only a byte scale's width: read as a u8, a
ue8m0 code converts as the integer it is.

ue4m3 scales stay declined: read as e4m3, the fused decode reconstructs a
different value than the unfused one.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@nathanielsimard @claude
perf(cubecl): slice a quantized tensor as a view of its words and scales
74730d9 
A slice that leaves the packed axis whole and cuts no axis a block spans
narrows the values' region and the scales' region of the allocation, as a
float slice narrows its one, where both still start on a boundary the device
binds at. Splitting a quantized projection's heads no longer decodes the
whole of it and quantizes each part again.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@nathanielsimard