rlx_compile/
scaled_quant_insert.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Native low-precision GEMM insertion (AMP-FP8/FP4).
17//!
18//! Rewrites each 2-D compute `Op::MatMul` into a native `Op::ScaledMatMul`:
19//! both operands are dynamically quantized (`ScaledQuantScale` +
20//! `ScaledQuantize`) and fed straight into the tensor-core GEMM with f32
21//! accumulation — the path that wins on Hopper/Ada/Blackwell and CDNA3/CDNA4.
22//!
23//! `ScaledMatMul` is TN (`out = lhs · rhsᵀ`), so the rhs is transposed to
24//! K-last; on a `Param` weight that `Transpose` const-folds at compile time, so
25//! only activations pay a runtime quantize. Run this **before** `ConstantFolding`.
26//!
27//! Unlike [`crate::precision::AutoMixedPrecision`] (a dtype-relabel + `Cast`
28//! engine), this is op-replacement + scale plumbing — a different transform, so
29//! it lives in its own pass. It changes numerics, so it is **opt-in**: the
30//! caller enables it per graph; nothing turns it on by default.
31
32use rlx_ir::*;
33use std::collections::HashMap;
34
35/// Which element formats + scale layout to quantize matmul operands to.
36#[derive(Debug, Clone, Copy)]
37pub struct ScaledQuantConfig {
38    pub lhs_format: ScaledFormat,
39    pub rhs_format: ScaledFormat,
40    pub scale_layout: ScaleLayout,
41}
42
43impl ScaledQuantConfig {
44    /// Per-tensor FP8 E4M3 for both operands (Hopper / Ada default).
45    pub fn fp8_e4m3() -> Self {
46        Self {
47            lhs_format: ScaledFormat::F8E4M3,
48            rhs_format: ScaledFormat::F8E4M3,
49            scale_layout: ScaleLayout::PerTensor,
50        }
51    }
52    /// OCP microscaling MXFP8 (E4M3 elements, E8M0 block scales).
53    pub fn mxfp8_e4m3() -> Self {
54        Self {
55            lhs_format: ScaledFormat::F8E4M3,
56            rhs_format: ScaledFormat::F8E4M3,
57            scale_layout: ScaleLayout::mx(),
58        }
59    }
60}
61
62/// Rewrite every 2-D `Op::MatMul` in `graph` into a native `Op::ScaledMatMul`.
63/// Batched (rank > 2) matmuls and all other ops are copied unchanged.
64pub fn insert_scaled_matmul(graph: Graph, cfg: ScaledQuantConfig) -> Graph {
65    let mut out = Graph::new(&graph.name);
66    let mut id_map: HashMap<NodeId, NodeId> = HashMap::new();
67    // Quantize-once: a tensor feeding two matmuls is quantized a single time.
68    let mut qcache: HashMap<(NodeId, ScaledFormat), (NodeId, NodeId)> = HashMap::new();
69    // Transpose-once for shared rhs weights.
70    let mut tcache: HashMap<NodeId, NodeId> = HashMap::new();
71
72    for node in graph.nodes() {
73        let new_inputs: Vec<NodeId> = node.inputs.iter().map(|i| id_map[i]).collect();
74
75        let new_id = match &node.op {
76            Op::MatMul if node.shape.rank() == 2 => {
77                let lhs = new_inputs[0];
78                let rhs = new_inputs[1];
79                let lhs_shape = out.node(lhs).shape.clone();
80                let rhs_shape = out.node(rhs).shape.clone();
81                let m = lhs_shape.dim(0).unwrap_static();
82                let k = lhs_shape.dim(1).unwrap_static();
83                let n = rhs_shape.dim(1).unwrap_static();
84
85                // rhs [k,n] → [n,k] so ScaledMatMul's TN layout sees K-last.
86                let rhs_t = *tcache.entry(rhs).or_insert_with(|| {
87                    out.add_node(
88                        Op::Transpose { perm: vec![1, 0] },
89                        vec![rhs],
90                        Shape::new(&[n, k], rhs_shape.dtype()),
91                    )
92                });
93
94                let (lq, ls) = quantize_operand(
95                    &mut out,
96                    &mut qcache,
97                    lhs,
98                    m,
99                    k,
100                    cfg.lhs_format,
101                    cfg.scale_layout,
102                );
103                let (rq, rs) = quantize_operand(
104                    &mut out,
105                    &mut qcache,
106                    rhs_t,
107                    n,
108                    k,
109                    cfg.rhs_format,
110                    cfg.scale_layout,
111                );
112
113                out.add_node(
114                    Op::ScaledMatMul {
115                        lhs_format: cfg.lhs_format,
116                        rhs_format: cfg.rhs_format,
117                        scale_layout: cfg.scale_layout,
118                        has_bias: false,
119                    },
120                    vec![lq, rq, ls, rs],
121                    Shape::new(&[m, n], DType::F32),
122                )
123            }
124            _ => out.add_node(node.op.clone(), new_inputs, node.shape.clone()),
125        };
126        id_map.insert(node.id, new_id);
127    }
128
129    let new_outputs: Vec<NodeId> = graph.outputs.iter().map(|&id| id_map[&id]).collect();
130    out.set_outputs(new_outputs);
131    out
132}
133
134/// Insert (cached) `ScaledQuantScale` + `ScaledQuantize` for one operand of
135/// logical shape `[rows, cols]` (blocks along cols = contraction). Returns
136/// `(codes_id, scale_id)`.
137fn quantize_operand(
138    out: &mut Graph,
139    qcache: &mut HashMap<(NodeId, ScaledFormat), (NodeId, NodeId)>,
140    operand: NodeId,
141    rows: usize,
142    cols: usize,
143    fmt: ScaledFormat,
144    layout: ScaleLayout,
145) -> (NodeId, NodeId) {
146    if let Some(&hit) = qcache.get(&(operand, fmt)) {
147        return hit;
148    }
149    let scale_shape = match layout {
150        ScaleLayout::PerTensor => Shape::new(&[1], layout.scale_dtype()),
151        _ => Shape::new(
152            &[rows, cols.div_ceil(layout.block() as usize)],
153            layout.scale_dtype(),
154        ),
155    };
156    let scale = out.add_node(
157        Op::ScaledQuantScale {
158            format: fmt,
159            scale_layout: layout,
160        },
161        vec![operand],
162        scale_shape,
163    );
164    let codes = out.add_node(
165        Op::ScaledQuantize {
166            format: fmt,
167            scale_layout: layout,
168        },
169        vec![operand, scale],
170        Shape::new(&[rows, cols], DType::U8),
171    );
172    qcache.insert((operand, fmt), (codes, scale));
173    (codes, scale)
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179
180    fn count_kind(g: &Graph, k: OpKind) -> usize {
181        g.nodes().iter().filter(|n| n.op.kind() == k).count()
182    }
183
184    #[test]
185    fn rewrites_matmul_to_scaled() {
186        let mut g = Graph::new("mm");
187        let x = g.input("x", Shape::new(&[2, 4], DType::F32));
188        let w = g.param("w", Shape::new(&[4, 3], DType::F32));
189        let mm = g.matmul(x, w, Shape::new(&[2, 3], DType::F32));
190        g.set_outputs(vec![mm]);
191
192        let out = insert_scaled_matmul(g, ScaledQuantConfig::fp8_e4m3());
193
194        assert_eq!(count_kind(&out, OpKind::MatMul), 0, "plain matmul remains");
195        assert_eq!(count_kind(&out, OpKind::ScaledMatMul), 1);
196        // one ScaledQuantize + one ScaledQuantScale per operand (2 operands)
197        assert_eq!(count_kind(&out, OpKind::ScaledQuantize), 2);
198        assert_eq!(count_kind(&out, OpKind::ScaledQuantScale), 2);
199        // rhs transposed to K-last
200        assert_eq!(count_kind(&out, OpKind::Transpose), 1);
201
202        // Output is the ScaledMatMul, shape [2,3] f32.
203        let out_node = out.node(out.outputs[0]);
204        assert!(matches!(out_node.op, Op::ScaledMatMul { .. }));
205        assert_eq!(out_node.shape.dim(0).unwrap_static(), 2);
206        assert_eq!(out_node.shape.dim(1).unwrap_static(), 3);
207        assert_eq!(out_node.shape.dtype(), DType::F32);
208    }
209
210    #[test]
211    fn quantize_once_for_shared_activation() {
212        // x feeds two matmuls — its quantization must be shared (cached).
213        let mut g = Graph::new("shared");
214        let x = g.input("x", Shape::new(&[2, 4], DType::F32));
215        let w1 = g.param("w1", Shape::new(&[4, 3], DType::F32));
216        let w2 = g.param("w2", Shape::new(&[4, 5], DType::F32));
217        let mm1 = g.matmul(x, w1, Shape::new(&[2, 3], DType::F32));
218        let mm2 = g.matmul(x, w2, Shape::new(&[2, 5], DType::F32));
219        g.set_outputs(vec![mm1, mm2]);
220
221        let out = insert_scaled_matmul(g, ScaledQuantConfig::fp8_e4m3());
222
223        assert_eq!(count_kind(&out, OpKind::ScaledMatMul), 2);
224        // 2 weights + 1 shared activation = 3 quantizations, not 4.
225        assert_eq!(count_kind(&out, OpKind::ScaledQuantize), 3);
226        assert_eq!(count_kind(&out, OpKind::ScaledQuantScale), 3);
227    }
228
229    #[test]
230    fn coexists_with_f16_amp() {
231        // fp8-ify the matmul, then run the f16 AutoMixedPrecision pass on top.
232        // The f16 pass must leave the fp8 subgraph intact (no U8→f16 casts, no
233        // dtype relabel of ScaledMatMul) while still lowering the surrounding
234        // elementwise op.
235        use crate::precision::{AutoMixedPrecision, PrecisionPolicy};
236        use rlx_fusion::pass::Pass;
237
238        let mut g = Graph::new("mix");
239        let x = g.input("x", Shape::new(&[2, 4], DType::F32));
240        let w = g.param("w", Shape::new(&[4, 3], DType::F32));
241        let mm = g.matmul(x, w, Shape::new(&[2, 3], DType::F32));
242        let relu = g.activation(
243            rlx_ir::op::Activation::Relu,
244            mm,
245            Shape::new(&[2, 3], DType::F32),
246        );
247        g.set_outputs(vec![relu]);
248
249        let g = insert_scaled_matmul(g, ScaledQuantConfig::fp8_e4m3());
250        let g = AutoMixedPrecision::new(PrecisionPolicy::AutoMixed).run(g);
251
252        let sm = g
253            .nodes()
254            .iter()
255            .find(|n| matches!(n.op, Op::ScaledMatMul { .. }))
256            .expect("ScaledMatMul survived the f16 AMP pass");
257        assert_eq!(
258            sm.shape.dtype(),
259            DType::F32,
260            "ScaledMatMul output stayed f32"
261        );
262        for &i in &sm.inputs {
263            let d = g.node(i).shape.dtype();
264            assert!(
265                d == DType::U8 || d == DType::F32,
266                "ScaledMatMul input wrongly relabeled to {d}"
267            );
268        }
269        // The graph must still pass shape verification end-to-end.
270        let errs = rlx_ir::verify::verify(&g);
271        assert!(errs.is_empty(), "post-AMP graph verifies: {errs:?}");
272    }
273
274    #[test]
275    fn leaves_batched_matmul_alone() {
276        let mut g = Graph::new("batched");
277        let a = g.input("a", Shape::new(&[8, 2, 4], DType::F32));
278        let b = g.param("b", Shape::new(&[8, 4, 3], DType::F32));
279        let mm = g.matmul(a, b, Shape::new(&[8, 2, 3], DType::F32));
280        g.set_outputs(vec![mm]);
281
282        let out = insert_scaled_matmul(g, ScaledQuantConfig::fp8_e4m3());
283        assert_eq!(
284            count_kind(&out, OpKind::MatMul),
285            1,
286            "batched matmul untouched"
287        );
288        assert_eq!(count_kind(&out, OpKind::ScaledMatMul), 0);
289    }
290}
rlx_compile/scaled_quant_insert.rs

rlx_compile/
scaled_quant_insert.rs
