Diffusion: auto TE/VAE precision, video DiT quant correctness, and HunyuanVideo-1.5 attention speedup

scripts/build_prequant_checkpoint.py

@@ -83,19 +83,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)
-    # Mirror the runtime path EXACTLY (offline == runtime, LPIPS-0 invariant): for int8 also skip
-    # the M=1 AdaLN-modulation / conditioning-embedder projections, else the 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 ().
-    exclude_name_tokens = exclude_tokens_for_scheme(scheme)
-    # fp8 and mxfp8 assert a bf16 weight, so their filter must skip any non-bf16 Linear the
-    # transformer keeps: a mixed-precision DiT (Wan / Hunyuan) keeps its _keep_in_fp32_modules in
-    # fp32 even under torch_dtype=bf16, so quantising one raises inside quantize_ and aborts the
-    # pass. nvfp4 quantises fp32 fine, so it isn't gated. Runtime quantize_transformer gates on
-    # scheme membership; mirror it here so the offline checkpoint quantises the same layer set.
+    # Mirror the runtime path EXACTLY (offline == runtime, LPIPS-0 invariant): int8 skips the
+    # M=1 modulation / conditioning-embedder projections (else the checkpoint bakes them int8 and
+    # crashes at the first denoise step); the scaled_mm schemes exclude only a family's
+    # padded-conditioning embedder. Pass the family so the offline set matches runtime.
+    exclude_name_tokens = exclude_tokens_for_scheme(scheme, fam.name)
+    # fp8 / mxfp8 assert a bf16 weight, so skip any non-bf16 Linear the transformer keeps in fp32
+    # (a mixed-precision DiT's _keep_in_fp32_modules), else quantize_ raises. nvfp4 handles fp32.
+    # Mirrors the runtime scheme gate so the offline layer set matches.
     require_bf16 = scheme in _REQUIRE_BF16_SCHEMES
-    # fp8 bakes the accumulate mode into the saved kernels; record the resolved choice so the
-    # loader can refuse a checkpoint whose baked value contradicts an explicit runtime request.
+    # fp8 bakes the accumulate mode into the kernels; record the resolved choice so the loader
+    # can refuse a checkpoint whose baked value contradicts an explicit runtime request.
     fast_accum = _resolve_fast_accum(None) if scheme == TQ_FP8 else None
     quantize_(
         transformer,
@@ -117,10 +115,8 @@ def main(argv = None) -> int:
         "family": fam.name,
         "scheme": scheme,
         "min_features": args.min_features,
-        # The layers skipped for this scheme (int8's M=1 modulation projections; () for the
-        # scaled_mm schemes), whether non-bf16 Linears were skipped (the scaled_mm bf16 gate), and,
-        # for fp8, the baked accumulate mode. All let the loader reject a checkpoint that wouldn't
-        # match the runtime path.
+        # Skipped layers, the non-bf16 gate, and (fp8) the baked accumulate mode -- all let the
+        # loader reject a checkpoint that would not match the runtime path.
         "exclude_name_tokens": list(exclude_name_tokens),
         "require_bf16": require_bf16,
         "fast_accum": fast_accum,