ruda_tensor/
packed_experts.rs

1//! Actual original AWQ/NF4 selected expert payloads and their native first-order chains.
2use crate::{Backend,grouped_nf4::Nf4ExpertPayload,tensor::{FloatTensor,IntTensor}};
3use core::fmt;
4
5/// Explicit resident AWQ geometry and global expert range, without guessed source metadata.
6#[derive(Clone,Copy,Debug,PartialEq,Eq)]
7pub struct AwqExpertOptions {
8    /// Actual expert matrices represented by original rank-three words/scales.
9    pub experts:usize,
10    /// First global expert ID represented by this resident cube.
11    pub expert_start:usize,
12    /// Original per-expert input width.
13    pub input_features:usize,
14    /// Original per-expert output width, divisible by eight.
15    pub output_features:usize,
16    /// Original complete input-channel group size.
17    pub group_size:usize,
18}
19/// Original AWQ expert cube, with scale storage independent of floating activation storage.
20#[derive(Clone,Debug)]
21pub struct AwqExpertPayload<B:Backend> {
22    /// Original permuted I32 words `[E,K,N/8]`.
23    pub qweight:IntTensor<B>,
24    /// Original I32 zero-point words `[E,K/group,N/8]`.
25    pub qzeros:IntTensor<B>,
26    /// Original FP32/FP16/BF16 scales `[E,K/group,N]`.
27    pub scales:FloatTensor<B>,
28    /// Actual optional original frozen scale-dtype bias `[E,N]`.
29    pub bias:Option<FloatTensor<B>>,
30    /// Exact source geometry and explicit original global expert range.
31    pub options:AwqExpertOptions,
32}
33/// Source-selected original packed expert format, not a quantizer or dense surrogate.
34#[derive(Clone,Debug)]
35pub enum PackedExpertPayload<B:Backend> {
36    /// Original RUDA high-nibble-first flat-block NF4.
37    Nf4(Nf4ExpertPayload<B>),
38    /// Original AWQ per-input-group permuted eight-code I32 words.
39    Awq(AwqExpertPayload<B>),
40    /// Actual owned NF4 bytes beginning at the first original scale-block boundary.
41    /// The first owned logical coefficient can start within that retained block.
42    Nf4Window {payload:Nf4ExpertPayload<B>,element_offset:usize},
43}
44impl<B:Backend> PackedExpertPayload<B> {
45    /// Actual resident expert count and explicitly declared first global expert ID.
46    pub fn expert_range(&self) -> (usize,usize) {match self {Self::Nf4(value)=>(value.options.experts,value.options.expert_start),
47        Self::Nf4Window {payload,..}=>(payload.options.experts,payload.options.expert_start),
48        Self::Awq(value)=>(value.options.experts,value.options.expert_start)}}
49}
50/// Original native failure or unsupported differentiation of frozen packed metadata/VJPs.
51#[derive(Debug)]
52pub enum PackedExpertAutodiffError<E:fmt::Debug> {
53    /// Original native projection/routing error.
54    Native(E),
55    /// Packed scales/codebook/bias must remain frozen; QAT is not this operation.
56    TrainableMetadata,
57    /// Original native input VJPs provide first-order derivatives only.
58    HigherDerivativeUnsupported,
59}
60impl<E:fmt::Debug> fmt::Display for PackedExpertAutodiffError<E> {
61    fn fmt(&self,f:&mut fmt::Formatter<'_>) -> fmt::Result {match self {Self::Native(error)=>write!(f,"native packed expert: {error:?}"),
62        Self::TrainableMetadata=>f.write_str("packed expert quantization metadata and bias must be frozen"),
63        Self::HigherDerivativeUnsupported=>f.write_str("native packed expert input VJPs provide first-order derivatives only")}}
64}
65impl<E:fmt::Debug> core::error::Error for PackedExpertAutodiffError<E> {}
66/// Native grouped projections and SwiGLU with independently selected original gate/up/down formats.
67pub trait FrozenPackedExpertOps:Backend {
68    /// Original native or first-order derivative contract failure.
69    type PackedExpertError:fmt::Debug;
70    /// Actual packed operand and native private row mapping.
71    type PackedProjectionState:Clone+Send+fmt::Debug+'static;
72    /// Actual original row mapping and optional real input VJP intermediates.
73    type PackedSwiGluState:Clone+Send+fmt::Debug+'static;
74    /// Evaluate only actual assigned experts and restore exact incoming row order.
75    fn packed_expert_forward(input:FloatTensor<Self>,global_ids:IntTensor<Self>,payload:PackedExpertPayload<Self>)
76        -> Result<(FloatTensor<Self>,Self::PackedProjectionState),Self::PackedExpertError>;
77    /// First-order packed projection input VJP in original activation storage.
78    fn packed_expert_input_backward(state:Self::PackedProjectionState,gradient:FloatTensor<Self>) -> Result<FloatTensor<Self>,Self::PackedExpertError>;
79    /// Complete native selected gate/up/down and original storage-rounded SwiGLU.
80    /// AD preserves required actual input caches even when `retain_input` is false.
81    fn packed_swiglu_forward(input:FloatTensor<Self>,global_ids:IntTensor<Self>,gate:PackedExpertPayload<Self>,up:PackedExpertPayload<Self>,down:PackedExpertPayload<Self>,retain_input:bool)
82        -> Result<(FloatTensor<Self>,Self::PackedSwiGluState),Self::PackedExpertError>;
83    /// Original first-order input VJP through the actual selected packed chain.
84    fn packed_swiglu_input_backward(state:Self::PackedSwiGluState,gradient:FloatTensor<Self>) -> Result<FloatTensor<Self>,Self::PackedExpertError>;
85}
ruda_tensor/packed_experts.rs

ruda_tensor/
packed_experts.rs
