Reports · 2026-09-18

Collection is bound by one attention kernel, and the fix is a flag

The profile

One collection step moves 1,024 environments forward once. The stage table below times one iteration (202 steps) of each path with synchronize brackets around every stage; the full table with the Rust calls and the host-device copies separated is in the run record.

The hypothesis named one stage above 40% of wall time; the forward pass sits at 96%. Inside it, the profiler shows the math attention fallback costing 136 ms per step on its own: the Radeon 8060S gates its fused attention kernels behind an experimental flag, and the default path materialises the 1024 by 6 by 93 by 93 attention matrices. With the flag set, the same attention measures 1.0 ms per call against 19.2 ms, and the environment, the sampling over 10,000 actions and every copy turn out to be rounding errors next to the network.

What was measured and rejected

The brief listed six candidate fixes; measurement kept one.

• torch.compile of the forward reached 60.9 ms against 86.1 ms eager with the flag, a further 1.42x, with a 36.6 s cold compile. It was rejected on correctness grounds: the compiled graph returns NaN rows for scattered batch entries on real data, on 44 to 99 of 202 steps across every attention backend, both input dtypes, an encoder-only graph and with split reductions disabled, while the eager network is clean on every batch tested. A per-step NaN guard with eager fallback caught every case, but a quarter of the steps then paid for two forwards and the win was gone. This is a bug in this inductor build for gfx1151, recorded here so the next attempt starts from evidence.
• Gumbel-max sampling on the device measured 2.9 ms per step against 2.8 ms for multinomial once both run eagerly; the fusion that would have made it faster needs torch.compile. The loop keeps multinomial.
• More environments do not help: the compiled forward served 16.8k env forwards per second at batch 1,024, 15.8k at 2,048 and 13.4k at 4,096. The GPU is saturated at the preset's batch size.
• Pinned-memory copies, uint8 planes and Rust thread tuning target stages that together cost 4.4 ms of a 225 ms step.

The curves match

Ten iterations from seed 0 on both paths, with seed 1 on the old path as the noise yardstick. The plot shows collection speed; the table shows the largest per-iteration difference between the fast run and the reference run against the largest difference between the two reference seeds.

Bit-exact equality was never expected across attention kernels: the worst logit difference between the two backends on real batches is 0.0078, one to two bf16 ulps, but one different sampled action changes the rest of a trajectory. The kept counter is the most sensitive metric because the advantage filter's floor at 0.01 turns small value differences into set-membership differences; the runs rejoin by iteration eight (36,980 kept fast against 39,055 and 36,958 reference). The metrics that carry the learning signal, policy loss, entropy, games and plies, track inside or at the edge of the two-seed envelope, and iteration 1 is indistinguishable (entropy 3.030 on both paths, mean plies 86 on both).

The environment share

The world-model study's premise was that environment steps are nearly free here. Measured: env.step, env.observe and env.legal_mask, including their host-device copies, took 2.0% of collection wall time before the fix and take 5.4% after. The engine alone steps the same 1,024 games at 0.6 to 0.85 million steps per second with no network in the loop, two orders of magnitude above what the training loop asks of it. Imagination cannot pay for itself on this machine; the premise holds with margin.

What remains

The forward pass is still 89% of the step and runs at about 8.5 TFLOP/s of effective bf16 compute, memory-bound on activations. The paths that would go further are the ones the evidence rejected (a working inductor build for this GPU, or hand-fused kernels), or a smaller or cheaper network, which is a learning decision rather than an engineering one. The scaled run should plan on about 10,000 collection steps per second per GPU process, and the update phase now costs roughly as much wall time as collection, so it is the next profiling target.

Run record

PYTHONPATH=training .venv/bin/python scripts/profile_collect.py --path reference --out runs/profile-reference.json
PYTHONPATH=training .venv/bin/python scripts/profile_collect.py --path fast --out runs/profile-fast.json
PYTHONPATH=training ARMIES_COLLECT_PATH=reference .venv/bin/python -m armies_train train --preset small --iterations 10 --out runs/throughput-ref-seed0 --seed 0
PYTHONPATH=training .venv/bin/python -m armies_train train --preset small --iterations 10 --out runs/throughput-fast-seed0 --seed 0
python3 scripts/build_throughput_assets.py