rlx_compile/
scaled_quant_insert.rs

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

rlx_compile/
scaled_quant_insert.rs
