From 4aa9f6fe9d8262cb914a299d39cb655522888abb Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 18 Jul 2026 12:31:30 +0000
Subject: [PATCH] Pin the fp8 weight-quantize kernel against silent MSLK
 switching

torchao's Float8Tensor KernelPreference defaults to AUTO, which switches
the weight-quantize kernel to MSLK whenever an mslk package is importable
on sm90+. Measured on B200: that changes fp8 scale rounding bitwise (8/8
FLUX matrices differ, scales ~55 percent of bytes), so a box that merely
gains mslk would break the hosted-prequant bit-identity invariant; the
mslk path is also slower under torch.compile (opaque extern call blocks
inductor's quantize fusion, FLUX.1 fp8 e2e 1.149 to 1.624 s). Pin
KernelPreference.TORCH explicitly, matching current no-mslk behaviour
bit for bit; signature-gated for older torchao. GPU-smoked (finite,
rel err 0.037) and pinned by test.
---
 .../inference/diffusion_transformer_quant.py  | 21 +++++++++++++++----
 .../tests/test_diffusion_transformer_quant.py | 16 ++++++++++++++
 2 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/studio/backend/core/inference/diffusion_transformer_quant.py b/studio/backend/core/inference/diffusion_transformer_quant.py
index 656edb61ef..5c3afc780f 100644
--- a/studio/backend/core/inference/diffusion_transformer_quant.py
+++ b/studio/backend/core/inference/diffusion_transformer_quant.py
@@ -460,11 +460,24 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any:
         from torchao.quantization import PerRow
 
         fp8_kwargs: dict = {"granularity": PerRow()}
-        if (
-            "activation_value_lb"
-            in inspect.signature(Float8DynamicActivationFloat8WeightConfig).parameters
-        ):
+        config_params = inspect.signature(Float8DynamicActivationFloat8WeightConfig).parameters
+        if "activation_value_lb" in config_params:
             fp8_kwargs["activation_value_lb"] = 1e-12
+        # Pin the plain-torch quantize kernel. The default AUTO silently switches to the MSLK
+        # kernel whenever an mslk package is importable (sm90+), which changes fp8 scale
+        # rounding BITWISE (measured: 8/8 FLUX matrices differ, scales ~55% of bytes) -- so a
+        # box that merely gains mslk would break the hosted-prequant bit-identity invariant.
+        # Measured on B200 the mslk path is also SLOWER compiled (opaque extern call blocks
+        # inductor's quantize fusion: FLUX.1 fp8 e2e 1.149 -> 1.624 s), so the pin costs nothing.
+        if "kernel_preference" in config_params:
+            try:
+                from torchao.quantization.quantize_.common.kernel_preference import (
+                    KernelPreference,
+                )
+
+                fp8_kwargs["kernel_preference"] = KernelPreference.TORCH
+            except Exception:  # noqa: BLE001 — enum moved: keep the library default
+                pass
         try:
             from torchao.float8 import Float8MMConfig
             return Float8DynamicActivationFloat8WeightConfig(
diff --git a/studio/backend/tests/test_diffusion_transformer_quant.py b/studio/backend/tests/test_diffusion_transformer_quant.py
index e925cf14eb..31236394b2 100644
--- a/studio/backend/tests/test_diffusion_transformer_quant.py
+++ b/studio/backend/tests/test_diffusion_transformer_quant.py
@@ -516,6 +516,22 @@ def test_fp8_config_uses_per_row_granularity():
     assert grans and all(isinstance(g, per_row) for g in grans), f"expected all PerRow, got {gran}"
 
 
+def test_fp8_config_pins_torch_kernel_preference():
+    """FP8 must pin KernelPreference.TORCH. The AUTO default silently switches the weight
+    quantize to the MSLK kernel whenever an mslk package is importable, which changes fp8
+    scale rounding bitwise (measured 8/8 FLUX matrices differ) and would break the hosted
+    prequant bit-identity invariant; the mslk path is also slower under torch.compile."""
+    pytest.importorskip("torchao.quantization")
+    try:
+        from torchao.quantization.quantize_.common.kernel_preference import KernelPreference
+    except Exception:
+        pytest.skip("torchao build without KernelPreference")
+    cfg = tq._make_quant_config(TQ_FP8)
+    if not hasattr(cfg, "kernel_preference"):
+        pytest.skip("torchao config without kernel_preference")
+    assert cfg.kernel_preference == KernelPreference.TORCH
+
+
 def test_quantize_transformer_applies_and_marks(monkeypatch):
     monkeypatch.setattr(
         tq, "select_transformer_quant_scheme", lambda target, mode, family = None: TQ_FP8