IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /max/get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo struct

OnlineLogSumExp

struct OnlineLogSumExp[dtype: DType, W: Int = simd_width_of[dtype]()]

Online (flash-style) log-sum-exp monoid (Milakov & Gimelshein, 2018) — the reduction half of softmax.

Holds {m, l} — the sufficient statistics for log(sum_i exp(x_i)), with m = max_i x_i and l = sum_i exp(x_i - m). LSE = m + log(l). Downstream normalization (exp(x - m) / l for softmax, x - m - log(l) for log-softmax) lives in the kernel's pass-2 body.

Single field per statistic: m: SIMD[dtype, W] and l: SIMD[dtype, W]. During accumulation each lane runs an independent flash state over its slice of the axis; join_parallel flash-combines the W lanes into lane 0. Bodies read m[0] and l[0].

Parameters​

  • ​dtype (DType): The accumulator dtype (must be floating-point).
  • ​W (Int): SIMD width of the per-lane flash accumulators. Defaults to the target's simd_width_of[dtype].

Fields​

  • ​m (SIMD[dtype, W]): W-lane running maxima. Final scalar at m[0] post-join_parallel.
  • ​l (SIMD[dtype, W]): W-lane running sum exp(x - m). Final scalar at l[0] post-join_parallel.

Implemented traits​

AnyType, Copyable, ImplicitlyCopyable, ImplicitlyDeletable, Movable, ReduceOp, RegisterPassable, TrivialRegisterPassable

comptime members​

Single​

comptime Single = OnlineLogSumExp[dtype, Int(1)]

width​

comptime width = W

Methods​

__init__​

def __init__() -> Self

Identity: every lane at (-inf, 0).

__getitem__​

def __getitem__(self, j: Int) -> Self.Single

Returns lane j as a width-1 monoid.

Returns:

Self.Single

__setitem__​

def __setitem__(mut self, j: Int, s: OnlineLogSumExp[dtype, Int(1)])

Writes width-1 monoid s into lane j.

accumulate​

def accumulate[val_dtype: DType, w: Int](mut self, val: SIMD[val_dtype, w], idx: SIMD[DType.int64, w] = 0)

Folds a SIMD tile via lane-wise flash update. val is padded to width Self.W via Self.pad with -inf tail lanes; padded lanes contribute nothing because exp(-inf - new_m) = 0, and a lane-wise active-mask zeroes the correction factor where both m and val are -inf (avoids exp(NaN)).

Parameters:

  • ​val_dtype (DType): The tile's source dtype.
  • ​w (Int): SIMD width of the tile.

Args:

join​

def join(mut self, other: Self)

Sequential combine: lane-wise flash combine.

Args:

  • ​other (Self): The state to combine into self.

reduce​

def reduce(self) -> Self.Single

Flash-combines the W lane partials into one (m, l) via the reduce_max / reduce_add intrinsics.

Explicit horizontal reduce: global max of m, then the max-corrected sum sum_j l[j] * exp(m[j] - lane_max). Lanes with m[j] = -inf carry l[j] = 0, so their corrected term is 0.

Returns:

Self.Single: A width-1 OnlineLogSumExp with the combined (m, l) in lane 0.

join_parallel​

def join_parallel[R: Reducer](mut self, reducer: R)

Cross-thread flash combine: global max on m[0], per-participant correction l *= exp(m[0] - new_m), global sum on l, both splatted across all lanes. Runs after reduce, so m[0]/l[0] already hold the within-thread flash state.

Parameters:

  • ​R (Reducer): The parallel scalar reducer.

Args:

  • ​reducer (R): The reducer instance.

Models

Cookbook

Puzzles

Blog

Community

Contact