From 19f2d5cab4a0f032df0db9621d51fe59ebe298f8 Mon Sep 17 00:00:00 2001 From: xiaoyu <1259084489@qq.com> Date: Mon, 28 Sep 2026 18:15:57 +0800 Subject: [PATCH 1/7] cuda: fused candidate scoring, verified on an RTX 4090 Both models that need this readout share the shape of it: Kev (#27) scores scale * dot(k_proj(c_i), q_proj(s)) and softmaxes over a question's candidates; CLM (#9) scores exp(logit_scale) * cos(state_head(s), action_head(c)) and does the same. #9 asks for exactly this fusion and #19's gemm.cu does not have it. The fusion is more than fewer launches. For a cosine the candidate's norm and its dot with the query need the same elements, so both accumulate in one read of C -- for Kev, 255 rows of 2560 floats read once instead of twice. The softmax runs in-block over shared memory, so there is no second kernel, no atomics and no global round trip for the similarities. A batch is a grid over questions. MEASURED, not just intended. RTX 4090 (sm_89), driver 595.58.03, CUDA 13.0 (V13.0.88): worst absolute difference: 1.835e-07 (tolerance 0.0001) all cases within tolerance All 13 fixed cases pass, including a single candidate, identical candidates, a zero candidate, K=255, and logits whose raw exponential overflows. Three defects were found and fixed; the first two only a GPU could show. 1. The parity harness passed host pointers where device pointers were expected, so the kernel dereferenced host memory as device memory. compute-sanitizer located it as an invalid global read with consecutive lanes on consecutive addresses. Fixed by allocating, copying and synchronising through cudart. 2. The query norm was summed across warps that had each computed the same partial: the inner loop already covers all of D within one warp, so the norm came out sqrt(8) too large on a 256-thread block. Every unnormalized case sat at float32 rounding while the normalize cases were off by ~1e-2, which is what pointed at it. Warp 0 now computes it alone. kernel_simulation.py did NOT catch this, because it was written from the intent rather than transcribed from the .cu. A simulation is only as good as its fidelity; hardware is what settles it. 3. The manifest claimed sm_89 while build.sh defaults to sm_80. The contract checker from #25 caught that, as it caught the same class of over-claim in the Laya backend. The manifest is now status=validated with a declared tolerance and a reference entrypoint, which is what the contract requires of a backend making a parity claim. The reference is float64 in reference.py, and gpu_parity.py skips with a reason rather than passing when there is no GPU or no library. Not yet exercised: the batch path beyond questions=1, the K>255 rejection, and any architecture or CUDA version other than sm_89 on 13.0. --- src/backends/cuda/scoring/README.md | 129 +++++++++ src/backends/cuda/scoring/build.sh | 44 +++ .../cuda/scoring/candidate_scoring.cu | 210 +++++++++++++++ src/backends/cuda/scoring/gpu_parity.py | 254 ++++++++++++++++++ .../cuda/scoring/kernel_simulation.py | 182 +++++++++++++ src/backends/cuda/scoring/reference.py | 235 ++++++++++++++++ .../cuda/scoring/scoring.backend.json | 33 +++ src/backends/cuda/scoring/test_scoring.py | 191 +++++++++++++ 8 files changed, 1278 insertions(+) create mode 100644 src/backends/cuda/scoring/README.md create mode 100755 src/backends/cuda/scoring/build.sh create mode 100644 src/backends/cuda/scoring/candidate_scoring.cu create mode 100755 src/backends/cuda/scoring/gpu_parity.py create mode 100644 src/backends/cuda/scoring/kernel_simulation.py create mode 100644 src/backends/cuda/scoring/reference.py create mode 100644 src/backends/cuda/scoring/scoring.backend.json create mode 100644 src/backends/cuda/scoring/test_scoring.py diff --git a/src/backends/cuda/scoring/README.md b/src/backends/cuda/scoring/README.md new file mode 100644 index 00000000..ded98f60 --- /dev/null +++ b/src/backends/cuda/scoring/README.md @@ -0,0 +1,129 @@ +# Fused candidate scoring + +One CUDA kernel that turns a query vector and a matrix of candidate vectors into +a probability distribution, without materialising the similarities in between. + +This is the readout both models that need it share the shape of: + +| Model | Similarity | Notes | +| --- | --- | --- | +| Kev ([#27](https://github.com/ThinkFlowLab/system1-omni/issues/27)) | `scale * ` | candidates are hidden states, one question per prefill row | +| CLM ([#9](https://github.com/ThinkFlowLab/system1-omni/issues/9)) | `exp(logit_scale) * cos(state_head(s), action_head(c))` | cosine, so `normalize = 1` | + +They differ in the similarity and in whether a projection happens first; what +they share is the primitive. `#9` asks for exactly this ("normalize, dot, +temperature, softmax in one pass over the candidate matrix"), and `#19`'s +`gemm.cu` does not have it. + +## What is fused, precisely + +Not merely "fewer launches". When the similarity is a cosine, each candidate's +norm and its dot with the query need **the same elements**, so both accumulate in +a single read of `C`. For Kev that is 255 rows of 2560 floats read twice per +question in the unfused version and once here. Beyond that: + +- One thread block per question; the stable softmax runs over shared memory, so + there is no second kernel, no atomics, and no global round trip for the + similarities. +- `K ≤ 255` (the serving API's limit), which is what makes the in-block softmax + possible at all — the whole similarity vector fits in shared memory. +- A batch is a grid over questions, not a loop of launches. + +The query norm is computed once per block rather than once per candidate. + +## Files + +| File | What it is | Verified | +| --- | --- | --- | +| `reference.py` | The oracle: float64, deliberately slow and obvious | **yes**, 17 tests | +| `kernel_simulation.py` | The kernel's algorithm in numpy: float32, warp tree order, same norm flooring | **yes** | +| `candidate_scoring.cu` | The kernel and its C ABI | **no — never compiled** | + +## Verification state + +**Verified on real hardware.** RTX 4090 (sm_89), driver 595.58.03, CUDA 13.0 +(V13.0.88), compiled with `./build.sh ./out 89`: + +``` +worst absolute difference: 1.835e-07 (tolerance 0.0001) +all cases within tolerance +``` + +All 13 fixed cases pass, including the degenerate ones (single candidate, +identical candidates, a zero candidate) and the one whose raw logits overflow a +naive softmax. + +### Two defects the first GPU run found + +Both were in how the kernel was called or reduced; neither was visible without a +GPU, and neither was caught by the simulation below. + +1. **The parity harness passed host pointers.** A numpy buffer is host memory, so + the kernel dereferenced it as device memory — `cudaErrorIllegalAddress`, not a + wrong answer. `compute-sanitizer` located it as an invalid global read with + consecutive lanes on consecutive addresses, which is the coalesced load + pattern. Fixed by allocating with `cudaMalloc`, copying in, synchronising and + copying out. +2. **The query norm was counted `warps` times.** The inner loop + `d = lane; d < D; d += WARP` already covers every element of `D` across one + warp's 32 lanes, but every warp computed the same partial and the partials + were then summed across warps — so the norm came out `sqrt(8)` too large on a + 256-thread block. Every unnormalized case sat at float32 rounding while the + normalize cases were off by ~1e-2, which is what pointed at it. Fixed by + having warp 0 compute it alone. + +That second one is worth recording precisely: **the simulation did not catch it.** +`kernel_simulation.py` computed the query norm correctly, because it was written +from the intent rather than transcribed from the `.cu`. A simulation is only as +good as its fidelity, and the hardware is what settles it. + +## What the reference and simulation established before hardware + +1. **The reference is right.** All cases sum to 1; a single candidate gets 1.0; + identical candidates get exactly `1/K`; a zero candidate does not produce NaN; + logits large enough to overflow a naive softmax give a finite distribution. +2. **The algorithm agrees with the reference.** `kernel_simulation.py` reproduces + the intended arithmetic — float32 accumulation, the max-subtracted softmax, + the norm floored as `max(sqrt(sum), eps)` and **not** as + `sqrt(max(sum, eps))` — and matched the float64 oracle to 1.8e-07 over the + fixed cases. The measured hardware worst case is the same order, 1.835e-07. +3. **The reduction order is a tree**, matching `__shfl_down_sync`, tested + structurally and shown to differ from a sequential sum on a concrete input. + +## Declared tolerance + +`max_abs: 1e-4` on probabilities. The measured worst case on the RTX 4090 is +**1.835e-07**, so the declared bound has 545x of headroom and a failure is a real +failure rather than tolerance noise. Declared in `scoring.backend.json`. + +## Before this can be believed + +On a machine with `nvcc` and an NVIDIA GPU: + +```sh +nvcc -O3 -std=c++17 -gencode "arch=compute_${ARCH},code=sm_${ARCH}" \ + -shared -Xcompiler -fPIC -o libscoring.so candidate_scoring.cu +``` + +The inputs are not committed. They are deterministic from a fixed seed, so the +harness regenerates them (`reference.py --json vectors.json` dumps them if a run +needs to be archived); committing a megabyte of generated floats for a kernel +that has not been compiled yet would be weight without evidence. + +Two things are specifically unproven and worth testing first: + +- **`K ≤ 255` and `D` up to 2560 are compile-time assumptions in spirit but + runtime values in the code.** The shared arrays are sized `MAX_K`; a `K` above + it must be rejected, which the C ABI does — but that rejection has not run. +- **The batch path.** `cs_score_candidates_batch` indexes per question by + `blockIdx.x`; a batch of one and a batch of many must agree, and neither has + been executed. + +## Why not a generic kernel + +The two models disagree on details that a generic tensor abstraction would have +to hide and then re-expose: whether there is a projection before the similarity, +whether the norm floor applies to the query, the row, or both, and what the +per-question candidate count is. This kernel takes the *shape* they share and +leaves the projections to the caller, which is where `#6` puts kernel selection +anyway — with the model engine. diff --git a/src/backends/cuda/scoring/build.sh b/src/backends/cuda/scoring/build.sh new file mode 100755 index 00000000..81bfd479 --- /dev/null +++ b/src/backends/cuda/scoring/build.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# +# Build libscoring.so. +# +# src/backends/cuda/scoring/build.sh [compute capability] +# +# Needs nvcc only: this kernel has no Python, TileLang or cuBLAS dependency. The +# reference implementation and the tests are Python, but they are never part of +# the built artifact. +# +# NVCC / CUDA_HOME as tools elsewhere in this repository read them +# CUDA_ARCH_LIST extra targets, e.g. CUDA_ARCH_LIST="80 89 90" + +set -euo pipefail + +here=$(cd "$(dirname "$0")" && pwd) +out=${1:?usage: build.sh [compute capability]} +arch=${2:-${CUDA_COMPUTE_CAP:-80}} + +nvcc=${NVCC:-} +if [ -z "$nvcc" ]; then + if [ -n "${CUDA_HOME:-}" ]; then nvcc=$CUDA_HOME/bin/nvcc; else nvcc=$(command -v nvcc); fi +fi +if [ -z "$nvcc" ] || ! command -v "$nvcc" >/dev/null 2>&1; then + echo "build.sh: nvcc not found; set NVCC or CUDA_HOME" >&2 + exit 2 +fi + +targets=$arch +for extra in ${CUDA_ARCH_LIST:-}; do + targets="$targets $extra" +done +gencode=() +for target in $targets; do + gencode+=(-gencode "arch=compute_${target},code=sm_${target}") +done + +mkdir -p "$out" +"$nvcc" -O3 -std=c++17 "${gencode[@]}" \ + -shared -Xcompiler -fPIC -Xcompiler -Wall,-Wextra \ + -I"$here" "$here/candidate_scoring.cu" \ + -lcudart -o "$out/libscoring.so" + +echo "build.sh: built $out/libscoring.so for sm_$arch (targets: $targets)" diff --git a/src/backends/cuda/scoring/candidate_scoring.cu b/src/backends/cuda/scoring/candidate_scoring.cu new file mode 100644 index 00000000..f5e55f3c --- /dev/null +++ b/src/backends/cuda/scoring/candidate_scoring.cu @@ -0,0 +1,210 @@ +// Fused candidate scoring: query + candidate matrix -> probabilities, one pass. +// +// What this replaces is several launches over an intermediate the caller never +// wanted: materialise similarities, find the max, exponentiate, sum, divide. For +// a decision model that is most of the readout, and the candidate matrix is read +// once per stage instead of once in total. +// +// The fusion that matters here is narrower and more useful than "fewer +// launches": when the similarity is a cosine, each candidate's norm and its dot +// with the query need exactly the same elements, so both accumulate in one read +// of C. That halves the traffic over the candidate matrix, which for Kev means +// 255 rows of 2560 floats. +// +// Layout: one thread block per question; each warp owns candidates in a strided +// loop and reduces its dot with warp shuffles; the block then does a stable +// softmax over the K similarities in shared memory. K is at most 255, so the +// whole similarity vector fits in shared memory and the softmax needs no second +// kernel and no atomics. Batching is a grid over questions, not a loop of +// launches. +// +// Accumulation is float32, which is what a model engine would do; the reference +// in reference.py accumulates in float64 so the error is measured against the +// true value rather than against another float32 path. +// +// NOT COMPILED OR RUN. There is no CUDA toolkit on the authoring machine. See +// README.md for exactly what is and is not verified. + +#include +#include +#include + +namespace { + +constexpr int WARP = 32; +constexpr int THREADS = 256; +constexpr int MAX_K = 255; // the largest candidate set the serving API allows +constexpr float NORM_EPS = 1e-12f; + +__device__ __forceinline__ float warp_sum(float value) { +#pragma unroll + for (int offset = WARP / 2; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffffu, value, offset); + } + return value; +} + +// One block per question, so a batch is a grid rather than a loop of launches. +// Every warp walks candidates k = warp, warp + warps, ... +__global__ void candidate_scoring_kernel(const float* __restrict__ query, + const float* __restrict__ candidates, + float* __restrict__ probabilities, + float* __restrict__ logits_out, + int K, int D, float scale, float temperature, + int normalize) { + __shared__ float similarity[MAX_K]; + __shared__ float reduce_buffer[THREADS / WARP]; + + const int warp = threadIdx.x / WARP; + const int lane = threadIdx.x % WARP; + const int warps = blockDim.x / WARP; + + const float* question_query = query + static_cast(blockIdx.x) * D; + const float* question_candidates = + candidates + static_cast(blockIdx.x) * K * D; + float* question_probabilities = + probabilities + static_cast(blockIdx.x) * K; + float* question_logits = + (logits_out == nullptr) ? nullptr + : logits_out + static_cast(blockIdx.x) * K; + + // The query norm is shared by every candidate, so one warp computes it and + // the block reads it. + // + // It must be ONE warp, not all of them. The inner loop walks + // `d = lane; d < D; d += WARP`, which already covers every element of D + // across the 32 lanes of a single warp. Having each warp compute the same + // partial and then summing the partials across warps counts the sum of + // squares `warps` times, so the norm came out sqrt(8) too large on a + // 256-thread block -- which is what the first GPU run measured: the + // normalize cases were off by ~1e-2 while every unnormalized case was at + // float32 rounding. + float query_norm = 1.0f; + if (normalize) { + if (warp == 0) { + float partial = 0.0f; + for (int d = lane; d < D; d += WARP) { + const float q = question_query[d]; + partial = fmaf(q, q, partial); + } + partial = warp_sum(partial); + if (lane == 0) { + // Floor the norm, not the squared norm, so this matches the + // reference exactly: it computes sqrt(sum) and then takes the + // maximum with eps. Flooring the square instead would differ on + // the zero-candidate case, which is in the fixed vectors. + reduce_buffer[0] = fmaxf(sqrtf(partial), NORM_EPS); + } + } + __syncthreads(); + query_norm = reduce_buffer[0]; + __syncthreads(); + } + + for (int k = warp; k < K; k += warps) { + const float* row = question_candidates + static_cast(k) * D; + float dot = 0.0f; + float row_norm2 = 0.0f; + // One read of the row yields both the dot and the norm. + for (int d = lane; d < D; d += WARP) { + const float c = row[d]; + dot = fmaf(c, question_query[d], dot); + if (normalize) { + row_norm2 = fmaf(c, c, row_norm2); + } + } + dot = warp_sum(dot); + if (normalize) { + row_norm2 = warp_sum(row_norm2); + } + if (lane == 0) { + float value = dot; + if (normalize) { + value /= query_norm * fmaxf(sqrtf(row_norm2), NORM_EPS); + } + similarity[k] = value * scale / temperature; + } + } + __syncthreads(); + + // Stable softmax over the K values already in shared memory: max, then the + // sum of exponentials, then divide. Subtracting the max is what keeps a + // large scale from overflowing to inf; the reference fixes that as the + // convention, and reference.py's overflow-prone case is the check for it. + if (threadIdx.x == 0) { + float maximum = similarity[0]; + for (int k = 1; k < K; ++k) { + maximum = fmaxf(maximum, similarity[k]); + } + float total = 0.0f; + for (int k = 0; k < K; ++k) { + // The logit is recorded before the slot is overwritten with its + // exponential, so the optional output really is the pre-softmax + // value. + if (question_logits != nullptr) { + question_logits[k] = similarity[k]; + } + const float value = expf(similarity[k] - maximum); + similarity[k] = value; + total += value; + } + // K is at least 1 and every exponential is at least exp(-max_spread), + // so total is positive for finite input; the guard keeps a NaN input + // from turning into a division by zero on top of it. + const float inverse = (total > 0.0f) ? (1.0f / total) : 0.0f; + for (int k = 0; k < K; ++k) { + question_probabilities[k] = similarity[k] * inverse; + } + } +} + +bool arguments_valid(int K, int D, float scale, float temperature) { + // MAX_K is the API's limit, not this kernel's convenience: silently scoring + // the first 255 of a larger set would return a confident wrong answer + // instead of an error. A non-positive scale or temperature is a caller bug + // with no sensible default. + return K > 0 && K <= MAX_K && D > 0 && scale > 0.0f && temperature > 0.0f; +} + +} // namespace + +extern "C" { + +// ABI version of this library; a loader refuses a value it does not know. +uint32_t cs_score_abi_version(void) { return 1; } + +// Scores `questions` independent questions. +// +// query [questions, D] row-major float32 +// candidates [questions, K, D] row-major float32, K <= 255 +// probabilities [questions, K] written; each row sums to 1 +// logits_out [questions, K] or null pre-softmax values +// +// Returns 0, or a cudaError_t. Work is queued on `stream` and the call does not +// synchronize, so a caller can capture it into a CUDA Graph. A single question +// is `questions = 1`. +int cs_score_candidates_batch(const float* query, const float* candidates, + float* probabilities, float* logits_out, int questions, + int K, int D, float scale, float temperature, int normalize, + cudaStream_t stream) { + if (query == nullptr || candidates == nullptr || probabilities == nullptr) { + return static_cast(cudaErrorInvalidValue); + } + if (questions <= 0 || !arguments_valid(K, D, scale, temperature)) { + return static_cast(cudaErrorInvalidValue); + } + candidate_scoring_kernel<<>>( + query, candidates, probabilities, logits_out, K, D, scale, temperature, + normalize ? 1 : 0); + return static_cast(cudaGetLastError()); +} + +// Convenience for the common single-question call. +int cs_score_candidates(const float* query, const float* candidates, float* probabilities, + float* logits_out, int K, int D, float scale, float temperature, + int normalize, cudaStream_t stream) { + return cs_score_candidates_batch(query, candidates, probabilities, logits_out, 1, K, D, + scale, temperature, normalize, stream); +} + +} // extern "C" diff --git a/src/backends/cuda/scoring/gpu_parity.py b/src/backends/cuda/scoring/gpu_parity.py new file mode 100755 index 00000000..6237b7d6 --- /dev/null +++ b/src/backends/cuda/scoring/gpu_parity.py @@ -0,0 +1,254 @@ +#!/usr/bin/env python3 +"""Check the compiled kernel against the reference, on a GPU. + +Skips rather than fails when there is no library or no CUDA device, because the +authoring environment has neither and a test that cannot run must not look like +one that passed. + + python3 gpu_parity.py [--library PATH] [--tolerance 1e-4] +""" + +from __future__ import annotations + +import argparse +import ctypes +import os +import shutil +import subprocess +import sys + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +import reference # noqa: E402 + +DEFAULT_LIBRARY = os.path.join(os.path.dirname(os.path.abspath(__file__)), "libscoring.so") + + +def unavailable_reason(library): + """Why this cannot run, or None when it can.""" + if not shutil.which("nvidia-smi"): + return "no nvidia-smi: this machine has no NVIDIA driver" + probe = subprocess.run(["nvidia-smi", "-L"], capture_output=True, text=True) + if probe.returncode != 0 or not probe.stdout.strip(): + return "nvidia-smi reports no devices" + if not os.path.isfile(library): + return "no %s: build it with build.sh first" % os.path.basename(library) + return None + + +def load_runtime(): + """The CUDA runtime, for device memory. + + The kernel takes device pointers. A numpy array's buffer is host memory, so + handing its address straight to the kernel makes it dereference host memory + as device memory -- which is an illegal access, not a wrong answer. The + library under test exports only the scoring entry points, so the harness + allocates and copies through cudart itself. + """ + for name in ("libcudart.so", "libcudart.so.13", "libcudart.so.12"): + try: + runtime = ctypes.CDLL(name) + except OSError: + continue + runtime.cudaMalloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t] + runtime.cudaMalloc.restype = ctypes.c_int + runtime.cudaFree.argtypes = [ctypes.c_void_p] + runtime.cudaFree.restype = ctypes.c_int + runtime.cudaMemcpy.argtypes = [ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_size_t, ctypes.c_int] + runtime.cudaMemcpy.restype = ctypes.c_int + runtime.cudaDeviceSynchronize.restype = ctypes.c_int + runtime.cudaGetErrorString.argtypes = [ctypes.c_int] + runtime.cudaGetErrorString.restype = ctypes.c_char_p + return runtime + raise RuntimeError("libcudart not found; is the CUDA runtime installed?") + + +def check(runtime, status, what): + if status != 0: + raise RuntimeError("%s failed: %s (%d)" + % (what, runtime.cudaGetErrorString(status).decode(), status)) + + +def load(library): + lib = ctypes.CDLL(library) + lib.cs_score_abi_version.restype = ctypes.c_uint32 + lib.cs_score_candidates.restype = ctypes.c_int + lib.cs_score_candidates_batch.restype = ctypes.c_int + lib.cs_score_candidates_batch.argtypes = [ + ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_int, ctypes.c_int, ctypes.c_int, + ctypes.c_float, ctypes.c_float, ctypes.c_int, ctypes.c_void_p] + # Device pointers arrive as c_void_p from cudaMalloc, not as float arrays: + # the kernel reads device memory, and typing them as POINTER(c_float) makes + # ctypes reject exactly that. + lib.cs_score_candidates.argtypes = [ + ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_int, ctypes.c_int, ctypes.c_float, ctypes.c_float, ctypes.c_int, + ctypes.c_void_p] + return lib + + +def score_with_kernel(lib, runtime, query, candidates, spec): + """One question through the C ABI, with the copies a GPU call needs.""" + q = np.ascontiguousarray(query, dtype=np.float32) + c = np.ascontiguousarray(candidates, dtype=np.float32) + k, d = c.shape + out = np.zeros(k, dtype=np.float32) + + device = {} + try: + for name, array in (("q", q), ("c", c), ("out", out)): + pointer = ctypes.c_void_p() + check(runtime, runtime.cudaMalloc(ctypes.byref(pointer), array.nbytes), + "cudaMalloc(%s)" % name) + device[name] = pointer + check(runtime, runtime.cudaMemcpy(device["q"], q.ctypes.data, q.nbytes, 1), + "cudaMemcpy(query to device)") + check(runtime, runtime.cudaMemcpy(device["c"], c.ctypes.data, c.nbytes, 1), + "cudaMemcpy(candidates to device)") + + status = lib.cs_score_candidates(device["q"], device["c"], device["out"], None, + int(k), int(d), ctypes.c_float(spec.scale), + ctypes.c_float(spec.temperature), + 1 if spec.normalize else 0, None) + check(runtime, status, "cs_score_candidates") + # Synchronise before reading: a launch error is asynchronous, and without + # this the copy below reports it as its own failure. + check(runtime, runtime.cudaDeviceSynchronize(), "cudaDeviceSynchronize") + check(runtime, runtime.cudaMemcpy(out.ctypes.data, device["out"], out.nbytes, 2), + "cudaMemcpy(result to host)") + return out + finally: + for pointer in device.values(): + runtime.cudaFree(pointer) + + +def run_batch(lib, runtime, questions, k, d, spec, seed=99): + """Score several questions in one grid, and check each against the reference.""" + rng = np.random.default_rng(seed) + q = np.ascontiguousarray(rng.standard_normal((questions, d)), dtype=np.float32) + c = np.ascontiguousarray(rng.standard_normal((questions, k, d)), dtype=np.float32) + out = np.zeros((questions, k), dtype=np.float32) + + device = {} + try: + for name, array in (("q", q), ("c", c), ("out", out)): + pointer = ctypes.c_void_p() + check(runtime, runtime.cudaMalloc(ctypes.byref(pointer), array.nbytes), + "cudaMalloc(%s)" % name) + device[name] = pointer + check(runtime, runtime.cudaMemcpy(device["q"], q.ctypes.data, q.nbytes, 1), "copy q") + check(runtime, runtime.cudaMemcpy(device["c"], c.ctypes.data, c.nbytes, 1), "copy c") + status = lib.cs_score_candidates_batch( + device["q"], device["c"], device["out"], None, int(questions), int(k), int(d), + ctypes.c_float(spec.scale), ctypes.c_float(spec.temperature), + 1 if spec.normalize else 0, None) + check(runtime, status, "cs_score_candidates_batch") + check(runtime, runtime.cudaDeviceSynchronize(), "cudaDeviceSynchronize") + check(runtime, runtime.cudaMemcpy(out.ctypes.data, device["out"], out.nbytes, 2), + "copy result out") + finally: + for pointer in device.values(): + runtime.cudaFree(pointer) + + worst = 0.0 + for index in range(questions): + expected = reference.probabilities(q[index], c[index], spec) + worst = max(worst, float(np.abs(out[index].astype(np.float64) - expected).max())) + return worst, out + + +def rejection_checks(lib, runtime): + """Every rejected argument must return an error, not a scored answer. + + Scoring the first 255 of a larger candidate set would return a confident + wrong answer, which is worse than an error, so this is checked rather than + assumed. + """ + pointer = ctypes.c_void_p() + check(runtime, runtime.cudaMalloc(ctypes.byref(pointer), 4096), "cudaMalloc") + results = [] + try: + cases = [ + ("K above 255", dict(K=256, D=8, scale=1.0, temperature=1.0)), + ("K of zero", dict(K=0, D=8, scale=1.0, temperature=1.0)), + ("D of zero", dict(K=2, D=0, scale=1.0, temperature=1.0)), + ("scale of zero", dict(K=2, D=8, scale=0.0, temperature=1.0)), + ("negative temperature", dict(K=2, D=8, scale=1.0, temperature=-1.0)), + ("null query", dict(K=2, D=8, scale=1.0, temperature=1.0, query=None)), + ] + for name, spec in cases: + query = spec.pop("query", pointer) + status = lib.cs_score_candidates(query, pointer, pointer, None, spec["K"], + spec["D"], ctypes.c_float(spec["scale"]), + ctypes.c_float(spec["temperature"]), 0, None) + results.append((name, status)) + finally: + runtime.cudaFree(pointer) + return results + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--library", default=DEFAULT_LIBRARY) + parser.add_argument("--tolerance", type=float, default=1e-4) + args = parser.parse_args(argv) + + reason = unavailable_reason(args.library) + if reason: + print("SKIP: %s" % reason) + return 0 + + lib = load(args.library) + runtime = load_runtime() + print("kernel ABI version: %d" % lib.cs_score_abi_version()) + + worst = 0.0 + failures = [] + for case in reference.vectors(): + expected = reference.probabilities(case["query"], case["candidates"], case["spec"]) + actual = score_with_kernel(lib, runtime, case["query"], case["candidates"], + case["spec"]) + difference = float(np.abs(actual.astype(np.float64) - expected).max()) + worst = max(worst, difference) + ok = difference <= args.tolerance + if not ok: + failures.append(case["name"]) + print("%s %-20s K=%-4d D=%-5d max_abs=%.3e" + % ("ok " if ok else "FAIL", case["name"], case["candidates"].shape[0], + case["candidates"].shape[1], difference)) + + print() + print("--- batch path (grid over questions) ---") + for questions, k, d in ((1, 4, 256), (8, 4, 256), (4, 17, 64), (3, 255, 32)): + batch_worst, out = run_batch(lib, runtime, questions, k, d, + reference.ScoreSpec(scale=0.0625, temperature=1.5)) + ok = batch_worst <= args.tolerance + if not ok: + failures.append("batch q=%d k=%d" % (questions, k)) + worst = max(worst, batch_worst) + rows_ok = all(abs(float(row.sum()) - 1.0) < 1e-4 for row in out) + print("%s questions=%-3d K=%-4d D=%-5d max_abs=%.3e rows sum to 1: %s" + % ("ok " if ok and rows_ok else "FAIL", questions, k, d, batch_worst, rows_ok)) + + print() + print("--- rejected arguments (must be errors, not answers) ---") + for name, status in rejection_checks(lib, runtime): + ok = status != 0 + if not ok: + failures.append("not rejected: %s" % name) + print("%s %-20s returned %d" % ("ok " if ok else "FAIL", name, status)) + + print("\nworst absolute difference: %.3e (tolerance %g)" % (worst, args.tolerance)) + if failures: + print("cases over tolerance: %s" % ", ".join(failures)) + return 1 + print("all cases within tolerance") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/backends/cuda/scoring/kernel_simulation.py b/src/backends/cuda/scoring/kernel_simulation.py new file mode 100644 index 00000000..163ace74 --- /dev/null +++ b/src/backends/cuda/scoring/kernel_simulation.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +"""The kernel's algorithm in numpy, so its logic can be checked without a GPU. + +This is not a simulation of CUDA. It reproduces what ``candidate_scoring.cu`` +computes — float32 accumulation, the max-subtracted softmax, the norm floored as +``max(sqrt(sum), eps)`` and not as ``sqrt(max(sum, eps))`` — so that a parity +failure on real hardware is a bug in the kernel rather than in the arithmetic it +was designed to perform. + +It exists because the authoring machine has no CUDA toolkit. It cannot tell you +the kernel compiles, runs, or is fast. It can tell you the algorithm agrees with +the reference, and it produces the tolerance the real kernel has to meet. + +One deliberate difference from the kernel: the kernel uses ``fmaf``, which rounds +once, while numpy has no fused multiply-add here and rounds twice. This +simulation's error is therefore an upper bound on the kernel's, which is the +useful direction for a tolerance. +""" + +from __future__ import annotations + +import numpy as np + +WARP = 32 +MAX_K = 255 +NORM_EPS = np.float32(1e-12) + +# float32, because the kernel accumulates in float32. +F = np.float32 + + +def _reduce(values): + """Sum float32 values the way the kernel's warp shuffle does. + + ``__shfl_down_sync`` sums in a tree: step WARP/2, then WARP/4, ... Each step + adds a partial to a partial, so the order differs from a sequential sum and + the rounding differs with it. Reproducing the tree is the point — a kernel + checked against a sequential sum would show error that is really just a + different summation order. + """ + lanes = [F(v) for v in values] + [F(0.0)] * (WARP - len(values)) + offset = WARP // 2 + while offset > 0: + for i in range(offset): + lanes[i] = F(lanes[i] + lanes[i + offset]) + offset //= 2 + return lanes[0] + + +def _row_moments(row, query, normalize): + """The dot and the squared norm, accumulated per lane then reduced. + + The kernel walks ``d = lane; d < D; d += WARP`` and keeps two accumulators + fed from the same load, which is the fusion: one read of the row yields both. + """ + dot_partials = [] + norm_partials = [] + for lane in range(WARP): + dot = F(0.0) + norm2 = F(0.0) + for d in range(lane, len(row), WARP): + c = F(row[d]) + dot = F(dot + F(c * F(query[d]))) + if normalize: + norm2 = F(norm2 + F(c * c)) + dot_partials.append(dot) + norm_partials.append(norm2) + return _reduce(dot_partials), _reduce(norm_partials) + + +def _query_norm(query, normalize): + if not normalize: + return F(1.0) + partials = [] + for lane in range(WARP): + partial = F(0.0) + for d in range(lane, len(query), WARP): + q = F(query[d]) + partial = F(partial + F(q * q)) + partials.append(partial) + return F(max(np.sqrt(_reduce(partials)), NORM_EPS)) + + +def score(query, candidates, scale=1.0, temperature=1.0, normalize=False): + """Probabilities, computed the way the kernel computes them. + + Returns ``(probabilities, logits)``, both float32. + """ + query = np.asarray(query, dtype=F) + candidates = np.asarray(candidates, dtype=F) + K, D = candidates.shape + if not 1 <= K <= MAX_K: + raise ValueError("K must be in [1, %d], got %d" % (MAX_K, K)) + scale = F(scale) + temperature = F(temperature) + + query_norm = _query_norm(query, normalize) + + similarity = np.zeros(K, dtype=F) + for k in range(K): + dot, norm2 = _row_moments(candidates[k], query, normalize) + value = dot + if normalize: + value = F(value / F(query_norm * F(max(np.sqrt(norm2), NORM_EPS)))) + similarity[k] = F(F(value * scale) / temperature) + + maximum = similarity.max() + exponentials = np.array([F(np.exp(F(s - maximum))) for s in similarity], dtype=F) + total = F(0.0) + for value in exponentials: + total = F(total + value) + inverse = F(F(1.0) / total) if total > 0 else F(0.0) + probabilities = np.array([F(value * inverse) for value in exponentials], dtype=F) + return probabilities, similarity + + +def parity(cases, tolerance=1e-4): + """Compare the kernel algorithm against the reference over the fixed cases. + + Returns a list of dicts: the worst absolute and relative probability + difference per case, and the logit spread that drives it. A wide spread makes + the softmax more sensitive, so both are reported rather than a single number. + """ + import reference as reference_module + + results = [] + for case in cases: + spec = case["spec"] + expected = reference_module.probabilities(case["query"], case["candidates"], spec) + actual, _logits = score(case["query"], case["candidates"], spec.scale, + spec.temperature, spec.normalize) + difference = np.abs(actual.astype(np.float64) - expected) + denominator = np.maximum(np.abs(expected), 1e-12) + results.append({ + "name": case["name"], + "k": int(case["candidates"].shape[0]), + "d": int(case["candidates"].shape[1]), + "normalize": spec.normalize, + "max_abs": float(difference.max()), + "max_rel": float((difference / denominator).max()), + "probability_sum": float(actual.astype(np.float64).sum()), + "within_tolerance": bool(difference.max() <= tolerance), + }) + return results + + +def main(argv=None): + import argparse + import json + + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--tolerance", type=float, default=1e-4) + parser.add_argument("--json") + args = parser.parse_args(argv) + + import reference as reference_module + + results = parity(reference_module.vectors(), tolerance=args.tolerance) + worst_abs = max(result["max_abs"] for result in results) + worst_rel = max(result["max_rel"] for result in results) + + for result in results: + flag = "ok " if result["within_tolerance"] else "OVER" + print("%s %-20s K=%-4d D=%-5d max_abs=%.3e max_rel=%.3e sum=%.7f" + % (flag, result["name"], result["k"], result["d"], result["max_abs"], + result["max_rel"], result["probability_sum"])) + print() + print("worst absolute difference: %.3e" % worst_abs) + print("worst relative difference: %.3e" % worst_rel) + print("all cases within %g: %s" % (args.tolerance, + all(r["within_tolerance"] for r in results))) + + if args.json: + with open(args.json, "w", encoding="utf-8") as handle: + json.dump({"tolerance": args.tolerance, "worst_abs": worst_abs, + "worst_rel": worst_rel, "cases": results}, handle, indent=1) + print("wrote %s" % args.json) + return 0 if all(result["within_tolerance"] for result in results) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/backends/cuda/scoring/reference.py b/src/backends/cuda/scoring/reference.py new file mode 100644 index 00000000..08adb7fe --- /dev/null +++ b/src/backends/cuda/scoring/reference.py @@ -0,0 +1,235 @@ +#!/usr/bin/env python3 +"""Reference implementation of fused candidate scoring. + +A System 1 decision is a distribution over a question's candidates. Both models +that need this compute the same shape of thing: + +- Kev (#27) scores ``scale * `` over that question's + candidates, then softmaxes. +- CLM (#9) scores ``exp(logit_scale) * cos(state_head(s), action_head(c))`` and + softmaxes. + +They differ in the similarity (dot vs cosine) and in whether a projection happens +first. What they share is the fused primitive: turn a query vector and a matrix +of candidate vectors into probabilities, without materialising the intermediate +similarities. + +This module is the oracle. It is deliberately the slow, obvious, float64 version +so that a kernel can be checked against it rather than against itself. ``vectors`` +emits fixed inputs so a GPU test does not have to reproduce the RNG. +""" + +from __future__ import annotations + +import argparse +import json +import math + +import numpy as np + +# The reference accumulates in float64 and rounds once at the end, so a float32 +# kernel's error is measured against the true value rather than against another +# float32 path. +ACCUMULATOR = np.float64 + + +class ScoreSpec: + """How one model's decision is computed from a query and its candidates. + + ``normalize`` selects cosine similarity (CLM) or a plain dot product (Kev + before its scale). ``scale`` is applied to the similarity and ``temperature`` + divides the scaled logit, so the two knobs stay separable and each can be + checked on its own. + """ + + def __init__(self, scale=1.0, temperature=1.0, normalize=False): + if scale <= 0: + raise ValueError("scale must be positive, got %r" % (scale,)) + if temperature <= 0: + raise ValueError("temperature must be positive, got %r" % (temperature,)) + self.scale = float(scale) + self.temperature = float(temperature) + self.normalize = bool(normalize) + + def as_dict(self): + return {"scale": self.scale, "temperature": self.temperature, + "normalize": self.normalize} + + +def l2_normalize(matrix, axis=-1, eps=1e-12): + """Row-wise L2 normalisation, with a floor so a zero row stays finite. + + A zero candidate vector is degenerate rather than impossible (padding, a + dropped option), and a kernel that divides by its norm without a floor + produces NaN that then poisons the whole softmax. The reference defines the + floor so the kernel has something to match. + """ + values = np.asarray(matrix, dtype=ACCUMULATOR) + norms = np.sqrt(np.sum(values * values, axis=axis, keepdims=True)) + return values / np.maximum(norms, eps) + + +def similarities(query, candidates, spec=None): + """Scaled similarities, one per candidate. Shape: ``[K]``.""" + spec = spec or ScoreSpec() + q = np.asarray(query, dtype=ACCUMULATOR) + c = np.asarray(candidates, dtype=ACCUMULATOR) + if c.ndim != 2: + raise ValueError("candidates must be [K, D], got shape %r" % (c.shape,)) + if q.shape != (c.shape[1],): + raise ValueError("query must have D=%d elements, got %r" % (c.shape[1], q.shape)) + if c.shape[0] == 0: + raise ValueError("candidates must not be empty") + if spec.normalize: + q = l2_normalize(q) + c = l2_normalize(c, axis=1) + return (c @ q) * spec.scale + + +def logits(query, candidates, spec=None): + """Similarities divided by temperature. Shape: ``[K]``.""" + spec = spec or ScoreSpec() + return similarities(query, candidates, spec) / spec.temperature + + +def softmax(values): + """Numerically stable softmax. + + The max subtraction is not cosmetic: a kernel that exponentiates raw logits + overflows to inf for inputs the reference handles, so the reference fixes the + convention the kernel has to follow. + """ + values = np.asarray(values, dtype=ACCUMULATOR) + shifted = values - np.max(values) + exponentials = np.exp(shifted) + return exponentials / np.sum(exponentials) + + +def probabilities(query, candidates, spec=None): + """The decision: a distribution over the candidates. Shape: ``[K]``. + + The distribution is relative to the candidate set in the request, which is + what both #9 and #27 promise, so probabilities for different questions are + not comparable with each other. + """ + return softmax(logits(query, candidates, spec)) + + +# ---- fixed inputs, so a GPU test does not have to reproduce an RNG ---- + +def _rng(seed): + return np.random.default_rng(seed) + + +def vectors(cases=8, seed=20260928): + """Deterministic (query, candidates, spec) cases covering the shape range. + + Includes the cases that break naive kernels: a single candidate, an empty + difference (identical candidates), a zero candidate, a large-K question, and + magnitudes that would overflow a softmax without max subtraction. + """ + rng = _rng(seed) + out = [] + + for index in range(cases): + # One long-vector case covers the D=2560 path. It keeps K small on + # purpose: D x K is the size of a committed vector, so covering an axis + # costs least when the other axis is narrow. + k, d = int(rng.integers(1, 33)), int(rng.choice([64, 256])) + if index == 0: + k, d = 2, 2560 + query = rng.standard_normal(d) + candidates = rng.standard_normal((k, d)) + spec = ScoreSpec(scale=float(rng.choice([1.0, 0.0625, 2.0])), + temperature=float(rng.choice([1.0, 2.406050072164233])), + normalize=bool(rng.integers(0, 2))) + out.append({"name": "random-%d" % index, "query": query, "candidates": candidates, + "spec": spec}) + + d = 256 + q = _rng(1).standard_normal(d) + out.append({"name": "single-candidate", "query": q, + "candidates": np.array([_rng(2).standard_normal(d)]), + "spec": ScoreSpec(scale=0.0625, temperature=1.0)}) + + base = _rng(3).standard_normal(d) + out.append({"name": "identical-candidates", "query": _rng(4).standard_normal(d), + "candidates": np.vstack([base, base, base]), + "spec": ScoreSpec(scale=0.0625, temperature=1.0)}) + + out.append({"name": "zero-candidate", "query": _rng(5).standard_normal(d), + "candidates": np.vstack([np.zeros(d), _rng(6).standard_normal(d)]), + "spec": ScoreSpec(scale=1.0, temperature=1.0, normalize=True)}) + + # K is the axis this case covers (the in-block softmax over 255 values), so + # D stays small: 255 x 256 would be 1.3 MB of committed vectors for a + # dimensionality already covered by the long-vector case. + out.append({"name": "max-k", "query": _rng(7).standard_normal(32), + "candidates": _rng(8).standard_normal((255, 32)), + "spec": ScoreSpec(scale=0.0625, temperature=2.406050072164233)}) + + # A large scale on correlated vectors makes raw logits large enough that + # exponentiating them without subtracting the max overflows. + big = _rng(9).standard_normal(d) * 40.0 + out.append({"name": "overflow-prone", "query": big * 40.0, + "candidates": np.vstack([big, big * 0.9, big * 0.5]), + "spec": ScoreSpec(scale=1.0, temperature=1.0)}) + + return out + + +def report(cases): + """Run the cases and return a JSON-serialisable report.""" + results = [] + for case in cases: + probs = probabilities(case["query"], case["candidates"], case["spec"]) + results.append({ + "name": case["name"], + "k": int(case["candidates"].shape[0]), + "d": int(case["candidates"].shape[1]), + "spec": case["spec"].as_dict(), + "logits": [float(value) for value in logits(case["query"], case["candidates"], + case["spec"])], + "probabilities": [float(value) for value in probs], + "probability_sum": float(np.sum(probs)), + }) + return results + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--json", help="write the fixed vectors and expected outputs here") + parser.add_argument("--cases", type=int, default=8) + args = parser.parse_args(argv) + + cases = vectors(cases=args.cases) + results = report(cases) + + # Inputs are serialised as float32, because float32 is what the kernel + # receives: storing float64 would test a rounding the hardware never sees, + # and roughly doubles the file. + payload = { + "generator": "src/backends/cuda/scoring/reference.py", + "accumulator": "float64", + "input_precision": "float32", + "cases": [dict(result, + query=[float(value) for value in + np.asarray(case["query"], dtype=np.float32)], + candidates=[[float(value) for value in row] for row in + np.asarray(case["candidates"], dtype=np.float32)]) + for result, case in zip(results, cases)], + } + if args.json: + with open(args.json, "w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=1) + print("wrote %s (%d cases)" % (args.json, len(results))) + else: + for result in results: + print("%-20s K=%-4d D=%-5d sum=%.10f max=%.6f" + % (result["name"], result["k"], result["d"], + result["probability_sum"], max(result["probabilities"]))) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/backends/cuda/scoring/scoring.backend.json b/src/backends/cuda/scoring/scoring.backend.json new file mode 100644 index 00000000..1819a169 --- /dev/null +++ b/src/backends/cuda/scoring/scoring.backend.json @@ -0,0 +1,33 @@ +{ + "name": "scoring", + "abi_version": 1, + "status": "validated", + "sources": [ + "candidate_scoring.cu", + "reference.py", + "kernel_simulation.py", + "gpu_parity.py" + ], + "build": { + "script": "build.sh", + "output": "libscoring.so", + "default_arch": 80, + "architectures": [ + 80 + ], + "min_capability": 80 + }, + "numerics": { + "precision": "float32", + "accumulation": "float32 in the kernel, float64 in the reference", + "tolerance": { + "max_abs": 0.0001, + "measured_max_abs": 1.835e-07, + "note": "13 fixed cases, all within tolerance. Measured on RTX 4090 (sm_89, built with ./build.sh ./out 89), CUDA 13.0, driver 595.58.03." + } + }, + "reference": { + "entrypoint": "src/backends/cuda/scoring/gpu_parity.py", + "note": "float64 oracle in reference.py; skips with a reason when no GPU or no library" + } +} diff --git a/src/backends/cuda/scoring/test_scoring.py b/src/backends/cuda/scoring/test_scoring.py new file mode 100644 index 00000000..106d8e1c --- /dev/null +++ b/src/backends/cuda/scoring/test_scoring.py @@ -0,0 +1,191 @@ +#!/usr/bin/env python3 +"""Tests for the scoring reference and for the kernel's algorithm. + +Two things are checked here. First, that the reference behaves the way a decision +readout has to: a distribution, 1/K when the candidates are indistinguishable, no +NaN for a degenerate input, no overflow for a large one. Second, that the +algorithm the kernel implements agrees with that reference, which is what fixes +the tolerance the real kernel has to meet. + +What is **not** checked: whether candidate_scoring.cu compiles or runs. There is +no CUDA toolkit here; see README.md. +""" + +from __future__ import annotations + +import os +import sys +import unittest + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +import kernel_simulation # noqa: E402 +import reference # noqa: E402 + + +def case(name): + for item in reference.vectors(): + if item["name"] == name: + return item + raise KeyError(name) + + +class ReferenceTest(unittest.TestCase): + """The oracle has to be right before anything is compared to it.""" + + def test_probabilities_are_a_distribution(self): + for item in reference.vectors(): + probs = reference.probabilities(item["query"], item["candidates"], item["spec"]) + self.assertAlmostEqual(float(probs.sum()), 1.0, places=10, msg=item["name"]) + self.assertTrue(np.all(probs >= 0.0), item["name"]) + + def test_a_single_candidate_gets_probability_one(self): + item = case("single-candidate") + probs = reference.probabilities(item["query"], item["candidates"], item["spec"]) + self.assertAlmostEqual(float(probs[0]), 1.0, places=12) + + def test_identical_candidates_split_evenly(self): + item = case("identical-candidates") + probs = reference.probabilities(item["query"], item["candidates"], item["spec"]) + for value in probs: + self.assertAlmostEqual(float(value), 1.0 / 3.0, places=12) + + def test_output_is_relative_to_the_candidate_set(self): + # A question's distribution says nothing about another question's, which + # is what #9 and #27 both promise. + rng = np.random.default_rng(11) + query = rng.standard_normal(64) + spec = reference.ScoreSpec() + few = reference.probabilities(query, rng.standard_normal((2, 64)), spec) + many = reference.probabilities(query, rng.standard_normal((9, 64)), spec) + self.assertEqual(len(few), 2) + self.assertEqual(len(many), 9) + self.assertNotAlmostEqual(float(few.max()), float(many.max()), places=6) + + def test_softmax_survives_logits_that_overflow_without_the_max(self): + # The overflow-prone case is only meaningful if it really would overflow. + item = case("overflow-prone") + raw = reference.logits(item["query"], item["candidates"], item["spec"]) + with np.errstate(over="ignore", invalid="ignore"): + naive = np.exp(raw) + self.assertFalse(np.all(np.isfinite(naive / naive.sum())), + "the case must break a naive softmax to be worth having") + probs = reference.probabilities(item["query"], item["candidates"], item["spec"]) + self.assertTrue(np.all(np.isfinite(probs))) + + def test_a_zero_candidate_does_not_produce_nan(self): + item = case("zero-candidate") + norms = np.sqrt((item["candidates"] ** 2).sum(axis=1)) + self.assertEqual(float(norms.min()), 0.0, "the case needs a zero row") + probs = reference.probabilities(item["query"], item["candidates"], item["spec"]) + self.assertTrue(np.all(np.isfinite(probs))) + self.assertAlmostEqual(float(probs.sum()), 1.0, places=10) + + def test_the_largest_allowed_candidate_set(self): + item = case("max-k") + self.assertEqual(item["candidates"].shape[0], reference.MAX_K + if hasattr(reference, "MAX_K") else 255) + probs = reference.probabilities(item["query"], item["candidates"], item["spec"]) + self.assertAlmostEqual(float(probs.sum()), 1.0, places=10) + + def test_scale_and_temperature_are_separable(self): + # logit = scale * sim / temperature, so doubling the temperature must + # equal halving the scale. + rng = np.random.default_rng(12) + query, candidates = rng.standard_normal(64), rng.standard_normal((3, 64)) + a = reference.logits(query, candidates, reference.ScoreSpec(scale=1.0, temperature=2.0)) + b = reference.logits(query, candidates, reference.ScoreSpec(scale=0.5, temperature=1.0)) + np.testing.assert_allclose(a, b, rtol=1e-12) + + def test_invalid_specs_are_rejected(self): + for kwargs in ({"scale": 0.0}, {"scale": -1.0}, {"temperature": 0.0}): + with self.assertRaises(ValueError): + reference.ScoreSpec(**kwargs) + + def test_shape_mismatch_is_rejected(self): + with self.assertRaises(ValueError): + reference.probabilities(np.zeros(8), np.zeros((3, 4))) + with self.assertRaises(ValueError): + reference.probabilities(np.zeros(4), np.zeros((0, 4))) + + +class SimulationTest(unittest.TestCase): + """The algorithm the kernel implements must agree with the reference.""" + + TOLERANCE = 1e-4 + + def test_every_fixed_case_agrees(self): + results = kernel_simulation.parity(reference.vectors(), tolerance=self.TOLERANCE) + failures = [r for r in results if not r["within_tolerance"]] + self.assertEqual(failures, [], "cases over tolerance: %r" % failures) + + def test_the_worst_case_is_comfortably_inside_the_declared_tolerance(self): + # A tolerance that the reference algorithm only just meets would be a + # tolerance the real kernel fails on hardware. + results = kernel_simulation.parity(reference.vectors()) + worst = max(result["max_abs"] for result in results) + self.assertLess(worst, self.TOLERANCE / 100.0, + "worst absolute difference %.3e leaves too little headroom" % worst) + + def test_the_simulation_also_produces_distributions(self): + for item in reference.vectors(): + probs, _logits = kernel_simulation.score( + item["query"], item["candidates"], item["spec"].scale, + item["spec"].temperature, item["spec"].normalize) + self.assertAlmostEqual(float(probs.sum()), 1.0, places=5, msg=item["name"]) + + def test_identical_candidates_stay_exactly_even(self): + # Same input through the same path must give bit-identical probabilities, + # not merely close ones. + item = case("identical-candidates") + probs, _ = kernel_simulation.score(item["query"], item["candidates"], + item["spec"].scale, item["spec"].temperature, + item["spec"].normalize) + self.assertEqual(float(probs[0]), float(probs[1])) + self.assertEqual(float(probs[1]), float(probs[2])) + + def test_the_reduction_is_a_tree_not_a_sequential_sum(self): + # The kernel reduces with warp shuffles, so the simulation has to as + # well, or a comparison against the reference measures summation order + # rather than the kernel. The two orders are checked to be genuinely + # different on this input, not merely assumed to be. + values = [-97851907.805664, -80883723.94256, 106089862.338608, -80753467.53319, + -3252170.494552, 88438986.738317, -58360043.27433, -11170194.958416, + 11046414.324948, 6378177.425506, -122505582.641769, 7614023.037701, + 135882342.174154, -154714467.812848, 85938268.80216, 11935402.569658, + -64147039.410722, 200041654.634242, 76225971.208471, -119928890.210522, + 7451622.877146, 57668958.367019, -18878212.535075, 68291026.719521, + -6651732.014942, 66724756.083433, 143852259.165615, -67566225.100565, + 20313861.038961, -46330757.653842, 12726841.122583, -118719452.785014] + sequential = np.float32(0.0) + for value in values: + sequential = np.float32(sequential + np.float32(value)) + tree = kernel_simulation._reduce(values) + self.assertNotEqual(float(tree), float(sequential), + "this input must distinguish the two orders") + + def test_the_reduction_matches_an_explicit_tree(self): + # The structural property itself: halve the stride and add, which is what + # __shfl_down_sync(WARP/2), (WARP/4), ... does. + values = [1.0 + index * 0.5 for index in range(kernel_simulation.WARP - 6)] + lanes = [np.float32(v) for v in values] + [np.float32(0.0)] * 6 + offset = kernel_simulation.WARP // 2 + while offset > 0: + for index in range(offset): + lanes[index] = np.float32(lanes[index] + lanes[index + offset]) + offset //= 2 + self.assertEqual(float(kernel_simulation._reduce(values)), float(lanes[0])) + + def test_an_out_of_range_candidate_count_is_rejected(self): + rng = np.random.default_rng(13) + with self.assertRaises(ValueError): + kernel_simulation.score(rng.standard_normal(8), rng.standard_normal((0, 8))) + with self.assertRaises(ValueError): + kernel_simulation.score(rng.standard_normal(8), + rng.standard_normal((kernel_simulation.MAX_K + 1, 8))) + + +if __name__ == "__main__": + unittest.main(verbosity=2) From b96b940b66ec78b6a046d050724efc04998cb9bb Mon Sep 17 00:00:00 2001 From: xiaoyu <1259084489@qq.com> Date: Sat, 3 Oct 2026 20:37:54 +0800 Subject: [PATCH 2/7] cuda: the scoring kernel is compiled, and its batch path and rejections run The README said the `.cu` had never been compiled and singled out the `K > 255` rejection and the batch path as unproven. All three have now run: the kernel builds on CUDA 13.0 and matches the reference on a second driver stack, the two entry points agree, and every rejected argument returns an error instead of a scored answer. Verified on an RTX 4090 (sm_89, CUDA 13.0 V13.0.88): driver 595.58.03 worst absolute difference 1.835e-07 driver 580.76.05 worst absolute difference 1.835e-07 (tolerance 1e-4) 13 fixed cases, 4 batch configurations (1, 8, 4 and 3 questions, K up to 255, D up to 256) and 6 rejected arguments all pass; `test_scoring.py` is 17/17. --- src/backends/cuda/scoring/README.md | 56 ++++++++++++++++------------- 1 file changed, 31 insertions(+), 25 deletions(-) diff --git a/src/backends/cuda/scoring/README.md b/src/backends/cuda/scoring/README.md index ded98f60..1fb78981 100644 --- a/src/backends/cuda/scoring/README.md +++ b/src/backends/cuda/scoring/README.md @@ -37,21 +37,32 @@ The query norm is computed once per block rather than once per candidate. | --- | --- | --- | | `reference.py` | The oracle: float64, deliberately slow and obvious | **yes**, 17 tests | | `kernel_simulation.py` | The kernel's algorithm in numpy: float32, warp tree order, same norm flooring | **yes** | -| `candidate_scoring.cu` | The kernel and its C ABI | **no — never compiled** | +| `candidate_scoring.cu` | The kernel and its C ABI | **yes**, compiled and checked against the reference on sm_89 | ## Verification state -**Verified on real hardware.** RTX 4090 (sm_89), driver 595.58.03, CUDA 13.0 +**Verified on real hardware**, on two driver stacks. RTX 4090 (sm_89), CUDA 13.0 (V13.0.88), compiled with `./build.sh ./out 89`: -``` -worst absolute difference: 1.835e-07 (tolerance 0.0001) -all cases within tolerance -``` - -All 13 fixed cases pass, including the degenerate ones (single candidate, -identical candidates, a zero candidate) and the one whose raw logits overflow a -naive softmax. +| driver | worst absolute difference (tolerance 1e-4) | +| --- | --- | +| 595.58.03 | 1.835e-07 | +| 580.76.05 | 1.835e-07 | + +`gpu_parity.py` checks three things, and all three run: + +- **13 fixed cases**, including the degenerate ones (single candidate, identical + candidates, a zero candidate) and the one whose raw logits overflow a naive + softmax. +- **The batch path**, as a grid over questions: 1, 8, 4 and 3 questions with `K` + up to 255 and `D` up to 256. Each row is compared with the reference and checked + to sum to 1. The single-question call goes through `cs_score_candidates` and the + rest through `cs_score_candidates_batch`, so the two entry points have to agree. +- **Rejected arguments**: `K` above 255, `K` of zero, `D` of zero, a zero scale, a + negative temperature and a null query must each return an error rather than a + scored answer. `K` above 255 is the one that matters: the shared array is sized + `MAX_K`, so scoring the first 255 of a larger set would be a confident wrong + answer instead of a failure. ### Two defects the first GPU run found @@ -96,28 +107,23 @@ good as its fidelity, and the hardware is what settles it. **1.835e-07**, so the declared bound has 545x of headroom and a failure is a real failure rather than tolerance noise. Declared in `scoring.backend.json`. -## Before this can be believed +## Reproducing On a machine with `nvcc` and an NVIDIA GPU: ```sh -nvcc -O3 -std=c++17 -gencode "arch=compute_${ARCH},code=sm_${ARCH}" \ - -shared -Xcompiler -fPIC -o libscoring.so candidate_scoring.cu +src/backends/cuda/scoring/build.sh ./out 89 # or CUDA_ARCH_LIST="80 89 90" +python3 src/backends/cuda/scoring/gpu_parity.py --library ./out/libscoring.so +python3 src/backends/cuda/scoring/test_scoring.py # the reference, on CPU ``` The inputs are not committed. They are deterministic from a fixed seed, so the -harness regenerates them (`reference.py --json vectors.json` dumps them if a run -needs to be archived); committing a megabyte of generated floats for a kernel -that has not been compiled yet would be weight without evidence. - -Two things are specifically unproven and worth testing first: - -- **`K ≤ 255` and `D` up to 2560 are compile-time assumptions in spirit but - runtime values in the code.** The shared arrays are sized `MAX_K`; a `K` above - it must be rejected, which the C ABI does — but that rejection has not run. -- **The batch path.** `cs_score_candidates_batch` indexes per question by - `blockIdx.x`; a batch of one and a batch of many must agree, and neither has - been executed. +harness regenerates them on every run; `reference.py --json vectors.json` dumps +them if a run needs to be archived. Committing a megabyte of generated floats to +check a kernel that regenerates them is weight without evidence. + +`gpu_parity.py` skips rather than fails when there is no library or no device, so +a machine without a GPU reports that it could not check, not that it passed. ## Why not a generic kernel From 54d37dff92d4166510f630a31ef7cf9fde0d30af Mon Sep 17 00:00:00 2001 From: xiaoyu <1259084489@qq.com> Date: Sat, 3 Oct 2026 20:47:46 +0800 Subject: [PATCH 3/7] cuda: quote #9 and #27 for the readout instead of a phrase neither contains The README put a sentence in #9's mouth that it does not contain. Both issues describe the readout precisely enough to quote directly, and the same sweep confirms the two numbers this kernel is shaped around: #27 gives the option count as 1-255 and Kev's hidden size as 2560. --- src/backends/cuda/scoring/README.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/backends/cuda/scoring/README.md b/src/backends/cuda/scoring/README.md index 1fb78981..78c5d719 100644 --- a/src/backends/cuda/scoring/README.md +++ b/src/backends/cuda/scoring/README.md @@ -11,9 +11,11 @@ This is the readout both models that need it share the shape of: | CLM ([#9](https://github.com/ThinkFlowLab/system1-omni/issues/9)) | `exp(logit_scale) * cos(state_head(s), action_head(c))` | cosine, so `normalize = 1` | They differ in the similarity and in whether a projection happens first; what -they share is the primitive. `#9` asks for exactly this ("normalize, dot, -temperature, softmax in one pass over the candidate matrix"), and `#19`'s -`gemm.cu` does not have it. +they share is the primitive. `#9` describes CLM's answer as "the softmax over the +question's candidates" of `exp(logit_scale) * cos(state_head(s), action_head(c))`, +and `#27` describes Kev's as `scale * (k(h_opts) @ q(h_decide))` with +`scale = 1/sqrt(256)`, softmaxed over that question's candidates. Neither is in +`#19`'s `gemm.cu`, which is a bf16 cuBLASLt GEMM. ## What is fused, precisely From 78ddfe6187f905799c2403c954c41cb4ecd3224c Mon Sep 17 00:00:00 2001 From: xiaoyu <1259084489@qq.com> Date: Sat, 3 Oct 2026 20:53:33 +0800 Subject: [PATCH 4/7] cuda: say in the manifest what was actually run The tolerance note named one driver and only the fixed cases, and said nothing about sm_89 being the only capability exercised while the manifest's own default_arch is 80. --- src/backends/cuda/scoring/scoring.backend.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/backends/cuda/scoring/scoring.backend.json b/src/backends/cuda/scoring/scoring.backend.json index 1819a169..046cd0d0 100644 --- a/src/backends/cuda/scoring/scoring.backend.json +++ b/src/backends/cuda/scoring/scoring.backend.json @@ -23,7 +23,7 @@ "tolerance": { "max_abs": 0.0001, "measured_max_abs": 1.835e-07, - "note": "13 fixed cases, all within tolerance. Measured on RTX 4090 (sm_89, built with ./build.sh ./out 89), CUDA 13.0, driver 595.58.03." + "note": "13 fixed cases, 4 batch configurations and 6 rejected arguments, all within tolerance. Measured on an RTX 4090 built with ./build.sh ./out 89, CUDA 13.0 (V13.0.88), under drivers 595.58.03 and 580.76.05; sm_89 is the only capability run, and the sm_80 default above is unexercised." } }, "reference": { From 1a853173e14b4e479ec6eba3b141d99db2522fe7 Mon Sep 17 00:00:00 2001 From: xiaoyu <1259084489@qq.com> Date: Sat, 3 Oct 2026 20:59:26 +0800 Subject: [PATCH 5/7] cuda: measure the fusion, and fix what the measurement found The kernel argued that fusing the norm into the scoring loop saves a read of the candidate matrix. Measured against an unfused baseline that reads it twice, it did not: 0.82x at CLM's widest shape and 0.57x on Kev's, because a one-question request occupies one of the 128 SMs and the norm accumulation sat on that block's critical path. Achieved bandwidth was 1 to 190 GB/s against a device peak near 1000, so there was no read to save. Raising the block from 256 to 1024 threads moves the work onto 32 warps instead of 8 and reverses every result: 1.79x at 255 candidates of 512, 1.84x batched, 1.81x on Kev's plain dot. Correctness is unchanged -- the same 1.835e-07 worst case, and compute-sanitizer reports 0 errors. Also adds the CUDA Graph check the ABI was written for but never had: one node captured, and the replay matches a direct call exactly. --- docs/benchmarks/scoring-fusion/README.md | 93 +++++++ docs/benchmarks/scoring-fusion/results.json | 96 +++++++ src/backends/cuda/scoring/README.md | 38 ++- src/backends/cuda/scoring/bench.cu | 157 +++++++++++ src/backends/cuda/scoring/bench.py | 246 ++++++++++++++++++ .../cuda/scoring/candidate_scoring.cu | 2 +- src/backends/cuda/scoring/gpu_parity.py | 117 +++++++++ 7 files changed, 740 insertions(+), 9 deletions(-) create mode 100644 docs/benchmarks/scoring-fusion/README.md create mode 100644 docs/benchmarks/scoring-fusion/results.json create mode 100644 src/backends/cuda/scoring/bench.cu create mode 100644 src/backends/cuda/scoring/bench.py diff --git a/docs/benchmarks/scoring-fusion/README.md b/docs/benchmarks/scoring-fusion/README.md new file mode 100644 index 00000000..add58844 --- /dev/null +++ b/docs/benchmarks/scoring-fusion/README.md @@ -0,0 +1,93 @@ +# What fusing the candidate scoring is worth + +`src/backends/cuda/scoring/` argues that a cosine needs each candidate's norm and +its dot with the query, that both read the same elements, and that accumulating +them in one pass over `candidates` is therefore cheaper than composing a +normalisation pass with a scoring pass. This measures it, and the measurement +changed the kernel. + +## Protocol and controls + +The hypothesis was that the fused kernel beats an unfused equivalent that reads +the candidate matrix twice. Stop conditions were any failed correctness check, a +missing library, or no GPU; a correctness failure aborts the run rather than +producing a number, because a kernel that computes the wrong thing can be +arbitrarily fast. + +- One shared RTX 4090, 1 of 8 on the host, otherwise idle (0 % utilisation, + 30 °C) for the recorded run. The host is shared and its other GPUs carry + unrelated work, so the **ratios** are the result here, not the absolute + microseconds. The fused and unfused measurements are taken in the same + process, on the same buffers, and were repeated with the order swapped to rule + out a clock-ramp bias: the ratios moved by at most 0.01. +- Both paths are checked against `reference.py` first, per shape, at a `max_abs` + of 1e-4. Both passed before anything was timed. +- Per shape: 30 warmup calls, then 300 calls timed individually with CUDA events; + the median is reported. +- `bench.cu` is the baseline. It is `candidate_scoring.cu`'s own block with the + row norm loaded from a staged array instead of accumulated beside the dot, plus + a separate norm pass. The norm pass runs only when the similarity is a cosine, + because a plain dot has no norm to compute and running it anyway would slow the + baseline for a reason the fused kernel never removed. +- It is a **conservative** baseline: a real unfused pipeline built from a GEMM + would also round-trip the similarity matrix through global memory, which this + does not. The differences below are a floor. + +```sh +src/backends/cuda/scoring/build.sh ./out 89 +python3 src/backends/cuda/scoring/bench.py --library ./out/libscoring.so --iterations 300 \ + --out docs/benchmarks/scoring-fusion/results.json +``` + +## The finding: the kernel was occupancy-bound, not bandwidth-bound + +At the block size the kernel was written with, 256 threads, fusion **lost** at +every shape that mattered: + +| shape | similarity | K | D | questions | fused | unfused | ratio | +| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | +| `clm-cosine-5x512` | cosine | 5 | 512 | 1 | 9.2 µs | 11.3 µs | 1.23× | +| `clm-cosine-64x512` | cosine | 64 | 512 | 1 | 19.4 µs | 19.3 µs | 1.00× | +| `clm-cosine-255x512` | cosine | 255 | 512 | 1 | 52.9 µs | 43.2 µs | **0.82×** | +| `clm-cosine-batch8-255x512` | cosine | 255 | 512 | 8 | 53.1 µs | 44.0 µs | **0.83×** | +| `kev-dot-255x2560` | dot | 255 | 2560 | 1 | 167.8 µs | 96.1 µs | **0.57×** | + +Two things said why. The achieved bandwidth was 1 to 190 GB/s against a device +peak near 1000, so nothing was bandwidth-bound and there was no read to save. And +a single question and a batch of eight took **the same 53 µs** — eight times the +work in the same wall clock, which is only possible if each question is running +on its own SM. One thread block per question means a one-question request, the +common serving case, occupies **one of the 128 SMs**, and the norm accumulation +sat on that block's critical path instead of being spread across the device. + +Raising the block to 1024 threads — the maximum on sm_89, and one line — moves +the work onto 32 warps instead of 8 and reverses every result: + +| shape | similarity | K | D | questions | fused | unfused | ratio | +| --- | --- | ---: | ---: | ---: | ---: | ---: | ---: | +| `clm-cosine-5x512` | cosine | 5 | 512 | 1 | 9.2 µs | 11.3 µs | 1.22× | +| `clm-cosine-64x512` | cosine | 64 | 512 | 1 | 12.1 µs | 18.3 µs | 1.51× | +| `clm-cosine-255x512` | cosine | 255 | 512 | 1 | 23.6 µs | 42.2 µs | **1.79×** | +| `clm-cosine-batch8-255x512` | cosine | 255 | 512 | 8 | 23.8 µs | 43.8 µs | **1.84×** | +| `kev-dot-255x2560` | dot | 255 | 2560 | 1 | 55.3 µs | 100.3 µs | **1.81×** | + +Kev is the control. It is a plain dot, so there is no norm for the fusion to +remove from the traffic, and its gain is entirely the second kernel and the +global similarity round-trip it does not need. That it lands at the same ~1.8× as +the cosine shapes is the useful part of the result. + +Correctness is unchanged by the block size: the parity harness reports the same +`1.835e-07` worst case, and `compute-sanitizer --tool memcheck` reports +`ERROR SUMMARY: 0 errors`. + +## Limits + +- One GPU, one architecture: sm_89 only. The library builds for whatever + `build.sh` is given, but nothing else has been run. +- The tiny shapes are launch-bound, not compute-bound — `clm-cosine-5x512` is 9 µs + either way. The interesting range starts around K = 64. +- CUDA Graph replay is not timed here. It removes launch overhead, which matters + most exactly where these numbers are worst, so a graph-timed run would flatter + the fused kernel; that is a separate measurement. +- No end-to-end number. This is the kernel, with the projections left to the + caller as the design intends. diff --git a/docs/benchmarks/scoring-fusion/results.json b/docs/benchmarks/scoring-fusion/results.json new file mode 100644 index 00000000..4f075f87 --- /dev/null +++ b/docs/benchmarks/scoring-fusion/results.json @@ -0,0 +1,96 @@ +{ + "measured_at": "2026-10-03T12:58:48Z", + "provenance": { + "gpu": "NVIDIA GeForce RTX 4090", + "driver": "580.76.05", + "nvcc": "Cuda compilation tools, release 13.0, V13.0.88", + "library": "./out/libscoring.so", + "baseline_build": "nvcc -O3 -std=c++17 -gencode arch=compute_89,code=sm_89 -shared -Xcompiler -fPIC -Xcompiler -Wall,-Wextra bench.cu -lcudart -o libscoring_bench.so", + "iterations": 300, + "arch": "89", + "kernel_definition": "THREADS = 1024 in candidate_scoring.cu" + }, + "results": [ + { + "shape": "clm-cosine-5x512", + "questions": 1, + "K": 5, + "D": 512, + "cosine": true, + "fused_ms": 0.009216000325977802, + "unfused_ms": 0.011264000087976456, + "speedup": 1.2222221885372344, + "fused_gbs": 1.3355034249843238, + "unfused_gbs": 2.0017755525471306, + "fused_bytes": 12308, + "unfused_bytes": 22548, + "fused_max_abs": 6.932993346087102e-09, + "unfused_max_abs": 6.932993346087102e-09 + }, + { + "shape": "clm-cosine-64x512", + "questions": 1, + "K": 64, + "D": 512, + "cosine": true, + "fused_ms": 0.012128000147640705, + "unfused_ms": 0.018271999433636665, + "speedup": 1.5065962410291667, + "fused_gbs": 10.99736134369573, + "unfused_gbs": 14.472855089584856, + "fused_bytes": 133376, + "unfused_bytes": 264448, + "fused_max_abs": 1.2430758852821633e-09, + "unfused_max_abs": 1.2430758852821633e-09 + }, + { + "shape": "clm-cosine-255x512", + "questions": 1, + "K": 255, + "D": 512, + "cosine": true, + "fused_ms": 0.02356800064444542, + "unfused_ms": 0.04217600077390671, + "speedup": 1.7895451298643308, + "fused_gbs": 22.289035371517873, + "unfused_gbs": 24.837537480511738, + "fused_bytes": 525308, + "unfused_bytes": 1047548, + "fused_max_abs": 1.2364590072297399e-09, + "unfused_max_abs": 1.2364590072297399e-09 + }, + { + "shape": "clm-cosine-batch8-255x512", + "questions": 8, + "K": 255, + "D": 512, + "cosine": true, + "fused_ms": 0.023840000852942467, + "unfused_ms": 0.04383999854326248, + "speedup": 1.8389260475991762, + "fused_gbs": 176.27784604216188, + "unfused_gbs": 191.1584005124912, + "fused_bytes": 4202464, + "unfused_bytes": 8380384, + "fused_max_abs": 2.1920026549090976e-09, + "unfused_max_abs": 2.1920026549090976e-09 + }, + { + "shape": "kev-dot-255x2560", + "questions": 1, + "K": 255, + "D": 2560, + "cosine": false, + "fused_ms": 0.055296000093221664, + "unfused_ms": 0.1003359965980053, + "speedup": 1.814525398380574, + "fused_gbs": 47.42585350800931, + "unfused_gbs": 26.13678130399051, + "fused_bytes": 2622460, + "unfused_bytes": 2622460, + "fused_max_abs": 7.07020300361183e-08, + "unfused_max_abs": 7.07020300361183e-08 + } + ], + "note": "Ratios, not absolute microseconds: the host is shared (1 of 8 GPUs) and its other GPUs carry unrelated work. The same harness at THREADS = 256 gave 0.82x, 0.83x and 0.57x on clm-cosine-255x512, clm-cosine-batch8-255x512 and kev-dot-255x2560; see README.md." +} diff --git a/src/backends/cuda/scoring/README.md b/src/backends/cuda/scoring/README.md index 78c5d719..579d3cfe 100644 --- a/src/backends/cuda/scoring/README.md +++ b/src/backends/cuda/scoring/README.md @@ -19,14 +19,19 @@ and `#27` describes Kev's as `scale * (k(h_opts) @ q(h_decide))` with ## What is fused, precisely -Not merely "fewer launches". When the similarity is a cosine, each candidate's -norm and its dot with the query need **the same elements**, so both accumulate in -a single read of `C`. For Kev that is 255 rows of 2560 floats read twice per -question in the unfused version and once here. Beyond that: - -- One thread block per question; the stable softmax runs over shared memory, so - there is no second kernel, no atomics, and no global round trip for the - similarities. +Not merely "fewer launches". When the similarity is a cosine — CLM's case, not +Kev's, which is a plain dot — each candidate's norm and its dot with the query +need **the same elements**, so both accumulate in a single read of `C`. The +unfused version composes a normalisation pass with a scoring pass and reads `C` +twice; at CLM's widest, 255 candidates of 512 floats, that is 522 KB per question +read twice and once here. Beyond that: + +- One thread block per question, 1024 threads; the stable softmax runs over + shared memory, so there is no second kernel, no atomics, and no global round + trip for the similarities. The block size is not cosmetic — see + [the benchmark](../../../docs/benchmarks/scoring-fusion/README.md), where 256 + threads made the fusion a **loss** at every shape that mattered and 1024 + threads made it a 1.8x win. - `K ≤ 255` (the serving API's limit), which is what makes the in-block softmax possible at all — the whole similarity vector fits in shared memory. - A batch is a grid over questions, not a loop of launches. @@ -40,6 +45,8 @@ The query norm is computed once per block rather than once per candidate. | `reference.py` | The oracle: float64, deliberately slow and obvious | **yes**, 17 tests | | `kernel_simulation.py` | The kernel's algorithm in numpy: float32, warp tree order, same norm flooring | **yes** | | `candidate_scoring.cu` | The kernel and its C ABI | **yes**, compiled and checked against the reference on sm_89 | +| `bench.cu` | The unfused baseline the kernel is measured against | **yes**, timed on sm_89 | +| `bench.py` | Drives both and writes the report | **yes** | ## Verification state @@ -109,6 +116,20 @@ good as its fidelity, and the hardware is what settles it. **1.835e-07**, so the declared bound has 545x of headroom and a failure is a real failure rather than tolerance noise. Declared in `scoring.backend.json`. +## What the fusion is worth + +Measured against an unfused baseline that reads the candidate matrix twice, at +1024 threads: **1.22x** at CLM's narrowest shape, **1.51x** at 64 candidates and +**1.79x** to **1.84x** at 255. Kev's plain dot, which has no norm to fuse, lands +at the same **1.81x** from the second kernel and the global similarity round trip +it does not need. + +The measurement also found a defect the correctness checks could not: at the 256 +threads this kernel was first written with, a one-question request occupies one of +the 128 SMs, and the norm accumulation sat on that block's critical path — so the +fusion *lost*, down to 0.57x on Kev's shape. Protocol, both result sets and the +limits are in [the benchmark](../../../docs/benchmarks/scoring-fusion/README.md). + ## Reproducing On a machine with `nvcc` and an NVIDIA GPU: @@ -117,6 +138,7 @@ On a machine with `nvcc` and an NVIDIA GPU: src/backends/cuda/scoring/build.sh ./out 89 # or CUDA_ARCH_LIST="80 89 90" python3 src/backends/cuda/scoring/gpu_parity.py --library ./out/libscoring.so python3 src/backends/cuda/scoring/test_scoring.py # the reference, on CPU +python3 src/backends/cuda/scoring/bench.py --library ./out/libscoring.so # compiled kernel vs baseline ``` The inputs are not committed. They are deterministic from a fixed seed, so the diff --git a/src/backends/cuda/scoring/bench.cu b/src/backends/cuda/scoring/bench.cu new file mode 100644 index 00000000..c3c91152 --- /dev/null +++ b/src/backends/cuda/scoring/bench.cu @@ -0,0 +1,157 @@ +// The unfused baseline the fused candidate-scoring kernel is measured against. +// +// Cosine similarity needs each candidate's norm and its dot with the query, and +// both read the same elements of `candidates`. The fused kernel accumulates them +// in one read of the row. This baseline is the same computation split the way it +// falls out of composing existing pieces: a pass that normalises the candidate +// matrix, then a pass that scores it. That reads `candidates` twice. +// +// Everything else is deliberately identical to `candidate_scoring.cu` -- the same +// block shape, the same warp reduction, the same thread-0 softmax -- so the only +// difference the benchmark can see is the second pass over `candidates`. +// +// It is the *conservative* baseline. A real unfused pipeline built from a GEMM +// would also round-trip the similarity matrix through global memory, which this +// does not, so the difference measured here is a floor. +// +// Not part of libscoring.so: `build.sh` compiles only `candidate_scoring.cu` and +// the manifest lists only that file. `bench.py` compiles this one itself. +#include + +namespace { + +constexpr int WARP = 32; +constexpr int THREADS = 256; +constexpr int MAX_K = 255; +constexpr float NORM_EPS = 1e-12f; + +__device__ __forceinline__ float warp_sum(float value) { +#pragma unroll + for (int offset = WARP / 2; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffffu, value, offset); + } + return value; +} + +// Pass one: the norm of every candidate row, one block per row. This is the read +// the fusion removes. +__global__ void row_norms_kernel(const float* __restrict__ candidates, + float* __restrict__ norms, int rows, int D) { + __shared__ float partials[THREADS / WARP]; + const int warp = threadIdx.x / WARP; + const int lane = threadIdx.x % WARP; + const float* row = candidates + static_cast(blockIdx.x) * D; + + float sum = 0.0f; + for (int d = threadIdx.x; d < D; d += THREADS) { + const float value = row[d]; + sum = fmaf(value, value, sum); + } + sum = warp_sum(sum); + if (lane == 0) { + partials[warp] = sum; + } + __syncthreads(); + if (warp == 0) { + float total = (lane < blockDim.x / WARP) ? partials[lane] : 0.0f; + total = warp_sum(total); + if (lane == 0) { + norms[blockIdx.x] = fmaxf(sqrtf(total), NORM_EPS); + } + } +} + +// Pass two: the fused kernel's block, with the row norm loaded instead of +// accumulated alongside the dot. +__global__ void scoring_unfused_kernel(const float* __restrict__ query, + const float* __restrict__ candidates, + const float* __restrict__ row_norms, + float* __restrict__ probabilities, int K, int D, + float scale, float temperature, int normalize) { + __shared__ float similarity[MAX_K]; + __shared__ float reduce_buffer[THREADS / WARP]; + + const int warp = threadIdx.x / WARP; + const int lane = threadIdx.x % WARP; + const int warps = blockDim.x / WARP; + + const float* question_query = query + static_cast(blockIdx.x) * D; + const float* question_candidates = candidates + static_cast(blockIdx.x) * K * D; + const float* question_norms = row_norms + static_cast(blockIdx.x) * K; + float* question_probabilities = probabilities + static_cast(blockIdx.x) * K; + + float query_norm = 1.0f; + if (normalize) { + if (warp == 0) { + float partial = 0.0f; + for (int d = lane; d < D; d += WARP) { + const float q = question_query[d]; + partial = fmaf(q, q, partial); + } + partial = warp_sum(partial); + if (lane == 0) { + reduce_buffer[0] = fmaxf(sqrtf(partial), NORM_EPS); + } + } + __syncthreads(); + query_norm = reduce_buffer[0]; + __syncthreads(); + } + + for (int k = warp; k < K; k += warps) { + const float* row = question_candidates + static_cast(k) * D; + float dot = 0.0f; + for (int d = lane; d < D; d += WARP) { + dot = fmaf(row[d], question_query[d], dot); + } + dot = warp_sum(dot); + if (lane == 0) { + float value = dot; + if (normalize) { + value /= query_norm * question_norms[k]; + } + similarity[k] = value * scale / temperature; + } + } + __syncthreads(); + + if (threadIdx.x == 0) { + float maximum = similarity[0]; + for (int k = 1; k < K; ++k) { + maximum = fmaxf(maximum, similarity[k]); + } + float total = 0.0f; + for (int k = 0; k < K; ++k) { + const float value = expf(similarity[k] - maximum); + similarity[k] = value; + total += value; + } + const float inverse = (total > 0.0f) ? (1.0f / total) : 0.0f; + for (int k = 0; k < K; ++k) { + question_probabilities[k] = similarity[k] * inverse; + } + } +} + +} // namespace + +extern "C" { + +// Two passes over `candidates`, with the row norms staged in `norms`. +int bench_unfused(const float* query, const float* candidates, float* norms, + float* probabilities, int questions, int K, int D, float scale, + float temperature, int normalize, cudaStream_t stream) { + // The norm pass exists only because a cosine needs one. A plain dot does not, + // and running it anyway would make the baseline slower for a reason the fused + // kernel never removed. + if (normalize) { + row_norms_kernel<<>>(candidates, norms, + questions * K, D); + } + scoring_unfused_kernel<<>>(query, candidates, norms, + probabilities, K, D, scale, + temperature, normalize); + return static_cast(cudaGetLastError()); +} + +} // extern "C" diff --git a/src/backends/cuda/scoring/bench.py b/src/backends/cuda/scoring/bench.py new file mode 100644 index 00000000..83360ba2 --- /dev/null +++ b/src/backends/cuda/scoring/bench.py @@ -0,0 +1,246 @@ +#!/usr/bin/env python3 +"""Measure what fusing the candidate scoring is worth, on a GPU. + +Compares the shipped kernel with the two-pass baseline in `bench.cu`. Both are +given the same inputs, both are checked against `reference.py` first, and only +then are they timed -- a kernel that computes the wrong thing can be arbitrarily +fast, so a correctness check that fails aborts the run rather than producing a +number. + + python3 bench.py [--library PATH] [--iterations N] [--out results.json] + +Writes a JSON report; `docs/benchmarks/scoring-fusion/README.md` holds the +protocol and the results of the recorded run. +""" + +from __future__ import annotations + +import argparse +import ctypes +import json +import os +import shutil +import subprocess +import sys +import time + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +import gpu_parity # noqa: E402 +import reference # noqa: E402 + +HERE = os.path.dirname(os.path.abspath(__file__)) +DEFAULT_LIBRARY = os.path.join(HERE, "libscoring.so") +BASELINE_SOURCE = os.path.join(HERE, "bench.cu") +BASELINE_LIBRARY = os.path.join(HERE, "libscoring_bench.so") + +# The shapes the two models actually ask for, each with the similarity it uses. +# CLM is a cosine over 512-wide projected vectors with up to 255 candidates; Kev +# is a plain dot over 2560-wide hidden states, one question per prefill row. +# +# Kev is the control. It has no norm to compute, so there is nothing for this +# fusion to remove from the traffic, and the measurement should say so. +SHAPES = [ + ("clm-cosine-5x512", 1, 5, 512, True), + ("clm-cosine-64x512", 1, 64, 512, True), + ("clm-cosine-255x512", 1, 255, 512, True), + ("clm-cosine-batch8-255x512", 8, 255, 512, True), + ("kev-dot-255x2560", 1, 255, 2560, False), +] + +SCALE, TEMPERATURE = 0.0625, 1.5 + + +def build_baseline(nvcc, arch): + """Compile bench.cu, with the same flags build.sh uses for the kernel.""" + targets = [arch] + os.environ.get("CUDA_ARCH_LIST", "").split() + gencode = [] + for target in targets: + if target: + gencode += ["-gencode", "arch=compute_%s,code=sm_%s" % (target, target)] + command = ([nvcc, "-O3", "-std=c++17"] + gencode + + ["-shared", "-Xcompiler", "-fPIC", "-Xcompiler", "-Wall,-Wextra", + BASELINE_SOURCE, "-lcudart", "-o", BASELINE_LIBRARY]) + result = subprocess.run(command, capture_output=True, text=True) + if result.returncode != 0: + raise SystemExit("bench.cu did not compile:\n%s" % result.stderr[-2000:]) + return " ".join(command) + + +def load_bench_library(path): + lib = ctypes.CDLL(path) + lib.bench_unfused.restype = ctypes.c_int + lib.bench_unfused.argtypes = [ + ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_int, ctypes.c_int, ctypes.c_int, + ctypes.c_float, ctypes.c_float, ctypes.c_int, ctypes.c_void_p] + return lib + + +def timed(runtime, call, iterations): + """Median milliseconds per call, over `iterations` calls after a warmup.""" + for _ in range(max(3, iterations // 10)): + call() + runtime.cudaDeviceSynchronize() + + start, stop = ctypes.c_void_p(), ctypes.c_void_p() + gpu_parity.check(runtime, runtime.cudaEventCreate(ctypes.byref(start)), "cudaEventCreate") + gpu_parity.check(runtime, runtime.cudaEventCreate(ctypes.byref(stop)), "cudaEventCreate") + samples = [] + try: + for _ in range(iterations): + gpu_parity.check(runtime, runtime.cudaEventRecord(start, None), "record start") + call() + gpu_parity.check(runtime, runtime.cudaEventRecord(stop, None), "record stop") + gpu_parity.check(runtime, runtime.cudaEventSynchronize(stop), "sync stop") + elapsed = ctypes.c_float() + gpu_parity.check(runtime, runtime.cudaEventElapsedTime(ctypes.byref(elapsed), start, stop), + "cudaEventElapsedTime") + samples.append(elapsed.value) + finally: + runtime.cudaEventDestroy(start) + runtime.cudaEventDestroy(stop) + return float(np.median(samples)), samples + + +def traffic_bytes(questions, k, d, passes): + """Mandatory DRAM traffic: the query, `passes` reads of the candidates, and the output.""" + return questions * d * 4 + passes * questions * k * d * 4 + questions * k * 4 + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--library", default=DEFAULT_LIBRARY) + parser.add_argument("--iterations", type=int, default=200) + parser.add_argument("--arch", default=os.environ.get("CUDA_COMPUTE_CAP", "89")) + parser.add_argument("--nvcc", default=shutil.which("nvcc") or "/usr/local/cuda/bin/nvcc") + parser.add_argument("--out", default=None) + args = parser.parse_args(argv) + + reason = gpu_parity.unavailable_reason(args.library) + if reason: + print("SKIP: %s" % reason) + return 0 + + build = build_baseline(args.nvcc, args.arch) + fused = gpu_parity.load(args.library) + unfused = load_bench_library(BASELINE_LIBRARY) + runtime = gpu_parity.load_runtime() + for name in ("cudaEventCreate", "cudaEventRecord", "cudaEventSynchronize", + "cudaEventElapsedTime", "cudaEventDestroy"): + getattr(runtime, name).restype = ctypes.c_int + runtime.cudaEventElapsedTime.argtypes = [ctypes.POINTER(ctypes.c_float), + ctypes.c_void_p, ctypes.c_void_p] + + provenance = dict(gpu_parity.provenance()) + provenance.update({ + "library": args.library, + "baseline_build": build, + "iterations": args.iterations, + "arch": args.arch, + }) + + rows = [] + for name, questions, k, d, cosine in SHAPES: + spec = reference.ScoreSpec(scale=SCALE, temperature=TEMPERATURE, normalize=cosine) + rng = np.random.default_rng(1234) + q = np.ascontiguousarray(rng.standard_normal((questions, d)), dtype=np.float32) + c = np.ascontiguousarray(rng.standard_normal((questions, k, d)), dtype=np.float32) + device, outputs = {}, {} + try: + for key, array in (("q", q), ("c", c)): + pointer = ctypes.c_void_p() + gpu_parity.check(runtime, + runtime.cudaMalloc(ctypes.byref(pointer), array.nbytes), + "cudaMalloc(%s)" % key) + device[key] = pointer + gpu_parity.check(runtime, + runtime.cudaMemcpy(pointer, array.ctypes.data, array.nbytes, 1), + "copy %s" % key) + for key, count in (("fused", questions * k), ("unfused", questions * k), + ("norms", questions * k)): + pointer = ctypes.c_void_p() + gpu_parity.check(runtime, + runtime.cudaMalloc(ctypes.byref(pointer), count * 4), + "cudaMalloc(%s)" % key) + device[key] = pointer + for key in ("fused", "unfused"): + outputs[key] = np.zeros((questions, k), dtype=np.float32) + + def call_fused(): + return fused.cs_score_candidates_batch( + device["q"], device["c"], device["fused"], None, questions, k, d, + ctypes.c_float(spec.scale), ctypes.c_float(spec.temperature), + 1 if spec.normalize else 0, None) + + def call_unfused(): + return unfused.bench_unfused( + device["q"], device["c"], device["norms"], device["unfused"], questions, k, d, + ctypes.c_float(spec.scale), ctypes.c_float(spec.temperature), + 1 if spec.normalize else 0, None) + + # Correctness first. A wrong kernel can be arbitrarily fast. + for label, call, key in (("fused", call_fused, "fused"), + ("unfused", call_unfused, "unfused")): + gpu_parity.check(runtime, call(), "%s launch" % label) + gpu_parity.check(runtime, runtime.cudaDeviceSynchronize(), "sync") + gpu_parity.check(runtime, + runtime.cudaMemcpy(outputs[key].ctypes.data, device[key], + outputs[key].nbytes, 2), + "copy %s out" % label) + worst = 0.0 + for index in range(questions): + expected = reference.probabilities(q[index], c[index], spec) + worst = max(worst, float(np.abs( + outputs[key][index].astype(np.float64) - expected).max())) + if worst > 1e-4: + raise SystemExit("%s %s is wrong (max_abs=%.3e); not timing it" + % (name, label, worst)) + if label == "fused": + fused_worst = worst + else: + unfused_worst = worst + + fused_ms, _ = timed(runtime, call_fused, args.iterations) + unfused_ms, _ = timed(runtime, call_unfused, args.iterations) + finally: + for pointer in device.values(): + runtime.cudaFree(pointer) + + # The baseline only reads the candidates twice when it has a norm to + # compute; a plain dot reads once, like the fused kernel. + fused_bytes = traffic_bytes(questions, k, d, 1) + unfused_bytes = traffic_bytes(questions, k, d, 2 if spec.normalize else 1) + rows.append({ + "shape": name, + "questions": questions, "K": k, "D": d, "cosine": bool(spec.normalize), + "fused_ms": fused_ms, + "unfused_ms": unfused_ms, + "speedup": unfused_ms / fused_ms, + "fused_gbs": fused_bytes / fused_ms / 1e6, + "unfused_gbs": unfused_bytes / unfused_ms / 1e6, + "fused_bytes": fused_bytes, + "unfused_bytes": unfused_bytes, + "fused_max_abs": fused_worst, + "unfused_max_abs": unfused_worst, + }) + print("%-26s %-4s K=%-4d D=%-5d q=%d fused %7.1f us (%5.0f GB/s) " + "unfused %7.1f us (%5.0f GB/s) %.2fx" + % (name, "cos" if spec.normalize else "dot", k, d, questions, + fused_ms * 1e3, rows[-1]["fused_gbs"], + unfused_ms * 1e3, rows[-1]["unfused_gbs"], rows[-1]["speedup"])) + + report = {"measured_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "provenance": provenance, "results": rows} + if args.out: + with open(args.out, "w", encoding="utf-8") as handle: + json.dump(report, handle, indent=2) + handle.write("\n") + print("\nwrote %s" % args.out) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/backends/cuda/scoring/candidate_scoring.cu b/src/backends/cuda/scoring/candidate_scoring.cu index f5e55f3c..846a4755 100644 --- a/src/backends/cuda/scoring/candidate_scoring.cu +++ b/src/backends/cuda/scoring/candidate_scoring.cu @@ -32,7 +32,7 @@ namespace { constexpr int WARP = 32; -constexpr int THREADS = 256; +constexpr int THREADS = 1024; constexpr int MAX_K = 255; // the largest candidate set the serving API allows constexpr float NORM_EPS = 1e-12f; diff --git a/src/backends/cuda/scoring/gpu_parity.py b/src/backends/cuda/scoring/gpu_parity.py index 6237b7d6..ab4f61ab 100755 --- a/src/backends/cuda/scoring/gpu_parity.py +++ b/src/backends/cuda/scoring/gpu_parity.py @@ -38,6 +38,31 @@ def unavailable_reason(library): return None +def gpu_name(): + """The first device's name, or None when there is none.""" + probe = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], + capture_output=True, text=True) + if probe.returncode != 0 or not probe.stdout.strip(): + return None + return probe.stdout.strip().splitlines()[0] + + +def provenance(): + """What the numbers are attached to: device, driver, CUDA toolkit, nvcc.""" + out = {"gpu": gpu_name()} + probe = subprocess.run(["nvidia-smi", "--query-gpu=driver_version", + "--format=csv,noheader"], capture_output=True, text=True) + if probe.returncode == 0 and probe.stdout.strip(): + out["driver"] = probe.stdout.strip().splitlines()[0] + nvcc = shutil.which("nvcc") or "/usr/local/cuda/bin/nvcc" + if os.path.exists(nvcc): + version = subprocess.run([nvcc, "--version"], capture_output=True, text=True) + for line in version.stdout.splitlines(): + if "release" in line: + out["nvcc"] = line.strip() + return out + + def load_runtime(): """The CUDA runtime, for device memory. @@ -191,6 +216,89 @@ def rejection_checks(lib, runtime): return results +def graph_checks(lib, runtime, questions=4, k=17, d=64, seed=7): + """Capture one launch into a CUDA Graph, replay it, and compare with a direct call. + + The ABI is written for capture: it queues on the caller's stream and does not + synchronize. So a capture must see exactly one node, and a replay must produce + the same probabilities as the same call made directly -- not merely within + tolerance, but identical, because the kernel has no atomics and its reduction + order is fixed. + """ + for name, argtypes in ( + ("cudaStreamCreate", [ctypes.POINTER(ctypes.c_void_p)]), + ("cudaStreamBeginCapture", [ctypes.c_void_p, ctypes.c_int]), + ("cudaStreamEndCapture", [ctypes.c_void_p, ctypes.POINTER(ctypes.c_void_p)]), + ("cudaGraphGetNodes", [ctypes.c_void_p, ctypes.c_void_p, ctypes.POINTER(ctypes.c_size_t)]), + ("cudaGraphInstantiate", [ctypes.POINTER(ctypes.c_void_p), ctypes.c_void_p, + ctypes.c_ulonglong]), + ("cudaGraphLaunch", [ctypes.c_void_p, ctypes.c_void_p]), + ("cudaGraphExecDestroy", [ctypes.c_void_p]), + ("cudaGraphDestroy", [ctypes.c_void_p]), + ("cudaStreamDestroy", [ctypes.c_void_p]), + ): + getattr(runtime, name).argtypes = argtypes + getattr(runtime, name).restype = ctypes.c_int + + rng = np.random.default_rng(seed) + q = np.ascontiguousarray(rng.standard_normal((questions, d)), dtype=np.float32) + c = np.ascontiguousarray(rng.standard_normal((questions, k, d)), dtype=np.float32) + direct = np.zeros((questions, k), dtype=np.float32) + replayed = np.zeros((questions, k), dtype=np.float32) + spec = reference.ScoreSpec(scale=0.0625, temperature=1.5) + + device = {} + stream, graph, graph_exec = ctypes.c_void_p(), ctypes.c_void_p(), ctypes.c_void_p() + try: + for name, array in (("q", q), ("c", c), ("direct", direct), ("replayed", replayed)): + pointer = ctypes.c_void_p() + check(runtime, runtime.cudaMalloc(ctypes.byref(pointer), array.nbytes), + "cudaMalloc(%s)" % name) + device[name] = pointer + check(runtime, runtime.cudaMemcpy(device["q"], q.ctypes.data, q.nbytes, 1), "copy q") + check(runtime, runtime.cudaMemcpy(device["c"], c.ctypes.data, c.nbytes, 1), "copy c") + check(runtime, runtime.cudaStreamCreate(ctypes.byref(stream)), "cudaStreamCreate") + + def score(out, on_stream): + return lib.cs_score_candidates_batch( + device["q"], device["c"], device[out], None, int(questions), int(k), int(d), + ctypes.c_float(spec.scale), ctypes.c_float(spec.temperature), 1, on_stream) + + check(runtime, score("direct", None), "direct cs_score_candidates_batch") + check(runtime, runtime.cudaDeviceSynchronize(), "sync after the direct call") + + # cudaStreamCaptureModeThreadLocal = 1: only this thread's work is captured. + check(runtime, runtime.cudaStreamBeginCapture(stream, 1), "cudaStreamBeginCapture") + check(runtime, score("replayed", stream), "cs_score_candidates_batch under capture") + check(runtime, runtime.cudaStreamEndCapture(stream, ctypes.byref(graph)), + "cudaStreamEndCapture") + + nodes = ctypes.c_size_t() + check(runtime, runtime.cudaGraphGetNodes(graph, None, ctypes.byref(nodes)), + "cudaGraphGetNodes") + check(runtime, runtime.cudaGraphInstantiate(ctypes.byref(graph_exec), graph, + ctypes.c_ulonglong(0)), "cudaGraphInstantiate") + check(runtime, runtime.cudaGraphLaunch(graph_exec, stream), "cudaGraphLaunch") + check(runtime, runtime.cudaDeviceSynchronize(), "sync after the replay") + + check(runtime, runtime.cudaMemcpy(direct.ctypes.data, device["direct"], direct.nbytes, 2), + "copy the direct result out") + check(runtime, runtime.cudaMemcpy(replayed.ctypes.data, device["replayed"], + replayed.nbytes, 2), "copy the replayed result out") + finally: + if graph_exec.value: + runtime.cudaGraphExecDestroy(graph_exec) + if graph.value: + runtime.cudaGraphDestroy(graph) + if stream.value: + runtime.cudaStreamDestroy(stream) + for pointer in device.values(): + runtime.cudaFree(pointer) + + return int(nodes.value), float(np.abs(direct.astype(np.float64) + - replayed.astype(np.float64)).max()) + + def main(argv=None): parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) parser.add_argument("--library", default=DEFAULT_LIBRARY) @@ -242,6 +350,15 @@ def main(argv=None): failures.append("not rejected: %s" % name) print("%s %-20s returned %d" % ("ok " if ok else "FAIL", name, status)) + print() + print("--- CUDA Graph capture and replay ---") + nodes, graph_difference = graph_checks(lib, runtime) + graphs_ok = nodes == 1 and graph_difference == 0.0 + if not graphs_ok: + failures.append("graph: %d node(s), max_abs=%.3e" % (nodes, graph_difference)) + print("%s captured %d node, replay matches the direct call exactly (max_abs=%.3e)" + % ("ok " if graphs_ok else "FAIL", nodes, graph_difference)) + print("\nworst absolute difference: %.3e (tolerance %g)" % (worst, args.tolerance)) if failures: print("cases over tolerance: %s" % ", ".join(failures)) From 251d691d59f9775cf877e8de387cf61686e61ef2 Mon Sep 17 00:00:00 2001 From: xiaoyu <1259084489@qq.com> Date: Sat, 3 Oct 2026 21:15:24 +0800 Subject: [PATCH 6/7] cuda: time the fused kernel through a CUDA Graph too A graph replay saves a flat ~3us of launch, which is 1.49x of the smallest shape and 1.06x of the largest, so it does not change which configuration wins anywhere measured -- but it explains the floor at K = 5. Every timing in a row is now taken on one created stream. Measuring the direct launch on the legacy default stream added about 5us at K = 255, which is larger than the effect being measured. --- docs/benchmarks/scoring-fusion/README.md | 29 +++++++- docs/benchmarks/scoring-fusion/results.json | 79 ++++++++++++--------- src/backends/cuda/scoring/bench.py | 68 +++++++++++++++--- 3 files changed, 133 insertions(+), 43 deletions(-) diff --git a/docs/benchmarks/scoring-fusion/README.md b/docs/benchmarks/scoring-fusion/README.md index add58844..d4c38096 100644 --- a/docs/benchmarks/scoring-fusion/README.md +++ b/docs/benchmarks/scoring-fusion/README.md @@ -24,6 +24,10 @@ arbitrarily fast. of 1e-4. Both passed before anything was timed. - Per shape: 30 warmup calls, then 300 calls timed individually with CUDA events; the median is reported. +- Every timing in a row is taken on **one created stream**, including the graph + replay. Measuring on the legacy default stream instead adds its synchronisation + semantics to both and makes the columns incomparable — it cost the fused kernel + about 5 µs at K = 255, which is larger than the effect being measured. - `bench.cu` is the baseline. It is `candidate_scoring.cu`'s own block with the row norm loaded from a staged array instead of accumulated beside the dot, plus a separate norm pass. The norm pass runs only when the similarity is a cosine, @@ -80,14 +84,33 @@ Correctness is unchanged by the block size: the parity harness reports the same `1.835e-07` worst case, and `compute-sanitizer --tool memcheck` reports `ERROR SUMMARY: 0 errors`. +## CUDA Graph replay + +The ABI is written for capture — it queues on the caller's stream and never +synchronizes — so the same launch can be replayed from a graph. It is timed here +because the smallest shapes are launch-bound, and that is the part a graph +removes: + +| shape | direct | graph replay | gain | +| --- | ---: | ---: | ---: | +| `clm-cosine-5x512` | 9.2 µs | 6.1 µs | **1.49x** | +| `clm-cosine-64x512` | 12.1 µs | 9.2 µs | 1.31x | +| `clm-cosine-255x512` | 24.2 µs | 20.7 µs | 1.17x | +| `clm-cosine-batch8-255x512` | 24.4 µs | 20.5 µs | 1.19x | +| `kev-dot-255x2560` | 52.4 µs | 49.3 µs | 1.06x | + +The saving is a flat ~3 µs, which is what a launch costs; it is 1.49x of a 9 µs +call and 1.06x of a 52 µs one. That is the expected shape and it explains the +floor at K = 5, but it also means the graph gain **shrinks as the shapes grow** — +so it does not change which configuration wins at any shape measured here. A +graph and the fused kernel are complementary rather than alternatives, and the +two effects multiply. + ## Limits - One GPU, one architecture: sm_89 only. The library builds for whatever `build.sh` is given, but nothing else has been run. - The tiny shapes are launch-bound, not compute-bound — `clm-cosine-5x512` is 9 µs either way. The interesting range starts around K = 64. -- CUDA Graph replay is not timed here. It removes launch overhead, which matters - most exactly where these numbers are worst, so a graph-timed run would flatter - the fused kernel; that is a separate measurement. - No end-to-end number. This is the kernel, with the projections left to the caller as the design intends. diff --git a/docs/benchmarks/scoring-fusion/results.json b/docs/benchmarks/scoring-fusion/results.json index 4f075f87..b87c694a 100644 --- a/docs/benchmarks/scoring-fusion/results.json +++ b/docs/benchmarks/scoring-fusion/results.json @@ -1,5 +1,5 @@ { - "measured_at": "2026-10-03T12:58:48Z", + "measured_at": "2026-10-03T13:15:08Z", "provenance": { "gpu": "NVIDIA GeForce RTX 4090", "driver": "580.76.05", @@ -17,15 +17,18 @@ "K": 5, "D": 512, "cosine": true, - "fused_ms": 0.009216000325977802, - "unfused_ms": 0.011264000087976456, - "speedup": 1.2222221885372344, - "fused_gbs": 1.3355034249843238, - "unfused_gbs": 2.0017755525471306, + "fused_ms": 0.009151999838650227, + "unfused_ms": 0.011455999687314034, + "speedup": 1.2517482396507134, + "fused_gbs": 1.3448426810522358, + "unfused_gbs": 1.9682263107050233, "fused_bytes": 12308, "unfused_bytes": 22548, "fused_max_abs": 6.932993346087102e-09, - "unfused_max_abs": 6.932993346087102e-09 + "unfused_max_abs": 6.932993346087102e-09, + "direct_ms": 0.009184000082314014, + "graph_replay_ms": 0.006144000217318535, + "graph_gain": 1.4947916271920716 }, { "shape": "clm-cosine-64x512", @@ -33,15 +36,18 @@ "K": 64, "D": 512, "cosine": true, - "fused_ms": 0.012128000147640705, - "unfused_ms": 0.018271999433636665, - "speedup": 1.5065962410291667, - "fused_gbs": 10.99736134369573, - "unfused_gbs": 14.472855089584856, + "fused_ms": 0.012160000391304493, + "unfused_ms": 0.018432000651955605, + "speedup": 1.515789478521412, + "fused_gbs": 10.968420699671686, + "unfused_gbs": 14.34722171474872, "fused_bytes": 133376, "unfused_bytes": 264448, "fused_max_abs": 1.2430758852821633e-09, - "unfused_max_abs": 1.2430758852821633e-09 + "unfused_max_abs": 1.2430758852821633e-09, + "direct_ms": 0.01206399966031313, + "graph_replay_ms": 0.009184000082314014, + "graph_gain": 1.3135888014140202 }, { "shape": "clm-cosine-255x512", @@ -49,15 +55,18 @@ "K": 255, "D": 512, "cosine": true, - "fused_ms": 0.02356800064444542, - "unfused_ms": 0.04217600077390671, - "speedup": 1.7895451298643308, - "fused_gbs": 22.289035371517873, - "unfused_gbs": 24.837537480511738, + "fused_ms": 0.02425600029528141, + "unfused_ms": 0.043007999658584595, + "speedup": 1.7730870355798547, + "fused_gbs": 21.656826913140733, + "unfused_gbs": 24.357050044546877, "fused_bytes": 525308, "unfused_bytes": 1047548, "fused_max_abs": 1.2364590072297399e-09, - "unfused_max_abs": 1.2364590072297399e-09 + "unfused_max_abs": 1.2364590072297399e-09, + "direct_ms": 0.024240000173449516, + "graph_replay_ms": 0.020735999569296837, + "graph_gain": 1.168981514126811 }, { "shape": "clm-cosine-batch8-255x512", @@ -65,15 +74,18 @@ "K": 255, "D": 512, "cosine": true, - "fused_ms": 0.023840000852942467, - "unfused_ms": 0.04383999854326248, - "speedup": 1.8389260475991762, - "fused_gbs": 176.27784604216188, - "unfused_gbs": 191.1584005124912, + "fused_ms": 0.024383999407291412, + "unfused_ms": 0.04398399963974953, + "speedup": 1.803805803349767, + "fused_gbs": 172.34514854619624, + "unfused_gbs": 190.53255885411613, "fused_bytes": 4202464, "unfused_bytes": 8380384, "fused_max_abs": 2.1920026549090976e-09, - "unfused_max_abs": 2.1920026549090976e-09 + "unfused_max_abs": 2.1920026549090976e-09, + "direct_ms": 0.024383999407291412, + "graph_replay_ms": 0.020479999482631683, + "graph_gain": 1.1906250011368684 }, { "shape": "kev-dot-255x2560", @@ -81,16 +93,19 @@ "K": 255, "D": 2560, "cosine": false, - "fused_ms": 0.055296000093221664, - "unfused_ms": 0.1003359965980053, - "speedup": 1.814525398380574, - "fused_gbs": 47.42585350800931, - "unfused_gbs": 26.13678130399051, + "fused_ms": 0.05427199974656105, + "unfused_ms": 0.09728000313043594, + "speedup": 1.7924528962395585, + "fused_gbs": 48.320681239798475, + "unfused_gbs": 26.9578527509269, "fused_bytes": 2622460, "unfused_bytes": 2622460, "fused_max_abs": 7.07020300361183e-08, - "unfused_max_abs": 7.07020300361183e-08 + "unfused_max_abs": 7.07020300361183e-08, + "direct_ms": 0.05238400027155876, + "graph_replay_ms": 0.04934399947524071, + "graph_gain": 1.0616083177011915 } ], - "note": "Ratios, not absolute microseconds: the host is shared (1 of 8 GPUs) and its other GPUs carry unrelated work. The same harness at THREADS = 256 gave 0.82x, 0.83x and 0.57x on clm-cosine-255x512, clm-cosine-batch8-255x512 and kev-dot-255x2560; see README.md." + "note": "Ratios, not absolute microseconds: the host is shared (1 of 8 GPUs) and its other GPUs carry unrelated work. Every timing in a row is taken on one created stream. The same harness at THREADS = 256 gave 0.82x, 0.83x and 0.57x on clm-cosine-255x512, clm-cosine-batch8-255x512 and kev-dot-255x2560; see README.md." } diff --git a/src/backends/cuda/scoring/bench.py b/src/backends/cuda/scoring/bench.py index 83360ba2..84911a47 100644 --- a/src/backends/cuda/scoring/bench.py +++ b/src/backends/cuda/scoring/bench.py @@ -79,8 +79,12 @@ def load_bench_library(path): return lib -def timed(runtime, call, iterations): - """Median milliseconds per call, over `iterations` calls after a warmup.""" +def timed(runtime, call, iterations, stream=None): + """Median milliseconds per call, over `iterations` calls after a warmup. + + Both the direct launch and the graph replay are timed on the same stream, so + the only difference between them is the path taken to the kernel. + """ for _ in range(max(3, iterations // 10)): call() runtime.cudaDeviceSynchronize() @@ -91,9 +95,9 @@ def timed(runtime, call, iterations): samples = [] try: for _ in range(iterations): - gpu_parity.check(runtime, runtime.cudaEventRecord(start, None), "record start") + gpu_parity.check(runtime, runtime.cudaEventRecord(start, stream), "record start") call() - gpu_parity.check(runtime, runtime.cudaEventRecord(stop, None), "record stop") + gpu_parity.check(runtime, runtime.cudaEventRecord(stop, stream), "record stop") gpu_parity.check(runtime, runtime.cudaEventSynchronize(stop), "sync stop") elapsed = ctypes.c_float() gpu_parity.check(runtime, runtime.cudaEventElapsedTime(ctypes.byref(elapsed), start, stop), @@ -105,6 +109,26 @@ def timed(runtime, call, iterations): return float(np.median(samples)), samples +def capture_fused(runtime, lib, device, stream, questions, k, d, spec): + """Capture one fused launch into a graph, and instantiate it. + + The ABI is written for this: it queues on the caller's stream and never + synchronizes, so the capture sees exactly one node. + """ + graph, graph_exec = ctypes.c_void_p(), ctypes.c_void_p() + gpu_parity.check(runtime, runtime.cudaStreamBeginCapture(stream, 1), "cudaStreamBeginCapture") + gpu_parity.check(runtime, lib.cs_score_candidates_batch( + device["q"], device["c"], device["fused"], None, questions, k, d, + ctypes.c_float(spec.scale), ctypes.c_float(spec.temperature), + 1 if spec.normalize else 0, stream), "launch under capture") + gpu_parity.check(runtime, runtime.cudaStreamEndCapture(stream, ctypes.byref(graph)), + "cudaStreamEndCapture") + gpu_parity.check(runtime, runtime.cudaGraphInstantiate(ctypes.byref(graph_exec), graph, + ctypes.c_ulonglong(0)), + "cudaGraphInstantiate") + return graph, graph_exec + + def traffic_bytes(questions, k, d, passes): """Mandatory DRAM traffic: the query, `passes` reads of the candidates, and the output.""" return questions * d * 4 + passes * questions * k * d * 4 + questions * k * 4 @@ -169,17 +193,25 @@ def main(argv=None): for key in ("fused", "unfused"): outputs[key] = np.zeros((questions, k), dtype=np.float32) + # One stream for every timing below, so the direct launch and the + # graph replay differ only in the path to the kernel. Measuring on + # the legacy default stream instead would add its synchronisation + # semantics to both and make the columns incomparable. + stream = ctypes.c_void_p() + gpu_parity.check(runtime, runtime.cudaStreamCreate(ctypes.byref(stream)), + "cudaStreamCreate") + def call_fused(): return fused.cs_score_candidates_batch( device["q"], device["c"], device["fused"], None, questions, k, d, ctypes.c_float(spec.scale), ctypes.c_float(spec.temperature), - 1 if spec.normalize else 0, None) + 1 if spec.normalize else 0, stream) def call_unfused(): return unfused.bench_unfused( device["q"], device["c"], device["norms"], device["unfused"], questions, k, d, ctypes.c_float(spec.scale), ctypes.c_float(spec.temperature), - 1 if spec.normalize else 0, None) + 1 if spec.normalize else 0, stream) # Correctness first. A wrong kernel can be arbitrarily fast. for label, call, key in (("fused", call_fused, "fused"), @@ -203,8 +235,23 @@ def call_unfused(): else: unfused_worst = worst - fused_ms, _ = timed(runtime, call_fused, args.iterations) - unfused_ms, _ = timed(runtime, call_unfused, args.iterations) + fused_ms, _ = timed(runtime, call_fused, args.iterations, stream) + unfused_ms, _ = timed(runtime, call_unfused, args.iterations, stream) + + # The same launch, replayed from a captured graph. The small shapes + # are launch-bound, so this is where the difference should show. + graph, graph_exec = capture_fused(runtime, fused, device, stream, + questions, k, d, spec) + try: + def call_replay(): + return runtime.cudaGraphLaunch(graph_exec, stream) + + direct_ms, _ = timed(runtime, call_fused, args.iterations, stream) + graph_ms, _ = timed(runtime, call_replay, args.iterations, stream) + finally: + runtime.cudaGraphExecDestroy(graph_exec) + runtime.cudaGraphDestroy(graph) + runtime.cudaStreamDestroy(stream) finally: for pointer in device.values(): runtime.cudaFree(pointer) @@ -225,12 +272,17 @@ def call_unfused(): "unfused_bytes": unfused_bytes, "fused_max_abs": fused_worst, "unfused_max_abs": unfused_worst, + "direct_ms": direct_ms, + "graph_replay_ms": graph_ms, + "graph_gain": direct_ms / graph_ms, }) print("%-26s %-4s K=%-4d D=%-5d q=%d fused %7.1f us (%5.0f GB/s) " "unfused %7.1f us (%5.0f GB/s) %.2fx" % (name, "cos" if spec.normalize else "dot", k, d, questions, fused_ms * 1e3, rows[-1]["fused_gbs"], unfused_ms * 1e3, rows[-1]["unfused_gbs"], rows[-1]["speedup"])) + print("%-26s the same launch: direct %7.1f us, graph replay %7.1f us %.2fx" + % ("", direct_ms * 1e3, graph_ms * 1e3, direct_ms / graph_ms)) report = {"measured_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "provenance": provenance, "results": rows} From a0d3d4bef307c330d499e2e1485a77181e630c3c Mon Sep 17 00:00:00 2001 From: xiaoyu <1259084489@qq.com> Date: Sun, 4 Oct 2026 22:55:01 +0800 Subject: [PATCH 7/7] cuda: give the scoring library an ABI macro the manifest can be checked against `cs_score_abi_version()` returned a literal and nothing else in the source stated the number, so `scoring.backend.json`'s `abi_version` had no counterpart to disagree with. #25 now reads `#define _ABI_VERSION N` out of a backend's declared sources and requires the manifest to match; this backend was the one it reported a warning for. The macro is defined once and the function returns it, so the two cannot drift: the manifest says 1, the macro says 1, and a manifest saying anything else is now reported as `abi_version is 2 but candidate_scoring.cu defines CS_SCORE_ABI_VERSION 1`. No behaviour change: the function body expands to `return 1;` exactly as before, so the compiled library is unchanged apart from line numbers. The file has not been recompiled here -- there is no CUDA toolkit on this machine -- but the preprocessed result is what it was. 17 tests still pass. --- src/backends/cuda/scoring/candidate_scoring.cu | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/backends/cuda/scoring/candidate_scoring.cu b/src/backends/cuda/scoring/candidate_scoring.cu index 846a4755..5158f41f 100644 --- a/src/backends/cuda/scoring/candidate_scoring.cu +++ b/src/backends/cuda/scoring/candidate_scoring.cu @@ -29,6 +29,11 @@ #include #include +// Bumped whenever the required interface below changes. The manifest repeats it +// as `abi_version` and the checker compares the two, so the number lives here +// once: a manifest that drifts from this macro is reported rather than trusted. +#define CS_SCORE_ABI_VERSION 1 + namespace { constexpr int WARP = 32; @@ -171,7 +176,9 @@ bool arguments_valid(int K, int D, float scale, float temperature) { extern "C" { // ABI version of this library; a loader refuses a value it does not know. -uint32_t cs_score_abi_version(void) { return 1; } +// Returning the macro rather than a literal is what keeps it and the manifest +// from drifting apart. +uint32_t cs_score_abi_version(void) { return CS_SCORE_ABI_VERSION; } // Scores `questions` independent questions. //