ruda_tensor/
grouped_nf4.rs

1//! Frozen original RUDA NF4 expert cubes, with discrete native U32 row assignments.
2use crate::{Backend,frozen_nf4::Nf4ProjectionOptions,tensor::{FloatTensor,IntTensor}};
3use core::fmt;
4
5/// Actual logical cube and explicitly owned global expert range, without per-expert reblocking.
6#[derive(Clone,Copy,Debug,PartialEq,Eq)]
7pub struct Nf4GroupedOptions {
8    /// Number of actual resident expert matrices.
9    pub experts:usize,
10    /// First global expert ID represented by this actual local cube.
11    pub expert_start:usize,
12    /// Per-expert geometry and original execution choices.
13    pub projection:Nf4ProjectionOptions,
14}
15/// Actual source expert payload; frozen FP32 scales/book never become activation storage.
16#[derive(Clone,Debug)]
17pub struct Nf4ExpertPayload<B:Backend> {
18    /// Original high-nibble-first U8 bytes for flattened `[E,N,K]`.
19    pub packed:IntTensor<B>,
20    /// Original FP32 flat-block scales, including blocks crossing expert boundaries.
21    pub scales:FloatTensor<B>,
22    /// Original sixteen FP32 codebook values.
23    pub codebook:FloatTensor<B>,
24    /// Explicit source geometry, ownership and execution choices.
25    pub options:Nf4GroupedOptions,
26}
27/// Native selected expert projection and its real first-order input derivative.
28pub trait FrozenNf4GroupedOps:Backend {
29    /// Native projection/routing failure or unsupported derivative contract.
30    type Nf4GroupedError:fmt::Debug;
31    /// Validated original row permutation and actual frozen packed operands.
32    type Nf4GroupedState:Clone+Send+fmt::Debug+'static;
33    /// Execute one actual assigned expert per incoming row and preserve original row order.
34    fn frozen_nf4_grouped_forward(input:FloatTensor<Self>,global_ids:IntTensor<Self>,payload:Nf4ExpertPayload<Self>)
35        -> Result<(FloatTensor<Self>,Self::Nf4GroupedState),Self::Nf4GroupedError>;
36    /// Original first-order input VJP, returned in original activation storage.
37    fn frozen_nf4_grouped_input_backward(state:Self::Nf4GroupedState,gradient:FloatTensor<Self>) -> Result<FloatTensor<Self>,Self::Nf4GroupedError>;
38}
39/// Selected original NF4 gate/up/down plus source-native storage-rounded SwiGLU.
40pub trait FrozenNf4SwiGluOps:FrozenNf4GroupedOps {
41    /// Original native row mapping and optional actual input VJP intermediates.
42    type Nf4SwiGluState:Clone+Send+fmt::Debug+'static;
43    /// Execute the original selected frozen expert chain. AD always retains the
44    /// cache required by actual tracked input, independently of `retain_input`.
45    fn frozen_nf4_swiglu_forward(input:FloatTensor<Self>,global_ids:IntTensor<Self>,gate:Nf4ExpertPayload<Self>,up:Nf4ExpertPayload<Self>,down:Nf4ExpertPayload<Self>,retain_input:bool)
46        -> Result<(FloatTensor<Self>,Self::Nf4SwiGluState),Self::Nf4GroupedError>;
47    /// Real native first-order input VJP; no base, discrete selection or higher-order gradients.
48    fn frozen_nf4_swiglu_input_backward(state:Self::Nf4SwiGluState,gradient:FloatTensor<Self>) -> Result<FloatTensor<Self>,Self::Nf4GroupedError>;
49}
ruda_tensor/grouped_nf4.rs

ruda_tensor/
grouped_nf4.rs
