teeny_kernels/nn/optim/
muon.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#![allow(non_snake_case)]
18
19use teeny_macros::kernel;
20use teeny_triton::triton::{
21    types::{AddOffsets, Comparison},
22    *,
23};
24
25// ── muon_frob_norm_sq ─────────────────────────────────────────────────────────
26
27/// Parallel Frobenius squared-norm via atomic reduction.
28///
29/// Each block sums `x[block]²` and atomically adds to `out_ptr[0]`.
30/// Caller must **zero** `out_ptr[0]` before launch.
31/// After launch: `out_ptr[0] = ||X||_F²`.
32///
33/// Grid: `[ceil(n_elements / BLOCK_SIZE), 1, 1]`.
34#[kernel]
35pub fn muon_frob_norm_sq<T: Triton, const BLOCK_SIZE: i32>(
36    x_ptr: T::Pointer<f32>,
37    out_ptr: T::Pointer<f32>,
38    n_elements: i32,
39) where
40    T::I32Tensor: types::Tensor<i32, 1>,
41    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
42    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
43{
44    let pid = T::program_id(Axis::X);
45    let offsets = T::arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE;
46    let mask = offsets.lt(n_elements);
47
48    let x = T::load(
49        x_ptr.add_offsets(offsets),
50        Some(mask),
51        Some(T::zeros::<f32>(&[BLOCK_SIZE])),
52        &[],
53        None,
54        None,
55        None,
56        false,
57    );
58    // Reduce x² within this block → scalar, then expand to tensor<1xf32> for atomic_add
59    let partial = T::expand_dims(T::sum(x * x, Some(0), false), 0);
60
61    // Atomic-add the block's partial sum to the single output accumulator
62    let out_off: T::I32Tensor = T::arange(0, 1);
63    let _ = T::atomic_add(out_ptr.add_offsets(out_off), partial, None, None, None);
64}
65
66// ── muon_ns_xtx ───────────────────────────────────────────────────────────────
67
68/// Gram matrix: `T = X·Xᵀ` (!TRANSPOSE) or `T = Xᵀ·X` (TRANSPOSE).
69///
70/// Both cases are expressed as `A @ Aᵀ` with different views of X:
71/// - `!TRANSPOSE` → A = X `[M, N]`,   output T `[M, M]`, contraction K = N.
72/// - `TRANSPOSE`  → A = Xᵀ `[N, M]`, output T `[N, N]`, contraction K = M.
73///
74/// `stride_xm` is the **row stride of X** (typically `= N` for row-major `[M, N]`).
75///
76/// Grid: `[ceil(R/BLOCK_R)² grouped by GROUP_R, 1, 1]` where R = M or N.
77#[kernel]
78pub fn muon_ns_xtx<
79    T: Triton,
80    const TRANSPOSE: bool,
81    const BLOCK_R: i32,
82    const BLOCK_K: i32,
83    const GROUP_R: i32,
84>(
85    x_ptr: T::Pointer<f32>,
86    t_ptr: T::Pointer<f32>,
87    M: i32,
88    N: i32,
89    stride_xm: i32,
90) {
91    // R: dimension of the square output; K: contraction dim.
92    // For !TRANSPOSE: A = X[M, N],   stride(N, 1),   R=M, K=N.
93    // For  TRANSPOSE: A = Xᵀ[N, M], stride(1, N),   R=N, K=M.
94    let R = if TRANSPOSE { N } else { M };
95    let K = if TRANSPOSE { M } else { N };
96    let a_stride_row = if TRANSPOSE { 1 } else { stride_xm };
97    let a_stride_col = if TRANSPOSE { stride_xm } else { 1 };
98
99    let pid = T::program_id(Axis::X);
100    let num_pid_r = T::cdiv(R, BLOCK_R);
101    let num_pid_in_group = GROUP_R * num_pid_r;
102    let group_id = pid / num_pid_in_group;
103    let first_pid_r = group_id * GROUP_R;
104    let remaining = num_pid_r - first_pid_r;
105    let group_size = if remaining < GROUP_R {
106        remaining
107    } else {
108        GROUP_R
109    };
110    let pid_in_group = pid % num_pid_in_group;
111    let pid_rm = first_pid_r + (pid_in_group % group_size);
112    let pid_rn = pid_in_group / group_size;
113
114    // A and B both view the same matrix (symmetric Gram product)
115    let a_desc = T::make_tensor_descriptor(
116        x_ptr,
117        &[R, K],
118        &[a_stride_row, a_stride_col],
119        &[BLOCK_R, BLOCK_K],
120        Some(PaddingOption::Zero),
121    );
122    let b_desc = T::make_tensor_descriptor(
123        x_ptr,
124        &[R, K],
125        &[a_stride_row, a_stride_col],
126        &[BLOCK_R, BLOCK_K],
127        Some(PaddingOption::Zero),
128    );
129
130    let mut acc = T::zeros::<f32>(&[BLOCK_R, BLOCK_R]);
131    let k_tiles = T::cdiv(K, BLOCK_K);
132    for k in 0..k_tiles {
133        let a = T::load_tensor_descriptor(a_desc, &[pid_rm * BLOCK_R, k * BLOCK_K]);
134        let b = T::load_tensor_descriptor(b_desc, &[pid_rn * BLOCK_R, k * BLOCK_K]);
135        let b_t = T::trans(b, &[1, 0]);
136        // IEEE, not TF32: this is Newton-Schulz orthogonalization, an iterative
137        // process where TF32's reduced mantissa would compound error across
138        // iterations. Explicit rather than relying on a default.
139        acc = T::dot::<f32, f32>(a, b_t, Some(acc), InputPrecision::IEEE, None);
140    }
141
142    // Output T: [R × R], row-major
143    let t_desc = T::make_tensor_descriptor(
144        t_ptr,
145        &[R, R],
146        &[R, 1],
147        &[BLOCK_R, BLOCK_R],
148        Some(PaddingOption::Zero),
149    );
150    T::store_tensor_descriptor(t_desc, &[pid_rm * BLOCK_R, pid_rn * BLOCK_R], acc);
151}
152
153// ── muon_ns_step ──────────────────────────────────────────────────────────────
154
155/// One Newton-Schulz step (in-place): `X ← a·X + b·(T·X)` or `X ← a·X + b·(X·T)`.
156///
157/// Both cases are expressed as `A @ B` (no explicit transpose on B) with unit
158/// inner strides so TMA descriptors remain valid:
159///
160/// - `!TRANSPOSE` (`T·X`): A = T `[M, K]`, B = X `[K, N]`, K = M.
161/// - `TRANSPOSE`  (`X·T`): A = X `[M, K]`, B = T `[K, N]`, K = N.
162///
163/// The GEMM result is fused with the elementwise update in one store.
164///
165/// - `stride_tm`: row stride of T (= M for !TRANSPOSE, = N for TRANSPOSE).
166/// - `stride_xm = N` (row-major X, always `[M, N]`).
167///
168/// Grid: `[ceil(M/BLOCK_M) × ceil(N/BLOCK_N) grouped by GROUP_M, 1, 1]`.
169#[kernel]
170pub fn muon_ns_step<
171    T: Triton,
172    const TRANSPOSE: bool,
173    const BLOCK_M: i32,
174    const BLOCK_N: i32,
175    const BLOCK_K: i32,
176    const GROUP_M: i32,
177>(
178    t_ptr: T::Pointer<f32>,
179    x_ptr: T::Pointer<f32>,
180    M: i32,
181    N: i32,
182    stride_tm: i32,
183    stride_xm: i32,
184    a: f32,
185    b: f32,
186) {
187    // Inner (contraction) dimension
188    let K = if TRANSPOSE { N } else { M };
189
190    let pid = T::program_id(Axis::X);
191    let num_pid_m = T::cdiv(M, BLOCK_M);
192    let num_pid_n = T::cdiv(N, BLOCK_N);
193    let num_pid_in_group = GROUP_M * num_pid_n;
194    let group_id = pid / num_pid_in_group;
195    let first_pid_m = group_id * GROUP_M;
196    let remaining = num_pid_m - first_pid_m;
197    let group_size = if remaining < GROUP_M {
198        remaining
199    } else {
200        GROUP_M
201    };
202    let pid_in_group = pid % num_pid_in_group;
203    let pid_m = first_pid_m + (pid_in_group % group_size);
204    let pid_n = pid_in_group / group_size;
205
206    // Descriptors for the GEMM (A @ B, inner dim K, both with unit inner stride):
207    //   !TRANSPOSE: A = T[M, K] strides(stride_tm,1),  B = X[K, N] strides(stride_xm,1), K=M
208    //    TRANSPOSE: A = X[M, K] strides(stride_xm,1),  B = T[K, N] strides(stride_tm,1), K=N
209    let (a_desc, b_desc) = if TRANSPOSE {
210        let ad = T::make_tensor_descriptor(
211            x_ptr,
212            &[M, K],
213            &[stride_xm, 1],
214            &[BLOCK_M, BLOCK_K],
215            Some(PaddingOption::Zero),
216        );
217        let bd = T::make_tensor_descriptor(
218            t_ptr,
219            &[K, N],
220            &[stride_tm, 1],
221            &[BLOCK_K, BLOCK_N],
222            Some(PaddingOption::Zero),
223        );
224        (ad, bd)
225    } else {
226        let ad = T::make_tensor_descriptor(
227            t_ptr,
228            &[M, K],
229            &[stride_tm, 1],
230            &[BLOCK_M, BLOCK_K],
231            Some(PaddingOption::Zero),
232        );
233        let bd = T::make_tensor_descriptor(
234            x_ptr,
235            &[K, N],
236            &[stride_xm, 1],
237            &[BLOCK_K, BLOCK_N],
238            Some(PaddingOption::Zero),
239        );
240        (ad, bd)
241    };
242
243    let mut acc = T::zeros::<f32>(&[BLOCK_M, BLOCK_N]);
244    let k_tiles = T::cdiv(K, BLOCK_K);
245    for k in 0..k_tiles {
246        let av = T::load_tensor_descriptor(a_desc, &[pid_m * BLOCK_M, k * BLOCK_K]);
247        let bv = T::load_tensor_descriptor(b_desc, &[k * BLOCK_K, pid_n * BLOCK_N]);
248        // IEEE, not TF32: see the note on the other dot in this file — Newton-Schulz
249        // iteration accuracy depends on it.
250        acc = T::dot::<f32, f32>(av, bv, Some(acc), InputPrecision::IEEE, None);
251    }
252
253    // Fused elementwise: X_new = a * X_tile + b * GEMM_result
254    let x_desc = T::make_tensor_descriptor(
255        x_ptr,
256        &[M, N],
257        &[stride_xm, 1],
258        &[BLOCK_M, BLOCK_N],
259        Some(PaddingOption::Zero),
260    );
261    let x_tile = T::load_tensor_descriptor(x_desc, &[pid_m * BLOCK_M, pid_n * BLOCK_N]);
262
263    let a_t = T::full::<f32>(&[BLOCK_M, BLOCK_N], a);
264    let b_t = T::full::<f32>(&[BLOCK_M, BLOCK_N], b);
265    let result = a_t * x_tile + b_t * acc;
266    T::store_tensor_descriptor(x_desc, &[pid_m * BLOCK_M, pid_n * BLOCK_N], result);
267}
268
269// ── muon_update ───────────────────────────────────────────────────────────────
270
271/// Parameter update: `W -= lr · G_orth`.
272///
273/// Grid: `[ceil(n_elements / BLOCK_SIZE), 1, 1]`.
274#[kernel]
275pub fn muon_update<T: Triton, const BLOCK_SIZE: i32>(
276    params_ptr: T::Pointer<f32>,
277    grad_ptr: T::Pointer<f32>,
278    n_elements: i32,
279    lr: f32,
280) where
281    T::I32Tensor: types::Tensor<i32, 1>,
282    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
283    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
284{
285    let pid = T::program_id(Axis::X);
286    let offsets = T::arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE;
287    let mask = offsets.lt(n_elements);
288
289    let p = T::load(
290        params_ptr.add_offsets(offsets),
291        Some(mask),
292        None,
293        &[],
294        None,
295        None,
296        None,
297        false,
298    );
299    let g = T::load(
300        grad_ptr.add_offsets(offsets),
301        Some(mask),
302        None,
303        &[],
304        None,
305        None,
306        None,
307        false,
308    );
309    let lr_t = T::full::<f32>(&[BLOCK_SIZE], lr);
310
311    T::store(
312        params_ptr.add_offsets(offsets),
313        p - lr_t * g,
314        Some(mask),
315        &[],
316        None,
317        None,
318    );
319}
teeny_kernels/nn/optim/muon.rs

teeny_kernels/nn/optim/
muon.rs
