fix(te-quant): probe the weight-only NVFP4 kernel for explicit TE nvfp4

The explicit text_encoder_quant=nvfp4 path gated on the transformer
smoke probe, which builds the dynamic-activation NVFP4 config, while the
TE caster _cast_nvfp4 applies weight-only NVFP4WeightOnlyConfig. On a
Blackwell build that carries the weight-only FP4 path but not the
dynamic FP4 GEMM, the probe would fail and the encoder would silently
stay dense even though the caster would run. Add a dedicated weight-only
NVFP4 smoke probe (mirroring _cast_nvfp4's config) and route TE nvfp4
through it; int8 / fp8_dynamic keep the dynamic transformer probe since
their TE casters are also dynamic-activation.

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)

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))