Skip to content

Add stochastic rounding and philox RNG under rocdl - #917

Closed
kashif wants to merge 1 commit into
ROCm:pr/enh-cast-rounding-modefrom
kashif:sr-rocdl
Closed

Add stochastic rounding and philox RNG under rocdl#917
kashif wants to merge 1 commit into
ROCm:pr/enh-cast-rounding-modefrom
kashif:sr-rocdl

Conversation

@kashif

@kashif kashif commented Jul 28, 2026

Copy link
Copy Markdown

Stacked on #916. 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 adds 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.

Note for review: this targets the #916 branch so the diff shows only the stochastic-rounding work. It should be merged after #916, and GitHub will retarget it to main automatically.

@coderfeli
coderfeli deleted the branch ROCm:pr/enh-cast-rounding-mode July 28, 2026 16:04
@coderfeli coderfeli closed this Jul 28, 2026
@kashif

kashif commented Jul 30, 2026

Copy link
Copy Markdown
Author

Reopened as #927 against main now that #916 has merged. GitHub would not let me reopen this one because its base branch was deleted on merge. Continuing there.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants