Studio diffusion (Phase 14) review round 2: apply int8 M=1 exclusion in the builder

Codex review: the M=1 modulation/embedder exclusion was wired only into the dense
runtime quantiser; the offline builder scripts/build_prequant_checkpoint.py called
make_filter_fn(min_features) with no exclusion. So an int8 prequant checkpoint
quantised the AdaLN modulation and conditioning-embedder linears, and loading it
via transformer_prequant_path (the load path only loads already-quantised tensors,
it can't re-skip them) reintroduced the torch._int_mm M=1 crash this phase fixes
for the runtime path.

Extracted int8_exclude_name_tokens(scheme) as the single source of truth (int8 ->
the M=1 exclusion, every other scheme -> none) and use it in both the runtime
quantiser and the builder, so a prequant artifact's quantised-layer set always
matches the runtime. fp8/fp4/mx artifacts are byte-identical (empty exclusion).

Test: int8_exclude_name_tokens returns the exclusion for int8 and () for
fp8/nvfp4/mxfp8.

scripts/build_prequant_checkpoint.py

@@ -57,6 +57,7 @@ def main(argv = None) -> int:
     from core.inference.diffusion_transformer_quant import (
         TQ_SCHEMES,
         _make_quant_config,
+        int8_exclude_name_tokens,
         make_filter_fn,
     )
     from torchao.quantization import quantize_
@@ -78,7 +79,17 @@ def main(argv = None) -> int:
         args.base, subfolder = "transformer", torch_dtype = torch.bfloat16, token = args.hf_token
     ).to("cuda")
     print(f"  quantising in place ({scheme}) ...", flush = True)
-    quantize_(transformer, _make_quant_config(scheme), filter_fn = make_filter_fn(args.min_features))
+    # Use the SAME int8 M=1 exclusion as the runtime quantiser (single source of truth):
+    # otherwise an int8 prequant checkpoint quantises the AdaLN modulation / conditioning
+    # embedders and reintroduces the torch._int_mm M=1 crash when loaded via
+    # transformer_prequant_path. fp8/fp4/mx get an empty exclusion (artifacts unchanged).
+    quantize_(
+        transformer,
+        _make_quant_config(scheme),
+        filter_fn = make_filter_fn(
+            args.min_features, exclude_name_tokens = int8_exclude_name_tokens(scheme)
+        ),
+    )
 
     # Move the state dict to CPU for a portable, GPU-free artifact.
     state_dict = {

studio/backend/core/inference/diffusion_transformer_quant.py

@@ -62,6 +62,15 @@ _INT8_EXCLUDE_NAME_TOKENS = (
     "pooled",
 )
 
+
+def int8_exclude_name_tokens(scheme: str) -> tuple[str, ...]:
+    """Filter exclusions for ``scheme``: the M=1 AdaLN-modulation / conditioning-embedder
+    linears for int8 (they crash ``torch._int_mm``, which needs M > 16), empty for every
+    other scheme (fp8 / fp4 / mx use ``scaled_mm``, no M limit). The single source of truth
+    shared by the runtime quantiser and the offline prequant builder, so a prequant artifact's
+    quantised-layer set matches the runtime exactly (no reintroduced M=1 crash)."""
+    return _INT8_EXCLUDE_NAME_TOKENS if scheme == TQ_INT8 else ()
+
 # 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
@@ -359,7 +368,7 @@ def quantize_transformer(
 
         # 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 ()
+        exclude = int8_exclude_name_tokens(scheme)
         quantize_(
             transformer,
             _make_quant_config(scheme, fast_accum = fast_accum),

studio/backend/tests/test_diffusion_transformer_quant.py

@@ -357,6 +357,20 @@ def test_make_filter_fn_int8_excludes_modulation_and_embedders(monkeypatch):
     assert keep(big(), "") is True
 
 
+def test_int8_exclude_name_tokens_shared_by_runtime_and_builder():
+    # The runtime quantiser and the offline prequant builder must apply the SAME int8
+    # exclusion, or an int8 prequant artifact quantises the M=1 modulation/embedder linears
+    # and reintroduces the torch._int_mm crash. int8 gets the exclusion; others get none.
+    from core.inference.diffusion_transformer_quant import (
+        _INT8_EXCLUDE_NAME_TOKENS,
+        int8_exclude_name_tokens,
+    )
+
+    assert int8_exclude_name_tokens(TQ_INT8) == _INT8_EXCLUDE_NAME_TOKENS
+    for scheme in (TQ_FP8, TQ_NVFP4, TQ_MXFP8):
+        assert int8_exclude_name_tokens(scheme) == ()
+
+
 # ── apply ───────────────────────────────────────────────────────────────────────