ruda_tensor/
frozen_awq.rs

1//! Optional backend extension for frozen, directly packed AWQ projections.
2use crate::{Backend, tensor::{FloatTensor, IntTensor}};
3use core::fmt::{Debug, Display, Formatter};
4
5/// A native backend failure or a violation of the explicitly frozen base contract.
6#[derive(Debug)]
7pub enum FrozenAwqError<E: Debug> {
8    /// Original backend error, without dense fallback or requantization.
9    Native(E),
10    /// The caller supplied a trainable base scale or bias to a frozen operator.
11    TrainableBase,
12}
13impl<E: Debug> Display for FrozenAwqError<E> {
14    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
15        match self {
16            Self::Native(error) => write!(f, "native AWQ: {error:?}"),
17            Self::TrainableBase => f.write_str("frozen AWQ scales and bias must not require gradients"),
18        }
19    }
20}
21impl<E: Debug> core::error::Error for FrozenAwqError<E> {}
22
23/// Direct AWQ `[K,N/8]` I32 words, `[K/group,N/8]` I32 zero points and
24/// `[K/group,N]` FP16/BF16/FP32 scales. The AWQ nibble order and scale-dtype
25/// coefficient rounding are preserved. This is not NF4 or a generic QAT STE.
26/// Backends implement this explicitly; there is no dequantized default.
27pub trait FrozenAwqOps: Backend {
28    /// Native validation/launch failure, including autodiff frozen-base validation.
29    type AwqError: Debug;
30
31    /// Native packed projection, optionally biased. Output retains input dtype
32    /// and leading axes; input storage may differ from original scale storage.
33    fn frozen_awq_forward(
34        input: FloatTensor<Self>, qweight: IntTensor<Self>, qzeros: IntTensor<Self>,
35        scales: FloatTensor<Self>, bias: Option<FloatTensor<Self>>, group_size: usize,
36    ) -> Result<FloatTensor<Self>, Self::AwqError>;
37
38    /// Apply the transpose of the same rounded packed coefficients, retaining
39    /// gradient dtype/leading axes. No weight/scales/bias gradients are implied.
40    fn frozen_awq_input_backward(
41        gradient: FloatTensor<Self>, qweight: IntTensor<Self>, qzeros: IntTensor<Self>,
42        scales: FloatTensor<Self>, group_size: usize,
43    ) -> Result<FloatTensor<Self>, Self::AwqError>;
44}
ruda_tensor/frozen_awq.rs

ruda_tensor/
frozen_awq.rs
