Add stochastic rounding and philox RNG under rocdl

[kashif]

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 through arith.truncf): it needs entropy and the hardware form is a target-specific instruction (CDNA v_cvt_*_sr, PTX cvt.rs). So it lives under rocdl/ as a target-specific function. The f32 to bf16 path here is a portable software emulation; a CDNA hardware sr variant can slot in alongside it later.

New API:

• fx.rocdl.philox_4x32(counter, key) is the Random123 Philox 4x32 counter RNG (same generator as numpy.random.Philox, cuRAND Philox4_32_10, Triton tl.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 rounding option, default "rn" so nothing changes. With rounding="rs" the f32 accumulator converts to bf16 through stochastic rounding, keyed per element from a runtime sr_seed launch 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.

[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?

[sjfeng1999]

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?

I add standard rounding mode support in #916. I found that RoundingMode.STOCHASTIC can't be implemented as a target-neutral mode like others rounding modes. If we want to use target-specific instructions to do rounding, like PTX cvt.rs on Blackwell, the sr cvt on CDNA. It should be a target-specific function under rocdl/.

[kashif]

thanks checking

[kashif]

Makes sense, thanks for adding #916. Agreed that STOCHASTIC does not fit the target-neutral arith.truncf roundingmode path, since it needs entropy and the hardware form is cvt.rs / sr cvt.

I moved stochastic_round_bf16 and philox_4x32 out of numeric.py and under rocdl/ as target-specific functions (fx.rocdl.stochastic_round_bf16, fx.rocdl.philox_4x32), so the standard IEEE modes stay in #916 and stochastic rounding lives with the target instructions it maps to. This is up as #917, stacked on your #916 branch so the diff shows only the stochastic-rounding work. Happy to close this PR in favor of #917.

[kashif]

Superseded by #917, which moves these helpers under rocdl/ (fx.rocdl.philox_4x32, fx.rocdl.stochastic_round_bf16) and stacks on #916. Continuing there.

[kashif]

@sjfeng1999 ready for a review

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (3)

python/flydsl/expr/rocdl/random.py:71

• stochastic_round_bf16 currently perturbs the low 16 bits for all inputs; for NaNs this can carry into higher mantissa/exponent bits, so NaN payload (and in extreme cases the exponent) is not preserved. This contradicts the docstring claim that NaNs pass through unchanged and can make NaN handling non-deterministic.

    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)

python/flydsl/expr/rocdl/random.py:27

• The PR description says this matches Random123 / NumPy / cuRAND Philox4_32_10, but the default here is 7 rounds. With the default, callers using philox_4x32(counter, key) will not reproduce those reference generators.

This issue also appears on line 67 of the same file.

def philox_4x32(counter, key, rounds: int = 7):

tests/unit/test_stochastic_rounding.py:41

• This test (and the PR description) mentions NaN pass-through, but the current parameterization excludes NaN and the assertions assume the usual round-up count property. If stochastic_round_bf16 is meant to keep NaNs deterministic, add an explicit NaN case with a pass-through assertion.

@pytest.mark.parametrize("x", [1.0, 1.001953125, 3.14159, 1e-3, float("inf"), float("-inf")])
def test_stochastic_round_bf16_unbiased(x):

[sjfeng1999]

Main thing in the latest commit: random is a standalone library now (expr/random/, universal impl plus an optional per-target override).
fx.random.<fn> picks the implementation based on the current compilation backend target
fx.random.universal.<fn> always runs the portable one

We should settle on what the complete API should look like before adding more one-off functions. Are these six triton.rand functions enough for your needs? https://triton-lang.org/main/python-api/triton.language.html#random-number-generation

[kashif]

Thanks, this is much nicer than what I had. The universal + per-target override split is clean, and the Library dispatch is a nice reusable piece.

I pushed 37bfa99 on top: it amortizes the Philox draw across the epilogue block (one call per 8 outputs using all four words, instead of one draw per element that used a single word), which was most of the RS-path RNG cost.

On the six functions: for stochastic rounding I only need raw bits, so randint4x (our philox_4x32) is the only one I actually use. No need for rand/randn today. But I'm happy to adopt all six as the target shape and match Triton's names/signatures (heads up: Triton is (seed, offset), the reverse of our (counter, key)). I'd just implement the uniform/normal ones lazily when something actually needs them.

Two smaller thoughts. I'd keep cvt_f32_to_bf16_sr out of fx.random, since it consumes RNG rather than generates it, maybe a cast/round namespace that calls into random. And I looked at making the cvt saturate (Inf to max finite, NVIDIA style) but left it alone: your non-saturating behavior matches AMD hardware and the gfx950 v_cvt_sr path.