ruda_tensor/
frozen_nf4.rs

1//! Backend extension for the original RUDA high-nibble-first NF4 format.
2use crate::{Backend,FloatDType,tensor::{FloatTensor,IntTensor}};
3use core::fmt;
4
5/// Explicit original geometry and execution choices; no quantizer or model family is inferred.
6#[derive(Clone,Copy,Debug,PartialEq,Eq)]
7pub struct Nf4ProjectionOptions {
8    /// Actual logical input width.
9    pub input_features:usize,
10    /// Actual logical output width, with `[output,input]` packed row-major storage.
11    pub output_features:usize,
12    /// Original positive even flat block size, including partial final blocks.
13    pub block_size:usize,
14    /// Maximum decoded output-row tile on the explicitly tiled path.
15    pub tile_rows:usize,
16    /// Half/BF16 uses the original fused cooperative kernel when true. FP32
17    /// remains tiled. Failed fused calls never select the tiled path as recovery.
18    pub use_tensor_core:bool,
19}
20
21/// Native failure or an unsupported derivative of the explicitly frozen NF4 base.
22#[derive(Debug)]
23pub enum FrozenNf4Error<E:fmt::Debug> {
24    /// Original native backend error.
25    Native(E),
26    /// Scales, codebook or bias must remain frozen, not learned QAT operands.
27    TrainableBase,
28    /// Original NF4 API provides first-order input gradients, not their derivatives.
29    InputGradientNotDifferentiable,
30}
31impl<E:fmt::Debug> fmt::Display for FrozenNf4Error<E> {
32    fn fmt(&self,f:&mut fmt::Formatter<'_>) -> fmt::Result {
33        match self {Self::Native(error)=>write!(f,"native NF4: {error:?}"),Self::TrainableBase=>f.write_str("NF4 scales/codebook/bias must be frozen"),
34            Self::InputGradientNotDifferentiable=>f.write_str("NF4 provides first-order input gradients only")}
35    }
36}
37impl<E:fmt::Debug> core::error::Error for FrozenNf4Error<E> {}
38
39/// Actual U8 byte-packed NF4 with FP32 scales and FP32 codebook. This is not
40/// AWQ, bitsandbytes double quantization or a floating fake-quantization STE.
41pub trait FrozenNf4Ops:Backend {
42    /// Original native or autodiff contract error.
43    type Nf4Error:fmt::Debug;
44    /// Native packed projection. Output retains incoming floating activation storage.
45    fn frozen_nf4_forward(input:FloatTensor<Self>,packed:IntTensor<Self>,scales:FloatTensor<Self>,codebook:FloatTensor<Self>,
46        bias:Option<FloatTensor<Self>>,options:Nf4ProjectionOptions) -> Result<FloatTensor<Self>,Self::Nf4Error>;
47    /// Original input VJP: cast incoming gradient to original activation storage,
48    /// accumulate FP32 and return that original activation dtype, without base gradients.
49    fn frozen_nf4_input_backward(gradient:FloatTensor<Self>,packed:IntTensor<Self>,scales:FloatTensor<Self>,codebook:FloatTensor<Self>,
50        options:Nf4ProjectionOptions,activation_dtype:FloatDType) -> Result<FloatTensor<Self>,Self::Nf4Error>;
51}
ruda_tensor/frozen_nf4.rs

ruda_tensor/
frozen_nf4.rs
