Softmax - Numerically Stable Implementation

Implements softmax activation for attention scores.

Mathematical Definition

Standard softmax:

softmax(x)ᵢ = exp(xᵢ) / Σⱼ exp(xⱼ)

Problem: Numerical instability when xᵢ is large (exp overflow)

Numerically Stable Softmax

1. Find maximum: M = max(x)
2. Subtract max:  y = x - M  (shifts values to ≤ 0)
3. Exponentiate:  z = exp(y)
4. Normalize:     softmax(x) = z / sum(z)

Why stable:

• Largest value becomes exp(0) = 1 (no overflow)
• All other values are exp(negative) ∈ (0, 1)
• Mathematically equivalent due to cancellation:

  exp(xᵢ - M) / Σⱼ exp(xⱼ - M)
  = [exp(xᵢ) / exp(M)] / [Σⱼ exp(xⱼ) / exp(M)]
  = exp(xᵢ) / Σⱼ exp(xⱼ)
  

Attention Masking

In causal (autoregressive) attention, we mask future positions:

mask[i,j] = -∞ if j > i  (can't attend to future)
mask[i,j] = 0   if j ≤ i  (can attend to past)

softmax_scores = softmax(attention_scores + mask)

After softmax, masked positions become ≈ 0.

Implementation Strategy

Two-pass approach (memory-efficient):

Pass 1: Parallel reduction to find max(x) and sum(exp(x - max))
Pass 2: Normalize each element: xᵢ / sum

Single-pass approach (simplified):

Each workgroup handles one row (one query token's attention to all keys)
Uses shared memory for max reduction and sum reduction

References

• Attention is All You Need (Vaswani et al., 2017)
• llama.cpp: ggml/src/ggml.c (ggml_soft_max_impl)
• Flash Attention: https://arxiv.org/abs/2205.14135

Configuration

structure Hesper.Layers.Softmax.Config : Type

Softmax configuration

• rowSize : Nat
• numRows : Nat
• useMask : Bool

instance Hesper.Layers.Softmax.instReprConfig : Repr Config

Equations

• Hesper.Layers.Softmax.instReprConfig = { reprPrec := Hesper.Layers.Softmax.instReprConfig.repr }

def Hesper.Layers.Softmax.instReprConfig.repr : Config → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.

GPU Kernel Implementation

def Hesper.Layers.Softmax.softmaxKernel (config : Config) : WGSL.Monad.ShaderM Unit

Softmax kernel with numerical stability

Input shape: [num_rows, row_size] Output shape: [num_rows, row_size]

Algorithm:

for each row:
  1. Find max value: M = max(row)
  2. Compute shifted exp: zᵢ = exp(xᵢ - M)
  3. Compute sum: S = Σᵢ zᵢ
  4. Normalize: yᵢ = zᵢ / S

Workgroup strategy:

• One workgroup per row (for small sequences)
• Parallel reduction for max/sum within workgroup
• Shared memory for intermediate results

@param config Softmax configuration

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Layers.Softmax.softmaxMaskedKernel (config : Config) : WGSL.Monad.ShaderM Unit

Softmax with causal masking

Applies causal mask before softmax: mask[i,j] = -∞ if j > i

@param config Softmax configuration (useMask should be true)

Equations

• One or more equations did not get rendered due to their size.

High-Level API

structure Hesper.Layers.Softmax.Softmax : Type

Softmax layer structure

• config : Config

def Hesper.Layers.Softmax.create (config : Config) : IO Softmax

Create Softmax layer (no learned parameters)

@param config Softmax configuration

Equations

• One or more equations did not get rendered due to their size.

@[inline]

def Hesper.Layers.Softmax.forward {β : Type} [GPUBackend β] (ctx : β) (layer : Softmax) (inputBuf outputBuf : GPUBackend.Buf β) : IO Unit

Apply softmax

@param device WebGPU device @param layer Softmax layer @param inputBuf GPU buffer [num_rows, row_size] @param outputBuf GPU buffer for output (same shape)

Equations

• One or more equations did not get rendered due to their size.

Optimized: Flash Softmax (Future)

def Hesper.Layers.Softmax.flashSoftmaxKernel (config : Config) : WGSL.Monad.ShaderM Unit

Flash Softmax: Fused softmax + attention matmul

This is part of Flash Attention optimization, where softmax is fused with the attention score computation to reduce memory bandwidth.

Key idea: Instead of:

1. Compute Q @ K^T (write to memory)
2. Softmax (read from memory, write back)
3. @ V (read from memory)

Do:

1. Compute chunks of Q @ K^T in registers
2. Apply softmax to chunks
3. Multiply by V chunk immediately
4. Accumulate result (never write intermediate scores)

Benefits:

• 2-4x speedup on long sequences
• O(N) memory instead of O(N²)
• Enables much longer context windows

This is marked for future implementation when integrating with Attention layer.

Equations

• Hesper.Layers.Softmax.flashSoftmaxKernel config = pure ()