teeny_kernels/nn/activation/
softmax.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_core::dtype::Float;
20use teeny_macros::kernel;
21use teeny_triton::triton::{
22    types::{AddOffsets, Comparison},
23    *,
24};
25
26/// Row-wise softmax forward pass.
27///
28/// Grid: one CTA per row — `pid = row index`.
29///
30/// Each CTA loads the entire row of `BLOCK_SIZE` elements, applies Triton's
31/// numerically-stable `softmax` builtin (`max`-subtraction + exp + normalise),
32/// and stores the result.
33///
34/// **Constraint**: `BLOCK_SIZE` must equal `n_cols` for this kernel; the caller
35/// is responsible for rounding `n_cols` up to the next power of two and passing
36/// that as `BLOCK_SIZE`.  No masking is needed when `BLOCK_SIZE == n_cols`.
37// ANCHOR: softmax_forward
38#[kernel]
39pub fn softmax_forward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
40    x_ptr: T::Pointer<D>,
41    y_ptr: T::Pointer<D>,
42    _n_rows: i32,
43    n_cols: i32,
44) where
45    T::I32Tensor: types::Tensor<i32, 1>,
46    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
47    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
48{
49    let pid = T::program_id(Axis::X);
50    let row_offset = pid * n_cols;
51    let col_offsets = T::arange(0, BLOCK_SIZE);
52    let offsets = col_offsets + row_offset;
53
54    let x = T::load(
55        x_ptr.add_offsets(offsets),
56        None,
57        None,
58        &[],
59        None,
60        None,
61        None,
62        false,
63    );
64
65    // Triton's builtin: numerically-stable softmax (max subtraction, exp, sum, div).
66    let y = T::softmax(x, None, false, false);
67
68    T::store(y_ptr.add_offsets(offsets), y, None, &[], None, None);
69}
70// ANCHOR_END: softmax_forward
71
72/// Row-wise softmax backward pass.
73///
74/// Given the saved softmax output `y = softmax(x)` and the upstream gradient
75/// `dy`, computes the input gradient:
76///
77/// ```text
78/// dx_i = y_i * (dy_i - sum_j(y_j * dy_j))
79/// ```
80///
81/// Grid: one CTA per row — `pid = row index`.
82///
83/// The dot product `sum(y * dy)` is a row-scalar that is broadcast back to the
84/// full row when computing `dy - dot`.
85///
86/// **Constraint**: `BLOCK_SIZE` must equal `n_cols` (same as the forward pass).
87#[kernel]
88pub fn softmax_backward<T: Triton, D: Float, const BLOCK_SIZE: i32>(
89    dy_ptr: T::Pointer<D>,
90    y_ptr: T::Pointer<D>,
91    dx_ptr: T::Pointer<D>,
92    _n_rows: i32,
93    n_cols: i32,
94) where
95    T::I32Tensor: types::Tensor<i32, 1>,
96    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
97    T::Pointer<D>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<D>>>,
98{
99    let pid = T::program_id(Axis::X);
100    let row_offset = pid * n_cols;
101    let col_offsets = T::arange(0, BLOCK_SIZE);
102    let offsets = col_offsets + row_offset;
103
104    let dy = T::load(
105        dy_ptr.add_offsets(offsets),
106        None,
107        None,
108        &[],
109        None,
110        None,
111        None,
112        false,
113    );
114    let y = T::load(
115        y_ptr.add_offsets(offsets),
116        None,
117        None,
118        &[],
119        None,
120        None,
121        None,
122        false,
123    );
124
125    // dot = sum_j(y_j * dy_j)  — a per-row scalar (0-D tensor after reduction).
126    let dot = T::sum(y * dy, Some(0), false);
127
128    // dx_i = y_i * (dy_i - dot)  — broadcast dot across the row.
129    let dx = y * (dy - dot);
130
131    T::store(dx_ptr.add_offsets(offsets), dx, None, &[], None, None);
132}
133
134impl<D: Float + Send + Sync + 'static> teeny_core::model::RuntimeOp for SoftmaxForward<D> {
135    fn n_activation_inputs(&self) -> usize {
136        1
137    }
138
139    fn param_shapes(&self, _input_shapes: &[&[usize]], _output_shape: &[usize]) -> Vec<Vec<usize>> {
140        Vec::new()
141    }
142
143    fn pack_args(
144        &self,
145        inputs: &[(teeny_core::model::RawPtr, &[usize])],
146        _params: &[teeny_core::model::RawPtr],
147        output: teeny_core::model::RawPtr,
148        output_shape: &[usize],
149        _output_row_stride: i32,
150        visitor: &mut dyn teeny_core::device::program::ArgVisitor,
151    ) {
152        // kernel args: x_ptr, y_ptr, n_rows, n_cols
153        let n_rows = output_shape[0] as i32;
154        let n_cols = output_shape[1] as i32;
155        visitor.visit_ptr(inputs[0].0);
156        visitor.visit_ptr(output);
157        visitor.visit_i32(n_rows);
158        visitor.visit_i32(n_cols);
159    }
160
161    fn block(&self) -> [u32; 3] {
162        [128, 1, 1]
163    }
164
165    // One CTA per row.
166    fn grid(&self, output_shape: &[usize]) -> [u32; 3] {
167        [output_shape[0] as u32, 1, 1]
168    }
169}
170
171pub struct SoftmaxOp<'a, T: Float> {
172    pub forward: SoftmaxForward<T>,
173    pub backward: SoftmaxBackward<T>,
174    _marker: core::marker::PhantomData<&'a ()>,
175}
teeny_kernels/nn/activation/softmax.rs

teeny_kernels/nn/activation/
softmax.rs
