Add stochastic rounding and philox RNG helpers

kernels/gemm/rdna3_f16_gemm.py

@@ -45,6 +45,7 @@ def create_wmma_gemm_module(
     in_dtype="bf16",
     out_dtype="bf16",
     *,
+    rounding="rn",  # "rn" (round to nearest) or "rs" (stochastic rounding)
     reg_m=4,
     reg_n=4,
     reg_k=2,
@@ -69,6 +70,9 @@ def create_wmma_gemm_module(
     THREADS_PER_BLOCK = NUM_WAVES * WAVE_SIZE  # 128
 
     assert reg_k >= 2 and reg_k % 2 == 0
+    assert rounding in ("rn", "rs"), f"rounding must be 'rn' or 'rs', got {rounding!r}"
+    if rounding == "rs":
+        assert out_dtype == "bf16", "stochastic rounding currently supports bf16 output only"
 
     LOAD_VEC = 8  # 8 bf16 = 128-bit GMEM/LDS load
     A_TILE_ELEMS = BLOCK_M * BLOCK_K
@@ -119,6 +123,7 @@ def wmma_gemm_kernel(
         arg_c: fx.Tensor,
         arg_a: fx.Tensor,
         arg_bt: fx.Tensor,
+        sr_seed: fx.Int32,  # runtime seed; only read on the stochastic-rounding path
     ):
         lds_storage = fx.SharedAllocator().allocate(_SharedStorage).peek()
         lds_ptr = lds_storage.lds.ptr  # i8-base aliased as elem_dtype*
@@ -331,11 +336,17 @@ def _do_compute_rk(accs_in, rk, buf_offset):
                     g_row = tile_m0 + wmma_m_off + 2 * si + klane
                     g_col = tile_n0 + wmma_n_off + lane16
                     val = accs[idx][si]
-                    if const_expr(out_dtype == "bf16"):
+                    elem_off = g_row * N + g_col
+                    if const_expr(rounding == "rs"):
+                        # Stochastic rounding: perturb the discarded bits with a
+                        # per-element Philox draw keyed by the element index, so
+                        # the f32 -> bf16 store is unbiased in expectation.
+                        rbits = fx.philox_4x32(fx.Uint32(elem_off), fx.Uint32(sr_seed))[0]
+                        val = fx.stochastic_round_bf16(val, rbits)
+                    elif const_expr(out_dtype == "bf16"):
                         val = val.to(fx.BFloat16)
                     elif const_expr(out_dtype == "f16"):
                         val = val.to(fx.Float16)
-                    elem_off = g_row * N + g_col
                     buffer_ops.buffer_store(val, c_rsrc, elem_off)
 
     @flyc.jit
@@ -344,12 +355,13 @@ def launch_gemm(
         arg_a: fx.Tensor,
         arg_bt: fx.Tensor,
         stream: fx.Stream,
+        sr_seed: fx.Int32 = 0,
     ):
         c1 = 1
         total_blocks = grid_m * grid_n
         bk = THREADS_PER_BLOCK
 
-        launcher = wmma_gemm_kernel(arg_c, arg_a, arg_bt)
+        launcher = wmma_gemm_kernel(arg_c, arg_a, arg_bt, sr_seed)
         launcher.launch(
             grid=(total_blocks, c1, c1),
             block=(bk, c1, c1),

kernels/gemm/rdna_f16_gemm.py

@@ -39,6 +39,7 @@ def create_wmma_gemm_module(
     in_dtype="bf16",
     out_dtype="bf16",
     *,
+    rounding="rn",  # "rn" (round to nearest) or "rs" (stochastic rounding)
     reg_m=4,  # M-repeats per warp
     reg_n=4,  # N-repeats per warp
     reg_k=2,  # K-steps per tile (32/16=2)
@@ -56,6 +57,9 @@ def create_wmma_gemm_module(
     THREADS_PER_BLOCK = NUM_WAVES * WAVE_SIZE  # 128
 
     assert reg_k >= 2 and reg_k % 2 == 0
+    assert rounding in ("rn", "rs"), f"rounding must be 'rn' or 'rs', got {rounding!r}"
+    if rounding == "rs":
+        assert out_dtype == "bf16", "stochastic rounding currently supports bf16 output only"
 
     # Loading: each thread loads 8 bf16 elements per load (128 bits = buffer_load_b128)
     LOAD_VEC = 8
@@ -96,6 +100,7 @@ def wmma_gemm_kernel(
         arg_c: fx.Tensor,
         arg_a: fx.Tensor,
         arg_bt: fx.Tensor,
+        sr_seed: fx.Int32,  # runtime seed; only read on the stochastic-rounding path
     ):
         lds = fx.SharedAllocator(static=False).allocate(fx.Array[lds_elem_dtype, LDS_TOTAL, 16]).peek()
 
@@ -319,11 +324,17 @@ def _load_a_single_from_lds(rk, rm_val, buf_offset):
                     g_row = tile_m0 + wmma_m_off + base8 + si
                     g_col = tile_n0 + wmma_n_off + lane16
                     val = accs[idx][si]
-                    if const_expr(out_dtype == "bf16"):
+                    elem_off = g_row * N + g_col
+                    if const_expr(rounding == "rs"):
+                        # Stochastic rounding: perturb the discarded bits with a
+                        # per-element Philox draw keyed by the element index, so
+                        # the f32 -> bf16 store is unbiased in expectation.
+                        rbits = fx.philox_4x32(fx.Uint32(elem_off), fx.Uint32(sr_seed))[0]
+                        val = fx.stochastic_round_bf16(val, rbits)
+                    elif const_expr(out_dtype == "bf16"):
                         val = val.to(fx.BFloat16)
                     elif const_expr(out_dtype == "f16"):
                         val = val.to(fx.Float16)
-                    elem_off = g_row * N + g_col
                     buffer_ops.buffer_store(val, c_rsrc, elem_off)
 
     # ── Host launcher ──────────────────────────────────────────────────────
@@ -333,12 +344,13 @@ def launch_gemm(
         arg_a: fx.Tensor,
         arg_bt: fx.Tensor,
         stream: fx.Stream,
+        sr_seed: fx.Int32 = 0,
     ):
         c1 = 1
         total_blocks = grid_m * grid_n
         bk = THREADS_PER_BLOCK
 
-        launcher = wmma_gemm_kernel(arg_c, arg_a, arg_bt)
+        launcher = wmma_gemm_kernel(arg_c, arg_a, arg_bt, sr_seed)
         launcher.launch(
             grid=(total_blocks, c1, c1),
             block=(bk, c1, c1),

python/flydsl/expr/numeric.py

@@ -894,3 +894,66 @@ def __init__(self, x):
         # x is now either: Python int, or index-typed ir.Value
         # Pass directly to Numeric.__init__ (bypass Integer conversion logic)
         Numeric.__init__(self, x)
+
+
+# ---------------------------------------------------------------------------
+# Stochastic rounding
+# ---------------------------------------------------------------------------
+# Rounding a float to a lower-precision type by adding a random value to the
+# discarded low bits before truncation, so the probability of rounding up equals
+# the fractional distance to the next representable value. Unlike round-to-
+# nearest it is unbiased in expectation, which matters for low-precision
+# accumulation and training. Randomness comes from a counter-based PRNG so the
+# stream is reproducible from a seed without carrying RNG state.
+
+_PHILOX_ROUND_A = 0xD2511F53
+_PHILOX_ROUND_B = 0xCD9E8D57
+_PHILOX_KEY_A = 0x9E3779B9
+_PHILOX_KEY_B = 0xBB67AE85
+
+
+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
+    state has to be threaded through a kernel. Each round does two widening
+    32x32 to 64 bit multiplies.
+    """
+    c0 = Uint32(counter)
+    c1 = Uint32(0)
+    c2 = Uint32(0)
+    c3 = Uint32(0)
+    k0 = Uint32(key)
+    k1 = Uint32(0)
+    # Materialize the round multipliers, key increments and shift once instead
+    # of rebuilding identical constants on every unrolled round.
+    round_a = Uint64(_PHILOX_ROUND_A)
+    round_b = Uint64(_PHILOX_ROUND_B)
+    key_a = Uint32(_PHILOX_KEY_A)
+    key_b = Uint32(_PHILOX_KEY_B)
+    shift32 = Uint64(32)
+    for _ in range(rounds):
+        prod_b = Uint64(c2) * round_b
+        prod_a = Uint64(c0) * round_a
+        c0 = Uint32(prod_b >> shift32) ^ c1 ^ k0
+        c2 = Uint32(prod_a >> shift32) ^ c3 ^ k1
+        c1 = Uint32(prod_b)
+        c3 = Uint32(prod_a)
+        k0 = k0 + key_a
+        k1 = k1 + key_b
+    return c0, c1, c2, c3
+
+
+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)

python/flydsl/expr/typing.py

@@ -49,6 +49,8 @@
     _resolve_numeric_type,
     _result_numeric_type_for_op,
     as_numeric,
+    philox_4x32,
+    stochastic_round_bf16,
 )
 from .primitive import *
 from .utils import lazy_classattr
@@ -356,6 +358,8 @@ def vec(self, n: int, elem: ir.Type) -> ir.Type:
     "as_dsl_value",
     "is_generic_address_space",
     "is_target_address_space",
+    "philox_4x32",
+    "stochastic_round_bf16",
     # DSL value types
     "Numeric",
     "as_numeric",

tests/kernels/test_rdna_gemm.py

@@ -103,6 +103,41 @@ def test_f16_gemm_correctness(M, N, K, in_dtype, out_dtype):
     assert verify_output(C.float(), C_ref, atol=0.05, rtol=0.05)
 
 
+def test_f16_gemm_stochastic_rounding():
+    """BF16 GEMM with the stochastic-rounding epilogue: bounded, seed-varying, reproducible.
+
+    Unbiasedness of the rounding itself is proven in tests/unit/test_stochastic_rounding.py;
+    here we only check the GEMM wiring.
+    """
+    _requires_rdna_wmma()
+
+    M = N = K = 256
+    torch.manual_seed(0)
+    A = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") * 0.1
+    B_T = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") * 0.1
+    C_ref = A.float() @ B_T.float().T
+    stream = torch.cuda.current_stream()
+
+    rn_gemm, _, _, _ = create_wmma_gemm_module(M, N, K, in_dtype="bf16", out_dtype="bf16", rounding="rn")
+    rs_gemm, _, _, _ = create_wmma_gemm_module(M, N, K, in_dtype="bf16", out_dtype="bf16", rounding="rs")
+
+    def run(launch_fn, *seed):
+        C = torch.zeros(M, N, dtype=torch.bfloat16, device="cuda")
+        launch_fn(C, A, B_T, stream, *seed)
+        torch.cuda.synchronize()
+        return C.float()
+
+    C_rn = run(rn_gemm)  # default 4-arg launch: seed is unused for round-to-nearest
+    C_rs1, C_rs2 = run(rs_gemm, 1), run(rs_gemm, 2)
+
+    # SR stays within about an ULP of round-to-nearest against the f32 reference
+    assert (C_rs1 - C_ref).abs().max() < 3 * (C_rn - C_ref).abs().max() + 5e-3
+    # the seed is a runtime argument: one compiled kernel, different seeds vary
+    # the rounding, and a repeated seed is reproducible
+    assert not torch.equal(C_rs1, C_rs2)
+    assert torch.equal(C_rs1, run(rs_gemm, 1))
+
+
 @pytest.mark.parametrize(
     "M, N, K",
     [

tests/unit/test_stochastic_rounding.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python3
+
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 FlyDSL Project Contributors
+
+"""Tests for stochastic rounding: stochastic_round_bf16 + philox_4x32.
+
+The rounding math is checked exactly in pure Python (L0): over the full 16-bit
+random range the number of round-ups equals the discarded low 16 bits, which is
+what makes stochastic rounding unbiased. The device path and the RNG are checked
+on GPU (L2): outputs land on the two nearest bf16 values with the right up
+probability, and the empirical mean matches the input.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+
+@pytest.mark.l0_backend_agnostic
+@pytest.mark.parametrize("x", [1.0, 1.001953125, 3.14159, 1e-3, float("inf"), float("-inf")])
+def test_stochastic_round_bf16_unbiased(x):
+    u = int(np.array([x], dtype=np.float32).view(np.uint32)[0])
+    low16 = u & 0xFFFF
+    base = u >> 16
+    r = np.arange(0, 1 << 16, dtype=np.uint64)
+    truncated = ((np.uint64(u) + r) >> np.uint64(16)).astype(np.uint64)
+    # exactly `low16` of the 65536 random draws round up
+    up_count = int(np.sum(truncated != np.uint64(base)))
+    assert up_count == low16
+    # every result is one of the two nearest bf16 values
+    assert set(truncated.tolist()) <= {base, base + 1}
+
+
+@pytest.mark.l2_device
+@pytest.mark.rocm_lower
+@pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires GPU")
+@pytest.mark.parametrize("frac", [0.0, 0.5, 0.9])
+def test_stochastic_round_bf16_device(frac):
+    BLOCK = 256
+    NBLOCKS = 4096
+    P = BLOCK * NBLOCKS
+
+    @flyc.kernel(known_block_size=[BLOCK, 1, 1])
+    def kernel(Out: fx.Tensor, x: fx.Float32, seed: fx.Int32):
+        g = fx.block_idx.x * BLOCK + fx.thread_idx.x
+        r = fx.philox_4x32(fx.Uint32(g), fx.Uint32(seed))[0]
+        fx.memref_store(fx.stochastic_round_bf16(x, r), Out, g)
+
+    @flyc.jit
+    def launch(Out: fx.Tensor, x: fx.Float32, seed: fx.Int32, stream: fx.Stream = fx.Stream(None)):
+        kernel(Out, x, seed).launch(grid=(NBLOCKS, 1, 1), block=(BLOCK, 1, 1), stream=stream)
+
+    step = 2.0**-7  # bf16 step above 1.0
+    x = 1.0 + frac * step
+    lo, hi = 1.0, 1.0 + step
+
+    out = torch.zeros(P, dtype=torch.bfloat16, device="cuda")
+    launch(out, x, 1234, stream=torch.cuda.Stream())
+    torch.cuda.synchronize()
+    f = out.float()
+
+    assert torch.all((f == lo) | (f == hi)), "outputs must be one of the two nearest bf16 values"
+    up_prob = (f == hi).float().mean().item()
+    assert abs(up_prob - frac) < 0.01, f"up_prob {up_prob} far from frac {frac}"
+    assert abs(f.mean().item() - x) < 1e-4
+
+
+@pytest.mark.l2_device
+@pytest.mark.rocm_lower
+@pytest.mark.skipif(torch is None or not torch.cuda.is_available(), reason="requires GPU")
+def test_philox_4x32_uniform():
+    BLOCK = 256
+    NBLOCKS = 1024
+    P = BLOCK * NBLOCKS
+
+    @flyc.kernel(known_block_size=[BLOCK, 1, 1])
+    def kernel(Out: fx.Tensor, seed: fx.Int32):
+        g = fx.block_idx.x * BLOCK + fx.thread_idx.x
+        fx.memref_store(fx.Int32(fx.philox_4x32(fx.Uint32(g), fx.Uint32(seed))[0]), Out, g)
+
+    @flyc.jit
+    def launch(Out: fx.Tensor, seed: fx.Int32, stream: fx.Stream = fx.Stream(None)):
+        kernel(Out, seed).launch(grid=(NBLOCKS, 1, 1), block=(BLOCK, 1, 1), stream=stream)
+
+    out = torch.zeros(P, dtype=torch.int32, device="cuda")
+    launch(out, 7, stream=torch.cuda.Stream())
+    torch.cuda.synchronize()
+
+    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)