rlx_ir/ops/linalg.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//! Linear-algebra builders: matmul, LoRA, dequant, fused
6//! matmul+bias+activation (plan #53).
7
8use crate::op::{Activation, SynthKind};
9use crate::quant::{QuantScheme, ScaleLayout, ScaledFormat};
10use crate::{DType, Graph, NodeId, Op, Shape};
11
12impl Graph {
13 /// Matrix multiply.
14 pub fn matmul(&mut self, lhs: NodeId, rhs: NodeId, out_shape: Shape) -> NodeId {
15 self.push(Op::MatMul, vec![lhs, rhs], out_shape, None)
16 }
17
18 /// Dynamically quantize `x` (logical `[rows, cols]`, blocks along the last
19 /// axis) to low-precision `fmt` codes plus a scale tensor, per `layout`.
20 /// Returns `(codes, scale)`: `codes` is `DType::U8` with `x`'s shape;
21 /// `scale`'s shape/dtype follow the layout (`[1]` f32 for per-tensor,
22 /// `[rows, cols/block]` u8 for block layouts). The building block of
23 /// [`scaled_matmul`](Self::scaled_matmul); `fmt` may be any
24 /// [`ScaledFormat`], including a parameterized [`ScaledFormat::Custom`].
25 pub fn scaled_quantize(
26 &mut self,
27 x: NodeId,
28 fmt: ScaledFormat,
29 layout: ScaleLayout,
30 ) -> (NodeId, NodeId) {
31 let xs = self.node(x).shape.clone();
32 let cols = xs.dim(xs.rank() - 1).unwrap_static();
33 let rows = xs.num_elements().unwrap() / cols.max(1);
34 let scale_shape = match layout {
35 ScaleLayout::PerTensor => Shape::new(&[1], layout.scale_dtype()),
36 _ => Shape::new(
37 &[rows, cols.div_ceil(layout.block() as usize)],
38 layout.scale_dtype(),
39 ),
40 };
41 let scale = self.push(
42 Op::ScaledQuantScale {
43 format: fmt,
44 scale_layout: layout,
45 },
46 vec![x],
47 scale_shape,
48 None,
49 );
50 let codes = self.push(
51 Op::ScaledQuantize {
52 format: fmt,
53 scale_layout: layout,
54 },
55 vec![x, scale],
56 xs.with_dtype(DType::U8),
57 None,
58 );
59 (codes, scale)
60 }
61
62 /// Reconstruct f32 from packed `codes` + `scale` — the inverse of
63 /// [`scaled_quantize`](Self::scaled_quantize).
64 pub fn scaled_dequantize(
65 &mut self,
66 codes: NodeId,
67 scale: NodeId,
68 fmt: ScaledFormat,
69 layout: ScaleLayout,
70 ) -> NodeId {
71 let shape = self.node(codes).shape.clone().with_dtype(DType::F32);
72 self.push(
73 Op::ScaledDequantize {
74 format: fmt,
75 scale_layout: layout,
76 },
77 vec![codes, scale],
78 shape,
79 None,
80 )
81 }
82
83 /// Native low-precision GEMM (TN: `lhs [m,k] · rhs [n,k]ᵀ → [m,n]` f32).
84 /// Both operands are dynamically quantized to `fmt`/`layout` and fed
85 /// straight into the scaled matmul with f32 accumulation — no hand-wiring of
86 /// [`Op::ScaledQuantScale`]/[`Op::ScaledQuantize`]. `rhs` must already be
87 /// K-last (`[n, k]`); transpose a `[k, n]` weight first. `fmt` may be any
88 /// [`ScaledFormat`], including a parameterized [`ScaledFormat::Custom`]
89 /// (e.g. `ScaledFormat::custom(3, 0)` for `f4e3m0`).
90 pub fn scaled_matmul(
91 &mut self,
92 lhs: NodeId,
93 rhs: NodeId,
94 fmt: ScaledFormat,
95 layout: ScaleLayout,
96 ) -> NodeId {
97 self.scaled_matmul_bias(lhs, rhs, None, fmt, layout)
98 }
99
100 /// [`scaled_matmul`](Self::scaled_matmul) with an optional f32 bias `[n]`
101 /// added to each output row.
102 pub fn scaled_matmul_bias(
103 &mut self,
104 lhs: NodeId,
105 rhs: NodeId,
106 bias: Option<NodeId>,
107 fmt: ScaledFormat,
108 layout: ScaleLayout,
109 ) -> NodeId {
110 let m = self.node(lhs).shape.dim(0).unwrap_static();
111 let n = self.node(rhs).shape.dim(0).unwrap_static();
112 let (lq, ls) = self.scaled_quantize(lhs, fmt, layout);
113 let (rq, rs) = self.scaled_quantize(rhs, fmt, layout);
114 let mut inputs = vec![lq, rq, ls, rs];
115 if let Some(b) = bias {
116 inputs.push(b);
117 }
118 self.push(
119 Op::ScaledMatMul {
120 lhs_format: fmt,
121 rhs_format: fmt,
122 scale_layout: layout,
123 has_bias: bias.is_some(),
124 },
125 inputs,
126 Shape::new(&[m, n], DType::F32),
127 None,
128 )
129 }
130
131 /// Native low-precision *grouped* (MoE) GEMM — the expert-indexed
132 /// [`scaled_matmul`](Self::scaled_matmul). Dynamically quantizes f32
133 /// `input [M,K]` and the per-expert f32 weight stack `weight [E,N,K]`
134 /// (K-last) to `fmt`/`layout` codes, then wires
135 /// [`Op::ScaledGroupedMatMul`]. `expert_idx [M]` is the f32-encoded
136 /// expert id per token. Output `[M,N]` f32, with f32 accumulation.
137 pub fn scaled_grouped_matmul(
138 &mut self,
139 input: NodeId,
140 weight: NodeId,
141 expert_idx: NodeId,
142 fmt: ScaledFormat,
143 layout: ScaleLayout,
144 ) -> NodeId {
145 self.scaled_grouped_matmul_bias(input, weight, expert_idx, None, fmt, layout)
146 }
147
148 /// [`scaled_grouped_matmul`](Self::scaled_grouped_matmul) with an optional
149 /// per-expert f32 bias `[E, N]` added to each routed output row.
150 pub fn scaled_grouped_matmul_bias(
151 &mut self,
152 input: NodeId,
153 weight: NodeId,
154 expert_idx: NodeId,
155 bias: Option<NodeId>,
156 fmt: ScaledFormat,
157 layout: ScaleLayout,
158 ) -> NodeId {
159 let m = self.node(input).shape.dim(0).unwrap_static();
160 let wshape = self.node(weight).shape.clone();
161 // weight is [E, N, K]; N is the output dim, K the contraction axis.
162 let e = wshape.dim(0).unwrap_static();
163 let n = wshape.dim(wshape.rank() - 2).unwrap_static();
164 let k = wshape.dim(wshape.rank() - 1).unwrap_static();
165 // Activations [M,K] quantize with the generic (2-D) helper.
166 let (iq, is) = self.scaled_quantize(input, fmt, layout);
167 // The per-expert weight stack keeps its natural rank: the scale is
168 // `[E, N, ⌈K/block⌉]` (or `[1]` per-tensor), matching `ScaledQuantScale`
169 // shape inference so the verifier is happy. Byte layout is identical to
170 // the flattened `[E·N, ⌈K/block⌉]` the oracle indexes.
171 let w_scale_shape = match layout {
172 ScaleLayout::PerTensor => Shape::new(&[1], layout.scale_dtype()),
173 _ => Shape::new(
174 &[e, n, k.div_ceil(layout.block() as usize)],
175 layout.scale_dtype(),
176 ),
177 };
178 let ws = self.push(
179 Op::ScaledQuantScale {
180 format: fmt,
181 scale_layout: layout,
182 },
183 vec![weight],
184 w_scale_shape,
185 None,
186 );
187 let wq = self.push(
188 Op::ScaledQuantize {
189 format: fmt,
190 scale_layout: layout,
191 },
192 vec![weight, ws],
193 wshape.with_dtype(DType::U8),
194 None,
195 );
196 let mut inputs = vec![iq, wq, is, ws, expert_idx];
197 if let Some(b) = bias {
198 inputs.push(b);
199 }
200 self.push(
201 Op::ScaledGroupedMatMul {
202 lhs_format: fmt,
203 rhs_format: fmt,
204 scale_layout: layout,
205 has_bias: bias.is_some(),
206 },
207 inputs,
208 Shape::new(&[m, n], DType::F32),
209 None,
210 )
211 }
212
213 /// Dense linear solve `x = A⁻¹·b`. `A` must be `[N, N]`; `b` is
214 /// `[N]` for a single right-hand side or `[N, K]` for multiple.
215 /// `out_shape` matches `b`'s shape.
216 pub fn dense_solve(&mut self, a: NodeId, b: NodeId, out_shape: Shape) -> NodeId {
217 self.push(Op::DenseSolve, vec![a, b], out_shape, None)
218 }
219
220 /// Cholesky factorization `A = L·Lᵀ`. `a` is `[n, n]` symmetric
221 /// positive-definite; the output is the lower-triangular `L` (same shape,
222 /// strict upper triangle zeroed).
223 pub fn cholesky(&mut self, a: NodeId, out_shape: Shape) -> NodeId {
224 self.push(Op::Cholesky, vec![a], out_shape, None)
225 }
226
227 /// Triangular solve `op(A)·X = B`. `a` is `[n, n]` triangular (lower/upper
228 /// per `lower`), `op(A)` is `Aᵀ` when `transpose`; `b`/output are `[n]` or
229 /// `[n, nrhs]`.
230 pub fn triangular_solve(
231 &mut self,
232 a: NodeId,
233 b: NodeId,
234 lower: bool,
235 transpose: bool,
236 out_shape: Shape,
237 ) -> NodeId {
238 self.push(
239 Op::TriangularSolve { lower, transpose },
240 vec![a, b],
241 out_shape,
242 None,
243 )
244 }
245
246 /// Determinant of the square matrix `a` `[n, n]` → scalar (`out_shape` `[]`).
247 pub fn det(&mut self, a: NodeId, out_shape: Shape) -> NodeId {
248 self.push(Op::Det, vec![a], out_shape, None)
249 }
250
251 /// `log|det(a)|` of the square matrix `a` `[n, n]` → scalar (`out_shape` `[]`).
252 pub fn logdet(&mut self, a: NodeId, out_shape: Shape) -> NodeId {
253 self.push(Op::LogDet, vec![a], out_shape, None)
254 }
255
256 /// Sort `x` along `axis` (`descending` = largest-first). Output same shape.
257 pub fn sort(&mut self, x: NodeId, axis: usize, descending: bool, out_shape: Shape) -> NodeId {
258 self.push(Op::Sort { axis, descending }, vec![x], out_shape, None)
259 }
260
261 /// Indices that would sort `x` along `axis` (as f32). Output same shape.
262 pub fn argsort(
263 &mut self,
264 x: NodeId,
265 axis: usize,
266 descending: bool,
267 out_shape: Shape,
268 ) -> NodeId {
269 self.push(Op::ArgSort { axis, descending }, vec![x], out_shape, None)
270 }
271
272 /// One factor of the thin SVD `a = U·diag(S)·Vᵀ` (`a` `[m,n]`,
273 /// `k=min(m,n)`). `out_shape`: `U` `[m,k]`, `S` `[k]`, `Vt` `[k,n]`.
274 pub fn svd(&mut self, a: NodeId, part: crate::op::SvdPart, out_shape: Shape) -> NodeId {
275 self.push(Op::Svd { part }, vec![a], out_shape, None)
276 }
277
278 /// One factor of the thin QR `a = Q·R` (`a` `[m,n]`, `k=min(m,n)`).
279 /// `out_shape`: `Q` `[m,k]`, `R` `[k,n]`.
280 pub fn qr(&mut self, a: NodeId, part: crate::op::QrPart, out_shape: Shape) -> NodeId {
281 self.push(Op::Qr { part }, vec![a], out_shape, None)
282 }
283
284 /// Batched dense linear solve. `A` is `[B, N, N]`; `b` is
285 /// `[B, N]` (single-RHS) or `[B, N, K]` (multi-RHS). Per-batch
286 /// independent — each slice solved as a separate `dense_solve`.
287 /// Typically constructed by `vmap` of `dense_solve`.
288 pub fn batched_dense_solve(&mut self, a: NodeId, b: NodeId, out_shape: Shape) -> NodeId {
289 self.push(Op::BatchedDenseSolve, vec![a, b], out_shape, None)
290 }
291
292 /// Fused LoRA matmul: out = x·W + scale * (x·A)·B.
293 /// Inputs: x [m, k], w [k, n], a [k, r], b [r, n]. r is the
294 /// LoRA rank; scale is the alpha/rank coefficient.
295 pub fn lora_matmul(
296 &mut self,
297 x: NodeId,
298 w: NodeId,
299 a: NodeId,
300 b: NodeId,
301 scale: f32,
302 shape: Shape,
303 ) -> NodeId {
304 self.push(Op::LoraMatMul { scale }, vec![x, w, a, b], shape, None)
305 }
306
307 /// Fused dequant + matmul. See [`Op::DequantMatMul`] for per-scheme
308 /// input layout (4 inputs for legacy/NVFP4, 2 for GGUF).
309 pub fn dequant_matmul(
310 &mut self,
311 x: NodeId,
312 w_q: NodeId,
313 scale: NodeId,
314 zp: NodeId,
315 scheme: QuantScheme,
316 shape: Shape,
317 ) -> NodeId {
318 self.push(
319 Op::DequantMatMul { scheme },
320 vec![x, w_q, scale, zp],
321 shape,
322 None,
323 )
324 }
325
326 /// GGUF / K-quant packed weights — `[x, packed_w_bytes]` only.
327 pub fn dequant_matmul_packed(
328 &mut self,
329 x: NodeId,
330 packed_w: NodeId,
331 scheme: QuantScheme,
332 shape: Shape,
333 ) -> NodeId {
334 debug_assert!(
335 scheme.is_gguf(),
336 "dequant_matmul_packed requires a GGUF QuantScheme"
337 );
338 self.push(Op::DequantMatMul { scheme }, vec![x, packed_w], shape, None)
339 }
340
341 /// GGUF / K-quant packed **expert bank** — `[x, packed_bank_bytes,
342 /// expert_idx]`, the grouped analogue of [`Self::dequant_matmul_packed`].
343 ///
344 /// `packed_bank` is one U8 tensor of `num_experts` contiguous slabs, each
345 /// the `[out_dim, in_dim]` blob of one expert — which is GGUF's native
346 /// order for `ffn_*_exps.weight` (GGML `ne = [in, out, experts]`), so no
347 /// transpose is needed. The op contracts along `in_dim` per expert, exactly
348 /// as [`Op::DequantMatMul`] does for a single linear, and the expert count
349 /// is recovered from the blob size.
350 ///
351 /// This is the packed counterpart of `Self::grouped_matmul`: an MoE layer
352 /// whose banks stay quantized never materializes them as F32, which for a
353 /// fine-grained MoE is the difference between a model fitting and not.
354 pub fn dequant_grouped_matmul_packed(
355 &mut self,
356 x: NodeId,
357 packed_bank: NodeId,
358 expert_idx: NodeId,
359 scheme: QuantScheme,
360 shape: Shape,
361 ) -> NodeId {
362 debug_assert!(
363 scheme.is_gguf(),
364 "dequant_grouped_matmul_packed requires a GGUF QuantScheme"
365 );
366 self.push(
367 Op::DequantGroupedMatMul { scheme },
368 vec![x, packed_bank, expert_idx],
369 shape,
370 None,
371 )
372 }
373
374 /// NVFP4 (E2M1) block matmul — group size 16, FP8 block scales,
375 /// optional f32 global scale (defaults to 1.0 when unset at runtime).
376 pub fn dequant_matmul_nvfp4(
377 &mut self,
378 x: NodeId,
379 w_q: NodeId,
380 block_scales: NodeId,
381 global_scale: NodeId,
382 shape: Shape,
383 ) -> NodeId {
384 self.dequant_matmul(
385 x,
386 w_q,
387 block_scales,
388 global_scale,
389 QuantScheme::Nvfp4Block,
390 shape,
391 )
392 }
393
394 /// On-chip codebook weight-synthesis matmul. See [`Op::SynthMatMul`].
395 /// `indices` is `[n, k/entry_dim]` u8; `codebook` is
396 /// `[num_entries, entry_dim]` f32; output is `[m, n]`. The weight is
397 /// reconstructed inside the matmul inner loop, never materialized.
398 pub fn synth_matmul(
399 &mut self,
400 x: NodeId,
401 indices: NodeId,
402 codebook: NodeId,
403 kind: SynthKind,
404 shape: Shape,
405 ) -> NodeId {
406 self.push(
407 Op::SynthMatMul { kind },
408 vec![x, indices, codebook],
409 shape,
410 None,
411 )
412 }
413
414 /// [`synth_matmul`](Self::synth_matmul) with a LOW-PRECISION codebook (fp8 /
415 /// fp4 / nvf4 / custom `fpXmYeZ`). The centroids are stored as `fmt`/`layout`
416 /// codes (`codebook_codes` U8 + `codebook_scale` from
417 /// [`scaled_quantize`](Self::scaled_quantize)) and decoded to f32 via
418 /// [`Op::ScaledDequantize`] before the synth matmul — so ANY [`ScaledFormat`],
419 /// including a parameterized [`ScaledFormat::Custom`] (e.g. `custom(3,0)` =
420 /// `f4e3m0`) or an [`ScaleLayout::Nvfp4`] layout, works on every backend with
421 /// no new kernel. The codebook is tiny, so the decode is negligible.
422 #[allow(clippy::too_many_arguments)]
423 pub fn synth_matmul_qcodebook(
424 &mut self,
425 x: NodeId,
426 indices: NodeId,
427 codebook_codes: NodeId,
428 codebook_scale: NodeId,
429 kind: SynthKind,
430 fmt: ScaledFormat,
431 layout: ScaleLayout,
432 shape: Shape,
433 ) -> NodeId {
434 let cb_f32 = self.scaled_dequantize(codebook_codes, codebook_scale, fmt, layout);
435 self.synth_matmul(x, indices, cb_f32, kind, shape)
436 }
437
438 /// KAN learnable spline activation. See [`Op::SplineActivation`]. `x` is
439 /// `[.., C]`, `coeff` is `[C, num_basis]`; the output matches `x`. Each
440 /// channel's univariate function is a Gaussian-RBF expansion with learned
441 /// per-channel coefficients.
442 pub fn spline_activation(
443 &mut self,
444 x: NodeId,
445 coeff: NodeId,
446 num_basis: u32,
447 grid_min: f32,
448 grid_max: f32,
449 ) -> NodeId {
450 let shape = self.node(x).shape.clone();
451 self.push(
452 Op::SplineActivation {
453 num_basis,
454 grid_min,
455 grid_max,
456 },
457 vec![x, coeff],
458 shape,
459 None,
460 )
461 }
462
463 /// Fused matmul + bias + activation (created by optimization passes).
464 pub fn fused_matmul_bias_act(
465 &mut self,
466 input: NodeId,
467 weight: NodeId,
468 bias: NodeId,
469 activation: Option<Activation>,
470 shape: Shape,
471 ) -> NodeId {
472 self.push(
473 Op::FusedMatMulBiasAct { activation },
474 vec![input, weight, bias],
475 shape,
476 None,
477 )
478 }
479
480 /// Real INT8-arithmetic matmul: i8 inputs, i32 bias, i8 output.
481 /// `mult = x_scale · w_scale / out_scale`. Caller's responsible
482 /// for asserting the input dtypes — the builder just plumbs the
483 /// shape with `dtype = I8` since that's what the kernel writes.
484 pub fn q_matmul(
485 &mut self,
486 x: NodeId,
487 w: NodeId,
488 bias: NodeId,
489 x_zp: i32,
490 w_zp: i32,
491 out_zp: i32,
492 mult: f32,
493 out_shape: Shape,
494 ) -> NodeId {
495 debug_assert_eq!(
496 out_shape.dtype(),
497 crate::DType::I8,
498 "q_matmul output dtype must be I8"
499 );
500 self.push(
501 Op::QMatMul {
502 x_zp,
503 w_zp,
504 out_zp,
505 mult,
506 },
507 vec![x, w, bias],
508 out_shape,
509 None,
510 )
511 }
512
513 /// Real INT8-arithmetic 2-D convolution. NCHW layout matching
514 /// `Op::Conv`. `mult = x_scale · w_scale / out_scale`.
515 #[allow(clippy::too_many_arguments)]
516 pub fn q_conv2d(
517 &mut self,
518 x: NodeId,
519 w: NodeId,
520 bias: NodeId,
521 kernel_size: Vec<usize>,
522 stride: Vec<usize>,
523 padding: Vec<usize>,
524 dilation: Vec<usize>,
525 groups: usize,
526 x_zp: i32,
527 w_zp: i32,
528 out_zp: i32,
529 mult: f32,
530 out_shape: Shape,
531 ) -> NodeId {
532 debug_assert_eq!(
533 out_shape.dtype(),
534 crate::DType::I8,
535 "q_conv2d output dtype must be I8"
536 );
537 self.push(
538 Op::QConv2d {
539 kernel_size,
540 stride,
541 padding,
542 dilation,
543 groups,
544 x_zp,
545 w_zp,
546 out_zp,
547 mult,
548 },
549 vec![x, w, bias],
550 out_shape,
551 None,
552 )
553 }
554}