gpu: three of the four remaining backend-layer gaps

docs/backends.md

@@ -60,6 +60,15 @@ buffer, so an output into a buffer whose live bytes are the host's has to bring
 them up or lose the rest, while an output into a fresh buffer, which is nearly
 every output, has nothing to bring.
 
+The table that holds those copies — host pointer, device handle, size,
+`residency` — and the size-keyed free list `alloc`/`release` recycle buffers
+through are `gpu::mirror_table<Handle>` (`gpu_abi.h`), one instantiation per
+mirrored backend (`mirror_table<CUdeviceptr>` in `cuda.h`,
+`mirror_table<wgpu::Buffer>` in `webgpu.h`). Only `Handle` — the backend's own
+device buffer type — varies; the map, the pool and the copying policy do not.
+A backend still does the actual copy (`before_kernel_`, `sync_to_host`) since
+that call differs by driver.
+
 ## The kernel ABI
 
 An op hands the backend a kernel id, an ordered list of views, a params
@@ -144,23 +153,27 @@ differ, not to avoid reordering a params struct.
 
 The ops fall in two tiers. Tier 0 — elementwise, broadcast, reduction, GEMM,
 copy and index — closes the array surface: a backend with those kernels runs
-every graph. Tier 1 is the fused ops of the model path — `rmsnorm`, `swiglu`,
+every graph. Tier 1 is the fused ops — the model path's `rmsnorm`, `swiglu`,
 the decode GEMV, the cache writes, `rope`, the attention, `split_heads` /
-`merge_heads`, `argmax` — each a kernel a backend *may* have. Under each of
-them `gpu_ops.h` holds one generic composition out of tier 0 (`gpu::generic`),
-and the op takes it when the shared launch declines or `own` has no member:
-several launches and a scratch buffer or two where the kernel is one launch,
-so a backend that cares for the decode loop writes the kernel, and a backend
-that has only tier 0 still runs the whole model path. The compositions are
-f32: an operand the tier has no reader for — a bf16 cache or weight, int4
-weights — still declines, and a model keeps to the array ops there
+`merge_heads`, `argmax`, plus two training ops off that path, `xent_bwd` and
+`adam_step` — each a kernel a backend *may* have. Under each of them
+`gpu_ops.h` holds one generic composition out of tier 0 (`gpu::generic`), and
+the op
+takes it when the shared launch declines or `own` has no member: several
+launches and a scratch buffer or two where the kernel is one launch, so a
+backend that cares writes the kernel, and a backend that has only tier 0
+still runs the whole model path — and, for `adam_step`, still updates its
+state on the device rather than a host round trip. The compositions are f32:
+an operand the tier has no reader for — a bf16 cache or weight, int4 weights
+— still declines, and a model keeps to the array ops there
 (`caps::row_gemv`, `caps::bf16_gemm`). A composition names no backend and no
 array: spans, tier-0 ops, and the device core's `alloc` / `release` /
 `cpu_barrier` / `sync_to_host`.
 
-Two tests hold the two routes together: the model-path test checks the op,
-however it ran, against the array oracle, and a second runs each composition
-beside the backend's kernel and requires the same numbers.
+Two tests hold the two routes together: the model-path test (and, for the two
+training ops off that path, their own eager tests) checks the op, however it
+ran, against the array oracle, and a second runs each composition beside the
+backend's kernel and requires the same numbers.
 
 ## Launch policy
 
@@ -192,6 +205,11 @@ why CUDA's fill is a constant rather than a device query.
 Which of its kernels a backend launches — CUDA's f32 tile and its wave plan,
 Metal's STEEL bands — stays with the backend: a tile is that kernel family's
 ABI. What such a choice measures its grid against is the same `fill_groups`.
+The wave plan itself (`sgemm_wave_chunk_`) has no Metal counterpart to share
+it with (Metal's GEMM dispatch only ever picks a tile, never splits K), and
+its own two-stage layering — a floor'd fill share, then spare slots folded
+into a shorter tail — doesn't fit `gpu::policy::split_rule`'s single ceil
+formula either, so it stays a CUDA-only function.
 
 ## Profiling
 
@@ -250,14 +268,8 @@ asks; a new backend edits no test.
 
 ## What is not shared yet
 
-- The mirror table (handle to host copy, device copy, size, `residency`) and
-  the buffer pool exist twice, in `cuda.h` and `webgpu.h`. The state machine
-  itself is shared.
-- `kop` still lists kernel ids only one backend has (Metal's GEMM tiles and
-  attention variants).
-- CUDA's f32 GEMM wave plan (`sgemm_wave_chunk_`: layers, spare slots and
-  rounds over a two-blocks-per-SM wave) is a split policy of its own, written
-  against `traits::fill_groups` but not yet a shared function.
+- WebGPU's kernels still predate the canonical ABI and go through `marshal_`
+  (see above).
 
 ## Verifying a change
 

include/array.h

@@ -3534,10 +3534,13 @@ struct graph {
                          inv_bc2)) {
         return true;
       }
-      // No kernel for it on this backend: the host loop below. The caller
-      // updates in place and cannot compose its way out, so this has to be
-      // total -- and Metal's unified memory makes the round trip a flush and a
-      // memcpy, the same bargain its other CPU fallbacks make.
+      // gpu::adam_step already falls to its tier-0 composition when the
+      // backend has no fused kernel, so this is the device itself declining
+      // -- a scratch it could not allocate, or a tier-0 launch it refused:
+      // the host loop below. The caller updates in place and cannot compose
+      // its way out, so this has to be total -- and Metal's unified memory
+      // makes the round trip a flush and a memcpy, the same bargain its
+      // other CPU fallbacks make.
     }
     // data() brings any device copy home first, the same as every other host
     // path. A CUDA build hands every buffer a mirror key, so "has a native

include/cuda.h

@@ -226,13 +226,6 @@ inline const char* kernel_name_(kop op) {
     case kop::merge_heads_: return "tl_merge_heads";
     case kop::split_heads_: return "tl_split_heads";
     case kop::argmax_: return "tl_argmax";
-    // Every f32 GEMM id is the one general kernel here (the tiled fast path
-    // has its own names: sgemm_tiles below).
-    case kop::sgemm32: case kop::sgemm32x64: case kop::sgemm64x32:
-    case kop::sgemm64: case kop::steel: case kop::steel32x64:
-    case kop::steel_ta: case kop::steel_tb: case kop::steel32x64_ta:
-    case kop::steel32x64_tb:
-      return "tl_sgemm";
     default: return nullptr;  // no kernel here: dispatch declines
   }
 }
@@ -289,29 +282,19 @@ struct context {
   std::vector<timed_launch> timed;
   std::vector<CUevent> spare_events;
 
-  // Host/device mirror per allocation, keyed by the device pointer (== the
-  // `native` handle stored in storage). Views sharing a storage share the key,
-  // so one state serves every view. When to copy is gpu::residency's decision
-  // (gpu_abi.h); the copies are made here.
-  struct mirror {
-    float* host = nullptr;  // CPU-side buffer (storage.contents/ptr)
-    CUdeviceptr dev = 0;    // device buffer (storage.native)
-    size_t bytes = 0;
-    gpu::residency live;
-  };
-  std::unordered_map<CUdeviceptr, mirror> mirrors;
-
-  // Size-keyed free list (like Metal's MTLBuffer pool). Released buffers are
-  // recycled, not cuMemFree'd — repeated large alloc/free otherwise fragments
-  // the driver allocator (decode benches, training that churns activations).
-  // Buffers persist until the (leaked) context tears down. Keyed by exact byte
-  // size; the workloads that churn reuse identical shapes.
-  std::unordered_map<size_t, std::vector<std::pair<CUdeviceptr, float*>>> pool;
-
-  mirror* mirror_(void* native) {
-    auto it = mirrors.find(reinterpret_cast<CUdeviceptr>(native));
-    return it == mirrors.end() ? nullptr : &it->second;
-  }
+  // Host/device mirror per allocation and its size-keyed free list (shared
+  // shape with webgpu.h; `gpu::mirror_table`, gpu_abi.h). Keyed by the device
+  // pointer, reinterpreted as the `native` handle stored in storage. Views
+  // sharing a storage share the key, so one state serves every view. When to
+  // copy is gpu::residency's decision; the copies are made here. Released
+  // buffers are recycled, not cuMemFree'd — repeated large alloc/free
+  // otherwise fragments the driver allocator (decode benches, training that
+  // churns activations); they persist until the (leaked) context tears down.
+  using mirror = gpu::mirror_table<CUdeviceptr>::entry;
+  gpu::mirror_table<CUdeviceptr> mt;
+
+  mirror* mirror_(void* native) { return mt.find(native); }
+
   // A kernel is about to touch this buffer as `a`: bring the host copy up if
   // residency says so. Async on the stream like the meta uploads (a blocking
   // copy would wait out every kernel already queued and stall the pipeline
@@ -517,6 +500,11 @@ struct context {
   CUfunction pad_() { return cached_(pad_fn, "tl_pad"); }
   CUfunction fold_() { return cached_(fold_fn, "tl_fold"); }
 
+  // The f32 gemm's general fallback (the tiled fast path above, sgemm_(), has
+  // its own names).
+  CUfunction sgemm_fallback_fn = nullptr;
+  CUfunction sgemm_fallback_() { return cached_(sgemm_fallback_fn, "tl_sgemm"); }
+
   // Embedding-table lookup (index_select/index_add) and pooling-style
   // one-hot scatter (scatter_to_axis), cached the same way.
   CUfunction index_add_fn = nullptr, scatter_axis_fn = nullptr;
@@ -1143,35 +1131,26 @@ inline void* alloc(int64_t bytes, float** contents, bool host_fill = false) {
   size_t nb = bytes > 0 ? (size_t)bytes : 4;
   CUdeviceptr dev = 0;
   float* host = nullptr;
-  auto it = c.pool.find(nb);  // reuse a recycled buffer of this exact size
-  if (it != c.pool.end() && !it->second.empty()) {
-    dev = it->second.back().first;
-    host = it->second.back().second;
-    it->second.pop_back();
-  } else {
+  if (!c.mt.take(nb, dev, host)) {  // no recycled buffer of this exact size
     if (c.d.MemAlloc(&dev, nb) != 0) return nullptr;
     host = static_cast<float*>(std::malloc(nb));
     if (!host) {
       c.d.MemFree(dev);
       return nullptr;
     }
   }
-  c.mirrors[dev] = context::mirror{host, dev, nb, gpu::residency(host_fill)};
+  void* native = reinterpret_cast<void*>(dev);
+  c.mt.insert(native, dev, host, nb, host_fill);
   if (contents) *contents = host;
-  return reinterpret_cast<void*>(dev);
+  return native;
 }
 
 inline void release(void* buf, int64_t, float*) {
   auto& c = context::get();
   if (!c.ready || !buf) return;
-  CUdeviceptr dev = reinterpret_cast<CUdeviceptr>(buf);
-  auto it = c.mirrors.find(dev);
-  if (it == c.mirrors.end()) {
-    c.d.MemFree(dev);  // untracked (shouldn't happen); free outright
-    return;
+  if (!c.mt.release(buf)) {
+    c.d.MemFree(reinterpret_cast<CUdeviceptr>(buf));  // untracked (shouldn't happen); free outright
   }
-  c.pool[it->second.bytes].push_back({dev, it->second.host});  // recycle
-  c.mirrors.erase(it);
 }
 
 // Reconcile a buffer for a CPU access: when the device holds the live copy,
@@ -2075,8 +2054,7 @@ inline bool own::gemm_batched(gpu::span a, int64_t lda, bool ta, int64_t sa,
   unsigned gx = (un + bx - 1) / bx, gy = (um + by - 1) / by;
   if (gx == 0) gx = 1;
   if (gy == 0) gy = 1;
-  // kop::sgemm32 is routed to tl_sgemm by kernel_name_.
-  return c.launch_(c.fn_(kop::sgemm32), {gx, gy}, {bx, by}, 0, pa, pb, po,
+  return c.launch_(c.sgemm_fallback_(), {gx, gy}, {bx, by}, 0, pa, pb, po,
                    pbias, um, un, uk, ula, ulb, uta, utb, scale, offset);
 }
 

include/gpu_abi.h

@@ -19,20 +19,22 @@
 #include <cstddef>
 #include <cstdint>
 #include <string_view>
+#include <unordered_map>
+#include <utility>
+#include <vector>
 
 #include "profile.h"
 
 namespace tl {
 namespace gpu {
 
-// Kernel ids. The shared ops name the ones every backend may implement; the
-// rest are one backend's own (Metal's GEMM tiles, its attention variants),
-// kept here only until that backend keys them privately.
+// Kernel ids. The shared ops name the ones every backend may implement.
+// Metal's own kernels -- the GEMM tiles and the attention variants, neither
+// touched by another backend's dispatch -- are keyed privately in metal.h's
+// own `gtile` and `attn_op` instead, not here.
 enum class kop {
   add, sub, mul, div, pow_, exp_, log_, sqrt_, sigmoid, relu, affine,
   badd, bsub, bmul, bdiv, bpow,  // rank-2 broadcast binary (strided operands)
-  sgemm32, sgemm32x64, sgemm64x32, sgemm64,
-  steel, steel32x64, steel_ta, steel_tb, steel32x64_ta, steel32x64_tb,
   softmax, row_sum, row_max, pad, fold,
   index_select, index_add, scatter_axis,
   badd_nd, bsub_nd, bmul_nd, bdiv_nd, bpow_nd,  // N-D broadcast binary
@@ -44,14 +46,6 @@ enum class kop {
   pow_s_, gt_s_, lt_s_, ge_s_, le_s_, eq_s_, ne_s_,  // scalar_op maps onto these
   layer_norm_,                                       // the fused layer norm
   layer_norm_bwd_dx_, layer_norm_bwd_gb_, layer_norm_bwd_gb_fold_,  // its pullback
-  attn_prefill_64_, attn_prefill_128_,  // causal prefill attention, per D
-  attn_prefill_bf16_64_, attn_prefill_bf16_128_,  // over a bf16 KV cache
-  attn_bwd_dq_64_, attn_bwd_dq_128_, attn_bwd_dkv_64_, attn_bwd_dkv_128_,
-  attn_decode_64_, attn_decode_128_,  // fused decode attention, per D
-  attn_decode_split_64_, attn_decode_split_128_,  // its split-KV pass
-  attn_combine_64_, attn_combine_128_,            // and their partials
-  attn_decode_bf16_64_, attn_decode_bf16_128_,    // the same over a bf16 cache
-  attn_decode_split_bf16_64_, attn_decode_split_bf16_128_,
   kv_append_, kv_append_bf16_, kv_fill_, kv_fill_bf16_,  // the KV cache's writes
   argmax_, rmsnorm_, add_rmsnorm_, swiglu_, split_heads_, merge_heads_,  // decode's rest
   gemv_f32_, gemv_bf16_, gemv_q4_,   // decode GEMVs, per weight dtype
@@ -65,22 +59,15 @@ inline constexpr size_t kKopCount = static_cast<size_t>(kop::adam_step_) + 1;
 // The ids' names, for what reports by kernel (the census, tl::profile).
 inline constexpr const char* kKopNames[] = {
     "add", "sub", "mul", "div", "pow_", "exp_", "log_", "sqrt_", "sigmoid",
-    "relu", "affine", "badd", "bsub", "bmul", "bdiv", "bpow", "sgemm32",
-    "sgemm32x64", "sgemm64x32", "sgemm64", "steel", "steel32x64", "steel_ta",
-    "steel_tb", "steel32x64_ta", "steel32x64_tb", "softmax", "row_sum",
-    "row_max", "pad", "fold", "index_select", "index_add", "scatter_axis",
+    "relu", "affine", "badd", "bsub", "bmul", "bdiv", "bpow", "softmax",
+    "row_sum", "row_max", "pad", "fold", "index_select", "index_add",
+    "scatter_axis",
     "badd_nd", "bsub_nd", "bmul_nd", "bdiv_nd", "bpow_nd", "where_nd",
     "copy_nd", "gt_", "lt_", "ge_", "le_", "eq_", "ne_", "tanh_", "sin_",
     "cos_", "clamp_", "sum_to_", "sum_to_blocked_", "concat_part_", "rope_",
     "pow_s_", "gt_s_", "lt_s_", "ge_s_", "le_s_", "eq_s_", "ne_s_",
     "layer_norm_", "layer_norm_bwd_dx_", "layer_norm_bwd_gb_",
-    "layer_norm_bwd_gb_fold_", "attn_prefill_64_", "attn_prefill_128_",
-    "attn_prefill_bf16_64_", "attn_prefill_bf16_128_", "attn_bwd_dq_64_",
-    "attn_bwd_dq_128_", "attn_bwd_dkv_64_", "attn_bwd_dkv_128_",
-    "attn_decode_64_", "attn_decode_128_", "attn_decode_split_64_",
-    "attn_decode_split_128_", "attn_combine_64_", "attn_combine_128_",
-    "attn_decode_bf16_64_", "attn_decode_bf16_128_",
-    "attn_decode_split_bf16_64_", "attn_decode_split_bf16_128_", "kv_append_",
+    "layer_norm_bwd_gb_fold_", "kv_append_",
     "kv_append_bf16_", "kv_fill_", "kv_fill_bf16_", "argmax_", "rmsnorm_",
     "add_rmsnorm_", "swiglu_", "split_heads_", "merge_heads_", "gemv_f32_",
     "gemv_bf16_", "gemv_q4_", "gemv_combine_", "gemv_bf16_row_",
@@ -176,6 +163,54 @@ struct residency {
   void uploaded() { where = both; }
 };
 
+// The host/device mirror and size-keyed buffer pool a mirrored backend (CUDA,
+// WebGPU) keeps per allocation. `Handle` is the backend's own device buffer
+// type (`CUdeviceptr`, `wgpu::Buffer`); the map, the pool and the entry shape
+// are otherwise identical between them, so only `Handle` varies. A backend
+// still does its own copying (`before_kernel_`, `sync_to_host`) and still
+// keys lookups by whatever `void*` it hands out as `native`; this only owns
+// where the state lives.
+template <class Handle>
+struct mirror_table {
+  struct entry {
+    float* host = nullptr;
+    Handle dev{};
+    size_t bytes = 0;
+    residency live;
+  };
+  std::unordered_map<void*, entry> mirrors;
+  std::unordered_map<size_t, std::vector<std::pair<Handle, float*>>> pool;
+
+  entry* find(void* native) {
+    auto it = mirrors.find(native);
+    return it == mirrors.end() ? nullptr : &it->second;
+  }
+
+  // A released buffer of this exact size, if the pool has one.
+  bool take(size_t bytes, Handle& dev, float*& host) {
+    auto it = pool.find(bytes);
+    if (it == pool.end() || it->second.empty()) return false;
+    dev = it->second.back().first;
+    host = it->second.back().second;
+    it->second.pop_back();
+    return true;
+  }
+
+  void insert(void* key, Handle dev, float* host, size_t bytes, bool host_fill) {
+    mirrors[key] = entry{host, std::move(dev), bytes, residency(host_fill)};
+  }
+
+  // Moves `native`'s entry into the free list, keyed by its size. False if
+  // `native` is not tracked (nothing to release).
+  bool release(void* native) {
+    auto it = mirrors.find(native);
+    if (it == mirrors.end()) return false;
+    pool[it->second.bytes].push_back({it->second.dev, it->second.host});
+    mirrors.erase(it);
+    return true;
+  }
+};
+
 // A launch's extent: how many groups, how many threads in each, and the bytes
 // of per-group scratch a kernel's reduction needs where the backend sizes it at
 // launch (CUDA's shared memory; Metal and WGSL size theirs in the kernel).