spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// Embedding gather + scatter-grad. ids are i32. Forward: one thread per4// (token, channel). Backward: one thread per (row, channel) accumulating5// with device atomic_float — repeated tokens collide, so order is6// nondeterministic (sum of floats); deterministic mode routes embedding7// backward through the CPU reference instead.8#include <metal_stdlib>9using namespace metal;1011kernel void embedding_fwd_f32(device const float* W   [[buffer(0)]],12                              device const int*   ids [[buffer(1)]],13                              device float*       out [[buffer(2)]],14                              constant uint&      C   [[buffer(3)]],15                              uint2 gid [[thread_position_in_grid]]) {16    // gid.x = channel, gid.y = token index17    out[ulong(gid.y) * C + gid.x] = W[ulong(ids[gid.y]) * C + gid.x];18}1920kernel void embedding_bwd_f32(device const int*     ids  [[buffer(0)]],21                              device const float*   dout [[buffer(1)]],22                              device atomic_float*  dW   [[buffer(2)]],23                              constant uint&        C    [[buffer(3)]],24                              uint2 gid [[thread_position_in_grid]]) {25    const float g = dout[ulong(gid.y) * C + gid.x];26    atomic_fetch_add_explicit(&dW[ulong(ids[gid.y]) * C + gid.x], g,27                              memory_order_relaxed);28}29