diff --git a/bench/models/check_qwen.cpp b/bench/models/check_qwen.cpp index c06607f..47d5843 100644 --- a/bench/models/check_qwen.cpp +++ b/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 && diff --git a/docs/backends.md b/docs/backends.md index 269ce4c..79c8805 100644 --- a/docs/backends.md +++ b/docs/backends.md @@ -123,23 +123,45 @@ template 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::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 diff --git a/include/cuda.h b/include/cuda.h index daf3600..59123a7 100644 --- a/include/cuda.h +++ b/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(traits::fill_groups); + if (tiles >= fill) return k; const long slabs = k / t.bk, min_slabs = kWaveMinK / t.bk; - const long full = std::max(1, std::min(kFillBlocks / tiles, slabs / min_slabs)); + const long full = std::max(1, std::min(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,6 +1501,8 @@ 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(); @@ -1509,17 +1510,12 @@ inline bool gemv_run_(CUfunction f, float* pa, float* pB, float* py, 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(bx) < target && uk >= 512) { - unsigned g = static_cast((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(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(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((target + n_heads - 1) / n_heads); - unsigned max_s = static_cast(ctx / 128); // >=128 keys/split - if (want > max_s) want = max_s; - return want > 1 ? want : 1; + return static_cast( + 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(ctx) + S - 1) / S; - chunk = (chunk + 3u) & ~3u; - return chunk ? chunk : 4u; + return static_cast(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] diff --git a/include/gpu.h b/include/gpu.h index c011e2a..7c884a7 100644 --- a/include/gpu.h +++ b/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 diff --git a/include/gpu_abi.h b/include/gpu_abi.h index 41036c9..ea83003 100644 --- a/include/gpu_abi.h +++ b/include/gpu_abi.h @@ -18,6 +18,9 @@ #include #include +#include + +#include "profile.h" namespace tl { namespace gpu { @@ -250,6 +253,18 @@ struct merge_heads_params { uint32_t T, H, D; }; +// The launch record. Every kernel launch on every backend is one row under +// tl::profile, made here from the kernel's name: a backend's launch primitive +// (the one place its launches funnel through, shared ops and own alike) calls +// this once per launch and stamps the row with a device time where it has one +// (profile::detail::device_time). Under TL_PROFILE=1 the first launch also +// starts the profile — a decoder on the model path never reaches the +// evaluator, so this is where its profile begins. +inline profile::row* launched(std::string_view kernel) { + profile::detail::env_autostart(); + return profile::active() ? profile::detail::launch(kernel) : nullptr; +} + // Launch policy: the shapes ops launch in, in one place. Host code shared by // every backend; what differs between devices will come in as traits. namespace policy { @@ -310,6 +325,38 @@ inline grid cells_2d(int64_t rows, int64_t cols) { static_cast((rows + 7) / 8), 1, 32, 8, 1, 0}; } +// Splitting a reduction over more groups. `groups` groups each walk k units +// (a GEMV's k, attention's keys, a GEMM tile's k); when they alone leave the +// device short of busy groups, each is cut into parts that run as extra groups +// and are combined after. The rule is the kernel's side of it — what its part +// must be a multiple of, and how short a part or a k is still worth the +// combine — and the target is the device's (traits::fill_groups, or a multiple +// of it for a kernel whose groups are small). A backend keeps to this shape +// of decision rather than its own arithmetic, so the split every backend's +// GEMV and attention make is one function of the device. +struct split_rule { + int64_t target; // groups that keep the device busy + int64_t min_k; // a k shorter than this is not split at all + int64_t min_part; // a part is at least this many units + int64_t granule; // and a multiple of this (the kernel's step) +}; + +// The parts the k units are cut into; 1 when not split. Monotone in k. +inline int64_t split_parts(int64_t groups, int64_t k, const split_rule& r) { + if (groups <= 0 || groups >= r.target || k < r.min_k) return 1; + int64_t parts = (r.target + groups - 1) / groups; + if (r.min_part > 0 && parts > k / r.min_part) parts = k / r.min_part; + return parts > 1 ? parts : 1; +} + +// The units each of `parts` parts takes, rounded up to the granule. Rounding +// can leave fewer parts than asked, so the count to launch is (k + chunk - 1) +// / chunk. +inline int64_t split_chunk(int64_t k, int64_t parts, const split_rule& r) { + const int64_t chunk = ((k + parts - 1) / parts + r.granule - 1) / r.granule * r.granule; + return chunk > 0 ? chunk : r.granule; +} + } // namespace policy } // namespace gpu diff --git a/include/gpu_host.h b/include/gpu_host.h index 6c9a5e9..1ee1a6a 100644 --- a/include/gpu_host.h +++ b/include/gpu_host.h @@ -355,10 +355,9 @@ inline bool merge_heads(const gpu::arg* v, const void* params) { // ---- launch: the kernel table and the kernels in one switch. An id with no // case declines, and the op above falls back to the CPU. -inline bool dispatch(kop k, const gpu::arg* v, size_t /*n*/, const void* params, - size_t /*params_bytes*/, const gpu::grid& g) { +inline bool kernel_(kop k, const gpu::arg* v, const void* params, + const gpu::grid& g) { namespace d = detail_; - d::pending_ = true; // declining leaves it set too: a flush then waits on nothing switch (k) { case kop::add: return d::binary(v, params, [](float a, float b) { return a + b; }); case kop::sub: return d::binary(v, params, [](float a, float b) { return a - b; }); @@ -430,6 +429,15 @@ inline bool dispatch(kop k, const gpu::arg* v, size_t /*n*/, const void* params, } } +// A kernel that ran is a row under tl::profile, as on any backend. +inline bool dispatch(kop k, const gpu::arg* v, size_t /*n*/, const void* params, + size_t /*params_bytes*/, const gpu::grid& g) { + detail_::pending_ = true; // declining leaves it set too: a flush then waits on nothing + if (!kernel_(k, v, params, g)) return false; + gpu::launched(kop_name(k)); + return true; +} + // ---- the ops this backend runs its own way: the ones with no single-kernel // form in gpu_ops.h that the conformance test expects of every backend. Each // is the definition of its op, as a loop. @@ -451,6 +459,7 @@ struct own { C[i * n + j] = static_cast(acc) * scale + offset; } } + gpu::launched("gemm"); return true; } @@ -468,6 +477,7 @@ struct own { const int64_t outer = i / (a_axis * inner), rest = i % (a_axis * inner); O[outer * o_axis * inner + before * inner + rest] = A[i]; } + gpu::launched("pad"); return true; } @@ -481,6 +491,7 @@ struct own { const int64_t row = static_cast(I[i] + 0.5f); for (int64_t c = 0; c < row_size; c++) O[row * row_size + c] += V[i * row_size + c]; } + gpu::launched("index_add"); return true; } @@ -491,6 +502,7 @@ struct own { float* O = view(out); std::fill(O, O + n * size, 0.0f); for (int64_t i = 0; i < n; i++) O[i * size + static_cast(I[i] + 0.5f)] = V[i]; + gpu::launched("scatter_to_axis"); return true; } @@ -504,15 +516,10 @@ struct own { // ---- what the shared layer may assume. struct traits { static constexpr bool cells_2d = false; // a cell is read from a flat index - // Whether this backend records its own launches under tl::profile (else the - // shared layer does), and whether each carries a device time. - static constexpr bool profiles_launches = false; - static constexpr bool times_launches = false; + static constexpr bool times_launches = false; // a row is counted, not timed + static constexpr int64_t fill_groups = 1; // a loop: a split is only more passes }; struct caps { - // The decoder's single-kernel ops are here, but not attention, rope or the - // GEMVs a model also needs, so a model keeps to the array ops. - static constexpr bool model_path = false; static constexpr bool graph_capture = false; static constexpr bool row_gemv = true; static constexpr bool bf16_gemm = false; diff --git a/include/gpu_null.h b/include/gpu_null.h index 5ab255a..59d3266 100644 --- a/include/gpu_null.h +++ b/include/gpu_null.h @@ -34,7 +34,8 @@ inline void upload(void*, const float*, int64_t) {} // stage host floats in // ---- launch: the one way a shared op (gpu_ops.h) runs a kernel. View i is the // kernel's i-th buffer at its byte offset, `params` a block of 4-byte fields // in the kernel's argument order (gpu_abi.h). False for a kernel id this -// backend has no kernel for. +// backend has no kernel for. Every launch a backend makes, here or in an own +// op, is one gpu::launched(kernel name) call: its row under tl::profile. inline bool dispatch(gpu::kop, const gpu::arg*, size_t, const void* /*params*/, size_t /*params_bytes*/, const gpu::grid&) { return false; @@ -47,16 +48,17 @@ struct own {}; // ---- what the shared launch policy may assume of this backend's kernels. struct traits { static constexpr bool cells_2d = false; - // Whether this backend records its own launches under tl::profile (else the - // shared layer does), and whether each carries a device time. - static constexpr bool profiles_launches = false; + // Whether a launch's tl::profile row (gpu::launched, called from the + // backend's launch primitive) carries a device time. static constexpr bool times_launches = false; + // Groups that keep the device busy: the target a split (policy::split_parts) + // measures its grid against. + static constexpr int64_t fill_groups = 1; }; // ---- what a model may assume, and the graph-capture plumbing behind // caps::graph_capture. struct caps { - static constexpr bool model_path = false; static constexpr bool graph_capture = false; static constexpr bool row_gemv = false; static constexpr bool bf16_gemm = false; diff --git a/include/gpu_ops.h b/include/gpu_ops.h index b94af6d..98dd0c9 100644 --- a/include/gpu_ops.h +++ b/include/gpu_ops.h @@ -9,12 +9,12 @@ // Included by gpu.h, after the backend is selected. #include +#include #include #include #include #include "gpu_abi.h" -#include "profile.h" namespace tl { namespace gpu { @@ -27,18 +27,8 @@ namespace gpu { namespace detail { inline std::array census_counts{}; inline uint64_t census_ops_run = 0; -// A launch the backend did not record under tl::profile itself is recorded -// here, by name, so a backend is profiled from its first kernel; one that -// stamps its launches with a device time says so (traits::profiles_launches) -// and records its own. -inline void profile_launch(const char* name) { - if (!traits::profiles_launches && profile::active()) profile::detail::launch(name); -} -inline bool ran(const char* op, bool ok) { - if (ok) { - census_ops_run++; - profile_launch(op); - } +inline bool ran(bool ok) { + if (ok) census_ops_run++; return ok; } } // namespace detail @@ -78,7 +68,6 @@ inline bool launch(kop k, std::initializer_list args, const P& params, } detail::census_counts[static_cast(k)]++; detail::census_ops_run++; - detail::profile_launch(kop_name(k)); return true; } @@ -237,126 +226,6 @@ inline bool adam_step(span p, span m, span v, span g, int64_t n, float beta1, policy::flat(n)); } -// ---- the model path: what a decoder runs on raw device buffers between its -// GEMVs and attention. No CPU fallback sits under these, so a model checks the -// return and keeps to the array ops where it is false. - -// out = x * rsqrt(mean(x^2) + eps) * w per row of [rows, n]. out may alias x. -inline bool rmsnorm(span x, span w, span o, int64_t n, float eps, - int64_t rows = 1) { - if (n <= 0 || rows <= 0) return false; - return launch(kop::rmsnorm_, {in(x), in(w), out(o)}, - rmsnorm_params{static_cast(n), eps}, - policy::one_group_per_row(rows)); -} - -// xout = x + delta and hout = rmsnorm(xout) * w: a residual add folded into -// the norm that follows it. xout may alias x. -inline bool rmsnorm_res(span x, span delta, span w, span xout, span hout, - int64_t n, float eps, int64_t rows = 1) { - if (n <= 0 || rows <= 0) return false; - return launch(kop::add_rmsnorm_, - {in(x), in(delta), in(w), out(xout), out(hout)}, - rmsnorm_params{static_cast(n), eps}, - policy::one_group_per_row(rows)); -} - -// out[rows, ff] = silu(gate) * up out of the fused gate|up buffer [rows, 2ff]. -inline bool swiglu(span gu, span o, int64_t ff, int64_t rows = 1) { - if (ff <= 0 || rows <= 0) return false; - return launch(kop::swiglu_, {in(gu), out(o)}, - swiglu_params{static_cast(ff)}, - policy::flat_rows(ff, rows)); -} - -// y[1,N] = a[1,K] . W[N,K]^T with the bf16 weight row-major (GGML-native): one -// group an output row. Requires k % 8 == 0. -inline bool gemv_bf16_row(span a, span W, span y, int64_t n, int64_t k) { - if (n <= 0 || k <= 0 || k % 8 != 0) return false; - return launch(kop::gemv_bf16_row_, {in(a), in(W), out(y)}, - gemv_row_params{static_cast(n), static_cast(k)}, - policy::row_reduce(n, k)); -} - -// The same over int4 weights: qw is [N][K/8] packed words and scales -// [N][K/group] floats, two views that may share a buffer. K % group == 0 and -// group % 8 == 0. -inline bool gemv_q4(span a, span qw, span scales, span y, int64_t N, int64_t K, - int64_t group) { - if (N <= 0 || K <= 0 || group <= 0 || K % group != 0 || group % 8 != 0) { - return false; - } - return launch(kop::gemv_q4_, {in(a), in(qw), in(scales), out(y)}, - gemv_q4_params{static_cast(N), static_cast(K), - static_cast(group)}, - policy::row_reduce(N, K)); -} - -// One decode step's k, v (each [n_kv_heads, D]) into row `pos` of a -// [n_kv_heads, kv_max, D] cache, f32 or bf16. -inline bool kv_append(span Kc, span Vc, span k_new, span v_new, int64_t pos, - int64_t kv_max, int64_t n_kv_heads, int64_t D, - bool kv_bf16 = false) { - if ((D != 64 && D != 128) || n_kv_heads <= 0) return false; - return launch(kv_bf16 ? kop::kv_append_bf16_ : kop::kv_append_, - {out(Kc), out(Vc), in(k_new), in(v_new)}, - kv_append_params{static_cast(pos), - static_cast(kv_max * D)}, - policy::per_head(n_kv_heads, 1, D)); -} - -// A prefill's k, v (each [n_kv_heads, T, D]) into cache rows [pos0, pos0 + T). -inline bool kv_fill(span Kc, span Vc, span K, span V, int64_t T, int64_t kv_max, - int64_t n_kv_heads, int64_t D, bool kv_bf16 = false, - int64_t pos0 = 0) { - if ((D != 64 && D != 128) || n_kv_heads <= 0 || T <= 0) return false; - return launch(kv_bf16 ? kop::kv_fill_bf16_ : kop::kv_fill_, - {out(Kc), out(Vc), in(K), in(V)}, - kv_fill_params{static_cast(T), - static_cast(kv_max * D), - static_cast(pos0)}, - policy::per_head(n_kv_heads, T, D)); -} - -// Head-major [H, T, D] -> token-major [T, H*D]: split_heads' inverse. -inline bool merge_heads(span src, span dst, int64_t T, int64_t H, int64_t D) { - if (T <= 0 || H <= 0 || D <= 0) return false; - return launch(kop::merge_heads_, {in(src), out(dst)}, - merge_heads_params{static_cast(T), - static_cast(H), - static_cast(D)}, - policy::per_head(H, T, D)); -} - -// Token-major [T, ld] -> head-major [H, T, D] from column block `off`, adding -// the optional per-head bias [H, D] (a null view: none). Backend-own: the -// kernels disagree on how "no bias" is said. -TL_GPU_DETECT_OWN(split_heads) -template -inline bool split_heads(span src, span bias, span dst, int64_t T, int64_t ld, - int64_t off, int64_t H, int64_t D) { - if (T <= 0 || H <= 0 || D <= 0) return false; - if constexpr (detail::owns_split_heads::value) { - return detail::ran("split_heads", Own::split_heads(src, bias, dst, T, ld, off, H, D)); - } else { - return false; - } -} - -// The argmax of a length-n vector, the smallest index on ties: greedy decoding -// reads one int back rather than the logits. Drains the queue. Backend-own: -// the result's staging buffer and its read-back are the backend's. -TL_GPU_DETECT_OWN(argmax) -template -inline bool argmax(span a, int64_t n, int64_t* out_idx) { - if (!a || n <= 0 || !out_idx) return false; - if constexpr (detail::owns_argmax::value) { - return detail::ran("argmax", Own::argmax(a, n, out_idx)); - } else { - return false; - } -} - // tanh / sin / cos: unary's shape under their own vocabulary (gpu_abi.h). inline bool unary_ext(unary_ext_op op, span a, span o, int64_t n, float scale, float offset) { @@ -384,7 +253,7 @@ inline bool binary_bcast_nd(kop op, span a, const int64_t* a_strides, span b, const int64_t* out_shape, int rank, int64_t n, float scale, float offset) { if constexpr (detail::owns_binary_bcast_nd::value) { - return detail::ran("binary_bcast_nd", Own::binary_bcast_nd(op, a, a_strides, b, b_strides, o, out_shape, rank, n, scale, offset)); + return detail::ran(Own::binary_bcast_nd(op, a, a_strides, b, b_strides, o, out_shape, rank, n, scale, offset)); } else { return false; } @@ -397,7 +266,7 @@ inline bool where_nd(span cond, const int64_t* c_strides, span a, const int64_t* a_strides, span b, const int64_t* b_strides, span o, const int64_t* out_shape, int rank, int64_t n) { if constexpr (detail::owns_where_nd::value) { - return detail::ran("where_nd", Own::where_nd(cond, c_strides, a, a_strides, b, b_strides, o, out_shape, rank, n)); + return detail::ran(Own::where_nd(cond, c_strides, a, a_strides, b, b_strides, o, out_shape, rank, n)); } else { return false; } @@ -409,7 +278,7 @@ template inline bool copy_nd(span a, const int64_t* a_strides, span o, const int64_t* out_shape, int rank, int64_t n) { if constexpr (detail::owns_copy_nd::value) { - return detail::ran("copy_nd", Own::copy_nd(a, a_strides, o, out_shape, rank, n)); + return detail::ran(Own::copy_nd(a, a_strides, o, out_shape, rank, n)); } else { return false; } @@ -423,7 +292,7 @@ inline bool sum_to(span a, const int64_t* a_shape, const int64_t* a_strides, const int64_t* acc, int rank, int64_t out_n, int64_t reduced_n, span o) { if constexpr (detail::owns_sum_to::value) { - return detail::ran("sum_to", Own::sum_to(a, a_shape, a_strides, acc, rank, out_n, reduced_n, o)); + return detail::ran(Own::sum_to(a, a_shape, a_strides, acc, rank, out_n, reduced_n, o)); } else { return false; } @@ -437,7 +306,7 @@ inline bool pad(span a, span o, const int64_t* a_shape, const int64_t* out_shape, int rank, int axis, int64_t before, int64_t n, int64_t out_n) { if constexpr (detail::owns_pad::value) { - return detail::ran("pad", Own::pad(a, o, a_shape, out_shape, rank, axis, before, n, out_n)); + return detail::ran(Own::pad(a, o, a_shape, out_shape, rank, axis, before, n, out_n)); } else { return false; } @@ -450,7 +319,7 @@ inline bool fold(span a, span o, const int64_t* a_shape, const int64_t* out_shape, int rank, int axis, int64_t step, int64_t n, int64_t out_n) { if constexpr (detail::owns_fold::value) { - return detail::ran("fold", Own::fold(a, o, a_shape, out_shape, rank, axis, step, n, out_n)); + return detail::ran(Own::fold(a, o, a_shape, out_shape, rank, axis, step, n, out_n)); } else { return false; } @@ -463,7 +332,7 @@ inline bool concat_part(span a, span o, const int64_t* a_shape, const int64_t* out_shape, int rank, int axis, int64_t before, int64_t n) { if constexpr (detail::owns_concat_part::value) { - return detail::ran("concat_part", Own::concat_part(a, o, a_shape, out_shape, rank, axis, before, n)); + return detail::ran(Own::concat_part(a, o, a_shape, out_shape, rank, axis, before, n)); } else { return false; } @@ -475,7 +344,7 @@ template inline bool index_add(span idx, span values, span o, int64_t row_size, int64_t k, int64_t out_n) { if constexpr (detail::owns_index_add::value) { - return detail::ran("index_add", Own::index_add(idx, values, o, row_size, k, out_n)); + return detail::ran(Own::index_add(idx, values, o, row_size, k, out_n)); } else { return false; } @@ -487,7 +356,7 @@ template inline bool scatter_to_axis(span idx, span values, span o, int64_t n, int64_t size) { if constexpr (detail::owns_scatter_to_axis::value) { - return detail::ran("scatter_to_axis", Own::scatter_to_axis(idx, values, o, n, size)); + return detail::ran(Own::scatter_to_axis(idx, values, o, n, size)); } else { return false; } @@ -501,7 +370,7 @@ inline bool gemm(span a, int64_t lda, bool ta, span b, int64_t ldb, bool tb, span o, int64_t m, int64_t n, int64_t k, float scale, float offset) { if constexpr (detail::owns_gemm::value) { - return detail::ran("gemm", Own::gemm(a, lda, ta, b, ldb, tb, o, m, n, k, scale, offset)); + return detail::ran(Own::gemm(a, lda, ta, b, ldb, tb, o, m, n, k, scale, offset)); } else { return false; } @@ -516,7 +385,7 @@ inline bool gemm_batched(span a, int64_t lda, bool ta, int64_t sa, span b, int64_t n, int64_t k, int64_t batch, float scale, float offset, span bias = {}) { if constexpr (detail::owns_gemm_batched::value) { - return detail::ran("gemm_batched", Own::gemm_batched(a, lda, ta, sa, b, ldb, tb, sb, o, m, n, k, batch, scale, offset, bias)); + return detail::ran(Own::gemm_batched(a, lda, ta, sa, b, ldb, tb, sb, o, m, n, k, batch, scale, offset, bias)); } else { return false; } @@ -529,20 +398,7 @@ inline bool gemm_bias(span a, int64_t lda, bool ta, span b, int64_t ldb, bool tb, span bias, span o, int64_t m, int64_t n, int64_t k, float scale, float offset) { if constexpr (detail::owns_gemm_bias::value) { - return detail::ran("gemm_bias", Own::gemm_bias(a, lda, ta, b, ldb, tb, bias, o, m, n, k, scale, offset)); - } else { - return false; - } -} - -// Rotary embedding over [rows, T, D] at position `pos`, adding the optional -// per-row bias first. -TL_GPU_DETECT_OWN(rope) -template -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::value) { - return detail::ran("rope", Own::rope(x, o, rows, T, D, pos, base, bias)); + return detail::ran(Own::gemm_bias(a, lda, ta, b, ldb, tb, bias, o, m, n, k, scale, offset)); } else { return false; } @@ -556,12 +412,410 @@ inline bool layer_norm_bwd(span x, span g, span dy, span dx, span dg, span db, int64_t cols, int64_t per_chunk, int64_t chunks, float eps) { if constexpr (detail::owns_layer_norm_bwd::value) { - return detail::ran("layer_norm_bwd", Own::layer_norm_bwd(x, g, dy, dx, dg, db, stats, partials, rows, cols, per_chunk, chunks, eps)); + return detail::ran(Own::layer_norm_bwd(x, g, dy, dx, dg, db, stats, partials, rows, cols, per_chunk, chunks, eps)); } else { return false; } } +// ---- tier 1 by composition. Each fused op below is a kernel a backend may +// have; where it does not — the shared launch declines, or `own` has no such +// member — the same result comes from the tier-0 ops above, written once here. +// Several launches and a scratch buffer or two rather than one kernel, so a +// backend that cares for the decode loop writes the kernel; a new backend has +// the whole model path from its tier 0 alone. f32 only: a bf16 or int4 operand +// needs a reader the tier does not have, so those ops still decline. Nothing +// here names a backend or the array layer: spans, tier-0 ops, and the device +// core's alloc / release / cpu_barrier / sync_to_host. +namespace generic { + +// A device buffer for the extent of one composition. With host_fill the host +// writes `contents` before the first kernel reads it (a table, a mask), and +// the backend uploads it where memory is not unified. +struct scratch { + void* buf = nullptr; + float* contents = nullptr; + int64_t bytes = 0; + explicit scratch(int64_t floats, bool host_fill = false) : bytes(floats * 4) { + buf = alloc(bytes, &contents, host_fill); + } + ~scratch() { + if (buf) release(buf, bytes, contents); + } + scratch(const scratch&) = delete; + scratch& operator=(const scratch&) = delete; + explicit operator bool() const { return buf != nullptr; } + operator span() const { return {buf, 0}; } +}; + +// A strided rank-2 view copied into a contiguous [m, n]: the broadcast add of +// a zero, which `zero` holds (one float the host filled). +inline bool copy_2d(span a, int64_t ars, int64_t acs, span zero, span o, + int64_t m, int64_t n) { + return binary_bcast(kop::badd, a, ars, acs, zero, 0, 0, o, m, n, 1.0f, 0.0f); +} + +inline bool rmsnorm(span x, span w, span o, int64_t n, float eps, + int64_t rows) { + scratch sq(rows * n), ms(rows), r(rows); + if (!sq || !ms || !r) return false; + // mean(x^2) + eps in the reduction's epilogue, then its inverse root, then + // two broadcasts: by the row's scale and by w. (No launch here writes a + // buffer it reads — a backend may forbid that — and o, which may alias x, + // is written last, after x's last read.) + return binary(kop::mul, x, x, sq, rows * n, 1.0f, 0.0f) && + row_op(kop::row_sum, sq, ms, rows, n, 1.0f / static_cast(n), eps) && + scalar_binary(scalar_op::pow, ms, r, rows, -0.5f, 1.0f, 0.0f) && + binary_bcast(kop::bmul, x, n, 1, r, 1, 0, sq, rows, n, 1.0f, 0.0f) && + binary_bcast(kop::bmul, sq, n, 1, w, 0, 1, o, rows, n, 1.0f, 0.0f); +} + +inline bool rmsnorm_res(span x, span delta, span w, span xout, span hout, + int64_t n, float eps, int64_t rows) { + if (xout.buf != x.buf) { + return binary(kop::add, x, delta, xout, rows * n, 1.0f, 0.0f) && + rmsnorm(xout, w, hout, n, eps, rows); + } + // xout aliases x: the sum goes through a buffer of this layer's own. + scratch sum(rows * n); + return sum && binary(kop::add, x, delta, sum, rows * n, 1.0f, 0.0f) && + unary(kop::affine, sum, xout, rows * n, 1.0f, 0.0f) && + rmsnorm(sum, w, hout, n, eps, rows); +} + +inline bool swiglu(span gu, span o, int64_t ff, int64_t rows) { + scratch zero(1, true), g(rows * ff), s(rows * ff), silu(rows * ff); + if (!zero || !g || !s || !silu) return false; + zero.contents[0] = 0.0f; + // gate out of the fused [rows, 2ff] buffer, silu(gate) = gate * sigmoid(gate), + // times up read in place. + return copy_2d(gu, 2 * ff, 1, zero, g, rows, ff) && + unary(kop::sigmoid, g, s, rows * ff, 1.0f, 0.0f) && + binary(kop::mul, g, s, silu, rows * ff, 1.0f, 0.0f) && + binary_bcast(kop::bmul, silu, ff, 1, gu.at(ff * 4), 2 * ff, 1, o, rows, + ff, 1.0f, 0.0f); +} + +// A GEMV is a GEMM with one row. +inline bool gemv_f32(span a, span B, span y, int64_t n, int64_t k) { + return gemm(a, k, false, B, n, false, y, 1, n, k, 1.0f, 0.0f); +} + +// A row into the cache is a contiguous copy per head. +inline bool kv_append(span Kc, span Vc, span k_new, span v_new, int64_t pos, + int64_t kv_max, int64_t n_kv_heads, int64_t D) { + for (int64_t h = 0; h < n_kv_heads; h++) { + const int64_t at = (h * kv_max + pos) * D * 4, from = h * D * 4; + if (!unary(kop::affine, k_new.at(from), Kc.at(at), D, 1.0f, 0.0f) || + !unary(kop::affine, v_new.at(from), Vc.at(at), D, 1.0f, 0.0f)) { + return false; + } + } + return true; +} + +inline bool kv_fill(span Kc, span Vc, span K, span V, int64_t T, int64_t kv_max, + int64_t n_kv_heads, int64_t D, int64_t pos0) { + for (int64_t h = 0; h < n_kv_heads; h++) { + const int64_t at = (h * kv_max + pos0) * D * 4, from = h * T * D * 4; + if (!unary(kop::affine, K.at(from), Kc.at(at), T * D, 1.0f, 0.0f) || + !unary(kop::affine, V.at(from), Vc.at(at), T * D, 1.0f, 0.0f)) { + return false; + } + } + return true; +} + +// [H, T, D] -> [T, H*D]: token t's row is the [H, D] gather of the heads' +// row t, one strided copy a token. +inline bool merge_heads(span src, span dst, int64_t T, int64_t H, int64_t D) { + scratch zero(1, true); + if (!zero) return false; + zero.contents[0] = 0.0f; + for (int64_t t = 0; t < T; t++) { + if (!copy_2d(src.at(t * D * 4), T * D, 1, zero, dst.at(t * H * D * 4), H, D)) { + return false; + } + } + return true; +} + +// [T, ld] -> [H, T, D] from column block `off`, plus the head's bias row: one +// broadcast add a head (of zero, when there is no bias). +inline bool split_heads(span src, span bias, span dst, int64_t T, int64_t ld, + int64_t off, int64_t H, int64_t D) { + scratch zero(1, true); + if (!zero) return false; + zero.contents[0] = 0.0f; + for (int64_t h = 0; h < H; h++) { + const span b = bias ? bias.at(h * D * 4) : span(zero); + if (!binary_bcast(kop::badd, src.at((off + h * D) * 4), ld, 1, b, 0, + bias ? 1 : 0, dst.at(h * T * D * 4), T, D, 1.0f, 0.0f)) { + return false; + } + } + return true; +} + +// Rotary embedding over [rows, D], rows = H*T, row r at position pos + r % T: +// out = x * C + swap(x) * S, with C and S the [T, D] tables of cos and of +// (-sin | +sin) over the two halves, and swap(x) the halves exchanged (a row +// gather of x seen as [2T, D/2]). Four launches a head over shared tables. +inline bool rope(span x, span o, int64_t rows, int64_t T, int64_t D, + int64_t pos, float base, span bias) { + if (rows <= 0 || T <= 0 || D <= 0 || D % 2 || rows % T) return false; + const int64_t half = D / 2, H = rows / T; + scratch C(T * D, true), S(T * D, true), idx(2 * T, true), + xb(bias ? rows * D : 1), sw(T * D), xc(T * D), xs(T * D); + if (!C || !S || !idx || (bias && !xb) || !sw || !xc || !xs) return false; + for (int64_t t = 0; t < T; t++) { + const double position = static_cast(pos + t); + for (int64_t j = 0; j < half; j++) { + const double ang = position * std::pow(static_cast(base), + -2.0 * static_cast(j) / D); + const float c = static_cast(std::cos(ang)), + s = static_cast(std::sin(ang)); + C.contents[t * D + j] = C.contents[t * D + j + half] = c; + S.contents[t * D + j] = -s; + S.contents[t * D + j + half] = s; + } + idx.contents[2 * t] = static_cast(2 * t + 1); + idx.contents[2 * t + 1] = static_cast(2 * t); + } + span src = x; + if (bias) { + if (!binary(kop::add, x, bias, xb, rows * D, 1.0f, 0.0f)) return false; + src = xb; + } + for (int64_t h = 0; h < H; h++) { + const span xh = src.at(h * T * D * 4), oh = o.at(h * T * D * 4); + if (!index_select(xh, idx, sw, half, 2 * T) || + !binary(kop::mul, xh, C, xc, T * D, 1.0f, 0.0f) || + !binary(kop::mul, sw, S, xs, T * D, 1.0f, 0.0f) || + !binary(kop::add, xc, xs, oh, T * D, 1.0f, 0.0f)) { + return false; + } + } + return true; +} + +// One decode step, a head at a time: scores = q . K^T over the cached +// prefix, softmax, times V. An f32 cache only. +inline bool attn_decode(span q, span K, span V, span o, int64_t n_q_heads, + int64_t n_kv_heads, int64_t ctx, int64_t kv_max, + int64_t D, float scale, bool kv_bf16) { + if (kv_bf16 || n_kv_heads <= 0 || n_q_heads % n_kv_heads || ctx <= 0 || + D <= 0) { + return false; + } + scratch s(ctx), p(ctx); + if (!s || !p) return false; + const int64_t group = n_q_heads / n_kv_heads; + for (int64_t h = 0; h < n_q_heads; h++) { + const int64_t kv = (h / group) * kv_max * D * 4; + if (!gemm(q.at(h * D * 4), D, false, K.at(kv), D, true, s, 1, ctx, D, scale, + 0.0f) || + !row_op(kop::softmax, s, p, 1, ctx, 1.0f, 0.0f) || + !gemm(p, ctx, false, V.at(kv), D, false, o.at(h * D * 4), 1, D, ctx, + 1.0f, 0.0f)) { + return false; + } + } + return true; +} + +// The causal prefill, a head at a time: scores [T, ctx] over the cache rows +// [0, pos0 + T), a mask that closes the keys past each query's position, +// softmax by row, times V. +inline bool attn_prefill(span q, span K, span V, span o, int64_t n_q_heads, + int64_t n_kv_heads, int64_t T, int64_t kv_max, + int64_t D, float scale, bool kv_bf16, int64_t pos0) { + if (kv_bf16 || n_kv_heads <= 0 || n_q_heads % n_kv_heads || T <= 0 || + D <= 0 || pos0 < 0 || pos0 + T > kv_max) { + return false; + } + const int64_t ctx = pos0 + T; + scratch mask(T * ctx, true), s(T * ctx), p(T * ctx); + if (!mask || !s || !p) return false; + for (int64_t t = 0; t < T; t++) { + for (int64_t j = 0; j < ctx; j++) { + mask.contents[t * ctx + j] = j <= pos0 + t ? 0.0f : -1e30f; + } + } + const int64_t group = n_q_heads / n_kv_heads; + for (int64_t h = 0; h < n_q_heads; h++) { + const int64_t kv = (h / group) * kv_max * D * 4, qh = h * T * D * 4; + if (!gemm(q.at(qh), D, false, K.at(kv), D, true, s, T, ctx, D, scale, 0.0f) || + !binary(kop::add, s, mask, p, T * ctx, 1.0f, 0.0f) || + !row_op(kop::softmax, p, s, T, ctx, 1.0f, 0.0f) || + !gemm(s, ctx, false, V.at(kv), D, false, o.at(qh), T, D, ctx, 1.0f, + 0.0f)) { + return false; + } + } + return true; +} + +// The vector copied into a buffer of this layer's own, brought to the host +// and scanned: the one host round trip the model path makes. +inline bool argmax(span a, int64_t n, int64_t* out_idx) { + scratch c(n); + if (!c || !unary(kop::affine, a, c, n, 1.0f, 0.0f)) return false; + cpu_barrier(); + sync_to_host(c.buf, false); + int64_t best = 0; + for (int64_t i = 1; i < n; i++) { + if (c.contents[i] > c.contents[best]) best = i; + } + *out_idx = best; + return true; +} + +} // namespace generic + +// ---- the model path: what a decoder runs on raw device buffers between its +// GEMVs and attention. Each is the backend's kernel where it has one and the +// generic composition where it does not, so every backend has the f32 model +// path; an op still answers false for an operand its backend cannot read (a +// bf16 cache or weight, int4 weights), and a model keeps to the array ops +// there. + +// out = x * rsqrt(mean(x^2) + eps) * w per row of [rows, n]. out may alias x. +inline bool rmsnorm(span x, span w, span o, int64_t n, float eps, + int64_t rows = 1) { + if (n <= 0 || rows <= 0) return false; + return launch(kop::rmsnorm_, {in(x), in(w), out(o)}, + rmsnorm_params{static_cast(n), eps}, + policy::one_group_per_row(rows)) || + generic::rmsnorm(x, w, o, n, eps, rows); +} + +// xout = x + delta and hout = rmsnorm(xout) * w: a residual add folded into +// the norm that follows it. xout may alias x. +inline bool rmsnorm_res(span x, span delta, span w, span xout, span hout, + int64_t n, float eps, int64_t rows = 1) { + if (n <= 0 || rows <= 0) return false; + return launch(kop::add_rmsnorm_, + {in(x), in(delta), in(w), out(xout), out(hout)}, + rmsnorm_params{static_cast(n), eps}, + policy::one_group_per_row(rows)) || + generic::rmsnorm_res(x, delta, w, xout, hout, n, eps, rows); +} + +// out[rows, ff] = silu(gate) * up out of the fused gate|up buffer [rows, 2ff]. +inline bool swiglu(span gu, span o, int64_t ff, int64_t rows = 1) { + if (ff <= 0 || rows <= 0) return false; + return launch(kop::swiglu_, {in(gu), out(o)}, + swiglu_params{static_cast(ff)}, + policy::flat_rows(ff, rows)) || + generic::swiglu(gu, o, ff, rows); +} + +// y[1,N] = a[1,K] . W[N,K]^T with the bf16 weight row-major (GGML-native): one +// group an output row. Requires k % 8 == 0. +inline bool gemv_bf16_row(span a, span W, span y, int64_t n, int64_t k) { + if (n <= 0 || k <= 0 || k % 8 != 0) return false; + return launch(kop::gemv_bf16_row_, {in(a), in(W), out(y)}, + gemv_row_params{static_cast(n), static_cast(k)}, + policy::row_reduce(n, k)); +} + +// The same over int4 weights: qw is [N][K/8] packed words and scales +// [N][K/group] floats, two views that may share a buffer. K % group == 0 and +// group % 8 == 0. +inline bool gemv_q4(span a, span qw, span scales, span y, int64_t N, int64_t K, + int64_t group) { + if (N <= 0 || K <= 0 || group <= 0 || K % group != 0 || group % 8 != 0) { + return false; + } + return launch(kop::gemv_q4_, {in(a), in(qw), in(scales), out(y)}, + gemv_q4_params{static_cast(N), static_cast(K), + static_cast(group)}, + policy::row_reduce(N, K)); +} + +// One decode step's k, v (each [n_kv_heads, D]) into row `pos` of a +// [n_kv_heads, kv_max, D] cache, f32 or bf16. +inline bool kv_append(span Kc, span Vc, span k_new, span v_new, int64_t pos, + int64_t kv_max, int64_t n_kv_heads, int64_t D, + bool kv_bf16 = false) { + if ((D != 64 && D != 128) || n_kv_heads <= 0) return false; + return launch(kv_bf16 ? kop::kv_append_bf16_ : kop::kv_append_, + {out(Kc), out(Vc), in(k_new), in(v_new)}, + kv_append_params{static_cast(pos), + static_cast(kv_max * D)}, + policy::per_head(n_kv_heads, 1, D)) || + (!kv_bf16 && + generic::kv_append(Kc, Vc, k_new, v_new, pos, kv_max, n_kv_heads, D)); +} + +// A prefill's k, v (each [n_kv_heads, T, D]) into cache rows [pos0, pos0 + T). +inline bool kv_fill(span Kc, span Vc, span K, span V, int64_t T, int64_t kv_max, + int64_t n_kv_heads, int64_t D, bool kv_bf16 = false, + int64_t pos0 = 0) { + if ((D != 64 && D != 128) || n_kv_heads <= 0 || T <= 0) return false; + return launch(kv_bf16 ? kop::kv_fill_bf16_ : kop::kv_fill_, + {out(Kc), out(Vc), in(K), in(V)}, + kv_fill_params{static_cast(T), + static_cast(kv_max * D), + static_cast(pos0)}, + policy::per_head(n_kv_heads, T, D)) || + (!kv_bf16 && + generic::kv_fill(Kc, Vc, K, V, T, kv_max, n_kv_heads, D, pos0)); +} + +// Head-major [H, T, D] -> token-major [T, H*D]: split_heads' inverse. +inline bool merge_heads(span src, span dst, int64_t T, int64_t H, int64_t D) { + if (T <= 0 || H <= 0 || D <= 0) return false; + return launch(kop::merge_heads_, {in(src), out(dst)}, + merge_heads_params{static_cast(T), + static_cast(H), + static_cast(D)}, + policy::per_head(H, T, D)) || + generic::merge_heads(src, dst, T, H, D); +} + +// Token-major [T, ld] -> head-major [H, T, D] from column block `off`, adding +// the optional per-head bias [H, D] (a null view: none). Backend-own: the +// kernels disagree on how "no bias" is said. +TL_GPU_DETECT_OWN(split_heads) +template +inline bool split_heads(span src, span bias, span dst, int64_t T, int64_t ld, + int64_t off, int64_t H, int64_t D) { + if (T <= 0 || H <= 0 || D <= 0) return false; + if constexpr (detail::owns_split_heads::value) { + if (detail::ran(Own::split_heads(src, bias, dst, T, ld, off, H, D))) { + return true; + } + } + return generic::split_heads(src, bias, dst, T, ld, off, H, D); +} + +// The argmax of a length-n vector, the smallest index on ties: greedy decoding +// reads one int back rather than the logits. Drains the queue. Backend-own: +// the result's staging buffer and its read-back are the backend's. +TL_GPU_DETECT_OWN(argmax) +template +inline bool argmax(span a, int64_t n, int64_t* out_idx) { + if (!a || n <= 0 || !out_idx) return false; + if constexpr (detail::owns_argmax::value) { + if (detail::ran(Own::argmax(a, n, out_idx))) return true; + } + return generic::argmax(a, n, out_idx); +} + +// Rotary embedding over [rows, T, D] at position `pos`, adding the optional +// per-row bias first. +TL_GPU_DETECT_OWN(rope) +template +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::value) { + if (detail::ran(Own::rope(x, o, rows, T, D, pos, base, bias))) return true; + } + return generic::rope(x, o, rows, T, D, pos, base, bias); +} + // ---- the LLM path. // y[1,n] = a[1,k] . B[k,n], the weight column-major, f32 or bf16. @@ -569,17 +823,16 @@ TL_GPU_DETECT_OWN(gemv_f32) template inline bool gemv_f32(span a, span B, span y, int64_t n, int64_t k) { if constexpr (detail::owns_gemv_f32::value) { - return detail::ran("gemv_f32", Own::gemv_f32(a, B, y, n, k)); - } else { - return false; + if (detail::ran(Own::gemv_f32(a, B, y, n, k))) return true; } + return generic::gemv_f32(a, B, y, n, k); } TL_GPU_DETECT_OWN(gemv_bf16) template inline bool gemv_bf16(span a, span B, span y, int64_t n, int64_t k) { if constexpr (detail::owns_gemv_bf16::value) { - return detail::ran("gemv_bf16", Own::gemv_bf16(a, B, y, n, k)); + return detail::ran(Own::gemv_bf16(a, B, y, n, k)); } else { return false; } @@ -591,7 +844,7 @@ template inline bool gemm_bf16_nt(span A, span B, span C, int64_t M, int64_t N, int64_t K) { if constexpr (detail::owns_gemm_bf16_nt::value) { - return detail::ran("gemm_bf16_nt", Own::gemm_bf16_nt(A, B, C, M, N, K)); + return detail::ran(Own::gemm_bf16_nt(A, B, C, M, N, K)); } else { return false; } @@ -605,10 +858,12 @@ inline bool attn_decode(span q, span K, span V, span o, int64_t n_q_heads, int64_t n_kv_heads, int64_t ctx, int64_t kv_max, int64_t D, float scale, bool kv_bf16 = false) { if constexpr (detail::owns_attn_decode::value) { - return detail::ran("attn_decode", Own::attn_decode(q, K, V, o, n_q_heads, n_kv_heads, ctx, kv_max, D, scale, kv_bf16)); - } else { - return false; + if (detail::ran(Own::attn_decode(q, K, V, o, n_q_heads, n_kv_heads, ctx, kv_max, D, scale, kv_bf16))) { + return true; + } } + return generic::attn_decode(q, K, V, o, n_q_heads, n_kv_heads, ctx, kv_max, D, + scale, kv_bf16); } // Causal prefill: q, out [n_q_heads, T, D] against the cache rows @@ -620,10 +875,12 @@ inline bool attn_prefill(span q, span K, span V, span o, int64_t n_q_heads, int64_t D, float scale, bool kv_bf16 = false, int64_t pos0 = 0) { if constexpr (detail::owns_attn_prefill::value) { - return detail::ran("attn_prefill", Own::attn_prefill(q, K, V, o, n_q_heads, n_kv_heads, T, kv_max, D, scale, kv_bf16, pos0)); - } else { - return false; + if (detail::ran(Own::attn_prefill(q, K, V, o, n_q_heads, n_kv_heads, T, kv_max, D, scale, kv_bf16, pos0))) { + return true; + } } + return generic::attn_prefill(q, K, V, o, n_q_heads, n_kv_heads, T, kv_max, D, + scale, kv_bf16, pos0); } // The causal prefill's pullback, query half then key/value half. @@ -633,7 +890,7 @@ inline bool attn_prefill_dq(span q, span K, span V, span dO, span O, span dq, span stats, int64_t H, int64_t T, int64_t D, float scale) { if constexpr (detail::owns_attn_prefill_dq::value) { - return detail::ran("attn_prefill_dq", Own::attn_prefill_dq(q, K, V, dO, O, dq, stats, H, T, D, scale)); + return detail::ran(Own::attn_prefill_dq(q, K, V, dO, O, dq, stats, H, T, D, scale)); } else { return false; } @@ -645,7 +902,7 @@ inline bool attn_prefill_dkv(span q, span K, span V, span dO, span stats, span dK, span dV, int64_t H, int64_t T, int64_t D, float scale) { if constexpr (detail::owns_attn_prefill_dkv::value) { - return detail::ran("attn_prefill_dkv", Own::attn_prefill_dkv(q, K, V, dO, stats, dK, dV, H, T, D, scale)); + return detail::ran(Own::attn_prefill_dkv(q, K, V, dO, stats, dK, dV, H, T, D, scale)); } else { return false; } @@ -659,7 +916,7 @@ template inline bool rope_dpos(span x, span o, int64_t rows, int64_t T, int64_t D, span d_pos, float base, span bias = {}) { if constexpr (detail::owns_rope_dpos::value) { - return detail::ran("rope_dpos", Own::rope_dpos(x, o, rows, T, D, d_pos, base, bias)); + return detail::ran(Own::rope_dpos(x, o, rows, T, D, d_pos, base, bias)); } else { return false; } @@ -670,7 +927,7 @@ template inline bool kv_append_dpos(span Kc, span Vc, span k_new, span v_new, span d_pos, int64_t kv_max, int64_t n_kv_heads, int64_t D) { if constexpr (detail::owns_kv_append_dpos::value) { - return detail::ran("kv_append_dpos", Own::kv_append_dpos(Kc, Vc, k_new, v_new, d_pos, kv_max, n_kv_heads, D)); + return detail::ran(Own::kv_append_dpos(Kc, Vc, k_new, v_new, d_pos, kv_max, n_kv_heads, D)); } else { return false; } @@ -682,7 +939,7 @@ inline bool attn_decode_dpos(span q, span K, span V, span o, int64_t n_q_heads, int64_t n_kv_heads, span d_pos, int64_t kv_max, int64_t D, float scale, span partials) { if constexpr (detail::owns_attn_decode_dpos::value) { - return detail::ran("attn_decode_dpos", Own::attn_decode_dpos(q, K, V, o, n_q_heads, n_kv_heads, d_pos, kv_max, D, scale, partials)); + return detail::ran(Own::attn_decode_dpos(q, K, V, o, n_q_heads, n_kv_heads, d_pos, kv_max, D, scale, partials)); } else { return false; } @@ -693,7 +950,7 @@ TL_GPU_DETECT_OWN(incr_u32) template inline bool incr_u32(span d_pos) { if constexpr (detail::owns_incr_u32::value) { - return detail::ran("incr_u32", Own::incr_u32(d_pos)); + return detail::ran(Own::incr_u32(d_pos)); } else { return false; } diff --git a/include/metal.h b/include/metal.h index 2ce1835..b2b80fa 100644 --- a/include/metal.h +++ b/include/metal.h @@ -280,6 +280,19 @@ struct context { inline bool available() { return context::get().device != nullptr; } +// 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 2-D thread + // position rather than a flat index. + static constexpr bool cells_2d = true; + // Each launch's tl::profile row carries a device time. + static constexpr bool times_launches = true; + // Threadgroups that keep the GPU busy: ~4 per core of a 16-core Apple GPU, + // the target the GEMV and attention splits below measure their grid against. + static constexpr int64_t fill_groups = 64; +}; + // The ops this backend runs its own way: a different algorithm, several // kernels, or a kernel whose ABI is its own. gpu_ops.h forwards to whichever of // these exist (TL_GPU_DETECT_OWN) and answers false for the rest, so a backend @@ -446,9 +459,11 @@ inline void dispatch_grid_(objc::id enc, mtl_size grid, mtl_size tg) { reinterpret_cast(objc_msgSend)( enc, sel_registerName("dispatchThreadgroups:threadsPerThreadgroup:"), grid, tg); - if (profile::active()) { - auto& c = context::get(); - c.commit_(profile::detail::launch(c.bound ? c.bound : "?")); + // Under a profile the launch is its own command buffer, owed the row's + // GPU time at the flush. + auto& c = context::get(); + if (profile::row* r = gpu::launched(c.bound ? c.bound : "?")) { + c.commit_(r); profile::detail::drain_hook = &flush; } } @@ -1057,16 +1072,14 @@ struct attn_combine_params { }; // Keys per split, or 0 for the single-pass kernel: one threadgroup a head -// leaves most of a 16-core GPU idle when a model has few heads, so cut the -// keys until there are enough threadgroups (cuda's attn_split_count). +// leaves most of the GPU idle when a model has few heads, so cut the keys +// until there are enough threadgroups. The kernel's rule: no split under 256 +// keys, a split at least 128 keys and a multiple of 4. inline unsigned long attn_split_chunk_(int64_t heads, int64_t ctx) { - constexpr long kWantGroups = 64, kMinKeys = 128; - if (heads <= 0 || heads >= kWantGroups || ctx < 2 * kMinKeys) return 0; - long want = (kWantGroups + heads - 1) / heads; - const long most = ctx / kMinKeys; - if (want > most) want = most; - if (want <= 1) return 0; - return static_cast((ctx + want - 1) / want); + constexpr gpu::policy::split_rule rule{traits::fill_groups, 256, 128, 4}; + const int64_t parts = gpu::policy::split_parts(heads, ctx, rule); + if (parts <= 1) return 0; + return static_cast(gpu::policy::split_chunk(ctx, parts, rule)); } inline void attn_dispatch_(objc::id enc, const attn_params& p, @@ -1108,17 +1121,18 @@ inline bool own::attn_prefill(gpu::span q, gpu::span K, gpu::span V, // y[1,N] = a[1,K] · B[K,N] with f32 or bf16 weights, all contiguous: the // decode projection. A narrow layer's N/256 threadgroups leave the GPU idle, -// so K is split across the grid's y and a combine pass sums the slices. +// so K is split across the grid's y and a combine pass sums the slices. The +// kernel's rule: a slice is whole 256-wide `a` tiles. inline bool gemv_(kop op, gpu::span a, gpu::span B, gpu::span y, int64_t n, int64_t k) { auto& c = context::get(); if (!c.device || n <= 0 || k <= 0) return false; - constexpr unsigned long NT = 256, kGroupsWanted = 64; + constexpr unsigned long NT = 256; + constexpr gpu::policy::split_rule rule{traits::fill_groups, 0, 0, NT}; const unsigned long cols = (static_cast(n) + NT - 1) / NT; - unsigned long parts = cols >= kGroupsWanted ? 1 : kGroupsWanted / cols; - unsigned long chunk = (static_cast(k) + parts - 1) / parts; - chunk = (chunk + NT - 1) / NT * NT; // whole `a` tiles - parts = (static_cast(k) + chunk - 1) / chunk; + unsigned long chunk = static_cast(gpu::policy::split_chunk( + k, gpu::policy::split_parts(cols, k, rule), rule)); + unsigned long parts = (static_cast(k) + chunk - 1) / chunk; gpu::span out = y; if (parts > 1) { out = {detail_::scratch_(static_cast(parts) * n * 4), 0}; @@ -1300,23 +1314,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 2-D thread - // position rather than a flat index. - static constexpr bool cells_2d = true; - // 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 = false; static constexpr bool row_gemv = true; // gemv_bf16_row: weights as [N,K] static constexpr bool bf16_gemm = true; // gemm_bf16_nt: a bf16-weight GEMM diff --git a/include/profile.h b/include/profile.h index 45f6681..6d5413b 100644 --- a/include/profile.h +++ b/include/profile.h @@ -317,10 +317,11 @@ inline void report(FILE* out) { } namespace detail { -// TL_PROFILE=1: profile the whole process from the first evaluation and -// print the table to stderr at exit. Called by the evaluator; the first call -// does the work (the exit handler reads the calling thread's state, which is -// the evaluating one in every consumer). +// TL_PROFILE=1: profile the whole process from the first evaluation or kernel +// launch (gpu::launched — a decoder on the model path never reaches the +// evaluator) and print the table to stderr at exit. The first call does the +// work (the exit handler reads the calling thread's state, which is the +// evaluating one in every consumer). inline void env_autostart() { static bool once = false; if (once) return; diff --git a/include/webgpu.h b/include/webgpu.h index 11ad0bc..6264257 100644 --- a/include/webgpu.h +++ b/include/webgpu.h @@ -417,7 +417,7 @@ struct context { pending = true; dispatch_counts[entry]++; - if (profile::active()) profile::detail::launch(entry); // counted, untimed + gpu::launched(entry); // counted, untimed return true; } @@ -1228,17 +1228,15 @@ struct traits { // A [rows, cols] elementwise kernel reads its cell from a 2-D thread // position rather than a flat index. static constexpr bool cells_2d = true; - // Launches are recorded under tl::profile by this backend itself, with - // (times_launches) a device time on each. - static constexpr bool profiles_launches = true; + // A launch's tl::profile row is counted, not timed: WebGPU has no per-launch + // device time to stamp it with. static constexpr bool times_launches = false; + // Workgroups that keep the device busy. Untuned: the browser hides the + // device, and no op here splits against it yet. + static constexpr int64_t fill_groups = 64; }; 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 = false; static constexpr bool graph_capture = false; static constexpr bool row_gemv = false; static constexpr bool bf16_gemm = false; diff --git a/test/test_array.cpp b/test/test_array.cpp index 7171f2a..512acd7 100644 --- a/test/test_array.cpp +++ b/test/test_array.cpp @@ -2896,6 +2896,198 @@ TEST_CASE("profile: an eager op that declines on its operands leaves no row") { } } +// The shared split policy, at the rules the backends state. The CUDA numbers +// are the ones its decode attention's device twin (attn_dpos_chunk) reproduces +// and its launch trace records, so they are pinned here as such. +TEST_CASE("policy: a split measures its groups against the device's fill") { + using tl::gpu::policy::split_chunk; + using tl::gpu::policy::split_parts; + using tl::gpu::policy::split_rule; + + // CUDA's decode GEMV: 164 blocks fill, no split under K=512, 32-wide steps. + constexpr split_rule gemv{164, 512, 0, 32}; + CHECK(split_parts(4, 896, gemv) == 41); // 896 = Qwen's NE, 4 column blocks + CHECK(split_chunk(896, 41, gemv) == 32); // 22 rounded up to the step + CHECK(split_parts(594, 4096, gemv) == 1); // the vocab projection fills + CHECK(split_parts(4, 511, gemv) == 1); // too short a K + CHECK(split_parts(0, 4096, gemv) == 1); + + // CUDA's decode attention: twice the fill, 256 keys to split, 128 a part, + // parts in whole warps of 4. + constexpr split_rule attn{328, 256, 128, 4}; + CHECK(split_parts(14, 1000, attn) == 7); // 24 wanted, 1000 / 128 allow 7 + CHECK(split_chunk(1000, 7, attn) == 144); // 143 rounded up + CHECK((1000 + 144 - 1) / 144 == 7); + CHECK(split_parts(14, 255, attn) == 1); + CHECK(split_parts(400, 4096, attn) == 1); + CHECK(split_parts(14, 4096, attn) == 24); + // parts never fall as k grows: the captured graph's grid is sized at max_ctx + for (int64_t ctx = 1; ctx < 4096; ctx++) { + CHECK(split_parts(14, ctx, attn) <= split_parts(14, ctx + 1, attn)); + } + + // CUDA's bf16 GEMM: four times the fill, slices at least 448 deep in + // 16-deep slabs. + constexpr split_rule gemm{656, 0, 448, 16}; + CHECK(split_parts(12, 4864, gemm) == 10); // 55 wanted, 4864 / 448 allow 10 + CHECK(split_chunk(4864, 10, gemm) == 496); + CHECK(split_parts(700, 4864, gemm) == 1); + + // A part is never zero: the shortest k still rounds up to one granule. + CHECK(split_chunk(1, 1, gemv) == 32); +} + +// The generic compositions (gpu_ops.h, tier 1 out of tier 0) against the op +// as the backend runs it. Where the backend has the fused kernel the two are +// different programs that must agree; where it does not, the op is the +// composition and this is the composition against itself, and the model-path +// test above holds it to the array oracle. +TEST_CASE("generic compositions agree with the fused kernels") { + if (!tl::gpu_available()) return; + auto prev = tl::device_; + tl::use_gpu(); + namespace gpu = tl::gpu; + namespace gen = tl::gpu::generic; + auto dev = [](const array& a) { + array c = a.clone(); + c.eval(); + return c; + }; + auto same = [](const array& got, const array& want, float tol) { + return tl::allclose(got, want, tol, tol); + }; + + SUBCASE("rmsnorm, rmsnorm_res and swiglu") { + const int64_t rows = 3, n = 896, ff = 512; + array x = dev(random_array({rows, n}, 950)), d = dev(random_array({rows, n}, 951)); + array w = dev(random_array({n}, 952)); + array a = array::empty({rows, n}), b = array::empty({rows, n}); + REQUIRE(gpu::rmsnorm(x.device_span(), w.device_span(), a.device_span(), n, 1e-6f, rows)); + REQUIRE(gen::rmsnorm(x.device_span(), w.device_span(), b.device_span(), n, 1e-6f, rows)); + tl::gpu::flush(); + CHECK(same(a, b, 1e-5f)); + // Rows small enough that eps carries the reciprocal (mean(x^2) ~ 1e-8 + // against eps 1e-6): a composition that dropped eps would be off tenfold. + array tiny = dev(x * 1e-4f); + REQUIRE(gpu::rmsnorm(tiny.device_span(), w.device_span(), a.device_span(), n, 1e-6f, rows)); + REQUIRE(gen::rmsnorm(tiny.device_span(), w.device_span(), b.device_span(), n, 1e-6f, rows)); + tl::gpu::flush(); + CHECK(same(a, b, 1e-5f)); + + array xa = array::empty({rows, n}), ha = array::empty({rows, n}); + array xb = array::empty({rows, n}), hb = array::empty({rows, n}); + REQUIRE(gpu::rmsnorm_res(x.device_span(), d.device_span(), w.device_span(), xa.device_span(), + ha.device_span(), n, 1e-6f, rows)); + REQUIRE(gen::rmsnorm_res(x.device_span(), d.device_span(), w.device_span(), xb.device_span(), + hb.device_span(), n, 1e-6f, rows)); + tl::gpu::flush(); + CHECK(same(xa, xb, 0.0f)); + CHECK(same(ha, hb, 1e-5f)); + + array gu = dev(random_array({rows, 2 * ff}, 953)); + array oa = array::empty({rows, ff}), ob = array::empty({rows, ff}); + REQUIRE(gpu::swiglu(gu.device_span(), oa.device_span(), ff, rows)); + REQUIRE(gen::swiglu(gu.device_span(), ob.device_span(), ff, rows)); + tl::gpu::flush(); + CHECK(same(oa, ob, 1e-5f)); + } + + SUBCASE("the f32 decode GEMV") { + const int64_t K = 896, N = 1152; + array a = dev(random_array({1, K}, 954)), B = dev(random_array({K, N}, 955)); + array ya = array::empty({1, N}), yb = array::empty({1, N}); + REQUIRE(gpu::gemv_f32(a.device_span(), B.device_span(), ya.device_span(), N, K)); + REQUIRE(gen::gemv_f32(a.device_span(), B.device_span(), yb.device_span(), N, K)); + tl::gpu::flush(); + CHECK(same(ya, yb, 1e-3f)); + } + + SUBCASE("the cache: fill, append, decode and prefill attention") { + const int64_t D = 64, HKV = 2, HQ = 14, MAXC = 64, T = 20, T2 = 7; + const float scale = 1.0f / std::sqrt((float)D); + array K = dev(random_array({HKV, T, D}, 960)), V = dev(random_array({HKV, T, D}, 961)); + array Ka = dev(array::zeros({HKV, MAXC, D})), Va = dev(array::zeros({HKV, MAXC, D})); + array Kb = dev(array::zeros({HKV, MAXC, D})), Vb = dev(array::zeros({HKV, MAXC, D})); + REQUIRE(gpu::kv_fill(Ka.device_span(), Va.device_span(), K.device_span(), V.device_span(), T, + MAXC, HKV, D)); + REQUIRE(gen::kv_fill(Kb.device_span(), Vb.device_span(), K.device_span(), V.device_span(), T, + MAXC, HKV, D, 0)); + array k1 = dev(random_array({HKV, D}, 962)), v1 = dev(random_array({HKV, D}, 963)); + REQUIRE(gpu::kv_append(Ka.device_span(), Va.device_span(), k1.device_span(), v1.device_span(), + T, MAXC, HKV, D)); + REQUIRE(gen::kv_append(Kb.device_span(), Vb.device_span(), k1.device_span(), v1.device_span(), + T, MAXC, HKV, D)); + tl::gpu::flush(); + CHECK(same(Ka, Kb, 0.0f)); + CHECK(same(Va, Vb, 0.0f)); + + // A decode step over the T + 1 cached rows, by both routes. + array q = dev(random_array({HQ, D}, 964)); + array oa = array::empty({HQ, D}), ob = array::empty({HQ, D}); + REQUIRE(gpu::attn_decode(q.device_span(), Ka.device_span(), Va.device_span(), oa.device_span(), + HQ, HKV, T + 1, MAXC, D, scale)); + REQUIRE(gen::attn_decode(q.device_span(), Ka.device_span(), Va.device_span(), ob.device_span(), + HQ, HKV, T + 1, MAXC, D, scale, false)); + tl::gpu::flush(); + CHECK(same(oa, ob, 1e-4f)); + + // A prefill of T2 more rows after them: the mask has to open the T + 1 + // rows already cached and close the keys past each query. + array K2 = dev(random_array({HKV, T2, D}, 965)), V2 = dev(random_array({HKV, T2, D}, 966)); + REQUIRE(gpu::kv_fill(Ka.device_span(), Va.device_span(), K2.device_span(), V2.device_span(), + T2, MAXC, HKV, D, false, T + 1)); + array qp = dev(random_array({HQ, T2, D}, 967)); + array pa = array::empty({HQ, T2, D}), pb = array::empty({HQ, T2, D}); + REQUIRE(gpu::attn_prefill(qp.device_span(), Ka.device_span(), Va.device_span(), + pa.device_span(), HQ, HKV, T2, MAXC, D, scale, false, T + 1)); + REQUIRE(gen::attn_prefill(qp.device_span(), Ka.device_span(), Va.device_span(), + pb.device_span(), HQ, HKV, T2, MAXC, D, scale, false, T + 1)); + tl::gpu::flush(); + CHECK(same(pa, pb, 1e-4f)); + } + + SUBCASE("rope, split_heads, merge_heads and argmax") { + const int64_t H = 14, T = 5, D = 64, ld = H * D; + array x = dev(random_array({H * T, D}, 970)), b = dev(random_array({H * T, D}, 971)); + array ra = array::empty({H * T, D}), rb = array::empty({H * T, D}); + REQUIRE(gpu::rope(x.device_span(), ra.device_span(), H * T, T, D, 37, 1e6f, b.device_span())); + REQUIRE(gen::rope(x.device_span(), rb.device_span(), H * T, T, D, 37, 1e6f, b.device_span())); + tl::gpu::flush(); + CHECK(same(ra, rb, 1e-4f)); + REQUIRE(gpu::rope(x.device_span(), ra.device_span(), H * T, T, D, 37, 1e6f)); + REQUIRE(gen::rope(x.device_span(), rb.device_span(), H * T, T, D, 37, 1e6f, {})); + tl::gpu::flush(); + CHECK(same(ra, rb, 1e-4f)); + + array src = dev(random_array({T, ld}, 972)), bias = dev(random_array({H, D}, 973)); + array sa = array::empty({H, T, D}), sb = array::empty({H, T, D}); + REQUIRE(gpu::split_heads(src.device_span(), bias.device_span(), sa.device_span(), T, ld, 0, H, D)); + REQUIRE(gen::split_heads(src.device_span(), bias.device_span(), sb.device_span(), T, ld, 0, H, D)); + tl::gpu::flush(); + CHECK(same(sa, sb, 0.0f)); + REQUIRE(gpu::split_heads(src.device_span(), {}, sa.device_span(), T, ld, 0, H, D)); + REQUIRE(gen::split_heads(src.device_span(), {}, sb.device_span(), T, ld, 0, H, D)); + array ma = array::empty({T, ld}), mb = array::empty({T, ld}); + REQUIRE(gpu::merge_heads(sa.device_span(), ma.device_span(), T, H, D)); + REQUIRE(gen::merge_heads(sa.device_span(), mb.device_span(), T, H, D)); + tl::gpu::flush(); + CHECK(same(sa, sb, 0.0f)); + CHECK(same(ma, mb, 0.0f)); + CHECK(same(ma, src, 0.0f)); + + std::vector v(10000, -1.0f); + v[4321] = v[4322] = 5.0f; + array a = dev(array::from(v, {(int64_t)v.size()})); + int64_t ia = -1, ib = -1; + REQUIRE(gpu::argmax(a.device_span(), (int64_t)v.size(), &ia)); + REQUIRE(gen::argmax(a.device_span(), (int64_t)v.size(), &ib)); + CHECK(ia == 4321); + CHECK(ib == 4321); + } + + tl::device_ = prev; +} + // A buffer released while the device still has work queued against it goes // straight back to the pool for device work, but not for the host to fill. // Here the logsumexp is encoded but not run when its temporary dies (xent_bwd @@ -2921,9 +3113,9 @@ TEST_CASE("the KV cache and the decode step's kernels match their array forms") // graph; each kernel here is checked against the array composition that // defines it. Shapes are Qwen2's head geometry (D=64, 14 q heads over 2 kv // heads) at a context past the split-KV cutoff, plus the D=128 - // instantiation. Nothing to compare where the backend answers false for the - // whole row (gpu::caps::model_path) or has no device at all. - if (!tl::gpu_available() || !tl::gpu::caps::model_path) return; + // instantiation. A backend without a kernel takes the generic composition, + // which has to agree too; a bf16 cache needs the backend's own attention. + if (!tl::gpu_available()) return; auto prev = tl::device_; tl::use_gpu(); namespace gpu = tl::gpu; @@ -2941,6 +3133,7 @@ TEST_CASE("the KV cache and the decode step's kernels match their array forms") SUBCASE("kv_cache: append + attn, prefill, in f32 and bf16") { for (int64_t D : {64, 128}) { for (tl::dtype kv : {tl::dtype::f32, tl::dtype::bf16}) { + if (kv == tl::dtype::bf16 && !gpu::has_attn_decode) continue; const int64_t HQ = 14, HKV = 2, MAXC = 300, T = 260; const float scale = 1.0f / std::sqrt((float)D); tl::kv_cache cache; @@ -3334,7 +3527,7 @@ TEST_CASE("gpu ops on views at non-zero offsets") { 0.0f), true, o, want, 1e-5f); } - const bool model = gpu::caps::model_path; + // Every backend has these: the kernel, or the generic composition. { staged xo = out_of(rows * cols), ho = out_of(rows * cols); std::vector wx(rows * cols), wh(rows * cols); @@ -3350,8 +3543,8 @@ TEST_CASE("gpu ops on views at non-zero offsets") { } const bool ran = gpu::rmsnorm_res(x.view(), d.view(), g.view(), xo.view(), ho.view(), cols, 1e-6f, rows); - check(ran, model, xo, wx, 1e-6f); - check(ran, model, ho, wh, 1e-5f); + check(ran, true, xo, wx, 1e-6f); + check(ran, true, ho, wh, 1e-5f); } { const int64_t ff = cols / 2; // x as [rows, 2*ff]: gate | up @@ -3362,7 +3555,7 @@ TEST_CASE("gpu ops on views at non-zero offsets") { const double gate = vx[r * cols + f], up = vx[r * cols + ff + f]; want[r * ff + f] = (float)(gate / (1.0 + std::exp(-gate)) * up); } - check(gpu::swiglu(x.view(), o.view(), ff, rows), model, o, want, 1e-6f); + check(gpu::swiglu(x.view(), o.view(), ff, rows), true, o, want, 1e-6f); } }