[mglyn]

This draft evaluates model-specific eager optimizations for the MiniMax-H3 video VAE decoder. The registered hardware entries target CUDA compute capabilities 9.0, 10.0, and 10.3, while hardware selection stays behind a flat operator-set dispatch so another backend can be added without changing the model integration. It is a follow-up and alternative optimization direction to #5979: instead of regionally compiling every VAE Transformer block, it accelerates the decoder's repeated eager operations.

The model-local placement is intentional: numerical correctness and performance depend on H3's exact math, dtype, layout, execution contract, and per-platform validation, so its dispatch and fallbacks remain beside the model integration. The adjacent ops/README.md records this boundary and the evidence expected when extending platform support. If #6305 establishes a suitable shared operator tier, migration can be a follow-up.

Every optimized operator now preserves the reference output bit for bit. Q/K normalization and RoPE use one fused Triton kernel per Q/K tensor while retaining the reference RMS reduction order and explicit FP16 multiply/add rounding. Every eligible call uses the same exact path. Unsupported devices, tensor contracts, spatial-parallel attention, and compile execution retain the original module operations.

Kernel development followed MIT HAN Lab Kernel Design Agents, the evidence-first workflow from KDA-Pilot.

Operations

Optimization	Implementation	Numerical result	Reference	Candidate	Speedup
Decoder-block Linear materialization	Materialize the FP16 weights already selected by decode autocast once at load time	Complete decoded tensor remains bitwise equal; isolated contribution not measured	—	—	Included in full decode
Q/K normalization + RoPE	Single fused Triton reduction and RoPE kernel per Q/K tensor; the validated SM90/SM103 warp layout preserves the RMS reduction tree and PTX preserves FP16 multiply/add rounding	torch.equal for Q and K	0.256031 ms	0.051836 ms	4.94x
SwiGLU	Reuse Omni SiluAndMul after the FP16 w1 output	torch.equal	0.084760 ms	0.025671 ms	3.30x
Scaled residual	Triton FP32 residual + FP16 branch * FP32 scale, with contraction disabled	torch.equal	0.028998 ms	0.022047 ms	1.32x

Only decoder Transformer-block Linear weights are materialized in FP16. Keyframe encoding and non-block decoder parameters, including proj_out, remain FP32.

Complete VAE decode

The exact fused Q/K path was checked on one NVIDIA H200 with the official 1344x768 output shape, 209 output frames, latent shape [1, 24, 62, 48, 84], and decoded tensor shape [1, 3, 209, 768, 1344]. The complete decoded tensor is bitwise equal to the reference Q/K implementation (torch.equal=True, max absolute error 0).

The same operator set was independently exercised on B300/SM103 before dispatch registration: the complete decoded tensor remained bitwise equal, 512 direct exactness comparisons passed, and complete VAE decode P50 improved from 8.7199 s to 4.4145 s. The current head explicitly registers SM90, SM100, and SM103; other architectures retain the reference fallback.

Turbo end-to-end performance

The controlled A/B used the MiniMax-H3 Turbo v1.0 FL2VA recipe on NVIDIA H200 GPUs: 768x1344 (width x height), 107 frames at 24 FPS, 4 denoiser evaluations, regional DiT compile, and the same prompt, image, and seed. The single-GPU case used USP1 and VAE tile PP1; the four-GPU case used USP4, ring 1, and VAE tile PP4. Both commits received two complete warmup requests followed by three measured requests. Stage 0 comes from the online benchmark response. VAE decode comes from the pipeline profiler; the four-GPU value is the per-request critical path, or the maximum across ranks.

GPUs	Metric	Main P50	This PR P50	Speedup	Latency reduction	Main P100	This PR P100
1	Stage 0 end to end	28.449 s	25.920 s	1.10x	8.89%	28.459 s	25.939 s
1	VAE decode	6.066 s	3.546 s	1.71x	41.54%	6.066 s	3.548 s
4	Stage 0 end to end	9.599 s	9.059 s	1.06x	5.63%	9.599 s	9.063 s
4	VAE decode	1.653 s	1.001 s	1.65x	39.40%	1.653 s	1.026 s

Diffuse P50 remained effectively unchanged: 21.578 s versus 21.593 s on one GPU and 5.960 s versus 5.964 s on four GPUs. This isolates the end-to-end improvement to VAE decode. The baseline is main commit c6df39e9f; the single-GPU candidate is f67759666, and the four-GPU candidate is 6030e5976.

Compile execution bypasses these operators and retains the original expressions for Inductor.

Integration and fallback

• The VAE operators expose one exact implementation.
• The remote-code VAE structure is validated before installation.
• Installation is idempotent and leaves unsupported targets untouched.
• Shape, dtype, layout, gradient, platform, and execution-mode guards fall back to the original operations.
• The loaded remote module's Q/K normalization behavior and FF execution state are checked before installation.
• The current operator set is selected only on compute capabilities 9.0, 10.0, and 10.3; adding hardware support means registering another complete operator set in the dispatch table.

Validation

The final validation was repeated after rebasing onto main commit c6df39e9f.

Scope	Configuration	Result
Changed files	Targeted pre-commit and git diff --check	PASS
MiniMax-H3 model tests	tests/diffusion/models/minimax_h3, excluding the external full-model case	257 passed, 17 skipped, 1 deselected
VAE operator tests	20 cases covering kernels, dispatch, installation, and fallback	PASS
Single-GPU eager A/B	NVIDIA H200, T2VA, 1344x768, 209 frames, 5 sigma points, seed 42	Main and optimized eager output signatures bitwise equal
Official I2VA E2E	4 GPUs, 1344x768, 192 frames, 50 steps, USP4 + HSDP4 + text-encoder TP4 + VAE tile PP4	PASS; official-reference thresholds SSIM >= 0.97 and PSNR >= 34 dB
Full VAE numerical check	Official-shape latent and complete decoded video tensor	Bitwise exact, max absolute error 0
Q/K stress matrix	256 shape, seed, and value-distribution cases	All bitwise exact

The operator tests cover production and regression token lengths, exact outputs, unsupported-input fallback, hardware dispatch, idempotent installation, official-semantics guards, incompatible remote-code structure, selective Linear conversion, and the post-Linear FP16 SwiGLU path.

[hsliuustc0106]

model specific only?

[mglyn]

This PR is an experiment to explore the practical boundary of eager, model-specific optimization. I intentionally do not plan to move these kernels into a broad common operator layer.
In diffusion models, preserving model quality while continuously supporting new checkpoints, request policies, and execution features requires model-specific ownership. A “generic” operator still needs dtype, shape, device, stride, execution-mode, and platform checks. Even after those checks, it cannot automatically guarantee the model’s exact mathematical semantics, nor that every platform receives a positive performance gain. Those guarantees still have to be validated and maintained model by model.
Moving the implementation into a common layer would therefore not remove the model-specific work; it would mostly hide that work behind a wider API and create stronger coupling between models.
The approach in this PR keeps the boundary explicit: the model owns its operator contract, numerical policy, hardware dispatch, fallbacks, and validation. This makes the path easier to reason about, easier to dispatch safely, and easier for the community to validate or extend for another platform or quality level.
The trade-off is intentional duplication. Source-code volume alone is also becoming a progressively smaller maintenance constraint with AI-assisted tooling; But I’m open to other views.

[hsliuustc0106]

Fair position — with #6305 (op-boundary RFC) still settling where shared ops live, an explicitly model-owned kernel with its own contract, numerical policy, and validation is a defensible point on that spectrum. Could you capture this rationale in the PR description or a module docstring, so the placement reads as a decision rather than an oversight? If a future #6305 tier ever makes these shareable, migration can be a follow-up.

[mglyn]

Captured this in the model-local ops/README.md and the PR description in d34b75bd4. The README makes the numerical/execution/platform contract the reason for local ownership, documents the evidence expected for additional platforms, and leaves migration to a future #6305 shared tier as a follow-up. Thanks.

[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

@mglyn, 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.

[david6666666]

B300 validation (SM103)

I tested base c6df39e9f and this PR at f67759666 on one NVIDIA B300 SXM6 AC (compute capability 10.3 / SM103, driver 610.43.02, CUDA 13.0, PyTorch 2.13.0+cu130, Triton 3.7.1).

Important: the PR does not currently enable the optimization on B300

resolve_h3_vae_operators(cuda:0) returns None because the only registered operator set requires capability exactly 9.0. Consequently, install_h3_vae_optimizations(...) returns False, decoder weights stay FP32, and the new VAE forwards are not installed.

The B300 run of tests/diffusion/models/minimax_h3/test_minimax_h3_vae_ops.py therefore reported 14 passed, 5 skipped; the five skipped cases are the selected-operator exactness/integration cases.

I also ran the complete official-shape VAE decode with seed 42, latent [1, 24, 62, 48, 84], output [1, 3, 209, 768, 1344], FP16 CUDA autocast, and five measured decodes after warmup:

Current code path	P50	P100	Relative to base
Base c6df39e9f	8.7155 s	8.7331 s	1.00x
PR f67759666 as-is (fallback)	8.7199 s	8.7209 s	0.9995x (+0.05%, noise)

So the current PR has no realized B300 performance gain, as expected from the fallback.

SM103 feasibility check (forced dispatch, not current PR behavior)

To evaluate the implementation itself, I forced the existing exact operator set to be selected on SM103 without changing any kernel code. This is a feasibility experiment only, not a claim that the PR currently supports B300.

On the same loaded VAE and fixed latent, using two warmups followed by five CUDA-event measurements per side:

Complete VAE decode	Reference	Forced operator set	Speedup	Reduction
P50	8.7199 s	4.4145 s	1.975x	49.37%
P100	8.7209 s	4.4170 s	1.974x	49.35%

Correctness passed for the complete decoded tensor: torch.equal=True, max absolute error 0.

Isolated production-shape operator measurements were also positive and bitwise exact:

Operator / input	Reference P50	Candidate P50	Speedup	Numerical result
Q/K RMSNorm + RoPE, Q/K [2, 1797, 32, 64]	0.420792 ms	0.028877 ms	14.57x	Q/K torch.equal=True, max abs 0
SwiGLU, input [1, 1797, 16384]	0.073302 ms	0.016706 ms	4.39x	torch.equal=True, max abs 0
Scaled residual, input [2, 1797, 2048]	0.036792 ms	0.012565 ms	2.93x	torch.equal=True, max abs 0

I additionally ran 256 combinations of shape, seed, and value distribution for direct Q/K+RoPE and scaled-residual checks (512 bitwise comparisons total); all passed on SM103.

Conclusion: the kernels appear both exact and beneficial on this B300, but the current hardware table prevents users from receiving the gain. B300 should get an explicitly registered and validated SM103 operator set (rather than being reported as supported by the current SM90-only dispatch). Once that is added, the five selected-operator tests should run instead of skip, and I recommend repeating the complete VAE A/B on the actual updated PR head.

[hsliuustc0106]

Ran a controlled local GPU A/B on the exact merge base c6df39e9f7ad8a12fd61bc2c5007103c6a62076a versus head f6775966619c4a7a3c4e88793b7bce3a023fb476 using the same official-size FP32 latent ([1, 24, 62, 48, 84], seed 42), FP16 CUDA autocast, one excluded warmup, and two measured decodes per SHA.

• Mean warmed decode latency: 12.4463 s → 7.5284 s (-39.51%, 1.653× speedup). Individual runs were 12.4439/12.4487 s and 7.5194/7.5375 s.
• Peak allocated VRAM: 17,402.8 → 12,793.0 MiB (-26.49%); peak reserved VRAM: 19,690 → 18,126 MiB (-7.94%).
• Correctness: both complete outputs were finite FP32 tensors with shape [1, 3, 209, 768, 1344] and identical full-tensor SHA-256 0e361e9b3dbe1cffd23720c18c8845ebda8b08ad80445a785eaa45f433a3cd98.
• Activation was verified: base used the original forwards with 146 FP32 decoder linears; head installed the optimized block/attention/feed-forward paths with 144 block linears in FP16 and the two non-block linears left FP32.
• The added test module passed all 19 cases across the GPU-free contract/fallback partition and reserved-GPU exact-kernel partition.

The one-time process-to-ready time increased from 35.9121 s to 36.9788 s (+1.07 s), reported separately from warmed decode latency.

Environment: NVIDIA L20X (SM90), Python 3.12.13, PyTorch 2.13.0+cu129, CUDA 12.9, Triton 3.7.1, vLLM 0.28.0, MiniMax-H3 snapshot 42ed227ee7df40d41602854ae760620d6eb651fe. This validates the local SM90 path, not H200-specific performance; layerwise-offload and multi-rank execution were not exercised end to end.

The tested path passed without a blocking correctness or integration issue.

[hsliuustc0106]

Thanks for the rationale on the model-specific boundary — accepted, especially while #6305 is still settling where shared ops live. The one remaining ask from my side is to capture that intent in the PR description or a module docstring so the placement reads as a decision rather than an oversight.

[mglyn]

This is already captured in the second paragraph of the PR description and in the model-local ops/README.md since d34b75bd4. The README records the ownership rationale, operator contract, platform-validation expectations, and possible migration to a future #6305 shared tier. The original review thread is resolved; please let me know if you specifically prefer this rationale in a module docstring as well.