Add stochastic rounding and philox RNG helpers

[kashif]

Adds stochastic rounding for f32 to bf16, plus a small counter-based RNG to feed it, and wires it into the RDNA bf16 GEMM epilogue as the first user. FlyDSL had no stochastic rounding of any kind, which matters for low-precision accumulation and training where round-to-nearest is biased.

Primitives:

stochastic_round_bf16(x, rand) perturbs the 16 mantissa bits that bf16 discards with the supplied random bits before truncating, so the probability of rounding up equals the fractional distance to the next bf16 value. bf16 shares the f32 exponent, so Inf and NaN pass through unchanged and no special casing is needed.

philox_4x32(counter, key) is a counter-based PRNG returning four random words, so a kernel can draw reproducible randomness from a global index and a seed without carrying RNG state.

Both are reachable as fx.stochastic_round_bf16 and fx.philox_4x32.

Consumer:

The RDNA WMMA GEMM (gfx11 and gfx12 variants) gains a rounding option, defaulting to "rn" so nothing changes. With rounding="rs" the f32 accumulator converts to bf16 through stochastic rounding, keyed per output element from sr_seed, instead of round-to-nearest.

Tests:

The rounding is proven exactly unbiased in pure Python (over the full 16-bit random range the number of round-ups equals the discarded low bits) 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-rounding path stays within tolerance of round-to-nearest, changes with the seed, and is reproducible for a fixed seed.

[Copilot]

Pull request overview

Adds stochastic rounding support for FP32→BF16 and a small Philox-based counter RNG helper to provide reproducible random bits inside kernels, along with unit/device tests validating correctness and statistical properties.

Changes:

• Implement philox_4x32(counter, key) and stochastic_round_bf16(x, rand) in python/flydsl/expr/numeric.py.
• Re-export these helpers via flydsl.expr (through python/flydsl/expr/typing.py) so they’re reachable as fx.philox_4x32 / fx.stochastic_round_bf16.
• Add tests covering unbiased rounding math (CPU) and device behavior/uniformity (GPU).

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
tests/unit/test_stochastic_rounding.py	Adds L0 unbiasedness checks and L2 GPU tests for stochastic rounding + Philox RNG
python/flydsl/expr/typing.py	Re-exports the new helpers from expr.numeric through the fx public surface
python/flydsl/expr/numeric.py	Implements Philox 4x32 RNG and BF16 stochastic rounding helper

Comments suppressed due to low confidence (2)

tests/unit/test_stochastic_rounding.py:58

• To actually exercise the public re-export, call the helpers via fx.* (not via a direct import). This ensures typing.__all__ / expr.__init__ re-exporting stays covered.

        g = fx.block_idx.x * BLOCK + fx.thread_idx.x
        r = philox_4x32(fx.Uint32(g), fx.Uint32(seed))[0]
        fx.memref_store(stochastic_round_bf16(x, r), Out, g)

tests/unit/test_stochastic_rounding.py:90

• Same as above: use fx.philox_4x32 so the test covers the documented fx.* entry point.

        g = fx.block_idx.x * BLOCK + fx.thread_idx.x
        fx.memref_store(fx.Int32(philox_4x32(fx.Uint32(g), fx.Uint32(seed))[0]), Out, g)

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[sjfeng1999]

Are stochastic_round and philox_4x32 part of any mature and widely used library API? Adding just those two in isolation feels rather ad hoc.

[kashif]

thanks @sjfeng1999 agree it does feel adhoc, Philox is a mature, widely used API: it is the generator, exposed as numpy.random.Philox, cuRAND's Philox4_32_10, and the generator behind Triton's tl.rand.

similarly, stochastic rounding is a standard technique rather than a single library function name. In practice it appears as a rounding mode of a narrowing conversion: PTX cvt.rs on Blackwell, the sr cvt on CDNA, and RoundingMode.RS in the CuTe DSL. So the right shape is a rounding mode, not a free function.

I could reorganize as:

• philox_4x32 moves under an fx.random namespace, since it is a general purpose RNG (dropout, init, and so on), not a rounding concern.
• stochastic_round becomes a rounding mode on the existing conversion, e.g. x.to(fx.BFloat16, rounding=RoundingMode.STOCHASTIC, seed=...), matching cvt.rs and RoundingMode.RS, so it folds into to() instead of adding a new top level function.

does that direction work for you?