[Model][MiniMax-H3] Multi-GPU serving fixes and layer-wise offload component selection

[heyuanliu-intel]

Purpose

Bringing MiniMax-H3 up on multi-card XPU serving surfaced a series of modeling-layer
defects: VAE residency fought between the pipeline and the offload backends, a padded
packed sequence that Ring attention rejects, reference images silently rescaled,
trust_remote_code never reaching the checkpoint VAEs, and a gather buffer that grew
device memory once per request. Fixing them also needed two capabilities the offloader
did not have: paging a declared encoder block stack in the single-process layer-wise
backend, and choosing which component families layer-wise offload manages at all.

1. [Bugfix] Keep MiniMax-H3 pipeline-staged VAEs out of offloader discovery

Problem: The MiniMax-H3 pipeline stages its video and audio VAEs itself: it
calls load_to_device() / offload_to_cpu() around every encode and decode so
that only one VAE is resident at a time. The pipeline nevertheless declared both
VAEs in _vae_modules, so ModuleDiscovery reported them as regular VAE
components. SupportsComponentDiscovery documents _vae_modules as "VAE(s)
(always on GPU)", and consumers act on that contract: the layer-wise offload
backend moves every discovered VAE to the device when offloading is enabled, and
the HSDP loader path moves them after sharding. Both fight the pipeline's own
staging and pull the VAEs back onto the accelerator, costing memory that
layer-wise offloading is supposed to save.

Fix: Declare no VAE components for MiniMax-H3. The pipeline-staged VAEs are
simply not visible to component discovery, so no offload backend and no model
loader relocates them, and the discovery contract holds again for every other
pipeline. Because the VAEs are no longer discovered, the on-demand component set
in the offload plan drops them as well (only the text encoder stages on demand),
and the model-level CPU offload entry point now lists its stages explicitly
instead of deriving them from discovery, which keeps the text encoder and both
VAEs in model-level offload exactly as before. The pipeline.vae alias is kept:
registry-side VAE patch parallelism reads it, and it is not a discovery
attribute for a pipeline that declares components explicitly.

The user guide sentence about resident VAE modules is reduced to one sentence
naming pipeline-staged components rather than describing a backend special case.

2. [Bugfix][MiniMax-H3] Release decoded outputs before the next seed

Problem: with model-level CPU offload, #6072 evicts the DiT before VAE decode, but the decoded output stays on the device until the request returns. For num_outputs_per_prompt > 1 the appended output of seed N is still resident when seed N+1 reloads the DiT, so the peak the eviction removed comes back; with a large output even the first request can OOM in decode.

Fix: on the executor reply rank, move each decoded output to the host inside the same seed iteration while active model-level offload is installed; non-reply ranks keep a metadata placeholder. Video is released after _prepare_minimax_h3_video_output quantizes to uint8, so #6824's transfer optimisation is kept. Other offload modes and the no-offload path are unchanged.

3. [Bugfix][Hardware][Intel] MiniMax-H3: reuse the tiled-VAE gather buffer across requests on XPU

Problem: the tiled video VAE gathers every rank's decoded tiles with a variable-shape all-gather whose output is allocated afresh each call. After the per-request empty_cache() the next gather lands at a new device address. On XPU the collective backend registers a device-side resource per new receive address and does not reclaim it, so device memory outside torch.xpu.memory_allocated() grows every request until the card is full. (Mechanism inferred from intervention experiments, not observed directly.)

Fix: keep one landing buffer per (shape, dtype, device) on the adapter and gather into it with all_gather_into_tensor; tiles are cloned out because the buffer is reused. Installed only when the adapter's device is XPU, bound on the checkpoint instance; other backends keep the checkpoint's own method. The buffer cache is uncapped on purpose: the key is the tile geometry, and evicting would re-register a new address.

4. [Bugfix] Fix MiniMax-H3 packed padding with Ring attention

Problem: MiniMax-H3 pads its packed sequence to an aligned length and passes valid_kv_length to the local attention backends. On the Ring path it instead builds an attn_mask at the global padded length (minimax_h3_transformer.py:514), which the Ring backends reject, so any padded request fails under Ring (the failure recorded in recipes/MiniMaxAI/MiniMax-H3.md).

Fix: both Ring backends accept an optional valid_kv_length and trim each K/V block to the valid prefix of its source rank (ring_utils.ring_kv_block_valid_length); a fully padded block launches no attention. H3 publishes valid_kv_length on the Ring path instead of the mask. causal=True together with valid_kv_length raises ValueError: on a rectangular tile the causal diagonal would silently shift. With valid_kv_length=None every path is byte-for-byte the previous behaviour.

5. [Bugfix] Honor runtime SP state in H3 Ring padding

Problem: MiniMaxH3Attention._run_packed_attention reads the static use_ring flag in three places to decide between the Ring and local packed paths. Whether Ring actually runs is a runtime decision (Attention._get_active_parallel_strategy also honours skip_sequence_parallel and forward_context.sp_active), so a layer that declares Ring but is not running it this step takes the Ring branch with the wrong inputs.

Fix: decide from the active strategy through one local helper; a layer not running Ring this step takes the local packed path (including NPU npu_attn_varlen). A layer genuinely running Ring behaves as before. Depends on the previous commit for the Ring branch.

6. [Bugfix][MiniMax-H3] Preserve reference image geometry

Problem: serving_video rescales every reference image to the request canvas. MiniMax-H3 expects reference images at their native geometry, so ref2va silently replaces the reference aspect ratio with the canvas one.

Fix: register an H3-specific reference_image_size_resolver through the existing model_extras/registry.py extension point that keeps the source geometry. Pipelines without a resolver keep the serving-side resize.

7. [Bugfix][MiniMax-H3] Respect trust_remote_code when loading checkpoint VAEs

Problem: the pipeline never passes trust_remote_code to the two checkpoint-supplied VAE constructors in pipeline_minimax_h3.py, so a checkpoint that ships its own VAE class either fails to load or executes remote code the user did not opt into (the only existing use is encoder.py:1118, hard-coded False).

Fix: thread od_config.trust_remote_code into both VAE adapters in vae.py; they forward it and default to refusing, with an error naming the component path and class.

8. [Model][MiniMax-H3] Let a request pin the packed-sequence length

Problem: compiled regions are keyed by input shape. The t2va/fl2va packer rounds the used rows up to the next 64-row boundary, so requests whose prompt length, frame count or frame size differ land on different packed lengths, and a fixed-shape deployment recompiles per length. The ref2va packer already accepts an explicit seq_len; the t2va/fl2va packer has no such contract.

Fix: share one resolver (_resolve_padded_seq_len) between the two packers and thread extra_args["pad_seq_len"] into both, validated at the request boundary (integer, multiple of 64, at most MINIMAX_H3_MAX_PAD_SEQ_LEN, not below the used rows). Without the pin both packers round up as before.

9. [Bugfix] Honor declared encoder block stacks in layer-wise offload

Problem: LayerWiseOffloadBackend.enable() pins every encoder on the device for the whole request. A pipeline can declare which encoder submodules are streamable block stacks (OffloadPlan.encoder_block_attrs) and the distributed backend honours that, but the single-process backend never reads it, so a large vision-language encoder stays resident, which is most of the memory layer-wise offload exists to free. disable() also left encoder hooks installed.

Fix: read the declaration in the single-process backend and page the declared block stacks like DiT blocks; non-block state (norms, embeddings, projections) is still placed on the device. disable() tears the encoder hooks down and restores each block's weights to the host through LayerwiseOffloadHook.restore_next_block_to_host(), rebuilding them with torch.as_strided from the stored stride so transposed (non-contiguous) weights keep their element order; prefetch_layer() is not used because it would put every block on the accelerator at once.

10. [Core] Let layer-wise offload select which component families it manages

Problem: layer-wise offload is all-or-nothing. There is no way to say "manage the encoders and VAEs, keep the DiT whole", which is the topology to run whenever the DiT fits on the device but the encoders and VAEs do not.

Fix: layerwise_offload_components, a comma list from {dit, text_encoder, vae}. A family left out stays fully device-resident: excluded dit is placed whole with no block hooks; excluded text_encoder is placed with no streaming and no parking even if the pipeline declares block stacks or on-demand; excluded vae is placed even if declared on-demand. Omitting the option selects every family (byte-for-byte today's behaviour). Unknown or empty selections fail at config validation.

11. [Model][MiniMax-H3] Honour the layer-wise component selection when staging encoders and VAEs

Problem: layerwise_offload_components promises that an excluded family stays device-resident, and the backend keeps that promise. The MiniMax-H3 pipeline manages encoder and VAE residency itself and never reads the selection: _component_on_device() stages purely from enable_layerwise_offload, and the construction site builds both VAEs on the host. With --layerwise-offload-components dit,text_encoder the backend places the VAEs on the device and the pipeline pulls them back to the host after every decode, a full round trip per request for weights the operator asked to keep resident.

Fix: route the three sites through _stages_component_family(), which reads the selection under plain layer-wise offload and keeps returning True under distributed layer-wise offload (that backend shards every component).

12. [Bugfix][MiniMax-H3] Always stage the pipeline-owned VAEs

Problem: Two changes in this series meet on the MiniMax-H3 VAEs and cancel
each other out. Keeping the pipeline-staged VAEs out of component discovery
means no offload backend ever places them: the pipeline is the only owner of
their residency. Routing the pipeline's staging sites through
_stages_component_family() then makes that ownership conditional on
--layerwise-offload-components naming the vae family. With
--layerwise-offload-components dit,text_encoder the VAEs are consequently
neither placed by the backend nor staged by the pipeline: they stay wherever
they were loaded and run from there, which is neither the "fully
device-resident" contract the selection publishes for an excluded family nor
the low-residency behaviour the operator enabled layer-wise offload for.

Fix: Honour the selection only for the families component discovery
actually reports. _stages_component_family("vae") now returns True whenever
layer-wise offloading is active, so the pipeline keeps staging its own VAEs
around every encode and decode and keeps building them in host memory; the
text_encoder family is unchanged and still follows the selection. The user
guide and the --layerwise-offload-components help text say so: the vae
family does not change what the backend does with discovered VAEs, and
pipeline-staged VAEs are always staged by the pipeline.

Changes outside the MiniMax-H3 directory

Every file below is outside vllm_omni/diffusion/models/minimax_h3/. For each one:
what changed, whether the default behaviour of any other model changes, and the
evidence for that answer.

Attention (Ring path, shared by every model that runs Ring sequence parallelism)

File	Change	Default behaviour for other models	Evidence
vllm_omni/diffusion/attention/backends/abstract.py	Comment only: the valid_kv_length metadata key is documented as usable by Ring backends, not only by backends advertising supports_prefix_kv_slicing.	Unchanged.	The diff contains no executable line.
vllm_omni/diffusion/attention/backends/ring/ring_utils.py	Adds ring_kv_block_valid_length(), a pure function mapping a circulated block index to its valid prefix length.	Unchanged: new symbol, no existing call site rewritten.	tests/diffusion/attention/test_ring_padding.py covers the helper directly.
vllm_omni/diffusion/attention/backends/ring_flash_attn.py, ring_pytorch_attn.py	Both accept an optional valid_kv_length; each circulated K/V block is trimmed to the valid prefix of its source rank, and a fully padded block launches no attention. causal=True combined with valid_kv_length raises ValueError instead of silently shifting the causal diagonal.	Unchanged when valid_kv_length is None, which is the value for every caller that does not pass it.	The None path is the pre-existing code path; tests/diffusion/attention/test_ring_padding.py asserts the unpadded results are identical to the previous behaviour and that the causal combination raises.
vllm_omni/diffusion/attention/parallel/ring.py	Forwards valid_kv_length from attention metadata to the two backends.	Unchanged: metadata without the key forwards None.	Same tests; the key is only published by the MiniMax-H3 transformer.

Layer-wise offloader

File	Change	Default behaviour for other models	Evidence
vllm_omni/diffusion/offloader/layerwise_backend.py	Pages the encoder block stacks a pipeline declares in OffloadPlan.encoder_block_attrs (previously honoured only by the distributed backend), tears those hooks down in disable(), and gates the DiT / encoder / VAE handling on the new component selection. Extracts stream_declared_encoder_blocks() as the shared helper.	Unchanged for a pipeline that declares no encoder_block_attrs, and unchanged for every pipeline when the selection is left at its default (all families).	tests/diffusion/offloader/test_layerwise_backend.py and tests/diffusion/offloader/test_layerwise_component_selection.py; the declaration is opt-in per pipeline, and the default selection is the full family set (OffloadConfig.layerwise_components).
vllm_omni/diffusion/offloader/distributed_layerwise_backend.py	Refactor only: its private encoder-streaming block is replaced by a call to the extracted helper.	Unchanged.	The helper is the same logic moved verbatim; the distributed lifecycle tests cover it.
vllm_omni/diffusion/offloader/base.py	Adds OffloadConfig.layerwise_components, defaulting to all three families, populated from the diffusion config.	Unchanged by default.	The default is frozenset({"dit", "text_encoder", "vae"}), which reproduces the previous all-or-nothing behaviour.

Configuration and entrypoints (plumbing for one new option)

File	Change	Default behaviour for other models	Evidence
vllm_omni/diffusion/data.py	Adds layerwise_offload_components plus layerwise_component_selection() and validation that rejects unknown or empty selections.	Unchanged: None selects every family.	layerwise_component_selection() returns the full family set when the option is unset; covered by tests/diffusion/offloader/test_layerwise_component_selection.py.
vllm_omni/config/omni_config.py, vllm_omni/config/stage_config.py, vllm_omni/engine/arg_utils.py, vllm_omni/engine/async_omni_engine.py, vllm_omni/entrypoints/cli/serve.py	One optional field / argument each, defaulting to None, carrying --layerwise-offload-components to the diffusion config.	Unchanged when the flag is not passed.	Each diff is a single None-defaulted declaration plus the argument definition.

Serving-side model extras

File	Change	Default behaviour for other models	Evidence
vllm_omni/model_extras/registry.py	Registers a reference_image_size_resolver for the two MiniMax-H3 pipeline classes through the existing extras mechanism.	Unchanged: pipelines without a resolver keep the serving-side resize.	The new entries are keyed by MiniMaxH3Pipeline and MiniMaxH3ModularPipeline; tests/model_extras/test_model_extras.py and tests/entrypoints/openai_api/test_video_server.py assert that a pipeline without the resolver still rescales.

Tests and documentation outside the model directory

tests/diffusion/attention/test_ring_padding.py (new), tests/diffusion/offloader/test_layerwise_backend.py,
tests/diffusion/offloader/test_layerwise_component_selection.py, tests/entrypoints/openai_api/test_video_server.py
and tests/model_extras/test_model_extras.py cover the changes above.
docs/design/feature/offloader/layerwise_offload.md, docs/user_guide/diffusion/offloader/layerwise_offload.md,
docs/design/feature/sequence_parallel.md,
docs/user_guide/diffusion/regional_compilation.md and recipes/MiniMaxAI/MiniMax-H3.md document the new option,
the Ring padding contract and the packed-length pin.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[vllm-omni-review-bot]

This PR appears to be related to model: MinimaxH3.

Model owners: @david6666666

@heyuanliu-intel, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Resolved as of f8589e22ba83: the high-priority or low-quality signal noted on an earlier commit no longer applies.

[hsliuustc0106]

please help review PR #5929, we will land that PR for layerwise offload component selection first

[hsliuustc0106]

I also suggest to split this PR into serverl independent PRs

[xuechendi]

I also suggest to split this PR into serverl independent PRs

Oh, I see, I actually asked @heyuanliu-intel to create a merged PR because splits are quite many , here is his previous list

#6990 [Bugfix][MiniMax-H3] Respect trust_remote_code when loading checkpoints
#6992 [Bugfix][MiniMax-H3] Preserve reference image geometry
#6993 [Bugfix] Stop layer-wise offload from undoing pipeline-managed VAE staging (v2)
#6995 [Bugfix] Fix MiniMax-H3 packed padding with Ring attention
#6996 [Bugfix][Hardware][Intel] MiniMax-H3: reuse the tiled-VAE gather buffer across requests on XPU
#6997 [Bugfix][MiniMax-H3] Release decoded outputs before the next seed

[heyuanliu-intel]

Thanks both. Since #5929 landed, I have re-based this PR on top of it and dropped everything it supersedes:
the four layer-wise component-selection commits and the VAE-staging change from #6993 are gone (the merged
#5929 handles both cases — legacy mode now registers the pipeline-owned VAE as staged instead of moving it).
What remains is 7 commits, all confined to the MiniMax-H3 model directory and its tests: decoded-output
release, tiled-VAE gather-buffer reuse on XPU, the two Ring-attention padding fixes, reference-image geometry,
trust_remote_code plumbing, and the padded sequence-length fix. CPU suites and pre-commit are clean on the
new head; each commit still reverts independently. We will keep this as one PR as discussed with @xuechendi;
at 7 self-contained commits it reviews commit-by-commit.

[hsliuustc0106]

[P2] Bound persistent gather memory across request geometries

Each new tile shape permanently retains another full gather buffer on every XPU rank. VAE offload and empty_cache() cannot release these dictionary-held tensors, so varying request geometries accumulates device memory and can eventually OOM. Stable-address reuse fixes repeated identical shapes, but needs a bounded workspace or an enforced, memory-budgeted set of geometry buckets.

[hsliuustc0106]

intel ci failed

[heyuanliu-intel]

intel ci failed

Thanks for flagging. The Intel CI build for 6f2a58c (#8014) failed after 5 minutes, well short of the 30–50 minutes the same pipeline took on the neighbouring passing builds (#8013, #8015, #8016), and two unrelated builds in the same window failed the same short way (#8012 on main, #8020 on #7336). The AMD failure on that head (#11615) also coincides with failures on two unrelated PRs in the same hour (#11614 on #7155, #11618 on #6956). None of the tests the Intel job selects (-m "xpu and B60" over tests/diffusion, tests/dfx and tests/e2e) touch MiniMax-H3, and the commits here only change the MiniMax-H3 model directory, the Ring-attention padding path and one model_extras registry entry.

I cannot see the Buildkite job log from outside the org. I have just pushed 93f0323 (addressing the P2 comment above), which re-triggers both pipelines; if the failure reproduces on this head, could someone with access paste the failing step? I will fix it right away if it is real.

[xuechendi]

@heyuanliu-intel , please resolve the conflict. BTW, we have upgrade vllm-omni xpu version to 0.29 last friday, please help to test locally on XPU, thanks

[heyuanliu-intel]

@xuechendi Conflict resolved (rebased onto current main, merges clean).
Tested locally on XPU with vllm/vllm-openai-xpu:v0.29.0 (vllm 0.29.0, torch 2.13.0+xpu):
ran the same 4 pytest suites as the Intel CI lane plus the MiniMax-H3 unit tests,
on both main and this branch — identical results, zero failures, 472 H3 tests pass on this branch.

[xuechendi]

@heyuanliu-intel , can you add a before / after comparison (perf / memory footprint ) with e2e test on BMG as well; Also put output for quick exam