Add stochastic rounding and philox RNG under rocdl - #927
Conversation
There was a problem hiding this comment.
🟡 Not ready to approve
The current implementation/documentation has correctness and contract mismatches (Philox round count vs stated generator, and NaN handling in stochastic bf16 rounding) that should be fixed and covered by tests before approval.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Pull request overview
This PR adds target-specific stochastic rounding support and a counter-based Philox RNG under flydsl.expr.rocdl, then wires stochastic bf16 rounding into the RDNA WMMA GEMM epilogue behind an opt-in rounding="rs" mode with a runtime seed.
Changes:
- Introduce
fx.rocdl.philox_4x32(counter, key)andfx.rocdl.stochastic_round_bf16(x, rand)in a newexpr/rocdl/random.pymodule and re-export them fromfx.rocdl. - Add an optional
roundingmode andsr_seedlaunch argument to RDNA3/RDNA4 bf16 WMMA GEMM epilogues to enable stochastic rounding on stores. - Add unit + device tests for stochastic rounding and Philox behavior, plus a GEMM-level wiring test for the stochastic epilogue path.
File summaries
| File | Description |
|---|---|
| tests/unit/test_stochastic_rounding.py | Adds L0 math test + L2 GPU tests for stochastic rounding and Philox; proposes expanding NaN device coverage. |
| tests/kernels/test_rdna_gemm.py | Adds an RDNA WMMA GEMM test that validates stochastic epilogue wiring, seed variability, and reproducibility. |
| python/flydsl/expr/rocdl/random.py | Implements Philox RNG and stochastic bf16 rounding helper under expr/rocdl. |
| python/flydsl/expr/rocdl/init.py | Re-exports the new RNG/rounding helpers from fx.rocdl. |
| kernels/gemm/rdna3_f16_gemm.py | Adds rounding/sr_seed plumbing and stochastic bf16 epilogue option for gfx11 WMMA GEMM. |
| kernels/gemm/rdna_f16_gemm.py | Adds rounding/sr_seed plumbing and stochastic bf16 epilogue option for gfx120 WMMA GEMM. |
Review details
- Files reviewed: 6/6 changed files
- Comments generated: 3
- Review effort level: Low
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
| def philox_4x32(counter, key, rounds: int = 7): | ||
| """Philox 4x32 counter-based PRNG. Returns four ``Uint32`` random words. | ||
|
|
||
| ``counter`` and ``key`` are ``Uint32`` (e.g. a global element index and a | ||
| seed). The same ``(counter, key)`` always yields the same words, so no RNG |
| def stochastic_round_bf16(x, rand): | ||
| """Round a ``Float32`` to ``BFloat16`` with stochastic rounding. | ||
|
|
||
| ``rand`` is a ``Uint32`` supplying entropy; its low 16 bits are used. | ||
| bf16 keeps the top 16 bits of the f32, so adding the random value to the | ||
| 16 discarded mantissa bits turns truncation into a round-up with | ||
| probability equal to the fractional distance to the next bf16 value. | ||
| bf16 shares the f32 exponent field, so Inf and NaN pass through unchanged | ||
| and no special-case handling is needed. | ||
| """ | ||
| u = Float32(x).bitcast(Uint32) + (Uint32(rand) & Uint32(0xFFFF)) | ||
| return Uint16(u >> Uint32(16)).bitcast(BFloat16) |
| u = out.cpu().numpy().astype(np.uint32).astype(np.float64) / 2**32 | ||
| assert abs(u.mean() - 0.5) < 0.01, f"mean {u.mean()} not uniform" | ||
| # counter-based PRNG on distinct counters must not collide in bulk | ||
| assert np.unique(out.cpu().numpy()).size > int(0.999 * P) |
|
Consolidating on #907, which now carries this same rocdl-namespaced code (fx.rocdl.philox_4x32, fx.rocdl.stochastic_round_bf16) rebased on main, and keeps the earlier review discussion. Continuing there. |
Adds stochastic rounding and a Philox counter RNG, placed under the rocdl target package rather than as top-level helpers, following the discussion on #907.
Stochastic rounding does not fit the target-neutral IEEE rounding modes that #916 added to
.to()(lowered througharith.truncf): it needs entropy and the hardware form is a target-specific instruction (CDNAv_cvt_*_sr, PTXcvt.rs). So it lives underrocdl/as a target-specific function. The f32 to bf16 path here is a portable software emulation; a CDNA hardwaresrvariant can slot in alongside it later.New API:
fx.rocdl.philox_4x32(counter, key)is the Random123 Philox 4x32 counter RNG (same generator asnumpy.random.Philox, cuRANDPhilox4_32_10, Tritontl.rand), returning four reproducible random words from a global index and a seed.fx.rocdl.stochastic_round_bf16(x, rand)perturbs the 16 mantissa bits bf16 discards with the random bits before truncating, so the probability of rounding up equals the fractional distance to the next bf16 value. Inf and NaN pass through unchanged.Consumer:
The RDNA WMMA GEMM (gfx11 and gfx12) gains a
roundingoption, default"rn"so nothing changes. Withrounding="rs"the f32 accumulator converts to bf16 through stochastic rounding, keyed per element from a runtimesr_seedlaunch argument.Tests:
The rounding is proven exactly unbiased in pure Python and covers Inf and NaN. Device tests confirm outputs land on the two nearest bf16 values with the right up probability and a matching mean, and that the RNG is uniform and collision free. The GEMM test checks the stochastic path stays within tolerance of round-to-nearest, varies with the runtime seed, and is reproducible for a fixed seed.