Tolerate one-ULP BF16 variation in correctness tests

README.md

@@ -73,7 +73,7 @@ See [BENCHMARK_H20.md](BENCHMARK_H20.md).
 bash tests/test.sh
 ```
 
-- `tests/test_fwd.py` — correctness tests (exact match against the torch reference; compared with `flash-linear-attention`)
+- `tests/test_fwd.py` — correctness tests (within one BF16 ULP of the torch reference; compared with `flash-linear-attention`)
 
 
 ## Kernel API

tests/numerics.py

@@ -0,0 +1,64 @@
+import torch
+
+
+_COMPARISON_CHUNK_SIZE = 1 << 20
+
+
+def _ordered_bfloat16(values):
+    bits = values.contiguous().view(torch.int16).to(torch.int32) & 0xFFFF
+    magnitude = bits & 0x7FFF
+    return torch.where(bits & 0x8000 != 0, 0x8000 - magnitude, 0x8000 + magnitude)
+
+
+def assert_bfloat16_close(actual, expected, *, max_ulps=1, msg=None):
+    """Assert that BF16 tensors differ by at most ``max_ulps`` representable values."""
+    if actual.shape != expected.shape:
+        raise AssertionError(f"shape mismatch: {actual.shape} != {expected.shape}")
+    if actual.dtype != torch.bfloat16 or expected.dtype != torch.bfloat16:
+        raise AssertionError(
+            f"expected bfloat16 tensors, got {actual.dtype} and {expected.dtype}"
+        )
+    if max_ulps < 0:
+        raise ValueError("max_ulps must be non-negative")
+
+    actual_flat = actual.reshape(-1)
+    expected_flat = expected.reshape(-1)
+    max_distance = torch.zeros((), dtype=torch.int32, device=actual.device)
+    finite_mismatch_count = torch.zeros((), dtype=torch.int64, device=actual.device)
+    nonfinite_mismatch_count = torch.zeros((), dtype=torch.int64, device=actual.device)
+
+    for start in range(0, actual_flat.numel(), _COMPARISON_CHUNK_SIZE):
+        end = start + _COMPARISON_CHUNK_SIZE
+        actual_chunk = actual_flat[start:end]
+        expected_chunk = expected_flat[start:end]
+        equal = actual_chunk == expected_chunk
+        finite_pair = torch.isfinite(actual_chunk) & torch.isfinite(expected_chunk)
+        distance = (
+            _ordered_bfloat16(actual_chunk) - _ordered_bfloat16(expected_chunk)
+        ).abs()
+        finite_mismatch = finite_pair & (distance > max_ulps)
+
+        max_distance = torch.maximum(
+            max_distance,
+            distance.masked_fill(~finite_pair, 0).max(),
+        )
+        finite_mismatch_count += finite_mismatch.sum()
+        nonfinite_mismatch_count += (~equal & ~finite_pair).sum()
+
+    finite_mismatches = finite_mismatch_count.item()
+    nonfinite_mismatches = nonfinite_mismatch_count.item()
+    if finite_mismatches or nonfinite_mismatches:
+        details = (
+            f"bfloat16 tensors differ by up to {max_distance.item()} ULPs "
+            f"(allowed {max_ulps}); {finite_mismatches} finite and "
+            f"{nonfinite_mismatches} non-finite values are outside the tolerance"
+        )
+        raise AssertionError(f"{msg}: {details}" if msg else details)
+
+
+def assert_matches_reference(actual, expected, *, msg):
+    """Use a one-ULP BF16 comparison while keeping other dtypes bitwise exact."""
+    if actual.dtype == torch.bfloat16 and expected.dtype == torch.bfloat16:
+        assert_bfloat16_close(actual, expected, msg=msg)
+    elif not torch.equal(actual, expected):
+        raise AssertionError(msg)

tests/test_fwd.py

@@ -4,6 +4,7 @@
 import math
 
 from torch_ref import torch_ref
+from numerics import assert_matches_reference
 
 
 # ============================================================
@@ -220,7 +221,7 @@ def moving_avg(data, w):
 # ============================================================
 
 def test_fwd():
-    """Test: cutlass kernel vs torch ref, require exact match."""
+    """Test: cutlass kernel vs torch ref with BF16-aware comparison."""
     B, T, H, D = 1, 8192, 96, 128
     LOWER_BOUND = -5.0
 
@@ -257,13 +258,19 @@ def test_fwd():
     print(f"{torch.max(out_kernel)} {torch.max(out_ref)}")
     print_error_stats("output", out_kernel, out_ref)
 
-    assert torch.equal(out_kernel, out_ref), "output mismatch between kernel and torch ref"
-    assert torch.equal(final_state_kernel, final_state_ref), "final_state mismatch between kernel and torch ref"
-    print("Success: kernel == torch ref (exact match)")
+    assert_matches_reference(
+        out_kernel, out_ref, msg="output mismatch between kernel and torch ref"
+    )
+    assert_matches_reference(
+        final_state_kernel,
+        final_state_ref,
+        msg="final_state mismatch between kernel and torch ref",
+    )
+    print("Success: kernel matches torch ref within BF16 tolerance")
 
 
 def test_fwd_varlen():
-    """Test: varlen cutlass kernel vs torch ref, require exact match."""
+    """Test: varlen cutlass kernel vs torch ref with BF16-aware comparison."""
     H, D = 96, 128
     LOWER_BOUND = -5.0
     seq_lens = [1300, 547, 2048, 963, 271, 3063]
@@ -307,9 +314,15 @@ def test_fwd_varlen():
     print(f"{torch.max(out_kernel)} {torch.max(out_ref)}")
     print_error_stats("output", out_kernel, out_ref)
 
-    assert torch.equal(out_kernel, out_ref), "output mismatch between kernel and torch ref"
-    assert torch.equal(final_state_kernel, final_state_ref), "final_state mismatch between kernel and torch ref"
-    print("Success: varlen kernel == torch ref (exact match)")
+    assert_matches_reference(
+        out_kernel, out_ref, msg="output mismatch between kernel and torch ref"
+    )
+    assert_matches_reference(
+        final_state_kernel,
+        final_state_ref,
+        msg="final_state mismatch between kernel and torch ref",
+    )
+    print("Success: varlen kernel matches torch ref within BF16 tolerance")
 
 
 @torch.inference_mode()

tests/test_fwd_full.py

@@ -1,6 +1,6 @@
 """Full correctness tests for FlashKDA forward pass.
 
-Compares cutlass kernel vs torch_ref for exact match across:
+Compares cutlass kernel vs torch_ref within one BF16 ULP across:
   - state_in / state_out: 4 combinations of None/present
   - state dtype: bf16, fp32
   - various H values (up to 256)
@@ -18,6 +18,7 @@
 
 import flash_kda
 from torch_ref import torch_ref
+from numerics import assert_matches_reference
 
 D = 128
 LOWER_BOUND = -5.0
@@ -87,11 +88,17 @@ def test_fwd_fixed(T, H, state_dtype, has_in, has_out):
               A_log=A_log, dt_bias=dt_bias, lower_bound=LOWER_BOUND,
               initial_state=init_r, final_state=final_r)
 
-    assert torch.equal(out_kernel, out_ref), \
-        f"output mismatch: T={T} H={H} dtype={state_dtype} in={has_in} out={has_out}"
+    assert_matches_reference(
+        out_kernel,
+        out_ref,
+        msg=f"output mismatch: T={T} H={H} dtype={state_dtype} in={has_in} out={has_out}",
+    )
     if final_k is not None:
-        assert torch.equal(final_k, final_r), \
-            f"final_state mismatch: T={T} H={H} dtype={state_dtype} in={has_in} out={has_out}"
+        assert_matches_reference(
+            final_k,
+            final_r,
+            msg=f"final_state mismatch: T={T} H={H} dtype={state_dtype} in={has_in} out={has_out}",
+        )
 
 
 # ---------------------------------------------------------------------------
@@ -143,11 +150,17 @@ def test_fwd_varlen(seq_lens, H, state_dtype, has_in, has_out):
               A_log=A_log, dt_bias=dt_bias, lower_bound=LOWER_BOUND,
               initial_state=init_r, final_state=final_r, cu_seqlens=cu_seqlens)
 
-    assert torch.equal(out_kernel, out_ref), \
-        f"output mismatch: seqs={seq_lens} H={H} dtype={state_dtype} in={has_in} out={has_out}"
+    assert_matches_reference(
+        out_kernel,
+        out_ref,
+        msg=f"output mismatch: seqs={seq_lens} H={H} dtype={state_dtype} in={has_in} out={has_out}",
+    )
     if final_k is not None:
-        assert torch.equal(final_k, final_r), \
-            f"final_state mismatch: seqs={seq_lens} H={H} dtype={state_dtype} in={has_in} out={has_out}"
+        assert_matches_reference(
+            final_k,
+            final_r,
+            msg=f"final_state mismatch: seqs={seq_lens} H={H} dtype={state_dtype} in={has_in} out={has_out}",
+        )
 
 
 # ---------------------------------------------------------------------------
@@ -198,11 +211,17 @@ def test_fwd_batched(B, T, H, state_dtype, has_in, has_out):
               A_log=A_log, dt_bias=dt_bias, lower_bound=LOWER_BOUND,
               initial_state=init_r, final_state=final_r)
 
-    assert torch.equal(out_kernel, out_ref), \
-        f"output mismatch: B={B} T={T} H={H} dtype={state_dtype} in={has_in} out={has_out}"
+    assert_matches_reference(
+        out_kernel,
+        out_ref,
+        msg=f"output mismatch: B={B} T={T} H={H} dtype={state_dtype} in={has_in} out={has_out}",
+    )
     if final_k is not None:
-        assert torch.equal(final_k, final_r), \
-            f"final_state mismatch: B={B} T={T} H={H} dtype={state_dtype} in={has_in} out={has_out}"
+        assert_matches_reference(
+            final_k,
+            final_r,
+            msg=f"final_state mismatch: B={B} T={T} H={H} dtype={state_dtype} in={has_in} out={has_out}",
+        )
 
 
 # ---------------------------------------------------------------------------
@@ -237,8 +256,8 @@ def test_fwd_long(T):
               A_log=A_log, dt_bias=dt_bias, lower_bound=LOWER_BOUND,
               initial_state=init_r, final_state=final_r)
 
-    assert torch.equal(out_kernel, out_ref), f"output mismatch: T={T}"
-    assert torch.equal(final_k, final_r), f"final_state mismatch: T={T}"
+    assert_matches_reference(out_kernel, out_ref, msg=f"output mismatch: T={T}")
+    assert_matches_reference(final_k, final_r, msg=f"final_state mismatch: T={T}")
 
 
 @pytest.mark.parametrize("seq_lens", LONG_VARLEN_CASES,
@@ -270,5 +289,9 @@ def test_fwd_long_varlen(seq_lens):
               A_log=A_log, dt_bias=dt_bias, lower_bound=LOWER_BOUND,
               initial_state=init_r, final_state=final_r, cu_seqlens=cu_seqlens)
 
-    assert torch.equal(out_kernel, out_ref), f"output mismatch: seqs={seq_lens}"
-    assert torch.equal(final_k, final_r), f"final_state mismatch: seqs={seq_lens}"
+    assert_matches_reference(
+        out_kernel, out_ref, msg=f"output mismatch: seqs={seq_lens}"
+    )
+    assert_matches_reference(
+        final_k, final_r, msg=f"final_state mismatch: seqs={seq_lens}"
+    )

tests/test_numerics.py

@@ -0,0 +1,37 @@
+import pytest
+import torch
+
+from numerics import assert_bfloat16_close, assert_matches_reference
+
+
+def test_assert_bfloat16_close_accepts_adjacent_values():
+    expected = torch.tensor([1.0, -1.0, 0.0, 4.0], dtype=torch.bfloat16)
+    actual = torch.tensor(
+        [1.0078125, -1.0078125, -0.0, 4.03125], dtype=torch.bfloat16
+    )
+
+    assert_bfloat16_close(actual, expected)
+
+
+def test_assert_bfloat16_close_rejects_two_ulp_difference():
+    expected = torch.tensor([1.0], dtype=torch.bfloat16)
+    actual = torch.tensor([1.015625], dtype=torch.bfloat16)
+
+    with pytest.raises(AssertionError, match="2 ULPs"):
+        assert_bfloat16_close(actual, expected)
+
+
+def test_assert_bfloat16_close_rejects_nonmatching_nonfinite_values():
+    expected = torch.tensor([float("inf"), 0.0], dtype=torch.bfloat16)
+    actual = torch.tensor([float("-inf"), float("nan")], dtype=torch.bfloat16)
+
+    with pytest.raises(AssertionError, match="non-finite"):
+        assert_bfloat16_close(actual, expected)
+
+
+def test_assert_matches_reference_keeps_float32_comparison_exact():
+    expected = torch.tensor([1.0], dtype=torch.float32)
+    actual = torch.nextafter(expected, torch.tensor([float("inf")]))
+
+    with pytest.raises(AssertionError, match="state mismatch"):
+        assert_matches_reference(actual, expected, msg="state mismatch")