Draft: QuantizedDType for normal tensors #3174
Draft
+3,039
−108
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.
Disclaimer
This PR is still a work in progress and needs some polishing.
The main goal of opening it early is to gather feedback on whether this approach is heading in the right direction.
Overview
The primary goal of this PR is to introduce
QuantizedDTypesupport.These data types allow regular tensors to store and operate on quantized data, making it much easier to add new quantization schemes to Candle in the future without relying on custom ops.
What’s Included
1. New
candle-macroscratescandle-macros: contains the procedural macros that generate the dispatch code forQuantizedDTypes.candle-macros-types: defines the traits that quantized types can implement to provide backend-specific support.2.
QuantizedTypetraitEach quantization implements the
QuantizedTypetrait, which defines:NAMEQuantizations can optionally implement one or more of the following backend traits:
QuantizedCpuOpsQuantizedCudaOpsQuantizedMetalOpsThese traits define the de/quantization logic and backend-specific matmul implementations (e.g.,
f32 × quantized).3. The
register_quantized_types!macroThis macro generates:
QuantizedDTypeenumThe enum is then integrated into Candle Core as a new
DType::Quantized(QuantizedDType)variant.Tensors using this type:
f32)4. External Quantization Support
A
register_external_quantized_type!macro is also included.This will allow external crates to register their own quantization types without modifying Candle Core directly.
Current Limitations
.to_dtype()— I haven’t yet found a clean way to load them from files.f32as the intermediate type isn’t ideal for some backends (like CUDA) and may need refinement.