fuel_backend_contract/
quantized.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! Backend-agnostic quantized kernel registration traits.
3//!
4//! Per-backend quantized fast paths (CUDA, Metal, future Vulkan/CPU-SIMD)
5//! implement [`DynQuantizedStorage`] for their concrete storage type and
6//! [`QuantizedDeviceKernels`] for their device handle. fuel-core dispatches
7//! through these traits without naming concrete backend types.
8//!
9//! The CPU "fast path" is the bare ggml block math (avx/neon
10//! vec_dots inside `k_quants.rs`); it stays in fuel-core because
11//! `BlockQX` types and the file-format readers (gguf/ggml/imatrix) are a
12//! single cohesive numerics unit. The trait still covers the CPU case via
13//! a fuel-core-side adapter (see `fuel-core/src/quantized/mod.rs`).
14//!
15//! The `GgmlDType` block-format *data* tag lives in [`fuel_ir::quantized`].
16
17use crate::dyn_backend::{DynBackendDevice, DynBackendStorage};
18use fuel_ir::{Error, GgmlDType, Layout, Result, Shape};
19use std::any::Any;
20use std::borrow::Cow;
21
22/// Object-safe per-backend quantized storage. Each backend (CPU, CUDA,
23/// Metal, ...) supplies a concrete type that implements this trait;
24/// fuel-core holds them as `Box<dyn DynQuantizedStorage>` and dispatches
25/// without naming backends.
26///
27/// The `_src` arguments to `quantize*` are typed as `&dyn DynBackendStorage`
28/// so the implementor can downcast to its own concrete storage; the
29/// `_onto` variants take a `HostBuffer`-shaped CPU source.
30pub trait DynQuantizedStorage: Send + Sync + std::fmt::Debug {
31    fn dtype(&self) -> GgmlDType;
32    fn block_size(&self) -> usize;
33    fn storage_size_in_bytes(&self) -> usize;
34
35    /// Quantize an in-device source storage onto self.
36    fn quantize(&mut self, src: &dyn DynBackendStorage) -> Result<()>;
37
38    /// Quantize with importance-matrix weighting.
39    fn quantize_imatrix(
40        &mut self,
41        src: &dyn DynBackendStorage,
42        imatrix_weights: &[f32],
43        n_per_row: usize,
44    ) -> Result<()>;
45
46    /// Quantize a CPU source onto self (cross-device).
47    fn quantize_onto(&mut self, src: &dyn DynBackendStorage) -> Result<()>;
48
49    /// Quantize with importance matrix from a CPU source.
50    fn quantize_imatrix_onto(
51        &mut self,
52        src: &dyn DynBackendStorage,
53        imatrix_weights: &[f32],
54        n_per_row: usize,
55    ) -> Result<()>;
56
57    /// Dequantize to backend-native f32 storage.
58    fn dequantize(&self, elem_count: usize) -> Result<Box<dyn DynBackendStorage>>;
59
60    /// Dequantize to f16 (CUDA fast path; default impl rejects).
61    fn dequantize_f16(&self, _elem_count: usize) -> Result<Box<dyn DynBackendStorage>> {
62        Err(Error::Msg("dequantize_f16 not supported on this backend".into()).bt())
63    }
64
65    /// Raw bytes of the quantized data (host-readable copy).
66    fn data(&self) -> Result<Cow<'_, [u8]>>;
67
68    /// Device pointer for callers that need raw addressing (CUDA only today).
69    fn device_ptr(&self) -> Result<*const u8> {
70        Err(Error::Msg("device_ptr not supported on this backend".into()).bt())
71    }
72
73    /// QMatMul forward against an in-device input storage. Returns a fresh
74    /// device storage and its output shape. `self_shape` is the weight shape
75    /// (the QTensor's own shape).
76    fn fwd(
77        &self,
78        self_shape: &Shape,
79        input: &dyn DynBackendStorage,
80        layout: &Layout,
81    ) -> Result<(Box<dyn DynBackendStorage>, Shape)>;
82
83    /// Indexed MoE forward (CUDA-only today; default rejects).
84    fn indexed_moe_forward(
85        &self,
86        _self_shape: &Shape,
87        _input: &dyn DynBackendStorage,
88        _input_layout: &Layout,
89        _ids: &dyn DynBackendStorage,
90        _ids_layout: &Layout,
91    ) -> Result<(Box<dyn DynBackendStorage>, Shape)> {
92        Err(Error::Msg("indexed_moe_forward not supported on this backend".into()).bt())
93    }
94
95    /// Identity returns the host device tag this storage is on.
96    fn as_any(&self) -> &dyn Any;
97
98    /// Owning device handle. Lets fuel-core recover a `Device` from a
99    /// `QTensor` without keeping a parallel `device` field on QTensor.
100    fn device_arc_dyn(&self) -> std::sync::Arc<dyn DynBackendDevice>;
101}
102
103/// Backend device → quantized-kernel constructors. fuel-core looks this up
104/// per Device to allocate fresh QStorage of a given dtype, or to load
105/// pre-quantized bytes.
106pub trait QuantizedDeviceKernels: DynBackendDevice {
107    /// Allocate a zero-initialized quantized storage on this device.
108    fn qzeros(&self, elem_count: usize, dtype: GgmlDType) -> Result<Box<dyn DynQuantizedStorage>>;
109
110    /// Load pre-quantized block-format bytes onto this device. `dtype`
111    /// describes the block format; `data` is interpreted as a flat byte
112    /// slice of those blocks.
113    fn load_quantized(
114        &self,
115        dtype: GgmlDType,
116        data: Cow<'_, [u8]>,
117    ) -> Result<Box<dyn DynQuantizedStorage>>;
118}
fuel_backend_contract/quantized.rs

fuel_backend_contract/
quantized.rs
