diff --git a/python/sglang/kernels/ops/diffusion/bitexact_gate.py b/python/sglang/kernels/ops/diffusion/bitexact_gate.py index d0ac7badd..a7be00ec9 100644 --- a/python/sglang/kernels/ops/diffusion/bitexact_gate.py +++ b/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", "") + 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( - mismatch_msg - or ( - f"{self.name} fast path is not bit-exact against this " - "platform's reference dispatch; falling back to eager" - ) + message = 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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py b/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py index 77e2339d2..3165efaae 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ernie_image.py +++ b/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) diff --git a/test/registered/kernels/ops/diffusion/test_bitexact_gate.py b/test/registered/kernels/ops/diffusion/test_bitexact_gate.py index aa30fcfa7..8b98c653b 100644 --- a/test/registered/kernels/ops/diffusion/test_bitexact_gate.py +++ b/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"]))