diff --git a/batch_invariant_ops/batch_invariant_ops.py b/batch_invariant_ops/batch_invariant_ops.py index b9021bb..12d850c 100644 --- a/batch_invariant_ops/batch_invariant_ops.py +++ b/batch_invariant_ops/batch_invariant_ops.py @@ -101,7 +101,11 @@ def matmul_kernel_persistent( a = tl.load(a_ptrs, mask=offs_k_for_mask[None, :] < K - ki * BLOCK_SIZE_K, other=0.0) b = tl.load(b_ptrs, mask=offs_k_for_mask[:, None] < K - ki * BLOCK_SIZE_K, other=0.0) - accumulator = tl.dot(a, b, accumulator) + # input_precision="ieee": fp32 inputs must use full IEEE precision, not TF32. + # tl.dot defaults to TF32 for fp32 (10-bit mantissa; integers exact only to + # 2**11=2048), which silently downgrades fp32 matmuls vs torch.mm's default + # IEEE path. No-op for bf16/fp16 (TF32 only applies to fp32 inputs). + accumulator = tl.dot(a, b, accumulator, input_precision="ieee") tile_id_c += NUM_SMS pid_m, pid_n = _compute_pid(tile_id_c, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS) diff --git a/test_batch_invariance.py b/test_batch_invariance.py index 912e113..1b271b4 100644 --- a/test_batch_invariance.py +++ b/test_batch_invariance.py @@ -33,6 +33,29 @@ def run_iters(iters=10): print( f"Batch Deterministic: {is_deterministic} run-to-run max/min/diff {max(difflist)}/{min(difflist)}/{max(difflist)-min(difflist)} for {dtype} in {iters} iterations") +def test_fp32_matmul_precision(): + """fp32 matmul must keep IEEE precision, not silently fall back to TF32. + + tl.dot defaults to TF32 for fp32 inputs (10-bit mantissa; integers exact only to + 2**11 = 2048), whereas torch.mm uses IEEE fp32 by default. Without + input_precision="ieee" in matmul_kernel_persistent, every fp32 mm/addmm under + batch-invariant mode loses precision for values > 2048 -- e.g. it corrupts the + rotary-embedding phase (positions @ inv_freq) for long context. Returns the count + of mismatched elements against the exact integer reference (0 == correct). + """ + pos = torch.arange(8192, dtype=torch.float32) # every value exact in fp32 + ones = torch.ones(1, 1, dtype=torch.float32) + with set_batch_invariant_mode(True): + out = torch.mm(ones, pos[None, :]) + mismatches = int((out[0] != pos).sum()) + assert mismatches == 0, ( + f"fp32 matmul lost precision: {mismatches} elements wrong, " + f"e.g. position 2049 -> {out[0, 2049].item()} (expected 2049.0). " + f"matmul_kernel_persistent is using TF32 for fp32 inputs." + ) + return mismatches + + # Test with standard PyTorch (likely to show differences) print("Standard PyTorch:") with set_batch_invariant_mode(False): @@ -41,3 +64,7 @@ def run_iters(iters=10): print("\nBatch-Invariant Mode:") with set_batch_invariant_mode(True): run_iters() + +# fp32 precision: matmul must not silently downgrade to TF32 (see issue #23) +print("\nfp32 precision (TF32 regression):") +print(f" fp32 matmul mismatches vs IEEE reference: {test_fp32_matmul_precision()} (expected 0)")