diff --git a/kernels/gemm/rdna3_f16_gemm.py b/kernels/gemm/rdna3_f16_gemm.py index ef6a65f68..00869e524 100644 --- a/kernels/gemm/rdna3_f16_gemm.py +++ b/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), diff --git a/kernels/gemm/rdna_f16_gemm.py b/kernels/gemm/rdna_f16_gemm.py index ca8217cc3..6e55af704 100644 --- a/kernels/gemm/rdna_f16_gemm.py +++ b/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), diff --git a/python/flydsl/expr/numeric.py b/python/flydsl/expr/numeric.py index 6d1b826a3..0890f1db9 100644 --- a/python/flydsl/expr/numeric.py +++ b/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) diff --git a/python/flydsl/expr/typing.py b/python/flydsl/expr/typing.py index 4b6ebcbd9..17b72bf48 100644 --- a/python/flydsl/expr/typing.py +++ b/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", diff --git a/tests/kernels/test_rdna_gemm.py b/tests/kernels/test_rdna_gemm.py index 9744075c9..b6bbf5db4 100644 --- a/tests/kernels/test_rdna_gemm.py +++ b/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", [ diff --git a/tests/unit/test_stochastic_rounding.py b/tests/unit/test_stochastic_rounding.py new file mode 100644 index 000000000..4b5d21dfe --- /dev/null +++ b/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)