spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// Fused softmax cross-entropy (llm.c fused_classifier pattern): one4// threadgroup per row computes online (max, sumexp), the loss in5// logsumexp form, and — in the same kernel — the logit gradient6// (softmax − onehot) · inv_n, accumulated into dlogits. The full softmax7// is never materialized. targets use ignore_index = -1 (loss 0, grad 0).8//9// losses[row] receives the per-row loss; the host reduces (sum_f32) and10// scales by 1/n_valid.11#include <metal_stdlib>12using namespace metal;1314struct CEParams {15    uint V;16    float inv_n;    // 1 / n_valid17    uint want_grad; // 0: loss only18};1920kernel void cross_entropy_f32(device const float* logits  [[buffer(0)]],21                              device const int*   targets [[buffer(1)]],22                              device float*       losses  [[buffer(2)]],23                              device float*       dlogits [[buffer(3)]],24                              constant CEParams&  p       [[buffer(4)]],25                              uint row      [[threadgroup_position_in_grid]],26                              uint lid      [[thread_index_in_threadgroup]],27                              uint tg_size  [[threads_per_threadgroup]],28                              uint lane     [[thread_index_in_simdgroup]],29                              uint simd_idx [[simdgroup_index_in_threadgroup]],30                              uint n_simds  [[simdgroups_per_threadgroup]]) {31    const int tgt = targets[row];32    device const float* x = logits + ulong(row) * p.V;3334    if (tgt < 0) {35        if (lid == 0) losses[row] = 0.0f;36        return;37    }3839    // online (m, l) as in softmax.metal40    float m = -FLT_MAX;41    float l = 0.0f;42    for (uint j = lid; j < p.V; j += tg_size) {43        const float v = x[j];44        const float m_new = max(m, v);45        l = l * exp(m - m_new) + exp(v - m_new);46        m = m_new;47    }48    float m_simd = simd_max(m);49    float l_simd = simd_sum(l * exp(m - m_simd));5051    threadgroup float tg_m[32];52    threadgroup float tg_l[32];53    if (simd_is_first()) {54        tg_m[simd_idx] = m_simd;55        tg_l[simd_idx] = l_simd;56    }57    threadgroup_barrier(mem_flags::mem_threadgroup);58    float m_row, l_row;59    {60        const uint i = min(lane, n_simds - 1);61        const float mi = tg_m[i];62        const float li = tg_l[i];63        m_row = simd_max(mi);64        const float contrib = (lane < n_simds) ? li * exp(mi - m_row) : 0.0f;65        l_row = simd_sum(contrib);66    }6768    if (lid == 0) {69        // loss = logsumexp − logit[target]70        losses[row] = m_row + log(l_row) - x[uint(tgt)];71    }7273    if (p.want_grad) {74        const float inv_sum = 1.0f / l_row;75        device float* drow = dlogits + ulong(row) * p.V;76        for (uint j = lid; j < p.V; j += tg_size) {77            const float prob = exp(x[j] - m_row) * inv_sum;78            const float ind = (j == uint(tgt)) ? 1.0f : 0.0f;79            drow[j] += (prob - ind) * p.inv_n;80        }81    }82}83