Studio diffusion (Phase 14) review round 2: align helper name with the stack

Rename the int8 exclusion helper to exclude_tokens_for_scheme, matching the
identical helper already present higher in the diffusion stack (Phase 16). The
helper definition, the runtime quantiser call, and the offline builder are now
byte-identical to that version, so the two branches no longer introduce a
divergent name for the same single-source-of-truth and the stack merges without
a conflict on this fix. No behavior change.

scripts/build_prequant_checkpoint.py

@@ -57,7 +57,7 @@ def main(argv = None) -> int:
     from core.inference.diffusion_transformer_quant import (
         TQ_SCHEMES,
         _make_quant_config,
-        int8_exclude_name_tokens,
+        exclude_tokens_for_scheme,
         make_filter_fn,
     )
     from torchao.quantization import quantize_
@@ -79,15 +79,15 @@ 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)
-    # 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).
+    # Mirror the runtime path EXACTLY (the offline == runtime, LPIPS-0 invariant): for int8 also
+    # skip the M=1 AdaLN-modulation / conditioning-embedder projections, else the saved checkpoint
+    # bakes them as int8 and crashes (torch._int_mm needs M>16) at the first denoise step on
+    # Flux / Qwen. fp8 / fp4 / mx use scaled_mm (no M limit) -> exclude_tokens_for_scheme returns ().
     quantize_(
         transformer,
         _make_quant_config(scheme),
         filter_fn = make_filter_fn(
-            args.min_features, exclude_name_tokens = int8_exclude_name_tokens(scheme)
+            args.min_features, exclude_name_tokens = exclude_tokens_for_scheme(scheme)
         ),
     )
 

studio/backend/core/inference/diffusion_transformer_quant.py

@@ -63,12 +63,13 @@ _INT8_EXCLUDE_NAME_TOKENS = (
 )
 
 
-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)."""
+def exclude_tokens_for_scheme(scheme: str) -> tuple[str, ...]:
+    """Name tokens to exclude from quantisation for ``scheme``. int8 (torch._int_mm, M>16)
+    skips the M=1 modulation / conditioning-embedder projections (see _INT8_EXCLUDE_NAME_TOKENS);
+    every other scheme uses scaled_mm (no M limit) and excludes nothing. Shared by the runtime
+    quantise path and the offline prequant-checkpoint builder so the two never drift -- an int8
+    checkpoint built offline must skip exactly the layers the runtime path skips, or it bakes the
+    M=1 projections as int8 and crashes at the first denoise step on Flux / Qwen."""
     return _INT8_EXCLUDE_NAME_TOKENS if scheme == TQ_INT8 else ()
 
 
@@ -369,7 +370,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(scheme)
+        exclude = exclude_tokens_for_scheme(scheme)
         quantize_(
             transformer,
             _make_quant_config(scheme, fast_accum = fast_accum),

studio/backend/tests/test_diffusion_transformer_quant.py

@@ -357,17 +357,18 @@ 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():
+def test_exclude_tokens_for_scheme_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,
+        exclude_tokens_for_scheme,
     )
-    assert int8_exclude_name_tokens(TQ_INT8) == _INT8_EXCLUDE_NAME_TOKENS
+
+    assert exclude_tokens_for_scheme(TQ_INT8) == _INT8_EXCLUDE_NAME_TOKENS
     for scheme in (TQ_FP8, TQ_NVFP4, TQ_MXFP8):
-        assert int8_exclude_name_tokens(scheme) == ()
+        assert exclude_tokens_for_scheme(scheme) == ()
 
 
 # ── apply ───────────────────────────────────────────────────────────────────────