Studio diffusion (Phase 14): fix int8 dense quant on Flux / Qwen (skip M=1 modulation linears)

The opt-in dense int8 transformer path crashed on Flux.1 and Qwen-Image with
'torch._int_mm: self.size(0) needs to be greater than 16, but got 1'. int8 dynamic quant
goes through torch._int_mm, which requires the activation row count M > 16. A DiT's AdaLN
modulation projections (Flux norm1.linear 3072->18432, Qwen img_mod.1 / txt_mod.1, Flux.2
*_modulation.linear) and its timestep / guidance / pooled-text conditioning embedders are
computed once from the [batch, dim] conditioning vector (M = batch = 1), not per token, so
they hit _int_mm at M=1 and crash. Their feature dims are large, so the existing
min_features filter did not exclude them.

Fix: the int8 filter now also skips any Linear whose fully-qualified name matches a
modulation / conditioning-embedder token (norm, _mod, modulation, timestep_embed,
guidance_embed, time_text_embed, pooled). These layers run at M=1 once per block and are a
negligible share of the FLOPs, so int8 keeps the full speedup on the attention / FFN layers
(M = sequence length). fp8 / nvfp4 / mxfp8 use scaled_mm, which has no M>16 limit and
quantises these layers fine, so the exclusion is int8-only. Sequence embedders
(context_embedder / x_embedder / txt_in, M = seq) are deliberately not excluded -- note
'context_embedder' contains the substring 'text_embed', which is why the token is the
specific 'time_text_embed', not 'text_embed'.

Measured on a B200 (1024px, transformer_quant=int8 + speed=default), int8 now runs on every
supported model and is the fastest dense path on Flux/Qwen (int8 runs full-rate vs fp8's
FP32-accumulate): FLUX.1-dev 9.62s eager -> 1.98s (4.86x, vs fp8 2.15s), Qwen-Image -> 1.87s
(5.57x, vs fp8 2.09s), FLUX.1-schnell -> 0.41s (3.59x). Z-Image and Flux.2-klein (already
working) are unchanged.

- diffusion_transformer_quant.py: add _INT8_EXCLUDE_NAME_TOKENS; make_filter_fn takes
  exclude_name_tokens; quantize_transformer passes it for int8 only.
- hermetic test that the int8 filter excludes the modulation / embedder linears (and keeps
  attention / FFN / sequence-embedder linears), while fp8 keeps them.
- scripts/int8_linear_probe.py: the meta-device probe used to enumerate each transformer's
  Linear layers and derive the exclusion list.

scripts/int8_linear_probe.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Examine which Linear layers the int8 dense-quant filter would select, to find the large M=1
+modulation/embedder projections that crash torch._int_mm (M>16). Loads each transformer on the
+META device (no weights, no GPU) from its base-repo config, lists nn.Linear fqn/in/out, and marks
+those that pass min_features=512. CPU-only, fast."""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "studio" / "backend"))
+
+# (label, transformer_class, base_repo)
+MODELS = [
+    ("flux.1-dev", "FluxTransformer2DModel", "black-forest-labs/FLUX.1-dev"),
+    ("qwen-image", "QwenImageTransformer2DModel", "Qwen/Qwen-Image"),
+    ("z-image", "ZImageTransformer2DModel", "Tongyi-MAI/Z-Image-Turbo"),
+    ("flux.2-klein-4b", "Flux2Transformer2DModel", "black-forest-labs/FLUX.2-klein-4B"),
+]
+MIN = 512
+
+
+def main() -> int:
+    import diffusers
+    import torch
+    from accelerate import init_empty_weights
+
+    tok = os.environ.get("HF_TOKEN")
+    for label, cls_name, base in MODELS:
+        cls = getattr(diffusers, cls_name, None)
+        if cls is None:
+            print(f"\n### {label}: {cls_name} NOT in diffusers"); continue
+        try:
+            cfg = cls.load_config(base, subfolder="transformer", token=tok)
+            with init_empty_weights():
+                model = cls.from_config(cfg)
+        except Exception as e:  # noqa: BLE001
+            print(f"\n### {label}: load failed {type(e).__name__}: {e}"); continue
+        lins = [(n, m) for n, m in model.named_modules() if isinstance(m, torch.nn.Linear)]
+        selected = [(n, m) for n, m in lins if m.in_features >= MIN and m.out_features >= MIN]
+        print(f"\n### {label}: {len(lins)} Linear, {len(selected)} pass min_features={MIN}")
+        # Heuristic: a modulation/embedder Linear is one OUTSIDE the repeated transformer blocks,
+        # i.e. its fqn does not contain a numeric block index, OR out==k*in (k>=3) AdaLN shape.
+        sus = []
+        for n, m in selected:
+            depth_idx = any(p.isdigit() for p in n.split("."))
+            ratio = m.out_features / m.in_features if m.in_features else 0
+            tag = []
+            if not depth_idx:
+                tag.append("NO-BLOCK-IDX")
+            if ratio >= 3:
+                tag.append(f"out={ratio:.0f}xin")
+            if any(t in n.lower() for t in ("norm", "embed", "time", "guidance", "modulation", "adaln", "cond")):
+                tag.append("NAME")
+            if tag:
+                sus.append((n, m.in_features, m.out_features, ",".join(tag)))
+        # Print the distinct fqn shapes (collapse block indices to {i})
+        import re
+        seen = {}
+        for n, i, o, tag in sus:
+            key = re.sub(r"\.\d+\.", ".{i}.", n)
+            seen.setdefault((key, i, o, tag), 0)
+            seen[(key, i, o, tag)] += 1
+        print(f"  SUSPECT (M=1 risk) distinct patterns:")
+        for (key, i, o, tag), cnt in sorted(seen.items()):
+            print(f"    [{cnt:>3}x] {key:55s} {i:>6}->{o:<6} [{tag}]")
+        # Also show a few non-suspect selected names for contrast (the real FLOP linears)
+        good = [n for n, m in selected if (n, m.in_features, m.out_features) not in
+                {(s[0], s[1], s[2]) for s in sus}][:6]
+        print(f"  kept-for-int8 examples: {good}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

studio/backend/core/inference/diffusion_transformer_quant.py

@@ -37,13 +37,31 @@ TQ_AUTO = "auto"
 TQ_SCHEMES = (TQ_INT8, TQ_FP8, TQ_NVFP4, TQ_MXFP8)
 TQ_MODES = (TQ_AUTO,) + TQ_SCHEMES
 
-# Skip linears whose in/out features are below this. The int8 dynamic path uses
-# torch._int_mm, which requires the activation row count M > 16, and the DiT's tiny
-# timestep / pooled / modulation projections run at M=1 and crash it. They are a
-# negligible share of the FLOPs, so leaving them bf16 costs ~nothing (measured:
-# 239/276 Z-Image linears quantised, full speedup) and keeps quality a touch higher.
+# Skip linears whose in/out features are below this. A small share of the FLOPs, so leaving
+# them bf16 costs ~nothing and keeps quality a touch higher.
 DEFAULT_MIN_LINEAR_FEATURES = 512
 
+# int8-ONLY name exclusions. The int8 dynamic path uses torch._int_mm, which requires the
+# activation row count M > 16. A DiT's AdaLN *modulation* projections and its timestep /
+# guidance / pooled-text *conditioning embedders* are computed once from the [batch, dim]
+# conditioning vector (M = batch = 1), not per token -- so they crash _int_mm even though
+# their feature dims are large (e.g. Flux's norm1.linear 3072->18432, Qwen's img_mod.1, Flux.2's
+# *_modulation.linear). min_features does NOT catch them (they are big), so int8 also skips any
+# Linear whose fqn matches one of these tokens. They are a negligible share of the FLOPs (run at
+# M=1, once per block), so int8 keeps the full speedup on the attention/FFN layers (M = seq).
+# fp8 / nvfp4 / mxfp8 use scaled_mm (no M>16 limit) and quantise these layers fine, so the
+# exclusion is int8-only. Sequence embedders (context_embedder / x_embedder / txt_in, M = seq)
+# are deliberately NOT excluded.
+_INT8_EXCLUDE_NAME_TOKENS = (
+    "norm",  # AdaLN modulation .linear (norm1 / norm1_context / norm / norm_out)
+    "_mod",  # Qwen img_mod / txt_mod
+    "modulation",  # Flux.2 double/single_stream_modulation
+    "timestep_embed",
+    "guidance_embed",
+    "time_text_embed",  # Flux/Qwen time_text_embed.* (pooled-text + timestep); NOT context_embedder
+    "pooled",
+)
+
 # Per-architecture preference order for ``auto`` -- best (fastest, in-bar) first, with
 # the lower-precision schemes listed as fallbacks for that arch tier. On Blackwell, fp8
 # leads: measured on a B200, plain fp8 dynamic is both faster AND more accurate than the
@@ -286,9 +304,11 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any:
     raise ValueError(f"unknown transformer quant scheme '{scheme}'")
 
 
-def make_filter_fn(min_features: int):
-    """A torchao ``quantize_`` filter keeping only the FLOP-heavy linears: nn.Linear
-    with both in/out features >= ``min_features``. Hides the (module, fqn) callback arity."""
+def make_filter_fn(min_features: int, exclude_name_tokens: tuple[str, ...] = ()):
+    """A torchao ``quantize_`` filter keeping only the FLOP-heavy linears: nn.Linear with both
+    in/out features >= ``min_features`` AND whose fully-qualified name contains none of
+    ``exclude_name_tokens`` (used by int8 to skip the M=1 modulation / conditioning-embedder
+    projections that crash ``torch._int_mm``). Hides the (module, fqn) callback arity."""
 
     def filter_fn(module: Any, fqn: str = "") -> bool:
         try:
@@ -301,7 +321,13 @@ def make_filter_fn(min_features: int):
         out_features = getattr(module, "out_features", None)
         if in_features is None or out_features is None:
             return False
-        return in_features >= min_features and out_features >= min_features
+        if in_features < min_features or out_features < min_features:
+            return False
+        if exclude_name_tokens:
+            name = fqn.lower()
+            if any(tok in name for tok in exclude_name_tokens):
+                return False
+        return True
 
     return filter_fn
 
@@ -331,10 +357,13 @@ def quantize_transformer(
     try:
         from torchao.quantization import quantize_
 
+        # int8 (torch._int_mm, M>16) additionally skips the M=1 modulation / conditioning-embedder
+        # projections; fp8 / fp4 / mx (scaled_mm) have no such limit and quantise everything.
+        exclude = _INT8_EXCLUDE_NAME_TOKENS if scheme == TQ_INT8 else ()
         quantize_(
             transformer,
             _make_quant_config(scheme, fast_accum = fast_accum),
-            filter_fn = make_filter_fn(min_features),
+            filter_fn = make_filter_fn(min_features, exclude_name_tokens = exclude),
         )
         # Runtime-only marker (torchao tensors are not safetensors-serializable; this
         # backend is inference-only, so this is purely diagnostic).

studio/backend/tests/test_diffusion_transformer_quant.py

@@ -309,6 +309,50 @@ def test_make_filter_fn(monkeypatch):
     assert keep(types.SimpleNamespace(), "no_attrs") is False
 
 
+def test_make_filter_fn_int8_excludes_modulation_and_embedders(monkeypatch):
+    # The int8 path skips the large M=1 AdaLN modulation / conditioning-embedder projections
+    # (they crash torch._int_mm's M>16), while keeping the attention / FFN compute layers and
+    # the sequence embedders. fp8 (no exclusion) keeps everything.
+    from core.inference.diffusion_transformer_quant import _INT8_EXCLUDE_NAME_TOKENS
+
+    class _Lin:
+        def __init__(self, i, o):
+            self.in_features, self.out_features = i, o
+
+    torch = types.ModuleType("torch")
+    torch.nn = types.SimpleNamespace(Linear = _Lin)
+    monkeypatch.setitem(sys.modules, "torch", torch)
+
+    keep = make_filter_fn(512, exclude_name_tokens = _INT8_EXCLUDE_NAME_TOKENS)
+    big = lambda: _Lin(3072, 18432)  # noqa: E731 — large enough to pass min_features
+    # Excluded (M=1 modulation / conditioning embedders), despite large features:
+    for fqn in (
+        "transformer_blocks.0.norm1.linear",
+        "transformer_blocks.0.norm1_context.linear",
+        "single_transformer_blocks.0.norm.linear",
+        "norm_out.linear",
+        "transformer_blocks.0.img_mod.1",
+        "transformer_blocks.0.txt_mod.1",
+        "double_stream_modulation_img.linear",
+        "time_text_embed.timestep_embedder.linear_2",
+        "time_text_embed.guidance_embedder.linear_2",
+        "time_guidance_embed.timestep_embedder.linear_2",
+    ):
+        assert keep(big(), fqn) is False, fqn
+    # Kept (M=seq compute layers + sequence embedders), NOT matched by the modulation tokens:
+    for fqn in (
+        "transformer_blocks.0.attn.to_q",
+        "transformer_blocks.0.ff.net.0.proj",
+        "single_transformer_blocks.0.proj_mlp",
+        "single_transformer_blocks.0.attn.to_qkv_mlp_proj",
+        "context_embedder",  # "context" contains "text" -> must NOT be excluded
+        "txt_in",
+    ):
+        assert keep(big(), fqn) is True, fqn
+    # Without the exclusion (fp8 path), the modulation layer is kept.
+    assert make_filter_fn(512)(big(), "transformer_blocks.0.norm1.linear") is True
+
+
 # ── apply ───────────────────────────────────────────────────────────────────────