[diffusion] MiniMax-H3: tiered AdaLN plan cache (pinned-host tier + per-plan LRU) - #37266
Conversation
github-actions
Bot
added
documentation
BBuf
left a comment
BBuf
left a comment
There was a problem hiding this comment.
Change summary
This PR replaces MiniMax-H3's repeated 24.2-GiB AdaLN checkpoint scans with a tiered plan cache: fixed GPU plan slots with per-plan LRU, a pinned-host group cache, streamed native-checkpoint projection/all-gather, device-resident resolved slots for the denoise loop, and admission/configuration guards for online, sidecar, LoRA, and Diffusers-layout modes.
flowchart LR
A["Request timestep plans"] --> B["prepare_adaln_plans"]:::changed
B --> C{"GPU plan hit?"}:::changed
C -->|yes| D["Fixed GPU slab / per-plan LRU"]:::changed
C -->|no| E{"Pinned-host group hit?"}:::changed
E -->|yes| D
E -->|no| F["Stream checkpoint + TP projection"]:::changed
F --> D
D --> G["Denoise reads device plan slots"]:::changed
H["Weight update"]:::changed --> I["Retarget / invalidate derived cache"]:::changed
I --> B
classDef changed stroke-dasharray: 5 5,stroke-width: 2px;
At admission time, the full timestep plan is resolved into GPU slots. Misses can be refilled from pinned host pages or rebuilt by streaming native AdaLN tensors and applying the TP-sharded projection. The denoise loop then uses only device-side slot indices. Weight updates are intended to invalidate this derived state, but the two update sources below do not preserve that invariant. Dashed nodes are changed by this PR.
Findings
- [P1] A sidecar deployment keeps serving stale AdaLN conditioning after a weight update (
python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py:1538-1544). The transformer update has already succeeded when this hook runs, but the sidecar branch only warns and returns without invalidating or rejecting service. Subsequent requests combine the new transformer weights with a sidecar generated from the old checkpoint. This needs to fail before applying the update, atomically switch to a matching sidecar, or invalidate the cache and make requests fail until a compatible sidecar is installed. - [P1] Tensor-RPC updates invalidate the cache and then rebuild it from the startup checkpoint (
python/sglang/multimodal_gen/runtime/post_training/weights_updater.py:575).weights_path=Noneleavescache.weight_filespointing at the original checkpoint; the next miss rereads old AdaLN weights. Worse, AdaLN modules are pruned in cache mode, so incomingadaln_projtensors are logged as unknown and skipped while the API still returns success. Reject tensor updates in online-cache mode or retain a coherent updated AdaLN rebuild source/payload before publishing success.
The same source-version rule should be applied to disk updates whose target lacks native adaln_proj safetensors: warning and rebuilding from the original path is not fail-closed.
Review evidence
- Reviewed all 14 changed files, including GPU-slot ownership, host-tier refcounts/LRU, copy-stream barriers, TP capacity agreement, plan-key construction, admission, LoRA, and both weight-update APIs.
- Exhaustively swept 32,639 historical human-review threads for every touched path and broader cache/weight-update/LoRA/pinned-memory/concurrency terms. Exact MiniMax-H3 paths had no prior episodes; broader weight-update discussions consistently require a single coherent model version across an in-flight update.
- Consulted the MiniMax model history card. It covers M2/M2.5/M2.7 autoregressive serving and does not cover H3 diffusion or AdaLN-derived caches, so this review followed the changed H3 source paths directly.
git diff --checkand Python byte-compilation pass. The local checkout has notorch/pytest, so the new CPU-oriented tests could not be executed here. CI lint/check-change steps passed, but GPU/test lanes did not run because the PR lacks therun-cilabel. The PR also currently conflicts withmain.
Residual validation risk after the fixes: the author-listed multi-rank end-to-end test and fp32 trajectory comparison are still needed, plus explicit disk/sidecar/tensor update tests that prove the cache never serves a mixed weight version.
| "built against the previous weights; requests keep using its " | ||
| "stale conditioning until the sidecar is rebuilt" | ||
| ) | ||
| return |
There was a problem hiding this comment.
[P1] This warning does not fail closed: the transformer weight update has already succeeded, then this branch leaves the old sidecar active. Every later request mixes the new transformer weights with AdaLN conditioning generated from the previous checkpoint. Please reject the update before mutation, atomically install a matching sidecar, or invalidate the cache and make requests fail until a compatible sidecar is available. A regression test should update weights under sidecar mode and prove stale plans cannot be served.
There was a problem hiding this comment.
Fixed by moving the decision ahead of the mutation.
BaseDiT.validate_weight_update_source(weights_path=...) is a new hook that runs before a single weight is written — in update_weights_from_disk right after _validate_weight_files, and in update_weights_from_tensor before the first _load_weights_into_module. A rejection returns (False, msg) with the served model untouched, so the mixed-version window this warning left open no longer exists.
The H3 override rejects a sidecar deployment outright: the sidecar is generated offline from the startup checkpoint and the server cannot regenerate it, so the error names the two ways out (rebuild the sidecar and restart, or serve with --minimax-h3-adaln-online). refresh_weight_derived_caches keeps only the retarget, and re-asserts the same predicate so a broken call order raises instead of degrading silently.
Regression test: test_weights_updater_rejects_sidecar_update_before_writing_weights drives the real WeightsUpdater against a sidecar-mode DiT and asserts the update is refused, the module parameter still holds its pre-update value, and pipeline.model_path is unchanged — i.e. no request can observe new weights against the old sidecar.
| if isinstance(module, BaseDiT): | ||
| # Same invariant as the disk path; there is no on-disk source | ||
| # to retarget the rebuild at, so only the caches drop. | ||
| module.refresh_weight_derived_caches(weights_path=None) |
There was a problem hiding this comment.
[P1] For tensor-RPC updates, weights_path=None leaves cache.weight_files pointing at the startup checkpoint and only invalidates the plans, so the next request rebuilds AdaLN from the old weights. In cache mode the resident adaln_proj modules are also pruned, meaning incoming AdaLN tensors are logged as unknown/skipped while this API still returns success. Please reject tensor updates in online-cache mode or retain a coherent updated AdaLN payload/source before publishing success, and test that a post-update rebuild cannot read the original checkpoint.
There was a problem hiding this comment.
Tensor-RPC updates are now rejected in both cache modes, before any tensor is loaded.
validate_weight_update_source(weights_path=None) has no directory the rebuild could stream adaln_proj from, so online mode refuses the update and points at update_weights_from_disk; sidecar mode refuses it for the reason in the other thread. That also removes the second half of this finding — no adaln_proj tensor can arrive at a pruned module and be logged as unknown while the API reports success, because the request never gets that far.
The same source-version rule now covers the disk path you flagged in the review body: a target directory without native adaln_proj safetensors is rejected instead of warning and rebuilding from the original checkpoint. For a target that does carry them, the rebuild source is retargeted at the new directory, so the pruned resident modules are not a coherence gap on that path — the next miss reads the new checkpoint by construction.
Regression tests: test_weights_updater_rejects_tensor_update_in_online_cache_mode (real WeightsUpdater, asserts weights untouched), test_online_cache_rejects_update_source_without_native_adaln, and test_disk_update_retargets_rebuild_source_and_drops_plans, which asserts cache.weight_files points at the new directory and every previously built plan is gone.
Pure relocation out of the 2.6k-line minimax_h3.py ahead of the tiered-cache work; no behavior change.
…h the block gather prepare_adaln_plans returns each step's slab slot as a device tensor that forward consumes directly, replacing the per-step device-side lookup and its capture-breaking host sync. block_all() fetches all 50 blocks' params in a single gather; the derived slab buffers leave state_dict.
…ng the slab LRU over plan slots with in-flight protection: two alternating 50-step schedules both stay resident instead of triggering a full slab reset and a 24.2 GiB checkpoint re-read per request. Requests that cannot fit are rejected at admission, before the encode stages spend GPU time.
--minimax-h3-adaln-host-cache-gb (default 8/rank, 0 disables) keeps built plans in one pinned slab; a plan set evicted from the GPU slab swaps back in over PCIe instead of re-reading the checkpoint (repeated-schedule requests 94.8 -> 88.4 s on 1xH200, outputs bit-identical). Schedule-level LRU with recompute fallback; capacity is MIN-reduced across ranks and every state transition commits synchronously so ranks stay in lockstep.
Opt-in: compute the one-time projections in fp32 (TF32 pinned off) before the bf16 store. Off by default -- more accurate params, but not bit-comparable to resident adaln_proj weights, so it stays gated until an e2e trajectory check clears it.
…ight updates in cache modes adaln_proj LoRA deltas now raise instead of being silently dropped; the Diffusers-layout checkpoint is rejected at startup instead of a first-request KeyError; update_weights_from_disk invalidates every cache tier and retargets the rebuild source at the new shards.
…ache --minimax-h3-adaln-host-cache-gb is the only new flag; the GPU slot count and the fp32 rebuild are SGLANG_DIFFUSION_MINIMAX_H3_ADALN_* env escape hatches. Document the online mode and the host tier in cli.mdx and the cookbook.
Correctness: touch the GPU slab LRU on pure cache hits (hot schedules were evicted first), invalidate the cache on tensor-RPC weight updates via a BaseDiT hook (not just disk updates, and without importing the model zoo in the shared updater), check the native layout before retargeting the rebuild source, roll back reserved slots when a host swap-in fails, name the env var instead of a deleted CLI flag in the overflow error, raise the coverage error instead of a shape mismatch on over-width lookups, and downgrade the flag-misuse check to a warning (config-file construction marks every key explicit). Cleanup: drop the dead per-plan host timesteps copy (~50 stream syncs per stored group), key each plan once per request instead of twice (~98 D2H syncs), reuse the world-group MIN reduce and the shared host reserve constants, and deduplicate the unpin bookkeeping.
0d0d492 to
808c96d
Compare
…follow The cache holds values derived from adaln_proj, so an update it cannot follow leaves new transformer weights paired with the old checkpoint's conditioning. The previous hook ran after the weights were already written and only warned, so both the sidecar deployment and the tensor RPC path kept serving that mix. BaseDiT.validate_weight_update_source now runs before any weight is written and rejects a sidecar deployment, a tensor update, and a disk target without native adaln_proj safetensors; the refresh hook keeps only the retarget. BaseDiT.validate_lora_layers closes the same hole in the LoRA IPC path, where adaln_proj layers were logged as unknown, skipped, and reported as success.
808c96d to
fdb9cb1
Compare
|
Thanks for the review. Both P1s and the source-version rule from the review body are fixed, the branch is rebased on What changedThe coherence decision moved ahead of the mutation.
Seven CPU tests cover each rejection plus the accepted disk retarget, and two of them drive the real Multi-rank e2e4x H200,
Every output is bit-identical across arms and across repeats, and the host tier saves 6.8 s per repeated schedule at four ranks (single-rank was 6.4 s). The rank-0 AdaLN decisions were fp32 rebuild trajectory
Worth stating plainly: the fp32 projection changes bits by construction, and 50 sampling steps amplify that into a visibly different trajectory, not a "more accurate version of the same output". It stays off by default and the docs keep the experimental / validate-end-to-end wording, which these numbers now back with a measurement instead of an assertion. Test statusOn the rebased head: One thing I could not cover end to end: an actual weight-update e2e needs a second full checkpoint, so the "never serves a mixed weight version" property is covered by the The PR still needs the |
|
/tag-and-rerun-ci |
# Conflicts: # python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py
|
/tag-and-rerun-ci |
Conflict: prepare_lora_adapter gained _reject_non_lora_delta_tensors on main while this branch added the AdaLN cache-mode guard; both run, main's format check first.
prepare_lora_adapter reads _adaln_precomputed to decide whether adaln_proj deltas are servable, so the SimpleNamespace standing in for the model needs that attribute; without it the guard raises AttributeError before reaching the rejection the test asserts.
|
CI triage on
That merge needed one follow-up. The remaining failures are The Separate from this PR, worth a look: The AMD lanes are still running; I have not seen their logs yet and will follow up if they fail for a reason of their own. |
|
/tag-and-rerun-ci |
|
Following up on the AMD lanes, as promised. Both failures reproduce on
Evidence for the AMD claim: So the only thing standing between this PR and a green CUDA lane is the |
|
Update: merged current The merge was conflict-free, and MiniMax/VSA/FastH3 unit tests are 161 passed on the merged head. The AMD lane failures stand as described above — they reproduce on |
|
Full CI has now settled on The 10 remaining reds are all outside this PR:
So the only actionable item left on my side is none; a re-run of |
…er-plan LRU) (sgl-project#37266) Co-authored-by: mickqian <mickqian@users.noreply.github.com>
3 participants
Motivation
Under
--minimax-h3-adaln-online, any plan miss re-reads all 24.2 GiB ofadaln_projweights from the checkpoint (nsys-measured 5.8-6.7 s of GPU idle per request), and slab overflow triggers a full reset, so two alternating 50-step schedules re-read the checkpoint on every request. This PR adds per-plan LRU eviction and a pinned-host tier so previously built plans swap back in over PCIe instead. Repeated-schedule requests drop 94.8 -> 88.4 s on 1x H200 with bit-identical outputs.The schedule space (task shape,
num_inference_steps,flow_shift,audio_flow_shift) is request-controlled in serving, so the combination count is unbounded; the cache keys stay "the exact fp32 bit patterns of one step's deduped timesteps", which every combination collapses into.Modifications
--minimax-h3-adaln-host-cache-gb(default 8 GB per rank, 0 disables): one pinned slab in 9.25 MiB pages, LRU at schedule granularity because a rebuild pass streams every layer regardless of how few plans it misses; over-capacity schedules recompute instead of raising. Capacity is MIN-reduced across ranks and state transitions commit synchronously inside prepare, keeping multi-rank collectives in lockstep.prepare_adaln_plansand passed to forward as a device tensor, removing the per-step device-side lookup and its capture-breaking host sync;block_all()replaces 50 per-block gathers with one.BaseDiT.validate_weight_update_sourceandvalidate_lora_layersrun before a single weight is written, so an update the cache cannot follow is refused with the served model untouched instead of pairing new weights with the old checkpoint's conditioning.MiniMaxH3AdalnCachemoves into its own module (pure relocation). Expert knobs are env vars (SGLANG_DIFFUSION_MINIMAX_H3_ADALN_GPU_PLANS,SGLANG_DIFFUSION_MINIMAX_H3_ADALN_FP32); the host-cache size is the only new CLI flag.Behavior changes: none in the default numerics path (the rebuild stays bit-exact with resident
adaln_projweights: per-plan-M GEMM, TP-sharded N + all-gather). The env-gated fp32 rebuild changes bits by construction and is off by default. LoRA adapters touchingadaln_projnow raise in cache modes instead of silently using base weights. Runtime weight updates are constrained in cache modes: online mode accepts a disk update whose target carries nativeadaln_projsafetensors and rejects everything else (tensor RPC, and directories without those tensors); a sidecar deployment rejects weight updates outright, since the sidecar is an offline artifact the server cannot regenerate.Accuracy Tests
A/B (host tier 8 GB vs disabled), 1x H200, t2va,
--minimax-h3-adaln-online, 50 steps, seed 1101, two alternatingflow_shiftvalues (12.0 / 10.0), 4 requests per arm in one server process: all 8 MP4s are sha256-identical across arms, and repeated requests are sha256-identical to their first occurrence, i.e. both the rebuild and the pinned-host round-trip are bit-faithful. Unit tests cover LRU order, alternating-schedule residency, host-cap accounting and pinned-allocation degradation, recompute fallback, bit-exact swap-in, the fp32 projection, the LoRA/Diffusers/weight-update guards, and admission-time rejection; the full multimodal_gen unit suite passes (the one remaining failure reproduces on main).Speed Tests and Profiling
Same A/B setup as above (
sglang generate-equivalent requests throughDiffGenerator, warmup excluded):The 6.4 s saved per repeated request matches the nsys-measured rebuild cost (5.8-6.7 s). An 8-step run with a deliberately small slab (
SGLANG_DIFFUSION_MINIMAX_H3_ADALN_GPU_PLANS=8) shows the same behavior for distilled schedules. Multi-rank (4x H200,ulysses=4, tp=1, same t2va schedule, 50 steps) reproduces the single-rank result: 31.95 / 31.88 s for the two first-seen schedules, 25.16 / 25.13 s when each repeats, every output bit-identical across the host-tier-on and host-tier-off arms and across repeats, no rank divergence or hang. The env-gated fp32 rebuild was measured against the default path on the same requests: 29.10 dB / 24.08 dB video PSNR (mean abs diff 3.42 / 6.09 on u8), i.e. it is a different sampling trajectory rather than a refinement of the same one -- it stays off by default and end-to-end validation remains required before using it.Conclusion
Mixed-schedule serving no longer pays a 24.2 GiB checkpoint read for any schedule it has seen before, at the cost of one pinned-host slab per rank, with the default path bit-exact and one new CLI knob.
Checklist
CI States
Latest PR Test (Base): ✅ Run #33694613256
Latest PR Test (Extra): ❌ Run #33694612785
Latest PR Test (AMD ROCm 7.2): ❌ Run #33694613236