Use input_precision="ieee" in matmul_kernel_persistent for fp32 inputs

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)

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)")