diff --git a/docs/benchmarks/scoring-fusion/README.md b/docs/benchmarks/scoring-fusion/README.md new file mode 100644 index 00000000..d4c38096 --- /dev/null +++ b/docs/benchmarks/scoring-fusion/README.md @@ -0,0 +1,116 @@ +# 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. +- 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, + 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`. + +## 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. +- 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..b87c694a --- /dev/null +++ b/docs/benchmarks/scoring-fusion/results.json @@ -0,0 +1,111 @@ +{ + "measured_at": "2026-10-03T13:15:08Z", + "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.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, + "direct_ms": 0.009184000082314014, + "graph_replay_ms": 0.006144000217318535, + "graph_gain": 1.4947916271920716 + }, + { + "shape": "clm-cosine-64x512", + "questions": 1, + "K": 64, + "D": 512, + "cosine": true, + "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, + "direct_ms": 0.01206399966031313, + "graph_replay_ms": 0.009184000082314014, + "graph_gain": 1.3135888014140202 + }, + { + "shape": "clm-cosine-255x512", + "questions": 1, + "K": 255, + "D": 512, + "cosine": true, + "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, + "direct_ms": 0.024240000173449516, + "graph_replay_ms": 0.020735999569296837, + "graph_gain": 1.168981514126811 + }, + { + "shape": "clm-cosine-batch8-255x512", + "questions": 8, + "K": 255, + "D": 512, + "cosine": true, + "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, + "direct_ms": 0.024383999407291412, + "graph_replay_ms": 0.020479999482631683, + "graph_gain": 1.1906250011368684 + }, + { + "shape": "kev-dot-255x2560", + "questions": 1, + "K": 255, + "D": 2560, + "cosine": false, + "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, + "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. 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/README.md b/src/backends/cuda/scoring/README.md new file mode 100644 index 00000000..579d3cfe --- /dev/null +++ b/src/backends/cuda/scoring/README.md @@ -0,0 +1,159 @@ +# 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` 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 + +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. + +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 | **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 + +**Verified on real hardware**, on two driver stacks. RTX 4090 (sm_89), CUDA 13.0 +(V13.0.88), compiled with `./build.sh ./out 89`: + +| 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 + +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`. + +## 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: + +```sh +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 +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 + +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/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..84911a47 --- /dev/null +++ b/src/backends/cuda/scoring/bench.py @@ -0,0 +1,298 @@ +#!/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, 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() + + 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, stream), "record start") + call() + 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), + "cudaEventElapsedTime") + samples.append(elapsed.value) + finally: + runtime.cudaEventDestroy(start) + runtime.cudaEventDestroy(stop) + 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 + + +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) + + # 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, 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, stream) + + # 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, 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) + + # 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, + "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} + 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/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..5158f41f --- /dev/null +++ b/src/backends/cuda/scoring/candidate_scoring.cu @@ -0,0 +1,217 @@ +// 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 + +// 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; +constexpr int THREADS = 1024; +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. +// 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. +// +// 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..ab4f61ab --- /dev/null +++ b/src/backends/cuda/scoring/gpu_parity.py @@ -0,0 +1,371 @@ +#!/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 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. + + 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 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) + 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() + 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)) + 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..046cd0d0 --- /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, 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": { + "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)