📦 WP-193 — perf(train): fuse post-backward state updates¶
Engineering log · Phase 17 — R1 pose-objective parity
WP-193 — post-backward state updates use fewer CUDA launches¶
The measured target is the complete interval, not either helper. A public COCO val2017 training harness measured from one completed batch to the next, including loader wait, host-to-device transfer, forward and loss, backward, clipping, MuSGD, EMA and synchronization. On an isolated g2-standard-8 NVIDIA L4 (driver 580.178.04, torch 2.9.1+cu129), variant n, batch 4, 640 px and BF16, five warm ABBA pairs moved from a 256.8 ms baseline median to 222.5 ms with this row. The paired reductions were 13.28%, 13.13%, 14.90%, 12.94% and 6.38%: median 13.13%, range 6.38–14.90%. The fifth candidate arm was the run's one timing outlier at 246.4 ms, but remained faster than the fastest baseline arm at 254.9 ms. Each arm discarded 15 warm-up intervals and measured the next 44. Peak allocated memory was 822.1 MiB in both arms. The paid VM was deleted after the run and an independent cleanup found no archbench instance alive.
Two launch-bound stages make one post-backward change. The existing optimizer already groups exact-shape CUDA matrices, but the torch.vmap layer around the scalar Newton--Schulz expression remains; it already lowered each product to one batched multiply, so the change drops that layer and folds the scale-and-add into the product epilogue. _orthogonalize_batched uses baddbmm for the same b*A + c*A@A and a*X + polynomial@X expressions, preserving each slice's normalization, coefficient triple, requested step count, gains, scaling and decay. EMA previously launched one multiply and one add for every floating parameter or buffer; grouping by (device, dtype) and calling torch._foreach_lerp_ once per bucket evaluates the same shadow + (1 - decay) * (live - shadow) recurrence. Singletons and non-CUDA optimizer paths remain scalar. Both transformations are scheduling changes derived from this tree's buckets and profiler counts; no training constant or imported implementation enters the diff.
Floating-point schedules are equivalent, not bit-identical. Native batching and interpolation reassociate operations, so the regression contract is a bounded trajectory against the prior scalar formulas rather than byte identity. The repository's multi-step trajectory tests run on CPU, over three (6, 4) float32 matrices for the optimizer and float32 parameters for the EMA. The 40-step fixed-gradient CUDA run on an L4 is a one-off measurement, not a repository test, and no multi-step trajectory test covers square, wide or tall shapes or a float64 bucket. A finite one-step public-COCO run starts from the same loss and checks model and momentum differences directly. This row does not claim an unchanged long-run stochastic trajectory: tiny per-step differences can change later discrete assignments, as the repository already records for the exact-shape batching it replaces.
The rejected queue was judged end to end. MuSGD's native batch alone reached only 7.48% median on its final five-pair L4 rerun and was rejected despite a larger exploratory result. EMA foreach was retained only after the combined five-pair run; no isolated repeat-level EMA claim is made. Five segmentation mask-loss compilation repeats had a 4.44% median but ranged from -10.88% to 6.05%, crossing zero. Packed host metadata was rejected from its structural stage bound, while loss logging and scaled-dot-product attention had only inconclusive single local A/B screens; none carries a quantified claim here. cuDNN benchmarking was slower in its exploratory complete-path run and also carries no repeat-level claim. None became a dossier. Combining the two adjacent post-backward launch reductions is the first candidate whose complete-pipeline median clears the requested 8% threshold.
The accelerator golden check is outstanding, not run clean. _orthogonalize_batched reassociates the same CUDA Newton--Schulz buckets the WP-177 entry (ESCALATION.md, 2026-09-15) already measured moving goldens/gpu/overfit_micro_kp.json's train_oks_ap 0.117 outside its stored band from float-scheduling alone, on a question that entry leaves unresolved. This row's own correctness evidence is the 40-step fixed-gradient CUDA trajectory and the one-step public-COCO check (both bounded-agreement, not golden-harness), plus the offline make gate (50/50, CPU-only). make gate-gpu / check_goldens.py --include-gpu was not run against this diff, so whether the additional reassociation here moves overfit_micro_kp or overfit_micro_obb further is unmeasured — the roadmap row stays 🔄 until it is.