src/OPTIMIZATION_IDEAS.md

YuE2 W4A4 optimization frontier

Goal: measured latency/memory Pareto improvements at maximal quality. Classification is a hypothesis until verified. A distribution-preserving change can change sampled outputs through floating-point ordering or RNG consumption; compare probabilities and same-seed outputs separately.

Exact or near-exact candidates, profile first

1. Semantic LM head restricted to EOS + 32768 codec classes (current full vocabulary: 184704).
2. ABC LM head restricted to legal text/ABC classes.
3. Top-p over selected top-k only, avoiding full-vocabulary sort; retain threshold ties correctly.
4. Sparse frequency penalty using window counts, avoiding full-size scatter and power buffers.
5. Fused mask + temperature + penalty + top-k kernel.
6. Fused softmax + categorical sampling; preserve documented RNG semantics or label divergence.
7. Capture sampler with decode graph; one GPU graph per token.
8. GPU token/history/EOS state, read CPU only every 8–32 tokens with cancellation bounds.
9. Conditional CUDA graph loop until EOS/budget, avoiding host iteration.
10. Chunked graph unrolling with EOS masking, no extra output after termination.
11. Device-side history ring and exact repeated-token count updates.
12. Cache immutable vocabulary masks and phase constants.
13. RMSNorm + RPBH rotation + activation quantization in one kernel.
14. Reuse one rotation for grouped QKV and gate/up (check current grouping before crediting).
15. Combine activation quantization with packed GEMV in a persistent CTA.
16. GEMV + row scaling + residual addition in one write.
17. Grouped gate/up GEMV + SwiGLU; avoid materializing two dense outputs.
18. Down projection fused with residual and next RMSNorm (cross-layer boundary).
19. Specialized single-query GQA attention, fused RoPE + KV scatter.
20. Flash decoding with split-K / multi-block KV reduction tuned to 2k–12k context.
21. Attend only live KV length without static-capacity padding overhead.
22. KV cache layout optimized for one query rather than prefill.
23. FP8/INT8 KV cache with quality checks; keep independent from W4A4 weights.
24. Persistent all-layer decode kernel or cooperative graph groups, limited by synchronization and register/shared-memory pressure.
25. Native INT4 MMA on supported hardware versus current INT8 surrogate math.
26. Blackwell tensor-memory/TMA/cp.async specialization for packed prefill/NAR GEMM.
27. GEMV warp mapping autotune over actual M/N/K; vectorized packed reads, register pressure limits.
28. Precompute codebook surrogates/rotation constants and use constant/shared memory.
29. Group offload transfers, pinned host buffers and stream overlap with immutable module ownership.
30. Keep AR-only weights in CPU during NAR and avoid unnecessary restore before VAE.
31. Avoid full-stage GPU->CPU->GPU movement when memory permits; expose memory/latency profiles.
32. Compile NAR pointwise islands around opaque packed kernels (fake/meta registration if needed).
33. Reuse NAR timestep embeddings/conditioning projections across 32 steps.
34. Capture complete NAR integration loop, compare earlier velocity-only graph result.
35. Fuse midpoint update arithmetic and velocity input preparation.
36. VAE conv + SnakeBeta fusion; test cuDNN layouts and channels-last where supported.
37. Pipeline VAE chunk decode with CPU audio copy and encoding on separate streams.
38. Tune VAE halo/core profile against exact full output, avoiding boundary artifacts.
39. Remove repeated integrity hashing from warm path only when content identity is already verified.
40. Cache score/prefix prefill for repeated prompts/style variations without mixing request state.

High-risk / approximation candidates

41. Exact speculative decoding with a small draft model and rejection correction; draft availability/training cost is a constraint.
42. Self-speculation via skipped layers or early exits, verified by full target model; acceptance and verification cost determine benefit.
43. N-gram / prompt lookup speculation for repetitive ABC syntax; exact target verification.
44. Jacobi/lookahead block proposals for semantic tokens, with exact verification where possible.
45. Multi-token prediction heads trained on teacher traces; evaluate draft cost and acceptance, not just tokens proposed.
46. Reduce or omit ABC planning (cot=melody/off); meaningful musical-structure change requiring listening.
47. Compress ABC syntax / learned latent planning; requires model adaptation.
48. Adaptive flow ODE steps with error estimate; compare same semantic/noise, preserve difficult transients.
49. Higher-order solver at equal function evaluations; training/sampler mismatch risk.
50. Distill 32-step acoustic generation into 8/4 steps; keep separate artifact and evaluate quality.
51. Cache/reuse acoustic velocities or residuals across adjacent ODE steps with error bounds.
52. Reduced precision VAE (BF16/FP16) and selective FP32 activations; test clipping/noise and boundaries.
53. Weight/codebook calibration for sensitive heads or layers while maintaining 4-bit transformer weight target.
54. Activation clipping calibration and per-layer rotation search for W4 quality, followed by kernel-compatible packaging.
55. Token-rate reduction or grouped semantic tokens; requires training and changes duration/codec semantics.
56. Streaming acoustic blocks with contextual overlap, optimizing time to first playable audio separately from total latency.
57. AR/NAR/VAE concurrent chunk pipeline with explicit context correctness and bounded memory.
58. Batch multiple songs to amortize weight reads; improves throughput, not single-song latency.
59. Parallel candidate generation / best-of quality selection; increases total compute, cannot claim cost reduction.
60. Alternate exact low-bit accumulation/codebook representation; verify per-layer distributions before audio.

Iteration rule

Profile -> rank by measured share and upper-bound speedup -> build a small falsifiable prototype -> verify numerical behavior -> unprofiled repeat benchmark -> listen/ASR when output changes -> retain nondominated profiles -> record rejected paths -> generate the next hypotheses from remaining bottlenecks. No claims of global optimality. No silent reduction of song length, quality settings or scope to improve headline numbers.

Primary references

• CUDA conditional WHILE graph nodes and restrictions: https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/cuda-graphs.html
• PyTorch CUDA generator graph-safe state APIs: https://docs.pytorch.org/docs/2.14/generated/torch.Generator.html (check pinned Torch 2.10 availability before use).

New hypotheses from the first GPU profile

61. Replace six separate Q/K norm, RoPE and KV scatter operations with one specialized decode kernel; preserve every BF16 rounding boundary. This follows the measured abundance of tiny AR kernels.
62. Tune the NAR wide-projection INT8 output chunk (currently 2048 columns) against 1024/4096/full. Integer accumulation should remain exact; transient accumulator memory must be measured.
63. Avoid restoring offloaded AR weights after NAR when the next stage immediately moves the whole transformer to CPU for VAE. Restore on the next AR request instead; preserve module ownership and cancellation behavior.
64. Validate SageAttention-style quantized QK / higher-precision PV for the measured NAR attention hotspot, as a separately labeled approximate profile.
65. Cache or batch all known flow timestep conditioning computations before the ODE loop, where independent of evolving latents; audit existing caching first.
66. Profile the optimized AR again: update upper bounds before investing in a persistent all-layer kernel, because the remaining bottleneck may have moved to attention/rotation rather than Python.

Observed initial microbench (RTX PRO 6000 Blackwell, same 256 forced tokens): baseline 200.56 tok/s, vectorized GEMV 204.55, semantic head 217.85, fused norms/RoPE/SwiGLU 290.43, fused+head 328.35. All 128 forced-step semantic logits matched baseline exactly in this sample. Full-song validation is pending; these numbers are not end-to-end generation speedups.

Subsequent evidence

• User superseded bitwise-parity priority: spectral distribution and asynchronous listening are the quality workflow; never wait for listening approval.
• OrbitQuant 0.9.3 and native kernels1.0.1 released (aff3d6e). Small-row DP4A GEMV is general; measured3.39–55.90x vs previous native tile path, not whole-model speedup.
• Compact top-k sampling + fast QKV +8 midpoint steps +FP16 VAE compute:22.85s/178.96s audio at7.33GiB PRO6000.8vs32steps on same semantics/noise changes normalized log-band energy only.053dB mean. More exact metric artifacts in reports.
• Graph-captured sampler improves token throughput slightly but full22.96s did not beat22.85s; not selected by current end-to-end evidence.
• Autocast outside VAE alone was a no-op because decode disables it. Real inner autocast FP16 measured.58s decoder vs1.06sFP32; preserve FP32 weights and label compute precision.
• NAR V-view backing-storage hypothesis is ALREADY covered by current runtime's contiguous cache materialization. Guard rejected stale-source patch before inference; no new optimization credited.
• Bidirectional expert staging with deferred NAR restore after acoustic AR prefix improves some memory at transfer cost. Allocator experiments running to test fragmentation directly.

Measured follow-up: byte-pair lookup and RMS fusion

• Byte-pair LUT: selected. Native1.0.2 halves codebook lookups for K1024..16384; shorter inputs use the original kernel. Actual native decoder3.148->2.943ms, saved logits exact. Full interleaved song31.82->30.34s, identical PCM.
• RMSNorm plus packed activation quantization: selected as an explicit fast/memory option. FP32 intermediate products can change musical arrangement. Early Hadamard stages use warp shuffles and the final reduction stages run in one warp. Combined with byte-pair GEMV, the same fused song29.81->28.28s in interleaved measurements.
• Warp-resident16-entry lookup: only1.8% model gain; short-K full-mask shuffle bug repaired. Slower than byte-pair lookup; retained as research evidence.
• 8/16warps perblock: roughly2.99/3.00ms, slower than4warps at2.94ms. Larger blocks amortize setup but reduce useful scheduling flexibility.
• UINT16 pair-table storage:2.921ms in one fixed-token probe. About0.7% below the4warpUINT32 result; requires interleaved confirmation before selection.
• Generic quantizer warp rewrite: dynamic branching blocked unrolling; static specialization removed most regression but gave negligible model improvement. Not selected.
• Future investigations: fuse device sampling/history maintenance if profiling justifies it; use per-shape table-layout selection; test attention-cache compression separately with spectral/listening milestones.