From 214d196bb0eca3d955ce0a2dd32864185a1eb500 Mon Sep 17 00:00:00 2001 From: yhirose Date: Mon, 21 Sep 2026 21:41:14 -0400 Subject: [PATCH 1/4] gpu: a launch is recorded in one place, and a backend only stamps its time MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every kernel launch on every backend is one row under tl::profile, made by gpu::launched(kernel name) in gpu_abi.h; a backend's launch primitive calls it once per launch — shared ops and own ops alike — and adds the device time where it has one. CUDA, Metal and WebGPU each had their own copy of the record, the shared layer a fourth for the backends that had none, and traits::profiles_launches said which was which; the flag is gone. TL_PROFILE=1 now starts at the first launch as well as the first evaluation, so a decoder on the model path, which never reaches the evaluator, is profiled from its first kernel (the chat decoder: 0 scopes, 2376 launches). The CUDA trace against master is identical; the host backend records a kernel that ran and its four own ops, so its suite sees the same rows. --- docs/backends.md | 23 ++++++++++----- include/cuda.h | 13 ++++----- include/gpu_abi.h | 15 ++++++++++ include/gpu_host.h | 23 ++++++++++----- include/gpu_null.h | 8 +++--- include/gpu_ops.h | 70 +++++++++++++++++++--------------------------- include/metal.h | 12 ++++---- include/profile.h | 9 +++--- include/webgpu.h | 7 ++--- 9 files changed, 99 insertions(+), 81 deletions(-) diff --git a/docs/backends.md b/docs/backends.md index 269ce4c..dab028f 100644 --- a/docs/backends.md +++ b/docs/backends.md @@ -148,9 +148,22 @@ 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`), and whether a launch's profile row +carries a device time (`traits::times_launches`). + +## 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,10 +213,6 @@ 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 diff --git a/include/cuda.h b/include/cuda.h index daf3600..50d9116 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_(); @@ -1045,9 +1044,7 @@ inline bool own::argmax(gpu::span a, int64_t n, int64_t* out_idx) { 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; + // Each launch's tl::profile row carries a device time. static constexpr bool times_launches = true; }; diff --git a/include/gpu_abi.h b/include/gpu_abi.h index 41036c9..32e5ff3 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 { diff --git a/include/gpu_host.h b/include/gpu_host.h index 6c9a5e9..b19e88f 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,10 +516,7 @@ 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 }; struct caps { // The decoder's single-kernel ops are here, but not attention, rope or the diff --git a/include/gpu_null.h b/include/gpu_null.h index 5ab255a..a8843a2 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,9 +48,8 @@ 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; }; diff --git a/include/gpu_ops.h b/include/gpu_ops.h index b94af6d..45af605 100644 --- a/include/gpu_ops.h +++ b/include/gpu_ops.h @@ -14,7 +14,6 @@ #include #include "gpu_abi.h" -#include "profile.h" namespace tl { namespace gpu { @@ -27,18 +26,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 +67,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; } @@ -337,7 +325,7 @@ 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)); + return detail::ran(Own::split_heads(src, bias, dst, T, ld, off, H, D)); } else { return false; } @@ -351,7 +339,7 @@ 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)); + return detail::ran(Own::argmax(a, n, out_idx)); } else { return false; } @@ -384,7 +372,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 +385,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 +397,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 +411,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 +425,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 +438,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 +451,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 +463,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 +475,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 +489,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 +504,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,7 +517,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)); + return detail::ran(Own::gemm_bias(a, lda, ta, b, ldb, tb, bias, o, m, n, k, scale, offset)); } else { return false; } @@ -542,7 +530,7 @@ 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::rope(x, o, rows, T, D, pos, base, bias)); } else { return false; } @@ -556,7 +544,7 @@ 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; } @@ -569,7 +557,7 @@ 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)); + return detail::ran(Own::gemv_f32(a, B, y, n, k)); } else { return false; } @@ -579,7 +567,7 @@ 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 +579,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,7 +593,7 @@ 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)); + return detail::ran(Own::attn_decode(q, K, V, o, n_q_heads, n_kv_heads, ctx, kv_max, D, scale, kv_bf16)); } else { return false; } @@ -620,7 +608,7 @@ 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)); + return detail::ran(Own::attn_prefill(q, K, V, o, n_q_heads, n_kv_heads, T, kv_max, D, scale, kv_bf16, pos0)); } else { return false; } @@ -633,7 +621,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 +633,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 +647,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 +658,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 +670,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 +681,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..100b948 100644 --- a/include/metal.h +++ b/include/metal.h @@ -446,9 +446,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; } } @@ -1306,9 +1308,7 @@ 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; + // Each launch's tl::profile row carries a device time. static constexpr bool times_launches = true; }; 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..e877bd4 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,9 +1228,8 @@ 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; }; From 7cc99d15d3381dd556557ef56dbf8ef79cfdfa58 Mon Sep 17 00:00:00 2001 From: yhirose Date: Mon, 21 Sep 2026 21:56:25 -0400 Subject: [PATCH 2/4] gpu: a split is one rule of the device's fill, stated beside its kernel A reduction that leaves the device short of busy groups is split over more of them and combined after: CUDA's decode GEMV, its bf16 GEMM tiles and its decode attention each did that arithmetic in their own words, and Metal's GEMV and attention copied it by hand with the numbers of a 16-core GPU. Now policy::split_parts and split_chunk (gpu_abi.h) are the arithmetic, under a split_rule that is the kernel's side of it (the target, the shortest k, the shortest part, the granule), and the target comes from traits::fill_groups (164 on CUDA, 64 on Metal) rather than a constant of one backend. CUDA's numbers are preserved exactly, since attn_dpos_chunk in the kernel source reproduces them for the captured-graph path: the launch trace against master is identical, and a test pins the rules' values. Metal's GEMV now rounds its part count up rather than down and its attention rounds a chunk to four keys, as CUDA's does; decode throughput is unchanged (121 tok/s either way on Qwen2.5-0.5B). CUDA's f32 GEMM wave plan stays its own, written against the trait. --- docs/backends.md | 27 +++++++++++--- include/cuda.h | 85 +++++++++++++++++++++------------------------ include/gpu_abi.h | 32 +++++++++++++++++ include/gpu_host.h | 1 + include/gpu_null.h | 3 ++ include/metal.h | 52 ++++++++++++++------------- include/webgpu.h | 3 ++ test/test_array.cpp | 41 ++++++++++++++++++++++ 8 files changed, 170 insertions(+), 74 deletions(-) diff --git a/docs/backends.md b/docs/backends.md index dab028f..50c7407 100644 --- a/docs/backends.md +++ b/docs/backends.md @@ -150,8 +150,26 @@ Metal and WGSL size theirs in the kernel). 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`), and whether a launch's profile row -carries a device time (`traits::times_launches`). +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 @@ -215,8 +233,9 @@ asks; a new backend edits no test. itself is shared. - `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. +- 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. - 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. diff --git a/include/cuda.h b/include/cuda.h index 50d9116..f693a46 100644 --- a/include/cuda.h +++ b/include/cuda.h @@ -895,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: @@ -910,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 @@ -930,11 +940,11 @@ 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; + if (tiles >= traits::fill_groups) 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(traits::fill_groups / tiles, slabs / min_slabs)); long chunk_slabs = (slabs + full - 1) / full; - if (const long spare = kFillBlocks - full * tiles; spare > 0) { + if (const long spare = traits::fill_groups - 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; @@ -1039,15 +1049,6 @@ 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; - // Each launch's tl::profile row carries a device time. - 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 @@ -1499,6 +1500,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(); @@ -1506,17 +1509,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 @@ -1604,17 +1602,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, @@ -1634,13 +1630,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 @@ -1669,10 +1666,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_abi.h b/include/gpu_abi.h index 32e5ff3..ea83003 100644 --- a/include/gpu_abi.h +++ b/include/gpu_abi.h @@ -325,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 b19e88f..9521d2c 100644 --- a/include/gpu_host.h +++ b/include/gpu_host.h @@ -517,6 +517,7 @@ struct own { struct traits { static constexpr bool cells_2d = false; // a cell is read from a flat index 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 diff --git a/include/gpu_null.h b/include/gpu_null.h index a8843a2..33a7094 100644 --- a/include/gpu_null.h +++ b/include/gpu_null.h @@ -51,6 +51,9 @@ struct traits { // 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 diff --git a/include/metal.h b/include/metal.h index 100b948..7c6d5a3 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 @@ -1059,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, @@ -1110,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}; @@ -1302,16 +1314,6 @@ 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; - // Each launch's tl::profile row carries a device time. - 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 diff --git a/include/webgpu.h b/include/webgpu.h index e877bd4..f55f0f7 100644 --- a/include/webgpu.h +++ b/include/webgpu.h @@ -1231,6 +1231,9 @@ struct traits { // 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 { diff --git a/test/test_array.cpp b/test/test_array.cpp index 7171f2a..c0c61e0 100644 --- a/test/test_array.cpp +++ b/test/test_array.cpp @@ -2896,6 +2896,47 @@ 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); +} + // 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 From e3b449dba5393553ba7718ca4fe6e8db75cbc69f Mon Sep 17 00:00:00 2001 From: yhirose Date: Mon, 21 Sep 2026 22:30:37 -0400 Subject: [PATCH 3/4] gpu: the model path out of tier 0, once, under every fused kernel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A fused op — rmsnorm, swiglu, the decode GEMV, the cache writes, rope, the attention, split_heads / merge_heads, argmax — was a kernel a backend had or a false the model had to route around, and caps::model_path said which backend was which. Now gpu_ops.h holds one generic composition of each out of the tier-0 ops (gpu::generic), and the op takes it when the shared launch declines or `own` has no member. A backend with only tier 0 runs the whole f32 model path: the KV cache, the decoder's imperative step, the model-path tests, all of it now runs on WebGPU and on the host backend, and caps::model_path is gone. An operand the tier has no reader for (a bf16 cache or weight, int4 weights) still declines, under the caps that name it. A composition names no backend and no array — spans, tier-0 ops, and the device core's alloc / release / cpu_barrier / sync_to_host — and no launch in it writes a buffer it reads, which WebGPU forbids within a dispatch (and invalidates the whole encoder over). A second test runs each composition beside the backend's kernel, where it has one, and requires the same numbers; the eps in rmsnorm is the kind of thing it catches. The CUDA trace against master adds the new test's launches and changes nothing else. --- bench/models/check_qwen.cpp | 13 +- docs/backends.md | 37 ++- include/cuda.h | 4 - include/gpu.h | 14 +- include/gpu_host.h | 3 - include/gpu_null.h | 1 - include/gpu_ops.h | 553 +++++++++++++++++++++++++++--------- include/metal.h | 4 - include/webgpu.h | 4 - test/test_array.cpp | 166 ++++++++++- 10 files changed, 612 insertions(+), 187 deletions(-) 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 50c7407..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: @@ -236,9 +258,6 @@ asks; a new backend edits no test. - 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. -- 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. ## Verifying a change diff --git a/include/cuda.h b/include/cuda.h index f693a46..9d25966 100644 --- a/include/cuda.h +++ b/include/cuda.h @@ -1050,10 +1050,6 @@ inline bool own::argmax(gpu::span a, int64_t n, int64_t* out_idx) { } 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 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_host.h b/include/gpu_host.h index 9521d2c..1ee1a6a 100644 --- a/include/gpu_host.h +++ b/include/gpu_host.h @@ -520,9 +520,6 @@ struct traits { 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 33a7094..59d3266 100644 --- a/include/gpu_null.h +++ b/include/gpu_null.h @@ -59,7 +59,6 @@ struct traits { // ---- 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 45af605..98dd0c9 100644 --- a/include/gpu_ops.h +++ b/include/gpu_ops.h @@ -9,6 +9,7 @@ // Included by gpu.h, after the backend is selected. #include +#include #include #include #include @@ -225,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(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(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) { @@ -523,19 +404,6 @@ inline bool gemm_bias(span a, int64_t lda, bool ta, span b, int64_t ldb, } } -// 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(Own::rope(x, o, rows, T, D, pos, base, bias)); - } else { - return false; - } -} - // Layer norm's pullback: dx, and dg / db reduced over rows in chunks. TL_GPU_DETECT_OWN(layer_norm_bwd) template @@ -550,6 +418,404 @@ inline bool layer_norm_bwd(span x, span g, span dy, span dx, span dg, span db, } } +// ---- 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. @@ -557,10 +823,9 @@ 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(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) @@ -593,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(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 @@ -608,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(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. diff --git a/include/metal.h b/include/metal.h index 7c6d5a3..b2b80fa 100644 --- a/include/metal.h +++ b/include/metal.h @@ -1315,10 +1315,6 @@ inline bool own::argmax(gpu::span a, int64_t n, int64_t* out_idx) { } 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/webgpu.h b/include/webgpu.h index f55f0f7..6264257 100644 --- a/include/webgpu.h +++ b/include/webgpu.h @@ -1237,10 +1237,6 @@ struct traits { }; 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 c0c61e0..512acd7 100644 --- a/test/test_array.cpp +++ b/test/test_array.cpp @@ -2937,6 +2937,157 @@ TEST_CASE("policy: a split measures its groups against the device's fill") { 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 @@ -2962,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; @@ -2982,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; @@ -3375,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); @@ -3391,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 @@ -3403,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); } } From 6abecc6363a3e5d17c067713689b1b5ff896f558 Mon Sep 17 00:00:00 2001 From: yhirose Date: Mon, 21 Sep 2026 23:08:40 -0400 Subject: [PATCH 4/4] cuda: the wave plan's fill stays `long`, the one width every arch agrees on sgemm_wave_chunk_ divided traits::fill_groups (int64_t) by variables the function keeps as long, and passed the mismatch straight into std::min with no explicit template argument. Every platform this library builds on has long == int64_t except Windows (LLP64, long is 32-bit), where deduction had no common type and the build failed. One cast at the top of the function keeps the rest of it in long, as it was before the trait existed. --- include/cuda.h | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/include/cuda.h b/include/cuda.h index 9d25966..59123a7 100644 --- a/include/cuda.h +++ b/include/cuda.h @@ -940,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 >= traits::fill_groups) 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(traits::fill_groups / 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 = traits::fill_groups - 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;