Apple GPU: Gumbel-max inference sampler (#18-safe, reproducible)

benchmarks/apple_gpu/benchmark_gumbel_sampler.py

@@ -0,0 +1,111 @@
+"""apple_gpu Gumbel-max sampler benchmark.
+
+Times the GPU Gumbel-max sampler (`tessera_apple_gpu_gumbel_argmax_f32` —
+per-row vocab argmax on-GPU) against the equivalent host numpy argmax, over a
+sweep of (batch, vocab) shapes. The GPU path's win grows with the batch size
+(many concurrent decode streams sampling at once). Same JSON schema as
+``benchmarks/benchmark_gemm.py``.
+
+Shape spec: ``BxV`` (batch × vocab).
+
+Usage:
+    python benchmarks/apple_gpu/benchmark_gumbel_sampler.py \\
+        --shapes 1x128000 8x128000 64x128000 256x32000 --reps 50
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import statistics
+import sys
+import time
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from tessera import runtime as R
+from tessera import rng as TR
+
+
+def _parse(spec: str):
+    parts = spec.lower().split("x")
+    if len(parts) != 2:
+        raise ValueError(f"shape must be BxV, got {spec!r}")
+    return int(parts[0]), int(parts[1])
+
+
+def _time(fn, reps):
+    fn()
+    s = []
+    for _ in range(reps):
+        t0 = time.perf_counter_ns()
+        fn()
+        s.append((time.perf_counter_ns() - t0) / 1e6)
+    return statistics.median(s), statistics.stdev(s) if reps > 1 else 0.0
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--shapes", nargs="+",
+                        default=["1x128000", "8x128000", "64x128000", "256x32000"])
+    parser.add_argument("--reps", type=int, default=50)
+    parser.add_argument("--output", type=Path, default=None)
+    args = parser.parse_args(argv)
+
+    if sys.platform != "darwin":
+        if args.output is not None:
+            args.output.write_text(json.dumps(
+                {"runs": [], "skipped_apple_gpu": "non-Darwin host"},
+                indent=2, sort_keys=True))
+        print("apple_gpu gumbel benchmark: skipping (non-Darwin)", file=sys.stderr)
+        return 0
+
+    version = "dev"
+    try:
+        import importlib.metadata
+        version = importlib.metadata.version("tessera")
+    except Exception:
+        pass
+
+    rows: list[dict[str, Any]] = []
+    for shape in args.shapes:
+        B, V = _parse(shape)
+        rng = np.random.RandomState(0)
+        logits = rng.randn(B, V).astype(np.float32)
+        key = TR.RNGKey.from_seed(0)
+        gumbel = R._gumbel_noise_from_key((B, V), key, np)
+
+        def gpu():
+            return R._apple_gpu_gumbel_sample(logits, np, key=key, temperature=1.0)
+
+        def host():
+            return np.argmax(logits + gumbel, axis=-1)
+
+        for mode, fn in (("gpu", gpu), ("host_numpy", host)):
+            ms, stdev_ms = _time(fn, args.reps)
+            rows.append({
+                "backend": "apple_gpu",
+                "op": "gumbel_sample",
+                "shape": shape,
+                "dtype": "f32",
+                "mode": mode,
+                "reps": args.reps,
+                "latency_ms": ms,
+                "stdev_ms": stdev_ms,
+                "tflops": 0.0,
+                "memory_bw_gb_s": (B * V * 4 / (ms / 1000.0)) / 1e9 if ms > 0 else 0.0,
+                "device": "apple_silicon_metal",
+                "tessera_version": version,
+            })
+
+    output = json.dumps({"runs": rows}, indent=2, sort_keys=True)
+    if args.output is not None:
+        args.output.write_text(output)
+    else:
+        print(output)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

docs/apple_gpu_tier2_tier3_plan.md

@@ -43,7 +43,7 @@ MPSGraph nodes cover most of what the original framing assumed needed bespoke MS
 | Item | Mapping | Effort | Recommendation |
 |---|---|---|---|
 | **Reductions** (`sum`/`mean`/`var`/`std`/`amax`/`amin`/`prod`/`argmax`/`argmin`/`cumsum`/`cumprod`) | **DONE** — `tessera_apple_gpu_mpsgraph_{reduce,argreduce,scan}_f32`; `runtime.py` normalizes arbitrary axis/keepdims/ddof by folding reduced axes to the last dim. `tests/unit/test_apple_gpu_reductions.py` (51). | Low–Med | **Done** |
-| **`dropout`, `rng_normal`, `rng_uniform`** | (a) MPSGraph random nodes — quick but **won't bit-match the CPU Philox stream** (breaks Decision #18); (b) hand-written Philox MSL for bit-exactness. | Med / Low | **Defer** (training-side); MSL if pursued |
+| **`dropout`, `rng_normal`, `rng_uniform`** | (a) MPSGraph random nodes — quick but **won't bit-match the CPU Philox stream** (breaks Decision #18); (b) hand-written Philox MSL for bit-exactness. **Inference sampler shipped** (separate from training RNG): `tessera_apple_gpu_gumbel_argmax_f32` + `runtime._apple_gpu_gumbel_sample` — Gumbel-max categorical draw `argmax(logits/T + g)` with the Gumbel noise taken from the canonical Philox stream (so it's deterministic / reproducible / **#18-safe**, no on-GPU RNG), supporting temperature / greedy / top-k / top-p. The per-row vocab argmax runs on-GPU; **honest benchmark finding (`benchmark_gumbel_sampler.py`): it is upload-bound today — host numpy argmax is faster until the logits stay GPU-resident (fully-fused decode) or a Philox-MSL noise generator removes the noise upload.** `tests/unit/test_apple_gpu_gumbel_sampler.py` (10, incl. a 20k-sample distribution-convergence check). | Med / Low | **Defer** (training-side); MSL if pursued |
 | **`conv2d`** | **DONE** — `tessera_apple_gpu_conv2d_{f32,f16}` via MPSGraph `convolution2DWithSourceTensor:weightsTensor:descriptor:` (NHWC source / HWIO weights, full stride/pad/dilation/groups, optional bias, fp32 internal accumulation; bf16 via host fp32 round-trip). Wired into the metadata op-loop + `_APPLE_GPU_CONV_OPS` envelope + driver gating; reference fallback in the stub. `tests/unit/test_apple_gpu_conv2d.py` (13). | Med | **Done** |
 | **`conv3d`** | **DONE** — MPSGraph has no 3-D conv node, so `tessera_apple_gpu_conv3d_{f32,f16}` lower to im2col + a single GPU MPSGraph **batched matmul** (batch = groups, fp32 accumulation): patches gathered to a per-group `[groups, rows, K]` column matrix, weights regrouped to `[groups, K, Cout/groups]`, the dominant GEMM on-GPU; bias + scatter on the host. NDHWC source / DHWIO weights, full stride/pad/dilation/groups, optional bias; bf16 via host fp32 round-trip. Wired into the metadata op-loop + `_APPLE_GPU_CONV_OPS` envelope + driver gating; reference fallback in the stub. `tests/unit/test_apple_gpu_conv3d.py` (12). | High | **Done** |
 

docs/audit/generated/runtime_abi.md

@@ -4,9 +4,9 @@ Generated from `python/tessera/compiler/runtime_abi_audit.py`.  Don't edit by ha
 
 ## Headline
 
-- **124** unique `extern "C" tessera_*` C ABI symbols across all backends.
+- **125** unique `extern "C" tessera_*` C ABI symbols across all backends.
 - **6 / 6** core runtime headers present.
-- **58** Apple GPU kernel families with per-dtype variants.
+- **59** Apple GPU kernel families with per-dtype variants.
 
 ## Core runtime headers
 
@@ -23,7 +23,7 @@ Generated from `python/tessera/compiler/runtime_abi_audit.py`.  Don't edit by ha
 
 | Backend | Unique tessera_* symbols |
 |---------|-------------------------:|
-| `apple` | 113 |
+| `apple` | 114 |
 | `nvidia` | 3 |
 | `x86` | 8 |
 
@@ -65,6 +65,7 @@ Generated from `python/tessera/compiler/runtime_abi_audit.py`.  Don't edit by ha
 | `flash_attn` | `bf16`, `f16`, `f32` |
 | `flash_attn_gqa` | `bf16`, `f16`, `f32` |
 | `gelu` | `bf16`, `f16`, `f32` |
+| `gumbel_argmax` | `f32` |
 | `layer_norm` | `f16`, `f32` |
 | `linear_attn` | `f32` |
 | `log_softmax` | `f16`, `f32` |

python/tessera/runtime.py

@@ -2482,6 +2482,97 @@ def _apple_gpu_mpsgraph_reduce_f32() -> Any:
     return sym
 
 
+def _apple_gpu_gumbel_argmax_f32() -> Any:
+    runtime = _load_apple_gpu_runtime()
+    sym = getattr(runtime, "tessera_apple_gpu_gumbel_argmax_f32", None)
+    if sym is None:
+        return None
+    sym.argtypes = [ctypes.POINTER(ctypes.c_float), ctypes.POINTER(ctypes.c_float),
+                    ctypes.POINTER(ctypes.c_int32), ctypes.c_int32,
+                    ctypes.c_int32, ctypes.c_float]
+    sym.restype = None
+    return sym
+
+
+def _gumbel_noise_from_key(shape: tuple, key: Any, np: Any) -> Any:
+    """Gumbel(0,1) noise g = -log(-log(u)) from the canonical Philox stream, so
+    sampling is deterministic + reproducible (and bit-exact vs a CPU reference)
+    without an on-GPU RNG. ``key`` is a ``tessera.rng.RNGKey``; if None, a
+    seed-0 key is used."""
+    from . import rng as _rng
+    if key is None:
+        key = _rng.RNGKey.from_seed(0)
+    u = np.asarray(_rng.uniform(key, shape, dtype="fp32"))
+    u = np.clip(u, 1e-9, 1.0 - 1e-7).astype(np.float32)
+    return (-np.log(-np.log(u))).astype(np.float32)
+
+
+def _apply_topk_topp_mask(logits: Any, top_k: int, top_p: float, np: Any) -> Any:
+    """Mask logits to -inf outside the top-k / top-p (nucleus) set, per row.
+    Operates on a [rows, vocab] f32 copy."""
+    out = logits.astype(np.float32, copy=True)
+    neg_inf = np.float32(-1e30)
+    if top_k and top_k > 0 and top_k < out.shape[-1]:
+        # keep the top_k largest per row; threshold = k-th largest
+        kth = np.partition(out, -top_k, axis=-1)[:, -top_k][:, None]
+        out = np.where(out < kth, neg_inf, out)
+    if top_p and 0.0 < top_p < 1.0:
+        order = np.argsort(-out, axis=-1)
+        sorted_logits = np.take_along_axis(out, order, axis=-1)
+        m = sorted_logits.max(-1, keepdims=True)
+        probs = np.exp(sorted_logits - m)
+        probs /= probs.sum(-1, keepdims=True)
+        cum = np.cumsum(probs, axis=-1)
+        # keep tokens up to and including the one that crosses top_p
+        keep = cum - probs <= top_p
+        keep[:, 0] = True  # always keep the most probable token
+        mask_sorted = np.where(keep, sorted_logits, neg_inf)
+        out = np.empty_like(out)
+        np.put_along_axis(out, order, mask_sorted, axis=-1)
+    return out
+
+
+def _apple_gpu_gumbel_sample(logits: Any, np: Any, *, key: Any = None,
+                             temperature: float = 1.0, top_k: int = 0,
+                             top_p: float = 0.0, greedy: bool = False) -> Any:
+    """GPU Gumbel-max categorical sampler — draws one token id per row of
+    ``logits`` ``[..., vocab]``.
+
+    ``argmax(logits/T + g)`` with Gumbel noise ``g`` (from the Philox ``key``)
+    is an exact draw from ``softmax(logits/T)``; the per-row argmax over the
+    vocab runs on-GPU (the throughput win for batched sampling). ``greedy=True``
+    (or ``temperature==0``) returns the plain argmax. ``top_k`` / ``top_p``
+    restrict the candidate set (host-side mask). Reproducible: same ``key`` +
+    logits ⇒ same tokens. Returns int64 ids shaped like the leading dims of
+    ``logits``; falls back to numpy when the GPU symbol is unavailable."""
+    arr = np.asarray(logits, dtype=np.float32)
+    lead = arr.shape[:-1]
+    vocab = int(arr.shape[-1])
+    rows2d = arr.reshape(-1, vocab)
+    rows = int(rows2d.shape[0])
+
+    masked = _apply_topk_topp_mask(rows2d, top_k, top_p, np)
+    if greedy or temperature == 0.0:
+        gumbel = np.zeros((rows, vocab), np.float32)
+        inv_temp = 1.0
+    else:
+        gumbel = _gumbel_noise_from_key((rows, vocab), key, np)
+        inv_temp = 1.0 / float(temperature)
+
+    sym = _apple_gpu_gumbel_argmax_f32()
+    masked = np.ascontiguousarray(masked, np.float32)
+    gumbel = np.ascontiguousarray(gumbel, np.float32)
+    if sym is not None:
+        out = np.zeros(rows, np.int32)
+        fp = lambda a: a.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
+        sym(fp(masked), fp(gumbel), out.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
+            ctypes.c_int32(rows), ctypes.c_int32(vocab), ctypes.c_float(inv_temp))
+        ids = out.astype(np.int64)
+    else:
+        ids = np.argmax(masked * inv_temp + gumbel, axis=-1).astype(np.int64)
+    return ids.reshape(lead) if lead else ids.reshape(())
+
+
 def _apple_gpu_mpsgraph_argreduce_f32() -> Any:
     runtime = _load_apple_gpu_runtime()
     sym = getattr(runtime, "tessera_apple_gpu_mpsgraph_argreduce_f32", None)

src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm

@@ -8443,8 +8443,79 @@ static void reference_reduce(int op, const float *x, float *out, int32_t rows,
   }
 }
 
+// Gumbel-max categorical sampler: ids = argmax(logits/T + gumbel) per row.
+// The Gumbel noise is supplied by the caller (generated from the Philox stream
+// on the host) so sampling is deterministic + reproducible without an on-GPU
+// RNG — argmax(z + g) with g_i = -log(-log(u_i)) draws from softmax(z). The
+// per-row reduction over the vocab runs on-GPU (the throughput win for batched
+// sampling of many concurrent sequences).
+static bool mpsg_run_gumbel_argmax(MetalDeviceContext &ctx, const float *logits,
+                                   const float *gumbel, int32_t *out,
+                                   int32_t rows, int32_t cols, float invT) {
+  if (rows <= 0 || cols <= 0) return true;
+  @autoreleasepool {
+    size_t xbytes = (size_t)rows * cols * 4;
+    TS_METAL_BUF_ACQUIRE_WITH_BYTES(bufL, ctx, logits, xbytes);
+    TS_METAL_BUF_ACQUIRE_WITH_BYTES(bufG, ctx, gumbel, xbytes);
+    if (!bufL || !bufG) return false;
+    NSArray<NSNumber *> *xs = @[ @(rows), @(cols) ];
+    NSString *key = [NSString stringWithFormat:@"gumbel:%d:%d:%a", rows, cols, invT];
+    NSArray *entry = mpsg_cache_get(key);
+    MPSGraph *g;
+    MPSGraphTensor *pl, *pg, *y;
+    if (entry) {
+      g = entry[0];
+      pl = ((NSArray *)entry[1])[0];
+      pg = ((NSArray *)entry[1])[1];
+      y = entry[2];
+    } else {
+      g = [MPSGraph new];
+      pl = [g placeholderWithShape:xs dataType:MPSDataTypeFloat32 name:nil];
+      pg = [g placeholderWithShape:xs dataType:MPSDataTypeFloat32 name:nil];
+      MPSGraphTensor *scaled = [g multiplicationWithPrimaryTensor:pl
+                                  secondaryTensor:[g constantWithScalar:(double)invT dataType:MPSDataTypeFloat32]
+                                             name:nil];
+      MPSGraphTensor *scores = [g additionWithPrimaryTensor:scaled secondaryTensor:pg name:nil];
+      MPSGraphTensor *idx = [g reductionArgMaximumWithTensor:scores axis:1 name:nil];
+      y = [g castTensor:idx toType:MPSDataTypeInt32 name:nil];
+      mpsg_cache_put(key, @[ g, @[ pl, pg ], y ]);
+    }
+    MPSGraphTensorData *ld = [[MPSGraphTensorData alloc] initWithMTLBuffer:bufL shape:xs dataType:MPSDataTypeFloat32];
+    MPSGraphTensorData *gd = [[MPSGraphTensorData alloc] initWithMTLBuffer:bufG shape:xs dataType:MPSDataTypeFloat32];
+    NSDictionary *res = [g runWithMTLCommandQueue:ctx.queue
+                                            feeds:@{pl : ld, pg : gd}
+                                    targetTensors:@[ y ]
+                                 targetOperations:nil];
+    MPSGraphTensorData *od = res[y];
+    if (!od) return false;
+    [[od mpsndarray] readBytes:out strideBytes:nil];
+    return true;
+  }
+}
+
 }  // namespace
 
+extern "C" void tessera_apple_gpu_gumbel_argmax_f32(const float *logits,
+                                                    const float *gumbel,
+                                                    int32_t *out, int32_t rows,
+                                                    int32_t cols,
+                                                    float inv_temp) {
+  MetalDeviceContext &ctx = deviceContext();
+  if (ctx.ok && mpsg_run_gumbel_argmax(ctx, logits, gumbel, out, rows, cols, inv_temp))
+    return;
+  for (int32_t r = 0; r < rows; ++r) {
+    const float *L = logits + (size_t)r * cols;
+    const float *G = gumbel + (size_t)r * cols;
+    int32_t best = 0;
+    float bs = L[0] * inv_temp + G[0];
+    for (int32_t c = 1; c < cols; ++c) {
+      float s = L[c] * inv_temp + G[c];
+      if (s > bs) { bs = s; best = c; }
+    }
+    out[r] = best;
+  }
+}
+
 extern "C" void tessera_apple_gpu_mpsgraph_reduce_f32(int32_t op, const float *x,
                                                       float *out, int32_t rows,
                                                       int32_t cols) {

src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp

@@ -1405,6 +1405,23 @@ extern "C" void tessera_apple_gpu_mpsgraph_argreduce_f32(int32_t op, const float
     out[r] = best;
   }
 }
+extern "C" void tessera_apple_gpu_gumbel_argmax_f32(const float* logits,
+                                                    const float* gumbel,
+                                                    int32_t* out, int32_t rows,
+                                                    int32_t cols,
+                                                    float inv_temp) {
+  for (int32_t r = 0; r < rows; ++r) {
+    const float* L = logits + static_cast<std::size_t>(r) * cols;
+    const float* G = gumbel + static_cast<std::size_t>(r) * cols;
+    int32_t best = 0;
+    float bs = L[0] * inv_temp + G[0];
+    for (int32_t c = 1; c < cols; ++c) {
+      float s = L[c] * inv_temp + G[c];
+      if (s > bs) { bs = s; best = c; }
+    }
+    out[r] = best;
+  }
+}
 extern "C" void tessera_apple_gpu_mpsgraph_scan_f32(int32_t op, const float* x,
                                                     float* out, int32_t rows,
                                                     int32_t cols) {