feat(minimax-remover): TeaCache, universal FP8 scale cache, parallel frame output#147
Open
chenping9999 wants to merge 32 commits into
Open
feat(minimax-remover): TeaCache, universal FP8 scale cache, parallel frame output#147chenping9999 wants to merge 32 commits into
chenping9999 wants to merge 32 commits into
Conversation
Performance highlights (RTX 5060 Ti, SM120, CUDA 13): - End-to-end: 2.6× speedup (30.7s → 11.9s for 123 frames, 3 segments) - RTF improvement: 6.0 → 2.3 (processing-time / clip-duration) - Per-layer GEMM: 1.14–1.30× speedup vs fp16 matmul - FP4 GEMM with quantize overhead: 4–9× faster than fp16 on large FFN projections - Precision: PSNR 52.0 dB (mean) / 45.2 dB (worst frame) vs fp16 Optimizations: - NVFP4 W4A4 quantization with dynamic per-call activation (no offline calibration) - Fused LayerNorm + adaLN + gate-residual in single fp32-stat Triton kernel - Fused FFN-up GEMM + bias + GELU → FP4 output (skip re-quantization for FFN-down) - QKV quantize-once optimization (reuse quantized norm output for Q/K/V projections) - Manual graph-capturable denoise loop with CUDA Graph support - BF16 transformer (NVFP4-native, eliminates fp16↔bf16 casts) - SageAttention / FA2 backend integration
- Add FlashRTFp8Linear with static calibration for FP8 GEMM acceleration - Add MiniMaxRemoverPipelineFP8 for full-frame inpainting with FP8 - Add fused kernel blocks (adaLN, gate-residual, GELU) - Add Triton Flash Attention (fp16/fp8 variants) - Add Triton fused LayerNorm with fp32 statistics - Add Triton RoPE for native [B,S,H,D] layout - Add shared utilities for kernel loading and validation
…precision paths - Add MiniMaxRemoverPipelineFP8 as default full-frame pipeline (FP8 W8A8) - Clarify NVFP4 is only suitable for small cropped regions - Add performance comparison table and correctness metrics - Update quickstart example with --use-fp4 flag - Refactor kernel loading into shared _utils module
…riton kernels - Make _attention the single source of truth for FLASHRT_ATTN_MODE routing (sage_*/triton_fp8/triton_fp16/fa2); _kern_block now reuses it via install_attention (with install_fa2_attention back-compat alias) instead of duplicating the dispatch and FA2 processor. - Consolidate _triton_fused_norm.py and _triton_rope.py into _kernels.py. - Add triton_fp8/triton_fp16 modes to the shared attention_forward dispatch. - Fail fast when FLASHRT_NORM_MODE=fp16 but the ada_layer_norm_fp16 kernel symbol is absent in the current flash_rt_kernels build. - Fix _fp8_linear docstring to match the set_calibration/freeze_calibration API.
Add a fp16-native fused RMS_norm kernel that replaces diffusers'
WanRMS_norm.forward (4 full-tensor fp32 passes) with a single-pass
kernel keeping fp16 in/out with fp32 internal stats (~6x faster per
call, preserving ~40 dB PSNR vs fp16 reference).
- csrc/quantize/fp16_rms_norm_ncdhw.{cu,cuh}: new CUDA kernel
- csrc/bindings.cpp + CMakeLists.txt: register fp16_rms_norm_ncdhw
- flash_rt/models/minimax_remover/_kernels.py: Triton wan_rms_norm_ncdhw
- flash_rt/models/minimax_remover/_vae_opt.py: VAE optimization patcher
- examples/minimax_remover_quickstart.py: add --vae-opt flag
…lone module Add fp16_rms_silu_ncdhw CUDA kernel that fuses RMSNorm + SiLU into a single pass for WanResidualBlock, eliminating one full tensor read/write and one kernel launch per site (~522 sites/decode). Single-kernel cosine >= 0.9999999 vs fp32 reference. Move both fp16 VAE kernels (fp16_rms_norm_ncdhw, fp16_rms_silu_ncdhw) from csrc/quantize/ into a dedicated csrc/kernels/minimax_remover/ folder and build them as a standalone flash_rt_minimax_remover pybind module (FLASHRT_ENABLE_MINIMAX_REMOVER=ON, OFF by default), following the flash_rt_omnivoice opt-in pattern. The kernels are removed from the default flash_rt_kernels target. Also eliminate the redundant fp32 cast in WanUpsample (nearest-exact upsample is index-only, so fp16 == fp32 bit-for-bit) and enable --vae-opt by default in the quickstart. Performance (RTX 5060 Ti, tennis 70 frames 432x240): fp16 reference: 17.33s FP8 (no VAE opt): 11.78s (1.47x) FP8 + VAE opt: 10.93s (1.59x) PSNR 40.9 dB
…cceleration
Add fp16_rms_norm_ndhwc / fp16_rms_silu_ndhwc CUDA kernels that
operate natively in channels-last (NDHWC) memory format, where the C
dimension is contiguous — enabling coalesced warp-level reduction
without shared memory.
Wire a full channels-last pipeline in install_vae_optimizations():
1. Convert all 61 WanCausalConv3d weights to channels_last_3d.
2. Patch WanCausalConv3d.forward to preserve CL through cat/pad/conv.
3. Swap residual-block and attention-block norm kernels to NDHWC
variants so norm output stays CL — no format break before conv.
This eliminates 97% of cuDNN's per-conv nchw<->nhwc conversion kernels
(287 ms -> 9 ms per decode) and gives ~1.3x per-conv speedup from
cuDNN's preferred CL algorithm variant.
End-to-end (tennis 70 frames 432x240, RTX 5060 Ti):
fp16 reference: 15.60 s
FlashRT default: 10.03 s (1.55x speedup, PSNR 40.8 dB, VRAM 2.49 GB)
Zero precision loss: PSNR 40.8 dB vs fp16 reference, identical to the
NCDHW path within fp16 rounding.
…celeration Hand-rolled FP8 e4m3 implicit-GEMM 3x3x3 causal conv3d kernel that replaces cuDNN fp16 conv3d for applicable layers (Ci % 32 == 0). The kernel computes im2col indices on-the-fly inside the MMA loop via cp.async, avoiding the 268MB intermediate matrix that made the naive im2col + _scaled_mm approach 3.5x slower than cuDNN. Adapted from the motus v17 kernel with three key changes: fp16 output/bias (VAE stays fp16), per-output-channel alpha vector (act_scale * w_scale[co] for higher PSNR), and fused activation quantization (2-pass device-side amax + quantize, no host sync, shared scale across cache+new frames). New files (all in csrc/kernels/minimax_remover/): - fp8_conv3d_mm_ndhwc_fp16out.cu/.cuh: 128/128/32 tile, 8 warps, 2-stage cp.async, virtual cache concat, direct causal output, bias-fused epilogue with per-channel dequant - fp16_quant_fp8_per_tensor.cu/.cuh: amax_fp16 + atomicMax (multi- tensor accumulation) and quantize_fp16_fp8_with_amax (standalone quantize pass reading pre-computed amax from device) Performance (RTX 5060 Ti, 70 frames 432x240): baseline (no flashrt): 17.33s (1.00x) channels-last only: 10.01s (1.73x, PSNR 40.85 dB) + FP8 conv3d (default): 8.75s (1.98x, PSNR 39.34 dB) 49 conv3d layers quantized (encoder + decoder). Non-applicable layers (1x1x1, 3x1x1, Ci%32!=0, Co<8) fall back to cuDNN fp16. Toggle via --fp8-conv / --no-fp8-conv.
…x amax pipeline Add three new fused NDHWC kernels and a dual-quantize launcher to further accelerate the FP8 conv3d VAE pipeline: New kernels (flash_rt_minimax_remover module): - fp16_rms_silu_amax_ndhwc: fuses amax computation into the norm+silu pass via warp-level reduction + atomicMax. Eliminates a separate full-tensor read for amax per conv layer. - fp16_rms_silu_quant_fp8_ndhwc: fuses norm+silu with FP8 e4m3 quantization (reads pre-computed amax from device, no fp16 write). - fp16_rms_silu_amax_quant_fp8_ndhwc: 2-pass combined launcher that produces only fp8 output + scale (no fp16 intermediate). - quantize_fp16_fp8_with_amax_dual: quantizes two fp16 buffers (cache + new) with a shared amax in a single kernel launch. Pipeline optimizations (_vae_opt.py): - Running-max amax: the norm module accumulates amax into a persistent device-side buffer (atomicMax) shared with the sister conv module. Since cache_x was a previous output of the same norm, its amax is already covered -- the conv skips the cache amax pass entirely, saving one full read of the 2-frame cache tensor per layer. - Sister-norm linkage: each FP8-enabled conv gets a reference to its norm module to access the pre-computed amax (set up during _install_fp8_conv3d_pipeline). - Dual-quantize: cache + new quantized in a single launch. Result (RTX 5060 Ti, tennis 70 frames 432x240): Previous: 8.75s, 39.3 dB PSNR median, 1.98x vs baseline Current: 8.58s, 39.9 dB PSNR median, 2.02x vs baseline Both speed and precision improved simultaneously -- the running-max scale is temporally consistent across frames, reducing quantization noise compared to per-frame dynamic amax.
…for transformer denoise Collapse the 3-kernel FFN glue sequence (add_bias_fp16 + gelu_inplace_fp16 + quantize_fp8_static_fp16) into a single-pass fp16->fp8 epilogue kernel (bias_gelu_quant_fp16_fp8). The fp8 output is the pre-quantised input of the next FP8 Linear, which skips its own activation quantise. All bias+gelu arithmetic is done in fp32 before the fp8 cast, so the result is slightly more accurate than the original path (which rounds to fp16 twice). End-to-end PSNR improves 39.9 -> 40.0 dB vs fp16 reference. Add mid-inference calibration freeze: a one-shot transformer forward hook freezes FP8 scales after the first denoise step, so steps 2..N run with static scales and the fused epilogue active -- a single-call invocation benefits without a separate warm-up pass. Denoise GPU time: 3.73s -> 3.16s (-15%). End-to-end: 8.58s -> 7.56s (2.29x vs fp16 baseline). Peak VRAM unchanged at 2.51 GB. Files: - fp16_bias_gelu_quant_fp8.cu/.cuh: fused kernel (gelu + identity variants) - _fp8_linear.py: FlashRTFp8Linear.gemm_no_bias / forward_from_fp8 helpers - _kern_block.py: FFN path dispatches fused epilogue (falls back during calibration) - _fp8_pipeline.py: mid-inference freeze via one-shot forward hook - minimax_remover_extra_bindings.cpp + CMakeLists.txt: pybind + build
Add a manual graph-capturable denoise loop for the FP8 (W8A8) path that mirrors the NVFP4 _manual_denoise design but calls the installed FP8 block forwards. The diffusers denoise loop cannot be captured directly because condition_embedder uses torch.arange (CPU op) and FlowMatchEulerDiscreteScheduler.step mutates CPU state. The manual loop pre-computes time embeddings + norm_out modulation + RoPE + dt outside the graph, then runs pre-allocated buffer copies + transformer forward + in-place Triton euler step -- all kernel launches, fully capturable. Gated behind FLASHRT_FP8_GRAPH=1 (default off). The first call uses the diffusers path for calibration; the graph is captured on the second call and replayed thereafter. Requires a graph-safe attention backend (triton_fp8/triton_fp16); the default sage_fp8 is not graph-safe. Also fixes a transformer dtype resolution bug (next(parameters()) hits an fp32 scale_shift_table) and uses per-channel latent normalize/denormalize in the manual path (the shared Triton helper collapses latents_std to a scalar via .max(), which is wrong for per-channel stats). Performance note: graph capture is technically correct but not a net win here -- the denoise loop is GPU-bound (graph saves only ~20ms), and the graph-safe attention backends are ~1.1s slower than sage_fp8. Retained for when a fast graph-safe attention backend (e.g. flash_rt_fa2) is built.
…dd kernels Collapse the O-projection / FFN-down block tails (add_bias_fp16 + gate_mul_residual_bcast) into a single fp16x8 (uint4) kernel that reads the GEMM output once and writes straight into the residual, eliminating the intermediate fp16 RMW pass across 720 slots per denoise. Also add a vectorised fp16_add_bias_vec8 for the Q/K/V projections (8x fewer memory transactions than the scalar path). Denoise GPU kernel time 3.10 s -> 3.03 s; end-to-end 7.57 s at 2.30x vs the fp16 reference (17.42 s). PSNR 40.0 dB median, 36.2 dB worst. Wall time is launch-bound at 12 steps x 30 blocks. Toggle via FLASHRT_DISABLE_BIAS_GATE=1 for A/B without a rebuild.
…used RMSNorm+RoPE kernels Two new CUDA kernels for the MiniMax-Remover transformer denoise: * fp16_ada_layernorm_quant_fp8: one-pass fp32-stat LayerNorm + adaLN modulation + per-tensor FP8 e4m3 quantise. Feeds a shared-scale FP8 tensor into Q/K/V (one quantise for three Linears, three descales) via FlashRTFp8Linear.gemm_from_fp8_ext, eliminating the three per-Linear activation-quant passes. Shared max scale also suppresses per-Linear outlier saturation (PSNR 40.05 -> 40.81 dB). * fp16_rmsnorm_rope_bshd: one-pass per-token RMSNorm (fp32 stats + fp16 affine) + interleaved RoPE on the native [B,S,H,Dd] fp16 layout. Replaces the Triton rms_norm_fp32stat + rope_apply_bshd pair and eliminates one full fp16 R/W of Q and K per attention block. End-to-end tennis clip (70 frames, 432x240, RTX 5060 Ti): 7.57 s -> 7.28 s (-4%, 2.30x -> 2.38x vs fp16 ref), PSNR mean 40.81 dB. Toggles: FLASHRT_DISABLE_ADA_QKV=1, FLASHRT_DISABLE_RMSNORM_ROPE=1.
… quant variant Transformer denoise: * norm2 now uses fp16_ada_layernorm_quant_fp8 to produce fp8 directly for FFN proj0, eliminating the intermediate fp16 [S,D] write+read+ quantize (4×S×D bytes/call × 360 calls). Toggle: FLASHRT_DISABLE_NORM2_FFN=1. VAE kernels: * Add fp16_rms_silu_amax_quant_fp8_ndhwc_nozero variant that skips zeroing amax_buf for running-max mode (retained for future use; direct VAE fusion blocked by WanResidualBlock cache slice needing fp16). End-to-end tennis clip (70 frames, 432×240, RTX 5060 Ti): 7.28 s → 7.25 s (2.38x → 2.39x vs fp16 ref), PSNR mean 40.86 dB.
Add a WanVAE-specific NVFP4 (W4A4) conv3d path that accelerates VAE
encode/decode by ~16% on top of the existing FP8 transformer path.
New CUDA kernels (in csrc/kernels/minimax_remover/):
- nvfp4_conv3d_ndhwc_fp16out.cu: Purpose-built NVFP4 implicit-GEMM
conv3d using mma.sync.kind::mxf4nvf4 (e2m1 x e2m1, UE4M3 block
scales). Outputs fp16 NDHWC directly (vs motus kernel's bf16 NCDHW),
eliminating two conversion passes per layer. Tile: BLOCK_M=128,
BLOCK_N=128, BLOCK_K=64, 8 warps, 2-stage cp.async pipeline.
Constraint: Ci % 64 == 0 (covers WanVAE Ci=192/384).
- fp16_quant_nvfp4_ndhwc.cu: Fused fp16 -> NVFP4 quantization kernel
with per-16-element UE4M3 block scales. Three variants:
(1) plain quant (NCDHW input),
(2) plain quant (channels-last 3D input — eliminates contiguous()
copy, coalesced channel reads),
(3) fused RMS_norm + SiLU + NVFP4 quant (channels-last input, single
kernel replaces 3 separate passes). kThreadsY=6 design supports
all WanVAE channel sizes (96/192/384).
Python integration (flash_rt/models/minimax_remover/_vae_nvfp4.py):
- install_vae_nvfp4(): pre-quantizes eligible conv3d weights to NVFP4
at install time; hooks WanCausalConv3d.forward to dispatch to the
NVFP4 kernel for eligible layers (Ci >= 192, 3x3x3, Ci % 64 == 0).
38 layers quantized (encode + decode).
- Rolling 2-frame FP4 cache: stores quantized FP4+SF from the previous
call and reuses it as the causal-conv cache, eliminating per-call
cache re-quantization. Handles T_new=1 (decode) by combining
[prev_frame, current_frame] in a sliding window.
- Channels-last direct input: detects channels-last 3D tensors and
uses the CL quant kernel variant, avoiding a contiguous() copy.
Quickstart integration: NVFP4 VAE is enabled by default; --no-nvfp4-vae
disables it.
Performance (RTX 5060 Ti, 70-frame tennis clip, 432x240):
Non-FlashRT (fp16): 17.33s 1.00x (reference)
FlashRT FP8: 7.18s 2.41x PSNR 40.0 dB vs fp16
FlashRT FP8+NVFP4 VAE: 6.71s 2.58x PSNR 34.7 dB vs fp16
VAE encode: 1829 -> 1564 ms (-14.5%)
VAE decode: 1664 -> 1438 ms (-13.6%)
VAE total: 3493 -> 3002 ms (-16.4%)
The NVFP4 VAE path works correctly on full-frame inputs (unlike the
NVFP4 transformer path which is broken on full-frame due to FP4 error
accumulation over the 12-step denoise loop). The VAE is a single-pass
encoder/decoder with no iterative error accumulation.
Documentation (docs/minimax_remover_usage.md) updated to:
- Distinguish NVFP4 VAE (works, default ON) from NVFP4 transformer
(broken on full-frame, small-region only)
- Add new kernel entries, performance row, CLI flag, env vars
- Add usage example for the NVFP4 VAE path
Three quality-neutral fusion points on the default FP8+NVFP4-VAE path (steady-state 5.86s -> 5.73s; end-to-end 6.71s -> 6.56s, 2.64x vs fp16 ref). PSNR unchanged at 35.2 dB mean vs the fp16 reference. * NVFP4 fused norm+silu+NVFP4-quant + FP4 cache reuse (flashrt-project#1): the sister norm of each fully-NVFP4 WanResidualBlock emits the FP4 activation directly via fp16_rms_silu_quant_nvfp4_cl_ndhwc (no fp16 round-trip) and the conv reuses Direction-2's rolling 2-frame FP4 cache. Mirrors diffusers' [prev, current] cache padding required by the WanVAE per-frame decode streaming loop. Gated by FLASHRT_NVFP4_FUSED_NORMQUANT (default 1). * FP8 fused norm+silu+running-amax+FP8-quant (flashrt-project#2): same idea for the FP8-conv residual blocks via fp16_rms_silu_amax_quant_fp8_ndhwc_nozero (accumulates into the running amax so the new-x scale stays consistent with the causal cache). Gated by FLASHRT_FP8_FUSED_NORMQUANT (default 1). * Q/K bias fused into rmsnorm+RoPE+int8-quant (flashrt-project#3): cuBLASLt's CUBLASLT_EPILOGUE_BIAS is NOT_SUPPORTED for FP8 (e4m3) GEMMs, so the Q/K projection bias is added pre-norm inside the downstream fp16_rmsnorm_rope_quant_int8 kernel (fp32 add, slightly more precise than the fp16 add it replaces). Eliminates 720 of 1092 add_bias_vec8 calls per denoise; V and proj_out keep their bias. Adds gemm_from_fp8_ext_nobias to feed the no-bias Q/K GEMM output. Also: steady-state uses the eager manual denoise path by default (FLASHRT_FP8_EAGER_MANUAL=1, avoids per-step torch.cat + scheduler sync); remove dead Triton wan_rms_norm_ncdhw kernel; update docs/minimax_remover_usage.md with the new numbers, env vars and fusion notes.
Clean up the fused norm+quant / Q/K bias path so it satisfies
docs/pr_review_checklist.md (hot-path cleanliness, fail-fast, doc/code
consistency, no dead code, green default test suite).
* Hot-path allocation: fp16_rmsnorm_rope_quant_int8_{q,k} did a
cudaMallocAsync + cudaFreeAsync of the per-token rstd buffer on every
Q/K call (~720x/denoise). Add a caller-owned rstd_buf scratch
parameter (binding default 0 -> transient fallback) and reuse a
persistent [B*S] fp32 buffer in FlashRTFA2Processor keyed by (B,S),
shared by Q and K on the same stream. Zero allocation in steady state.
* Fail-fast: check the kernel return code in the attention fused-quant
path and raise a clear RuntimeError (was silently ignored, unlike the
VAE paths).
* Docs/code consistency: remove the dead _USE_FUSED_NORMQUANT module
variable in _vae_nvfp4.py whose "default OFF / diverges ~27 dB"
comment contradicted the actual default-ON behaviour (35.2 dB) and the
usage docs; fix the stale kernel name in gemm_from_fp8_ext_nobias
docstring (fp16_rmsnorm_rope_bias_quant_int8 -> ..._quant_int8_q/k).
* Dead code: remove the unused warp_reduce_max_rq device helper (caused
an nvcc "declared but never referenced" warning) and the never-called
_cache_fp16_for_next nested function in the NVFP4 fused residual path.
* Tests: repair the long-failing
test_fp8_pipeline_call_does_not_patch_pipe_class by completing the
pipe/transformer stubs (config.eps, parameters(), vae,
register_forward_hook that fires on the stub forward) and exercising
the delegation path via FLASHRT_FP8_EAGER_MANUAL=0; the class-isolation
guarantee under test is independent of the steady-state dispatch mode.
Verified on RTX 5060 Ti (SM120, CUDA 13): smoke suite 18 passed / 1
skipped; quickstart default stack 6.53s (2.66x vs fp16 ref), PSNR 35.15
dB mean / 31.63 dB worst vs the fp16 reference (matches docs 35.2/31.7).
Sync upstream changes (qwen36 DFlash, VLASH, RTC temporal fusion, CPP runtime, SM89 Qwen3-VL FP8, Qwen3 prefill fast path) into fork. All minimax-remover conflicts resolved in favour of the fork's advanced fused-kernel implementation. Conflicts resolved: - CMakeLists.txt: keep both FLASHRT_ENABLE_MINIMAX_REMOVER (ours) and FLASHRT_ENABLE_SM120_DEV_KERNELS (upstream) - docs/minimax_remover_usage.md: keep fork version - examples/minimax_remover_quickstart.py: keep fork version - flash_rt/models/minimax_remover/_fp8_linear.py: keep fork version - flash_rt/models/minimax_remover/_fp8_pipeline.py: keep fork version - flash_rt/models/minimax_remover/_kern_block.py: keep fork version - tests/test_minimax_remover_smoke.py: keep fork version
… vs fp16)
Add zeroth-order TeaCache step caching to the MiniMax-Remover denoise
loop, mirroring the Motus/Cosmos3/Wan2.2 mechanism: at skip steps the
cached noise prediction is reused and only the cheap euler step runs.
MiniMax-Remover's flow-matching velocity is highly cacheable, so the
default schedule '3,5,7,9' (4 of 12 interior steps, step 0 and the last
step never skipped) is quality-neutral on full-frame inpainting:
tennis 70 frames 432x240, RTX 5060 Ti (back-to-back A/B):
fp16 reference (--no-flashrt): 17.33 s 1.00x 3.67 GB
FlashRT default (FP8+VAE+NVFP4-VAE+TeaCache): 5.66 s 3.06x 2.57 GB
PSNR vs fp16 ref: 35.1 dB mean (within noise of the 35.2 dB no-skip path)
Plumbed through all three denoise paths: the diffusers reference __call__
(calibration / first call), the FP8 eager manual denoise (steady state),
and the captured CUDA Graph (skip schedule baked in at capture time).
Also make --no-flashrt the master 'pure reference' switch: it now
disables ALL FlashRT optimisations (VAE fused kernels, FP8/NVFP4 conv,
NVFP4 VAE, TeaCache) so users need not also pass --no-vae-opt to get the
vanilla fp16 ground truth for PSNR/timing A/B.
…r clarity The old --use-fp4 name was misleading: the default path already uses NVFP4 for the VAE (single-pass, error-free). The renamed flag makes explicit that it switches the TRANSFORMER to NVFP4 W4A4 -- the 12-step iterative denoise where FP4 error accumulates and breaks full-frame outputs (cosine ~0.0). This distinguishes it from the always-on NVFP4 VAE path. - examples/minimax_remover_quickstart.py: rename flag + attr (args.use_fp4 -> args.nvfp4_transformer), update help text and precision note to spell out the transformer-vs-VAE distinction - docs/minimax_remover_usage.md: update CLI table, command examples, performance table, and takeaway to the new name - tests/test_minimax_remover_smoke.py: fix the TeaCache guard assertion (now excludes --nvfp4-transformer), add two new tests verifying the rename and the NVFP4 VAE exclusion
… ops, SageAttention auto-fallback
Three optimizations for one-shot / arbitrary-resolution cold-call speedup
(~25% faster, PSNR-neutral at 34.8 dB vs fp16 reference):
1. Universal-scale disk cache (_fp8_pipeline.py):
- FP8 act_amax_max calibrated once, persisted to
~/.flash_rt/calibration/minimax_remover_fp8_{fingerprint}.json
- Reused across ALL resolutions with enlarged margin (1.3) that
absorbs cross-resolution activation variance (<5% median, ~3%
outlier layers)
- First call skips the dynamic-FP8 calibration step entirely
- New CLI flags: --universal-scale / --universal-margin
- --no-universal-scale falls back to per-resolution calibration
2. PyTorch ops replace Triton JIT kernels (cold-start fix):
- mask_mul -> images_t * (1.0 - masks_t) in _fp8_pipeline
- euler_step_inplace -> lat_buf.add_(noise, alpha=dt) in
_fp8_manual_denoise
- Eliminates ~0.3s Triton JIT compilation on the first call
- Verified opt-vs-Triton PSNR >= 46 dB (negligible difference)
3. SageAttention crash auto-fallback (_attention.py):
- sm89 fused-quant kernel crashes (query_scale shape mismatch) at
certain sequence lengths (e.g. 80-frame clips, S=8100)
- Catches RuntimeError, blacklists the seq len, falls back to
standard rmsnorm+rope path (sage high-level API handles any S)
- First block fails fast (Python-level check); subsequent blocks
skip fused-quant entirely (zero overhead)
- Tennis (S=7290) unaffected; bmx-trees (S=8100) now works with
default settings (6.24s vs previous crash)
…me ~20x Replace the sequential PIL save loop with a thread-pool fan-out that pre-builds per-frame output arrays and encodes/writes them concurrently. zlib and jpeg release the GIL while compressing, so this scales near-linearly with cores. Prefer OpenCV's faster encoders (~20% on PNG) with a PIL fallback; cv2 is imported lazily on first save so its ~0.19s import does not tax process startup. PNG compress level lowered 6 -> 3 (lossless, pixels unchanged). Save phase: ~2.3s -> ~0.1s on a 70-frame 432x240 clip; end-to-end wall time 11.6s -> 9.6s. Encoder equivalence verified (cv2 vs PIL pixel- identical on the same array, 10/10).
The universal-scale cache (773b13d) calls _model_fingerprint() during pipeline construction, which assumes a dict-like transformer.config and writes to ~/.flash_rt/calibration/. Two unit-test stubs used a SimpleNamespace config and relied on the default construction path, so they raised TypeError; a third assertion went stale when the TeaCache default was factored into TEACACHE_SKIP_DEFAULT. - test_fp8_pipeline_call_does_not_patch_pipe_class / test_fp8_pipeline_call_forwards_skip_steps_to_pipe: construct with use_universal_scale=False. These tests cover class isolation and skip_steps forwarding, not the scale cache, so staying off the disk-backed path keeps them hermetic and avoids needing a realistic diffusers config stub. - test_quickstart_teacache_default_and_reference_guard: update the stale literal assertion to match the TEACACHE_SKIP_DEFAULT constant. Full suite now passes: 26 passed, 1 skipped (build-config).
…project/FlashRT Brings in upstream commit d0db114 (sm_120 INT4 W4A4 GEMV kernel + OMMA unlock tool). All minimax_remover conflicts resolved in favor of our fork's enhanced version (TeaCache, universal FP8 scale cache, parallel frame output, etc.). Conflicts resolved (--ours): - docs/minimax_remover_usage.md - examples/minimax_remover_quickstart.py - flash_rt/models/minimax_remover/_attention.py - flash_rt/models/minimax_remover/_fp8_manual_denoise.py - flash_rt/models/minimax_remover/_fp8_pipeline.py - tests/test_minimax_remover_smoke.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR contributes a set of
minimax_removerenhancements developed on thechenping9999fork, now rebased/synced against the latestflashrt-project/main(including the FP8/NVFP4 kernelized pipeline in #142 and the sm_120 INT4 W4A4 kernel in #146).The branch is fully up to date with
upstream/main, so the effective diff is small (~640 lines across 6 files): all shared FP8/NVFP4 kernel work already landed upstream identically and cancels out. Only the new enhancements below are introduced.Highlights
cv2-based parallel frame writing, cutting video save time by ~20x.--nvfp4-transformerflag — clearer naming replacing the old--use-fp4switch.tests/test_minimax_remover_smoke.py.Sync details
upstream/main(commitsffbcf87,d0db114) into the fork.docs/int4_blackwell.md+tools/patch_int4_omma_sm120.pywere taken from upstream unmodified.CMakeLists.txtrequired no changes (the fork already carried identical minimax FP8/NVFP4 build entries).minimax_removercontent conflicts, the fork's enhanced version was kept (the shared FP8/NVFP4 work already matches upstream; only the enhancement deltas differ):docs/minimax_remover_usage.mdexamples/minimax_remover_quickstart.pyflash_rt/models/minimax_remover/_attention.pyflash_rt/models/minimax_remover/_fp8_manual_denoise.pyflash_rt/models/minimax_remover/_fp8_pipeline.pytests/test_minimax_remover_smoke.pyFiles changed
flash_rt/models/minimax_remover/_fp8_pipeline.py--nvfp4-transformerflash_rt/models/minimax_remover/_attention.pyflash_rt/models/minimax_remover/_fp8_manual_denoise.pyexamples/minimax_remover_quickstart.pycv2frame output, new flagstests/test_minimax_remover_smoke.pydocs/minimax_remover_usage.mdChecklist
upstream/main; no stale conflictsminimax_removersmoke tests updatedFeedback welcome, especially on the TeaCache integration points and the universal-scale cache API.