diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index a44de06fb5..0b87dbdd6f 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/studio/backend/core/inference/diffusion_precision.py @@ -140,12 +140,51 @@ def _te_family_denied(family: Optional[str], scheme: str) -> bool: return scheme in _TE_FAMILY_SCHEME_DENY.get((family or "").strip().lower(), frozenset()) +# nvfp4 TE casts WEIGHT-ONLY (see _cast_nvfp4), a different torchao kernel from the transformer's +# dynamic-activation NVFP4 GEMM, so it gets its own cached smoke probe. +_TE_NVFP4_PROBE_CACHE: dict[str, bool] = {} + + +def _te_nvfp4_weightonly_probe(device: str) -> bool: + """True iff weight-only NVFP4 (the config ``_cast_nvfp4`` applies) runs one forward on this + build. The transformer ``_smoke_probe`` tests the DYNAMIC-activation NVFP4 GEMM instead, a + different kernel: a Blackwell build can carry the weight-only FP4 path without the dynamic + one, so an explicit TE ``nvfp4`` request needs this dedicated probe to avoid falsely staying + dense when the caster would in fact run. Cached per device.""" + if device in _TE_NVFP4_PROBE_CACHE: + return _TE_NVFP4_PROBE_CACHE[device] + ok = False + try: + import torch + from torchao.prototype.mx_formats import NVFP4WeightOnlyConfig + from torchao.quantization import quantize_ + + from .diffusion_transformer_quant import make_filter_fn + + lin = torch.nn.Linear(512, 512, bias = False).to(device = device, dtype = torch.bfloat16) + quantize_(lin, NVFP4WeightOnlyConfig(), filter_fn = make_filter_fn(0)) + x = torch.randn(32, 512, device = device, dtype = torch.bfloat16) + with torch.no_grad(): + lin(x) + torch.cuda.synchronize() + ok = True + except Exception: # noqa: BLE001 -- an unavailable kernel just means stay dense + ok = False + _TE_NVFP4_PROBE_CACHE[device] = ok + return ok + + def _te_scheme_probe(scheme: str, device: str) -> bool: """True iff ``scheme`` actually runs on this build. Layerwise fp8 (no torchao GEMM) always - passes; the torchao TE modes reuse the transformer module's cached quantise+matmul smoke test.""" + passes; nvfp4 probes its weight-only kernel (``_cast_nvfp4``); the other torchao TE modes + reuse the transformer module's cached quantise+matmul smoke test (same dynamic config).""" tq = _TE_SMOKE_SCHEME.get(scheme) if tq is None: return True + # nvfp4 TE casts weight-only, whereas the transformer smoke probe tests the dynamic-activation + # NVFP4 GEMM: probe the kernel that will actually run so a weight-only-only build is not skipped. + if scheme == TE_QUANT_NVFP4: + return _te_nvfp4_weightonly_probe(device) try: from .diffusion_transformer_quant import _smoke_probe return _smoke_probe(tq, device) diff --git a/studio/backend/tests/test_diffusion_precision.py b/studio/backend/tests/test_diffusion_precision.py index 7a64139618..e80fa4a898 100644 --- a/studio/backend/tests/test_diffusion_precision.py +++ b/studio/backend/tests/test_diffusion_precision.py @@ -88,6 +88,8 @@ def _stub_casters(monkeypatch, recorder): # by default so these caster tests exercise the cast, not a broken-kernel fallback. dtq._smoke_probe = lambda tq, device: True monkeypatch.setitem(sys.modules, "core.inference.diffusion_transformer_quant", dtq) + # nvfp4 TE probes its own weight-only kernel (not the dynamic _smoke_probe); pass it too. + monkeypatch.setattr(dp, "_te_nvfp4_weightonly_probe", lambda device: True) # ── normalisation ───────────────────────────────────────────────────────────── @@ -288,6 +290,23 @@ def test_te_scheme_probe_bypasses_layerwise_fp8(): assert scheme in dp._TE_SMOKE_SCHEME +def test_te_scheme_probe_nvfp4_uses_weightonly_kernel(monkeypatch): + # nvfp4 TE casts weight-only (_cast_nvfp4 -> NVFP4WeightOnlyConfig), a different torchao kernel + # from the transformer's dynamic-activation NVFP4 probe. On a build where the dynamic FP4 GEMM + # is unavailable but weight-only FP4 works, the nvfp4 TE probe must consult its own weight-only + # probe, not the transformer dynamic one, or an explicit request would falsely stay dense. + dp._TE_NVFP4_PROBE_CACHE.clear() + dtq = types.ModuleType("core.inference.diffusion_transformer_quant") + dtq._smoke_probe = lambda scheme, device: False # every dynamic-activation GEMM "unavailable" + monkeypatch.setitem(sys.modules, "core.inference.diffusion_transformer_quant", dtq) + monkeypatch.setattr(dp, "_te_nvfp4_weightonly_probe", lambda device: True) + # nvfp4 follows its own weight-only probe (True), not the transformer dynamic probe (False). + assert dp._te_scheme_probe(TE_QUANT_NVFP4, "cuda") is True + # int8 / fp8_dynamic still follow the (dynamic) transformer probe -> False here. + assert dp._te_scheme_probe(TE_QUANT_INT8, "cuda") is False + assert dp._te_scheme_probe(TE_QUANT_FP8_DYNAMIC, "cuda") is False + + def test_quantize_int8_unsupported_hw_is_noop(monkeypatch): # int8 on pre-Ampere silicon (no int8 tensor cores) applies nothing. _stub_torch(monkeypatch, cc = (7, 5))