[Diffusion] Improve bit-exact fusion fallback diagnostics (#34412)

python/sglang/kernels/ops/diffusion/bitexact_gate.py

@@ -10,6 +10,54 @@ import torch
 
 T = TypeVar("T")
 EqualFn = Callable[[Any, Any], bool]
+DiagnosticHintFn = Callable[[], str | None]
+
+
+def flashinfer_rmsnorm_diagnostic_hint() -> str:
+    """Describe the live FlashInfer RMSNorm backend after an exactness miss.
+
+    Keep the imports and metadata lookups inside this function: callers pass it
+    as a callback, so none of this work runs on the verified steady-state path.
+    """
+    import importlib
+    import importlib.metadata
+    import os
+
+    try:
+        flashinfer_norm = importlib.import_module("flashinfer.norm")
+        use_cuda_norm = getattr(flashinfer_norm, "_USE_CUDA_NORM", None)
+    except Exception:
+        backend = "unavailable"
+    else:
+        if use_cuda_norm is True:
+            backend = "CUDA JIT"
+        elif use_cuda_norm is False:
+            backend = "CuTe DSL"
+        else:
+            backend = "legacy or unknown (no _USE_CUDA_NORM flag)"
+
+    versions = []
+    for package in (
+        "flashinfer-python",
+        "flashinfer-cubin",
+        "flashinfer-jit-cache",
+    ):
+        try:
+            package_version = importlib.metadata.version(package)
+        except importlib.metadata.PackageNotFoundError:
+            package_version = "not installed"
+        except Exception:
+            package_version = "unknown"
+        versions.append(f"{package}={package_version}")
+
+    env_backend = os.environ.get("FLASHINFER_USE_CUDA_NORM", "<unset>")
+    return (
+        "RMSNorm exactness can change when FlashInfer selects a different "
+        f"reduction backend. Detected backend={backend}, "
+        f"FLASHINFER_USE_CUDA_NORM={env_backend}, {', '.join(versions)}. "
+        "Check that the FlashInfer packages are version-aligned and that the "
+        "expected RMSNorm backend is selected"
+    )
 
 
 class BitExactFusionGate:
@@ -91,6 +139,7 @@ class BitExactFusionGate:
         equal: EqualFn | None = None,
         logger: logging.Logger | None = None,
         mismatch_msg: str | None = None,
+        diagnostic_hint: DiagnosticHintFn | None = None,
     ) -> T:
         """Return ``out`` when bit-exact; otherwise disable and return ``ref``."""
         if self.is_verified(sig):
@@ -100,13 +149,23 @@ class BitExactFusionGate:
             self.mark_verified(sig)
             return out
         if logger is not None:
-            logger.warning_once(
+            message = mismatch_msg or (
-                mismatch_msg
-                or (
                 f"{self.name} fast path is not bit-exact against this "
                 "platform's reference dispatch; falling back to eager"
             )
+            details = (
+                "Correctness is preserved because the eager reference output "
+                "is used. A platform-specific reference kernel or reduction-order "
+                "change may have caused this fallback"
             )
+            if diagnostic_hint is not None:
+                try:
+                    diagnostic_details = diagnostic_hint()
+                except Exception:
+                    diagnostic_details = None
+                if diagnostic_details:
+                    details = f"{details}. {diagnostic_details.rstrip('.')}"
+            logger.warning_once(f"{message.rstrip('.')}. {details}.")
         self.disable()
         return ref
 

python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py

@@ -24,6 +24,7 @@ from sglang.kernels.ops.activation.activation import (
 )
 from sglang.kernels.ops.diffusion.bitexact_gate import (
     BitExactFusionGate,
+    flashinfer_rmsnorm_diagnostic_hint,
     tensors_equal,
 )
 from sglang.kernels.ops.diffusion.residual_gate_add import residual_gate_add
@@ -111,6 +112,7 @@ def _ernie_norm_scale_shift(
                     "ERNIE fused-norm fast path is not bit-exact against this "
                     "platform's rmsnorm dispatch; falling back to eager"
                 ),
+                diagnostic_hint=flashinfer_rmsnorm_diagnostic_hint,
             )
 
     return _eager_norm_scale_shift(norm, x, scale, shift)
@@ -165,6 +167,7 @@ def _ernie_gated_norm_scale_shift(
                     "ERNIE fused gated-norm fast path is not bit-exact against "
                     "this platform's rmsnorm dispatch; falling back to eager"
                 ),
+                diagnostic_hint=flashinfer_rmsnorm_diagnostic_hint,
             )
 
     res = residual_gate_add(residual, update, gate)

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

@@ -1,13 +1,17 @@
 import sys
+from types import ModuleType
+from unittest.mock import MagicMock, patch
 
 import pytest
 import torch
 
 from sglang.kernels.ops.diffusion.bitexact_gate import (
     BitExactFusionGate,
+    flashinfer_rmsnorm_diagnostic_hint,
     tensors_equal,
 )
 from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
 
 register_cpu_ci(est_time=2, suite="base-a-test-cpu")
 
@@ -76,5 +80,80 @@ def test_tensors_equal_supports_sequences():
     )
 
 
+class TestBitExactFallbackDiagnostics(CustomTestCase):
+    def test_mismatch_warning_is_actionable_and_diagnostic_is_lazy(self):
+        logger = MagicMock()
+        diagnostic = MagicMock(return_value="backend=CuTe DSL")
+        gate = BitExactFusionGate("diagnostic")
+
+        matched = gate.accept_or_fallback(
+            torch.tensor([1.0]),
+            torch.tensor([1.0]),
+            logger=logger,
+            diagnostic_hint=diagnostic,
+        )
+        self.assertTrue(torch.equal(matched, torch.tensor([1.0])))
+        diagnostic.assert_not_called()
+        logger.warning_once.assert_not_called()
+
+        gate = BitExactFusionGate("diagnostic")
+        fallback = gate.accept_or_fallback(
+            torch.tensor([1.0]),
+            torch.tensor([2.0]),
+            logger=logger,
+            diagnostic_hint=diagnostic,
+        )
+
+        self.assertTrue(torch.equal(fallback, torch.tensor([2.0])))
+        diagnostic.assert_called_once_with()
+        warning = logger.warning_once.call_args.args[0]
+        self.assertIn("Correctness is preserved", warning)
+        self.assertIn("reference kernel or reduction-order change", warning)
+        self.assertIn("backend=CuTe DSL", warning)
+
+    def test_diagnostic_failure_cannot_break_the_eager_fallback(self):
+        logger = MagicMock()
+
+        def broken_diagnostic():
+            raise RuntimeError("diagnostics unavailable")
+
+        gate = BitExactFusionGate("diagnostic")
+        fallback = gate.accept_or_fallback(
+            torch.tensor([1.0]),
+            torch.tensor([2.0]),
+            logger=logger,
+            diagnostic_hint=broken_diagnostic,
+        )
+
+        self.assertTrue(torch.equal(fallback, torch.tensor([2.0])))
+        self.assertTrue(gate.disabled)
+        self.assertIn("Correctness is preserved", logger.warning_once.call_args.args[0])
+
+    def test_flashinfer_rmsnorm_hint_reports_backend_and_versions(self):
+        flashinfer = ModuleType("flashinfer")
+        flashinfer_norm = ModuleType("flashinfer.norm")
+        flashinfer_norm._USE_CUDA_NORM = False
+        versions = {
+            "flashinfer-python": "0.6.12",
+            "flashinfer-cubin": "0.6.12",
+            "flashinfer-jit-cache": "0.6.12+cu130",
+        }
+
+        with (
+            patch.dict(
+                sys.modules,
+                {"flashinfer": flashinfer, "flashinfer.norm": flashinfer_norm},
+            ),
+            patch("importlib.metadata.version", side_effect=versions.__getitem__),
+            patch.dict("os.environ", {"FLASHINFER_USE_CUDA_NORM": "0"}),
+        ):
+            hint = flashinfer_rmsnorm_diagnostic_hint()
+
+        self.assertIn("backend=CuTe DSL", hint)
+        self.assertIn("FLASHINFER_USE_CUDA_NORM=0", hint)
+        for package, version in versions.items():
+            self.assertIn(f"{package}={version}", hint)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v"]))