[mickqian]

Motivation

FLUX.2 Klein's first-request regression surfaced in #45919 but reproduces without its progress-reporting changes. A GC-instrumented H100 run captured a 449.5 ms generation-2 collection inside denoising. The CI bisection identified #46047 as the conditional trigger, not a Klein kernel regression.

This PR collects at initialization/warmup boundaries and prevents automatic cyclic GC from interrupting protected worker execution. With server warmup the extra cleanup is paid before readiness; with request-based warmup the first caller pays it. It does not modify #45919, performance thresholds, or GT.

Changes

• Collect after worker initialization and the final warmup, including compile prewarms. Handle request-based warmup, sequential outputs, and recoverable warmup failures.
• Share a thread-safe GC guard across model forward, deferred finalization, and output replies. Nested/overlapping execution restores the original GC state only after the last owner exits, including exceptions and forked children.
• With async output saving, drain continuous overlapping finalizers at an admission boundary after 30 seconds without a GC opportunity. Preserve ordinary overlap otherwise.
• Document the behavior and tradeoffs; extend existing worker, async-output, and OOM tests.

Reference counting and explicit gc.collect() still work. There is no gc.freeze(), threshold change, new CLI flag, or GPU allocator flush. GC can still run outside protected regions. Cyclic garbage may survive until a request finishes; the 30-second rule bounds continuous overlap, not a single long request. The drain can reduce overlap at that boundary.

Validation

• Merged main (fd5c52588ff) at 3f4a75b2809, retaining [Feat] diffusion/openai: Report per-step progress for async video jobs #45919's progress reporter alongside the GC guard. The only conflict was the import block.
• Merged head: 195 passed, 14 subtests passed, 1 skipped on an isolated H200 devbox, including video-progress/health tests. The skipped test requires two CUDA devices.
• H100 validation before the main merge: 143 passed, including single-GPU Klein and two-GPU Z-Image E2E with existing performance, memory and consistency checks. Merged-head validation is dispatched.
• Pre-commit and git diff --check pass.
• Same H200, Python 3.12.15, PyTorch 2.14.1+cu130, FLUX.2-klein-4B, 1024x1024, four steps, seed 42, quality=exact, server warmup, no torch.compile. Both arms use identical GC callback instrumentation and the same editable environment.

The comparison below isolates the new request guard from the already-present startup collection (a8b84ce6). The fixed measurement uses 264d0863; the subsequent commit only fixes sequential warmup handling and its tests.

80 HTTP requests per concurrency	Startup-only	Request protection
Concurrency 1 throughput, req/s	3.110	3.118
Concurrency 1 P99, ms	324.95	323.18
Concurrency 4 throughput, req/s	3.139	3.144
Concurrency 4 P99, ms	1280.86	1277.39
GC collections inside forward, including warmup	45	0

Outputs are byte-identical. The throughput differences are below 0.3% and are not evidence of a meaningful steady-state speedup. The direct result is removal of GC from the protected path, not faster GPU kernels.

With async output saving, a further 600 requests per arm at concurrency 4 give 3.158 vs 3.165 req/s. Forward GC, including warmup, drops from 38 to 0. Peak reported GPU memory is 12,996 MiB in both arms; sampled whole-device usage is 14,789 MiB in both. Between the first and sixth 100-request batches, aggregate worker RSS grows by only 0.24/0.27 MiB. These are several-minute tests, not a production-duration leak proof.

Request-based and disabled warmup also pass real HTTP generation with unchanged output hashes and zero observed forward GC.

Z-Image-Turbo, 1024x1024, nine steps, 20 prompts at each concurrency, explicitly --component-residency all=resident: current main fd5c52588ff vs merged head 3f4a75b2809 gives 1.087/1.087 req/s at concurrency 1 and 1.112/1.112 req/s at concurrency 4. P99 is 923.70/926.37 ms and 3613.48/3607.57 ms, respectively. Per-request hashes match, and both report 23,714 MiB peak GPU memory (25,493 MiB sampled whole-device usage). Forward GC including warmup drops from 34 to 0; the main trace includes a full collection. This is neutral steady-state performance, not a speedup.

The initial Z-Image auto-placement comparison is excluded from performance claims: downloading its metadata changed TE offload to resident. Explicit placement removes that confound.

Original Regression Evidence

• Regressed H100 + GC tracing: denoise 545.7 ms, including a 449.5 ms full GC.
• Startup fix + GC tracing: denoise 228.1 ms, no full GC during the request.
• Startup-only final validation: three fresh-server requests passed unchanged performance, VRAM, and consistency checks, E2E 374.72 / 377.02 / 445.03 ms. GPU reserved/allocated peaks remained 13,620/10,791 MiB, with byte-identical JPEGs.

The historical H100 tracing runs used separate runners with six ancillary package-version differences. They identify the GC stall, but are not a perfectly controlled same-host throughput comparison. The newer H200 comparison above uses one isolated device and unchanged packages.

---

CI States

Latest PR Test (Base): ✅ Run #40265313152
Latest PR Test (Extra): ❌ Run #40265313171
Latest PR Test (AMD ROCm 10): ❌ Run #40265313155

[mickqian]

/tag-and-rerun-ci

[mintlify]

Preview deployment for your docs. Learn more about Mintlify Previews.

Project	Status	Preview	Updated
sglang-doc	🟢 Ready	View Preview	Oct 26, 2026, 3:46 AM

💡 Tip: Enable Automations to automatically generate PRs for you.