teeny_kernels/nn/optim/
nadam.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/// NAdam step (Nesterov-accelerated Adam).
26///
27/// All bias correction terms and Nesterov coefficients are precomputed on the
28/// host and passed as scalars:
29///
30/// ```text
31/// exp_avg    = beta1 * exp_avg    + (1 - beta1) * g
32/// exp_avg_sq = beta2 * exp_avg_sq + (1 - beta2) * g²
33/// denom      = sqrt(exp_avg_sq) / bias_corr2_sqrt + eps
34/// p         -= lr * (coeff_g * g + coeff_m * exp_avg) / denom
35/// ```
36///
37/// Precomputed on host:
38/// - `bias_corr2_sqrt = sqrt(1 - beta2^t)`
39/// - `coeff_g  = (1 - mu_t)  / (1 - mu_product)`  (grad contribution)
40/// - `coeff_m  = mu_t1       / (1 - mu_product_next)` (moment contribution)
41///
42/// Grid: `[ceil(n_elements / BLOCK_SIZE), 1, 1]`.
43#[kernel]
44pub fn nadam_step<T: Triton, const BLOCK_SIZE: i32>(
45    params_ptr: T::Pointer<f32>,
46    grad_ptr: T::Pointer<f32>,
47    exp_avg_ptr: T::Pointer<f32>,
48    exp_avg_sq_ptr: T::Pointer<f32>,
49    n_elements: i32,
50    lr: f32,
51    beta1: f32,
52    beta2: f32,
53    eps: f32,
54    weight_decay: f32,
55    bias_corr2_sqrt: f32,
56    coeff_g: f32,
57    coeff_m: f32,
58) where
59    T::I32Tensor: types::Tensor<i32, 1>,
60    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
61    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
62{
63    let pid = T::program_id(Axis::X);
64    let block_start = pid * BLOCK_SIZE;
65    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
66    let mask = offsets.lt(n_elements);
67
68    let p = T::load(
69        params_ptr.add_offsets(offsets),
70        Some(mask),
71        None,
72        &[],
73        None,
74        None,
75        None,
76        false,
77    );
78    let g = T::load(
79        grad_ptr.add_offsets(offsets),
80        Some(mask),
81        None,
82        &[],
83        None,
84        None,
85        None,
86        false,
87    );
88    let exp_avg = T::load(
89        exp_avg_ptr.add_offsets(offsets),
90        Some(mask),
91        None,
92        &[],
93        None,
94        None,
95        None,
96        false,
97    );
98    let exp_avg_sq = T::load(
99        exp_avg_sq_ptr.add_offsets(offsets),
100        Some(mask),
101        None,
102        &[],
103        None,
104        None,
105        None,
106        false,
107    );
108
109    let beta1_t = T::full(&[BLOCK_SIZE], beta1);
110    let beta2_t = T::full(&[BLOCK_SIZE], beta2);
111    let one_m_beta1 = T::full(&[BLOCK_SIZE], 1.0_f32 - beta1);
112    let one_m_beta2 = T::full(&[BLOCK_SIZE], 1.0_f32 - beta2);
113    let eps_t = T::full(&[BLOCK_SIZE], eps);
114    let lr_t = T::full(&[BLOCK_SIZE], lr);
115    let bc2sqrt_t = T::full(&[BLOCK_SIZE], bias_corr2_sqrt);
116    let cg_t = T::full(&[BLOCK_SIZE], coeff_g);
117    let cm_t = T::full(&[BLOCK_SIZE], coeff_m);
118    let wd_t = T::full(&[BLOCK_SIZE], weight_decay);
119
120    let g_eff = g + wd_t * p;
121
122    let exp_avg_new = beta1_t * exp_avg + one_m_beta1 * g_eff;
123    let exp_avg_sq_new = beta2_t * exp_avg_sq + one_m_beta2 * g_eff * g_eff;
124
125    let denom = T::sqrt_rn(exp_avg_sq_new) / bc2sqrt_t + eps_t;
126    let p_new = p - lr_t * (cg_t * g_eff + cm_t * exp_avg_new) / denom;
127
128    T::store(
129        params_ptr.add_offsets(offsets),
130        p_new,
131        Some(mask),
132        &[],
133        None,
134        None,
135    );
136    T::store(
137        exp_avg_ptr.add_offsets(offsets),
138        exp_avg_new,
139        Some(mask),
140        &[],
141        None,
142        None,
143    );
144    T::store(
145        exp_avg_sq_ptr.add_offsets(offsets),
146        exp_avg_sq_new,
147        Some(mask),
148        &[],
149        None,
150        None,
151    );
152}
teeny_kernels/nn/optim/nadam.rs

teeny_kernels/nn/optim/
nadam.rs
