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