cubek_std/tile/ops/
softmax.rs

1use cubecl::prelude::*;
2
3use crate::StageIdent;
4use crate::tile::mask::Mask;
5use crate::tile::variants::InnerLayout;
6use crate::tile::{Plane, RowWise, Tile, TileExpand, TileKind, TileKindExpand};
7
8/// Logits below this are considered masked (effectively -inf).
9/// Value chosen to fit within f16 range (~-65,504 max).
10pub const LOGIT_MASKED: f32 = -6e4;
11
12/// Row-shape descriptor for online softmax.
13#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
14pub enum SoftmaxKind {
15    /// `Tile::Unit` / `Tile::Register`: each unit owns its full tile.
16    Direct { num_rows_per_unit: u32 },
17    /// `Tile::WhiteboxFragment` / `Tile::Bounce`: plane-fragmented.
18    Plane { inner_layout: InnerLayout },
19}
20
21impl SoftmaxKind {
22    pub const fn num_rows_per_unit(&self) -> u32 {
23        match self {
24            SoftmaxKind::Direct { num_rows_per_unit } => *num_rows_per_unit,
25            SoftmaxKind::Plane { inner_layout } => match inner_layout {
26                InnerLayout::Contiguous => 1,
27                InnerLayout::SplitRows => 2,
28            },
29        }
30    }
31}
32
33/// Initial `(m, l)` running state for online softmax.
34#[cube]
35pub fn softmax_init_state<E: Float>(
36    #[comptime] num_rows_per_unit: u32,
37) -> (RowWise<E>, RowWise<E>) {
38    (
39        RowWise::<E>::new_min_value(num_rows_per_unit as usize),
40        RowWise::<E>::new_zero(num_rows_per_unit as usize),
41    )
42}
43
44#[cube]
45impl<Acc: Float> Tile<Acc, Plane> {
46    /// Online softmax update fused with the precision-cast write into the
47    /// value-matmul lhs tile.
48    pub fn softmax<Lhs: Float, M: Mask>(
49        &mut self,
50        mask: &M,
51        softmaxed_tile: &mut Tile<Lhs, Plane>,
52        state: &mut (RowWise<Acc>, RowWise<Acc>),
53        head_dim_factor: Acc,
54    ) -> RowWise<Acc> {
55        match &mut self.kind {
56            TileKind::Bounce(s) => {
57                s.softmax::<Lhs, M>(mask, softmaxed_tile, state, head_dim_factor)
58            }
59            TileKind::WhiteboxFragment(s) => {
60                s.softmax::<Lhs, M>(mask, softmaxed_tile, state, head_dim_factor)
61            }
62            TileKind::Unit(s) => s.softmax::<Lhs, M>(mask, softmaxed_tile, state, head_dim_factor),
63            TileKind::Register(s) => {
64                s.softmax::<Lhs, M>(mask, softmaxed_tile, state, head_dim_factor)
65            }
66            _ => panic!("softmax: unsupported score variant"),
67        }
68    }
69
70    /// Copy `self` into `dest`.
71    pub fn write_results<DE: Float, DS: Size>(&self, dest: &mut Tile<DE, Plane>) {
72        dest.copy_from::<Acc, DS, Acc, Acc, Acc>(self, StageIdent::Out);
73    }
74}
cubek_std/tile/ops/softmax.rs

cubek_std/tile/ops/
softmax.rs
