cublaslt: gate bf16/f16 cast-F32 fusion on ?beta = 0.0

The mixed-dtype rules in `cublaslt_mixed_dtype_rewrite.egg` rewrote
`Cast(matmul, F32)` over a low-precision matmul into a single cuBLASLt
op that emits F32 directly. The action kept the original `?inputs`
list. When a residual block had already been beta-fused into the matmul
(so `?beta = 1.0` and `?c` is a bf16/f16 residual tensor in the inputs),
the rule produced a sibling op declaring `c_dtype = F32` over that same
low-precision buffer. cuBLASLt read bf16/f16 bytes through F32 alignment,
scrambling the residual contribution and silently collapsing the block
output to ~residual-only magnitudes (~0.5 vs ~24 on Llama-3.2-1B layer 0).

The bug fired non-deterministically (≈40 % of fresh Python processes on
a 2-block chain reproducer) because egglog's `FxHashMap` iteration in
`SerializedEGraph::new` randomized which e-node the extractor picked.

Restricting both rules to `?beta = 0.0` rules out the failing case
exactly: with beta=0 cuBLASLt skips the C read entirely, so the
c_dtype/c_buffer mismatch is harmless. For beta!=0 matmuls a proper
fix would need to also cast `?c` to F32 in the new input list, which
is a deeper rewrite — not done here.

Also flips the five full-sized Llama tests in `tests/test_llama3.py`
from `torch_dtype=torch.float32` back to the checkpoint's native bf16,
relaxing their `atol` from 1e-5 / 1e-4 / 1e-3 to bf16-appropriate
0.5 (1B) / 2.0 (8B). All five pass cleanly on the bf16 path.

New regression test `tests/test_bf16_chain_block.py` is the minimal
two-residual-block module that fired the bug.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

crates/luminal_cuda_lite/src/host/cublaslt/cublaslt_mixed_dtype_rewrite.egg

@@ -7,6 +7,9 @@
 ; Luminal graphs express this today as a Cast(F32) around a low-precision
 ; matmul. cuBLASLt can write the f32 output directly, so expose that candidate
 ; before beta fusion tries to consume an f32 C input.
+;
+; `?beta = 0.0` guard: with non-zero beta the same `?inputs` C is read at
+; F32 over a low-precision buffer. Repro: tests/test_bf16_chain_block.py.
 
 (rule
     (
@@ -19,7 +22,7 @@
             ?stride_a ?stride_b ?stride_c ?stride_d
             (F16) (F16) (F16) (F16)
             ?compute_type ?scale_dtype
-            ?alpha ?beta ?epilogue)
+            ?alpha 0.0 ?epilogue)
             ?inputs))
         (= ?cast (Op (Cast ?size (F32)) (ICons ?matmul (INil))))
     )
@@ -32,7 +35,7 @@
             ?stride_a ?stride_b ?stride_c ?stride_d
             (F16) (F16) (F32) (F32)
             ?compute_type ?scale_dtype
-            ?alpha ?beta ?epilogue)
+            ?alpha 0.0 ?epilogue)
             ?inputs))
         (union ?cast ?fused)
         (set (dtype ?fused) (F32))
@@ -52,7 +55,7 @@
             ?stride_a ?stride_b ?stride_c ?stride_d
             (Bf16) (Bf16) (Bf16) (Bf16)
             ?compute_type ?scale_dtype
-            ?alpha ?beta ?epilogue)
+            ?alpha 0.0 ?epilogue)
             ?inputs))
         (= ?cast (Op (Cast ?size (F32)) (ICons ?matmul (INil))))
     )
@@ -65,7 +68,7 @@
             ?stride_a ?stride_b ?stride_c ?stride_d
             (Bf16) (Bf16) (F32) (F32)
             ?compute_type ?scale_dtype
-            ?alpha ?beta ?epilogue)
+            ?alpha 0.0 ?epilogue)
             ?inputs))
         (union ?cast ?fused)
         (set (dtype ?fused) (F32))

crates/luminal_python/tests/test_bf16_chain_block.py

@@ -0,0 +1,96 @@
+"""Regression: bf16 matmul → cast(F32) → norm → bf16 matmul → bf16 residual
+add. Two residual blocks chained, with the f32-internal RMSNorm in between.
+
+This is the smallest module that reproduces the silent-wrong-output bug
+caused by the now-disabled `cublaslt bf16 matmul cast f32 output` rule
+in `crates/luminal_cuda_lite/src/host/cublaslt/cublaslt_mixed_dtype_rewrite.egg`.
+Before the rule was disabled, ~40 % of fresh Python processes produced
+output magnitudes ~10× smaller than the eager reference because the bf16
+residual add was reading an F32-laid-out buffer through a bf16 dtype lens.
+
+If this test starts failing intermittently again, the rule (or an
+equivalent one) has been re-enabled without the missing num-users guard.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from luminal.main import luminal_backend
+
+
+class RMSNorm(torch.nn.Module):
+    """Mirrors `LlamaRMSNorm`: f32-internal compute around a bf16 weight."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(dim, dtype=torch.bfloat16))
+        self.eps = 1e-5
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = x.to(torch.float32)
+        var = h.pow(2).mean(-1, keepdim=True)
+        h = h * torch.rsqrt(var + self.eps)
+        h = h.to(torch.bfloat16)
+        return self.weight * h
+
+
+class ChainBlock(torch.nn.Module):
+    """Two residual blocks. The intermediate `residual + lin1(norm1(x))` has
+    two downstream consumers: `norm2`'s cast-to-f32 and the second residual
+    add. That dual-consumer-with-cast pattern is the trigger.
+    """
+
+    def __init__(self, dim: int = 2048) -> None:
+        super().__init__()
+        self.norm1 = RMSNorm(dim)
+        self.lin1 = torch.nn.Linear(dim, dim, bias=False, dtype=torch.bfloat16)
+        self.norm2 = RMSNorm(dim)
+        self.lin2 = torch.nn.Linear(dim, dim, bias=False, dtype=torch.bfloat16)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = self.norm1(x)
+        x = self.lin1(x)
+        x = residual + x
+        residual = x
+        x = self.norm2(x)
+        x = self.lin2(x)
+        return residual + x
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="CUDA-only — exercises the cuBLASLt rule that produced the bug",
+)
+def test_bf16_residual_chain_block(device: torch.device) -> None:
+    torch.manual_seed(0)
+    dim = 2048
+    x = torch.randn(1, 6, dim, dtype=torch.bfloat16, device=device) * 0.1
+
+    block = ChainBlock(dim).eval().to(device)
+    with torch.no_grad():
+        # Initialize linears to Llama-scale weights so the bf16 path
+        # exercises realistic magnitudes (not the unit-norm default).
+        block.lin1.weight.normal_(0.0, 0.02)
+        block.lin2.weight.normal_(0.0, 0.02)
+        ref = block(x)
+
+    torch._dynamo.reset()
+    compiled = torch.compile(block, backend=luminal_backend)
+    with torch.no_grad():
+        out = compiled(x)
+
+    # bf16 precision is ~1/128 relative; allow generous absolute slack so
+    # the test catches structural regressions (the bug we hit was ~40× off,
+    # not a few-percent precision drift).
+    md = torch.max(torch.abs(out.float() - ref.float())).item()
+    ref_max = torch.abs(ref.float()).max().item()
+    rel = md / max(ref_max, 1e-6)
+    assert rel < 5e-2, (
+        f"bf16 chain block diverged: max_diff={md:.4e} "
+        f"rel={rel:.4e} out_max={out.abs().max():.4f} ref_max={ref_max:.4f} "
+        "(suspect the cublaslt bf16-cast-F32 fusion has been re-enabled "
+        "without a num-users guard)"
+    )

crates/luminal_python/tests/test_llama3.py

@@ -246,7 +246,6 @@ def test_hf_llama3_1b_decode_loop_dynamic(device: torch.device):
         LlamaForCausalLM.from_pretrained(
             "NousResearch/Llama-3.2-1B",
             config=config,
-            torch_dtype=torch.float32,
         )
         .eval()
         .to(device)
@@ -265,7 +264,8 @@ def test_hf_llama3_1b_decode_loop_dynamic(device: torch.device):
         with torch.no_grad():
             ref = model(input_ids)
             out = compiled(input_ids)
-        assert torch.allclose(out.logits, ref.logits, atol=1e-5), (
+        # bf16-grade tolerance: 16 layers × ~1/128 per-layer precision.
+        assert torch.allclose(out.logits, ref.logits, atol=0.5), (
             f"step {step}: max_diff={torch.max(torch.abs(out.logits - ref.logits)).item():.2e}"
         )
         next_token = ref.logits[0, -1, :].argmax().item()
@@ -305,14 +305,13 @@ def test_hf_llama3_full(device: torch.device):
         LlamaForCausalLM.from_pretrained(
             "NousResearch/Llama-3.2-1B",
             config=config,
-            torch_dtype=torch.float32,
         )
         .eval()
         .to(device)
     )
     n_params = sum(p.numel() for p in model.parameters())
     print(
-        f"[MODEL] Total parameters: {n_params:,} ({n_params * 4 / 1024**3:.3f} GiB in f32)"
+        f"[MODEL] Total parameters: {n_params:,} ({n_params * 2 / 1024**3:.3f} GiB in bf16)"
     )
     _gpu_mem("after model load")
 
@@ -330,7 +329,8 @@ def test_hf_llama3_full(device: torch.device):
         out = compiled(input_ids)
         _gpu_mem("after compiled forward (includes compilation)")
 
-    assert torch.allclose(out.logits, ref.logits, atol=1e-5), (
+    # bf16-grade tolerance: 16 layers × ~1/128 per-layer precision.
+    assert torch.allclose(out.logits, ref.logits, atol=0.5), (
         f"max_diff={torch.max(torch.abs(out.logits - ref.logits)).item():.2e}"
     )
 
@@ -352,7 +352,6 @@ def test_hf_llama3_large_full(device: torch.device):
         LlamaForCausalLM.from_pretrained(
             "NousResearch/Meta-Llama-3.1-8B-Instruct",
             config=config,
-            torch_dtype=torch.float32,
         )
         .eval()
         .to(device)
@@ -362,7 +361,8 @@ def test_hf_llama3_large_full(device: torch.device):
     with torch.no_grad():
         ref = model(input_ids)
         out = compiled(input_ids)
-    assert torch.allclose(out.logits, ref.logits, atol=1e-4), (
+    # bf16-grade tolerance: 32 layers × ~1/128 per-layer precision.
+    assert torch.allclose(out.logits, ref.logits, atol=2.0), (
         f"max_diff={torch.max(torch.abs(out.logits - ref.logits)).item():.2e}"
     )
 
@@ -433,7 +433,6 @@ def test_hf_llama38b_full(device: torch.device):
         LlamaForCausalLM.from_pretrained(
             "NousResearch/Meta-Llama-3.1-8B-Instruct",
             config=config,
-            torch_dtype=torch.float32,
         )
         .eval()
         .to(device)
@@ -443,7 +442,8 @@ def test_hf_llama38b_full(device: torch.device):
     with torch.no_grad():
         ref = model(input_ids)
         out = compiled(input_ids)
-    assert torch.allclose(out.logits, ref.logits, atol=1e-4), (
+    # bf16-grade tolerance: 32 layers × ~1/128 per-layer precision.
+    assert torch.allclose(out.logits, ref.logits, atol=2.0), (
         f"max_diff={torch.max(torch.abs(out.logits - ref.logits)).item():.2e}"
     )
 
@@ -500,7 +500,6 @@ def test_hf_llama38b_mark_dynamic_seq_dim_before_compile(device: torch.device):
             LlamaForCausalLM.from_pretrained(
                 "NousResearch/Meta-Llama-3.1-8B-Instruct",
                 config=config,
-                torch_dtype=torch.float32,
             )
             .eval()
             .to(device)
@@ -551,7 +550,7 @@ def test_hf_llama38b_mark_dynamic_seq_dim_before_compile(device: torch.device):
         )
 
         first_diff = torch.max(torch.abs(first_out.logits - first_ref.logits)).item()
-        assert torch.allclose(first_out.logits, first_ref.logits, atol=1e-3, rtol=0), (
+        assert torch.allclose(first_out.logits, first_ref.logits, atol=2.0, rtol=0), (
             f"seq_len=4: max_diff={first_diff:.2e}"
         )
 
@@ -568,7 +567,7 @@ def test_hf_llama38b_mark_dynamic_seq_dim_before_compile(device: torch.device):
                     config.vocab_size,
                 )
             ), f"seq_len={seq_len}: got {out.logits.shape}, expected {ref.logits.shape}"
-            assert torch.allclose(out.logits, ref.logits, atol=1e-3, rtol=0), (
+            assert torch.allclose(out.logits, ref.logits, atol=2.0, rtol=0), (
                 f"seq_len={seq_len}: "
                 f"max_diff={torch.max(torch.abs(out.logits - ref.logits)).item():.2e}"
             )