gpu: launch recording, split policy and the model path, all in one place

bench/models/check_qwen.cpp

@@ -97,11 +97,12 @@ int main(int argc, char** argv) {
 
   // The same sequence again, on the imperative path: the fused model-path
   // kernels (rmsnorm/rmsnorm_res/swiglu, the decode GEMVs, kv_append,
-  // attn_decode, argmax) instead of the array compositions the checkpoints
-  // above validated. They are meant to compute the same thing, so at F32 the
-  // greedy sequence must be identical — which makes every one of them gated
-  // against the numpy reference too, without a second oracle.
-  if (tl::gpu::caps::model_path) {
+  // attn_decode, argmax), or their generic compositions on a backend without
+  // them, instead of the array compositions the checkpoints above validated.
+  // They are meant to compute the same thing, so at F32 the greedy sequence
+  // must be identical — which makes every one of them gated against the numpy
+  // reference too, without a second oracle.
+  {
     qm::reset_cache(M);
     int64_t p = 0, tok = 0;
     for (int64_t i = 0; i < NP; i++) tok = qm::step_imperative(M, qwenoracle::prompt_ids[i], p++);
@@ -112,8 +113,6 @@ int main(int argc, char** argv) {
       tok = qm::step_imperative(M, tok, p++);
     }
     std::printf("\n  greedy %s\n", imp_ok ? "MATCH" : "DIVERGE");
-  } else {
-    std::printf("\nimperative path: backend has no model path — skipped\n");
   }
 
   ok_f32 = emb_mr < 1e-3 && l0_mr < 5e-3 && fn_mr < 5e-3 && logit_mr < 5e-3 &&

docs/backends.md

@@ -123,23 +123,45 @@ template <class Own = own>
 inline bool rope(span x, span o, int64_t rows, int64_t T, int64_t D,
                  int64_t pos, float base, span bias = {}) {
   if constexpr (detail::owns_rope<Own>::value) {
-    return Own::rope(x, o, rows, T, D, pos, base, bias);
-  } else {
-    return false;
+    if (Own::rope(x, o, rows, T, D, pos, base, bias)) return true;
   }
+  return generic::rope(x, o, rows, T, D, pos, base, bias);
 }
 
 // metal.h — declared in `struct own`, defined among its helpers
 inline bool own::rope(gpu::span x, gpu::span out, ...) { ... }
 ```
 
-`gpu_ops.h` detects the member and forwards to it, or answers false. A backend
-declares what it has and nothing else: there are no stubs. A member whose
-signature drifts from the shared one is a compile error, not a silent fallback.
+`gpu_ops.h` detects the member and forwards to it, or takes the generic
+composition (below). A backend declares what it has and nothing else: there
+are no stubs. A member whose signature drifts from the shared one is a compile
+error, not a silent fallback.
 
 Prefer the single-kernel form. Reach for `own` when the kernels genuinely
 differ, not to avoid reordering a params struct.
 
+## Tiers
+
+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`,
+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
+(`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.
+
 ## Launch policy
 
 The shapes ops launch in live in `gpu::policy` (`gpu_abi.h`), shared host code:
@@ -148,9 +170,40 @@ A `grid` is groups x threads-per-group plus the bytes of per-group scratch a
 reduction needs where the backend sizes it at launch (CUDA's shared memory;
 Metal and WGSL size theirs in the kernel).
 
-What differs between backends' kernels comes in through the backend's `traits`.
-Today that is one fact: whether a rank-2 elementwise kernel reads its cell from
-a 2-D thread position or from a flat index (`traits::cells_2d`).
+What differs between backends' kernels comes in through the backend's `traits`:
+whether a rank-2 elementwise kernel reads its cell from a 2-D thread position
+or from a flat index (`traits::cells_2d`), whether a launch's profile row
+carries a device time (`traits::times_launches`), and how many groups keep
+the device busy (`traits::fill_groups`: 164 on CUDA, two per SM of an RTX
+3090; 64 on Metal).
+
+A reduction that leaves the device short of that many groups is split over
+more of them and combined after: a decode GEMV's K, decode attention's keys,
+a bf16 GEMM tile's K. Every such split is `policy::split_parts` and
+`policy::split_chunk` (`gpu_abi.h`) under a `split_rule` — the kernel's side
+of the decision: the target (the fill, or a multiple of it for a kernel with
+small groups), the shortest k worth splitting, the shortest part, and what a
+part must be a multiple of. A backend states its rule next to its kernel and
+launches the count the rule gives; it does not do the arithmetic itself. The
+CUDA decode attention's device twin (`attn_dpos_chunk` in
+`tensorlib_cuda.cu`) reproduces its rule for the captured-graph path, which is
+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`.
+
+## Profiling
+
+Every kernel launch on every backend is one row under `tl::profile`, made in
+one place: `gpu::launched(kernel name)` (`gpu_abi.h`), which a backend's launch
+primitive — the one place its launches funnel through, shared ops and own ops
+alike — calls once per launch. What the backend adds is the device time, where
+it has one: CUDA brackets the launch with events and stamps the row when the
+stream drains, Metal commits the launch as its own command buffer and stamps
+the row at the flush, WebGPU and the host backend count. `TL_PROFILE=1` starts
+at the first evaluation or the first launch, so a decoder on the model path,
+which never reaches the evaluator, is profiled from its first kernel.
 
 ## Adding an op
 
@@ -200,17 +253,11 @@ asks; a new backend edits no test.
 - 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.
-- `tl::profile` hooks sit in each real backend's launch path, because each
-  stamps its launches with a device time its own way. A backend that records
-  none (`traits::profiles_launches = false`) gets a row per launch from the
-  shared layer, by kernel name, so it is profiled from its first kernel.
 - `kop` still lists kernel ids only one backend has (Metal's GEMM tiles and
   attention variants).
-- Launch policy inside the own ops (CUDA's split-K and tile choices, Metal's
-  GEMM ladder) is still the backend's.
-- There is no generic composition under the fused ops, so a backend without the
-  model-path kernels reports `caps::model_path = false` rather than running them
-  slowly. WebGPU is in that position.
+- 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.
 
 ## Verifying a change
 

include/cuda.h

@@ -743,11 +743,10 @@ struct context {
                     void** argv) {
     if (!f) return false;
     pending = true;
-    // Profiling: the launch under the open scope, and — outside a graph
-    // capture, where an event record would become a graph node — an event on
-    // each side of it for the elapsed time.
-    profile::row* pr =
-        profile::active() ? profile::detail::launch(name_(f)) : nullptr;
+    // Profiling: the launch's row, and — outside a graph capture, where an
+    // event record would become a graph node — an event on each side of it
+    // for the elapsed time.
+    profile::row* pr = gpu::launched(name_(f));
     CUevent begin = nullptr;
     if (pr && d.timing_ok() && !stream) {  // null = the default stream
       begin = event_();
@@ -896,9 +895,19 @@ struct own {
                              float eps);
 };
 
-// Blocks that keep the GPU busy: ~2 per SM on the 82-SM RTX 3090. The
-// threshold every tile and split-K choice below measures its grid against.
-constexpr long kFillBlocks = 164;
+// What the shared launch policy (gpu_abi.h) may assume of this backend's
+// kernels.
+struct traits {
+  // A [rows, cols] elementwise kernel reads its cell from a flat index.
+  static constexpr bool cells_2d = false;
+  // Each launch's tl::profile row carries a device time.
+  static constexpr bool times_launches = true;
+  // Blocks that keep the GPU busy: ~2 per SM on the 82-SM RTX 3090. The
+  // target every tile and split choice below measures its grid against. A
+  // constant rather than a device query because the decode attention's
+  // device twin (attn_dpos_chunk in tensorlib_cuda.cu) bakes it in.
+  static constexpr int64_t fill_groups = 164;
+};
 
 // The bf16 gemm's tile choice: the 128² tile is the more arithmetically
 // efficient, but a few-hundred-row activation against a projection (256×768:
@@ -911,11 +920,11 @@ inline long blocks128_(int64_t m, int64_t n, int64_t batch) {
   return (long)((n + 127) / 128) * ((m + 127) / 128) * batch;
 }
 inline bool big_tile_(int64_t m, int64_t n, int64_t k, int64_t batch = 1) {
-  return blocks128_(m, n, batch) >= kFillBlocks || k % 16 != 0;
+  return blocks128_(m, n, batch) >= traits::fill_groups || k % 16 != 0;
 }
 
 // The wave plan. A launch places one block per SM up to kWaveSingles blocks
-// and two per SM past that, so a wave is kFillBlocks slots; a block takes time
+// and two per SM past that, so a wave is the fill's slots; a block takes time
 // in proportion to its slabs, and a launch lasts as long as its busiest slot.
 // So split K into `full` equal layers that fit the wave, each at least
 // kWaveMinK deep (a block's fixed cost is some twenty 128² slabs), and when
@@ -931,11 +940,16 @@ inline bool big_tile_(int64_t m, int64_t n, int64_t k, int64_t batch = 1) {
 // 256×768×256:nt 5.1k → 5.6k and 256×512×256:nt 3.9k → 4.5k.
 constexpr unsigned kWaveMinK = 192;
 inline unsigned sgemm_wave_chunk_(long tiles, unsigned k, const sgemm_tile& t) {
-  if (tiles >= kFillBlocks) return k;
+  // fill_groups is int64_t, tiles/slabs long — the same width everywhere this
+  // library builds except Windows (LLP64, long is 32-bit), where std::min/max
+  // over the two otherwise deduce to no common type. One cast here keeps the
+  // rest of the function in `long`, as it was before the trait.
+  const long fill = static_cast<long>(traits::fill_groups);
+  if (tiles >= fill) return k;
   const long slabs = k / t.bk, min_slabs = kWaveMinK / t.bk;
-  const long full = std::max<long>(1, std::min(kFillBlocks / tiles, slabs / min_slabs));
+  const long full = std::max<long>(1, std::min<long>(fill / tiles, slabs / min_slabs));
   long chunk_slabs = (slabs + full - 1) / full;
-  if (const long spare = kFillBlocks - full * tiles; spare > 0) {
+  if (const long spare = fill - full * tiles; spare > 0) {
     const long rounds = (tiles + spare - 1) / spare;
     const long parts = full * rounds + 1;  // a full layer is `rounds` tails
     if (slabs >= min_slabs * parts) chunk_slabs = (slabs * rounds + parts - 1) / parts;
@@ -1040,22 +1054,7 @@ inline bool own::argmax(gpu::span a, int64_t n, int64_t* out_idx) {
   return true;
 }
 
-// What the shared launch policy (gpu_ops.h) may assume of this backend's
-// kernels.
-struct traits {
-  // A [rows, cols] elementwise kernel reads its cell from a flat index.
-  static constexpr bool cells_2d = false;
-  // Launches are recorded under tl::profile by this backend itself, with
-  // (times_launches) a device time on each.
-  static constexpr bool profiles_launches = true;
-  static constexpr bool times_launches = true;
-};
-
 struct caps {
-  // Whether the model-path row is real here, or answers false: a decoder
-  // runs on raw buffers only where it is true, and keeps to the array ops
-  // otherwise (there is no CPU fallback under that row).
-  static constexpr bool model_path = true;
   static constexpr bool graph_capture = true;
   static constexpr bool row_gemv = true;   // gemv_bf16_row: weights as [N,K]
   static constexpr bool bf16_gemm = true;  // gemm_bf16_nt: the batched prefill
@@ -1502,24 +1501,21 @@ inline bool own::scatter_to_axis(gpu::span idx, gpu::span values, gpu::span out,
 // Split-K when the N/256 column-blocks alone underfill the SMs (small-N layers):
 // partition K over gridDim.y, atomicAdd into a pre-zeroed y, so the kernel stays
 // bandwidth-bound rather than occupancy-bound. gridDim.y==1 stores directly.
+// The kernel's rule: no split under K=512, and a part is a multiple of its
+// 32-wide K step.
 inline bool gemv_run_(CUfunction f, float* pa, float* pB, float* py,
                       unsigned un, unsigned uk, unsigned vcols = 1) {
   auto& c = context::get();
   unsigned per = 256u * vcols;  // output columns covered by one block
   unsigned bx = (un + per - 1) / per;
   if (bx == 0) bx = 1;
   unsigned gy = 1, ksplit = uk;
-  const long target = kFillBlocks;
-  if (!c.no_splitk && static_cast<long>(bx) < target && uk >= 512) {
-    unsigned g = static_cast<unsigned>((target + bx - 1) / bx);
-    unsigned chunk = (uk + g - 1) / g;
-    chunk = (chunk + 31u) & ~31u;
-    if (chunk == 0) chunk = 32;
-    unsigned s = (uk + chunk - 1) / chunk;
-    if (s > 1) {
-      gy = s;
-      ksplit = chunk;
-    }
+  constexpr gpu::policy::split_rule rule{traits::fill_groups, 512, 0, 32};
+  const int64_t parts =
+      c.no_splitk ? 1 : gpu::policy::split_parts(bx, uk, rule);
+  if (parts > 1) {
+    ksplit = static_cast<unsigned>(gpu::policy::split_chunk(uk, parts, rule));
+    gy = (uk + ksplit - 1) / ksplit;
   }
   if (gy > 1) {
     // Zero y for the split-K atomicAdd. Async on the stream (ordered before the
@@ -1607,17 +1603,15 @@ inline bool own::gemm_bf16_nt(gpu::span a, gpu::span B, gpu::span out,
   // (>= 448 K-elements), and stop once the grid is comfortably several waves
   // (~8 blocks/SM). Measured at M=512: wd 506 -> 312 us, wo 92 -> 75, while a
   // grid that already fills (gateup, 1216 blocks) correctly declines to split.
+  // As a rule: four times the fill, slices in the 64 tile's 16-deep slabs.
+  constexpr gpu::policy::split_rule rule{4 * traits::fill_groups, 0, 448, 16};
   unsigned z = 1;
   if (blocks < 128 && c.d.MemsetD8Async) {
-    unsigned by_k = (unsigned)(k / 448);
-    unsigned by_fill = (656 + blocks - 1) / blocks;
-    z = by_k < by_fill ? by_k : by_fill;
-    if (z < 1) z = 1;
+    z = (unsigned)gpu::policy::split_parts(blocks, k, rule);
   }
   if (z > 1) {
-    constexpr unsigned BK = 16;  // the 64 tile's K slab
-    unsigned ksplit = ((uK + z - 1) / z + BK - 1) / BK * BK;
-    z = (uK + ksplit - 1) / ksplit;  // recompute after rounding
+    unsigned ksplit = (unsigned)gpu::policy::split_chunk(k, z, rule);
+    z = (uK + ksplit - 1) / ksplit;  // recount after rounding
     // atomicAdd combine needs a zeroed C; async on the stream, so it is ordered
     // before the launch without a host sync (and stays capturable).
     c.d.MemsetD8Async(reinterpret_cast<CUdeviceptr>(po), 0,
@@ -1637,13 +1631,14 @@ inline bool own::gemm_bf16_nt(gpu::span a, gpu::span B, gpu::span out,
 // split needs >=128 keys to amortize its fixed cost). Shared by attn_decode
 // (evaluated at the live ctx) and attn_decode_dpos (evaluated at max_ctx, so
 // the CUDA-graph grid is pos-independent).
+// The kernel's rule: 128-thread blocks, so twice the fill (~4 per SM); no
+// split under 256 keys, a split at least 128 keys and a multiple of its 4
+// warps.
+inline constexpr gpu::policy::split_rule attn_split_rule{
+    2 * traits::fill_groups, 256, 128, 4};
 inline unsigned attn_split_count(unsigned n_heads, int64_t ctx) {
-  const long target = 2 * kFillBlocks;  // ~4 blocks per SM
-  if (n_heads == 0 || (long)n_heads >= target || ctx < 256) return 1;
-  unsigned want = static_cast<unsigned>((target + n_heads - 1) / n_heads);
-  unsigned max_s = static_cast<unsigned>(ctx / 128);  // >=128 keys/split
-  if (want > max_s) want = max_s;
-  return want > 1 ? want : 1;
+  return static_cast<unsigned>(
+      gpu::policy::split_parts(n_heads, ctx, attn_split_rule));
 }
 
 // Launch shape of the tiled prefill attention — the ONE place it lives. It is
@@ -1672,10 +1667,8 @@ inline constexpr unsigned attn_bwd_tile(int64_t D) {
 // stay in lockstep — the host/dpos bit-identity rests on it. Guarded by the
 // attn64 ctest's host-vs-dpos bit-equality sweep.
 inline unsigned attn_split_chunk(unsigned n_heads, int64_t ctx) {
-  unsigned S = attn_split_count(n_heads, ctx);
-  unsigned chunk = (static_cast<unsigned>(ctx) + S - 1) / S;
-  chunk = (chunk + 3u) & ~3u;
-  return chunk ? chunk : 4u;
+  return static_cast<unsigned>(gpu::policy::split_chunk(
+      ctx, attn_split_count(n_heads, ctx), attn_split_rule));
 }
 
 // Split-KV partials scratch: ONE buffer laid out pm[H*S] | pl[H*S] | pacc[H*S*D]

include/gpu.h

@@ -20,18 +20,20 @@
 //                          gpu_ops.h forwards to the ones that exist; an op a
 //                          backend does not declare has no stub to keep in step
 //               traits     what the launch policy may assume of its kernels
-//               caps       what a model may assume (model_path, graph_capture,
-//                          row_gemv, bf16_gemm), plus the graph-capture plumbing
+//               caps       what a model may assume (graph_capture, row_gemv,
+//                          bf16_gemm), plus the graph-capture plumbing
 //                          — graph_available / capture_begin / capture_end /
 //                          graph_launch / graph_destroy / upload_u32 /
 //                          attn_dpos_partials_bytes — as no-ops where absent
 //
 // An op answers false when the backend has no kernel for it, and the evaluator
 // falls back to the CPU. The model path (what a decoder runs on raw device
-// buffers between its GEMVs and attention: kv_cache.h, bench/models) has no CPU
-// fallback, so a model checks the return and keeps to the array ops where it
-// is false. gpu::census(kernel) counts launches, which is how a test tells a
-// kernel that ran from an op that quietly fell back.
+// buffers between its GEMVs and attention: kv_cache.h, bench/models) falls
+// back instead to gpu_ops.h's generic compositions of the tier-0 ops, so it
+// runs on every backend in f32; an op there still answers false for an
+// operand its backend cannot read (bf16, int4), and a model keeps to the
+// array ops where it does. gpu::census(kernel) counts launches, which is how
+// a test tells a kernel that ran from an op that quietly fell back.
 //
 // One backend is selected below, by the gate its header is written under:
 // webgpu.h under TENSORLIB_WEBGPU && __EMSCRIPTEN__, cuda.h under TENSORLIB_CUDA