teeny_kernels/nn/optim/
asgd.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/// ASGD step (Averaged SGD) — averaging-active phase (step > t0).
26///
27/// ```text
28/// p  = p - lr * g
29/// ax = ax + (p - ax) / d_ax        where d_ax = max(1, step - t0)
30/// ```
31///
32/// `d_ax = max(1.0, step - t0)` is precomputed on the host.
33///
34/// Grid: `[ceil(n_elements / BLOCK_SIZE), 1, 1]`.
35#[kernel]
36pub fn asgd_step<T: Triton, const BLOCK_SIZE: i32>(
37    params_ptr: T::Pointer<f32>,
38    grad_ptr: T::Pointer<f32>,
39    ax_ptr: T::Pointer<f32>,
40    n_elements: i32,
41    lr: f32,
42    weight_decay: f32,
43    d_ax: f32,
44) where
45    T::I32Tensor: types::Tensor<i32, 1>,
46    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
47    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
48{
49    let pid = T::program_id(Axis::X);
50    let block_start = pid * BLOCK_SIZE;
51    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
52    let mask = offsets.lt(n_elements);
53
54    let p = T::load(
55        params_ptr.add_offsets(offsets),
56        Some(mask),
57        None,
58        &[],
59        None,
60        None,
61        None,
62        false,
63    );
64    let g = T::load(
65        grad_ptr.add_offsets(offsets),
66        Some(mask),
67        None,
68        &[],
69        None,
70        None,
71        None,
72        false,
73    );
74    let ax = T::load(
75        ax_ptr.add_offsets(offsets),
76        Some(mask),
77        None,
78        &[],
79        None,
80        None,
81        None,
82        false,
83    );
84
85    let lr_t = T::full(&[BLOCK_SIZE], lr);
86    let wd_t = T::full(&[BLOCK_SIZE], weight_decay);
87    let d_ax_t = T::full(&[BLOCK_SIZE], d_ax);
88
89    let g_eff = g + wd_t * p;
90    let p_new = p - lr_t * g_eff;
91    let ax_new = ax + (p_new - ax) / d_ax_t;
92
93    T::store(
94        params_ptr.add_offsets(offsets),
95        p_new,
96        Some(mask),
97        &[],
98        None,
99        None,
100    );
101    T::store(
102        ax_ptr.add_offsets(offsets),
103        ax_new,
104        Some(mask),
105        &[],
106        None,
107        None,
108    );
109}
teeny_kernels/nn/optim/asgd.rs

teeny_kernels/nn/optim/
asgd.rs
