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}
41impl<B:Backend> PackedExpertPayload<B> {
42    /// Actual resident expert count and explicitly declared first global expert ID.
43    pub fn expert_range(&self) -> (usize,usize) {match self {Self::Nf4(value)=>(value.options.experts,value.options.expert_start),
44        Self::Awq(value)=>(value.options.experts,value.options.expert_start)}}
45}
46/// Original native failure or unsupported differentiation of frozen packed metadata/VJPs.
47#[derive(Debug)]
48pub enum PackedExpertAutodiffError<E:fmt::Debug> {
49    /// Original native projection/routing error.
50    Native(E),
51    /// Packed scales/codebook/bias must remain frozen; QAT is not this operation.
52    TrainableMetadata,
53    /// Original native input VJPs provide first-order derivatives only.
54    HigherDerivativeUnsupported,
55}
56impl<E:fmt::Debug> fmt::Display for PackedExpertAutodiffError<E> {
57    fn fmt(&self,f:&mut fmt::Formatter<'_>) -> fmt::Result {match self {Self::Native(error)=>write!(f,"native packed expert: {error:?}"),
58        Self::TrainableMetadata=>f.write_str("packed expert quantization metadata and bias must be frozen"),
59        Self::HigherDerivativeUnsupported=>f.write_str("native packed expert input VJPs provide first-order derivatives only")}}
60}
61impl<E:fmt::Debug> core::error::Error for PackedExpertAutodiffError<E> {}
62/// Native grouped projections and SwiGLU with independently selected original gate/up/down formats.
63pub trait FrozenPackedExpertOps:Backend {
64    /// Original native or first-order derivative contract failure.
65    type PackedExpertError:fmt::Debug;
66    /// Actual packed operand and native private row mapping.
67    type PackedProjectionState:Clone+Send+fmt::Debug+'static;
68    /// Actual original row mapping and optional real input VJP intermediates.
69    type PackedSwiGluState:Clone+Send+fmt::Debug+'static;
70    /// Evaluate only actual assigned experts and restore exact incoming row order.
71    fn packed_expert_forward(input:FloatTensor<Self>,global_ids:IntTensor<Self>,payload:PackedExpertPayload<Self>)
72        -> Result<(FloatTensor<Self>,Self::PackedProjectionState),Self::PackedExpertError>;
73    /// First-order packed projection input VJP in original activation storage.
74    fn packed_expert_input_backward(state:Self::PackedProjectionState,gradient:FloatTensor<Self>) -> Result<FloatTensor<Self>,Self::PackedExpertError>;
75    /// Complete native selected gate/up/down and original storage-rounded SwiGLU.
76    /// AD preserves required actual input caches even when `retain_input` is false.
77    fn packed_swiglu_forward(input:FloatTensor<Self>,global_ids:IntTensor<Self>,gate:PackedExpertPayload<Self>,up:PackedExpertPayload<Self>,down:PackedExpertPayload<Self>,retain_input:bool)
78        -> Result<(FloatTensor<Self>,Self::PackedSwiGluState),Self::PackedExpertError>;
79    /// Original first-order input VJP through the actual selected packed chain.
80    fn packed_swiglu_input_backward(state:Self::PackedSwiGluState,gradient:FloatTensor<Self>) -> Result<FloatTensor<Self>,Self::PackedExpertError>;
81}
ruda_tensor/packed_experts.rs

ruda_tensor/
packed_experts.rs
