MiniMax Music 3 modular pipeline: 2.9x consumer-GPU speedup - findings that may be worth upstreaming

Following the merged MiniMax Music 3 modular pipeline (#14456) — thank you for
it — we rebuilt its hot paths for single-GPU speed and open-sourced the result:

https://github.com/TheDutchRuler/minimax-music3-studio

Measured on an RTX 4090 (20s songs, warm, bf16): reference 50.5s/song →
17.7s/song (0.88x realtime) with quality-neutral changes only. Findings that
may be worth upstreaming or documenting:

1. **`torch.compile` has a silent no-op trap.** The semantic-generation block
calls `language_model.model(...)` (the submodule) in its decode loop, so
`torch.compile(pipe.language_model)` compiles a callable the loop never
invokes. Compiling the wrong target costs nothing and gains nothing —
easy to burn time on.

2. **A manually built `StaticCache` never gets CUDA graphs.** transformers'
compiled-generate path calls `torch._dynamo.mark_static_address` on the
cache tensors; building the cache by hand skips that, and cudagraph trees
silently refuse to capture the in-graph cache mutation (no warning at
default log levels). Marking the tensors after prefill fixes it.

3. **Cheap wins in the blocks themselves:**
- The DiT inner loop runs conditional/unconditional as two sequential
batch=1 forwards; one batch=2 forward is formula-identical (12.8s → 8.8s
for the stage in our measurements).
- `lm_head` reads a 1.38GB matrix per frame while every non-audio row is
masked to -inf; slicing to the sampleable rows is provably identical.
- Replacing the final softmax+multinomial with Gumbel-argmax (identical
distribution, unit-tested at L1≈0.015 over 200k draws) lets the whole
depth-decoder chain compile with zero eager glue.

4. **The big one — batched variation generation.** The AR stage is memory-
bandwidth bound (~23GB weight reads/frame), so K same-prompt variations
decoded in lockstep (batch 2K rows, per-song seeded noise, cross-song
batched CFG denoise grouped by exact frame count) make the marginal song
nearly free: 3 variations at 17.7s each vs 31s solo. Everything stays
row-independent; might be worth first-class support in the pipeline.

Also documented: negative results (per-layer group offload on the AR LLM
re-streams 16.4GB/frame — the model card's low-VRAM snippet; FP8 weight-only
via torchao 2.1x slower on Windows/torch 2.11), details and harnesses in the
repo's `labs/`.

This work was engineered end-to-end with Claude (Fable 5 Max) by Anthropic.
Happy to expand any of these into a proper issue/PR if useful.

Start with the MiniMax Music 3 modular pipeline from #14456 and the benchmarks and harnesses in the linked repository's labs/ directory; inspect the torch.compile, StaticCache, batched CFG, and decoder findings against the current implementation. Done requires selecting a specific optimization or documentation change, reproducing its performance and quality evidence, and defining the upstream scope.

Scritto dal modello di indicizzazione a partire dal testo della issue.