perf(quantization): stream int8 convrot weights as int8 under partial loading

[Pfannkuchensack]

Summary

Under partial loading, an Int8ConvrotLinear whose int8 buffers stayed in RAM streamed its weight with a single self.weight.to(device=device, dtype=dtype). For a CPU→GPU copy that changes dtype, PyTorch converts on the CPU first and then transfers the widened tensor: a host-side int8→bf16 pass plus twice the PCIe bytes, for every streamed layer on every forward.

The weight is now moved first and widened on the device (.to(device=device).to(dtype=dtype)). Output is bit-identical and the peak transient is unchanged: the chain frees the int8 copy before the scale multiply allocates, so dequant_transient_bytes stays as it is. The fully resident path is unaffected, since .to(device) is a no-op there.

Measured on NVIDIA (RTX 4090, CUDA) and AMD (RX 9060 XT, ROCm), with pixel-identical images on both. Partial-load generations got 57–70% faster on the 4090 and 6–18% faster on the RX 9060 XT; that card's gain is limited by its narrow PCIe link (see QA Instructions).

This covers every int8_tensorwise(+convrot) model, because they all share Int8ConvrotLinear: Krea-2, FLUX.1/FLUX.2, Z-Image, Qwen-Image, Anima, Ideogram 4, LTX-2, MiniMax H3, PiD and the int8 text encoders. No other per-forward streaming path had the pattern. nvfp4, scaled fp8, GGUF, SDNQ, bnb and the autocast wrappers already move to the device before casting or dequantizing.

Related Issues / Discussions

Closes #9773

QA Instructions

NVIDIA (CUDA)

Measured on an RTX 4090 with torch 2.13.0+cu130 on Windows. The A/B runs were interleaved in one server process (warm-up, then B A A B B A A B, use_cache=false), with the VRAM cache capped through INVOKEAI_MAX_CACHE_VRAM_GB to force partial loading. Both arms produced pixel-identical images.

Model (1024×1024)	transformer resident	before	after
Krea-2 Turbo int8 (official Comfy), 8 steps	24%	56.5 s	20.2 s (−64%)
Krea-2 Turbo int8 (official Comfy), 8 steps	56%	35.7 s	15.4 s (−57%)
FLUX.1 dev int8 convrot, 20 steps	26%	116.3 s	34.6 s (−70%)

Per streamed layer (microbenchmark, unpinned host memory, bf16), the decode is 4.8–6.5× faster across 3072×3072 to 10240×3840. Results are bit-identical and the peak allocation stays at 4.00 B/element.

Tests:

• pytest -m "slow or not slow" tests/backend/quantization/test_int8_convrot.py tests/backend/quantization/test_int8_convrot_lora_sidecar.py tests/backend/quantization/test_quantized_payloads.py: 39 passed, 1 skipped.
• New test_a_layer_left_in_ram_moves_its_int8_codes_before_widening_them (runs on CPU via the meta device): fails against the previous one-step .to(device=, dtype=).
• New slow CUDA test test_a_layer_left_in_ram_streams_its_int8_codes_within_the_transient_estimate: fails if the streamed int8 copy is kept alive through the decode (5.2 MB vs. 4.2 MB on its fixture). It skips under cudaMallocAsync, which does not track requested bytes.
• ruff check and ruff format --check pass on the changed files.

AMD (ROCm)

Measured on an RX 9060 XT (16 GB, torch 2.13.0+rocm10.0.0, Windows) with the same method and the same tests (39 passed, 1 unrelated skip):

Model (1024×1024)	transformer resident	before	after
Krea-2 Turbo int8, 8 steps (default cache, 16 GB card)	86%	47.1 s	44.1 s (−6%)
Krea-2 Turbo int8, 8 steps (cache capped at 8 GB)	32%	82.9 s	67.7 s (−18%)

On this card each streamed layer is 2.1× faster, against ~5× on the 4090. The card sits on an M.2 adapter (PCIe 4.0 ×4) and moves only 2.2 GiB/s from unpinned memory, compared with 12.1 GiB/s on the 4090. So the transfer dominates there, and halving its bytes is the whole gain. On the 4090, the host-side int8→bf16 cast was the larger cost: 12.6 ms for a 36 MiB layer, against 6.2 ms to copy the bf16 result. A card attached at full width should land between the two.

Review

No blocking findings. Changes made in response to review:

• Added a CPU-runnable guard for the transfer itself; the memory test alone cannot tell the old form from the new one.
• The CUDA memory test now skips instead of passing vacuously when the allocator backend does not report requested bytes.
• Made the comment and docstring wording correct for fp32 compute (PiD decodes in fp32).

Remaining limitation: with the native caching allocator, the extra int8-sized blocks raised peak reserved memory by about one int8 copy of the largest layer: 36 MiB over a 30-layer mixed-shape loop, with peak allocated unchanged. With cudaMallocAsync this does not occur.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Meaningful regression coverage added / updated where needed; obsolete tests/code removed
• Persisted-state and API changes include required migrations / compatibility validation
• Relevant performance/efficiency opportunities considered; material claims have evidence
• Material review findings resolved and relevant checks rerun
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)