diff --git a/FAST_CUDA.md b/FAST_CUDA.md new file mode 100644 index 0000000000..dfb6bf2d1d --- /dev/null +++ b/FAST_CUDA.md @@ -0,0 +1,280 @@ +# FAST_CUDA.md — GPU-native trainer speed investigation handoff + +Repo: `/home/keith/Git/ml/PufferLib-5.0`, branch `affine-5.0-port`, PR #645 +(`PufferAI/PufferLib`, base `5.0`, env-only scope: `ocean/affine_lock/*`). + +## Goal + +Get a real, measured speedup from the GPU-native (`.cu`) ocean env backend on +serious hardware (RTX 5090, RTX 4090, eventually multi-GPU), **without** +changing training results — checkpoints/metrics should match the CPU (`.h`) +backend, just run faster. The affine_lock env's own PR (#645) is env-scoped +only; anything touching `src/pufferl.cu` / `src/algo.cu` is core, shared by +every ocean env, and must live on a **separate branch**, not #645. + +## Hardware in play + +- **This box**: AMD Ryzen 9 9950X3D, 16 cores / 32 threads, RTX 5090 (32GB, + driver 580.105.08, compute cap 12.0 / sm_120). +- **Reference box** (a different agent, "the 5060 box"): RTX 5060, only + **4 CPU cores / 8 threads**. This matters a lot — see Finding 3. +- User also runs "Puffer boxes" on 4090s and wants this to generalize there + and to multi-GPU eventually. + +## TL;DR finding + +**On this hardware, the `.cu` backend currently gives no real speedup over +the CPU `.h` backend, at any network size tested.** The reference agent's +advertised "3.46x" was real on their machine but was measuring their CPU +backend being thread-starved (8 threads), not a GPU architecture win. The +actual bottleneck for a real win is in the **shared trainer** (CUDA graph +launch/dispatch overhead in `src/pufferl.cu`'s epoch loop and the Muon +optimizer, likely `src/algo.cu`), not in `ocean/affine_lock/affine_lock.cu` +itself. The env's own GPU kernels are already fast and small — see Finding 5. + +## Build verification (confirms build/env are not the problem) + +``` +source .venv/bin/activate +export CUDA_HOME=/usr/local/cuda +export NVCC_ARCH=sm_120 +unset NVCC_PREPEND_FLAGS +./build.sh affine_lock build/puffer_affine_lock_cu --cu +``` + +`cuobjdump --dump-resource-usage` on the resulting binary confirms real GPU +kernels embedded, `arch = sm_120`, and for +`gpu_affine_lock_shared_step_kernel`: `REG:36 SHARED:17920` (matches the +expected footprint from the other agent's recipe almost exactly — REG 36 vs +their 38, negligible compiler-version noise). **Build is healthy, not a +stale-CPU-binary problem.** + +GPU health during a run: P1 pstate, 100% util, ~558/600W (93%), SM clock +~2820MHz, 78°C, no throttling. **Not a hardware/thermal problem either.** + +## Finding 1 — runtime args don't move SPS at matched hypers + +All tests: `hidden_size=512, num_layers=3.93814397` (the known-best +`dainty-firefly-578` CPU-backend hypers, perf=0.0636), `vec.num_buffers=1` +(hard-required for GPU backend, see Finding 4), `train.total_timesteps= +67108864` canary. + +| config | overall SPS | +|---|---| +| default (agents=4096, minibatch=8192, horizon=64) | 2.586M | +| agents=8192 | 2.534M | +| agents=16384 | 2.575M | +| async=0 | 2.549M | +| horizon=128 | 2.420M | +| **minibatch=32768** | **3.556M** | + +Only minibatch size moved anything, and it changes the actual training +recipe (fewer, larger optimizer updates), not a free win. Reference agent's +correction on why: it's not that Muon gets more efficient per-update: it's +that **fewer total updates means fewer total kernel-graph-node launches**, +and launch/dispatch overhead is the real cost (Finding 5), not compute. + +`replay_ratio=0` (rollout-only) ceiling: **21.887M SPS** — confirms 3x over +578's 2.73M (target ~8.2M) is theoretically reachable; the bottleneck is +specifically in the train/optimizer phase, not rollout. + +## Finding 2 — 578's real CPU-backend baseline (for reference) + +From wandb (`kinvert-k/affine3`, run `ktlyegxl`, `dainty-firefly-578`): +`vec.gpu_env=0`, `vec.num_buffers=2`, same hypers as above. **SPS: 2.73M +mean/final, uptime 709.7s for 1.937B steps.** This is the number any GPU +win needs to beat by 3x (~8.2M) to matter. + +## Finding 3 — apples-to-apples backend comparison, THIS hardware + +Same exact matched-hypers method the reference agent used (same seed, +`async=0`, `cudagraphs=1`, `num_buffers=1`, BF16, identical everything else +between the two binaries — only backend changes): + +**Large net** (hidden_size=512, 3.93814397 layers~4, agents=4096, horizon=64, +minibatch=8192, replay_ratio≈1.87... — actually see raw log for exact args +used, effectively 578-style): +- CPU: 2.586M SPS (from Finding 1's baseline row) + +**Matched net** (hidden_size=256, num_layers=3, agents=4096, horizon=64, +minibatch=8192, replay_ratio=3, 100,139,008 steps — reference agent's exact +matched-comparison recipe): +- CPU: 3.004M SPS (33.335s) +- CUDA: 3.143M SPS (31.859s) +- **Speedup: 1.05x** + +**Tiny net** (hidden_size=64, num_layers=1, agents=8192, horizon=32, +minibatch=8192, replay_ratio=3, 100,139,008 steps — the regime where GPU +should win most, since env-stepping cost is a bigger fraction of total time +relative to a tiny policy): +- CPU: 10.902M SPS (9.185s) +- CUDA: 12.879M SPS (7.775s) +- **Speedup: 1.18x** + +**Reference agent's own matched result on their 4-core/8-thread box**, same +256-hidden recipe: CPU 0.311M SPS, CUDA 1.077M SPS, **3.46x**. Their CPU +number is ~10x slower than ours on the identical config — that gap is +entirely CPU-thread starvation (4 cores vs our 32 threads), not a GPU +architecture advantage. **The 3.46x does not reproduce on strong CPU +hardware.** + +Open question sent to the reference agent, answer still pending: their +best-tuned SPS after block/lane kernel tuning, at both the 512-hidden and +256-hidden recipes, for direct comparison against our numbers above. + +## Finding 4 — why `num_buffers` is capped at 1 for GPU-native + +`src/pufferl.cu:978`: `assert(vec->buffers == 1 && "GPU env: num_buffers +must be 1");` inside `env_setup()`. + +This is not an arbitrary limitation. CPU backend's `num_buffers=2` lets CPU +worker threads step the *next* rollout buffer's environments while the GPU +trains on the *current* buffer — genuine overlap because CPU cores and GPU +cores are separate hardware. GPU-native env-stepping happens **on the same +GPU** that does training, so naively enabling `num_buffers=2` there would +just interleave GPU work on one stream, not add real parallelism. A real fix +requires running env-step (for the next buffer) on a **separate CUDA +stream** concurrently with the training graph's stream, relying on the +GPU's own multi-stream concurrency (Hopper/Blackwell can do this if kernels +don't fight over the same SMs/queues) — genuine new engineering, not a flag. + +## Finding 5 — Nsight profiling: where the time actually goes + +Setup used (repeat exactly for any follow-up profiling): + +``` +BIN=./build/puffer_affine_lock_cu +mkdir -p /tmp/affine-nsys/{logs,checkpoints} + +for MB in 8192 32768; do + nsys profile \ + --force-overwrite=true \ + --capture-range=cudaProfilerApi \ + --capture-range-end=stop \ + --cuda-graph-trace=node \ + --trace=cuda,nvtx,cublas \ + --sample=none \ + --cpuctxsw=none \ + --output="/tmp/affine-nsys/mb${MB}" \ + "$BIN" train \ + --base.profile=1 --base.async=1 --base.cudagraphs=1 \ + --base.run_id="nsys-mb${MB}" \ + --base.log_dir=/tmp/affine-nsys/logs \ + --base.checkpoint_dir=/tmp/affine-nsys/checkpoints \ + --base.checkpoint_interval=0 --base.eval_episodes=0 --base.seed=73 \ + --vec.total_agents=4096 --vec.num_buffers=1 \ + --policy.hidden_size=512 --policy.num_layers=3.93814 \ + --train.gpus=1 --train.horizon=64 --train.replay_ratio=1.87009 \ + --train.minibatch_size="$MB" --train.total_timesteps=8388608 +done +``` + +(`--base.profile=1` gates real `cudaProfilerStart()`/`Stop()` calls at +`src/pufferl.cu:2064` / `:2081`, so `--capture-range=cudaProfilerApi` scopes +correctly. Exit code 1 with `missing key env/perf` at the end is harmless — +`eval_episodes=0` means nothing to log; the `.nsys-rep` is still complete.) + +Reports pulled: `nsys stats --report cuda_api_sum / cuda_kern_exec_sum:base +/ cuda_gpu_kern_sum / nvtx_gpu_proj_sum --force-export=true `. + +**Results, MB8192 (32 epochs, 8.1M steps, ~3.38s uptime):** +- `cudaStreamSynchronize`: 76.7% of *API* time (2.756s total, 96 calls) — + this is the CPU blocked waiting for GPU graph replay to finish. It is a + measurement artifact of CPU-side view, not GPU idle time — ignore the + dashboard's "0ms train" rows for the same reason (CUDA-event timing is + unreliable inside a captured graph). +- **Real total GPU kernel-busy time (sum of `cuda_gpu_kern_sum`): only + ~529ms out of ~3.38s wall-clock — about 15%.** The other ~85% is gaps + between kernel launches / graph-node dispatch overhead. +- No single kernel exceeds **1.6%** of total accounted GPU time. Nothing to + point at as "the one slow kernel." +- All six Muon-named kernels (`muon_weight_update`, `muon_store_update`, + `muon_l2_normalize`, `muon_sum_sq_reduce`, `muon_sum_sq_partials`, + `muon_clip_nesterov`) combined: **~142ms**, a minority of the 529ms. +- `affine_lock.cu`'s own kernels (`gpu_affine_lock_shared_step_kernel` + + friends): **~15.8ms** — about 3% of real kernel time. **The env kernels + are not the bottleneck; block/lane tuning there caps at ~1-5% per the + reference agent's own report and won't move the needle meaningfully.** +- High-launch-count kernels: `splitKreduce_kernel` (cuBLASLt internal + reduction) fires **69,856 times** in one run at 1-2µs each; many + `cutlass::Kernel2<...>` GEMM variants (the real Muon/forward/backward + matmuls) fire thousands of times at a few µs each. **The signature is + "thousands of tiny kernel launches, each too small to amortize dispatch + overhead," not "one kernel is slow."** This matches a network this size + (hidden_size 256-512, ~3 real layers — note `num_layers=3.93814397` in + config *logs* as that value but **executes as 3 actual layers**, rounds + down) being fundamentally launch-overhead-bound on a 5090, not + compute-bound. + +**MB32768 comparison**: real kernel time drops to ~275ms (vs 529ms), +`train_forward_backward` NVTX-projected GPU time drops from ~103ms/epoch to +~70.8ms/epoch (1.46x, not the 4.21x reduction in optimizer-update count — +confirms per-update cost went up as batches got bigger, and the net win +comes from fewer total launches, not cheaper Muon math). + +**Reference agent's correction worth preserving**: Muon's ~75 GEMMs/update +mostly show up under generic cuBLAS names (`splitKreduce_kernel`, the +`cutlass::Kernel2<...>` variants) rather than the six Muon-named kernels — +so the *node-count* attribution to Muon is likely higher than the +Muon-named-kernel total of 142ms suggests, even though the physical +diagnosis (launch/dispatch overhead, not compute) holds either way. + +## Finding 6 — proposed core fixes (not yet started, need a separate branch) + +Two concrete pieces of engineering, both **`src/pufferl.cu` / `src/algo.cu` +core**, both affect every ocean env on GPU, **out of scope for PR #645**: + +1. **Stream-overlap for GPU-native env-step vs train** (addresses Finding + 4). Run next-buffer env-step on a separate CUDA stream concurrent with + the training graph's stream, instead of the current single-stream + serialization that forces `num_buffers=1`. This is the GPU-native + equivalent of CPU's double-buffering win. +2. **cuBLASLt separate C/D output for Muon** (reference agent's suggestion, + addresses Finding 5). Muon currently does a copy-before-GEMM step; + using cuBLASLt's separate C/D outputs could eliminate that copy and + remove roughly 50 graph nodes per optimizer update, while keeping the + same five-step optimizer math (i.e. should be numerically + equivalent — verify with byte-identical-checkpoint testing like the + reference agent did for their 3.46x claim). + +Neither has been started. Correctness bar for both: checkpoints and eval +metrics should match the CPU `.h` backend at matched hypers (see the +reference agent's own validation method in Finding 3 — byte-identical +checkpoints, "effectively identical" eval metrics). + +## Artifacts preserved + +- `/tmp/affine-nsys/mb8192.nsys-rep`, `/tmp/affine-nsys/mb32768.nsys-rep` — + full Nsight captures backing Finding 5. +- `/tmp/affine5090-tune/logs/affine_lock/*.ini` — all canary run logs + backing Findings 1-3 (`baseline-5090*.ini`, `matched-*.ini`, + `matched-small-*.ini`). +- These are in `/tmp`, not durable — copy them out before they get cleared + if you want to keep the raw data rather than just this summary. + +## What NOT to do + +- Don't touch `ocean/affine_lock/affine_lock.cu` expecting a big win — it's + already fast (Finding 5), that's not where the time goes. +- Don't touch `src/pufferl.cu` / `src/algo.cu` on the `affine-5.0-port` + branch / PR #645 — that PR is env-scoped only. Use a separate branch. +- Don't accept a minibatch-size change as a "speedup" without flagging that + it changes the training recipe (fewer/larger optimizer updates) — it's + not metric-equivalent to the baseline it's compared against. +- Don't compare CPU-vs-GPU SPS on a CPU-thread-starved machine and + generalize the ratio to strong hardware — that's exactly the mistake in + the original "3.46x" pitch (real number, wrong causal attribution). + +## 2026-08-20 profile correction and core plan + +The raw Nsight SQLite data does not support Finding 5's claim that total GPU +kernel-busy time was only about 529 ms or that roughly 85% of wall time was +empty graph-node dispatch gaps. In the MB32768 capture, the union of kernel +execution intervals is 2,244.640 ms across a 2,397.508 ms kernel span (93.6%). +The GPU can still be underfilled while a kernel is active, but the optimization +target is kernel shape, concurrency, fusion, and memory traffic rather than +another layer of host graph submission. + +See `FAST_CUDA_CORE_PLAN.md` for the corrected measurements, separate-branch +plan, permissions, optimization roadmap, correctness tiers, and golden +environment benchmark gates. diff --git a/FAST_CUDA_CORE_PLAN.md b/FAST_CUDA_CORE_PLAN.md new file mode 100644 index 0000000000..218cd9f648 --- /dev/null +++ b/FAST_CUDA_CORE_PLAN.md @@ -0,0 +1,1206 @@ +# CUDA trainer throughput plan + +Date: 2026-08-20 + +Status: investigation and benchmark design only. No core implementation has +started, no branch has been created, and no benchmark has been run under this +plan. + +This document is the durable handoff for making the PufferLib CUDA trainer +substantially faster on high-end GPUs. It supplements `FAST_CUDA.md` and +supersedes that document's interpretation that roughly 85% of the profiled +wall time was empty CUDA graph dispatch gaps. + +## Objective + +Push the training Pareto front forward on the RTX 5090, with a target of at +least 2x end-to-end training throughput where technically achievable, while: + +- preserving training behavior at matched hyperparameters for implementation + speedup claims; +- measuring final training quality, not SPS alone; +- avoiding regressions across established golden environments; +- keeping core trainer work out of the Affine Lock environment PR; +- keeping `logs/*.ini` and other constellation inputs untouched; +- separating strict implementation wins from retuned Pareto-front results. + +A 2-4x absolute improvement may be possible as a combined trainer, kernel, +and workload-shaping program. A 2-4x gain attributable only to +`ocean/affine_lock/affine_lock.cu` is not realistic on the 5090 host because +environment stepping is a small part of total runtime. + +Shared changes to `src/pufferl.cu` and `src/algo.cu` can move the absolute +Pareto front, but they also benefit CPU-environment (`.h`) builds. Results must +not present a shared learner optimization as an Affine Lock CUDA-environment +speedup. + +## Branch and PR boundary + +This work is larger than Affine Lock and belongs on a separate core branch. + +Recommended branch structure, to be created only after explicit approval: + +1. `cuda-trainer-throughput` from a clean `5.0` base. This contains shared + trainer changes, benchmark tooling, and core documentation. +2. A local-only integration branch or worktree combining + `cuda-trainer-throughput` with `affine-5.0-port`. This exists only to measure + Affine Lock before PR #645 lands. +3. No core commits should be added to the Affine Lock PR. +4. Nothing should be pushed until the correctness and golden-environment gates + in this document pass. + +If Affine Lock lands first, the integration branch becomes unnecessary and the +core branch can be updated from the new `5.0` base before final validation. + +Creating, switching, merging, rebasing, committing, or pushing branches is a +Git state mutation and requires explicit user approval. + +## Corrected profile interpretation + +The preserved Nsight databases do not support the earlier conclusion that the +GPU was executing kernels for only about 15% of wall time. + +### MB8192 capture + +- CUTLASS `Kernel2`: 178,272 launches, 2,830.752 ms summed duration. +- `splitKreduce_kernel`: 69,856 launches, 137.571 ms summed duration. +- Muon optimizer updates: 1,888, inferred from the one-per-update + `muon_weight_update` count. + +The CUTLASS duration alone is much larger than the earlier claimed total of +about 529 ms. + +### MB32768 capture + +- Total kernel launches: 100,837. +- Summed kernel duration: 2,695.073 ms. This double-counts overlap across + streams and is not a wall-time share. +- Union of all kernel execution intervals: 2,244.640 ms. +- Span from first kernel start to last kernel end: 2,397.508 ms. +- At least one kernel was active for 93.6% of that span. +- CUTLASS `Kernel2`: 50,112 launches, 1,771.521 ms summed duration. +- `splitKreduce_kernel`: 16,576 launches, 36.644 ms summed duration. +- Muon optimizer updates: 448. + +Comparing the two captures gives an exact decomposition of the launch counts: + +- 10,240 fixed CUTLASS kernels per captured run outside optimizer updates; +- 89 main GEMM kernels per optimizer update; +- 37 split-K reduction kernels per optimizer update. + +The GPU being occupied does not imply that every kernel fills the SMs. Small +serial GEMMs, scans, and elementwise kernels can keep the device nominally busy +while using only part of its resources. The revised working hypothesis is +underfilled or inefficient kernels plus unnecessary intermediate traffic, not +large empty gaps between graph nodes. + +Graph-node tracing can perturb a capture, so the absolute timing must be +re-established without node tracing before implementation decisions are +finalized. + +## Existing trainer architecture + +- GPU environment observations, actions, rewards, terminals, and masks are + device-resident. +- GPU rollout captures the entire horizon in one CUDA graph. +- Training captures preprocessing and the complete minibatch loop in one CUDA + graph. +- Async mode already overlaps the next rollout with learner execution on + separate nonblocking streams. +- Actor parameters are snapshotted separately from learner parameters. +- GPU environments currently require `vec.num_buffers=1`, but the trainer has + two async rollout/train slots. +- `base.cudagraphs=0` does not disable graphs in the current implementation; + only a negative value does. Benchmark commands must account for this. + +Consequences: + +- Capturing another larger host graph is not the primary opportunity. +- Basic rollout/train overlap is already present. +- Host synchronization cleanup may help, but cannot explain a 2x target. +- The learner and model kernels are the primary target. + +## Optimization roadmap + +### Track 0: repair measurement and attribution + +Before changing core code: + +1. Capture a matched baseline without CUDA graph node tracing. +2. Use Nsight Compute on representative Muon GEMMs, split-K reductions, + MinGRU scans, and forward/backward GEMMs. +3. Record tensor-core utilization, achieved occupancy, memory throughput, + launch count, kernel duration, and workspace/algorithm choice. +4. Attribute wall time to rollout, train preprocessing, model + forward/backward, Muon, copies, and logging. +5. Keep all binaries, logs, checkpoints, profiles, and reports under a unique + `/tmp/puffer-fast//` directory. + +### Track 1: Muon, formula-preserving first + +Current Muon processes five matrix parameters serially. Each matrix performs: + +- a norm reduction and normalization; +- five Newton-Schulz rounds; +- three GEMMs and two copies per round; +- a final store/scale operation. + +Across five matrices this is 75 Muon GEMMs and 50 copies per optimizer update. + +Implement in this order: + +1. Replace legacy copy-before-`cublasGemmEx` operations with cuBLASLt matmuls + using separate `C` input and `D` output buffers. Cache layouts, descriptors, + algorithms, and workspace at initialization so graph capture stays static. +2. Benchmark whole-graph GEMM algorithms, including non-split-K candidates. + Removing a split-K reduction can win end-to-end even when its main GEMM is + slightly slower in isolation. +3. Allocate disjoint per-parameter scratch and process independent matrices + concurrently. Batch the identically shaped MinGRU weights where profitable. +4. Use one handle and user-owned workspace per concurrent stream, with an + explicit fan-out/fan-in dependency before the final update. +5. Batch per-parameter norm, reduction, normalization, and store kernels. +6. Fuse the final FP32 weight update with the BF16 parameter write when this can + preserve the same values and visibility rules. + +These operations can preserve the optimizer formula. Byte-identical results +are not assumed because a different cuBLAS algorithm may change floating-point +reduction order. That must be tested rather than claimed. + +### Track 2: MinGRU and model kernels + +In the MB32768 capture, named MinGRU kernels account for about 764 ms of summed +kernel duration: + +- forward and backward scans: about 581 ms; +- gate and add kernels: about 183 ms. + +Candidate work: + +1. Fuse gate preparation and residual/add traffic where dependencies permit. +2. Improve the fixed-length affine scan with vectorized, coalesced, + warp/block-level implementations. +3. Avoid materializing intermediate scan operands when a fused kernel can + consume them directly. +4. Combine compatible projections into wider GEMMs. +5. Overlap independent `dW` GEMMs with critical-path recurrent backward work, + joining before optimizer execution. +6. Benchmark cached cuBLASLt algorithms for the exact fixed shapes. + +Changing scan association, fusing across BF16 stores, or introducing atomics +can change numerical trajectories. Those changes belong in a separately +labeled numerical-equivalence tier. + +### Track 3: CUDA-environment-specific cleanup + +This track is important for correctly attributing a `.cu` backend benefit, but +it is secondary for the current Affine Lock workload. + +1. Remove the three unconditional timing event nodes captured per rollout + timestep from the production graph. +2. Keep an instrumented graph only for sampled profiling runs. +3. Fuse reward, terminal, mask, and recurrent-state preparation. +4. Eliminate avoidable observation copies and rollout materializations. +5. Let GPU environments write train-ready fields directly where ownership and + graph-fixed pointers permit it. +6. Consider a persistent multi-step policy-plus-environment kernel only after + simpler learner work is measured. + +Do not retune `ocean/affine_lock/affine_lock.cu` block or lane geometry expecting +a large end-to-end win. Its measured kernel time is too small. + +### Track 4: aggressive Pareto work + +Only after formula-preserving changes are measured: + +1. Rewrite the Newton-Schulz polynomial in Horner form, reducing Muon from 75 + to 50 GEMMs per update. This is algebraically equivalent but changes BF16 + intermediate rounding. +2. Build specialized grouped or persistent tensor-core Newton-Schulz kernels + for common matrix shapes. +3. Build persistent or more heavily fused MinGRU paths. +4. Increase minibatch size or reshape agents/horizon to fill the 5090 better. +5. Retune the full training recipe and report a new Pareto front. + +These are not strict backend-equivalence results and must never be mixed into +the implementation-only speedup table. + +## Correctness tiers + +Every result must state one of these tiers. + +### C0: implementation-identical + +- Same hyperparameters, seed, samples, update count, and operation semantics. +- Byte-identical checkpoints are the preferred gate on the same GPU, toolkit, + compiler, and build flags. +- If checkpoints are not identical, the result cannot be called C0 without a + documented reason and a stronger equivalence proof. + +### C1: numerically equivalent + +- Same optimizer and training recipe, but floating-point association or + library algorithm can differ. +- Compare per-update losses, gradients, parameter differences, logits, and + evaluation metrics. +- Require predeclared tolerances and multi-seed non-inferiority. + +### C2: Pareto-retuned + +- Minibatch, agents/horizon, kernel algebra, precision, or other recipe details + may change. +- Compare score versus samples, score versus wall time, time to threshold, and + final quality across multiple seeds. +- Report this as a new Pareto point, not a matched-backend speedup. + +## Golden benchmark suite + +The suite deliberately spans CUDA and CPU environment backends plus very +different learner shapes. + +| Environment | Env backends | Canonical learner/workload | Benchmark role | +|---|---|---|---| +| Affine Lock | `.h`, `.cu` | hidden 512, effective 3 layers, MB 8192, horizon 64 | Primary target and backend parity | +| Breakout | `.h`, `.cu` | hidden 32, 4 layers, MB 65536, horizon 32 | Tiny learner and second true CUDA env | +| G2048 | `.h` only | hidden 1024, effective 5 layers, MB 65536, horizon 64 | Large learner dominated control | +| Maze | `.h` only | hidden 512, 5 layers, MB 16384, horizon 256 | Mid-size learner, long horizon, CPU transfer/overlap | +| Boxoban | `.h` only | hidden 1024, 3 layers, MB 65536, horizon 128 | Large observations/model and puzzle quality control | +| Benchmark env | `.h` | synthetic | Learner-focused diagnostic only, not a quality gate | + +Important naming and backend facts: + +- The repository environment is `boxoban`; there is no separate `sokoban` + environment. Sokoban is the game family. +- A normal build uses the CPU `.h` environment plus the CUDA learner. +- `build.sh --cpu` is a standalone CPU play/eval binary, not the correct + CPU-environment trainer baseline. +- `build.sh ENV OUTPUT --cu` requires `ocean/ENV/ENV.cu` and exclusively uses + that GPU environment implementation. +- G2048, Maze, and Boxoban have no `.cu` backend. They validate shared learner + changes but cannot validate CUDA-environment-specific speedups. +- Breakout's CUDA backend should be tested without unsupported self-play + features. +- G2048's fractional configured layer count is stored in an integer and + executes as five layers. + +## Benchmark harness design + +Create one dedicated harness rather than modifying or running existing seed +scripts unchanged. + +### Artifact isolation + +- Default root: `/tmp/puffer-fast//`. +- Separate `bin`, `logs`, `checkpoints`, `profiles`, `maps`, and `results` + subdirectories. +- Use separately named baseline and candidate binaries with identical compiler + flags. +- Run from the repository root because configuration lookup uses relative + `config/` paths. +- Disable external logging and evaluation for throughput canaries. +- Never write to the repository's `logs/` tree. +- Do not delete old campaign directories without explicit approval. +- Hash binaries, configs, and staged datasets for reproducibility. + +Existing `profile.sh` must not be used unchanged. It can silently reuse the +wrong `./puffer` binary, writes artifacts in the repository, force-overwrites +outputs, traces graph nodes by default, and reports summed kernel durations as +percentages even when streams overlap. + +Existing multi-seed scripts must not be used unchanged because some write +directly into protected `logs/`, use old CLI syntax, or group conditions in an +order vulnerable to thermal drift. + +### Performance gate + +1. Build baseline and candidate binaries independently. +2. Use the exact same config, seed, compiler, architecture flags, and runtime + overrides. +3. Alternate paired order as `A/B`, then `B/A` to reduce thermal/order bias. +4. Use at least five pairs for screening and ten pairs before a public claim. +5. Run long enough to amortize initialization and graph capture. +6. Discard warmup and compute throughput as change in `agent_steps` divided by + change in `uptime`; do not average logged SPS samples. +7. Report per-environment median, MAD, paired speedup ratios, geometric-mean + speedup, and bootstrap 95% confidence intervals. +8. Record CUDA/compiler versions, GPU clocks, temperature, power state, CPU + thread count, buffer count, run order, config overrides, and binary hash. +9. Report native wall time for end-to-end claims and `perf/*` splits only for + diagnosis. +10. Include Affine Lock and Breakout in both `.h` and `.cu` modes. Include + G2048, Maze, and Boxoban in `.h` mode. + +### Quality gate + +1. Use paired fixed seeds. The existing Breakout convention + `11, 22, 33, 44, 55` is a reasonable initial common seed set. +2. Use representative or full training budgets, not short throughput canaries. +3. Compare final score, score-versus-step AUC, score-versus-time AUC, success + fraction, and time/steps to predefined thresholds. +4. Define non-inferiority tolerances before viewing candidate results. +5. Evaluate checkpoint/trajectory parity for `.h` versus `.cu` Affine Lock and + Breakout where the semantics support it. +6. Report environment-specific quality metrics. + +Environment-specific metrics: + +- Affine Lock: solve rate, maximum solved depth, depth-specific solve rates, + efficiency, and the existing `env/perf` objective. +- Breakout: score, episode return, episode length, and any established + score/survival thresholds. +- G2048: score, merge score, maximum tile, and tile reach rates. +- Maze: score, episode return, completion/success rate if exposed, and episode + length. +- Boxoban: targets hit, solve/success rate if exposed, episode return, and + episode length. + +### Boxoban dataset safety + +The current repository does not contain the medium map binary required by the +default Boxoban run. A normal run can download levels and generate files under +`resources/boxoban`, which violates artifact isolation. + +Before Boxoban benchmarking: + +1. Stage a valid map binary under the campaign's `/tmp` map directory. +2. Record its checksum. +3. Pass `--env.map_bin=/tmp/puffer-fast//maps/...` explicitly. +4. Use the identical staged binary for every baseline and candidate run. + +## Acceptance criteria + +Exact thresholds should be finalized from baseline variance before inspecting +candidate results. Initial policy: + +- Primary target: at least 2x Affine Lock CUDA end-to-end SPS at the strictest + correctness tier that can support it. +- No golden environment may regress by more than 3% median SPS unless the + paired confidence interval shows the change is noise or the code path is + explicitly gated away from that environment. +- A shared core optimization should improve the suite geometric mean and not + merely transfer time between rollout and training counters. +- No quality metric may fail its predeclared non-inferiority tolerance. +- No new nondeterminism, race, OOM, or unsupported backend behavior. +- Memory growth and workspace use must be reported; a throughput win that + prevents canonical configurations from fitting is a regression. +- Every speed claim names its correctness tier and includes matched baseline + data. + +## Expected gain envelope + +- Graph and host synchronization cleanup alone: likely low single digits. +- Muon copy removal, algorithm selection, reduction batching, and independent + matrix concurrency: plausible 1.3-1.8x end-to-end if profiling confirms + underfilled Muon work. +- Muon plus MinGRU/model fusion: the best credible route to a strict-recipe 2x. +- Persistent kernels plus Pareto retuning: potentially 2-4x, but a larger + research effort and not a C0 claim. +- Strict byte-identical 4x on the current recipe: unlikely based on the current + kernel occupancy evidence. + +These are hypotheses, not promised results. The paired benchmark suite is the +gate. + +## Operational permissions and protected state + +Standing rules for this work: + +Allowed without asking: + +- Read files required for the investigation. +- Edit or create files inside the PufferLib repository, including core source. +- Research official technical documentation. + +Ask first: + +- Any `rm`, unlink, truncation, destructive overwrite, or cleanup. +- Any Git state mutation, including add, commit, push, branch, switch, stash, + reset, rebase, merge, or clean. +- Builds, training runs, benchmark runs, or profiler runs that consume the GPU + or create artifacts. +- Package, compiler, driver, system, or privilege changes. +- Process termination or edits outside this repository. +- Any external write such as W&B, GitHub, upload, or network publishing. + +Protected: + +- Treat `logs/*.ini` and the repository `logs/` tree as read-only. +- Never overwrite, move, truncate, or delete constellation logs. +- Route all experimental outputs to an isolated `/tmp` campaign. +- Stop and ask if unexpected external changes appear. + +## Proposed execution sequence + +1. Obtain approval to create the separate core and local integration branches. +2. Obtain approval for isolated build, baseline, and profiling commands. +3. Build the unified benchmark harness with `/tmp` defaults and no external + logging. +4. Establish the full golden baseline matrix before core changes. +5. Correct profiling and attribute Muon, MinGRU, model GEMM, rollout, transfer, + and logging time. +6. Implement the smallest formula-preserving Muon change. +7. Run microbenchmarks, then the performance gate, then the quality gate. +8. Keep or revert the change based on measured suite-wide results. +9. Continue one independently attributable optimization at a time. +10. Start C1/C2 aggressive work only after C0 opportunities are exhausted and + explicitly authorized. + +No speedup should be accepted because one short Affine Lock run got faster. +The unit of success is a reproducible, correctly classified improvement across +the golden matrix with preserved training quality. + +## Locked Affine Lock best-run reference + +The following configuration is the current best-run quality reference supplied +by the user. Do not silently change its minibatch, replay ratio, horizon, +network, optimizer, or evaluation settings when making an +implementation-equivalent speed claim. + +Run identity: + +- `base.run_id=sweep_1787042600825_0515` +- `base.env_name=affine_lock` +- `base.seed=73` +- `env.seed=42` +- `train.seed=42` +- `base.wandb=false` + +Backend and vectorization: + +- CPU environment backend: `vec.gpu_env=0` +- `base.async=1` +- `base.cudagraphs=1` +- `vec.total_agents=4096` +- `vec.num_buffers=2` +- `vec.num_threads=16` +- `train.gpus=1` + +Environment: + +- `env.start_depth=2` +- `env.max_depth=16` +- `env.num_agents=1` +- `env.num_bots=0` +- `env.step_grace=0` +- `env.perf_weighting=1` + +Network: + +- `torch.network=MinGRU` +- `torch.encoder=DefaultEncoder` +- `torch.decoder=DefaultDecoder` +- `policy.hidden_size=512` +- `policy.num_layers=3.93814397` +- `policy.expansion_factor=1` + +The current native trainer stores `policy.num_layers` in an integer, so this +configuration executes three MinGRU layers. Preserve the current conversion +behavior for matched results. + +Training: + +- `train.horizon=64` +- `train.minibatch_size=8192` +- `train.replay_ratio=1.87008977` +- `train.total_timesteps=1936646020` +- `train.learning_rate=0.00326743093` +- `train.anneal_lr=1` +- `train.min_lr_ratio=0` +- `train.gamma=0.999899983` +- `train.gae_lambda=0.921891332` +- `train.clip_coef=0.366497159` +- `train.vf_coef=0.100000001` +- `train.vf_clip_coef=0.00100000005` +- `train.ent_coef=0.0405269228` +- `train.anneal_ent_coef=0` +- `train.min_ent_coef_ratio=0.1` +- `train.momentum=0.924182773` +- `train.max_grad_norm=0.881812513` +- `train.vtrace=0` +- `train.verb_eps=0` + +Evaluation: + +- `base.eval_episodes=10000` +- `base.burnin_games=0` +- `base.reset_every_horizon=0` + +This exact run is a CPU-environment quality reference and uses two vector +buffers. The current CUDA environment requires one vector buffer, so it cannot +be presented as a byte-for-byte backend comparison without addressing that +structural difference. Use two distinct comparisons: + +1. Baseline versus candidate within the `.h` backend, holding this + configuration fixed. +2. Baseline versus candidate within the `.cu` backend, holding its valid + one-buffer configuration fixed. + +The final Pareto comparison can compare the best valid `.h` and `.cu` systems, +but must disclose the buffer difference. Increasing minibatch size or reducing +optimizer-update count belongs only in C2 Pareto-retuned results. + +## Universal-first optimization policy + +The first changes should remove objectively redundant work or preserve an +existing operation and dependency order. Do not replace one global path with a +path that helps medium matrices but slows small or large matrices. + +Rules: + +1. Retain the existing implementation as a fallback for shape-sensitive work. +2. Select GEMM or concurrency paths per exact shape and architecture rather + than applying one algorithm globally. +3. Require paired golden-suite measurements before enabling a new path by + default. +4. Reject or gate a path that causes a reproducible regression greater than 3% + in any golden environment. +5. Prefer removing copies, redundant event nodes, redundant scalar work, and + unnecessary synchronization before changing floating-point association. +6. Keep operation-order-changing fusion, grouped algorithms, persistent + kernels, and workload retuning in later, separately labeled work. + +Universal does not require every environment to gain the same percentage. It +means a core change has no material regression on unsupported shapes and uses +the old path when the new path is not a measured win. + +## Proposed PR split + +Each PR should be independently benchmarkable and should not depend on a large +unreviewable stack. + +### PR 1: benchmark and measurement infrastructure + +- Add the isolated golden-environment harness. +- Correct overlapping-kernel accounting. +- Add paired A/B statistics and quality gates. +- Add `.h`/`.cu` parity coverage for Affine Lock and Breakout. +- Make no training-math changes. + +### PR 2: exact trainer cleanup + +- Remove unconditional rollout timing nodes from production graphs. +- Replace avoidable host synchronizations with explicit stream dependencies. +- Make logging and scalar transfers asynchronous where ordering is unchanged. +- Remove other demonstrably redundant work that preserves the same kernels and + arithmetic order. + +### PR 3: Muon copy and GEMM path + +- Add cuBLASLt separate `C`/`D` output support. +- Eliminate Newton-Schulz copy-before-GEMM operations. +- Cache deterministic per-shape algorithms and preserve the legacy fallback. +- Treat the PR as C0 only if checkpoint identity actually passes; otherwise + evaluate and label it as C1. + +### PR 4: Muon concurrency and reduction batching + +- Add disjoint per-parameter scratch. +- Batch or concurrently execute only shapes with a measured win. +- Batch norm and elementwise work while preserving reduction order where + possible. +- Keep large, already-saturating matrices on the legacy serial path if that is + faster. + +### PR 5: MinGRU/model kernels + +- Optimize scans, gates, model GEMMs, and backward overlap. +- Separate exact scheduling changes from numerically different scan/fusion + algorithms. +- Require all five golden environments because this code is shared broadly. + +### PR 6: CUDA environment pipeline + +- Improve direct device staging and remove duplicate rollout materialization. +- Cover Affine Lock and Breakout `.cu` plus their `.h` controls. +- Keep environment-specific kernels out of the shared learner PRs. + +### Separate research results: Pareto retuning + +- Horner-form Muon, persistent kernels, minibatch changes, and agents/horizon + reshaping are C1/C2 experiments. +- Do not bundle them with low-risk core cleanup. +- Report them as new Pareto points rather than transparent implementation wins. + +The first two PRs provide a bounded go/no-go point before investing in the +deeper Muon and MinGRU work. + +## C0 validation record (2026-08-20) + +The first exact cleanup candidate removes unused per-step rollout timing event +nodes when CUDA graphs are enabled. Baseline and candidate were built +independently from clean source into a unique mode-0700 campaign tree under +`/tmp/puffer-c0-golden-jCEYhkTm`; repository logs and resources were not used. + +- Affine run 578 `.h`/`.cu` and the 1024x4 stress `.h`/`.cu` produced + byte-identical baseline/candidate checkpoints at 8,388,608 steps. +- Breakout `.h` was not self-repeatable with its production async settings. + It also was not self-repeatable with async disabled while retaining multiple + vector threads. A one-thread, async-off clean baseline canary was exactly + repeatable. +- With that deterministic canary, Breakout `.h`, Breakout `.cu`, G2048 `.h`, + and Maze `.h` produced byte-identical baseline/candidate checkpoints at + 67,108,864 steps. The retained campaign is + `/tmp/puffer-c0-golden-jCEYhkTm/puffer-throughput-rndn2pa8`. +- The short Affine campaign geomean was 1.0019x with a noisy worst pair of + 0.9870x. The deterministic non-Affine campaign geomean was 0.9974x. These + one-pair C0 screens establish neither a gain nor a regression. +- Boxoban `.h` was self-repeatable and produced byte-identical + baseline/candidate checkpoints at 33,554,432 steps using an immutable staged + 450,000-puzzle map. The archive SHA256 is + `fbd7b1efb4e7dd77e06d390051d60fbcc61f11efe127f0c5edaf9ec1547a417b` and + map SHA256 is + `87bf4fc7c180895f4b3a75d0393df9909feb8f97de76494c978273ec72bddb53`. + The retained campaign is + `/tmp/puffer-c0-golden-jCEYhkTm/puffer-throughput-jm7mqcmo`. +- The Boxoban corpus came from mutable upstream `main`; the recorded hashes + freeze this campaign but do not create a repository-defined canonical + corpus. A future fixture should pin the upstream revision or release. + +Exact canaries and production performance runs are deliberately separate. +Production async runs use paired statistics and record checkpoint variation; +full-budget, multi-seed learning-quality results remain mandatory before a +core trainer PR can merge. + +## Exact optimization experiment log + +These candidates were tested against the committed exact-cleanup binaries. +They are retained here even when rejected so later work does not repeat a +failed optimization. + +### Muon scalar broadcast: rejected + +Computing the clip coefficient and matrix inverse norm once per CUDA block +instead of once per element preserved byte-identical Affine checkpoints. A +balanced two-pair production screen measured `0.9983x` on Affine 578 CUDA, +`1.0060x` on 1024x4 CUDA, and `1.0021x` combined. This is below a credible +signal and not a universal win, so the source change was removed. Artifacts: +`/tmp/puffer-muon-scalar-DOOmXEUF/puffer-throughput-dcmitjed`. + +### Batched Muon norm launches: rejected + +Batching the identical per-matrix partial reduction, final reduction, and +normalization work from `3P` graph nodes to three preserved byte-identical +checkpoints. A balanced two-pair production screen measured `1.0040x` on +Affine 578 CUDA, `0.9983x` on 1024x4 CUDA, and `1.0011x` combined; the +1024x4 pairs disagreed in direction. The source change was removed. Artifacts: +`/tmp/puffer-muon-batchnorm-pKIiyQFT/puffer-throughput-0t98qn5h`. + +### Actor snapshot host-wait removal: rejected as non-C0 + +Replacing the per-epoch host synchronization with explicit actor-ready +stream events passed the Affine 578 CUDA checkpoint but changed the 1024x4 +CUDA checkpoint under production async settings. The speed screen was skipped +and the source change was removed. Artifacts: +`/tmp/puffer-actor-ready-iHjn6riQ/puffer-throughput-4vrgbtor`. + +These results reinforce the corrected profile interpretation: fixed graph +bookkeeping is not the 5090 bottleneck. Further work must reduce or overlap +the dominant Muon GEMMs/copies or optimize MinGRU computation while continuing +to apply the exact-checkpoint gate first. + +### Concurrent per-matrix Muon: promoted + +Running each independent 2D parameter's unchanged Muon pipeline on a disjoint +stream, handle, workspace, norm scratch, and NS scratch produced the first +clear C0 gain. The main stream forks after global clip/Nesterov and rejoins all +lanes before the unchanged flat weight update. + +- All Affine, Breakout, G2048, Maze, and Boxoban deterministic canaries + produced byte-identical baseline/candidate checkpoints. +- The final balanced two-pair production build measured `1.1926x` on run 578 + CUDA and `1.0665x` on 1024x4 CUDA, for `1.1278x` combined. Every pair + improved and every checkpoint matched. Artifacts: + `/tmp/puffer-muon-final2-b6gIHRz7/puffer-throughput-9y3xw5y0`. +- A deterministic Breakout CUDA sentinel measured `1.1565x`. The broader + exact promotion screen measured `1.1963x` Breakout CUDA, `1.0243x` + Breakout CPU, `1.0995x` Maze, and `1.0095x` Boxoban in one pair each. +- G2048 initially measured `0.9973x` in a balanced two-pair production run. A + shape-based saturated-workload fallback now selects the legacy serial path + only when there are at least seven matrices, at least five heavy + `3072x1024`-class lanes, and at least 15 Mi matrix elements. With that gate, + G2048 produced byte-identical checkpoints and `1.0009x` across two balanced + pairs during selector qualification. Artifacts: + `/tmp/puffer-muon-concurrent-gated-Bp5B17Q6/puffer-throughput-35l7l3c4`. + +The final two-pair production golden screen measured `1.1129x` Breakout CUDA, +`1.0292x` Breakout CPU, and `1.0281x` Maze. Boxoban at `0.9972x` and the +G2048 serial fallback at `0.9953x` were statistically flat; all pairs remained +above `0.987x`. Production checkpoint variation was recorded rather than +treated as candidate divergence because those async configurations are not +self-repeatable. The deterministic final-build canaries were byte-identical in +all cases. Artifacts: +`/tmp/puffer-muon-final2-b6gIHRz7/puffer-throughput-_vlevnd0`. + +The selected implementation uses one private 32 MiB cuBLAS workspace per +enabled matrix. Concurrency is capped at eight matrices, bounding private +workspace use at 256 MiB; larger models use the allocation-free serial path. +The two legacy 32 MiB allocations were removed because their per-call +`cublasSetStream` reset them to the default pool before every GEMM. A fixed +three-lane experiment reduced memory but lost `5.5%` relative throughput on +Affine 512. An 8 MiB-per-lane experiment stayed exact but lost about `0.5-0.7%` +relative throughput. The 32 MiB per-matrix version remains the measured speed +Pareto point for this 5090 campaign. +## Concurrent Muon cleanup validation + +Removed the experimental CUDA/cuBLAS/host assertion-and-abort wrappers. The +optimization now follows the existing direct-call style; the max-eight lane +gate, saturated-workload serial fallback, workspace policy, scheduling, +fork/join topology, and optimizer arithmetic are unchanged. + +Fresh binaries: + +- `/tmp/puffer-clean-bin-dW4YM5Cg` + +Deterministic exact-checkpoint results against the immutable pre-optimization +baseline `/tmp/puffer-c0-golden-jCEYhkTm/candidate-bin`: + +- Affine 578 `.h` and `.cu`: exact. +- Affine 1024x4 `.h` and `.cu`: exact. +- Breakout `.h` and `.cu`: exact. +- G2048 `.h`: exact. +- Maze `.h`: exact. +- Boxoban `.h`: exact. +- Affine artifact (the combined campaign stopped after these successful cases + when the following short case lacked enough uptime samples): + `/tmp/puffer-clean-validate-kFHz6jz5/puffer-throughput-7jkeeli3`. +- Breakout artifact: + `/tmp/puffer-clean-breakout-v9mapnzs/puffer-throughput-6mh11ek9`. +- G2048/Maze/Boxoban artifact: + `/tmp/puffer-clean-heavy-goldens-DJVTxiOy/puffer-throughput-78q3o_ak`. + +Balanced production throughput: + +- Affine 578 CUDA: `1.176870159x` geomean, pair range + `1.153879229x` to `1.200319180x`. +- Affine 1024x4 CUDA: `1.076241413x` geomean, pair range + `1.072833497x` to `1.079660154x`. +- Combined Affine geomean: `1.125431651x`. +- All Affine checkpoints matched exactly. +- Artifact: + `/tmp/puffer-clean-affine-balanced-nqAzWvrl/puffer-throughput-50ps5sl8`. + +The single deterministic G2048 timing pair was noisy (`0.962639308x`), so the +unchanged serial-fallback case was repeated for three production pairs. The +repeat measured `1.000669099x` geomean with a `0.997054066x` worst pair and +matching checkpoints: +`/tmp/puffer-clean-g2048-repeat-lwKRoN0E/puffer-throughput-l4dudlth`. +## Minimum-code concurrent Muon ablations + +Goal for this pass: retain only code that is required for measured throughput, +bit-identical arithmetic, CUDA-graph fork/join, or bounded resources. New +defensive wrappers and unused generalizations are not part of the optimization. + +### Retained cleanup + +- Removed all custom CUDA/cuBLAS/allocation assert-and-abort wrappers. +- Restored the original legacy cuBLAS initializer and call sites. +- Removed redundant pointer initialization and optional-workspace plumbing. +- Removed the abandoned shared-lane load balancer, matrix-to-lane indices, + max-scratch aggregation, and duplicate warmup search. +- Replaced heap metadata/lane arrays with fixed eight-entry storage, matching + the measured and qualified eight-lane cap. +- Folded parameter discovery and matrix descriptor construction into one scan. +- Replaced duplicate concurrency state with `num_lanes` (`0` means serial). +- Removed the redundant total-element saturation threshold; five heavy + `3072x1024`-or-larger matrices already imply the same 15 Mi-element bound. +- Reused the existing maximum-dimension scan for the 4096 qualification gate. +- Replaced unused tensor-shaped lane scratch metadata with raw device pointers. +- Packed each lane's 256 norm partials and one norm scalar into one allocation. +- Kept the original shared serial scratch allocation/registration unchanged. +- Removed the separate lane cuBLAS initializer. Lane setup uses the original + initializer, binds the private stream, then restores the private workspace + because `cublasSetStream` resets it. + +### Ablation: remove lane GEMM warmup - kept removed + +The entire per-lane cuBLAS warmup helper and calls were deleted. Fresh lane +handles successfully captured and replayed without it. Affine 578 CUDA remained +byte-identical and measured `1.006785243x` versus the already-clean concurrent +candidate in the initial canary. + +Artifact: +`/tmp/puffer-min-canary-XSgjuhth/puffer-throughput-9bnq23jd`. + +Decision: keep the warmup deleted. + +### Ablation: remove largest-first scheduling - rejected + +Removing `MuonMatrix::work` and the stable insertion sort preserved exact +checkpoints but reduced throughput versus the sorted minimum-code candidate: + +- Affine 578 CUDA: `0.974554593x` (`-2.54%`). +- Affine 1024x4 CUDA: `0.987948112x` (`-1.21%`). +- Combined: `0.981228501x`. + +Artifact: +`/tmp/puffer-nosort-ab-dBHD6INB/puffer-throughput-hcr26xye`. + +Decision: restore the small work field and stable largest-first insertion sort. +Those lines have measured value and remain in the implementation. + +### Final minimum-code validation + +Fresh binaries: +`/tmp/puffer-min-final-bin-e985GY0q`. + +Deterministic checkpoint comparison against the immutable pre-optimization +baseline passed exactly for Affine 578 `.h/.cu`, Affine 1024x4 `.h/.cu`, +Breakout `.h/.cu`, G2048 `.h`, Maze `.h`, and Boxoban `.h`. + +- Affine/G2048/Maze/Boxoban artifact: + `/tmp/puffer-min-final-exact-Gwxuuy1H/puffer-throughput-zthgmrb4`. +- Breakout artifact: + `/tmp/puffer-min-final-breakout-JHe4038Y/puffer-throughput-_ib_t7di`. + +Three-pair production Affine result against the immutable pre-optimization +baseline: + +- Affine 578 CUDA: `1.167157783x` geomean, range `1.159652395x` to + `1.181321136x`. +- Affine 1024x4 CUDA: `1.073290530x` geomean, range `1.063182301x` to + `1.079850177x`. +- Combined Affine geomean: `1.119240544x`. +- Every paired checkpoint matched exactly. +- Artifact: + `/tmp/puffer-min-final-affine-yxtO2Tb8/puffer-throughput-ud3m_3ae`. + +## 2026-08-20: post-Muon hotspot profiling and strict-C0 ablations + +Baseline for this round: commit `9a2afdae` (`cuda: overlap independent Muon matrix updates`) on RTX 5090 with CUDA 13.1. All candidates below used the native throughput harness, isolated `/tmp` binaries/artifacts, and exact checkpoint comparison. No candidate was retained in `src/algo.cu`; the committed Muon implementation remains the source baseline. + +### Nsight Systems attribution + +| GPU work class | Affine 578 | Affine 1024x4 | +| --- | ---: | ---: | +| Model/non-Muon GEMMs | 43.21% | 46.95% | +| Muon lane GEMMs | 31.26% | 35.29% | +| MinGRU custom kernels | 14.24% | 10.09% | +| Muon lane D2D copies | 4.23% | 2.57% | +| Muon lane custom kernels | 2.46% | 1.61% | + +Reports: `/tmp/puffer-nsys-affine578-AeLPq6DN/affine578.nsys-rep` and `/tmp/puffer-nsys-affine1024-xrrO5zxs/affine1024.nsys-rep`. MinGRU scan backward was about 7.6% of summed kernel time on 578 and 4.7% on 1024x4; forward was about 3.8% and 2.4%. Nsight Compute hardware-counter collection was attempted but rejected by the driver with `ERR_NVGPUCTRPERM`; no permission or driver setting was changed. Log: `/tmp/puffer-ncu-gemm-4bbRUQRm/ncu.log`. + +### Ablation ledger + +| Candidate | Exact checkpoints | Result | Decision | +| --- | --- | --- | --- | +| Rollout MinGRU state update in place, deleting one D2D copy per layer | Yes | Affine suite `1.000075x`; isolated Boxoban at original 256-thread scans `0.992227x` over 3 pairs | Reject: not universally free | +| Ordinary MinGRU scans at 128 threads, initially stacked on in-place rollout | Yes | Short Affine test: 578 `1.002227x`, 1024x4 `1.010714x`, suite `1.006461x` | Promising short result, required clean isolation | +| Ordinary MinGRU scans at 512 threads, stacked on in-place rollout | Yes | 578 `1.009817x`, 1024x4 `0.999703x`, suite `1.004747x` | Reject: size-dependent | +| Ordinary MinGRU scans at 64 threads, stacked on in-place rollout | Yes | Versus in-place 256: 578 `1.013120x`, 1024x4 `1.013406x`; versus committed baseline over 3 pairs: suite `1.010901x` | Reject: Boxoban `0.985530x` over 3 pairs | +| Ordinary MinGRU scans at 128 threads with the in-place change removed | Yes | Affine 578 `0.999373x`, 1024x4 `0.997186x`, suite `0.998279x`; Boxoban `0.998644x` | Reject: no repeatable gain | + +The positive short 128/64 results were not accepted because the better-isolated repeats contradicted them or exposed an environment regression. This is why throughput changes need paired repeats and the golden suite even when arithmetic is trivially bit-preserving. + +Ordinary scan shapes were: Affine 578 `B*H=65,536,T=64`; Affine 1024x4 `131,072,64`; Breakout `65,536,32`; G2048 `1,048,576,64`; Boxoban `524,288,128`. Maze (`32,768,256`) uses the existing row-scan path and was unaffected. All ordinary totals divide both 128 and 256 exactly, so the rejected geometry experiments changed scheduling only, not arithmetic or tail behavior. + +### Golden results for the rejected 64-thread/in-place candidate + +All nine deterministic backend/config checkpoints matched exactly: Affine 578 `.h/.cu`, Affine 1024x4 `.h/.cu`, Breakout `.h/.cu`, G2048 `.h`, Maze `.h`, and Boxoban `.h`. The full one-pair artifact roots are `/tmp/puffer-hotspots-golden-UecOhctF/standard/puffer-throughput-f9kt97kw` and `/tmp/puffer-hotspots-golden-UecOhctF/breakout/puffer-throughput-l9sef9qv`. Repeat artifacts are `/tmp/puffer-hotspots-regression-check-c4BKFlvl/maze-boxoban/puffer-throughput-uwk0wf_t` and `/tmp/puffer-hotspots-regression-check-c4BKFlvl/breakout-h/puffer-throughput-kpyhrbuu`. + +### Muon copy-overlap result + +The private-copy-stream experiment preserved every GEMM, copy, coefficient, and checkpoint bit, but it did not clear the performance gate. Over two Affine pairs, 578 was noisy at `1.008304x`, 1024x4 consistently regressed to `0.997829x`, and the suite result was only `1.003053x`. The added stream/events were reverted. Artifact: `/tmp/puffer-muon-copy-overlap-ab-I2JJdU62/puffer-throughput-7p1jziev`. + +### Next direction + +Further work should target arithmetic-preserving work removal rather than more dW scheduling. Trace parsing found a median final dW tail of only `9.920 us` on Affine 578 (about `0.79%` of iteration time) and `7.745 us` on 1024x4 (about `0.12%`). The previous dW was still running when the next layer became ready only once in 7,552 opportunities on 578 and never in 5,577 opportunities on 1024x4, so neither a second dW stream nor stream-priority tuning is justified. cuBLASLt separate-C/D and GEMM autotuning remain deferred because different kernels/reduction orders are not guaranteed bit-identical. + +### MinGRU training-fusion ablations + +Inter-layer backward-add fusion was also rejected. Instead of materializing each upper-layer `BF16_round(FP32(dX) + FP32(highway))`, the next lower scan reconstructed that exact rounded value from the two BF16 sources before its unchanged arithmetic. The bottom add remained materialized for encoder backward, and both ordinary and row scans were supported. The cleaned implementation was roughly 25-30 net lines. A three-pair run measured Affine 578 `1.008689x`, 1024x4 `1.001612x`, suite `1.005144x`, with exact checkpoints and worst pair `0.995648x`. All nine final golden checkpoints matched exactly, but the repeated Breakout-h performance gate measured `0.995334x` with a `0.960713x` worst pair. That is not enough gain or stability to justify the code. Affine artifact: `/tmp/puffer-grad-add-clean-ab-MZ967oyD/puffer-throughput-s2g89pvp`; golden artifacts: `/tmp/puffer-grad-add-clean-golden-ObwQjjBP/standard/puffer-throughput-2lplq623` and `/tmp/puffer-grad-add-clean-golden-ObwQjjBP/breakout/puffer-throughput-mxvbxckj`; repeat artifact: `/tmp/puffer-grad-add-clean-regression-tFpJLK6m/breakout-h/puffer-throughput-tswg4v27`. + +An additional saved-input-copy fusion was rejected. It raw-copied the exact loaded `precision_t` input from each forward scan into the existing backward-save buffer, removing one serialized D2D node per layer while preserving allocator layout and bits. Combined with backward-add fusion it measured 578 `1.011174x` but 1024x4 `0.995644x`, including a `0.980087x` pair; the extra scan-store pressure outweighed the removed copy on the larger net. Artifact: `/tmp/puffer-scan-fusions-ab-HdiVwMZc/puffer-throughput-x68n0j3o`. + +## 2026-08-20: pinned cuBLASLt critical-path GEMMs + +Nsight attribution mapped recurrent train projection, rollout projection, and dX to `26.06%` of Affine 578 and `32.09%` of 1024x4 raw GPU work. A standalone CUDA 13.1 tuner searched the six exact signatures in native-row and legacy-column encodings, filtering every candidate by device-side BF16 bit equality against production `cublasGemmEx`. Source and complete results are `/tmp/puffer_cublaslt_tuner.cu` and `/tmp/puffer_cublaslt_tuner_results.txt`. + +The retained production path pins only two H512 signatures: train projection `(8192,1536,512,N/T)` and recurrent dX `(8192,512,1536,N/N)`. Both use native-row cuBLASLt algo 21 with tile 15, split-K 1, reduction 0, swizzle 0, custom 0, stages 12, and zero workspace. Qualification used `C == D`, cold CUDA graph capture/replay, two deterministic BF16 distributions, repeated bit comparisons, and production-equivalent 256-byte alignment. Graph medians were `1.258x` for projection and `1.124x` for dX. Qualification artifact: `/tmp/puffer_cublaslt_graph_aligned_results.txt`. + +Enablement is deliberately narrow: BF16 build, CUDA runtime 13.1, driver 13.0, cuBLAS 13.2.1, exact `NVIDIA GeForce RTX 5090` SM120/170-SM device, main cuBLAS handle, alpha `+1`, beta `+0`, exact signature/ops, and A/B/C pointers all aligned to 256 bytes. Every miss uses the original `cublasGemmEx` path. Float builds compile with the pinned path excluded. `build.sh` links cuBLASLt explicitly for trainer and profiler targets. + +The accepted five-pair production result is Affine 578 `1.015413x`, with every pair positive (`1.010564x` minimum), and 1024x4 fallback `1.003004x`. All checkpoints matched exactly. Artifact: `/tmp/puffer-lt-h512-ab-iY8D0W9e/puffer-throughput-y5j0vdxj`. + +All nine deterministic golden/backend checkpoints matched exactly. Affine 578 measured `1.038142x` on `.cu` and `1.032719x` on `.h`; unqualified 1024x4, G2048, Maze, Boxoban, and Breakout cases remained effectively flat on legacy fallback. Artifacts: `/tmp/puffer-lt-h512-golden-6l11xVS2/standard/puffer-throughput-g67kpoi0` and `/tmp/puffer-lt-h512-golden-6l11xVS2/breakout/puffer-throughput-nfpyrrfp`. + +Two qualified H1024 plans were rejected despite faster isolated kernels. Rollout projection was `1.169x` faster at kernel median in Nsight, and train projection was `1.026x` in aligned graph replay, but the four-plan end-to-end gate regressed 1024x4 to `0.990571x` over five pairs. Both entries were removed rather than trading large-net SPS for run-578 speed. Artifact: `/tmp/puffer-lt-four-plan-ab-ib0QWhyt/puffer-throughput-d0m8ianh`. A bounded broader H1024 dX search checked 49,326 configurations and found zero strict-bit, zero-workspace survivors; results are `/tmp/puffer_cublaslt_h1024_dx_search_results_v3.txt`. + +## 2026-08-20: agent-oriented terminal-state reset + +Accepted `zero_term_state_agents`: one block handles one agent, thread 0 loads +that agent's terminal once into shared memory, the block uniformly returns for +a nonterminal, and its threads coalescently zero the compact state across all +layers. The optimized launch is gated on +`num_layers * hidden_size > BLOCK_SIZE && count >= BLOCK_SIZE`; smaller shapes +retain the original flat kernel. State offsets, stream ordering, and +`from_float(0.0f)` writes are unchanged, so every written BF16 zero bit remains +exact. + +Cleaned three-pair Affine result: + +- Affine 578: `1.004842x`. +- Affine 1024x4: `1.004952x`. +- Combined suite: `1.004897x`. +- Artifact: `/tmp/puffer-agent-reset-clean-ab-sxhyHfJi/puffer-throughput-fu0uydwm`. + +All nine one-pair golden/backend cases matched exactly: Affine 578 `.h/.cu`, +Affine 1024x4 `.h/.cu`, Breakout `.h/.cu`, G2048 `.h`, Maze `.h`, and Boxoban +`.h`. Artifacts: +`/tmp/puffer-agent-reset-clean-golden-ogVgLH6U/standard/puffer-throughput-eoaa5404` +and +`/tmp/puffer-agent-reset-clean-golden-ogVgLH6U/breakout/puffer-throughput-_wa3tayd`. + +The repeated three-pair Affine 1024x4 CUDA gate was exact and measured +`1.004387x`, with every pair positive and a `1.003429x` minimum. Artifact: +`/tmp/puffer-agent-reset-repeat-WYhbJQq7/puffer-throughput-1oc4_66g`. + +### Ablation: Muon normalization dual-write/copy elision - rejected + +Writing the normalized Muon matrix directly to both required destinations +removed its following D2D copy and preserved exact checkpoints. The initial +five-pair gate measured Affine 578 `1.004766x`, Affine 1024x4 `1.007477x`, and +the combined suite `1.006121x`. Artifact: +`/tmp/puffer-muon-dualnorm-ab-96JJXWVX/puffer-throughput-2k0gyi23`. + +An independent three-pair confirmation contradicted that result: Affine 578 +was `1.002256x`, Affine 1024x4 was `0.992711x`, the suite was `0.997472x`, and +the worst pair was `0.979177x`. Artifact: +`/tmp/puffer-muon-dualnorm-confirm-P2RcsY8U/puffer-throughput-u5ceobdw`. +The pooled nominal result was about `1.00287x`, but the instability and paired +confidence interval did not establish a gain. Decision: reject the ablation; +the source change was reverted. + +### Ablation: Muon update-side endpoint fusion - rejected + +This candidate used an all-matrix exact-coverage gate and fused +`store_update` with `weight_update` while preserving the intervening BF16 +boundary. The implementation was subsequently hardened to advance its source +with each registration's allocator cursor rather than assuming packed tensor +sizes. + +The exact three-pair screen measured Affine 578 `1.009736x`, Affine 1024x4 +`1.007901x`, and the combined suite `1.008818x`. Artifact: +`/tmp/puffer-muon-weight-fuse-screen-1YXtRhOB/puffer-throughput-r4k10yja`. +An independent exact three-pair run of the hardened candidate fell to +`0.999061x`, `1.001957x`, and `1.000508x`, respectively. Artifact: +`/tmp/puffer-muon-weight-fuse-confirm-1Dr4Ocxe/puffer-throughput-6dw9n3u5`. + +A decisive exact three-pair run at four times the duration measured Affine 578 +`0.999569x`, Affine 1024x4 `0.998889x`, and the suite `0.999229x`. Artifact: +`/tmp/puffer-muon-weight-fuse-long-dyHU2XR0/puffer-throughput-xoytjwwf`. +Decision: reject and revert; the longer run establishes no SPS gain. + +### Accepted: Muon clip-side endpoint fusion + +The retained path has a strict all-matrix concurrent-local gate. It leaves the +global raw gradient norm unchanged, then fuses clip/Nesterov with each +matrix's norm partial while preserving the original BF16 round and reload and +the exact reduction mapping. Serial and mixed-coverage cases retain the +original path. Cleanup removed the persistent eligibility flag and derives the +condition locally without changing dispatch. + +All Affine checkpoints were exact across four independent rounds: + +- Three-pair screen: Affine 578 `1.001991x`, Affine 1024x4 `1.007118x`, suite + `1.004551x`; artifact + `/tmp/puffer-muon-clip-fuse-screen-yQdQ6Ddg/puffer-throughput-7661ejf1`. +- Three-pair long run: `1.005827x`, `1.001692x`, suite `1.003758x`; artifact + `/tmp/puffer-muon-clip-fuse-long-6CtCrAxy/puffer-throughput-gfnjx_vb`. +- Three-pair preconditioned long run: `1.007203x`, `1.000872x`, suite + `1.004032x`; artifact + `/tmp/puffer-muon-clip-fuse-final-XmgXUMzz/puffer-throughput-13zp5sba`. +- Cleaned five-pair long run: `1.004353x`, `1.018761x`, suite `1.011531x`; + artifact + `/tmp/puffer-muon-clip-clean-final-OFyMj0Ju/puffer-throughput-5c05c2uh`. + +The conservative equal-batch pool of the first three rounds was Affine 578 +`1.005005x`, Affine 1024x4 `1.003224x`, and suite `1.004114x`. All nine golden +backend/configuration checkpoints matched exactly. Artifacts: +`/tmp/puffer-muon-clip-clean-golden-MczA66kv/standard/puffer-throughput-y3eqf016` +and +`/tmp/puffer-muon-clip-clean-golden-MczA66kv/breakout/puffer-throughput-9il5s6g5`. +The negative one-pair Maze timing was repeated for five preconditioned exact +pairs and measured `0.999185x`, consistent with noise rather than a meaningful +fallback regression. Artifact: +`/tmp/puffer-muon-clip-maze-repeat-RVJS5MqK/puffer-throughput-w0_7udpj`. + +Decision: accept the cleaned clip-side fusion. + +### Accepted: H512 Muon separate-C/D cuBLASLt plans + +The final path keeps the legacy Gram GEMM and pins only two H512, +zero-workspace cuBLASLt plans with separate C and D. This removes ten +intermediate copies per Muon matrix while retaining the exact intermediate +BF16 boundaries. Enablement requires the qualified platform, exact +shape/operation, pointer alignment, and concurrent-Muon path; every miss uses +the complete legacy implementation. Cleanup reduced dispatch to direct plan +indexing. + +Standalone DAG qualification covered both homogeneous and mixed algorithms. +Intermediate and final outputs matched exactly in the mixed checks, with DAG +speedups of `1.057064x` for Affine 578 and `1.026971x` for Affine 1024x4. +Artifacts: `/tmp/puffer_muon_lt_dag_results.txt` and +`/tmp/puffer_muon_lt_mixed_dag_results.txt`. + +The initial four-plan production screen was exact and measured Affine 578 +`1.013303x` and Affine 1024x4 `0.999484x`. Artifact: +`/tmp/puffer-muon-lt-screen-PVUjMwcb/puffer-throughput-xnjxu1ia`. Its exact +long confirmation measured `1.004826x` and `1.003628x`. Artifact: +`/tmp/puffer-muon-lt-long-jWvXrcOw/puffer-throughput-5q9xs6cz`. The H1024 +plans were then ablated because their pooled SPS contribution was marginal; +removing them also deleted four lines from the production path. + +H512-only qualification remained exact: + +- Screen: Affine 578 `1.007749x`, Affine 1024x4 legacy fallback `1.002107x`; + artifact + `/tmp/puffer-muon-lt-h512-screen-SUJfPhvu/puffer-throughput-8mkep5z1`. +- Five-pair long run: Affine 578 `1.008373x` with every pair positive, Affine + 1024x4 fallback `0.998299x`; artifact + `/tmp/puffer-muon-lt-h512-final-bUnngM4U/puffer-throughput-44_3abuy`. +- Final-clean screen: Affine 578 `1.017608x`, Affine 1024x4 fallback + `0.999136x`; artifact + `/tmp/puffer-muon-lt-final-screen-xNW8oq08/puffer-throughput-mfexwvxh`. + +All nine golden backend/configuration checkpoints matched exactly. Artifacts: +`/tmp/puffer-muon-lt-final-golden-HuZIm3Hk/standard/puffer-throughput-dmm5i97r` +and +`/tmp/puffer-muon-lt-final-golden-HuZIm3Hk/breakout/puffer-throughput-cu9w63s5`. +The H512 Maze repeat was exact and measured `1.004782x`. Artifact: +`/tmp/puffer-muon-lt-maze-repeat-5OQIpmuE/puffer-throughput-baupfgty`. + +Decision: accept the minimized H512-only separate-C/D plans with unconditional +legacy fallback outside their strict gate. + +### Ablation: exhaustive H512 rollout cuBLASLt search - rejected + +The bounded search checked 49,326 configurations, including 1,398 legal +zero-workspace candidates and 1,147 that matched the legacy output directly. +The fastest 64 graph finalists all remained exact, but every one was slower: +legacy mean `0.034835 ms` versus best cuBLASLt `0.038926 ms` (`0.8949x`). +Artifact: `/tmp/puffer_cublaslt_h512_rollout_search_results.txt`. Decision: no +production integration. + +### Ablation: ordinary MinGRU backward compile-time T64 - rejected + +Making only the ordinary backward scan's `T=64` loop bound compile-time +constant preserved exact checkpoints but regressed both Affine cases: 578 +`0.995279x`, 1024x4 `0.998330x`, suite `0.996803x`, with a `0.988153x` worst +pair. Artifact: +`/tmp/puffer-mingru-bwd-t64-screen-EifWk50K/puffer-throughput-h6uw3oz2`. +Decision: reject the constant-only specialization. + +### Ablation: ordinary MinGRU backward T64 unroll-4 - rejected + +The isolated T64-only kernel with `#pragma unroll 4` preserved exact +checkpoints and measured Affine 578 `1.004573x`, Affine 1024x4 `0.997249x`, +and suite `1.000904x`, with a `0.995564x` worst pair. Artifact: +`/tmp/puffer-mingru-bwd-unroll4-screen-8SAFL9nf/puffer-throughput-8hkvlc00`. +Decision: reject because H1024 regressed and the temporary specialization cost +99 lines; restore the committed ordinary scan. + +### Accepted: master-weight BF16 cast fusion + +`muon_weight_update` already materializes the exact FP32 `new_weight`; the +retained change stores that value to the master weights and writes +`from_float(new_weight)` to the BF16 mirror in the same kernel. It removes only +the immediate post-Muon cast. Initialization and model-load casts remain, and +the float compile-time path is unchanged. The final implementation is one net +line across `algo.cu` and `pufferl.cu`. + +The exact three-pair screen measured Affine 578 `1.001907x`, Affine 1024x4 +`1.001135x`, and suite `1.001521x`. Artifact: +`/tmp/puffer-weight-cast-fuse-screen-qdBiX1E7/puffer-throughput-vorooe3a`. +The exact five-pair long run measured `1.001355x`, `1.002141x`, and +`1.001748x`, respectively. Artifact: +`/tmp/puffer-weight-cast-fuse-long-jK2weCV6/puffer-throughput-p0pwattw`. + +All nine golden backend/configuration checkpoints matched exactly. Artifacts: +`/tmp/puffer-weight-cast-fuse-golden-MaFJtGi7/standard/puffer-throughput-u7goii9y` +and +`/tmp/puffer-weight-cast-fuse-golden-MaFJtGi7/breakout/puffer-throughput-7b2vf3q5`. +Their one-pair SPS values were noisy; exactness is the authoritative golden +gate for this optimizer-only change. + +Decision: accept the one-line fusion and its repeatable marginal combined gain +of about `0.17%`. + +### Ablation: isolated H1024 train-projection cuBLASLt plan - rejected + +Adding the previously qualified fixed H1024 train-projection plan cost two +lines and preserved exact checkpoints, but its five-pair production screen +measured Affine 578 legacy fallback `0.997998x`, Affine 1024x4 `0.996658x`, and +suite `0.997328x`, with a `0.987755x` worst pair. Artifact: +`/tmp/puffer-h1024-train-lt-screen-OfsEmK2I/puffer-throughput-fr6o_0l5`. +Production overlap reverses the graph-local kernel gain. Decision: reject; +the source was reverted and the plan was not committed. + +### Ablation: Muon lane scheduling tuner - rejected + +All tuner variants preserved exact outputs. Endpoint-high scheduling measured +median DAG gains of `1.004146x` at H512 and `1.001934x` at H1024, but its +projected end-to-end value was well below `0.1%`. Recurrent-high scheduling +regressed H512 by `8.37%` and H1024 by `1.16%`. Restricting execution to pools +of two, three, or four lanes regressed H512 by `21.77%`, `19.89%`, and +`11.23%`, and H1024 by `5.57%`, `7.42%`, and `3.76%`, respectively. Artifact: +`/tmp/puffer_muon_lane_schedule_results.txt`. Decision: no production change. + +### Ablation: H512/H1024 rollout MinGRU gate mapping - rejected + +The 20-line specialization replaced per-thread divide/modulo and the tail check +with exact 2D agent/chunk mapping for H512 and H1024. All checkpoints matched, +but the five-pair screen measured Affine 578 `1.002590x`, Affine 1024x4 +`1.000172x`, and suite `1.001380x`, with a `0.987519x` worst pair. Artifact: +`/tmp/puffer-mingru-gate-map-screen-19VbjWvm/puffer-throughput-5j2zmj3j`. +Decision: reject because the gain was below the `0.2%` threshold and variance +was high; restore the committed gate. + +### Ablation: isolated H1024 Muon cuBLASLt re-addition - rejected + +Re-adding the H1024 Muon plans cost four lines and preserved exact checkpoints. +The first five-pair long batch measured `1.003779x`; artifact: +`/tmp/puffer-muon-lt-h1024-long-QM4p8zxG/puffer-throughput-hps4mgry`. An +independent exact five-pair batch measured `1.002997x`; artifact: +`/tmp/puffer-muon-lt-h1024-confirm-fM90dWXC/puffer-throughput-dikx4mst`. + +Across all ten pairs, the geomean was `1.003388x` and seven pairs were +positive, but the paired-log Student-t 95% interval was approximately +`[0.99914x, 1.00765x]`. Decision: reject because the interval includes parity; +the four lines were reverted and no commit was made. + +### Ablation: ordinary MinGRU forward compile-time T64 - rejected + +The net five-line specialization preserved exact checkpoints in every pair, +but measured Affine 578 `0.999050x`, Affine 1024x4 `0.994973x`, and suite +`0.997010x`, with a `0.991014x` worst pair. Artifact: +`/tmp/puffer-mingru-fwd-t64-screen-8w6xqioM/puffer-throughput-foljj5nm`. +Decision: reject; the specialization was reverted and not committed. + +### Ablation: exhaustive H1024 Muon cuBLASLt search - rejected + +The search covered 20 algorithm IDs and used the widened 200,000-evaluation +cap per shape, still reporting truncation. The best result used square algo 21, +tile 18, stage 12 and X algo 67, tile 29, stage 35, custom 30. Every +coefficient, intermediate node, final output, and graph replay check matched +exactly. + +The X operation was about `1.103x` faster, but the complete four-lane mixed DAG +reached only `1.023230x` median and `1.024480x` mean, below the `1.05x` +pre-integration threshold. Artifact: +`/tmp/puffer_h1024_muon_lt_exhaustive_results.txt`. Decision: no production +integration. + +### Ablation: exhaustive H512 Muon cuBLASLt search - rejected + +The accepted-relative search covered all 20 algorithm IDs and all 485,194 +tuples per shape without truncation, yielding 13,460 valid zero-workspace +configurations. The best result retained square algo 21, tile 15, stage 12 and +used X algo 67, tile 318, stage 35, custom 133. All coefficient, intermediate +node, final output, and graph replay checks matched exactly. + +Despite a raw X speedup of about `1.464x`, the complete three-lane DAG improved +from `0.220416 ms` accepted to `0.217664 ms` candidate: only `1.012643x` +median and `1.011507x` by means. This is below the `1.05x` pre-integration +threshold. Artifact: `/tmp/puffer_h512_muon_lt_exhaustive_results.txt`. +Decision: no source integration. diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 0000000000..4a30bbf3d5 --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,145 @@ +# Native trainer throughput benchmarks + +This harness compares separately built baseline and candidate binaries without +writing to repository logs, checkpoints, resources, or build outputs. + +The default `--checkpoint-policy exact` makes this harness the strict C0 +correctness gate. Every scored baseline/candidate pair must produce a +byte-identical final checkpoint. It stops at the first mismatch and retains +both checkpoints plus a mismatch report. + +Byte identity is appropriate only for changes claimed to preserve operation +order and arithmetic, such as removing unused event nodes. Optimizations that +change GEMM algorithms, reduction order, or floating-point association require +a separate full-budget, paired multi-seed quality gate across every golden +environment; this short harness cannot establish their learning equivalence. +Some production async configurations are not repeatable even when the same +baseline binary is run twice. Use `--checkpoint-policy record` for those +throughput runs so every mismatch is retained without being misattributed to +the candidate. This never turns a mismatch into correctness evidence. + +## Cases + +- `affine578-h`: the resolved run-578 recipe using the CPU environment. +- `affine578-cu`: run 578 with the CUDA backend's required one-buffer setting. +- `affine1024x4-h`: run 578 with hidden size 1024 and four effective layers. +- `affine1024x4-cu`: the same stress profile on the CUDA environment backend. +- `breakout-h` and `breakout-cu`: small-network controls. +- `g2048-h`: large learner-dominated control. +- `maze-h`: long-horizon recurrent control. +- `boxoban-h`: large-observation control using an explicitly staged map file. + +Run 578's raw sweep layer value was `3.93814397`. The native trainer resolves +that value through an integer field, so the fixture passes the effective value +`3`. The `1024x4` profile changes only hidden size and effective layer count. + +## Build contract + +The harness does not build or delete anything. Build baseline and candidate +binaries into separate directories with these names: + +```bash +./build.sh affine_lock /tmp/puffer-baseline/affine_lock_h +./build.sh affine_lock /tmp/puffer-baseline/affine_lock_cu --cu +./build.sh breakout /tmp/puffer-baseline/breakout_h +./build.sh breakout /tmp/puffer-baseline/breakout_cu --cu +./build.sh g2048 /tmp/puffer-baseline/g2048_h +./build.sh maze /tmp/puffer-baseline/maze_h +./build.sh boxoban /tmp/puffer-baseline/boxoban_h +``` + +Repeat from the candidate source tree with `/tmp/puffer-candidate` outputs. +Use identical compiler, CUDA, architecture, and environment settings. + +## Run + +Screen the four Affine cases with five paired trials: + +```bash +python3 benchmarks/native_throughput.py \ + --baseline-dir /tmp/puffer-baseline \ + --candidate-dir /tmp/puffer-candidate \ + --cases affine \ + --pairs 5 +``` + +Run a deterministic exact-checkpoint canary by disabling async collection and +using one vector thread: + +```bash +python3 benchmarks/native_throughput.py \ + --baseline-dir /tmp/puffer-baseline \ + --candidate-dir /tmp/puffer-candidate \ + --cases goldens \ + --pairs 1 \ + --override base.async=0 \ + --override vec.num_threads=1 \ + --checkpoint-policy exact \ + --boxoban-map /tmp/puffer-maps/boxoban_maps_medium.bin +``` + +Keep production async settings for the paired performance campaign and record, +rather than abort on, intrinsic checkpoint variation: + +```bash +python3 benchmarks/native_throughput.py \ + --baseline-dir /tmp/puffer-baseline \ + --candidate-dir /tmp/puffer-candidate \ + --cases all \ + --pairs 5 \ + --checkpoint-policy record \ + --boxoban-map /tmp/puffer-maps/boxoban_maps_medium.bin +``` + +Every `--override KEY=VALUE` is stored in campaign metadata and each exact +command. Harness-owned timestep, run ID, log, checkpoint, profiling, and +downsample settings are appended afterward and cannot be displaced by an +override. + +Inspect a generated plan without executing a binary: + +```bash +python3 benchmarks/native_throughput.py \ + --baseline-dir /tmp/puffer-baseline \ + --candidate-dir /tmp/puffer-candidate \ + --cases all \ + --pairs 5 \ + --boxoban-map /tmp/puffer-maps/boxoban_maps_medium.bin \ + --dry-run +``` + +Every invocation creates a new mode-0700 +`/tmp/puffer-throughput-XXXXXXXX/` campaign and prints its path immediately. +It never removes or overwrites a campaign. + +## Measurement + +- Baseline/candidate order alternates `A/B`, then `B/A`. +- One short unscored preconditioning run is performed for each case and binary. +- Scored Affine runs use 134,217,728 steps, exactly 512 horizons at 4096 agents + and horizon 64. +- Steady-state SPS is a least-squares slope of agent steps against native + uptime after removing the first 10% of steps. +- At least two distinct post-warmup metric samples are required. Exact padded + terminal duplicates are ignored; real non-monotonic samples still fail. +- Process-wall SPS is recorded separately and includes startup/finalization. +- Results include paired ratios, geometric means, median/MAD, deterministic + bootstrap 95% confidence intervals, binary hashes, commands, and machine/GPU + metadata. + +The initial performance gate rejects a reproducible regression worse than 3%. +Affine improvement confidence intervals must exclude 1.0. Exact canary +checkpoint hashes must match. Async production mismatches are recorded, and +performance results still require a separate full-budget, paired multi-seed +quality gate. + +## Safety + +- The harness uses explicit binary paths and never invokes `build.sh`. +- Repository `logs/`, checkpoints, resources, and configurations are not + modified. +- Boxoban requires `--boxoban-map` pointing to a pre-staged immutable map + binary; the harness never downloads or generates maps. +- Existing GPU compute processes cause an abort unless + `--allow-foreign-gpu-processes` is supplied. +- No artifact cleanup is automatic. diff --git a/benchmarks/native_cases.toml b/benchmarks/native_cases.toml new file mode 100644 index 0000000000..d52feafbf9 --- /dev/null +++ b/benchmarks/native_cases.toml @@ -0,0 +1,130 @@ +version = 1 + +[profiles.affine578] +args = [ + "base.seed=73", + "base.async=1", + "base.cudagraphs=1", + "base.reset_every_horizon=0", + "env.seed=42", + "env.start_depth=2", + "env.max_depth=16", + "env.num_agents=1", + "env.num_bots=0", + "env.step_grace=0", + "env.perf_weighting=1", + "train.gpus=1", + "train.horizon=64", + "train.minibatch_size=8192", + "train.replay_ratio=1.87008977", + "train.learning_rate=0.00326743093", + "train.anneal_lr=1", + "train.min_lr_ratio=0", + "train.gamma=0.999899983", + "train.gae_lambda=0.921891332", + "train.clip_coef=0.366497159", + "train.vf_coef=0.100000001", + "train.vf_clip_coef=0.00100000005", + "train.ent_coef=0.0405269228", + "train.anneal_ent_coef=0", + "train.min_ent_coef_ratio=0.1", + "train.momentum=0.924182773", + "train.max_grad_norm=0.881812513", + "train.vtrace=0", + "train.verb_eps=0", + "policy.hidden_size=512", + "policy.num_layers=3", + "selfplay.enabled=0", + "vec.total_agents=4096", + "vec.num_threads=16", +] + +[groups] +affine = ["affine578-h", "affine578-cu", "affine1024x4-h", "affine1024x4-cu"] +goldens = ["breakout-h", "breakout-cu", "g2048-h", "maze-h", "boxoban-h"] +all = ["affine", "goldens"] + +[cases."affine578-h"] +binary = "affine_lock_h" +env = "affine_lock" +backend = "h" +profile = "affine578" +timesteps = 134217728 +precondition_timesteps = 8388608 +description = "Run 578 resolved recipe with the CPU environment backend" +overrides = { "vec.num_buffers" = "2" } + +[cases."affine578-cu"] +binary = "affine_lock_cu" +env = "affine_lock" +backend = "cu" +profile = "affine578" +timesteps = 134217728 +precondition_timesteps = 8388608 +description = "Run 578 resolved recipe with the CUDA-required one-buffer adjustment" +overrides = { "vec.num_buffers" = "1" } + +[cases."affine1024x4-h"] +binary = "affine_lock_h" +env = "affine_lock" +backend = "h" +profile = "affine578" +timesteps = 134217728 +precondition_timesteps = 8388608 +description = "Run 578 with only hidden size doubled and one effective layer added" +overrides = { "vec.num_buffers" = "2", "policy.hidden_size" = "1024", "policy.num_layers" = "4" } + +[cases."affine1024x4-cu"] +binary = "affine_lock_cu" +env = "affine_lock" +backend = "cu" +profile = "affine578" +timesteps = 134217728 +precondition_timesteps = 8388608 +description = "1024x4 Affine stress profile with the CUDA-required one-buffer adjustment" +overrides = { "vec.num_buffers" = "1", "policy.hidden_size" = "1024", "policy.num_layers" = "4" } + +[cases."breakout-h"] +binary = "breakout_h" +env = "breakout" +backend = "h" +timesteps = 134217728 +precondition_timesteps = 8388608 +description = "Canonical small-network CPU-environment control" +args = ["base.seed=73", "base.async=1", "base.cudagraphs=1"] + +[cases."breakout-cu"] +binary = "breakout_cu" +env = "breakout" +backend = "cu" +timesteps = 134217728 +precondition_timesteps = 8388608 +description = "Canonical small-network CUDA-environment control" +args = ["base.seed=73", "base.async=1", "base.cudagraphs=1", "vec.num_buffers=1"] + +[cases."g2048-h"] +binary = "g2048_h" +env = "g2048" +backend = "h" +timesteps = 33554432 +precondition_timesteps = 8388608 +description = "Canonical large-network learner-dominated control" +args = ["base.seed=73", "base.async=1", "base.cudagraphs=1"] + +[cases."maze-h"] +binary = "maze_h" +env = "maze" +backend = "h" +timesteps = 33554432 +precondition_timesteps = 8388608 +description = "Canonical long-horizon MinGRU control" +args = ["base.seed=73", "base.async=1", "base.cudagraphs=1"] + +[cases."boxoban-h"] +binary = "boxoban_h" +env = "boxoban" +backend = "h" +timesteps = 33554432 +precondition_timesteps = 8388608 +description = "Canonical large-observation puzzle control with an explicitly staged map" +args = ["base.seed=73", "base.async=1", "base.cudagraphs=1"] diff --git a/benchmarks/native_throughput.py b/benchmarks/native_throughput.py new file mode 100644 index 0000000000..1f20f9f8dc --- /dev/null +++ b/benchmarks/native_throughput.py @@ -0,0 +1,740 @@ +#!/usr/bin/env python3 +"""Paired native-trainer throughput benchmarks with isolated artifacts.""" + +from __future__ import annotations + +import argparse +import configparser +import csv +import hashlib +import json +import math +import os +import platform +import random +import re +import shlex +import statistics +import subprocess +import tempfile +import time +import tomllib +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +REPO_ROOT = Path(__file__).resolve().parents[1] +DEFAULT_CASE_FILE = Path(__file__).with_name("native_cases.toml") +FLOAT_RE = re.compile( + r"[-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?" +) + + +@dataclass(frozen=True) +class Case: + name: str + binary: str + env: str + backend: str + timesteps: int + precondition_timesteps: int + args: tuple[str, ...] + description: str + + +def parse_cli() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--baseline-dir", required=True, type=Path) + parser.add_argument("--candidate-dir", required=True, type=Path) + parser.add_argument("--runtime-root", type=Path, default=REPO_ROOT) + parser.add_argument("--case-file", type=Path, default=DEFAULT_CASE_FILE) + parser.add_argument( + "--cases", + default="affine", + help="Comma-separated case or group names from native_cases.toml", + ) + parser.add_argument("--pairs", type=int, default=5) + parser.add_argument( + "--artifact-root", + type=Path, + default=Path("/tmp"), + help="A new puffer-throughput-* campaign is created below this path", + ) + parser.add_argument( + "--timesteps", + type=int, + help="Override the scored timestep budget for every selected case", + ) + parser.add_argument( + "--override", + action="append", + default=[], + metavar="KEY=VALUE", + help="Apply and record a native configuration override to every run", + ) + parser.add_argument( + "--checkpoint-policy", + choices=("exact", "record"), + default="exact", + help="Abort on checkpoint mismatch or only record it (default: exact)", + ) + parser.add_argument("--boxoban-map", type=Path) + parser.add_argument("--timeout", type=float, default=0.0) + parser.add_argument( + "--no-precondition", + action="store_true", + help="Skip one short unscored run per case and binary", + ) + parser.add_argument( + "--allow-foreign-gpu-processes", + action="store_true", + help="Record rather than reject pre-existing GPU compute processes", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Write the campaign plan but do not execute binaries", + ) + args = parser.parse_args() + if args.pairs < 1: + parser.error("--pairs must be positive") + if args.timesteps is not None and args.timesteps < 1: + parser.error("--timesteps must be positive") + if args.timeout < 0: + parser.error("--timeout cannot be negative") + for value in args.override: + try: + key, _ = split_override(value) + except ValueError as error: + parser.error(str(error)) + if not key: + parser.error(f"Invalid empty override key: {value!r}") + return args + + +def split_override(value: str) -> tuple[str, str]: + if "=" not in value: + raise ValueError(f"Invalid case override without '=': {value!r}") + return value.split("=", 1) + + +def merge_overrides(base: list[str], overrides: dict[str, Any]) -> tuple[str, ...]: + merged: dict[str, str] = {} + order: list[str] = [] + for value in base: + key, setting = split_override(value) + if key not in merged: + order.append(key) + merged[key] = setting + for key, setting in overrides.items(): + if key not in merged: + order.append(key) + merged[key] = str(setting) + return tuple(f"{key}={merged[key]}" for key in order) + + +def load_case_file(path: Path) -> tuple[dict[str, Case], dict[str, list[str]]]: + with path.open("rb") as handle: + data = tomllib.load(handle) + profiles = data.get("profiles", {}) + cases: dict[str, Case] = {} + for name, raw in data.get("cases", {}).items(): + profile_name = raw.get("profile") + profile_args: list[str] = [] + if profile_name is not None: + if profile_name not in profiles: + raise ValueError(f"Case {name!r} has unknown profile {profile_name!r}") + profile_args = list(profiles[profile_name].get("args", [])) + profile_args.extend(raw.get("args", [])) + cases[name] = Case( + name=name, + binary=str(raw["binary"]), + env=str(raw["env"]), + backend=str(raw["backend"]), + timesteps=int(raw["timesteps"]), + precondition_timesteps=int(raw.get("precondition_timesteps", 8_388_608)), + args=merge_overrides(profile_args, raw.get("overrides", {})), + description=str(raw.get("description", "")), + ) + groups = {name: list(values) for name, values in data.get("groups", {}).items()} + return cases, groups + + +def select_cases( + selection: str, cases: dict[str, Case], groups: dict[str, list[str]] +) -> list[Case]: + selected: list[str] = [] + expanding: set[str] = set() + + def add(name: str) -> None: + if name in cases: + if name not in selected: + selected.append(name) + return + if name not in groups: + raise ValueError(f"Unknown case or group: {name!r}") + if name in expanding: + raise ValueError(f"Recursive case group: {name!r}") + expanding.add(name) + for child in groups[name]: + add(child) + expanding.remove(name) + + for item in selection.split(","): + item = item.strip() + if item: + add(item) + if not selected: + raise ValueError("No benchmark cases selected") + return [cases[name] for name in selected] + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def capture_command(command: list[str]) -> dict[str, Any]: + try: + result = subprocess.run( + command, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + check=False, + timeout=10, + ) + return {"command": command, "returncode": result.returncode, "output": result.stdout} + except (FileNotFoundError, subprocess.TimeoutExpired) as error: + return {"command": command, "error": str(error)} + + +def foreign_gpu_processes() -> list[str]: + result = capture_command( + [ + "nvidia-smi", + "--query-compute-apps=pid,process_name,used_gpu_memory", + "--format=csv,noheader,nounits", + ] + ) + if result.get("returncode") != 0: + return [] + return [line.strip() for line in result.get("output", "").splitlines() if line.strip()] + + +def machine_metadata() -> dict[str, Any]: + return { + "hostname": platform.node(), + "platform": platform.platform(), + "python": platform.python_version(), + "nvidia_smi": capture_command(["nvidia-smi"]), + "gpu_query": capture_command( + [ + "nvidia-smi", + "--query-gpu=name,uuid,driver_version,pstate,temperature.gpu,clocks.sm,power.draw", + "--format=csv,noheader", + ] + ), + "nvcc": capture_command(["nvcc", "--version"]), + } + + +def numeric_series(value: str) -> list[float]: + return [float(match.group(0)) for match in FLOAT_RE.finditer(value)] + + +def find_run_ini(log_dir: Path, run_id: str) -> Path: + candidates = list(log_dir.rglob("*.ini")) + matching = [path for path in candidates if run_id in path.name] + if len(matching) == 1: + return matching[0] + if len(candidates) == 1: + return candidates[0] + if not candidates: + raise RuntimeError(f"No INI log produced under {log_dir}") + raise RuntimeError( + f"Ambiguous INI logs under {log_dir}: " + ", ".join(str(path) for path in candidates) + ) + + +def find_final_checkpoint(checkpoint_dir: Path) -> Path: + candidates = list(checkpoint_dir.rglob("*.bin")) + if not candidates: + raise RuntimeError(f"No checkpoint produced under {checkpoint_dir}") + if len(candidates) == 1: + return candidates[0] + numeric = [path for path in candidates if path.stem.isdigit()] + if numeric: + return max(numeric, key=lambda path: int(path.stem)) + raise RuntimeError( + f"Ambiguous checkpoints under {checkpoint_dir}: " + + ", ".join(str(path) for path in candidates) + ) + + +def parse_metric_series(path: Path) -> dict[str, list[float]]: + wanted = {"agent_steps", "uptime", "env/score", "perf", "SPS"} + parsed: dict[str, list[float]] = {} + parser = configparser.ConfigParser(interpolation=None, strict=False) + parser.optionxform = str + try: + with path.open("r", encoding="utf-8") as handle: + parser.read_file(handle) + for section in parser.sections(): + for key, value in parser.items(section): + normalized = key.strip() + if normalized in wanted: + parsed[normalized] = numeric_series(value) + except configparser.Error: + parsed = {} + + if "agent_steps" not in parsed or "uptime" not in parsed: + with path.open("r", encoding="utf-8") as handle: + for line in handle: + if "=" not in line: + continue + key, value = line.split("=", 1) + key = key.strip() + if key in wanted: + parsed[key] = numeric_series(value) + return parsed + + +def throughput_metrics(metrics: dict[str, list[float]]) -> dict[str, float | int]: + steps = metrics.get("agent_steps", []) + uptime = metrics.get("uptime", []) + count = min(len(steps), len(uptime)) + if count < 2: + raise RuntimeError(f"Need at least two step/uptime samples, found {count}") + steps = steps[-count:] + uptime = uptime[-count:] + + samples = [] + for elapsed, agent_steps in zip(uptime, steps): + if samples and (elapsed, agent_steps) == samples[-1]: + continue + samples.append((elapsed, agent_steps)) + + uptime = [elapsed for elapsed, _ in samples] + steps = [agent_steps for _, agent_steps in samples] + count = len(samples) + if count < 2: + raise RuntimeError("Fewer than two unique uptime samples were recorded") + + if any(b <= a for a, b in zip(uptime, uptime[1:])): + raise RuntimeError("Uptime samples are not strictly increasing") + if any(b < a for a, b in zip(steps, steps[1:])): + raise RuntimeError("Agent-step samples are not monotonic") + + cutoff = steps[0] + 0.10 * (steps[-1] - steps[0]) + retained = [(x, y) for x, y in zip(uptime, steps) if y >= cutoff] + if len(retained) < 2: + raise RuntimeError( + f"Need two samples after 10% warmup removal, found {len(retained)}" + ) + xs = [item[0] for item in retained] + ys = [item[1] for item in retained] + x_mean = statistics.fmean(xs) + y_mean = statistics.fmean(ys) + denominator = sum((x - x_mean) ** 2 for x in xs) + if denominator <= 0: + raise RuntimeError("Cannot fit throughput slope from identical uptime samples") + slope = sum((x - x_mean) * (y - y_mean) for x, y in retained) / denominator + full_delta = (steps[-1] - steps[0]) / (uptime[-1] - uptime[0]) + return { + "steady_sps": slope, + "native_full_sps": full_delta, + "metric_samples": count, + "retained_samples": len(retained), + } + + +def command_for_run( + binary: Path, + case: Case, + timesteps: int, + run_id: str, + log_dir: Path, + checkpoint_dir: Path, + boxoban_map: Path | None, + run_overrides: tuple[str, ...], +) -> list[str]: + overrides = list(case.args) + for value in run_overrides: + key, setting = split_override(value) + overrides = list(merge_overrides(overrides, {key: setting})) + overrides.extend( + [ + f"train.total_timesteps={timesteps}", + f"base.run_id={run_id}", + f"base.log_dir={log_dir}", + f"base.checkpoint_dir={checkpoint_dir}", + "base.profile=0", + "base.eval_episodes=1", + "base.checkpoint_interval=2147483647", + "sweep.downsample=64", + ] + ) + if case.env == "boxoban": + if boxoban_map is None: + raise ValueError("Boxoban requires --boxoban-map with a staged map binary") + return [str(binary), "train", *(f"--{value}" for value in overrides)] + + +def append_jsonl(path: Path, record: dict[str, Any]) -> None: + with path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(record, sort_keys=True) + "\n") + + +def execute_run( + *, + campaign: Path, + runtime_root: Path, + label: str, + binary: Path, + case: Case, + timesteps: int, + pair: int, + sequence: int, + phase: str, + timeout: float, + boxoban_map: Path | None, + run_overrides: tuple[str, ...], + allow_foreign: bool, + dry_run: bool, +) -> dict[str, Any]: + run_id = f"{case.name}_{phase}_p{pair:02d}_{label}_{sequence:03d}" + run_dir = campaign / "runs" / run_id + log_dir = run_dir / "logs" + checkpoint_dir = run_dir / "checkpoints" + log_dir.mkdir(parents=True) + checkpoint_dir.mkdir() + command = command_for_run( + binary, + case, + timesteps, + run_id, + log_dir, + checkpoint_dir, + boxoban_map, + run_overrides, + ) + foreign = foreign_gpu_processes() + if foreign and not allow_foreign: + raise RuntimeError( + "Foreign GPU compute processes detected before run: " + "; ".join(foreign) + ) + record: dict[str, Any] = { + "case": case.name, + "description": case.description, + "backend": case.backend, + "label": label, + "phase": phase, + "pair": pair, + "sequence": sequence, + "timesteps": timesteps, + "run_id": run_id, + "run_dir": str(run_dir), + "binary": str(binary), + "binary_sha256": sha256(binary), + "command": command, + "command_shell": shlex.join(command), + "foreign_gpu_processes_before": foreign, + "started_unix": time.time(), + } + run_environment = os.environ.copy() + if case.env == "boxoban": + assert boxoban_map is not None + run_environment["BOXOBAN_MAP_BIN"] = str(boxoban_map) + record["boxoban_map"] = str(boxoban_map) + record["boxoban_map_sha256"] = sha256(boxoban_map) + (run_dir / "command.json").write_text( + json.dumps(record, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + if dry_run: + record["dry_run"] = True + (run_dir / "result.json").write_text( + json.dumps(record, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + return record + + started = time.perf_counter() + result = subprocess.run( + command, + cwd=runtime_root, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=False, + timeout=None if timeout == 0 else timeout, + env=run_environment, + ) + wall_seconds = time.perf_counter() - started + (run_dir / "stdout.txt").write_text(result.stdout, encoding="utf-8") + (run_dir / "stderr.txt").write_text(result.stderr, encoding="utf-8") + record.update( + { + "returncode": result.returncode, + "wall_seconds": wall_seconds, + "process_wall_sps": timesteps / wall_seconds, + "finished_unix": time.time(), + "foreign_gpu_processes_after": foreign_gpu_processes(), + } + ) + if result.returncode != 0: + (run_dir / "result.json").write_text( + json.dumps(record, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + raise RuntimeError( + f"{case.name} {label} exited {result.returncode}; see {run_dir}" + ) + + ini_path = find_run_ini(log_dir, run_id) + metrics = parse_metric_series(ini_path) + record.update(throughput_metrics(metrics)) + record["ini_path"] = str(ini_path) + checkpoint_path = find_final_checkpoint(checkpoint_dir) + record["checkpoint_path"] = str(checkpoint_path) + record["checkpoint_sha256"] = sha256(checkpoint_path) + record["checkpoint_bytes"] = checkpoint_path.stat().st_size + (run_dir / "result.json").write_text( + json.dumps(record, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + return record + + +def percentile(sorted_values: list[float], fraction: float) -> float: + position = fraction * (len(sorted_values) - 1) + lower = math.floor(position) + upper = math.ceil(position) + if lower == upper: + return sorted_values[lower] + weight = position - lower + return sorted_values[lower] * (1.0 - weight) + sorted_values[upper] * weight + + +def summarize_ratios(ratios: list[float]) -> dict[str, float | int]: + logs = [math.log(value) for value in ratios] + rng = random.Random(73) + bootstrap: list[float] = [] + for _ in range(10_000): + sample = [logs[rng.randrange(len(logs))] for _ in logs] + bootstrap.append(math.exp(statistics.fmean(sample))) + bootstrap.sort() + median = statistics.median(ratios) + mad = statistics.median(abs(value - median) for value in ratios) + return { + "pairs": len(ratios), + "geomean_speedup": math.exp(statistics.fmean(logs)), + "median_speedup": median, + "minimum_speedup": min(ratios), + "mad": mad, + "bootstrap_95_low": percentile(bootstrap, 0.025), + "bootstrap_95_high": percentile(bootstrap, 0.975), + } + + +def write_summary(campaign: Path, pairs: list[dict[str, Any]]) -> dict[str, Any]: + by_case: dict[str, list[float]] = {} + for pair in pairs: + by_case.setdefault(pair["case"], []).append(pair["steady_speedup"]) + case_summary = { + name: summarize_ratios(ratios) for name, ratios in sorted(by_case.items()) + } + affine_summaries = [ + values for name, values in case_summary.items() if name.startswith("affine") + ] + all_logs = [math.log(pair["steady_speedup"]) for pair in pairs] + summary = { + "cases": case_summary, + "suite_geomean_speedup": math.exp(statistics.fmean(all_logs)), + "worst_pair_speedup": min(pair["steady_speedup"] for pair in pairs), + "acceptance": { + "no_pair_below_0.97": all(pair["steady_speedup"] >= 0.97 for pair in pairs), + "all_checkpoints_match": all(pair["checkpoint_match"] for pair in pairs), + "affine_ci_excludes_1": bool(affine_summaries) + and all(values["bootstrap_95_low"] > 1.0 for values in affine_summaries), + }, + } + (campaign / "summary.json").write_text( + json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + with (campaign / "pairs.csv").open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter( + handle, + fieldnames=[ + "case", + "pair", + "baseline_steady_sps", + "candidate_steady_sps", + "steady_speedup", + "baseline_wall_sps", + "candidate_wall_sps", + "wall_speedup", + "checkpoint_match", + "checkpoint_sha256", + ], + ) + writer.writeheader() + writer.writerows(pairs) + return summary + + +def main() -> int: + args = parse_cli() + cases_by_name, groups = load_case_file(args.case_file.resolve()) + cases = select_cases(args.cases, cases_by_name, groups) + baseline_dir = args.baseline_dir.resolve() + candidate_dir = args.candidate_dir.resolve() + runtime_root = args.runtime_root.resolve() + boxoban_map = args.boxoban_map.resolve() if args.boxoban_map else None + binaries: dict[str, dict[str, Path]] = {"baseline": {}, "candidate": {}} + for case in cases: + for label, directory in ( + ("baseline", baseline_dir), + ("candidate", candidate_dir), + ): + binary = directory / case.binary + if not binary.is_file(): + raise FileNotFoundError(f"Missing {label} binary for {case.name}: {binary}") + if not os.access(binary, os.X_OK): + raise PermissionError(f"Binary is not executable: {binary}") + binaries[label][case.name] = binary + if boxoban_map is not None and not boxoban_map.is_file(): + raise FileNotFoundError(f"Boxoban map does not exist: {boxoban_map}") + + args.artifact_root.mkdir(parents=True, exist_ok=True) + campaign = Path( + tempfile.mkdtemp(prefix="puffer-throughput-", dir=args.artifact_root.resolve()) + ) + campaign.chmod(0o700) + metadata = { + "created_unix": time.time(), + "runtime_root": str(runtime_root), + "case_file": str(args.case_file.resolve()), + "case_file_sha256": sha256(args.case_file.resolve()), + "selected_cases": [case.name for case in cases], + "pairs": args.pairs, + "timesteps_override": args.timesteps, + "run_overrides": args.override, + "checkpoint_policy": args.checkpoint_policy, + "precondition": not args.no_precondition, + "dry_run": args.dry_run, + "machine": machine_metadata(), + } + (campaign / "campaign.json").write_text( + json.dumps(metadata, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + print(f"Artifacts: {campaign}", flush=True) + + sequence = 0 + raw_path = campaign / "runs.jsonl" + if not args.no_precondition: + for case in cases: + for label in ("baseline", "candidate"): + sequence += 1 + record = execute_run( + campaign=campaign, + runtime_root=runtime_root, + label=label, + binary=binaries[label][case.name], + case=case, + timesteps=case.precondition_timesteps, + pair=-1, + sequence=sequence, + phase="precondition", + timeout=args.timeout, + boxoban_map=boxoban_map, + run_overrides=tuple(args.override), + allow_foreign=args.allow_foreign_gpu_processes, + dry_run=args.dry_run, + ) + append_jsonl(raw_path, record) + + paired_results: list[dict[str, Any]] = [] + for case in cases: + scored_timesteps = args.timesteps or case.timesteps + for pair_index in range(args.pairs): + order = ( + ("baseline", "candidate") + if pair_index % 2 == 0 + else ("candidate", "baseline") + ) + results: dict[str, dict[str, Any]] = {} + for label in order: + sequence += 1 + record = execute_run( + campaign=campaign, + runtime_root=runtime_root, + label=label, + binary=binaries[label][case.name], + case=case, + timesteps=scored_timesteps, + pair=pair_index, + sequence=sequence, + phase="scored", + timeout=args.timeout, + boxoban_map=boxoban_map, + run_overrides=tuple(args.override), + allow_foreign=args.allow_foreign_gpu_processes, + dry_run=args.dry_run, + ) + append_jsonl(raw_path, record) + results[label] = record + if args.dry_run: + continue + baseline = results["baseline"] + candidate = results["candidate"] + pair_result = { + "case": case.name, + "pair": pair_index, + "baseline_steady_sps": baseline["steady_sps"], + "candidate_steady_sps": candidate["steady_sps"], + "steady_speedup": candidate["steady_sps"] / baseline["steady_sps"], + "baseline_wall_sps": baseline["process_wall_sps"], + "candidate_wall_sps": candidate["process_wall_sps"], + "wall_speedup": candidate["process_wall_sps"] + / baseline["process_wall_sps"], + "checkpoint_match": baseline["checkpoint_sha256"] + == candidate["checkpoint_sha256"], + "checkpoint_sha256": baseline["checkpoint_sha256"], + } + paired_results.append(pair_result) + if not pair_result["checkpoint_match"]: + failure = { + "case": case.name, + "pair": pair_index, + "baseline": baseline["checkpoint_sha256"], + "candidate": candidate["checkpoint_sha256"], + "baseline_path": baseline["checkpoint_path"], + "candidate_path": candidate["checkpoint_path"], + } + append_jsonl(campaign / "checkpoint_mismatches.jsonl", failure) + if args.checkpoint_policy == "record": + continue + failure_path = campaign / "CHECKPOINT_MISMATCH.json" + failure_path.write_text( + json.dumps(failure, indent=2, sort_keys=True) + + "\n", + encoding="utf-8", + ) + raise RuntimeError( + f"Checkpoint mismatch for {case.name} pair {pair_index}; " + f"see {failure_path}" + ) + + if args.dry_run: + print(f"Dry-run plan retained at {campaign}") + return 0 + summary = write_summary(campaign, paired_results) + print(json.dumps(summary, indent=2, sort_keys=True)) + print(f"Artifacts retained at {campaign}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/build.sh b/build.sh index 115d8f4b73..b556c8ae37 100755 --- a/build.sh +++ b/build.sh @@ -385,7 +385,7 @@ if [ "$MODE" = "native" ]; then "${LINK_ARCHIVES[@]}" \ -L$CUDA_HOME/lib64 $NCCL_LFLAG \ "${EXTRA_LDFLAGS[@]}" \ - -lcudart -lnccl -lnvidia-ml -lcublas -lcusolver -lcurand \ + -lcudart -lnccl -lnvidia-ml -lcublas -lcublasLt -lcusolver -lcurand \ -lm -lpthread $OMP_LIB "${STANDALONE_LDFLAGS[@]}" \ -o "$TRAIN_BIN" echo "Built: ./$TRAIN_BIN" @@ -408,7 +408,7 @@ elif [ "$MODE" = "profile" ]; then tests/profile_kernels.cu \ "$RAYLIB_A" \ -L$CUDA_HOME/lib64 \ - -lnccl -lnvidia-ml -lcublas -lcusolver -lcurand \ + -lnccl -lnvidia-ml -lcublas -lcublasLt -lcusolver -lcurand \ -lGL -lm -lpthread $OMP_LIB \ -o "$PROFILE_BIN" echo "Built: ./$PROFILE_BIN" diff --git a/ocean/affine_lock/affine_lock.c b/ocean/affine_lock/affine_lock.c index 8d2ae8cb81..13e64c2aad 100644 --- a/ocean/affine_lock/affine_lock.c +++ b/ocean/affine_lock/affine_lock.c @@ -24,6 +24,9 @@ int main(void) { puf_render(&env); while (!WindowShouldClose()) { + if (IsKeyPressed(KEY_H)) { + show_hint(&env); + } puf_step(&env); puf_render(&env); } diff --git a/ocean/affine_lock/affine_lock.cu b/ocean/affine_lock/affine_lock.cu new file mode 100644 index 0000000000..adb550eb5f --- /dev/null +++ b/ocean/affine_lock/affine_lock.cu @@ -0,0 +1,808 @@ +// Vibe coded by OpenAI Codex. +// GPU Affine Lock environment. This is intentionally standalone from +// affine_lock.h: --gpu builds include this file instead of the CPU source. +#ifndef PUFFER_AFFINE_LOCK_GPU_CU +#define PUFFER_AFFINE_LOCK_GPU_CU + +#define PUF_BACKEND PUF_GPU + +#include +#include + +#include +#include +#include +#include +#include + +// Environment observations are fixed bf16. All bit observations are +/-1 +// (exact in bf16); only the timer is rounded. Keeping them bf16 halves the +// rollout bandwidth compared with the CPU float representation. +typedef __nv_bfloat16 obs_t; +#include "pufferenv.h" +#include "affine_lock_visible_targets.h" + +#define BITS 16 +#define TIMER_INDEX (2 * BITS) +#define OBS_SIZE (TIMER_INDEX + 1) +#define NUM_ATNS 1 +#define NUM_ACTIONS 8 +#define MAX_SOLUTION_DEPTH 16 +#define CURRICULUM_DEPTH_COUNT 6 +#define STEP_REWARD (-0.01f) +#ifndef VISIBLE_TARGET_TABLE_PATH +#define VISIBLE_TARGET_TABLE_PATH "ocean/affine_lock/generated/affine_lock_8action_visible_targets.bin" +#endif +#define ACT_SIZES {NUM_ACTIONS} +#define PUF_STEPS_PER_SEC 2 + +#define PERF_WEIGHTING_LINEAR 0 +#define PERF_WEIGHTING_QUADRATIC 1 + +#ifndef AFFINE_LOCK_GPU_SHARED_OBS +#define AFFINE_LOCK_GPU_SHARED_OBS 1 +#endif +#ifndef AFFINE_LOCK_GPU_SHARED_BLOCK +#define AFFINE_LOCK_GPU_SHARED_BLOCK 128 +#endif +#define AFFINE_LOCK_GPU_DEPTH_LUT_SIZE (MAX_SOLUTION_DEPTH + 1) + +static_assert(AFFINE_LOCK_GPU_SHARED_BLOCK >= 32 && + AFFINE_LOCK_GPU_SHARED_BLOCK <= 256 && + AFFINE_LOCK_GPU_SHARED_BLOCK % 32 == 0, + "AFFINE_LOCK_GPU_SHARED_BLOCK must contain whole warps"); + +#if !AFFINE_LOCK_GPU_SHARED_OBS +#ifndef AFFINE_LOCK_GPU_LANES +#define AFFINE_LOCK_GPU_LANES 4 +#endif +#define AFFINE_LOCK_GPU_BLOCK 256 +static_assert(AFFINE_LOCK_GPU_LANES == 4 || AFFINE_LOCK_GPU_LANES == 8 || + AFFINE_LOCK_GPU_LANES == 16 || AFFINE_LOCK_GPU_LANES == 32, + "AFFINE_LOCK_GPU_LANES must be a power-of-two subwarp"); +static_assert(AFFINE_LOCK_GPU_BLOCK % AFFINE_LOCK_GPU_LANES == 0, + "block size must contain whole environments"); +#endif + +struct Log { + float perf; + float score; + float solve_rate; + float max_depth_solve; + float episode_return; + float episode_length; + float solve_steps; + float timeout_rate; + float solve_efficiency; + float target_distance; + float solved_target_distance; + float d6_rate; + float d6_solve_rate; + float d8_rate; + float d8_solve_rate; + float d16_rate; + float d16_solve_rate; + float n; +}; + +static_assert(sizeof(Log) == 18 * sizeof(float), + "trainer log reduction requires a packed float-only Log"); + +// The trainer only reads Env::log for a GPU backend. Runtime state is kept in +// a separate compact array so log scans do not pull state into cache and state +// updates do not stride over the relatively large log payload. +struct Env { + Log log; + Agent agents[1]; + int num_agents; + int tag; + int boundary_reached; + unsigned int rng; +}; + +// Exactly 32 bytes: four adjacent environment records fit in one 128-byte +// transaction. The default shared-observation kernel reads one record per env. +typedef struct GpuAffineLockState { + uint32_t rng; + uint16_t state; + uint16_t target; + int step_count; + int max_steps; + int scramble_depth; + int curriculum_depth; + int target_distance; + float episode_return; +} GpuAffineLockState; + +static_assert(sizeof(GpuAffineLockState) == 32, + "GpuAffineLockState layout is performance-sensitive"); + +typedef struct GpuAffineLockConfig { + int start_depth; + int max_depth; + int step_grace; + int perf_weighting; + uint32_t depth_first[AFFINE_LOCK_GPU_DEPTH_LUT_SIZE]; + uint32_t depth_counts[AFFINE_LOCK_GPU_DEPTH_LUT_SIZE]; +} GpuAffineLockConfig; + +__constant__ GpuAffineLockConfig d_affine_lock_config; + +static struct { + Env* envs; + GpuAffineLockState* states; + uint32_t* target_pairs; + int n; + obs_t* observations; + float* actions; + float* rewards; + float* terminals; + cudaStream_t stream; + GpuAffineLockConfig config; +} g_gpu; + +static void gpu_affine_lock_check(cudaError_t status, const char* operation) { + if (status != cudaSuccess) { + std::fprintf(stderr, "Affine Lock CUDA: %s failed: %s\n", + operation, cudaGetErrorString(status)); + std::exit(1); + } +} + +#if !AFFINE_LOCK_GPU_SHARED_OBS +static int gpu_affine_lock_grid(int threads) { + return (threads + AFFINE_LOCK_GPU_BLOCK - 1) / AFFINE_LOCK_GPU_BLOCK; +} +#endif + +__device__ __forceinline__ uint32_t gpu_affine_lock_random_mixed_u32( + GpuAffineLockState* env) { + env->rng = env->rng * 1664525u + 1013904223u; + uint32_t x = env->rng; + x ^= x >> 16; + x *= 0x7feb352du; + x ^= x >> 15; + x *= 0x846ca68bu; + x ^= x >> 16; + return x; +} + +__device__ __forceinline__ int gpu_affine_lock_random_bounded( + GpuAffineLockState* env, int bound) { + uint32_t ubound = (uint32_t)bound; + uint32_t limit = UINT32_MAX - UINT32_MAX % ubound; + uint32_t value = gpu_affine_lock_random_mixed_u32(env); + while (value >= limit) { + value = gpu_affine_lock_random_mixed_u32(env); + } + return (int)(value % ubound); +} + +__device__ __forceinline__ void gpu_affine_lock_reset_state( + GpuAffineLockState* env, const uint32_t* target_pairs) { + env->scramble_depth = env->curriculum_depth; + env->step_count = 0; + env->episode_return = 0.0f; + int depth = env->scramble_depth; + uint32_t count = d_affine_lock_config.depth_counts[depth]; + int choice = gpu_affine_lock_random_bounded(env, (int)count); + uint32_t record_index = d_affine_lock_config.depth_first[depth] + + (uint32_t)choice; + uint32_t pair = target_pairs[record_index]; + env->state = (uint16_t)(pair & 0xffffu); + env->target = (uint16_t)(pair >> 16); + env->target_distance = env->scramble_depth; + env->max_steps = env->target_distance + d_affine_lock_config.step_grace; +} + +__device__ __forceinline__ uint16_t gpu_affine_lock_apply_action( + uint16_t state, int action) { + uint32_t next = state; + switch (action) { + case 0: + next = (state >> 1) | ((state & 1u) << 15); + break; + case 1: + next = ((state << 1) & 0xffffu) | ((state >> 15) & 1u); + break; + case 2: + next = state ^ 0xfe00u; + break; + case 3: + next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); + break; + case 4: + next = ((state & 0x3333u) << 2) | ((state & 0xccccu) >> 2); + break; + case 5: + next = ((state & 0x0f0fu) << 4) | ((state & 0xf0f0u) >> 4); + break; + case 6: + next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); + next = ((next & 0x3333u) << 2) | ((next & 0xccccu) >> 2); + break; + case 7: + next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); + next = ((next & 0x3333u) << 2) | ((next & 0xccccu) >> 2); + next = ((next & 0x0f0fu) << 4) | ((next & 0xf0f0u) >> 4); + break; + } + return (uint16_t)(next & 0xffffu); +} + +__device__ __forceinline__ int gpu_affine_lock_next_curriculum_depth( + int current_depth) { + constexpr int curriculum_depths[CURRICULUM_DEPTH_COUNT] = {2, 4, 5, 6, 8, 16}; +#pragma unroll + for (int i = 0; i < CURRICULUM_DEPTH_COUNT; i++) { + int depth = curriculum_depths[i]; + if (depth > current_depth) { + return depth < d_affine_lock_config.max_depth + ? depth : d_affine_lock_config.max_depth; + } + } + return d_affine_lock_config.max_depth; +} + +__device__ __forceinline__ void gpu_affine_lock_add_log( + Env* trainer_env, const GpuAffineLockState* env, int solved) { + int log_depth = env->target_distance; + int at_max_depth = log_depth == d_affine_lock_config.max_depth; + float ratio = log_depth / (float)d_affine_lock_config.max_depth; + float solve_credit = 0.0f; + if (solved) { + solve_credit = d_affine_lock_config.perf_weighting == PERF_WEIGHTING_QUADRATIC + ? ratio * ratio : ratio; + } + Log* log = &trainer_env->log; + log->perf += solve_credit; + log->score += solve_credit; + log->solve_rate += solved; + log->max_depth_solve += solved && at_max_depth; + log->episode_return += env->episode_return; + log->episode_length += env->step_count; + log->solve_steps += solved ? env->step_count : 0; + log->timeout_rate += !solved; + log->solve_efficiency += solved + ? env->step_count / (float)log_depth : 0.0f; + log->target_distance += env->target_distance; + log->solved_target_distance += solved ? env->target_distance : 0; + log->d6_rate += log_depth == 6; + log->d6_solve_rate += solved && log_depth == 6; + log->d8_rate += log_depth == 8; + log->d8_solve_rate += solved && log_depth == 8; + log->d16_rate += log_depth == 16; + log->d16_solve_rate += solved && log_depth == 16; + log->n += 1; +} + +__device__ __forceinline__ uint32_t gpu_affine_lock_step_one( + Env* trainer_env, GpuAffineLockState* env, + const uint32_t* target_pairs, float action, + float* reward_out, float* terminal_out, float* timer_out) { + float reward = STEP_REWARD; + float terminal = 0.0f; + int solved = 0; + env->step_count += 1; + int invalid = !isfinite(action) || action < 0.0f || action > NUM_ACTIONS - 1; + if (invalid) { + reward = -1.0f; + terminal = 1.0f; + } else { + env->state = gpu_affine_lock_apply_action(env->state, (int)action); + if (env->state == env->target) { + reward = 1.0f; + terminal = 1.0f; + solved = 1; + } else if (env->step_count >= env->max_steps) { + reward = -1.0f; + terminal = 1.0f; + } + } + env->episode_return += reward; + if (terminal != 0.0f) { + gpu_affine_lock_add_log(trainer_env, env, solved); + env->curriculum_depth = solved + ? gpu_affine_lock_next_curriculum_depth(env->scramble_depth) + : d_affine_lock_config.start_depth; + gpu_affine_lock_reset_state(env, target_pairs); + } + *reward_out = reward; + *terminal_out = terminal; + *timer_out = env->step_count / (float)env->max_steps; + return (uint32_t)env->state | ((uint32_t)env->target << 16); +} + +#if !AFFINE_LOCK_GPU_SHARED_OBS +__device__ __forceinline__ void gpu_affine_lock_write_observations( + obs_t* observations, uint32_t packed_bits, float timer, int lane) { +#pragma unroll + for (int bit = lane; bit < 2 * BITS; bit += AFFINE_LOCK_GPU_LANES) { + observations[bit] = __float2bfloat16( + (packed_bits & (1u << bit)) ? 1.0f : -1.0f); + } + if (lane == 0) { + observations[TIMER_INDEX] = __float2bfloat16(timer); + } +} + +__global__ __launch_bounds__(AFFINE_LOCK_GPU_BLOCK) +void gpu_affine_lock_reset_kernel(Env* envs, GpuAffineLockState* states, + const uint32_t* target_pairs, obs_t* observations, + float* rewards, float* terminals, int num_envs) { + int thread = blockIdx.x * blockDim.x + threadIdx.x; + int relative_env = thread / AFFINE_LOCK_GPU_LANES; + int lane = thread & (AFFINE_LOCK_GPU_LANES - 1); + int active = relative_env < num_envs; + uint32_t packed_bits = 0; + float timer = 0.0f; + if (active && lane == 0) { + GpuAffineLockState* env = &states[relative_env]; + gpu_affine_lock_reset_state(env, target_pairs); + rewards[relative_env] = 0.0f; + terminals[relative_env] = 0.0f; + packed_bits = (uint32_t)env->state | ((uint32_t)env->target << 16); + } + int leader = (threadIdx.x & 31) & ~(AFFINE_LOCK_GPU_LANES - 1); + packed_bits = __shfl_sync(0xffffffffu, packed_bits, leader); + timer = __shfl_sync(0xffffffffu, timer, leader); + if (active) { + gpu_affine_lock_write_observations( + observations + (size_t)relative_env * OBS_SIZE, + packed_bits, timer, lane); + } + (void)envs; +} + +__global__ __launch_bounds__(AFFINE_LOCK_GPU_BLOCK) +void gpu_affine_lock_step_kernel(Env* envs, GpuAffineLockState* states, + const uint32_t* target_pairs, const float* actions, + obs_t* observations, float* rewards, float* terminals, + int num_envs) { + int thread = blockIdx.x * blockDim.x + threadIdx.x; + int relative_env = thread / AFFINE_LOCK_GPU_LANES; + int lane = thread & (AFFINE_LOCK_GPU_LANES - 1); + int active = relative_env < num_envs; + uint32_t packed_bits = 0; + float timer = 0.0f; + if (active && lane == 0) { + packed_bits = gpu_affine_lock_step_one( + &envs[relative_env], &states[relative_env], target_pairs, + actions[(size_t)relative_env * NUM_ATNS], + &rewards[relative_env], &terminals[relative_env], &timer); + } + int leader = (threadIdx.x & 31) & ~(AFFINE_LOCK_GPU_LANES - 1); + packed_bits = __shfl_sync(0xffffffffu, packed_bits, leader); + timer = __shfl_sync(0xffffffffu, timer, leader); + if (active) { + gpu_affine_lock_write_observations( + observations + (size_t)relative_env * OBS_SIZE, + packed_bits, timer, lane); + } +} +#endif + +#if AFFINE_LOCK_GPU_SHARED_OBS +// One simulation thread per environment writes +// a conflict-free 33-float shared-memory row, then the whole block converts and +// stores a linear bf16 tile with fully coalesced global writes. +__global__ __launch_bounds__(AFFINE_LOCK_GPU_SHARED_BLOCK) +void gpu_affine_lock_shared_reset_kernel(Env* envs, + GpuAffineLockState* states, const uint32_t* target_pairs, + obs_t* observations, float* rewards, float* terminals, int num_envs) { + __shared__ float observation_tile[AFFINE_LOCK_GPU_SHARED_BLOCK * OBS_SIZE]; + int block_start = blockIdx.x * AFFINE_LOCK_GPU_SHARED_BLOCK; + int relative_env = block_start + threadIdx.x; + int active_count = num_envs - block_start; + if (active_count > AFFINE_LOCK_GPU_SHARED_BLOCK) { + active_count = AFFINE_LOCK_GPU_SHARED_BLOCK; + } + if (active_count < 0) { + active_count = 0; + } + if (threadIdx.x < active_count) { + GpuAffineLockState* env = &states[relative_env]; + gpu_affine_lock_reset_state(env, target_pairs); + rewards[relative_env] = 0.0f; + terminals[relative_env] = 0.0f; + uint32_t bits = (uint32_t)env->state | ((uint32_t)env->target << 16); + float* row = observation_tile + threadIdx.x * OBS_SIZE; +#pragma unroll + for (int bit = 0; bit < 2 * BITS; bit++) { + row[bit] = (bits & (1u << bit)) ? 1.0f : -1.0f; + } + row[TIMER_INDEX] = 0.0f; + } + __syncthreads(); + int tile_values = active_count * OBS_SIZE; + for (int value = threadIdx.x; value < tile_values; value += blockDim.x) { + observations[(size_t)block_start * OBS_SIZE + value] = + __float2bfloat16(observation_tile[value]); + } + (void)envs; +} + +__global__ __launch_bounds__(AFFINE_LOCK_GPU_SHARED_BLOCK) +void gpu_affine_lock_shared_step_kernel(Env* envs, + GpuAffineLockState* states, const uint32_t* target_pairs, + const float* actions, obs_t* observations, + float* rewards, float* terminals, int num_envs) { + __shared__ float observation_tile[AFFINE_LOCK_GPU_SHARED_BLOCK * OBS_SIZE]; + int block_start = blockIdx.x * AFFINE_LOCK_GPU_SHARED_BLOCK; + int relative_env = block_start + threadIdx.x; + int active_count = num_envs - block_start; + if (active_count > AFFINE_LOCK_GPU_SHARED_BLOCK) { + active_count = AFFINE_LOCK_GPU_SHARED_BLOCK; + } + if (active_count < 0) { + active_count = 0; + } + if (threadIdx.x < active_count) { + float timer = 0.0f; + uint32_t bits = gpu_affine_lock_step_one( + &envs[relative_env], &states[relative_env], target_pairs, + actions[(size_t)relative_env * NUM_ATNS], + &rewards[relative_env], &terminals[relative_env], &timer); + float* row = observation_tile + threadIdx.x * OBS_SIZE; +#pragma unroll + for (int bit = 0; bit < 2 * BITS; bit++) { + row[bit] = (bits & (1u << bit)) ? 1.0f : -1.0f; + } + row[TIMER_INDEX] = timer; + } + __syncthreads(); + int tile_values = active_count * OBS_SIZE; + for (int value = threadIdx.x; value < tile_values; value += blockDim.x) { + observations[(size_t)block_start * OBS_SIZE + value] = + __float2bfloat16(observation_tile[value]); + } +} +#endif + +void puf_log(Log* log, Dict* out) { + float nsolve = log->solve_rate; + float solved_min_win_moves = nsolve + ? log->solved_target_distance / nsolve : 0; + float conditional_solve_steps = nsolve ? log->solve_steps / nsolve : 0; + float conditional_solve_efficiency = nsolve + ? log->solve_efficiency / nsolve : 0; + + dict_set(out, "perf", log->perf); + dict_set(out, "score", log->score); + dict_set(out, "solve_rate", log->solve_rate); + dict_set(out, "max_depth_solve", log->max_depth_solve); + dict_set(out, "episode_return", log->episode_return); + dict_set(out, "episode_length", log->episode_length); + dict_set(out, "timeout_rate", log->timeout_rate); + dict_set(out, "min_win_moves", log->target_distance); + dict_set(out, "solved_min_win_moves", solved_min_win_moves); + dict_set(out, "conditional_solve_steps", conditional_solve_steps); + dict_set(out, "conditional_solve_efficiency", conditional_solve_efficiency); + dict_set(out, "d6_solve_rate", log->d6_rate + ? log->d6_solve_rate / log->d6_rate : 0); + dict_set(out, "d8_solve_rate", log->d8_rate + ? log->d8_solve_rate / log->d8_rate : 0); + dict_set(out, "d16_solve_rate", log->d16_rate + ? log->d16_solve_rate / log->d16_rate : 0); + dict_set(out, "n", log->n); +} + +static int gpu_affine_lock_host_has_depth( + const GpuAffineLockConfig* config, int depth) { + return depth >= 0 && depth < AFFINE_LOCK_GPU_DEPTH_LUT_SIZE + && config->depth_counts[depth] != 0; +} + +static int gpu_affine_lock_host_next_depth(int current_depth, int max_depth) { + static const int curriculum_depths[CURRICULUM_DEPTH_COUNT] = {2, 4, 5, 6, 8, 16}; + for (int i = 0; i < CURRICULUM_DEPTH_COUNT; i++) { + int depth = curriculum_depths[i]; + if (depth > current_depth) { + return depth < max_depth ? depth : max_depth; + } + } + return max_depth; +} + +static void gpu_affine_lock_validate_curriculum( + const GpuAffineLockConfig* config) { + if (config->start_depth <= 0 || + config->max_depth < config->start_depth || + config->max_depth > MAX_SOLUTION_DEPTH) { + std::fprintf(stderr, + "Affine Lock CUDA: invalid curriculum range start=%d max=%d\n", + config->start_depth, config->max_depth); + std::exit(1); + } + int depth = config->start_depth; + for (int i = 0; i <= CURRICULUM_DEPTH_COUNT; i++) { + if (!gpu_affine_lock_host_has_depth(config, depth)) { + std::fprintf(stderr, + "Affine Lock CUDA: target table has no depth %d section\n", depth); + std::exit(1); + } + if (depth + config->step_grace <= 0) { + std::fprintf(stderr, + "Affine Lock CUDA: depth %d with step_grace=%d has no valid steps\n", + depth, config->step_grace); + std::exit(1); + } + if (depth == config->max_depth) { + return; + } + int next = gpu_affine_lock_host_next_depth(depth, config->max_depth); + if (next == depth) { + break; + } + depth = next; + } + std::fprintf(stderr, "Affine Lock CUDA: curriculum does not reach max depth %d\n", + config->max_depth); + std::exit(1); +} + +Env* puf_vec_create(int n, Dict* env_kwargs, + obs_t* observations, float* actions, + float* rewards, float* terminals) { + if (n <= 0) { + std::fprintf(stderr, "Affine Lock CUDA: vector size must be positive\n"); + std::exit(1); + } + if (g_gpu.envs != nullptr) { + std::fprintf(stderr, "Affine Lock CUDA: vector already exists\n"); + std::exit(1); + } + + VisibleTargetTable table = {}; + if (visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table) != 0) { + std::fprintf(stderr, + "Affine Lock CUDA: failed to load visible target table %s\n", + VISIBLE_TARGET_TABLE_PATH); + std::exit(1); + } + if (table.num_actions != NUM_ACTIONS) { + std::fprintf(stderr, + "Affine Lock CUDA: target table has %u actions, expected %d\n", + table.num_actions, NUM_ACTIONS); + visible_targets_free(&table); + std::exit(1); + } + + GpuAffineLockConfig config = {}; + config.start_depth = (int)dict_get(env_kwargs, "start_depth"); + config.max_depth = (int)dict_get(env_kwargs, "max_depth"); + config.step_grace = (int)dict_get(env_kwargs, "step_grace"); + config.perf_weighting = (int)dict_get(env_kwargs, "perf_weighting"); + for (uint32_t i = 0; i < table.depth_count; i++) { + const VisibleTargetDepth* depth = &table.depths[i]; + if (depth->depth >= AFFINE_LOCK_GPU_DEPTH_LUT_SIZE || + depth->stored_count == 0 || + config.depth_counts[depth->depth] != 0) { + std::fprintf(stderr, + "Affine Lock CUDA: invalid target-table depth section %u\n", + depth->depth); + visible_targets_free(&table); + std::exit(1); + } + config.depth_first[depth->depth] = depth->first_record; + config.depth_counts[depth->depth] = depth->stored_count; + for (uint32_t record_offset = 0; + record_offset < depth->stored_count; record_offset++) { + const VisibleTargetRecord* record = + &table.records[depth->first_record + record_offset]; + if (record->depth != depth->depth) { + std::fprintf(stderr, + "Affine Lock CUDA: record depth %u does not match section %u\n", + (unsigned int)record->depth, depth->depth); + visible_targets_free(&table); + std::exit(1); + } + } + } + gpu_affine_lock_validate_curriculum(&config); + + uint32_t* host_pairs = (uint32_t*)std::malloc( + (size_t)table.record_count * sizeof(uint32_t)); + if (host_pairs == nullptr) { + std::perror("malloc"); + visible_targets_free(&table); + std::exit(1); + } + for (uint32_t i = 0; i < table.record_count; i++) { + host_pairs[i] = (uint32_t)table.records[i].start + | ((uint32_t)table.records[i].target << 16); + } + + GpuAffineLockState* host_states = (GpuAffineLockState*)std::calloc( + (size_t)n, sizeof(GpuAffineLockState)); + if (host_states == nullptr) { + std::perror("calloc"); + std::free(host_pairs); + visible_targets_free(&table); + std::exit(1); + } + unsigned int running_seed = (unsigned int)dict_get(env_kwargs, "seed"); + for (int i = 0; i < n; i++) { + host_states[i].rng = (uint32_t)rand_r(&running_seed); + host_states[i].curriculum_depth = config.start_depth; + } + + Env* device_envs = nullptr; + GpuAffineLockState* device_states = nullptr; + uint32_t* device_pairs = nullptr; + gpu_affine_lock_check(cudaMalloc((void**)&device_envs, + (size_t)n * sizeof(Env)), "cudaMalloc envs"); + gpu_affine_lock_check(cudaMalloc((void**)&device_states, + (size_t)n * sizeof(GpuAffineLockState)), "cudaMalloc states"); + gpu_affine_lock_check(cudaMalloc((void**)&device_pairs, + (size_t)table.record_count * sizeof(uint32_t)), "cudaMalloc target pairs"); + gpu_affine_lock_check(cudaMemset(device_envs, 0, + (size_t)n * sizeof(Env)), "clear env logs"); + gpu_affine_lock_check(cudaMemcpy(device_states, host_states, + (size_t)n * sizeof(GpuAffineLockState), cudaMemcpyHostToDevice), "copy states"); + gpu_affine_lock_check(cudaMemcpy(device_pairs, host_pairs, + (size_t)table.record_count * sizeof(uint32_t), cudaMemcpyHostToDevice), + "copy target pairs"); + gpu_affine_lock_check(cudaMemcpyToSymbol(d_affine_lock_config, + &config, sizeof(config)), "copy config"); + + std::free(host_pairs); + std::free(host_states); + visible_targets_free(&table); + + g_gpu.envs = device_envs; + g_gpu.states = device_states; + g_gpu.target_pairs = device_pairs; + g_gpu.n = n; + g_gpu.observations = observations; + g_gpu.actions = actions; + g_gpu.rewards = rewards; + g_gpu.terminals = terminals; + g_gpu.stream = nullptr; + g_gpu.config = config; + return device_envs; +} + +void puf_bind_stream(cudaStream_t stream) { + g_gpu.stream = stream; +} + +// GPU creation is vector-only; puf_init exists to satisfy the common API. +void puf_init(Env* env, Dict* kwargs) { + (void)env; + (void)kwargs; +} + +void puf_reset(Env* env) { + (void)env; +#if AFFINE_LOCK_GPU_SHARED_OBS + int blocks = (g_gpu.n + AFFINE_LOCK_GPU_SHARED_BLOCK - 1) + / AFFINE_LOCK_GPU_SHARED_BLOCK; + gpu_affine_lock_shared_reset_kernel<<< + blocks, AFFINE_LOCK_GPU_SHARED_BLOCK, 0, g_gpu.stream>>>( + g_gpu.envs, g_gpu.states, g_gpu.target_pairs, + g_gpu.observations, g_gpu.rewards, g_gpu.terminals, g_gpu.n); +#else + int threads = g_gpu.n * AFFINE_LOCK_GPU_LANES; + gpu_affine_lock_reset_kernel<<< + gpu_affine_lock_grid(threads), AFFINE_LOCK_GPU_BLOCK, 0, g_gpu.stream>>>( + g_gpu.envs, g_gpu.states, g_gpu.target_pairs, + g_gpu.observations, g_gpu.rewards, g_gpu.terminals, g_gpu.n); +#endif + gpu_affine_lock_check(cudaPeekAtLastError(), "launch reset kernel"); +} + +void puf_step(Env* env) { + (void)env; +#if AFFINE_LOCK_GPU_SHARED_OBS + int blocks = (g_gpu.n + AFFINE_LOCK_GPU_SHARED_BLOCK - 1) + / AFFINE_LOCK_GPU_SHARED_BLOCK; + gpu_affine_lock_shared_step_kernel<<< + blocks, AFFINE_LOCK_GPU_SHARED_BLOCK, 0, g_gpu.stream>>>( + g_gpu.envs, g_gpu.states, g_gpu.target_pairs, + g_gpu.actions, g_gpu.observations, + g_gpu.rewards, g_gpu.terminals, g_gpu.n); +#else + int threads = g_gpu.n * AFFINE_LOCK_GPU_LANES; + gpu_affine_lock_step_kernel<<< + gpu_affine_lock_grid(threads), AFFINE_LOCK_GPU_BLOCK, 0, g_gpu.stream>>>( + g_gpu.envs, g_gpu.states, g_gpu.target_pairs, + g_gpu.actions, g_gpu.observations, + g_gpu.rewards, g_gpu.terminals, g_gpu.n); +#endif + gpu_affine_lock_check(cudaPeekAtLastError(), "launch step kernel"); +} + +void puf_close(Env* env) { + (void)env; + if (IsWindowReady()) { + CloseWindow(); + } + if (g_gpu.envs != nullptr) { + gpu_affine_lock_check(cudaFree(g_gpu.envs), "cudaFree envs"); + } + if (g_gpu.states != nullptr) { + gpu_affine_lock_check(cudaFree(g_gpu.states), "cudaFree states"); + } + if (g_gpu.target_pairs != nullptr) { + gpu_affine_lock_check(cudaFree(g_gpu.target_pairs), "cudaFree target pairs"); + } + g_gpu = {}; +} + +void puf_render(Env* env) { + (void)env; + if (g_gpu.envs == nullptr || g_gpu.n < 1) { + return; + } + if (IsWindowReady() && (WindowShouldClose() || IsKeyPressed(KEY_ESCAPE))) { + puf_close(g_gpu.envs); + std::exit(0); + } + if (!IsWindowReady()) { + InitWindow(780, 360, "PufferLib AffineLock CUDA"); + SetTargetFPS(30); + } + if (g_gpu.stream != nullptr) { + gpu_affine_lock_check(cudaStreamSynchronize(g_gpu.stream), + "synchronize render stream"); + } + GpuAffineLockState state; + float reward = 0.0f; + float terminal = 0.0f; + gpu_affine_lock_check(cudaMemcpy(&state, g_gpu.states, sizeof(state), + cudaMemcpyDeviceToHost), "copy render state"); + gpu_affine_lock_check(cudaMemcpy(&reward, g_gpu.rewards, sizeof(reward), + cudaMemcpyDeviceToHost), "copy render reward"); + gpu_affine_lock_check(cudaMemcpy(&terminal, g_gpu.terminals, sizeof(terminal), + cudaMemcpyDeviceToHost), "copy render terminal"); + + uint32_t mismatches = (state.state ^ state.target) & 0xffffu; + const char* status = terminal == 0.0f + ? "running" : (reward > 0.0f ? "solved" : "failed"); + Color status_color = terminal == 0.0f + ? (Color){190, 198, 206, 255} + : (reward > 0.0f + ? (Color){80, 210, 140, 255} + : (Color){238, 88, 88, 255}); + + BeginDrawing(); + ClearBackground((Color){6, 24, 24, 255}); + DrawText("Affine Lock CUDA", 30, 24, 28, RAYWHITE); + DrawText(TextFormat("depth %d/%d step %d/%d last reward %.2f", + state.scramble_depth, g_gpu.config.max_depth, + state.step_count, state.max_steps, reward), + 30, 62, 20, (Color){180, 190, 200, 255}); + DrawText(TextFormat("status %s mismatches 0x%04x", + status, mismatches), 30, 90, 20, status_color); + + const char* row_label[2] = {"current", "target"}; + uint32_t row_value[2] = {state.state, state.target}; + int row_y[2] = {138, 220}; + for (int row = 0; row < 2; row++) { + DrawText(row_label[row], 30, row_y[row] + 9, 20, RAYWHITE); + for (int bit = 0; bit < BITS; bit++) { + int x = 145 + bit * 34; + int on = (row_value[row] >> bit) & 1u; + int mismatch = ((state.state ^ state.target) >> bit) & 1u; + Color fill = on + ? (Color){80, 210, 140, 255} + : (Color){38, 48, 58, 255}; + Color border = mismatch + ? (Color){238, 88, 88, 255} + : (Color){182, 196, 205, 255}; + DrawRectangle(x, row_y[row], 24, 34, fill); + DrawRectangleLinesEx((Rectangle){(float)x, (float)row_y[row], 24, 34}, + mismatch ? 3 : 1, border); + DrawText(TextFormat("%d", bit), x + 5, row_y[row] + 40, 10, + (Color){128, 140, 150, 255}); + } + } + DrawText("GPU-resident environment", 30, 310, 16, + (Color){160, 170, 178, 255}); + EndDrawing(); + puf_web_vsync(); +} + +#endif diff --git a/ocean/affine_lock/affine_lock.h b/ocean/affine_lock/affine_lock.h index aef57f422e..dc5419aab8 100644 --- a/ocean/affine_lock/affine_lock.h +++ b/ocean/affine_lock/affine_lock.h @@ -89,6 +89,8 @@ struct Env { int target_distance; float episode_return; int owns_shared; + int pending_human_action; + int hint_action; // -2 none, -1 already solved, else action to press AffineLockShared* shared; }; typedef Env AffineLock; @@ -165,6 +167,8 @@ static void init_env(AffineLock* env, AffineLockShared* shared, unsigned int see env->rng = seed; env->num_agents = 1; env->curriculum_depth = shared->start_depth; + env->pending_human_action = -1; + env->hint_action = -2; } void puf_init(Env* env, Dict* kwargs) { @@ -279,6 +283,7 @@ static void reset_state(AffineLock* env) { env->solution_actions[i] = (record->packed_actions >> (3 * i)) & 7; } env->max_steps = env->target_distance + shared->step_grace; + env->hint_action = -2; } static void compute_observations(AffineLock* env) { @@ -312,9 +317,9 @@ static int next_curriculum_depth( const AffineLockShared* shared, int current_de return shared->max_depth; } -static int human_controls(AffineLock *env) { +static void human_controls(AffineLock *env) { if (!IsWindowReady() || !IsKeyDown(KEY_LEFT_SHIFT)) { - return 0; + return; } static const int keys[NUM_ACTIONS] = { KEY_ONE, KEY_TWO, KEY_THREE, KEY_FOUR, @@ -322,17 +327,21 @@ static int human_controls(AffineLock *env) { }; for (int i = 0; i < NUM_ACTIONS; i++) { if (IsKeyPressed(keys[i])) { - env->agents[0].actions[0] = (float)i; - return 1; + env->pending_human_action = i; + return; } } - return -1; } void puf_step(AffineLock* env) { - if (human_controls(env) < 0) { - return; + if (IsWindowReady() && IsKeyDown(KEY_LEFT_SHIFT)) { + if (env->pending_human_action < 0) { + return; + } + env->agents[0].actions[0] = (float)env->pending_human_action; + env->pending_human_action = -1; } + env->hint_action = -2; AffineLockShared* shared = env->shared; float reward = STEP_REWARD; int terminal = 0; @@ -433,6 +442,20 @@ void my_vec_close(Env* envs) { free(envs[0].shared); } +static const char* action_name(int action) { + static const char* names[NUM_ACTIONS] = { + "shift_left", "shift_right", "invert_right_7", "swap_adjacent_bits", + "swap_adjacent_pairs", "swap_nibbles_each_byte", "reverse_each_nibble", + "reverse_each_byte", + }; + return names[action]; +} + +static void show_hint(AffineLock* env) { + env->hint_action = env->step_count < env->solution_length ? + env->solution_actions[env->step_count] : -1; +} + void puf_render(AffineLock* env) { if (IsWindowReady() && (WindowShouldClose() || IsKeyPressed(KEY_ESCAPE))) { puf_close(env); @@ -467,6 +490,14 @@ void puf_render(AffineLock* env) { 30, 62, 20, (Color){180, 190, 200, 255}); DrawText(TextFormat("status %s mismatches 0x%04x", status, rel), 30, 90, 20, status_color); + if (env->hint_action != -2) { + const char* hint = env->hint_action >= 0 ? + TextFormat("Hint: press %d (%s)", env->hint_action + 1, + action_name(env->hint_action)) : + "Hint: already solved"; + DrawText(hint, 780 - MeasureText(hint, 18) - 30, 90, + 18, (Color){245, 205, 92, 255}); + } const char* row_label[2] = {"current", "target"}; uint32_t row_value[2] = {env->state, env->target}; @@ -476,7 +507,7 @@ void puf_render(AffineLock* env) { for (int bit = 0; bit < BITS; bit++) { int x = 145 + bit * 34; int on = (row_value[row] >> bit) & 1u; - int mismatch = ((env->state ^ env->target) >> bit) & 1u; + int mismatch = row == 0 && (((env->state ^ env->target) >> bit) & 1u); Color fill = on ? (Color){80, 210, 140, 255} : (Color){38, 48, 58, 255}; Color border = mismatch ? (Color){238, 88, 88, 255} : (Color){182, 196, 205, 255}; DrawRectangle(x, row_y[row], 24, 34, fill); @@ -488,9 +519,9 @@ void puf_render(AffineLock* env) { } } - DrawText("1 shiftL 2 shiftR 3 inv7 4 bit-swap 5 pair-swap", + DrawText("shift+1 shiftL shift+2 shiftR shift+3 inv7 shift+4 bit-swap shift+5 pair-swap", 30, 300, 16, (Color){160, 170, 178, 255}); - DrawText("6 nib-swap 7 rev-nib 8 rev-byte R reset", + DrawText("shift+6 nib-swap shift+7 rev-nib shift+8 rev-byte H hint R reset", 30, 322, 16, (Color){160, 170, 178, 255}); EndDrawing(); puf_web_vsync(); diff --git a/ocean/affine_lock/tests/run_cuda.sh b/ocean/affine_lock/tests/run_cuda.sh new file mode 100755 index 0000000000..f7762be9a1 --- /dev/null +++ b/ocean/affine_lock/tests/run_cuda.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "$0")/../../.." && pwd)" +cd "$ROOT" + +OUT="${TMPDIR:-/tmp}/affine_lock_cuda_tests" +CUDA_ROOT="${CUDA_HOME:-${CUDA_PATH:-/usr/local/cuda}}" +NVCC_BIN="${NVCC:-$CUDA_ROOT/bin/nvcc}" +CUDA_ARCH="${NVCC_ARCH:-native}" + +RAYLIB_ROOT="$ROOT/raylib-5.5_linux_amd64" +if [ ! -d "$RAYLIB_ROOT/include" ]; then + echo "raylib-5.5_linux_amd64 not found" >&2 + exit 1 +fi + +"$NVCC_BIN" \ + -std=c++17 -O3 -lineinfo -arch="$CUDA_ARCH" \ + -Xcompiler=-Wall,-Wextra,-Werror,-Wno-unused-function,-Wno-unused-parameter,-Wno-missing-field-initializers \ + -Xcompiler=-ffunction-sections,-fdata-sections \ + -I"$ROOT" -I"$ROOT/src" -I"$ROOT/ocean/affine_lock" \ + -I"$ROOT/vendor" -I"$RAYLIB_ROOT/include" \ + "$ROOT/ocean/affine_lock/tests/test_affine_lock_cuda.cu" \ + "$RAYLIB_ROOT/lib/libraylib.a" \ + -Xlinker=--gc-sections \ + -lGL -lpthread -ldl -lrt -lm \ + -o "$OUT" + +"$OUT" diff --git a/ocean/affine_lock/tests/test_affine_lock_cuda.cu b/ocean/affine_lock/tests/test_affine_lock_cuda.cu new file mode 100644 index 0000000000..54cb06ad34 --- /dev/null +++ b/ocean/affine_lock/tests/test_affine_lock_cuda.cu @@ -0,0 +1,1109 @@ +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "../affine_lock.cu" + +#define EXPECT_TRUE(condition) do { \ + if (!(condition)) { \ + std::fprintf(stderr, "%s:%d: expected true: %s\n", \ + __FILE__, __LINE__, #condition); \ + std::exit(1); \ + } \ +} while (0) + +#define EXPECT_EQ(actual, expected) do { \ + auto actual_value = (actual); \ + auto expected_value = (expected); \ + if (actual_value != expected_value) { \ + std::fprintf(stderr, "%s:%d: expected %s == %s, got %lld != %lld\n", \ + __FILE__, __LINE__, #actual, #expected, \ + (long long)actual_value, (long long)expected_value); \ + std::exit(1); \ + } \ +} while (0) + +#define EXPECT_NEAR(actual, expected, tolerance) do { \ + float actual_value = (float)(actual); \ + float expected_value = (float)(expected); \ + if (!std::isfinite(actual_value) || !std::isfinite(expected_value) || \ + std::fabs(actual_value - expected_value) > (tolerance)) { \ + std::fprintf(stderr, "%s:%d: expected %s ~= %.9g, got %.9g\n", \ + __FILE__, __LINE__, #actual, expected_value, actual_value); \ + std::exit(1); \ + } \ +} while (0) + +static void check_cuda(cudaError_t status, const char* operation) { + if (status != cudaSuccess) { + std::fprintf(stderr, "%s failed: %s\n", operation, cudaGetErrorString(status)); + std::exit(1); + } +} + +typedef struct OracleState { + uint32_t rng; + uint16_t state; + uint16_t target; + int step_count; + int max_steps; + int scramble_depth; + int curriculum_depth; + int target_distance; + float episode_return; + Log log; +} OracleState; + +static uint32_t oracle_random_mixed_u32(OracleState* env) { + env->rng = env->rng * 1664525u + 1013904223u; + uint32_t x = env->rng; + x ^= x >> 16; + x *= 0x7feb352du; + x ^= x >> 15; + x *= 0x846ca68bu; + x ^= x >> 16; + return x; +} + +static int oracle_random_bounded(OracleState* env, int bound) { + uint32_t ubound = (uint32_t)bound; + uint32_t limit = UINT32_MAX - UINT32_MAX % ubound; + uint32_t value = oracle_random_mixed_u32(env); + while (value >= limit) { + value = oracle_random_mixed_u32(env); + } + return (int)(value % ubound); +} + +static const VisibleTargetDepth* oracle_depth( + const VisibleTargetTable* table, int requested_depth) { + for (uint32_t i = 0; i < table->depth_count; i++) { + if ((int)table->depths[i].depth == requested_depth) { + return &table->depths[i]; + } + } + return nullptr; +} + +static void oracle_reset_state(OracleState* env, + const VisibleTargetTable* table, int step_grace) { + env->scramble_depth = env->curriculum_depth; + env->step_count = 0; + env->episode_return = 0.0f; + const VisibleTargetDepth* depth = oracle_depth(table, env->scramble_depth); + EXPECT_TRUE(depth != nullptr); + int choice = oracle_random_bounded(env, (int)depth->stored_count); + const VisibleTargetRecord* record = + &table->records[depth->first_record + (uint32_t)choice]; + env->state = record->start; + env->target = record->target; + env->target_distance = record->depth; + env->max_steps = env->target_distance + step_grace; +} + +static uint16_t oracle_apply_action(uint16_t state, int action) { + uint32_t next = state; + switch (action) { + case 0: next = (state >> 1) | ((state & 1u) << 15); break; + case 1: next = ((state << 1) & 0xffffu) | ((state >> 15) & 1u); break; + case 2: next = state ^ 0xfe00u; break; + case 3: next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); break; + case 4: next = ((state & 0x3333u) << 2) | ((state & 0xccccu) >> 2); break; + case 5: next = ((state & 0x0f0fu) << 4) | ((state & 0xf0f0u) >> 4); break; + case 6: + next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); + next = ((next & 0x3333u) << 2) | ((next & 0xccccu) >> 2); + break; + case 7: + next = ((state & 0x5555u) << 1) | ((state & 0xaaaau) >> 1); + next = ((next & 0x3333u) << 2) | ((next & 0xccccu) >> 2); + next = ((next & 0x0f0fu) << 4) | ((next & 0xf0f0u) >> 4); + break; + } + return (uint16_t)(next & 0xffffu); +} + +static int oracle_next_curriculum_depth(int current_depth, int max_depth) { + static const int curriculum_depths[] = {2, 4, 5, 6, 8, 16}; + for (int depth : curriculum_depths) { + if (depth > current_depth) { + return depth < max_depth ? depth : max_depth; + } + } + return max_depth; +} + +static void oracle_add_log(OracleState* env, int solved, + int max_depth, int perf_weighting) { + int log_depth = env->target_distance; + int at_max_depth = log_depth == max_depth; + float ratio = log_depth / (float)max_depth; + float solve_credit = 0.0f; + if (solved) { + solve_credit = perf_weighting == PERF_WEIGHTING_QUADRATIC + ? ratio * ratio : ratio; + } + env->log.perf += solve_credit; + env->log.score += solve_credit; + env->log.solve_rate += solved; + env->log.max_depth_solve += solved && at_max_depth; + env->log.episode_return += env->episode_return; + env->log.episode_length += env->step_count; + env->log.solve_steps += solved ? env->step_count : 0; + env->log.timeout_rate += !solved; + env->log.solve_efficiency += solved + ? env->step_count / (float)log_depth : 0.0f; + env->log.target_distance += env->target_distance; + env->log.solved_target_distance += solved ? env->target_distance : 0; + env->log.d6_rate += log_depth == 6; + env->log.d6_solve_rate += solved && log_depth == 6; + env->log.d8_rate += log_depth == 8; + env->log.d8_solve_rate += solved && log_depth == 8; + env->log.d16_rate += log_depth == 16; + env->log.d16_solve_rate += solved && log_depth == 16; + env->log.n += 1; +} + +static void oracle_step(OracleState* env, float action, + const VisibleTargetTable* table, int start_depth, int max_depth, + int step_grace, int perf_weighting, float* reward, float* terminal) { + *reward = STEP_REWARD; + *terminal = 0.0f; + int solved = 0; + env->step_count += 1; + int invalid = !std::isfinite(action) || action < 0.0f || action > 7.0f; + if (invalid) { + *reward = -1.0f; + *terminal = 1.0f; + } else { + env->state = oracle_apply_action(env->state, (int)action); + if (env->state == env->target) { + *reward = 1.0f; + *terminal = 1.0f; + solved = 1; + } else if (env->step_count >= env->max_steps) { + *reward = -1.0f; + *terminal = 1.0f; + } + } + env->episode_return += *reward; + if (*terminal != 0.0f) { + oracle_add_log(env, solved, max_depth, perf_weighting); + env->curriculum_depth = solved + ? oracle_next_curriculum_depth(env->scramble_depth, max_depth) + : start_depth; + oracle_reset_state(env, table, step_grace); + } +} + +static uint16_t obs_bits(obs_t value) { + uint16_t bits = 0; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} + +static uint32_t float_bits(float value) { + uint32_t bits = 0; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; +} + +static void expect_log_equal(const Log& actual, const Log& expected) { + const float* a = (const float*)&actual; + const float* e = (const float*)&expected; + for (size_t i = 0; i < sizeof(Log) / sizeof(float); i++) { + EXPECT_EQ(float_bits(a[i]), float_bits(e[i])); + } +} + +static void expect_state_equal(const GpuAffineLockState& actual, + const OracleState& expected) { + EXPECT_EQ(actual.rng, expected.rng); + EXPECT_EQ(actual.state, expected.state); + EXPECT_EQ(actual.target, expected.target); + EXPECT_EQ(actual.step_count, expected.step_count); + EXPECT_EQ(actual.max_steps, expected.max_steps); + EXPECT_EQ(actual.scramble_depth, expected.scramble_depth); + EXPECT_EQ(actual.curriculum_depth, expected.curriculum_depth); + EXPECT_EQ(actual.target_distance, expected.target_distance); + EXPECT_EQ(float_bits(actual.episode_return), + float_bits(expected.episode_return)); +} + +static void expect_observation_equal(const obs_t* actual, + const OracleState& expected) { + uint32_t bits = (uint32_t)expected.state | ((uint32_t)expected.target << 16); + for (int bit = 0; bit < 32; bit++) { + float value = (bits & (1u << bit)) ? 1.0f : -1.0f; + EXPECT_EQ(obs_bits(actual[bit]), obs_bits(__float2bfloat16(value))); + } + float timer = expected.step_count / (float)expected.max_steps; + EXPECT_EQ(obs_bits(actual[TIMER_INDEX]), + obs_bits(__float2bfloat16(timer))); +} + +static void fill_kwargs(Dict* kwargs, int seed, int step_grace, + int perf_weighting) { + std::memset(kwargs, 0, sizeof(*kwargs)); + dict_set(kwargs, "seed", seed); + dict_set(kwargs, "start_depth", 2); + dict_set(kwargs, "max_depth", 16); + dict_set(kwargs, "step_grace", step_grace); + dict_set(kwargs, "perf_weighting", perf_weighting); +} + +static void test_deterministic_reset_and_step_parity() { + constexpr int n = 257; + constexpr int seed = 42; + constexpr int step_grace = 2; + constexpr int perf_weighting = PERF_WEIGHTING_QUADRATIC; + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, (size_t)n * OBS_SIZE * sizeof(obs_t)), "cudaMalloc observations"); + check_cuda(cudaMalloc(&actions, (size_t)n * sizeof(float)), "cudaMalloc actions"); + check_cuda(cudaMalloc(&rewards, (size_t)n * sizeof(float)), "cudaMalloc rewards"); + check_cuda(cudaMalloc(&terminals, (size_t)n * sizeof(float)), "cudaMalloc terminals"); + + Dict kwargs; + fill_kwargs(&kwargs, seed, step_grace, perf_weighting); + Env* envs = puf_vec_create(n, &kwargs, observations, actions, rewards, terminals); + EXPECT_TRUE(envs != nullptr); + puf_reset(envs); + check_cuda(cudaDeviceSynchronize(), "initial reset"); + + std::vector oracle(n); + unsigned int running_seed = seed; + for (int i = 0; i < n; i++) { + oracle[i].rng = (uint32_t)rand_r(&running_seed); + oracle[i].curriculum_depth = 2; + oracle_reset_state(&oracle[i], &table, step_grace); + } + + std::vector states(n); + std::vector host_envs(n); + std::vector host_obs((size_t)n * OBS_SIZE); + std::vector host_rewards(n), host_terminals(n), host_actions(n); + check_cuda(cudaMemcpy(states.data(), g_gpu.states, + n * sizeof(GpuAffineLockState), cudaMemcpyDeviceToHost), "copy reset states"); + check_cuda(cudaMemcpy(host_obs.data(), observations, + host_obs.size() * sizeof(obs_t), cudaMemcpyDeviceToHost), "copy reset observations"); + check_cuda(cudaMemcpy(host_rewards.data(), rewards, + n * sizeof(float), cudaMemcpyDeviceToHost), "copy reset rewards"); + check_cuda(cudaMemcpy(host_terminals.data(), terminals, + n * sizeof(float), cudaMemcpyDeviceToHost), "copy reset terminals"); + for (int i = 0; i < n; i++) { + expect_state_equal(states[i], oracle[i]); + expect_observation_equal(&host_obs[(size_t)i * OBS_SIZE], oracle[i]); + EXPECT_NEAR(host_rewards[i], 0.0f, 0.0f); + EXPECT_NEAR(host_terminals[i], 0.0f, 0.0f); + } + + for (int step = 0; step < 96; step++) { + for (int i = 0; i < n; i++) { + int selector = (step * 17 + i * 13) % 41; + if (selector == 0) host_actions[i] = std::numeric_limits::quiet_NaN(); + else if (selector == 1) host_actions[i] = -0.25f; + else if (selector == 2) host_actions[i] = 8.0f; + else host_actions[i] = (float)((step + 3 * i) & 7) + (selector == 3 ? 0.75f : 0.0f); + } + check_cuda(cudaMemcpy(actions, host_actions.data(), + n * sizeof(float), cudaMemcpyHostToDevice), "copy actions"); + puf_step(envs); + check_cuda(cudaDeviceSynchronize(), "step"); + + for (int i = 0; i < n; i++) { + oracle_step(&oracle[i], host_actions[i], &table, + 2, 16, step_grace, perf_weighting, + &host_rewards[i], &host_terminals[i]); + } + + check_cuda(cudaMemcpy(states.data(), g_gpu.states, + n * sizeof(GpuAffineLockState), cudaMemcpyDeviceToHost), "copy states"); + check_cuda(cudaMemcpy(host_envs.data(), envs, + n * sizeof(Env), cudaMemcpyDeviceToHost), "copy logs"); + check_cuda(cudaMemcpy(host_obs.data(), observations, + host_obs.size() * sizeof(obs_t), cudaMemcpyDeviceToHost), "copy observations"); + std::vector actual_rewards(n), actual_terminals(n); + check_cuda(cudaMemcpy(actual_rewards.data(), rewards, + n * sizeof(float), cudaMemcpyDeviceToHost), "copy rewards"); + check_cuda(cudaMemcpy(actual_terminals.data(), terminals, + n * sizeof(float), cudaMemcpyDeviceToHost), "copy terminals"); + for (int i = 0; i < n; i++) { + expect_state_equal(states[i], oracle[i]); + expect_log_equal(host_envs[i].log, oracle[i].log); + expect_observation_equal(&host_obs[(size_t)i * OBS_SIZE], oracle[i]); + EXPECT_NEAR(actual_rewards[i], host_rewards[i], 0.0f); + EXPECT_NEAR(actual_terminals[i], host_terminals[i], 0.0f); + } + } + + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree observations"); + check_cuda(cudaFree(actions), "cudaFree actions"); + check_cuda(cudaFree(rewards), "cudaFree rewards"); + check_cuda(cudaFree(terminals), "cudaFree terminals"); +} + +static void test_reset_rejection_sampling() { + constexpr uint32_t rejection_seed = 24481u; + constexpr uint32_t expected_final_rng = 3424986747u; + static const int depths[] = {2, 16}; + static const int expected_choices[] = {44338, 15778}; + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, OBS_SIZE * sizeof(obs_t)), + "cudaMalloc rejection observations"); + check_cuda(cudaMalloc(&actions, sizeof(float)), + "cudaMalloc rejection action"); + check_cuda(cudaMalloc(&rewards, sizeof(float)), + "cudaMalloc rejection reward"); + check_cuda(cudaMalloc(&terminals, sizeof(float)), + "cudaMalloc rejection terminal"); + + Dict kwargs; + fill_kwargs(&kwargs, 1, 0, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(1, &kwargs, + observations, actions, rewards, terminals); + + for (int case_index = 0; case_index < 2; case_index++) { + int depth = depths[case_index]; + GpuAffineLockState injected = {}; + injected.rng = rejection_seed; + injected.curriculum_depth = depth; + check_cuda(cudaMemcpy(g_gpu.states, &injected, sizeof(injected), + cudaMemcpyHostToDevice), "inject rejection state"); + + puf_reset(envs); + check_cuda(cudaDeviceSynchronize(), "rejection reset"); + + OracleState expected = {}; + expected.rng = rejection_seed; + expected.curriculum_depth = depth; + oracle_reset_state(&expected, &table, 0); + EXPECT_EQ(expected.rng, expected_final_rng); + const VisibleTargetDepth* table_depth = oracle_depth(&table, depth); + EXPECT_TRUE(table_depth != nullptr); + const VisibleTargetRecord* selected = &table.records[ + table_depth->first_record + (uint32_t)expected_choices[case_index]]; + EXPECT_EQ(expected.state, selected->start); + EXPECT_EQ(expected.target, selected->target); + + GpuAffineLockState actual = {}; + obs_t actual_obs[OBS_SIZE]; + float actual_reward = 123.0f; + float actual_terminal = 123.0f; + check_cuda(cudaMemcpy(&actual, g_gpu.states, sizeof(actual), + cudaMemcpyDeviceToHost), "copy rejection state"); + check_cuda(cudaMemcpy(actual_obs, observations, sizeof(actual_obs), + cudaMemcpyDeviceToHost), "copy rejection observations"); + check_cuda(cudaMemcpy(&actual_reward, rewards, sizeof(actual_reward), + cudaMemcpyDeviceToHost), "copy rejection reward"); + check_cuda(cudaMemcpy(&actual_terminal, terminals, sizeof(actual_terminal), + cudaMemcpyDeviceToHost), "copy rejection terminal"); + expect_state_equal(actual, expected); + expect_observation_equal(actual_obs, expected); + EXPECT_EQ(float_bits(actual_reward), float_bits(0.0f)); + EXPECT_EQ(float_bits(actual_terminal), float_bits(0.0f)); + } + + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree rejection observations"); + check_cuda(cudaFree(actions), "cudaFree rejection action"); + check_cuda(cudaFree(rewards), "cudaFree rejection reward"); + check_cuda(cudaFree(terminals), "cudaFree rejection terminal"); +} + +static void test_puf_log_exports_cpu_contract() { + Log log = {}; + log.perf = 1.25f; + log.score = 2.5f; + log.solve_rate = 2.0f; + log.max_depth_solve = 1.0f; + log.episode_return = 3.5f; + log.episode_length = 8.0f; + log.solve_steps = 5.0f; + log.timeout_rate = 1.0f; + log.solve_efficiency = 1.75f; + log.target_distance = 20.0f; + log.solved_target_distance = 12.0f; + log.d6_rate = 2.0f; + log.d6_solve_rate = 1.0f; + log.d8_rate = 4.0f; + log.d8_solve_rate = 3.0f; + log.d16_rate = 1.0f; + log.d16_solve_rate = 1.0f; + log.n = 3.0f; + Dict out = {}; + puf_log(&log, &out); + EXPECT_EQ(out.size, 15); + EXPECT_NEAR(dict_get(&out, "perf"), 1.25f, 0.0f); + EXPECT_NEAR(dict_get(&out, "score"), 2.5f, 0.0f); + EXPECT_NEAR(dict_get(&out, "solve_rate"), 2.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "max_depth_solve"), 1.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "episode_return"), 3.5f, 0.0f); + EXPECT_NEAR(dict_get(&out, "episode_length"), 8.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "timeout_rate"), 1.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "min_win_moves"), 20.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "solved_min_win_moves"), 6.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "conditional_solve_steps"), 2.5f, 0.0f); + EXPECT_NEAR(dict_get(&out, "conditional_solve_efficiency"), 0.875f, 0.0f); + EXPECT_NEAR(dict_get(&out, "d6_solve_rate"), 0.5f, 0.0f); + EXPECT_NEAR(dict_get(&out, "d8_solve_rate"), 0.75f, 0.0f); + EXPECT_NEAR(dict_get(&out, "d16_solve_rate"), 1.0f, 0.0f); + EXPECT_NEAR(dict_get(&out, "n"), 3.0f, 0.0f); + dict_clear(&out); + + Log zero_denominators = {}; + zero_denominators.solved_target_distance = 12.0f; + zero_denominators.solve_steps = 5.0f; + zero_denominators.solve_efficiency = 1.75f; + zero_denominators.d6_solve_rate = 1.0f; + zero_denominators.d8_solve_rate = 1.0f; + zero_denominators.d16_solve_rate = 1.0f; + Dict zero_out = {}; + puf_log(&zero_denominators, &zero_out); + EXPECT_EQ(zero_out.size, 15); + EXPECT_NEAR(dict_get(&zero_out, "solved_min_win_moves"), 0.0f, 0.0f); + EXPECT_NEAR(dict_get(&zero_out, "conditional_solve_steps"), 0.0f, 0.0f); + EXPECT_NEAR(dict_get(&zero_out, "conditional_solve_efficiency"), 0.0f, 0.0f); + EXPECT_NEAR(dict_get(&zero_out, "d6_solve_rate"), 0.0f, 0.0f); + EXPECT_NEAR(dict_get(&zero_out, "d8_solve_rate"), 0.0f, 0.0f); + EXPECT_NEAR(dict_get(&zero_out, "d16_solve_rate"), 0.0f, 0.0f); + dict_clear(&zero_out); +} + +static void test_exhaustive_action_transforms() { + constexpr int n = 1 << BITS; + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, (size_t)n * OBS_SIZE * sizeof(obs_t)), + "cudaMalloc exhaustive observations"); + check_cuda(cudaMalloc(&actions, (size_t)n * sizeof(float)), + "cudaMalloc exhaustive actions"); + check_cuda(cudaMalloc(&rewards, (size_t)n * sizeof(float)), + "cudaMalloc exhaustive rewards"); + check_cuda(cudaMalloc(&terminals, (size_t)n * sizeof(float)), + "cudaMalloc exhaustive terminals"); + + Dict kwargs; + fill_kwargs(&kwargs, 7, 100, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(n, &kwargs, + observations, actions, rewards, terminals); + std::vector states(n); + std::vector host_actions(n), host_rewards(n), host_terminals(n); + + for (int action = 0; action < NUM_ACTIONS; action++) { + for (int value = 0; value < n; value++) { + uint16_t expected = oracle_apply_action((uint16_t)value, action); + states[value] = {}; + states[value].rng = (uint32_t)(value + 1); + states[value].state = (uint16_t)value; + states[value].target = expected ^ 1u; + states[value].max_steps = 100; + states[value].scramble_depth = 16; + states[value].curriculum_depth = 16; + states[value].target_distance = 16; + host_actions[value] = (float)action; + } + check_cuda(cudaMemcpy(g_gpu.states, states.data(), + n * sizeof(GpuAffineLockState), cudaMemcpyHostToDevice), + "copy exhaustive states"); + check_cuda(cudaMemcpy(actions, host_actions.data(), + n * sizeof(float), cudaMemcpyHostToDevice), + "copy exhaustive actions"); + puf_step(envs); + check_cuda(cudaDeviceSynchronize(), "exhaustive action step"); + check_cuda(cudaMemcpy(states.data(), g_gpu.states, + n * sizeof(GpuAffineLockState), cudaMemcpyDeviceToHost), + "copy exhaustive results"); + check_cuda(cudaMemcpy(host_rewards.data(), rewards, + n * sizeof(float), cudaMemcpyDeviceToHost), + "copy exhaustive rewards"); + check_cuda(cudaMemcpy(host_terminals.data(), terminals, + n * sizeof(float), cudaMemcpyDeviceToHost), + "copy exhaustive terminals"); + for (int value = 0; value < n; value++) { + EXPECT_EQ(states[value].state, + oracle_apply_action((uint16_t)value, action)); + EXPECT_EQ(states[value].step_count, 1); + EXPECT_NEAR(states[value].episode_return, STEP_REWARD, 0.0f); + EXPECT_NEAR(host_rewards[value], STEP_REWARD, 0.0f); + EXPECT_NEAR(host_terminals[value], 0.0f, 0.0f); + } + } + + puf_close(envs); + dict_clear(&kwargs); + check_cuda(cudaFree(observations), "cudaFree exhaustive observations"); + check_cuda(cudaFree(actions), "cudaFree exhaustive actions"); + check_cuda(cudaFree(rewards), "cudaFree exhaustive rewards"); + check_cuda(cudaFree(terminals), "cudaFree exhaustive terminals"); +} + +static void test_action_boundaries() { + const float infinity = std::numeric_limits::infinity(); + const float actions_under_test[] = { + -infinity, + -0.25f, + -std::numeric_limits::denorm_min(), + -0.0f, + 0.0f, + 0.5f, + 0.999f, + 1.0f, + 1.5f, + 6.999f, + 7.0f, + std::nextafter(7.0f, infinity), + 8.0f, + infinity, + std::numeric_limits::quiet_NaN(), + }; + constexpr int n = sizeof(actions_under_test) / sizeof(actions_under_test[0]); + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, (size_t)n * OBS_SIZE * sizeof(obs_t)), + "cudaMalloc boundary observations"); + check_cuda(cudaMalloc(&actions, (size_t)n * sizeof(float)), + "cudaMalloc boundary actions"); + check_cuda(cudaMalloc(&rewards, (size_t)n * sizeof(float)), + "cudaMalloc boundary rewards"); + check_cuda(cudaMalloc(&terminals, (size_t)n * sizeof(float)), + "cudaMalloc boundary terminals"); + + Dict kwargs; + fill_kwargs(&kwargs, 11, 100, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(n, &kwargs, + observations, actions, rewards, terminals); + + std::vector injected(n); + std::vector expected(n); + std::vector expected_rewards(n), expected_terminals(n); + for (int i = 0; i < n; i++) { + injected[i] = {}; + injected[i].rng = (uint32_t)(1000 + i); + injected[i].state = 0x1234u; + injected[i].target = 0xbeefu; + injected[i].max_steps = 100; + injected[i].scramble_depth = 2; + injected[i].curriculum_depth = 2; + injected[i].target_distance = 2; + + expected[i].rng = injected[i].rng; + expected[i].state = injected[i].state; + expected[i].target = injected[i].target; + expected[i].max_steps = injected[i].max_steps; + expected[i].scramble_depth = injected[i].scramble_depth; + expected[i].curriculum_depth = injected[i].curriculum_depth; + expected[i].target_distance = injected[i].target_distance; + oracle_step(&expected[i], actions_under_test[i], &table, + 2, 16, 100, PERF_WEIGHTING_LINEAR, + &expected_rewards[i], &expected_terminals[i]); + } + check_cuda(cudaMemcpy(g_gpu.states, injected.data(), + (size_t)n * sizeof(GpuAffineLockState), cudaMemcpyHostToDevice), + "copy boundary states"); + check_cuda(cudaMemcpy(actions, actions_under_test, + sizeof(actions_under_test), cudaMemcpyHostToDevice), + "copy boundary actions"); + puf_step(envs); + check_cuda(cudaDeviceSynchronize(), "boundary step"); + + std::vector actual_states(n); + std::vector actual_envs(n); + std::vector actual_obs((size_t)n * OBS_SIZE); + std::vector actual_rewards(n), actual_terminals(n); + check_cuda(cudaMemcpy(actual_states.data(), g_gpu.states, + (size_t)n * sizeof(GpuAffineLockState), cudaMemcpyDeviceToHost), + "copy boundary results"); + check_cuda(cudaMemcpy(actual_envs.data(), envs, + (size_t)n * sizeof(Env), cudaMemcpyDeviceToHost), + "copy boundary logs"); + check_cuda(cudaMemcpy(actual_obs.data(), observations, + actual_obs.size() * sizeof(obs_t), cudaMemcpyDeviceToHost), + "copy boundary observations"); + check_cuda(cudaMemcpy(actual_rewards.data(), rewards, + (size_t)n * sizeof(float), cudaMemcpyDeviceToHost), + "copy boundary rewards"); + check_cuda(cudaMemcpy(actual_terminals.data(), terminals, + (size_t)n * sizeof(float), cudaMemcpyDeviceToHost), + "copy boundary terminals"); + + for (int i = 0; i < n; i++) { + bool invalid = !std::isfinite(actions_under_test[i]) + || actions_under_test[i] < 0.0f || actions_under_test[i] > 7.0f; + expect_state_equal(actual_states[i], expected[i]); + expect_log_equal(actual_envs[i].log, expected[i].log); + expect_observation_equal( + &actual_obs[(size_t)i * OBS_SIZE], expected[i]); + EXPECT_EQ(float_bits(actual_rewards[i]), + float_bits(expected_rewards[i])); + EXPECT_EQ(float_bits(actual_terminals[i]), + float_bits(expected_terminals[i])); + EXPECT_EQ(float_bits(actual_terminals[i]), + float_bits(invalid ? 1.0f : 0.0f)); + EXPECT_EQ(float_bits(actual_envs[i].log.n), + float_bits(invalid ? 1.0f : 0.0f)); + } + + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree boundary observations"); + check_cuda(cudaFree(actions), "cudaFree boundary actions"); + check_cuda(cudaFree(rewards), "cudaFree boundary rewards"); + check_cuda(cudaFree(terminals), "cudaFree boundary terminals"); +} + +static const VisibleTargetRecord* find_solution_record( + const VisibleTargetTable* table, const OracleState& state) { + const VisibleTargetDepth* depth = oracle_depth(table, state.scramble_depth); + EXPECT_TRUE(depth != nullptr); + for (uint32_t i = 0; i < depth->stored_count; i++) { + const VisibleTargetRecord* record = + &table->records[depth->first_record + i]; + if (record->start == state.state && record->target == state.target) { + return record; + } + } + return nullptr; +} + +static void test_solution_curriculum_and_logs() { + constexpr int seed = 69; + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, OBS_SIZE * sizeof(obs_t)), + "cudaMalloc solution observations"); + check_cuda(cudaMalloc(&actions, sizeof(float)), "cudaMalloc solution action"); + check_cuda(cudaMalloc(&rewards, sizeof(float)), "cudaMalloc solution reward"); + check_cuda(cudaMalloc(&terminals, sizeof(float)), "cudaMalloc solution terminal"); + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + Dict kwargs; + fill_kwargs(&kwargs, seed, 0, PERF_WEIGHTING_QUADRATIC); + Env* envs = puf_vec_create(1, &kwargs, + observations, actions, rewards, terminals); + puf_reset(envs); + check_cuda(cudaDeviceSynchronize(), "solution reset"); + + OracleState oracle = {}; + unsigned int running_seed = seed; + oracle.rng = (uint32_t)rand_r(&running_seed); + oracle.curriculum_depth = 2; + oracle_reset_state(&oracle, &table, 0); + static const int expected_depths[] = {2, 4, 5, 6, 8, 16}; + for (int expected_depth : expected_depths) { + EXPECT_EQ(oracle.scramble_depth, expected_depth); + const VisibleTargetRecord* record = find_solution_record(&table, oracle); + EXPECT_TRUE(record != nullptr); + for (int move = 0; move < record->solution_length; move++) { + float action = (float)((record->packed_actions >> (3 * move)) & 7u); + check_cuda(cudaMemcpy(actions, &action, sizeof(float), + cudaMemcpyHostToDevice), "copy solution action"); + puf_step(envs); + check_cuda(cudaDeviceSynchronize(), "solution step"); + + float expected_reward = 0.0f; + float expected_terminal = 0.0f; + oracle_step(&oracle, action, &table, 2, 16, 0, + PERF_WEIGHTING_QUADRATIC, + &expected_reward, &expected_terminal); + GpuAffineLockState actual_state; + Env actual_env; + obs_t actual_obs[OBS_SIZE]; + float actual_reward = 0.0f; + float actual_terminal = 0.0f; + check_cuda(cudaMemcpy(&actual_state, g_gpu.states, + sizeof(actual_state), cudaMemcpyDeviceToHost), + "copy solution state"); + check_cuda(cudaMemcpy(&actual_env, envs, + sizeof(actual_env), cudaMemcpyDeviceToHost), + "copy solution log"); + check_cuda(cudaMemcpy(actual_obs, observations, + sizeof(actual_obs), cudaMemcpyDeviceToHost), + "copy solution observations"); + check_cuda(cudaMemcpy(&actual_reward, rewards, + sizeof(float), cudaMemcpyDeviceToHost), + "copy solution reward"); + check_cuda(cudaMemcpy(&actual_terminal, terminals, + sizeof(float), cudaMemcpyDeviceToHost), + "copy solution terminal"); + expect_state_equal(actual_state, oracle); + expect_log_equal(actual_env.log, oracle.log); + expect_observation_equal(actual_obs, oracle); + EXPECT_NEAR(actual_reward, expected_reward, 0.0f); + EXPECT_NEAR(actual_terminal, expected_terminal, 0.0f); + EXPECT_NEAR(actual_terminal, + move + 1 == record->solution_length ? 1.0f : 0.0f, 0.0f); + } + } + EXPECT_NEAR(oracle.log.solve_rate, 6.0f, 0.0f); + EXPECT_NEAR(oracle.log.d6_solve_rate, 1.0f, 0.0f); + EXPECT_NEAR(oracle.log.d8_solve_rate, 1.0f, 0.0f); + EXPECT_NEAR(oracle.log.d16_solve_rate, 1.0f, 0.0f); + EXPECT_NEAR(oracle.log.max_depth_solve, 1.0f, 0.0f); + + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree solution observations"); + check_cuda(cudaFree(actions), "cudaFree solution action"); + check_cuda(cudaFree(rewards), "cudaFree solution reward"); + check_cuda(cudaFree(terminals), "cudaFree solution terminal"); +} + +static void test_linear_solve_scoring() { + constexpr int seed = 17; + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, OBS_SIZE * sizeof(obs_t)), + "cudaMalloc linear observations"); + check_cuda(cudaMalloc(&actions, sizeof(float)), + "cudaMalloc linear action"); + check_cuda(cudaMalloc(&rewards, sizeof(float)), + "cudaMalloc linear reward"); + check_cuda(cudaMalloc(&terminals, sizeof(float)), + "cudaMalloc linear terminal"); + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + Dict kwargs; + fill_kwargs(&kwargs, seed, 0, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(1, &kwargs, + observations, actions, rewards, terminals); + puf_reset(envs); + check_cuda(cudaDeviceSynchronize(), "linear reset"); + + OracleState oracle = {}; + unsigned int running_seed = seed; + oracle.rng = (uint32_t)rand_r(&running_seed); + oracle.curriculum_depth = 2; + oracle_reset_state(&oracle, &table, 0); + const VisibleTargetRecord* record = find_solution_record(&table, oracle); + EXPECT_TRUE(record != nullptr); + for (int move = 0; move < record->solution_length; move++) { + float action = (float)((record->packed_actions >> (3 * move)) & 7u); + check_cuda(cudaMemcpy(actions, &action, sizeof(action), + cudaMemcpyHostToDevice), "copy linear solution action"); + puf_step(envs); + } + check_cuda(cudaDeviceSynchronize(), "linear solution"); + + Env actual_env = {}; + float actual_reward = 0.0f; + float actual_terminal = 0.0f; + check_cuda(cudaMemcpy(&actual_env, envs, sizeof(actual_env), + cudaMemcpyDeviceToHost), "copy linear log"); + check_cuda(cudaMemcpy(&actual_reward, rewards, sizeof(actual_reward), + cudaMemcpyDeviceToHost), "copy linear reward"); + check_cuda(cudaMemcpy(&actual_terminal, terminals, sizeof(actual_terminal), + cudaMemcpyDeviceToHost), "copy linear terminal"); + EXPECT_NEAR(actual_env.log.perf, 2.0f / 16.0f, 0.0f); + EXPECT_NEAR(actual_env.log.score, 2.0f / 16.0f, 0.0f); + EXPECT_NEAR(actual_env.log.solve_rate, 1.0f, 0.0f); + EXPECT_NEAR(actual_env.log.n, 1.0f, 0.0f); + EXPECT_NEAR(actual_reward, 1.0f, 0.0f); + EXPECT_NEAR(actual_terminal, 1.0f, 0.0f); + + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree linear observations"); + check_cuda(cudaFree(actions), "cudaFree linear action"); + check_cuda(cudaFree(rewards), "cudaFree linear reward"); + check_cuda(cudaFree(terminals), "cudaFree linear terminal"); +} + +static void expect_device_canaries(const unsigned char* device_storage, + size_t prefix_bytes, size_t payload_bytes, size_t suffix_bytes, + unsigned char canary, const char* label) { + size_t total_bytes = prefix_bytes + payload_bytes + suffix_bytes; + std::vector host_storage(total_bytes); + check_cuda(cudaMemcpy(host_storage.data(), device_storage, total_bytes, + cudaMemcpyDeviceToHost), "copy canary storage"); + for (size_t i = 0; i < prefix_bytes; i++) { + if (host_storage[i] != canary) { + std::fprintf(stderr, + "%s prefix canary overwritten at byte %zu: 0x%02x != 0x%02x\n", + label, i, host_storage[i], canary); + std::exit(1); + } + } + size_t suffix_start = prefix_bytes + payload_bytes; + for (size_t i = suffix_start; i < total_bytes; i++) { + if (host_storage[i] != canary) { + std::fprintf(stderr, + "%s suffix canary overwritten at byte %zu: 0x%02x != 0x%02x\n", + label, i - suffix_start, host_storage[i], canary); + std::exit(1); + } + } +} + +static void run_io_canary_case(int n) { + constexpr size_t guard_bytes = 64; + constexpr size_t observation_prefix = guard_bytes + sizeof(obs_t); + constexpr unsigned char canary = 0xa5u; + size_t observation_bytes = (size_t)n * OBS_SIZE * sizeof(obs_t); + size_t scalar_bytes = (size_t)n * sizeof(float); + size_t observation_storage_bytes = observation_prefix + + observation_bytes + guard_bytes; + size_t scalar_storage_bytes = guard_bytes + scalar_bytes + guard_bytes; + + unsigned char* observation_storage = nullptr; + unsigned char* reward_storage = nullptr; + unsigned char* terminal_storage = nullptr; + float* actions = nullptr; + check_cuda(cudaMalloc(&observation_storage, observation_storage_bytes), + "cudaMalloc guarded observations"); + check_cuda(cudaMalloc(&reward_storage, scalar_storage_bytes), + "cudaMalloc guarded rewards"); + check_cuda(cudaMalloc(&terminal_storage, scalar_storage_bytes), + "cudaMalloc guarded terminals"); + check_cuda(cudaMalloc(&actions, scalar_bytes), + "cudaMalloc guarded actions"); + + obs_t* observations = reinterpret_cast( + observation_storage + observation_prefix); + float* rewards = reinterpret_cast(reward_storage + guard_bytes); + float* terminals = reinterpret_cast(terminal_storage + guard_bytes); + EXPECT_EQ((uintptr_t)observations & 1u, 0u); + EXPECT_EQ((uintptr_t)observations & 3u, 2u); + EXPECT_EQ((uintptr_t)rewards & 3u, 0u); + EXPECT_EQ((uintptr_t)terminals & 3u, 0u); + + check_cuda(cudaMemset(observation_storage, canary, + observation_storage_bytes), "initialize observation canaries"); + check_cuda(cudaMemset(reward_storage, canary, scalar_storage_bytes), + "initialize reward canaries"); + check_cuda(cudaMemset(terminal_storage, canary, scalar_storage_bytes), + "initialize terminal canaries"); + check_cuda(cudaMemset(actions, 0, scalar_bytes), + "initialize guarded actions"); + + Dict kwargs; + fill_kwargs(&kwargs, 31 + n, 1, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(n, &kwargs, + observations, actions, rewards, terminals); + puf_reset(envs); + check_cuda(cudaDeviceSynchronize(), "guarded reset"); + expect_device_canaries(observation_storage, observation_prefix, + observation_bytes, guard_bytes, canary, "reset observations"); + expect_device_canaries(reward_storage, guard_bytes, + scalar_bytes, guard_bytes, canary, "reset rewards"); + expect_device_canaries(terminal_storage, guard_bytes, + scalar_bytes, guard_bytes, canary, "reset terminals"); + + check_cuda(cudaMemset(observation_storage, canary, + observation_storage_bytes), "reinitialize observation canaries"); + check_cuda(cudaMemset(reward_storage, canary, scalar_storage_bytes), + "reinitialize reward canaries"); + check_cuda(cudaMemset(terminal_storage, canary, scalar_storage_bytes), + "reinitialize terminal canaries"); + puf_step(envs); + check_cuda(cudaDeviceSynchronize(), "guarded step"); + expect_device_canaries(observation_storage, observation_prefix, + observation_bytes, guard_bytes, canary, "step observations"); + expect_device_canaries(reward_storage, guard_bytes, + scalar_bytes, guard_bytes, canary, "step rewards"); + expect_device_canaries(terminal_storage, guard_bytes, + scalar_bytes, guard_bytes, canary, "step terminals"); + + puf_close(envs); + dict_clear(&kwargs); + check_cuda(cudaFree(observation_storage), + "cudaFree guarded observations"); + check_cuda(cudaFree(reward_storage), "cudaFree guarded rewards"); + check_cuda(cudaFree(terminal_storage), "cudaFree guarded terminals"); + check_cuda(cudaFree(actions), "cudaFree guarded actions"); +} + +static void test_io_canaries_and_observation_alignment() { +#if AFFINE_LOCK_GPU_SHARED_OBS + constexpr int environments_per_block = AFFINE_LOCK_GPU_SHARED_BLOCK; +#else + constexpr int environments_per_block = + AFFINE_LOCK_GPU_BLOCK / AFFINE_LOCK_GPU_LANES; +#endif + run_io_canary_case(environments_per_block); + run_io_canary_case(environments_per_block + 1); +} + +static void test_nondefault_stream_and_cuda_graph() { + constexpr int n = 4099; + obs_t* observations = nullptr; + float* actions = nullptr; + float* rewards = nullptr; + float* terminals = nullptr; + check_cuda(cudaMalloc(&observations, (size_t)n * OBS_SIZE * sizeof(obs_t)), + "cudaMalloc graph observations"); + check_cuda(cudaMalloc(&actions, (size_t)n * sizeof(float)), + "cudaMalloc graph actions"); + check_cuda(cudaMalloc(&rewards, (size_t)n * sizeof(float)), + "cudaMalloc graph rewards"); + check_cuda(cudaMalloc(&terminals, (size_t)n * sizeof(float)), + "cudaMalloc graph terminals"); + check_cuda(cudaMemset(actions, 0, (size_t)n * sizeof(float)), + "clear graph actions"); + + Dict kwargs; + fill_kwargs(&kwargs, 123, 3, PERF_WEIGHTING_LINEAR); + Env* envs = puf_vec_create(n, &kwargs, + observations, actions, rewards, terminals); + cudaStream_t stream; + check_cuda(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), + "create graph stream"); + puf_bind_stream(stream); + puf_reset(envs); + check_cuda(cudaStreamSynchronize(stream), "graph reset"); + + VisibleTargetTable table = {}; + EXPECT_EQ(visible_targets_load(VISIBLE_TARGET_TABLE_PATH, + VISIBLE_TARGET_8ACTION_V1_HASH, &table), 0); + std::vector oracle(n); + unsigned int running_seed = 123; + for (int i = 0; i < n; i++) { + oracle[i].rng = (uint32_t)rand_r(&running_seed); + oracle[i].curriculum_depth = 2; + oracle_reset_state(&oracle[i], &table, 3); + } + + GpuAffineLockState* actual_states = nullptr; + Env* actual_envs = nullptr; + obs_t* actual_obs = nullptr; + float* actual_rewards = nullptr; + float* actual_terminals = nullptr; + check_cuda(cudaMallocHost((void**)&actual_states, + (size_t)n * sizeof(GpuAffineLockState)), + "cudaMallocHost graph states"); + check_cuda(cudaMallocHost((void**)&actual_envs, + (size_t)n * sizeof(Env)), "cudaMallocHost graph envs"); + check_cuda(cudaMallocHost((void**)&actual_obs, + (size_t)n * OBS_SIZE * sizeof(obs_t)), + "cudaMallocHost graph observations"); + check_cuda(cudaMallocHost((void**)&actual_rewards, + (size_t)n * sizeof(float)), "cudaMallocHost graph rewards"); + check_cuda(cudaMallocHost((void**)&actual_terminals, + (size_t)n * sizeof(float)), "cudaMallocHost graph terminals"); + + cudaGraph_t graph; + cudaGraphExec_t graph_exec; + check_cuda(cudaStreamBeginCapture(stream, cudaStreamCaptureModeThreadLocal), + "begin graph capture"); + for (int i = 0; i < 3; i++) { + puf_step(envs); + } + check_cuda(cudaStreamEndCapture(stream, &graph), "end graph capture"); + check_cuda(cudaGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0), + "instantiate graph"); + check_cuda(cudaGraphLaunch(graph_exec, stream), "launch graph first"); + check_cuda(cudaGraphLaunch(graph_exec, stream), "launch graph second"); + + check_cuda(cudaMemcpyAsync(actual_states, g_gpu.states, + (size_t)n * sizeof(GpuAffineLockState), cudaMemcpyDeviceToHost, stream), + "queue graph states D2H"); + check_cuda(cudaMemcpyAsync(actual_envs, envs, + (size_t)n * sizeof(Env), cudaMemcpyDeviceToHost, stream), + "queue graph logs D2H"); + check_cuda(cudaMemcpyAsync(actual_obs, observations, + (size_t)n * OBS_SIZE * sizeof(obs_t), cudaMemcpyDeviceToHost, stream), + "queue graph observations D2H"); + check_cuda(cudaMemcpyAsync(actual_rewards, rewards, + (size_t)n * sizeof(float), cudaMemcpyDeviceToHost, stream), + "queue graph rewards D2H"); + check_cuda(cudaMemcpyAsync(actual_terminals, terminals, + (size_t)n * sizeof(float), cudaMemcpyDeviceToHost, stream), + "queue graph terminals D2H"); + check_cuda(cudaStreamSynchronize(stream), "synchronize graph and D2H"); + + std::vector expected_rewards(n), expected_terminals(n); + for (int step = 0; step < 6; step++) { + for (int i = 0; i < n; i++) { + oracle_step(&oracle[i], 0.0f, &table, + 2, 16, 3, PERF_WEIGHTING_LINEAR, + &expected_rewards[i], &expected_terminals[i]); + } + } + for (int i = 0; i < n; i++) { + expect_state_equal(actual_states[i], oracle[i]); + expect_log_equal(actual_envs[i].log, oracle[i].log); + expect_observation_equal( + &actual_obs[(size_t)i * OBS_SIZE], oracle[i]); + EXPECT_EQ(float_bits(actual_rewards[i]), + float_bits(expected_rewards[i])); + EXPECT_EQ(float_bits(actual_terminals[i]), + float_bits(expected_terminals[i])); + } + + check_cuda(cudaGraphExecDestroy(graph_exec), "destroy graph exec"); + check_cuda(cudaGraphDestroy(graph), "destroy graph"); + check_cuda(cudaFreeHost(actual_states), "cudaFreeHost graph states"); + check_cuda(cudaFreeHost(actual_envs), "cudaFreeHost graph envs"); + check_cuda(cudaFreeHost(actual_obs), "cudaFreeHost graph observations"); + check_cuda(cudaFreeHost(actual_rewards), "cudaFreeHost graph rewards"); + check_cuda(cudaFreeHost(actual_terminals), "cudaFreeHost graph terminals"); + check_cuda(cudaStreamDestroy(stream), "destroy graph stream"); + puf_bind_stream(nullptr); + puf_close(envs); + dict_clear(&kwargs); + visible_targets_free(&table); + check_cuda(cudaFree(observations), "cudaFree graph observations"); + check_cuda(cudaFree(actions), "cudaFree graph actions"); + check_cuda(cudaFree(rewards), "cudaFree graph rewards"); + check_cuda(cudaFree(terminals), "cudaFree graph terminals"); +} + +int main() { + test_deterministic_reset_and_step_parity(); + test_reset_rejection_sampling(); + test_exhaustive_action_transforms(); + test_action_boundaries(); + test_solution_curriculum_and_logs(); + test_linear_solve_scoring(); + test_io_canaries_and_observation_alignment(); + test_nondefault_stream_and_cuda_graph(); + test_puf_log_exports_cpu_contract(); + std::puts("affine_lock CUDA tests passed"); + return 0; +} diff --git a/src/algo.cu b/src/algo.cu index fd97203d19..2827487ce8 100644 --- a/src/algo.cu +++ b/src/algo.cu @@ -1,3 +1,7 @@ +#include +#include +#include + // PufferNet model API + architecture // Writing custom nets in 4.0+ requires a fair bit of code because you are // responsible for defining your own activation and gradient buffers. @@ -85,6 +89,160 @@ thread_local void* g_cublas_dw_workspace = NULL; thread_local cudaStream_t g_dw_stream = NULL; thread_local cudaEvent_t g_dw_done = NULL; +// CUDA 13.1 recommends a 32 MiB cuBLAS workspace for Hopper sm90 and +// Blackwell sm10x/sm12x, which includes the local RTX 5090. It recommends +// 4 MiB for other architectures, and these recommendations can change with +// the toolkit version. Workspace size affects which cuBLAS algorithms are +// eligible, their performance, and potentially their bitwise outputs. Every +// concurrent lane therefore needs a distinct workspace; the eight-lane cap +// bounds these allocations to 256 MiB. cublasSetStream resets a user-provided +// workspace, so lane handles restore their workspace after binding their stream +// and never change streams afterward. This path is qualified only on RTX 5090 +// with CUDA 13.1; before portable enablement it should select workspace by +// architecture/toolkit or fall back to serial Muon. +// https://docs.nvidia.com/cuda/archive/13.1.0/cublas/index.html#cublassetworkspace +static constexpr size_t CUBLAS_WORKSPACE_BYTES = 32 * 1024 * 1024; + +#if CUDART_VERSION == 13010 && !defined(PRECISION_FLOAT) +struct PinnedLtPlan { + int M, N, K; + cublasOperation_t op_b; + cublasLtMatmulDesc_t op; + cublasLtMatrixLayout_t a_layout, b_layout, c_layout; + cublasLtMatmulAlgo_t algo; +}; + +thread_local cublasLtHandle_t g_cublaslt_handle = NULL; +thread_local PinnedLtPlan g_cublaslt_plans[4] = {}; +thread_local bool g_cublaslt_enabled = false; +thread_local bool g_cublaslt_muon_enabled = false; + +static bool cublaslt_make_row_layout(cublasLtMatrixLayout_t* layout, + uint64_t rows, uint64_t cols, int64_t ld) { + cublasLtOrder_t order = CUBLASLT_ORDER_ROW; + return cublasLtMatrixLayoutCreate( + layout, CUDA_R_16BF, rows, cols, ld) == CUBLAS_STATUS_SUCCESS + && cublasLtMatrixLayoutSetAttribute(*layout, + CUBLASLT_MATRIX_LAYOUT_ORDER, &order, + sizeof(order)) == CUBLAS_STATUS_SUCCESS; +} + +static bool cublaslt_init_plan(PinnedLtPlan* plan, + int M, int N, int K, cublasOperation_t op_b, + uint32_t tile = 15, uint32_t stages = 12) { + plan->M = M; + plan->N = N; + plan->K = K; + plan->op_b = op_b; + cublasOperation_t op_a = CUBLAS_OP_N; + if (cublasLtMatmulDescCreate( + &plan->op, CUBLAS_COMPUTE_32F, + CUDA_R_32F) != CUBLAS_STATUS_SUCCESS + || cublasLtMatmulDescSetAttribute(plan->op, + CUBLASLT_MATMUL_DESC_TRANSA, &op_a, + sizeof(op_a)) != CUBLAS_STATUS_SUCCESS + || cublasLtMatmulDescSetAttribute(plan->op, + CUBLASLT_MATMUL_DESC_TRANSB, &op_b, + sizeof(op_b)) != CUBLAS_STATUS_SUCCESS + || !cublaslt_make_row_layout( + &plan->a_layout, M, K, K) + || !cublaslt_make_row_layout(&plan->b_layout, + op_b == CUBLAS_OP_T ? N : K, + op_b == CUBLAS_OP_T ? K : N, + op_b == CUBLAS_OP_T ? K : N) + || !cublaslt_make_row_layout( + &plan->c_layout, M, N, N) + || cublasLtMatmulAlgoInit(g_cublaslt_handle, + CUBLAS_COMPUTE_32F, CUDA_R_32F, + CUDA_R_16BF, CUDA_R_16BF, + CUDA_R_16BF, CUDA_R_16BF, 21, + &plan->algo) != CUBLAS_STATUS_SUCCESS) { + return false; + } + + uint32_t split_k = 1; + uint32_t reduction = 0; + uint32_t swizzle = 0; + uint32_t custom = 0; + return cublasLtMatmulAlgoConfigSetAttribute(&plan->algo, + CUBLASLT_ALGO_CONFIG_TILE_ID, &tile, + sizeof(tile)) == CUBLAS_STATUS_SUCCESS + && cublasLtMatmulAlgoConfigSetAttribute(&plan->algo, + CUBLASLT_ALGO_CONFIG_SPLITK_NUM, &split_k, + sizeof(split_k)) == CUBLAS_STATUS_SUCCESS + && cublasLtMatmulAlgoConfigSetAttribute(&plan->algo, + CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME, &reduction, + sizeof(reduction)) == CUBLAS_STATUS_SUCCESS + && cublasLtMatmulAlgoConfigSetAttribute(&plan->algo, + CUBLASLT_ALGO_CONFIG_CTA_SWIZZLING, &swizzle, + sizeof(swizzle)) == CUBLAS_STATUS_SUCCESS + && cublasLtMatmulAlgoConfigSetAttribute(&plan->algo, + CUBLASLT_ALGO_CONFIG_CUSTOM_OPTION, &custom, + sizeof(custom)) == CUBLAS_STATUS_SUCCESS + && cublasLtMatmulAlgoConfigSetAttribute(&plan->algo, + CUBLASLT_ALGO_CONFIG_STAGES_ID, &stages, + sizeof(stages)) == CUBLAS_STATUS_SUCCESS; +} + +static void cublaslt_init_pinned() { + int runtime_version = 0; + int driver_version = 0; + int cublas_version = 0; + int device = 0; + cudaDeviceProp prop = {}; + if (cudaRuntimeGetVersion(&runtime_version) != cudaSuccess + || runtime_version != 13010 + || cudaDriverGetVersion(&driver_version) != cudaSuccess + || driver_version != 13000 + || cublasGetVersion( + g_cublas_handle, &cublas_version) != CUBLAS_STATUS_SUCCESS + || cublas_version != 130201 + || cudaGetDevice(&device) != cudaSuccess + || cudaGetDeviceProperties(&prop, device) != cudaSuccess + || prop.major != 12 || prop.minor != 0 + || prop.multiProcessorCount != 170 + || std::strcmp(prop.name, "NVIDIA GeForce RTX 5090") != 0 + || cublasLtCreate( + &g_cublaslt_handle) != CUBLAS_STATUS_SUCCESS) { + return; + } + g_cublaslt_enabled = + cublaslt_init_plan(&g_cublaslt_plans[0], + 8192, 1536, 512, CUBLAS_OP_T) + && cublaslt_init_plan(&g_cublaslt_plans[1], + 8192, 512, 1536, CUBLAS_OP_N); + g_cublaslt_muon_enabled = + cublaslt_init_plan(&g_cublaslt_plans[2], + 512, 512, 512, CUBLAS_OP_N, 11, 19) + && cublaslt_init_plan(&g_cublaslt_plans[3], + 1536, 512, 512, CUBLAS_OP_N, 18, 15); +} + +static bool cublaslt_try_pinned(cublasHandle_t handle, + cublasOperation_t op_a, cublasOperation_t op_b, + int M, int N, int K, void* A, void* B, void* C, + cudaStream_t stream, float alpha, float beta) { + uint32_t beta_bits = 0; + std::memcpy(&beta_bits, &beta, sizeof(beta_bits)); + if (!g_cublaslt_enabled || handle != g_cublas_handle + || op_a != CUBLAS_OP_N || alpha != 1.0f || beta_bits != 0 + || (((uintptr_t)A | (uintptr_t)B | (uintptr_t)C) & 255) != 0) { + return false; + } + for (int i = 0; i < 2; i++) { + PinnedLtPlan& plan = g_cublaslt_plans[i]; + if (plan.M == M && plan.N == N && plan.K == K + && plan.op_b == op_b) { + return cublasLtMatmul(g_cublaslt_handle, plan.op, + &alpha, A, plan.a_layout, B, plan.b_layout, + &beta, C, plan.c_layout, C, plan.c_layout, + &plan.algo, NULL, 0, stream) == CUBLAS_STATUS_SUCCESS; + } + } + return false; +} +#endif + static void cublas_init_one(cublasHandle_t* handle, void** workspace) { const size_t ws_bytes = 32 * 1024 * 1024; cublasCreate(handle); @@ -96,6 +254,9 @@ static void cublas_init_one(cublasHandle_t* handle, void** workspace) { void cublas_init_handle() { cublas_init_one(&g_cublas_handle, &g_cublas_workspace); cublas_init_one(&g_cublas_dw_handle, &g_cublas_dw_workspace); +#if CUDART_VERSION == 13010 && !defined(PRECISION_FLOAT) + cublaslt_init_pinned(); +#endif cudaStreamCreateWithFlags(&g_dw_stream, cudaStreamNonBlocking); cudaEventCreateWithFlags(&g_dw_done, cudaEventDisableTiming); cudaEventCreateWithFlags(&g_main_ready, cudaEventDisableTiming); @@ -105,10 +266,19 @@ void cublas_init_handle() { static void cublasGemmExDense(cublasHandle_t handle, cublasOperation_t op_a, cublasOperation_t op_b, int M, int N, int K, void* A, void* B, void* C, - cudaStream_t stream, float alpha = 1.0f, float beta = 0.0f) { + cudaStream_t stream, float alpha = 1.0f, float beta = 0.0f, + bool handle_bound_to_stream = false) { +#if CUDART_VERSION == 13010 && !defined(PRECISION_FLOAT) + if (cublaslt_try_pinned(handle, op_a, op_b, + M, N, K, A, B, C, stream, alpha, beta)) { + return; + } +#endif int lda = (op_a == CUBLAS_OP_N) ? K : M; int ldb = (op_b == CUBLAS_OP_N) ? N : K; - cublasSetStream(handle, stream); + if (!handle_bound_to_stream) { + cublasSetStream(handle, stream); + } cublasGemmEx(handle, op_b, op_a, N, M, K, &alpha, B, CUBLAS_PRECISION, ldb, A, CUBLAS_PRECISION, lda, &beta, C, CUBLAS_PRECISION, N, CUBLAS_COMPUTE, CUBLAS_GEMM_DEFAULT); @@ -116,33 +286,41 @@ static void cublasGemmExDense(cublasHandle_t handle, // out(...,N) = alpha * a(...,K) @ b(N,K)^T + beta * out — leading dims folded into M void puf_mm(Prec* a, Prec* b, Prec* out, cudaStream_t stream, - float alpha = 1.0f, float beta = 0.0f) { + float alpha = 1.0f, float beta = 0.0f, + cublasHandle_t handle = g_cublas_handle, + bool handle_bound_to_stream = false) { int M = batch_size(a->shape) * a->shape[ndim(a->shape)-2]; int K = a->shape[ndim(a->shape)-1]; int N = b->shape[ndim(b->shape)-2]; - cublasGemmExDense(g_cublas_handle, CUBLAS_OP_N, CUBLAS_OP_T, M, N, K, - a->data, b->data, out->data, stream, alpha, beta); + cublasGemmExDense(handle, CUBLAS_OP_N, CUBLAS_OP_T, M, N, K, + a->data, b->data, out->data, stream, alpha, beta, + handle_bound_to_stream); } // out(M,N) = alpha * a(...,M)^T @ b(...,N) + beta * out — leading dims folded into K void puf_mm_tn(Prec* a, Prec* b, Prec* out, cudaStream_t stream, float alpha = 1.0f, float beta = 0.0f, - cublasHandle_t handle = g_cublas_handle) { + cublasHandle_t handle = g_cublas_handle, + bool handle_bound_to_stream = false) { int M = a->shape[ndim(a->shape)-1]; int K = batch_size(a->shape) * a->shape[ndim(a->shape)-2]; int N = b->shape[ndim(b->shape)-1]; cublasGemmExDense(handle, CUBLAS_OP_T, CUBLAS_OP_N, M, N, K, - a->data, b->data, out->data, stream, alpha, beta); + a->data, b->data, out->data, stream, alpha, beta, + handle_bound_to_stream); } // out(...,N) = alpha * a(...,K) @ b(K,N) + beta * out — leading dims folded into M void puf_mm_nn(Prec* a, Prec* b, Prec* out, cudaStream_t stream, - float alpha = 1.0f, float beta = 0.0f) { + float alpha = 1.0f, float beta = 0.0f, + cublasHandle_t handle = g_cublas_handle, + bool handle_bound_to_stream = false) { int M = batch_size(a->shape) * a->shape[ndim(a->shape)-2]; int K = a->shape[ndim(a->shape)-1]; int N = b->shape[ndim(b->shape)-1]; - cublasGemmExDense(g_cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, M, N, K, - a->data, b->data, out->data, stream, alpha, beta); + cublasGemmExDense(handle, CUBLAS_OP_N, CUBLAS_OP_N, M, N, K, + a->data, b->data, out->data, stream, alpha, beta, + handle_bound_to_stream); } // Queue dW (mm_tn) on side stream once inputs are ready on main. Per-layer @@ -1062,6 +1240,29 @@ __global__ void muon_clip_nesterov(float* __restrict__ mb, } } +// Preserves the BF16 round-trip and bit-identical result while fusing work; +// measured about 0.5% higher SPS. +__global__ void muon_clip_nesterov_sum_sq_partials( + float* __restrict__ partials, float* __restrict__ mb, + precision_t* __restrict__ gc, const float* __restrict__ sum_sq_ptr, + float max_norm, float eps, float mu, int n) { + __shared__ float sdata[256]; + int tid = threadIdx.x; + float clip_coef = fminf(max_norm / (sqrtf(*sum_sq_ptr) + eps), 1.0f); + float sum = 0.0f; + for (int i = blockIdx.x * blockDim.x + tid; i < n; i += blockDim.x * gridDim.x) { + float g = to_float(gc[i]) * clip_coef; + float m = mu * mb[i] + g; + mb[i] = m; + precision_t update = from_float(g + mu * m); + gc[i] = update; + float v = to_float(update); + sum += v * v; + } + sdata[tid] = sum; + block_reduce_sum(sdata, &partials[blockIdx.x], tid, blockDim.x, 1); +} + // x *= 1 / max(sqrt(sum_sq), eps) — NS input normalize __global__ void muon_l2_normalize(precision_t* __restrict__ dst, const float* __restrict__ sum_sq_ptr, float eps, int n) { @@ -1083,13 +1284,20 @@ __global__ void muon_store_update(precision_t* __restrict__ dst, // wb = wb * (1 - lr*wd) - lr * update (update already scaled; one call for all params) __global__ void muon_weight_update(float* __restrict__ wb, + precision_t* __restrict__ model_weights, const precision_t* __restrict__ update, const float* __restrict__ lr_ptr, float wd, int n) { float lr = *lr_ptr; float wd_scale = 1.0f - lr * wd; int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < n) { - wb[idx] = wb[idx] * wd_scale - lr * to_float(update[idx]); + float new_weight = wb[idx] * wd_scale - lr * to_float(update[idx]); + wb[idx] = new_weight; + // Reuse the exact FP32 update to remove the later cast: about 0.17% SPS, + // bit-identical across the golden environments. + if (USE_BF16) { + model_weights[idx] = from_float(new_weight); + } } } @@ -1101,6 +1309,50 @@ constexpr double ns_coeffs[5][3] = { {2.8366, -3.0525, 1.2012}, }; +#if CUDART_VERSION == 13010 && !defined(PRECISION_FLOAT) +// H512 separate-C/D removes Newton-Schulz copies: about 0.8% Affine578 SPS, +// bit-identical across the golden environments. +static bool cublaslt_try_muon_nn(int M, int N, int K, + void* A, void* B, void* C, void* D, + cudaStream_t stream, int round, bool square) { + if (!g_cublaslt_muon_enabled + || (((uintptr_t)A | (uintptr_t)B + | (uintptr_t)C | (uintptr_t)D) & 255) != 0) { + return false; + } + float alpha = square ? (float)ns_coeffs[round][2] : 1.0f; + float beta = square + ? (float)ns_coeffs[round][1] : (float)ns_coeffs[round][0]; + PinnedLtPlan& plan = g_cublaslt_plans[square ? 2 : 3]; + if (plan.M != M || plan.N != N || plan.K != K) { + return false; + } + cublasLtMatmul(g_cublaslt_handle, plan.op, + &alpha, A, plan.a_layout, B, plan.b_layout, + &beta, C, plan.c_layout, D, plan.c_layout, + &plan.algo, NULL, 0, stream); + return true; +} +#endif + +struct MuonMatrix { + long offset; + long rows; + long cols; + long work; +}; + +struct MuonLane { + float* norm_partials; + precision_t* gram; + precision_t* gram_buf; + precision_t* x_buf; + cudaStream_t stream; + cudaEvent_t done; + cublasHandle_t cublas_handle; + void* cublas_workspace; +}; + // Muon optimizer. Our benchmarks show this is a major // upgrade over Adam (weight decay not needed in RL). struct Muon { @@ -1113,6 +1365,10 @@ struct Muon { Float mb; // flat momentum buffer (param-sized) Prec gram, gram_buf, x_buf; Allocator* param_alloc; + int num_lanes; + cudaEvent_t matrices_ready; + MuonMatrix matrices[8]; + MuonLane lanes[8]; }; void muon_init(Muon* m, Allocator* param_alloc, double momentum, Allocator* alloc) { @@ -1125,13 +1381,28 @@ void muon_init(Muon* m, Allocator* param_alloc, double momentum, Allocator* allo m->mb = {.shape = {param_alloc->total_elems}}; alloc_register(alloc, &m->mb); long max_M = 0, max_N = 0; + long offset = 0; + int num_matrices = 0; + int heavy_ns_lanes = 0; for (int _i = 0; _i < param_alloc->num_regs; _i++) { AllocEntry& e = param_alloc->regs[_i]; + long ne = numel(e.shape); if (ndim(e.shape) >= 2) { - long R = e.shape[0], C = numel(e.shape) / R; - max_M = max(max_M, min(R, C)); + long R = e.shape[0], C = ne / R; + long M = min(R, C); + if (num_matrices < 8) { + MuonMatrix& matrix = m->matrices[num_matrices]; + matrix.offset = offset; + matrix.rows = R; + matrix.cols = C; + matrix.work = ne * M; + } + num_matrices++; + heavy_ns_lanes += min(R, C) >= 1024 && max(R, C) >= 3072; + max_M = max(max_M, M); max_N = max(max_N, max(R, C)); } + offset += ne; } m->gram = {.shape = {max_M, max_M}}; m->gram_buf = {.shape = {max_M, max_M}}; @@ -1139,75 +1410,189 @@ void muon_init(Muon* m, Allocator* param_alloc, double momentum, Allocator* allo alloc_register(alloc, &m->gram); alloc_register(alloc, &m->gram_buf); alloc_register(alloc, &m->x_buf); + + bool saturated_workload = num_matrices >= 7 && heavy_ns_lanes >= 5; + m->num_lanes = num_matrices >= 2 && num_matrices <= 8 + && max_N <= 4096 && !saturated_workload ? num_matrices : 0; + if (m->num_lanes == 0) { + return; + } + + // Stable largest-work-first schedule. Equal-work matrices retain their + // parameter registration order. + for (int i = 1; i < m->num_lanes; i++) { + MuonMatrix key = m->matrices[i]; + int j = i; + while (j > 0 && m->matrices[j - 1].work < key.work) { + m->matrices[j] = m->matrices[j - 1]; + j--; + } + m->matrices[j] = key; + } + + cudaEventCreateWithFlags(&m->matrices_ready, cudaEventDisableTiming); + for (int i = 0; i < m->num_lanes; i++) { + MuonMatrix& matrix = m->matrices[i]; + MuonLane& lane = m->lanes[i]; + long M = min(matrix.rows, matrix.cols); + long ne = matrix.rows * matrix.cols; + cudaMalloc((void**)&lane.norm_partials, 257 * sizeof(float)); + cudaMalloc((void**)&lane.gram, M * M * sizeof(precision_t)); + cudaMalloc((void**)&lane.gram_buf, M * M * sizeof(precision_t)); + cudaMalloc((void**)&lane.x_buf, ne * sizeof(precision_t)); + cudaStreamCreateWithFlags(&lane.stream, cudaStreamNonBlocking); + cudaEventCreateWithFlags(&lane.done, cudaEventDisableTiming); + cublas_init_one(&lane.cublas_handle, &lane.cublas_workspace); + cublasSetStream(lane.cublas_handle, lane.stream); + cublasSetWorkspace(lane.cublas_handle, lane.cublas_workspace, + CUBLAS_WORKSPACE_BYTES); + } } -void muon_step(Muon* m, Float weights, Prec grads, +static void muon_matrix_step(precision_t* gc_ptr, long R, long C, + float* ns_norm, float* norm_partials, + precision_t* gram_storage, precision_t* gram_buf_storage, + precision_t* x_buf_storage, + cudaStream_t stream, cublasHandle_t handle, + bool handle_bound_to_stream, bool norm_partials_ready) { + long ne = R * C; + long M = min(R, C); + bool tall = R > C; + Prec x = {.data = gc_ptr, .shape = {R, C}}; + Prec x_buf = {.data = x_buf_storage, .shape = {R, C}}; + Prec gram = {.data = gram_storage, .shape = {M, M}}; + Prec gram_buf = {.data = gram_buf_storage, .shape = {M, M}}; + + int nblk = min((int)grid_size(ne), 256); + if (!norm_partials_ready) { + muon_sum_sq_partials<<>>( + norm_partials, x.data, (int)ne); + } + muon_sum_sq_reduce<<<1, 256, 0, stream>>>( + ns_norm, norm_partials, nblk); + muon_l2_normalize<<>>( + x.data, ns_norm, 1e-7f, (int)ne); + + // 5 steps land in x_buf. 4 = you break it. + for (int i = 0; i < 5; ++i) { + Prec& src = (i % 2 == 0) ? x : x_buf; + Prec& dst = (i % 2 == 0) ? x_buf : x; + if (tall) { + puf_mm_tn(&src, &src, &gram, stream, 1.0f, 0.0f, + handle, handle_bound_to_stream); + } else { + puf_mm(&src, &src, &gram, stream, 1.0f, 0.0f, + handle, handle_bound_to_stream); + } + bool lt_square = false; +#if CUDART_VERSION == 13010 && !defined(PRECISION_FLOAT) + if (handle_bound_to_stream) { + lt_square = cublaslt_try_muon_nn(M, M, M, + gram.data, gram.data, gram.data, gram_buf.data, + stream, i, true); + } +#endif + if (!lt_square) { + puf_copy(&gram_buf, &gram, stream); + puf_mm_nn(&gram, &gram, &gram_buf, stream, + (float)ns_coeffs[i][2], (float)ns_coeffs[i][1], + handle, handle_bound_to_stream); + } + bool lt_x = false; +#if CUDART_VERSION == 13010 && !defined(PRECISION_FLOAT) + if (handle_bound_to_stream && tall) { + lt_x = cublaslt_try_muon_nn(R, C, C, + src.data, gram_buf.data, src.data, dst.data, + stream, i, false); + } +#endif + if (!lt_x) { + puf_copy(&dst, &src, stream); + if (tall) { + puf_mm_nn(&src, &gram_buf, &dst, + stream, 1.0f, (float)ns_coeffs[i][0], + handle, handle_bound_to_stream); + } else { + puf_mm_nn(&gram_buf, &src, &dst, + stream, 1.0f, (float)ns_coeffs[i][0], + handle, handle_bound_to_stream); + } + } + } + float scale = sqrtf(fmaxf(1.0f, (float)R / (float)C)); + muon_store_update<<>>( + gc_ptr, x_buf.data, scale, (int)ne); +} + +void muon_step(Muon* m, Float weights, Prec model_weights, Prec grads, float max_grad_norm, cudaStream_t stream = 0) { int n_grad = (int)numel(grads.shape); + bool fuse_clip_nesterov = m->num_lanes > 0 + && m->num_lanes == m->param_alloc->num_regs; int sum_blocks = min((int)grid_size(n_grad), 256); muon_sum_sq_partials<<>>( m->norm_partials, grads.data, n_grad); muon_sum_sq_reduce<<<1, 256, 0, stream>>>( m->grad_norm, m->norm_partials, sum_blocks); - muon_clip_nesterov<<>>( - m->mb.data, grads.data, m->grad_norm, - max_grad_norm, 1e-6f, (float)m->momentum, n_grad); + if (!fuse_clip_nesterov) { + muon_clip_nesterov<<>>( + m->mb.data, grads.data, m->grad_norm, + max_grad_norm, 1e-6f, (float)m->momentum, n_grad); + } // Per-param NS into workspace; write scaled update back into flat grads. // 1D params already hold their update in-place (scale 1). - long offset = 0; - for (int _i = 0; _i < m->param_alloc->num_regs; _i++) { - AllocEntry& e = m->param_alloc->regs[_i]; - precision_t* gc_ptr = grads.data + offset; - long ne = numel(e.shape); - offset += ne; - if (ndim(e.shape) < 2) { - continue; + if (m->num_lanes > 0) { + // Fork each matrix lane from the post-clip point on the caller stream. + cudaEventRecord(m->matrices_ready, stream); + for (int i = 0; i < m->num_lanes; i++) { + cudaStreamWaitEvent( + m->lanes[i].stream, m->matrices_ready, 0); } - - long R = e.shape[0], C = ne / R; - long M = min(R, C); - bool tall = R > C; - Prec x = {.data = gc_ptr, .shape = {R, C}}; - Prec x_buf = {.data = m->x_buf.data, .shape = {R, C}}; - Prec gram = {.data = m->gram.data, .shape = {M, M}}; - Prec gram_buf = {.data = m->gram_buf.data, .shape = {M, M}}; - - int nblk = min((int)grid_size(ne), 256); - muon_sum_sq_partials<<>>( - m->norm_partials, x.data, (int)ne); - muon_sum_sq_reduce<<<1, 256, 0, stream>>>( - m->ns_norm, m->norm_partials, nblk); - muon_l2_normalize<<>>( - x.data, m->ns_norm, 1e-7f, (int)ne); - - // 5 steps land in x_buf. 4 = you break it. - for (int i = 0; i < 5; ++i) { - Prec& src = (i % 2 == 0) ? x : x_buf; - Prec& dst = (i % 2 == 0) ? x_buf : x; - if (tall) { - puf_mm_tn(&src, &src, &gram, stream); - } else { - puf_mm(&src, &src, &gram, stream); + // Enqueue in the stable descending-work schedule. Each matrix has its + // own stream and private scratch. + for (int i = 0; i < m->num_lanes; i++) { + MuonMatrix& matrix = m->matrices[i]; + MuonLane& lane = m->lanes[i]; + long ne = matrix.rows * matrix.cols; + if (fuse_clip_nesterov) { + int nblk = min((int)grid_size(ne), 256); + muon_clip_nesterov_sum_sq_partials<<< + nblk, 256, 0, lane.stream>>>(lane.norm_partials, + m->mb.data + matrix.offset, grads.data + matrix.offset, + m->grad_norm, max_grad_norm, 1e-6f, + (float)m->momentum, (int)ne); } - puf_copy(&gram_buf, &gram, stream); - puf_mm_nn(&gram, &gram, &gram_buf, stream, - (float)ns_coeffs[i][2], (float)ns_coeffs[i][1]); - puf_copy(&dst, &src, stream); - if (tall) { - puf_mm_nn(&src, &gram_buf, &dst, - stream, 1.0f, (float)ns_coeffs[i][0]); - } else { - puf_mm_nn(&gram_buf, &src, &dst, - stream, 1.0f, (float)ns_coeffs[i][0]); + muon_matrix_step(grads.data + matrix.offset, + matrix.rows, matrix.cols, + lane.norm_partials + 256, lane.norm_partials, + lane.gram, lane.gram_buf, lane.x_buf, + lane.stream, lane.cublas_handle, true, + fuse_clip_nesterov); + } + for (int i = 0; i < m->num_lanes; i++) { + MuonLane& lane = m->lanes[i]; + cudaEventRecord(lane.done, lane.stream); + cudaStreamWaitEvent(stream, lane.done, 0); + } + } else { + long offset = 0; + for (int _i = 0; _i < m->param_alloc->num_regs; _i++) { + AllocEntry& e = m->param_alloc->regs[_i]; + long ne = numel(e.shape); + if (ndim(e.shape) >= 2) { + long R = e.shape[0], C = ne / R; + muon_matrix_step(grads.data + offset, R, C, + m->ns_norm, m->norm_partials, + m->gram.data, m->gram_buf.data, m->x_buf.data, + stream, g_cublas_handle, false, false); } + offset += ne; } - float scale = sqrtf(fmaxf(1.0f, (float)R / (float)C)); - muon_store_update<<>>( - gc_ptr, x_buf.data, scale, (int)ne); } muon_weight_update<<>>( - weights.data, grads.data, m->lr, 0.0f, n_grad); + weights.data, model_weights.data, grads.data, + m->lr, 0.0f, n_grad); } // Train layout is (B, T). Views are sliced each mb; scratch is allocated. diff --git a/src/pufferl.cu b/src/pufferl.cu index ed62658531..7d9f9190f1 100644 --- a/src/pufferl.cu +++ b/src/pufferl.cu @@ -761,6 +761,28 @@ __global__ void zero_term_state(Prec state, Float terminals, state.data[i] = from_float(0.0f); } +__global__ void zero_term_state_agents( + Prec state, const float* terminals) { + __shared__ float terminal; + int rel = blockIdx.x; + if (threadIdx.x == 0) { + terminal = terminals[rel]; + } + __syncthreads(); + if (terminal == 0.0f) { + return; + } + + int L = state.shape[0]; + int H = state.shape[2]; + for (int lh = threadIdx.x; lh < L * H; lh += blockDim.x) { + int h = lh % H; + int layer = lh / H; + long i = state_elem_idx(layer, (int)state.shape[1], rel, h, H); + state.data[i] = from_float(0.0f); + } +} + // Select time t, then agents [start, start+count). Rank-2 has F==0 (zero-term shape); // stride uses max(F, 1). Out shape {count, F} keeps ndim 1 when F==0. Prec puf_slice(Prec p, int t, int start, int count) { @@ -850,9 +872,15 @@ static void pufferl_forward_step(PuffeRL* pufferl, int buf, int t, Prec mask_b = puf_slice(rollouts.action_mask, t, sub, n); // Per-policy state is compact (n agents); local index 0..n-1. - int state_n = (int)st->shape[0] * n * (int)st->shape[2]; - zero_term_state<<>>( - *st, env->terminals, 0, sub, n); + int agent_state_n = (int)st->shape[0] * (int)st->shape[2]; + int state_n = agent_state_n * n; + if (agent_state_n > BLOCK_SIZE && n >= BLOCK_SIZE) { + zero_term_state_agents<<>>( + *st, env->terminals.data + sub); + } else { + zero_term_state<<>>( + *st, env->terminals, 0, sub, n); + } // Carry path: snapshot trainable policy state into per-slot initial_states. if (!pol->frozen && t == 0 && rollouts.initial_states.data != NULL) { @@ -1250,14 +1278,21 @@ static void rollout_start(PuffeRL* p) { && "cudaStreamBeginCapture failed"); } int H = p->hypers.horizon; - for (int t = 0; t < H; t++) { - int base = t * EV_T; + bool record_step_timing = !p->hypers.cudagraphs; + for (int t = 0; t < H; t++) { + int base = t * EV_T; + if (record_step_timing) { cudaEventRecord(ev[base + MODEL_START], stream); - pufferl_forward_step(p, 0, t, stream); + } + pufferl_forward_step(p, 0, t, stream); + if (record_step_timing) { cudaEventRecord(ev[base + MODEL_END], stream); - puf_step(p->vec->envs); + } + puf_step(p->vec->envs); + if (record_step_timing) { cudaEventRecord(ev[base + ENV_END], stream); } + } if (first) { cudaGraph_t graph; assert(cudaStreamEndCapture(stream, &graph) == cudaSuccess @@ -1533,13 +1568,8 @@ static void train_epoch_gpu(PuffeRL* pufferl, RolloutBuf src, int slot, numel(pufferl->grad.shape), NCCL_PRECISION, ncclAvg, pufferl->nccl_comm, stream); } - muon_step(&pufferl->muon, primary->master_weights, + muon_step(&pufferl->muon, primary->master_weights, primary->param, pufferl->grad, hypers->max_grad_norm, stream); - if (USE_BF16) { - int n = numel(primary->param.shape); - cast<<>>( - primary->param.data, primary->master_weights.data, n); - } } cudaEventRecord(ev[TE_FE], stream); }