Studio diffusion (Phase 15): build int8 pre-quantized checkpoints (skip M=1 modulation linears)

The prequant-checkpoint builder applied the dense quant filter without the int8-only
M=1 modulation / conditioning-embedder exclusion the runtime path uses, so a built int8
checkpoint baked those projections as int8 and crashed (torch._int_mm needs M>16) at the
first denoise step on Flux / Qwen. Factor the scheme->exclusion decision into a shared
exclude_tokens_for_scheme() used by both the runtime quantise path and the offline builder
so they can never drift, and apply it in build_prequant_checkpoint.py. int8 prequant now
produces a working checkpoint on every supported model, giving int8 (the consumer-preferred
scheme) the same ~2x load-VRAM and download reduction fp8 already had.

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,
+        exclude_tokens_for_scheme,
         make_filter_fn,
     )
     from torchao.quantization import quantize_
@@ -78,7 +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)
-    quantize_(transformer, _make_quant_config(scheme), filter_fn = make_filter_fn(args.min_features))
+    # 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 = exclude_tokens_for_scheme(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,16 @@ _INT8_EXCLUDE_NAME_TOKENS = (
     "pooled",
 )
 
+
+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 ()
+
 # 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 +369,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 = exclude_tokens_for_scheme(scheme)
         quantize_(
             transformer,
             _make_quant_config(scheme, fast_accum = fast_accum),

studio/backend/tests/test_diffusion_transformer_quant.py

@@ -353,6 +353,22 @@ def test_make_filter_fn_int8_excludes_modulation_and_embedders(monkeypatch):
     assert make_filter_fn(512)(big(), "transformer_blocks.0.norm1.linear") is True
 
 
+def test_exclude_tokens_for_scheme():
+    # The shared scheme->exclusion decision used by BOTH the runtime quantise path and the offline
+    # prequant-checkpoint builder, so an int8 checkpoint built ahead of time skips exactly the
+    # layers the runtime path skips (offline == runtime). int8 excludes the M=1 modulation /
+    # embedder tokens; every scaled_mm scheme excludes nothing.
+    from core.inference.diffusion_transformer_quant import (
+        _INT8_EXCLUDE_NAME_TOKENS,
+        exclude_tokens_for_scheme,
+    )
+
+    assert exclude_tokens_for_scheme(TQ_INT8) == _INT8_EXCLUDE_NAME_TOKENS
+    assert exclude_tokens_for_scheme(TQ_FP8) == ()
+    assert exclude_tokens_for_scheme(TQ_NVFP4) == ()
+    assert exclude_tokens_for_scheme(TQ_MXFP8) == ()
+
+
 # ── apply ───────────────────────────────────────────────────────────────────────