spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// Fused online softmax over the last dim. One threadgroup per row; each4// thread keeps a running (max, sum) over a strided slice of the row, lanes5// combine with simd reductions, simdgroups combine through a small6// threadgroup scratch (<= 32 simdgroups, so one simd-reduce finishes it).7// Row layout: X[row * C + j].8//9// Masked/padded values are the caller's problem; -FLT_MAX sentinels are safe10// here (never -INFINITY: fast-math exp(-inf) is undefined — RESEARCH.md §8).11#include <metal_stdlib>12using namespace metal;1314kernel void softmax_f32(device const float* X   [[buffer(0)]],15                        device float*       OUT [[buffer(1)]],16                        constant uint&      C   [[buffer(2)]],17                        uint row      [[threadgroup_position_in_grid]],18                        uint lid      [[thread_index_in_threadgroup]],19                        uint tg_size  [[threads_per_threadgroup]],20                        uint lane     [[thread_index_in_simdgroup]],21                        uint simd_idx [[simdgroup_index_in_threadgroup]],22                        uint n_simds  [[simdgroups_per_threadgroup]]) {23    device const float* x = X + ulong(row) * C;24    device float* out = OUT + ulong(row) * C;2526    // per-thread online (m, l)27    float m = -FLT_MAX;28    float l = 0.0f;29    for (uint j = lid; j < C; j += tg_size) {30        const float v = x[j];31        const float m_new = max(m, v);32        l = l * exp(m - m_new) + exp(v - m_new);33        m = m_new;34    }3536    // combine across the simdgroup37    float m_simd = simd_max(m);38    float l_simd = simd_sum(l * exp(m - m_simd));3940    // combine across simdgroups41    threadgroup float tg_m[32];42    threadgroup float tg_l[32];43    if (simd_is_first()) {44        tg_m[simd_idx] = m_simd;45        tg_l[simd_idx] = l_simd;46    }47    threadgroup_barrier(mem_flags::mem_threadgroup);4849    float m_row, l_row;50    {51        // Indexed by lane-within-simdgroup so EVERY simdgroup runs the same52        // reduction over the same scratch and lands on identical row stats.53        const uint i = min(lane, n_simds - 1); // lanes >= n_simds mirror the last entry54        const float mi = tg_m[i];55        const float li = tg_l[i];56        m_row = simd_max(mi);57        // lanes beyond n_simds would double-count: zero their contribution58        const float contrib = (lane < n_simds) ? li * exp(mi - m_row) : 0.0f;59        l_row = simd_sum(contrib);60    }6162    const float inv_l = 1.0f / l_row;63    for (uint j = lid; j < C; j += tg_size) {64        out[j] = exp(x[j] - m_row) * inv_l;65    }66}67