Skip to content

perf: end-to-end inference pipeline optimization - #38

Open
Kenshin9977 wants to merge 10 commits into
edenaion:mainfrom
Kenshin9977:perf/global-optimizations
Open

perf: end-to-end inference pipeline optimization#38
Kenshin9977 wants to merge 10 commits into
edenaion:mainfrom
Kenshin9977:perf/global-optimizations

Conversation

@Kenshin9977

@Kenshin9977 Kenshin9977 commented Mar 14, 2026 •
edited
Loading

Copy link
Copy Markdown

Summary

Rebased on upstream main (1.10.0). Adds inference perf and fixes four open issues. Auto-detects GPU capability, no manual flags needed.

Issues closed

Plus a [tool.uv] constraint-dependencies = ["opencv-python==99999"] to block the non-headless opencv that diffusers/imageio/PIMS pull in transitively.

Auto-detected accelerators

These all auto-enable at engine init based on installed packages and the GPU's compute capability. Override with CORRIDORKEY_USE_SAGE or CORRIDORKEY_NVFP4 set to 0 (force off) or 1 (force on).

SageAttention (thu-ml/SageAttention, NeurIPS 2025 Spotlight). Quantized 8-bit / 4-bit attention kernel, drop-in for F.scaled_dot_product_attention. Auto-on from Ampere upwards (SM 8+).

GPU Sage version Speedup vs FA2
RTX 3090 2++ INT8 ~1.5x
RTX 4090 2++ INT8 ~2x
RTX 5090 3 FP4 ~5x

The Hiera global-attention monkey-patch already collapses Q/K/V to a 4D contiguous layout, so the call routes through _attention_kernel(device) and Sage just plugs in. At engine init we cross-check Sage vs SDPA on a dummy tensor; if the diff exceeds 5e-2 we fall back silently.

NVFP4 weight quantization (torchao). Packs weights to 4-bit microscaling on Blackwell's 5th gen Tensor Cores. Auto-on at SM 12+ when torchao is installed. No-op on Ampere/Ada (no native FP4 hardware, emulation would be slower).

Inference engine, color_utils, refiner

  • model.half() after load. Lossless because autocast was already running fp16. Saves ~450 MB VRAM.
  • clean_matte_gpu for realtime/preview paths (max_pool erode/dilate + cached gaussian). CPU clean_matte stays the production default since it preserves hair strands.
  • Pure-tensor linear_to_srgb_tensor / srgb_to_linear_tensor without isinstance, safe inside torch.compile graphs.
  • create_checkerboard cached with @lru_cache, result marked read-only.
  • Refiner: when B=1 (always at inference), all tiles get stacked into one forward. _blend_weight_cache bounded at 64 entries.

I/O overlap with GPU work

Adapted from upstream PR #226. Sequential inference loop now runs:

  • Background prefetch thread feeds a depth-3 queue while the GPU processes the current frame
  • Single-worker ThreadPoolExecutor for output writes, parallel with the next GPU step
  • Pending writes drained at end-of-clip; write failures retroactively flip the matching FrameResult to (success=False, warning=...)

Big gain on 4K EXR. Smaller but still positive on PNG/JPG.

Rust native extension

corridorkey_native/ (PyO3) with three hot ops, ~3x faster than the NumPy versions:

  • gbr_planar_to_rgb: gbrpf32le planar to interleaved RGB
  • bgr_u8_to_rgb_f32 / rgb_f32_to_bgr_u8: channel swap and cast in one pass

CorridorKeyModule/core/native_ops.py wraps and falls back to NumPy when the crate isn't installed. Build wired into the install scripts.

FFmpegFrameReader / FFmpegFrameWriter

backend/ffmpeg_tools/streaming.py. Pipe-based subprocess wrappers, no intermediate files.

  • Reader: hardware decode auto-detected (NVDEC, VAAPI, VideoToolbox), gbrpf32le output decoded by the Rust crate
  • Writer: encoder auto-picked. NVENC, AMF (Windows), VAAPI (Linux), QSV (Intel), libx264 fallback. Encoder-specific flags wired correctly.

Install scripts

NVIDIA detected: install torch-tensorrt, sageattention, torchao. AMD ROCm detected: install torch_migraphx. If cargo isn't on PATH, install rustup non-interactively then build the Rust crate via maturin. Everything silent-fail to the Python+inductor fallback.

Tests

  • tests/test_color_utils.py: checkerboard cache, GPU vs CPU clean_matte, despill clamping, sRGB roundtrip, gaussian kernel cache
  • tests/test_native_ops.py: Rust vs NumPy equivalence and roundtrip
  • .gitignore allows tests/test_*.py (still ignores root-level scratch scripts)

scripts/profile_vram.py is a small per-stage VRAM breakdown for OOM diagnosis (model load, input upload, peak forward, post-proc, cleanup).

uv run pytest -> 515 passed, 29 skipped on my machine.

Dropped (already merged upstream)

  • BiRefNet wrapper (upstream version is more complete: more variants, FP16, macOS bundle handling)
  • BiRefNet UI wiring (job type, button, dropdown, label, worker)
  • Hiera FlashAttention patch and tiled refiner with compile_tile_kernel
  • TF32 / cudnn.benchmark = False (upstream moved them to app startup)
  • weights_only=True for torch.load
  • stderr capture in stitch_video (upstream keeps last 10 lines, I override to log everything)
  • Configurable model_resolution setting in preferences

Expected gains at 1080p

RTX 3080 (10 GB) RTX 4090 (24 GB) RTX 5090 (32 GB)
SageAttention +15-20% +25% +60-80%
I/O overlap +10-25% +10-25% +5-15%
NVFP4 n/a n/a +30-50%
fp16 + tile batching + Rust +25-35% +25-35% +25-35%

Cumulative: ~2x current on RTX 3080, 5-8x on RTX 5090.

@Kenshin9977
Kenshin9977 force-pushed the perf/global-optimizations branch 2 times, most recently from 5015483 to 47060aa Compare March 14, 2026 12:38
@Kenshin9977 Kenshin9977 changed the title Perf/global optimizations perf: end-to-end inference pipeline optimization Mar 14, 2026
@Kenshin9977
Kenshin9977 marked this pull request as ready for review March 14, 2026 18:05
@edenaion

edenaion commented Mar 14, 2026

Copy link
Copy Markdown
Owner

Thank you for your contribution. Going to dive in I hope tomorrow, currently dealing with my own update.

DrCheeseFace
DrCheeseFace reviewed Apr 13, 2026
View reviewed changes
Comment thread backend/ffmpeg_tools.py Outdated
@Kenshin9977
fix: resolve issues edenaion#16 edenaion#28 edenaion#34 edenaion#39 +…
9d65166 
… add VRAM profiler

- protobuf>=5.0 pinned (fix _SixMetaPathImporter on Python 3.12+)
- opencv-python==99999 constraint (block transitive non-headless variant)
- nvidia-smi search paths: System32 + CUDA_PATH\bin
- 3-update scripts: --reinstall-package corridorkey to refresh metadata
- stitch_video: include all stderr lines on error (per review feedback)
- scripts/profile_vram.py: per-step VRAM breakdown for diagnostics
@Kenshin9977
perf: fp16 weights, GPU clean_matte helper, batched refiner tiles
1b47f6b 
- model.half() after load saves ~450 MB VRAM (autocast already runs fp16)
- color_utils: linear_to_srgb_tensor, srgb_to_linear_tensor (torch.compile safe)
- color_utils: clean_matte_gpu (no CPU roundtrip, for preview/realtime)
- color_utils: lru_cache on create_checkerboard
- model_transformer: stack tiles into single forward pass when B=1
- model_transformer: bound _blend_weight_cache to 64 entries
@Kenshin9977
perf: add Rust native extension for plane shuffle and color conversion
3196d80 
corridorkey_native provides optional Rust-compiled fast paths for:
- gbr_planar_to_rgb: FFmpeg gbrpf32le -> interleaved RGB
- bgr_u8_to_rgb_f32 / rgb_f32_to_bgr_u8: fused channel swap + dtype

native_ops.py wrapper falls back to NumPy when not installed.
Build: cd corridorkey_native && maturin develop --release
@Kenshin9977
feat: add FFmpegFrameReader/Writer with hardware encode auto-detection
124dcca 
Pipe-based subprocess I/O for in-memory video processing:
- FFmpegFrameReader: stdout pipe + NVDEC/VAAPI/VideoToolbox hardware decode
- FFmpegFrameWriter: stdin pipe with auto-detected encoder:
  - NVENC (NVIDIA), AMF (AMD/Win), VAAPI (AMD/Linux), QSV (Intel)
  - Falls back to libx264

No intermediate disk I/O — useful for streaming/realtime pipelines.
@Kenshin9977
feat: auto-install torch-tensorrt, torch-migraphx and Rust extension
29cdc8b 
Install scripts now detect GPU vendor and install matching torch.compile
backend (torch-tensorrt for NVIDIA, torch_migraphx for AMD ROCm). Auto-
install Rust toolchain via rustup if cargo is missing, then build the
corridorkey_native extension via maturin. All steps silent-fail to
the Python/inductor fallback.
@Kenshin9977
test: add color_utils and native_ops tests
a21e6f8 
- test_color_utils: checkerboard cache, clean_matte (CPU + GPU),
  despill, pure-tensor sRGB, gaussian kernel cache
- test_native_ops: Rust/Python equivalence for all native ops + roundtrip
- Mark cached checkerboard read-only to prevent caller mutation
- gitignore: allow tests/test_*.py (only ignore test_*.py at root)
@Kenshin9977
Kenshin9977 force-pushed the perf/global-optimizations branch from fc4d29d to a21e6f8 Compare April 25, 2026 10:16
@Kenshin9977
perf: SageAttention drop-in for SDPA (opt-in via CORRIDORKEY_USE_SAGE)
1fc6337 
SageAttention is a quantized attention kernel that auto-selects the right
implementation per GPU generation: INT8 (Sage 2++) on Ampere/Ada giving
~2x vs FA2, FP4 (Sage 3) on Blackwell giving ~5x vs FA2 on RTX 5090.

Changes:
- pyproject.toml: new optional perf-attn extra (sageattention>=2.2.0)
- inference_engine.py: try-import sageattn, _attention_kernel() picks the
  active kernel based on env var, _patch_hiera_global_attention threads
  it through to the existing 4D Q/K/V layout.
- _validate_sage_attention(): runs sageattn vs SDPA on a dummy tensor
  shaped like the smallest Hiera global attn block; on divergence beyond
  5e-2 the env var is silently flipped off so users get SDPA fallback.
- 1-install.sh / .bat: opportunistic install on NVIDIA + CUDA12+. Same
  silent-fail pattern as torch-tensorrt and the Rust extension.

Default behaviour is byte-identical: CORRIDORKEY_USE_SAGE must be set
explicitly to opt in. 515 tests still pass.
@Kenshin9977
perf: overlap I/O with GPU work via prefetch + async writer
4830aeb 
Adapted from upstream PR nikopueringer#226 (nikopueringer/CorridorKey).

Sequential inference loop in backend/service/inference.py now runs:
- A background prefetch thread reading the next frame (input + alpha) into
  a depth-3 queue while the GPU processes the current one
- A single-worker ThreadPoolExecutor for output writes, so disk I/O runs
  in parallel with the next GPU step
- Pending writes are drained at end-of-clip; write failures retroactively
  flip their FrameResult to (success=False, warning=...) so the user sees
  the error in the queue panel.

Significant gain on 4K EXR sequences where disk I/O previously dominated.
PNG/JPG sequences see a smaller but still positive gain.

The boundary cleanups upstream removed (per-frame torch.cuda.empty_cache)
were never present in our service package — model-switch boundaries
already handle cache clearing.

515 tests still pass.
@Kenshin9977
perf: NVFP4 weight quantization on Blackwell + auto-detect accelerators
bf29ad2 
NVFP4 packs model weights to 4-bit microscaling format on RTX 50xx
(SM 12+) using torchao. ~halves memory bandwidth and compute time on
Blackwell's 5th gen Tensor Cores.

Refactors the SageAttention + NVFP4 enablement so users don't need to
flip env vars by hand:

- _auto_enable_sage(device) returns True when sageattention is installed
  and the GPU is Ampere or newer (SM >= 8)
- _auto_enable_nvfp4(device) returns True when torchao is installed and
  the GPU is Blackwell (SM >= 12)
- CORRIDORKEY_USE_SAGE and CORRIDORKEY_NVFP4 keep their meaning but are
  now overrides: 1 forces on, 0 forces off, unset = auto

DRY:
- _env_override(name) replaces three identical 1/true/yes parsing blocks
- _device_capability(device) replaces two get_device_capability try/except
- _attention_kernel(device) is the single source of truth for which
  attention function the patched Hiera blocks call into

Install scripts opportunistically install torchao alongside the
existing torch-tensorrt and sageattention. Silent-fail like the
others — no broken installs if the wheel can't build.

Also trims verbose docstrings and inline comments per review feedback,
and moves the torchao try-import to module level instead of inside
_try_nvfp4_quantize.

515 tests still pass.
@Kenshin9977
fix(install): unbreak ffmpeg check, surface install errors, skip inco…
0f7b716 
…mpatible accelerators

- scripts/check_ffmpeg.py: load backend.ffmpeg_tools as a normal package
  import (it's a directory now, not a single file). The importlib trick
  was for pre-deps lightweight loading but the script runs after step 4.
- 1-install.bat / .sh: stderr from each accelerator pip install goes to
  logs/install/<name>.log instead of /dev/null. Failure messages now
  point at the file so users can see the real error.
- maturin: install into the venv before running it, instead of bailing
  out with 'maturin not found'. Use the venv's python -m maturin.
- torch-tensorrt: removed from auto-install. The latest wheels pull
  torch 2.11 and break our 2.9.1 pin. Users with a compatible torch
  can run the install manually.
- sageattention / torchao: drop the version pins. PyPI ships sage 1.0.6
  only; >=2.2.0 fails to resolve. Sage 2.x/3.x require source build or
  GitHub-hosted wheels, document that path separately.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

1 more reviewer
@DrCheeseFace DrCheeseFace DrCheeseFace left review comments
Reviewers whose approvals may not affect merge requirements

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@Kenshin9977 @edenaion @DrCheeseFace