[diffusion] Clean up shared bitexact gates, helpers, and stale naming (#34180)

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

test/registered/kernels/ops/diffusion/test_bitexact_gate.py

@@ -0,0 +1,80 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion.bitexact_gate import (
+    BitExactFusionGate,
+    tensors_equal,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+def test_bitexact_gate_once_mode_verifies_then_reuses():
+    gate = BitExactFusionGate("once")
+    calls = {"fused": 0, "ref": 0}
+
+    def fused():
+        calls["fused"] += 1
+        return torch.tensor([1.0])
+
+    def ref():
+        calls["ref"] += 1
+        return torch.tensor([1.0])
+
+    assert torch.equal(gate.accept_or_fallback(fused(), ref()), torch.tensor([1.0]))
+    assert gate.verified and not gate.disabled and calls == {"fused": 1, "ref": 1}
+    assert torch.equal(fused(), torch.tensor([1.0]))
+    assert calls == {"fused": 2, "ref": 1}
+
+
+def test_bitexact_gate_mismatch_disables_permanently():
+    gate = BitExactFusionGate("mismatch")
+
+    out = gate.accept_or_fallback(
+        torch.tensor([1.0]),
+        torch.tensor([2.0]),
+        mismatch_msg="mismatch",
+    )
+    assert torch.equal(out, torch.tensor([2.0]))
+    assert gate.disabled and not gate.verified
+
+
+def test_bitexact_gate_per_signature_tracks_each_sig():
+    gate = BitExactFusionGate("sig", per_signature=True)
+    a = torch.tensor([1.0])
+    assert torch.equal(gate.accept_or_fallback(a, a, sig=("a",)), a)
+    assert gate.is_verified(("a",))
+    assert not gate.is_verified(("b",))
+    assert torch.equal(gate.accept_or_fallback(a, a, sig=("b",)), a)
+    assert gate.verified_sigs == {("a",), ("b",)}
+
+
+def test_bitexact_gate_skips_first_sight_during_graph_capture(monkeypatch):
+    # Negative-branch contract: an unverified gate must not attempt first-sight
+    # verification inside CUDA graph capture — the eager-reference host sync
+    # would abort the capture (and BCG would permanently block the signature).
+    gate = BitExactFusionGate("capture")
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
+    assert not gate.can_attempt_once()
+    # A verified gate replays the fused kernel alone, which is capture-safe.
+    gate.mark_verified()
+    assert gate.can_attempt_once()
+
+
+def test_tensors_equal_supports_sequences():
+    assert tensors_equal(
+        (torch.tensor([1.0]), torch.tensor([2.0])),
+        (torch.tensor([1.0]), torch.tensor([2.0])),
+    )
+    assert not tensors_equal(
+        (torch.tensor([1.0]), torch.tensor([2.0])),
+        (torch.tensor([1.0]), torch.tensor([3.0])),
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernels/ops/diffusion/test_ernie_norm_scale_shift.py

@@ -47,10 +47,10 @@ def test_fused_norm_scale_shift_is_bit_exact(shape):
         assert torch.equal(out2, ref2)
 
     # the fast paths must actually be in use (not silently disabled)
-    assert ernie_image._ERNIE_FUSED_NORM_VERIFIED
-    assert ernie_image._ERNIE_FUSED_GATED_NORM_VERIFIED
-    assert not ernie_image._ERNIE_FUSED_NORM_DISABLED
-    assert not ernie_image._ERNIE_FUSED_GATED_NORM_DISABLED
+    assert ernie_image._ERNIE_NORM.verified
+    assert ernie_image._ERNIE_GATED_NORM.verified
+    assert not ernie_image._ERNIE_NORM.disabled
+    assert not ernie_image._ERNIE_GATED_NORM.disabled
 
 
 if __name__ == "__main__":

test/registered/kernels/ops/diffusion/test_flux_ln_modulate.py

@@ -48,8 +48,8 @@ def test_flux_fused_ln_modulate_is_bit_exact(shape, chunks):
     out = _flux_fused_ln_modulate(norm, x, scale, shift)
     assert out is not None
     assert torch.equal(out, _eager(norm, x, scale, shift))
-    assert not flux._FLUX_FUSED_LN_MOD_DISABLED
-    assert flux._FLUX_FUSED_LN_MOD_VERIFIED
+    assert not flux._FLUX_LN_MOD.disabled
+    assert flux._FLUX_LN_MOD.verified
 
 
 def test_flux_norm_modulate_bitexact_supersedes_high_fold():

test/registered/kernels/ops/diffusion/test_glm_image_ln_modulate.py

@@ -28,8 +28,8 @@ def test_fused_ln_modulate_is_bit_exact(shape):
     shift, scale = chunks[0], chunks[2]
     out = _glm_ln_modulate(norm, x, scale, shift, x.dtype)
     assert torch.equal(out, _eager_ln_modulate(norm, x, scale, shift, x.dtype))
-    assert glm_image._GLM_FUSED_LN_MOD_VERIFIED
-    assert not glm_image._GLM_FUSED_LN_MOD_DISABLED
+    assert glm_image._GLM_LN_MOD.verified
+    assert not glm_image._GLM_LN_MOD.disabled
 
 
 @pytest.mark.parametrize("shape", [(1, 4360, 32, 128), (2, 37, 3, 40), (1, 129, 5, 64)])
@@ -45,8 +45,8 @@ def test_fused_qk_head_layernorm_is_bit_exact(shape):
     q_out, k_out = _glm_qk_layernorm(norm_q, norm_k, q, k, q.dtype)
     assert torch.equal(q_out, norm_q(q).to(q.dtype))
     assert torch.equal(k_out, norm_k(k).to(k.dtype))
-    assert glm_image._GLM_FUSED_QK_LN_VERIFIED
-    assert not glm_image._GLM_FUSED_QK_LN_DISABLED
+    assert glm_image._GLM_QK_LN.verified
+    assert not glm_image._GLM_QK_LN.disabled
 
 
 if __name__ == "__main__":

test/registered/kernels/ops/diffusion/test_sana_ln_modulate.py

@@ -36,17 +36,17 @@ def test_sana_fused_ln_modulate_is_bit_exact(shape, nmod, transposed):
     emb = torch.randn(batch, nmod, hidden, device="cuda").bfloat16()
     shift, scale = emb.chunk(nmod, dim=1)[0], emb.chunk(nmod, dim=1)[-1]
     # default-stream eager serving must stay on the untouched eager chain
-    n_sigs = len(sana._SANA_FUSED_LN_MOD_OK_SIGS)
+    n_sigs = len(sana._SANA_LN_MOD.verified_sigs)
     _sana_ln_modulate(norm, x, scale, shift)
-    assert len(sana._SANA_FUSED_LN_MOD_OK_SIGS) == n_sigs
+    assert len(sana._SANA_LN_MOD.verified_sigs) == n_sigs
     # the fusion engages on non-default streams (the BCG warmup/capture path)
     with torch.cuda.stream(torch.cuda.Stream()):
         out = _sana_ln_modulate(norm, x, scale, shift)
-        assert len(sana._SANA_FUSED_LN_MOD_OK_SIGS) == n_sigs + 1  # verified
+        assert len(sana._SANA_LN_MOD.verified_sigs) == n_sigs + 1  # verified
         out2 = _sana_ln_modulate(norm, x, scale, shift)  # verified-sig lane
     torch.cuda.synchronize()
     assert torch.equal(out, _eager_ln_modulate(norm, x, scale, shift))
-    assert torch.equal(out2, out) and not sana._SANA_FUSED_LN_MOD_DISABLED
+    assert torch.equal(out2, out) and not sana._SANA_LN_MOD.disabled
 
 
 if __name__ == "__main__":