logicaffeine_compile/spatial/decode_engine/
shim.rs

1//! What the decode engine runs as one unit: its answer to the shared optimizer.
2//!
3//! The engine reads a step block by block ([`super::lower_graph`]), and the merges it makes are the shared table's
4//! classes:
5//!   • the projections of one activation whose matrices are one checkpoint tensor — a Gated-DeltaNet block's query, key
6//!     and value as row ranges following one another from the tensor's first row, an attention block's queries and
7//!     output gate as the halves of each head — as one matvec (`merged_qkv`);
8//!   • a convolution and the SiLU after it, the conv unit's one pass over a Gated-DeltaNet block's channels
9//!     (`conv_bias_silu`);
10//!   • a Gated-DeltaNet block's write strength and forget gate, the sigmoid and the decay its delta rule reads, as the
11//!     gates unit (`gdn_gating`);
12//!   • an attention's output gate, its result times the sigmoid of the gate, as the sigmoid-gate unit (`glu_act_mul`).
13//! A norm, a rotation, a cache append and an add each run as a command of their own. An op the engine has no unit for is
14//! refused as having none, for what [`super::capability`] names when it refuses a whole checkpoint for it.
15
16use crate::model::bind::{self, Source, View};
17use crate::model::spec::ModelSpec;
18use crate::tensor_graph::front::TensorOpSpec;
19use crate::tensor_graph::fuse::{activation, leaves};
20use crate::tensor_graph::graph::{Access, Op, QuantFormat, Sym, TensorGraph};
21use crate::spatial::serving::{AttentionMath, VALUE_BLOCK};
22use crate::tensor_graph::numerics::{AttentionNumerics, Block, Bound, Contract, Keep, Scale, Scheme, Taken};
23use crate::tensor_graph::shim::{Shim, Support};
24use crate::tensor_graph::symbols::ClassId;
25use crate::tensor_recognize::qwen::Activation;
26
27use super::capability::Refusal;
28use super::model::{CacheFormat, EngineNumerics, StateFormat};
29use super::numerics::{ActScaling, Words, ACT_BLOCK, WIDE_BITS};
30
31/// How the shim begins every refusal of an op the engine has no unit for.
32pub const NO_UNIT: &str = "the engine has no unit for ";
33
34/// The engine's shim for one model at one knob point: its spec says which checkpoint tensor and view each constant
35/// reads, its numerics the contract its build keeps — which change the contract, never what the engine runs whole.
36pub struct EngineShim<'a> {
37    pub spec: &'a ModelSpec,
38    pub numerics: EngineNumerics,
39}
40
41impl EngineShim<'_> {
42    /// The checkpoint tensor and view `constant` reads, where the checkpoint binds it to a weight.
43    fn view(&self, constant: &str) -> Option<(String, View)> {
44        match bind::bind(self.spec, constant) {
45            Ok(Source::Weight(b)) => Some((b.tensor, b.view)),
46            _ => None,
47        }
48    }
49
50    /// Whether `parts`, products of one activation, multiply by one checkpoint matrix the engine streams as one matvec —
51    /// row ranges of one tensor following one another from its first row, or the two halves of each head of one tensor —
52    /// or by the start of one: a merge grows a product at a time.
53    fn one_matrix(&self, parts: &[Op]) -> bool {
54        let Some(views) = parts.iter().map(|p| p.ins.get(1).and_then(|w| self.view(w))).collect::<Option<Vec<_>>>() else {
55            return false;
56        };
57        if views.iter().any(|(t, _)| *t != views[0].0) {
58            return false;
59        }
60        let ranges = views.iter().try_fold(0usize, |next, (_, v)| match v {
61            View::RowRange { first, count, transpose: true } if *first == next => Some(first + count),
62            _ => None,
63        });
64        let halves = match views.as_slice() {
65            [(_, View::HeadPart { stride: s1, offset: 0, width: w1, transpose: true }), rest @ ..] if *s1 == 2 * w1 => match rest {
66                [] => true,
67                [(_, View::HeadPart { stride: s2, offset: o2, width: w2, transpose: true })] => s2 == s1 && o2 == w1 && w2 == w1,
68                _ => false,
69            },
70            _ => false,
71        };
72        ranges.is_some() || halves
73    }
74}
75
76impl Shim for EngineShim<'_> {
77    fn name(&self) -> String {
78        format!("d5005-engine:{}", self.spec.arch)
79    }
80
81    fn supports(&self, class: ClassId) -> Support {
82        use ClassId::*;
83        use Support::{Decompose, Native, Refuse};
84        match class {
85            MergedQkv | ConvBiasSilu | GdnGating | GluActMul => Native,
86            NormWeight | AddNorm | NormThenAdd | NormScale | NormWeightOffset | DualNorm | GatedNorm | EmbedNorm => Decompose,
87            GemvNormEpilogue | PrefillNormEpilogues | AttnResidualNorm | RecomputeFanout | ElementwiseChains | ElementwiseRegions => Decompose,
88            QkNormRope | QkNormRopeKvAppend | RopeKvAppend | RopeQkOneLaunch | QkNormRopeGateSplit | QkvHeadEpilogue => Decompose,
89            HorizontalNormRopeInsert | CrossLayerLaunch | ConvParts | ConvSiluL2norm | GdnGateProjections | GdnInputEpilogue => Decompose,
90            GdnPostConvPrep | GdnDecodeWhole | GemvGateUpGlu | MergedGateUp | GemmDequantInkernel | FusedSampling | LastLayerRows => Decompose,
91            NormQuant | AddNormQuant | GluActMulQuant | ElementwiseQuant | GdnOutputQuant | GemvQuant | AttnOutputQuant | GdnStep => Decompose,
92            AttnOutputGate | RowWriteQuant => Decompose,
93            GemmBias | GemmScale | AffineInputFold => Refuse("the engine's matvec unit adds no bias and scales by no factor"),
94            EmbedScale => Refuse("the engine reads its embedding unscaled"),
95            Mla | MlaPrologue | MlaUnabsorbInFold => Refuse("the engine has no unit for multi-head latent attention"),
96            RouterTopk | MoeTokenSort | MoeGateUpAct | MoeDownCombine | SharedExpertFusion | MoeGatherInOperand => {
97                Refuse("the engine has no unit for experts")
98            }
99            MoeDeferredFinalize | MoeDirectPairs | GroupedTopk | RouterSharedGateMerge | MaskedExpertActivation => {
100                Refuse("the engine has no unit for experts")
101            }
102            FlashDecode | FlashPrefill | AttnLseMerge | AttnGqaRowPacking => Refuse("the engine attends one query over its cache in its attention unit"),
103            AttentionVariantsInline | Softcap | MixedBatchAttention | SparseAttention => {
104                Refuse("the engine's attention takes no mask, window, sink or soft cap")
105            }
106            StridedRead | RowWrite | JoinInPlace | StridedValues => Refuse("the engine's units read and write whole slots"),
107            ReluSqr => Refuse("the engine has no ReLU unit"),
108            Sampling => Refuse("the engine picks the next token by arg-max"),
109            KvCacheQuant | HadamardRotation => Refuse("the engine's cache format is its numerics', not a fusion"),
110            CompensatedResidual => Refuse("the engine sums its residual uncompensated"),
111            HostResidentKv => Refuse("the engine keeps its caches in its card's memory"),
112            PersistentLayerMegakernel | CudaGraphs => Refuse("the engine runs a program of commands, not kernel launches"),
113            AllreduceNorm => Refuse("the engine runs on one card"),
114        }
115    }
116
117    fn fits(&self, op: &Op, _: &TensorGraph) -> Result<(), String> {
118        for leaf in leaves(op.clone()) {
119            if leaf.access.iter().any(|a| *a != Access::Whole) || leaf.dst.row.is_some() || leaf.dst.unit.is_some() {
120                return Err("the engine's units read and write whole slots".into());
121            }
122            let refused = match leaf.prim() {
123                Some(TensorOpSpec::Moe(_)) => Some(Refusal::Ffn { layer: 0, what: "experts" }.gap()),
124                Some(TensorOpSpec::Gelu(_)) => Some(Refusal::Ffn { layer: 0, what: "GELU" }.gap()),
125                Some(TensorOpSpec::Attention(a)) if a.window.is_some() => Some(Refusal::Attention { layer: 0, what: "a sliding window" }.gap()),
126                Some(TensorOpSpec::Attention(a)) if a.sinks.is_some() => Some(Refusal::Attention { layer: 0, what: "attention sinks" }.gap()),
127                Some(TensorOpSpec::Attention(a)) if a.softcap.is_some() => {
128                    Some(Refusal::Attention { layer: 0, what: "soft-capped scores" }.gap())
129                }
130                Some(TensorOpSpec::Softcap(_)) => Some(Refusal::LogitSoftcap.gap()),
131                Some(TensorOpSpec::HeadMatmul(_)) => Some(Refusal::Mixer { layer: 0, what: "multi-head latent attention" }.gap()),
132                Some(TensorOpSpec::Scale(_)) => Some(Refusal::EmbeddingScale.gap()),
133                Some(spec @ (TensorOpSpec::Relu(_)
134                | TensorOpSpec::IntRequant(_)
135                | TensorOpSpec::IntClamp(_)
136                | TensorOpSpec::LutAttention(_)
137                | TensorOpSpec::FullAttention(_)
138                | TensorOpSpec::Transpose(_)
139                | TensorOpSpec::HeadSlice(_)
140                | TensorOpSpec::Join(_)
141                | TensorOpSpec::RowsGather(_))) => Some(format!("a {}", spec.kind())),
142                _ => None,
143            };
144            if let Some(what) = refused {
145                return Err(format!("{NO_UNIT}{what}"));
146            }
147        }
148        let Sym::Fused(f) = &op.sym else { return Ok(()) };
149        match f.class {
150            ClassId::MergedQkv if !self.one_matrix(&f.parts) => {
151                Err("the engine streams one checkpoint matrix a matvec: row ranges of one tensor from its first row, or each head's two halves".into())
152            }
153            ClassId::GluActMul if f.parts.iter().find_map(activation) != Some(Activation::Sigmoid) => {
154                Err("the engine gates an attention's result by a sigmoid only".into())
155            }
156            _ => Ok(()),
157        }
158    }
159
160    fn activation(&self, _: &Op, _: usize, _: &TensorGraph) -> Option<QuantFormat> {
161        None
162    }
163
164    fn numerics(&self) -> Option<Contract> {
165        let n = self.numerics;
166        let (block, significand_bits) = match n.scaling {
167            ActScaling::Sub32 => (Block::Values(ACT_BLOCK), 4),
168            ActScaling::PerBlock => (Block::OfTheWeights, WIDE_BITS),
169        };
170        let words = match n.activations {
171            Words::One => 1,
172            Words::Two => 2,
173        };
174        let state = match n.state {
175            StateFormat::F32 => Keep::F32,
176            StateFormat::F16 => Keep::F16,
177        };
178        let cache = match n.cache {
179            CacheFormat::F16 => Keep::F16,
180            CacheFormat::Int8 => Keep::Int8Rows,
181            CacheFormat::Int8Hadamard => Keep::Int8RowsRotated,
182        };
183        let activations = Scheme { block, scale: Scale::LeastPowerOfTwo { significand_bits }, words };
184        // On the array an attention takes its head's query as an activation, a whole head a block where the products'
185        // blocks are the weights' (`quantize_head`), the cached rows as the cache keeps them, and its weights, each
186        // times its value row's scale, as unsigned codes under a power of two a block of positions (`weighted_int8`).
187        let attention = match n.attention {
188            AttentionMath::Vector => None,
189            AttentionMath::Array => Some(AttentionNumerics {
190                queries: Taken::Quantized(Scheme { block: if block == Block::OfTheWeights { Block::Row } else { block }, ..activations }),
191                keys_values: Taken::Kept(cache),
192                weights: Taken::Quantized(Scheme { block: Block::Values(VALUE_BLOCK as usize), scale: Scale::LeastPowerOfTwoOver255, words: 1 }),
193            }),
194        };
195        // Integer dots are exact in i64 (36 bits used); a block's or a group's sum is exact in a 64-place window over
196        // i128 and rounded to f32 once.
197        Some(Contract { point: format!("the engine: {n:?}"), activations: Some(activations), cache, state, attention, fused: Bound::EXACT, accumulator_bits: 64 })
198    }
199}
logicaffeine_compile/spatial/decode_engine/shim.rs

logicaffeine_compile/spatial/decode_engine/
shim.rs
