From 24741961ebdbdd0e87daf54737ceedb3d74dc1d1 Mon Sep 17 00:00:00 2001 From: Isotr0py Date: Mon, 3 Aug 2026 14:10:40 +0800 Subject: [PATCH] [Model] Add MiniMax H3 diffusion support (#5691) --- docs/models/supported_models.md | 1 + .../diffusion/attention_backends.md | 28 +- pyproject.toml | 6 + recipes/MiniMaxAI/MiniMax-H3.md | 329 +++++ recipes/README.md | 1 + tests/config/test_omni_config.py | 1 + tests/diffusion/attention/test_flash_attn.py | 43 +- tests/diffusion/attention/test_ring_fa4.py | 74 + tests/diffusion/models/minimax_h3/__init__.py | 2 + .../minimax_h3/test_minimax_h3_contract.py | 375 +++++ .../models/minimax_h3/test_minimax_h3_e2e.py | 69 + .../minimax_h3/test_minimax_h3_packing.py | 146 ++ .../minimax_h3/test_minimax_h3_parallel.py | 94 ++ .../openai_api/test_video_server.py | 47 + tests/model_tests/diffusion/test_alignment.py | 1 + vllm_omni/config/omni_config.py | 2 + vllm_omni/config/stage_config.py | 1 + .../diffusion/attention/backends/abstract.py | 4 + .../attention/backends/flash_attn.py | 57 + .../attention/backends/ring/ring_globals.py | 13 + .../attention/backends/ring/ring_kernels.py | 36 + .../attention/backends/ring/ring_selector.py | 5 + .../diffusion/attention/backends/utils/fa.py | 45 +- .../diffusion/attention/parallel/ring.py | 11 +- vllm_omni/diffusion/data.py | 10 + vllm_omni/diffusion/model_metadata.py | 4 + .../diffusion/models/minimax_h3/__init__.py | 12 + .../models/minimax_h3/condition_noise.py | 167 +++ .../models/minimax_h3/denoise_loop.py | 249 ++++ .../diffusion/models/minimax_h3/encoder.py | 1214 +++++++++++++++++ .../minimax_h3/minimax_h3_transformer.py | 1112 +++++++++++++++ .../models/minimax_h3/packed_sequence.py | 572 ++++++++ .../models/minimax_h3/packed_tokens.py | 114 ++ .../models/minimax_h3/pipeline_minimax_h3.py | 1197 ++++++++++++++++ .../models/minimax_h3/presentation.py | 297 ++++ .../models/minimax_h3/reference_video.py | 297 ++++ .../scheduling_minimax_h3_euler_ancestral.py | 179 +++ .../models/minimax_h3/time_request.py | 110 ++ vllm_omni/diffusion/models/minimax_h3/vae.py | 360 +++++ vllm_omni/diffusion/registry.py | 6 + vllm_omni/engine/arg_utils.py | 1 + vllm_omni/engine/async_omni_engine.py | 2 + vllm_omni/entrypoints/cli/serve.py | 9 + vllm_omni/entrypoints/openai/api_server.py | 81 +- vllm_omni/entrypoints/openai/serving_video.py | 3 +- 45 files changed, 7356 insertions(+), 31 deletions(-) create mode 100644 recipes/MiniMaxAI/MiniMax-H3.md create mode 100644 tests/diffusion/attention/test_ring_fa4.py create mode 100644 tests/diffusion/models/minimax_h3/__init__.py create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_e2e.py create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py create mode 100644 tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/__init__.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/condition_noise.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/denoise_loop.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/encoder.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/packed_sequence.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/packed_tokens.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/presentation.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/reference_video.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/scheduling_minimax_h3_euler_ancestral.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/time_request.py create mode 100644 vllm_omni/diffusion/models/minimax_h3/vae.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 021df65ce..99a16e87b 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -44,6 +44,7 @@ th { | `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | | `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | | `LingBotVideoPipeline` | LingBot-Video dense and MoE T2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | | +| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, FL2VA, and Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | | | | | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md index 63cba58e5..92a13aa74 100644 --- a/docs/user_guide/diffusion/attention_backends.md +++ b/docs/user_guide/diffusion/attention_backends.md @@ -14,8 +14,8 @@ The full set of backends and their platform defaults is in the **Backend Options | Value | Notes | |---|---| -| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention and **SAGE** quantized attention — see below. Requires `flashinfer`. | -| `FLASH_ATTN` | Wraps FlashAttention 2. Default on Hopper / Ada / Ampere when `flash-attn` is installed. | +| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. | +| `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. | | `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). | | `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. | | `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. | @@ -222,6 +222,30 @@ Same prompt and seed across runs. `Total generation time` from `text_to_video.py Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise. +## Known Limitations + +### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile + +LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121). + +**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0: + +```bash +DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \ + --model Lightricks/LTX-2 ... +``` + +### FlashAttention-4 on Blackwell + +Install the optional CUDA 13 extra to use the CuTe-based FA4 path: + +```bash +pip install 'vllm-omni[fa4]' +``` + +`FLASH_ATTN` prefers `flash_attn.cute` on Blackwell and falls back to FA3/FA2 +when it is unavailable. Version `4.0.0b18` is required; earlier beta +wheels had known JIT failures on Blackwell. ## Choosing a Backend Manually diff --git a/pyproject.toml b/pyproject.toml index dd3080d84..2559bcd5b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -96,6 +96,12 @@ quack = [ "quack-kernels>=0.3.11", ] +# Optional Blackwell-native FlashAttention-4 backend. Keep this opt-in because +# the CuTe/CUTLASS dependency is CUDA 13 and Blackwell specific. +fa4 = [ + "flash-attn-4[cu13]==4.0.0b18", +] + docs = [ "mkdocs>=1.5.0", "mkdocs-api-autonav", diff --git a/recipes/MiniMaxAI/MiniMax-H3.md b/recipes/MiniMaxAI/MiniMax-H3.md new file mode 100644 index 000000000..e94d2fb34 --- /dev/null +++ b/recipes/MiniMaxAI/MiniMax-H3.md @@ -0,0 +1,329 @@ +# MiniMax H3 + +> Joint video and audio generation with text, first-frame, image/audio, or +> multi-video conditions + +## Summary + +- Vendor: MiniMaxAI +- Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3) +- Tasks: T2VA, FL2VA, and Ref2VA +- Mode: OpenAI-compatible `/v1/videos` HTTP serving +- Maintainer: Community + +MiniMax H3 is a CFG-distilled joint video/audio diffusion transformer. Its +checkpoint has two independently served partitions: + +- `FL2VA`: text-to-video+audio (`t2va`) and first-frame-to-video+audio + (`fl2va`) +- `Ref2VA`: image+audio or one-or-more-video reference-to-video+audio + (`ref2va`) + +The generated MP4 contains H.264 video and synchronized stereo audio. + +## Prerequisites + +The checkpoint requires Hugging Face access approval. Download it and point +`MODEL_ROOT` at the directory containing the `FL2VA` and `Ref2VA` +subdirectories: + +```bash +hf auth login +export MODEL_ROOT=/path/to/MiniMax-H3 +hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}" +``` + +Install vLLM-Omni from the checkout containing MiniMax H3 support. On +Blackwell, install the optional FlashAttention-4 dependency: + +```bash +uv venv +source .venv/bin/activate +uv pip install -e '.[fa4]' +``` + +`ffmpeg` and `ffprobe` must be available on `PATH`. They are used for +reference-video preparation and MP4 output. + +## Start a server + +One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA +and FL2VA requests, or to `Ref2VA` for either Ref2VA request. + +### Single GPU: accuracy and memory first + +The single-GPU configuration uses model-level CPU offload. +This matches the accuracy-qualified reference path and prevents the Qwen3-VL +encoder and DiT from being resident on the GPU at the same time. + +```bash +export MODEL="${MODEL_ROOT}/FL2VA" +export PORT=8091 + +CUDA_VISIBLE_DEVICES=0 \ +VLLM_WORKER_MULTIPROC_METHOD=spawn \ +VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \ +vllm serve "${MODEL}" \ + --omni \ + --host 0.0.0.0 \ + --port "${PORT}" \ + --trust-remote-code \ + --num-gpus 1 \ + --enable-cpu-offload \ + --diffusion-attention-backend FLASH_ATTN +``` + +Use a GPU with enough memory for the active H3 component and enough system RAM +for the offloaded components. CPU offload reduces GPU memory pressure but adds +PCIe/NVLink transfer latency. + +### Four GPUs: measured best-practice throughput + +The best validated four-GPU configuration on four NVIDIA B300 GPUs is: + +- no CPU or layerwise offload; +- Ulysses sequence parallelism degree 4; +- native tiled VAE patch parallelism degree 4; +- regional `torch.compile` for the repeated DiT blocks; +- FlashAttention, with Ring and TP left at 1. + +```bash +export MODEL="${MODEL_ROOT}/FL2VA" +export PORT=8091 + +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +FLASHINFER_DISABLE_VERSION_CHECK=1 \ +VLLM_WORKER_MULTIPROC_METHOD=spawn \ +VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \ +vllm serve "${MODEL}" \ + --omni \ + --host 0.0.0.0 \ + --port "${PORT}" \ + --trust-remote-code \ + --num-gpus 4 \ + --usp 4 \ + --ring 1 \ + --vae-patch-parallel-size 4 \ + --vae-parallel-mode tile \ + --vae-use-tiling \ + --diffusion-attention-backend FLASH_ATTN +``` + +Do not add `--enforce-eager` to this performance configuration. The first +request includes regional compilation; warm the server once before measuring +steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain +1. The H3 VAE supports its native `tile` mode, not +`spatial_shard_height` or `spatial_shard_width`. + +### Text encoder tensor parallelism + +The Qwen3-VL text encoder (~51.5 GB in BF16 for the retained 50 layers) is by +default fully resident on the DiT main rank. On multi-GPU no-offload runs that +rank becomes the peak-memory hotspot. Add `--text-encoder-tp-size N` to shard +the encoder across the first `N` DiT ranks (the encoder is implemented with +vLLM-style tensor-parallel layers and runs with distributed collectives over +its own encoder process group): + +```bash +export MODEL="${MODEL_ROOT}/FL2VA" +export PORT=8091 + +CUDA_VISIBLE_DEVICES=0,1,2,3 \ +FLASHINFER_DISABLE_VERSION_CHECK=1 \ +VLLM_WORKER_MULTIPROC_METHOD=spawn \ +VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \ +vllm serve "${MODEL}" \ + --omni \ + --host 0.0.0.0 \ + --port "${PORT}" \ + --trust-remote-code \ + --num-gpus 4 \ + --usp 4 \ + --ring 1 \ + --text-encoder-tp-size 4 \ + --vae-patch-parallel-size 4 \ + --vae-parallel-mode tile \ + --vae-use-tiling \ + --diffusion-attention-backend FLASH_ATTN +``` + +`N` must divide the Qwen3-VL head counts (64 attention heads / 8 KV heads), so +valid values on a 4-GPU server are 1, 2, and 4 (1, 2, 4, 8 on 8 GPUs). The +encoder TP rank set is the first `N` DiT ranks; on 4 GPUs +`--text-encoder-tp-size 4` shards the encoder 4-way, dropping the DiT main +rank's no-offload peak by roughly `(N-1)/N` of the ~51.5 GB encoder while the +other ranks each gain `~51.5/N` GB. The encoder output remains identical to +the reference path within bf16 rounding: every encoder rank all-reduces the +row-parallel projections, so the full `[seq, 5120]` layer-50 hidden state is +replicated on every rank. + +To serve Ref2VA, stop the FL2VA server and restart the same one- or four-GPU +command with: + +```bash +export MODEL="${MODEL_ROOT}/Ref2VA" +``` + +## HTTP API examples + +The following requests use the synchronous endpoint so the returned body can +be saved directly as an MP4. The asynchronous `POST /v1/videos` endpoint can +also be used when job polling is preferred. + +All four tasks use 24 FPS, 50 sigma points, seed values from the validated +workloads, and the checkpoint-reference video/audio flow shifts of 12 and 3. +Decimal durations are passed through `extra_params`. + +Set the endpoint once: + +```bash +export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync" +``` + +### 1. T2VA: text to video and audio + +Run this request against the `FL2VA` partition: + +```bash +curl -sS -X POST "${API_URL}" \ + -F 'prompt=In a snowy blue-purple forest, Ori carefully walks past a sleeping giant; footsteps crunch in the snow while the creature breathes and softly snorts.' \ + -F 'width=1344' \ + -F 'height=768' \ + -F 'fps=24' \ + -F 'num_inference_steps=50' \ + -F 'flow_shift=12' \ + -F 'seed=1101' \ + -F 'extra_params={"task":"t2va","duration":8.7,"audio_flow_shift":3.0}' \ + -o t2va.mp4 +``` + +### 2. FL2VA: first frame to video and audio + +Run this request against the `FL2VA` partition. When width and height are +omitted, H3 preserves the first-frame aspect ratio and uses a 768-pixel short +edge. + +```bash +export FIRST_FRAME=/path/to/fl2va_first_frame.png + +curl -sS -X POST "${API_URL}" \ + -F 'prompt=A man stands beside a yellow car at night. The car drives away; he follows it with his eyes and begins singing sadly, with synchronized voice and city ambience.' \ + -F 'fps=24' \ + -F 'num_inference_steps=50' \ + -F 'flow_shift=12' \ + -F 'seed=2101' \ + -F 'extra_params={"task":"fl2va","duration":8.7,"audio_flow_shift":3.0}' \ + -F "input_reference=@${FIRST_FRAME};type=image/png" \ + -o fl2va.mp4 +``` + +### 3. Ref2VA: image and audio reference + +Run this request against the `Ref2VA` partition. `audio_reference` accepts an +HTTP(S) URL or a `data:` URL. In one terminal, expose the local reference +assets to the serving host: + +```bash +python -m http.server 8092 \ + --bind 127.0.0.1 \ + --directory /path/to/reference_assets +``` + +Then submit the image and audio request from another terminal: + +```bash +export REF_IMAGE=/path/to/reference_assets/ref2va_image.png +export AUDIO_URL=http://127.0.0.1:8092/ref2va_audio.mp3 + +curl -sS -X POST "${API_URL}" \ + -F 'prompt=A white cat with black mustache and eyebrow markings sits on a beige couch, lip-syncing precisely to the complete reference audio before shifting from confusion to deadpan speechlessness.' \ + -F 'width=1344' \ + -F 'height=768' \ + -F 'fps=24' \ + -F 'num_inference_steps=50' \ + -F 'flow_shift=12' \ + -F 'seed=3101' \ + -F 'extra_params={"task":"ref2va","duration":15.0,"audio_flow_shift":3.0}' \ + -F "input_reference=@${REF_IMAGE};type=image/png" \ + -F "audio_reference={\"audio_url\":\"${AUDIO_URL}\"}" \ + -o ref2va_image_audio.mp4 +``` + +The requested duration should cover the complete audio. If `duration` is +shorter, the reference soundtrack is truncated to the generated clip. + +### 4. Ref2VA: two video references + +Run this request against the `Ref2VA` partition. Repeat the +`input_references` multipart field once per source video. H3 consumes the +videos in form order and preserves their original soundtracks during +conditioning. + +```bash +export SUBJECT_VIDEO=/path/to/green_screen_subject.mp4 +export BACKGROUND_VIDEO=/path/to/fairytale_background.mov + +curl -sS -X POST "${API_URL}" \ + -F 'prompt=Remove the green screen background of Video 1 and replace it with the fairytale environment from Video 2. Match the background motion to the character actions and relight the character to fit the scene.' \ + -F 'width=1344' \ + -F 'height=768' \ + -F 'fps=24' \ + -F 'num_inference_steps=50' \ + -F 'flow_shift=12' \ + -F 'seed=3101' \ + -F 'extra_params={"task":"ref2va","duration":15.0,"audio_flow_shift":3.0}' \ + -F "input_references=@${SUBJECT_VIDEO};type=video/mp4" \ + -F "input_references=@${BACKGROUND_VIDEO};type=video/quicktime" \ + -o ref2va_video_video.mp4 +``` + +The server stores uploaded multi-video references only for the lifetime of the +request and deletes the temporary files after generation. Video Ref2VA uses +the source-video soundtracks and does not accept a separate +`audio_reference`. + +## Key parameters + +| Parameter | Recommended value | Notes | +|-----------|-------------------|-------| +| `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; must match the served partition | +| `duration` | Workload-specific | Decimal seconds in `extra_params`; converted to H3-compatible frame count | +| `fps` | `24` | H3 output FPS is fixed | +| `num_inference_steps` | `50` | Matches the reference accuracy workloads | +| `flow_shift` | `12` | Video sigma shift | +| `audio_flow_shift` | `3` | Audio sigma shift, passed in `extra_params` | +| `seed` | Task-specific | Use a fixed value for reproducibility | +| `width`, `height` | Multiples of 32 | Aspect ratio must be between 1:4 and 4:1 | + +## Validated four-GPU evidence + +The four-GPU recommendation was measured on four NVIDIA B300 GPUs with one +excluded warmup followed by three requests. + +| Workload | Configuration | Observed result | +|----------|---------------|-----------------| +| FL2VA, 209 frames, 1248x768 | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency | +| Two-video Ref2VA, 362 frames, 1344x768 | no offload, U4, VPP4 tile, regional compile | 784.394 s accounted model-stage mean | + +These measurements describe the validated shapes rather than a general +throughput guarantee. Multi-video Ref2VA is much slower because the two +reference videos expand both the Qwen3-VL vision sequence and the packed DiT +attention sequence. + +## Known limitations + +- Each server process loads only one checkpoint partition. +- H3 currently executes one generation request per diffusion batch. +- The first regional-compile request is a warmup and should not be included in + steady-state performance measurements. +- Image+audio Ref2VA accepts exactly one image and one audio reference. +- Video Ref2VA accepts one or more video files, but not an additional standalone + audio reference. +- VAE patch parallelism requires size 1 or the full DiT group size and supports + the H3 native `tile` mode only. + +## Additional resources + +- [Supported models](../../docs/models/supported_models.md) +- [Video API](../../docs/serving/videos_api.md) +- [Diffusion parallelism](../../docs/user_guide/diffusion/parallelism/overview.md) diff --git a/recipes/README.md b/recipes/README.md index 1d89f0ef1..f5e279a48 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -45,6 +45,7 @@ recipes/ | [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB | | [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) | | [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU | +| [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 | | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB | | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU | | [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom | diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index d477e05c1..9650471d6 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -440,6 +440,7 @@ def test_sub_config_fields_match_rfc_scopes(): "ulysses_mode", "cfg_parallel_size", "vae_patch_parallel_size", + "text_encoder_tp_size", "vae_parallel_mode", "use_hsdp", "mask_sp_padding", diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py index fbe037274..53673dd58 100644 --- a/tests/diffusion/attention/test_flash_attn.py +++ b/tests/diffusion/attention/test_flash_attn.py @@ -20,7 +20,6 @@ from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] - is_gpu = current_omni_platform.is_cuda_alike() or current_omni_platform.is_xpu() HAS_FLASH_ATTN = fa.HAS_FLASH_ATTN flash_attn_func = fa.flash_attn_func # noqa: N813 @@ -330,6 +329,48 @@ def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch): assert calls[0]["softmax_scale"] == 0.5 +def test_packed_varlen_metadata_bypasses_mask_unpadding(monkeypatch): + calls = [] + + def fake_varlen_func(q, k, v, **kwargs): + calls.append((q.shape, k.shape, v.shape, kwargs)) + return q + + monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True) + monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func) + + impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=False) + query = torch.randn(1, 8, 2, 4) + cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32) + metadata = AttentionMetadata( + attn_mask=torch.tensor([[True] * 6 + [False] * 2]), + extra={ + "cu_seqlens_q": cu_seqlens, + "cu_seqlens_k": cu_seqlens, + "max_seqlen_q": 6, + "max_seqlen_k": 6, + }, + ) + + output = impl.forward_cuda(query, query, query, metadata) + + assert torch.equal(output, query) + assert len(calls) == 1 + assert calls[0][0] == torch.Size([8, 2, 4]) + assert calls[0][3]["cu_seqlens_q"] is cu_seqlens + assert calls[0][3]["max_seqlen_q"] == 6 + + +def test_packed_varlen_metadata_must_be_complete(monkeypatch): + monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True) + impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=False) + query = torch.randn(1, 8, 2, 4) + metadata = AttentionMetadata(extra={"cu_seqlens_q": torch.tensor([0, 8], dtype=torch.int32)}) + + with pytest.raises(ValueError, match="Incomplete packed FlashAttention metadata"): + impl.forward_cuda(query, query, query, metadata) + + if __name__ == "__main__": print("Running FlashAttention Padding Tests...") print("=" * 60) diff --git a/tests/diffusion/attention/test_ring_fa4.py b/tests/diffusion/attention/test_ring_fa4.py new file mode 100644 index 000000000..b24ce2f20 --- /dev/null +++ b/tests/diffusion/attention/test_ring_fa4.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.attention.backends.ring import ring_kernels +from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType +from vllm_omni.diffusion.attention.parallel import ring as ring_parallel + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_fa4_wrapper_returns_lse_and_normalizes_unbounded_window(monkeypatch): + expected_out = object() + expected_lse = object() + captured = {} + + def fake_fa4(q, k, v, **kwargs): + captured.update(q=q, k=k, v=v, **kwargs) + return expected_out, expected_lse + + monkeypatch.setattr(ring_kernels, "HAS_FA4", True) + monkeypatch.setattr(ring_kernels, "fa4_attn_func", fake_fa4) + + out, lse = ring_kernels.flash_attn4_func_forward( + "q", + "k", + "v", + softmax_scale=0.125, + causal=True, + window_size=(-1, -1), + ) + + assert (out, lse) == (expected_out, expected_lse) + assert captured == { + "q": "q", + "k": "k", + "v": "v", + "softmax_scale": 0.125, + "causal": True, + "window_size": (None, None), + "softcap": 0.0, + "return_lse": True, + } + + +def test_ring_prefers_fa4_on_blackwell(monkeypatch): + captured = {} + + def fake_ring_flash_attn_func(*args, **kwargs): + captured.update(kwargs) + return "output" + + monkeypatch.setattr(ring_parallel, "HAS_FA4", True) + monkeypatch.setattr(ring_parallel, "HAS_FA3", True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda _device: (10, 3)) + + from vllm_omni.diffusion.attention.backends import ring_flash_attn + + monkeypatch.setattr(ring_flash_attn, "ring_flash_attn_func", fake_ring_flash_attn_func) + + query = SimpleNamespace( + dtype=torch.bfloat16, + shape=(1, 16, 4, 128), + is_cuda=True, + device=torch.device("cuda:0"), + ) + strategy = ring_parallel.RingParallelAttention(SimpleNamespace(ring_group=object())) + + assert strategy.run_attention(query, "key", "value", None) == "output" + assert captured["attn_type"] is AttnType.FA4 diff --git a/tests/diffusion/models/minimax_h3/__init__.py b/tests/diffusion/models/minimax_h3/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/tests/diffusion/models/minimax_h3/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py new file mode 100644 index 000000000..5afd8c027 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py @@ -0,0 +1,375 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from multiprocessing.reduction import ForkingPickler +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_pipeline_import_registry_and_component_discovery(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.registry import ( + _DIFFUSION_MODELS, + _DIFFUSION_POST_PROCESS_FUNCS, + ) + + assert _DIFFUSION_MODELS["MiniMaxH3Pipeline"] == ( + "minimax_h3", + "pipeline_minimax_h3", + "MiniMaxH3Pipeline", + ) + assert _DIFFUSION_POST_PROCESS_FUNCS["MiniMaxH3Pipeline"] == "get_minimax_h3_post_process_func" + assert MiniMaxH3Pipeline._dit_modules == ["transformer"] + assert MiniMaxH3Pipeline._encoder_modules == ["text_encoder"] + assert MiniMaxH3Pipeline._vae_modules == ["video_vae", "audio_vae"] + + +def test_joint_postprocess_is_multiprocessing_picklable(): + from vllm_omni.diffusion.models.minimax_h3 import ( + get_minimax_h3_post_process_func, + ) + + postprocess = get_minimax_h3_post_process_func(SimpleNamespace()) + postprocess = ForkingPickler.loads(ForkingPickler.dumps(postprocess)) + video = torch.linspace(0, 1, 2 * 3 * 2 * 4 * 5).reshape(2, 3, 2, 4, 5) + audio = torch.arange(12, dtype=torch.float32).reshape(1, 2, 6) + + result = postprocess((video, audio), output_type="np") + + assert isinstance(result["video"], list) + assert result["video"][0].shape == (2, 4, 5, 3) + np.testing.assert_array_equal(result["audio"], audio.numpy()) + assert result["audio_sample_rate"] == 32000 + assert result["fps"] == 24 + + +def test_cfg_parallel_is_rejected_for_distilled_checkpoint(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + od_config = SimpleNamespace( + parallel_config=SimpleNamespace(cfg_parallel_size=2), + ) + with pytest.raises(ValueError, match="CFG-distilled"): + MiniMaxH3Pipeline(od_config=od_config) + + +def test_shape_contract_matches_three_reference_tasks(): + from vllm_omni.diffusion.models.minimax_h3.time_request import ( + MINIMAX_H3_SHAPE_PLANNER, + minimax_h3_align_frame_count, + ) + + assert minimax_h3_align_frame_count(round(8.7 * 24)) == 209 + assert MINIMAX_H3_SHAPE_PLANNER.video_latent_t(209) == 62 + assert MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(8.7) == 348 + assert minimax_h3_align_frame_count(round(5.0 * 24)) == 124 + assert MINIMAX_H3_SHAPE_PLANNER.video_latent_t(124) == 37 + assert MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(5.0) == 200 + assert MINIMAX_H3_SHAPE_PLANNER.frame_count_from_video_latent_t(62) == 209 + + +def test_shifted_sigma_schedule_matches_reference_values(): + from vllm_omni.diffusion.models.minimax_h3.time_request import ( + minimax_h3_time_shift_sigmas, + ) + + sigmas = minimax_h3_time_shift_sigmas(num_steps=5, shift_scale=12.0) + + assert sigmas == pytest.approx( + [1.0, 0.9729729891, 0.9230769277, 0.8000000119, 0.0], + abs=1e-7, + ) + + +def test_reference_image_resize_contract(): + from PIL import Image + + from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import ( + _reference_image_shape, + ) + + assert _reference_image_shape(Image.new("RGB", (1080, 1440))) == ( + 2048, + 2720, + ) + with pytest.raises(ValueError, match="aspect ratio"): + _reference_image_shape(Image.new("RGB", (100, 501))) + + +def test_encoder_forward_uses_hook_compatible_encode_entrypoint(): + from vllm_omni.diffusion.models.minimax_h3.encoder import ( + MiniMaxH3Qwen3VLEncoder, + ) + + encoder = object.__new__(MiniMaxH3Qwen3VLEncoder) + torch.nn.Module.__init__(encoder) + expected = torch.ones(2, 3) + encoder.encode_ids = Mock(return_value=expected) + input_ids = torch.tensor([1, 2]) + pixel_values = torch.ones(1, 4) + image_grid_thw = torch.tensor([[1, 1, 1]]) + + actual = encoder( + input_ids, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + ) + + assert actual is expected + encoder.encode_ids.assert_called_once_with( + input_ids, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + ) + + +def test_encoder_forward_forwards_video_inputs(): + from vllm_omni.diffusion.models.minimax_h3.encoder import ( + MiniMaxH3Qwen3VLEncoder, + ) + + encoder = object.__new__(MiniMaxH3Qwen3VLEncoder) + torch.nn.Module.__init__(encoder) + expected = torch.ones(2, 3) + encoder.encode_ids = Mock(return_value=expected) + input_ids = torch.tensor([1, 2]) + pixel_values_videos = torch.ones(1, 4) + video_grid_thw = torch.tensor([[2, 1, 1]]) + + actual = encoder( + input_ids, + pixel_values_videos=pixel_values_videos, + video_grid_thw=video_grid_thw, + ) + + assert actual is expected + encoder.encode_ids.assert_called_once_with( + input_ids, + pixel_values=None, + image_grid_thw=None, + pixel_values_videos=pixel_values_videos, + video_grid_thw=video_grid_thw, + ) + + +def test_reference_video_shape_uses_h3_adapt_shape_policy(): + from vllm_omni.diffusion.models.minimax_h3.reference_video import ( + _reference_video_shape, + ) + + assert _reference_video_shape(1280, 720) == (1344, 768) + assert _reference_video_shape(3844, 2160) == (1344, 768) + + +def test_text_encoder_stub_constructs_without_group_or_weights(): + from vllm_omni.diffusion.models.minimax_h3.encoder import ( + MiniMaxH3Qwen3VLEncoder, + ) + + encoder = MiniMaxH3Qwen3VLEncoder( + "/nonexistent/text_encoder", + device=torch.device("cpu"), + load_model=False, + encoder_group=None, + ) + assert not encoder.is_loaded + assert encoder.tp_size == 1 + # The stub has no parameters, so it never contributes to the runner's + # strict missing-parameter check on non-encoder ranks. + assert list(encoder.named_parameters()) == [] + + +def test_no_offload_keeps_text_encoder_resident(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + ) + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + input_ids = torch.tensor([1, 2]) + + actual = pipeline._encode_text_hidden(input_ids, {}) + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_not_called() + + +def test_model_offload_uses_hooked_text_encoder_call(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=True, + enable_layerwise_offload=False, + ) + expected = torch.ones(2, 3) + pipeline.text_encoder = Mock(return_value=expected) + input_ids = torch.tensor([1, 2]) + vision_kwargs = {"pixel_values": torch.ones(1, 4)} + + actual = pipeline._encode_text_hidden(input_ids, vision_kwargs) + + assert actual is expected + pipeline.text_encoder.assert_called_once_with(input_ids, **vision_kwargs) + pipeline.text_encoder.load_to_device.assert_not_called() + + +def test_layerwise_offload_releases_text_encoder(): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=True, + ) + pipeline.text_encoder = Mock() + expected = torch.ones(2, 3) + pipeline.text_encoder.encode_ids.return_value = expected + + actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {}) + + assert actual is expected + pipeline.text_encoder.load_to_device.assert_called_once_with() + pipeline.text_encoder.offload_to_cpu.assert_called_once_with() + + +def test_video_vae_keeps_reference_fp32_weights(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import vae as vae_module + + class FakeRemote(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = torch.nn.Linear(1, 1).half() + + monkeypatch.setattr( + vae_module, + "_load_component_config", + lambda _path: { + "latent_channels": 1, + "latents_mean": [0.0], + "latents_std": [1.0], + }, + ) + monkeypatch.setattr( + vae_module, + "_load_remote_component", + lambda _path, _config: FakeRemote(), + ) + + video_vae = vae_module.MiniMaxH3VideoVAE( + "unused", + device=torch.device("cpu"), + ) + + assert next(video_vae.parameters()).dtype == torch.float32 + + +def test_video_vae_encode_uses_configured_parallel_tiling(): + from vllm_omni.diffusion.models.minimax_h3.vae import ( + MiniMaxH3VideoVAE, + ) + + class FakeModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(1)) + self.parallel_tiling = True + self.encode_calls = [] + + def encode_videos(self, frames, *, use_fp16_latent): + assert self.parallel_tiling + self.encode_calls.append((frames, use_fp16_latent)) + return [torch.ones(1, 1, 2, 2, 2)] + + video_vae = object.__new__(MiniMaxH3VideoVAE) + torch.nn.Module.__init__(video_vae) + video_vae.model = FakeModel() + video_vae.config_dict = { + "latent_channels": 1, + "latents_mean": [0.0], + "latents_std": [1.0], + } + + rows, shape = video_vae.encode_video("frames") + + assert video_vae.model.parallel_tiling + assert video_vae.model.encode_calls == [("frames", True)] + assert rows.shape == (2, 4) + assert shape == (2, 2, 2) + + +def test_distributed_video_vae_encodes_references_sequentially(monkeypatch): + from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline + from vllm_omni.diffusion.models.minimax_h3 import ( + pipeline_minimax_h3 as pipeline_module, + ) + + prepared = [ + {"prepared_path": "video-1.mp4"}, + {"prepared_path": "video-2.mp4"}, + ] + + class FakeVideoVAE: + def __init__(self): + self.calls = [] + + def is_distributed_enabled(self): + return True + + def encode_video(self, frames): + self.calls.append(frames) + index = len(self.calls) + return torch.full((1, 2), index, dtype=torch.float32), (index, 2, 3) + + pipeline = object.__new__(MiniMaxH3Pipeline) + torch.nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.video_vae = FakeVideoVAE() + + monkeypatch.setattr( + pipeline_module, + "_dit_rank_world", + lambda: ("dit-group", 1, 4), + ) + + def fake_broadcast_object_list(values, *, src, group, device): + assert values == [None] + assert (src, group, device) == (0, "dit-group", torch.device("cpu")) + values[0] = prepared + + monkeypatch.setattr( + pipeline_module.dist, + "broadcast_object_list", + fake_broadcast_object_list, + ) + monkeypatch.setattr( + pipeline_module, + "load_video_frames", + lambda path: f"frames:{path}", + ) + + rows, shapes = pipeline._encode_video_conditions(None, count=2) + + assert pipeline.video_vae.calls == [ + "frames:video-1.mp4", + "frames:video-2.mp4", + ] + torch.testing.assert_close( + rows, + torch.tensor([[1.0, 1.0], [2.0, 2.0]]), + ) + assert shapes == [(1, 2, 3), (2, 2, 3)] diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_e2e.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_e2e.py new file mode 100644 index 000000000..59b720df4 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_e2e.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os + +import numpy as np +import pytest + +pytestmark = [ + pytest.mark.local_model, + pytest.mark.cuda, + pytest.mark.diffusion, + pytest.mark.parallel, +] + +MODEL_ENV = "VLLM_TEST_MINIMAX_H3_FL2VA_MODEL" + + +@pytest.mark.skipif( + not os.environ.get(MODEL_ENV), + reason=f"set {MODEL_ENV} to an authorized FL2VA checkpoint path", +) +def test_minimax_h3_t2va_ulysses8_smoke(): + import torch + + from vllm_omni.diffusion.data import DiffusionParallelConfig + from vllm_omni.entrypoints.omni import Omni + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + if torch.accelerator.device_count() < 8: + pytest.skip("MiniMax H3 Ulysses smoke requires eight GPUs") + + engine = Omni( + model=os.environ[MODEL_ENV], + parallel_config=DiffusionParallelConfig(ulysses_degree=8), + trust_remote_code=True, + enable_cpu_offload=True, + enforce_eager=True, + ) + try: + outputs = engine.generate( + "A quiet cinematic night scene with matching ambient sound.", + OmniDiffusionSamplingParams( + height=256, + width=448, + num_frames=29, + fps=24, + num_inference_steps=2, + seed=42, + output_type="np", + extra_args={ + "task": "t2va", + "flow_shift": 12.0, + "audio_flow_shift": 3.0, + }, + ), + use_tqdm=False, + ) + finally: + engine.close() + + assert len(outputs) == 1 + frames = np.asarray(outputs[0].images[0]) + assert frames.shape == (39, 256, 448, 3) + multimodal = outputs[0].multimodal_output + assert multimodal is not None + assert np.asarray(multimodal["audio"]).shape[1] == 2 + assert multimodal["audio_sample_rate"] == 32000 + assert multimodal["fps"] == 24 diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py new file mode 100644 index 000000000..6e448ab52 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_packing.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_video_patchify_round_trip_preserves_values(): + from vllm_omni.diffusion.models.minimax_h3.packed_tokens import ( + minimax_h3_patchify_video_latent, + minimax_h3_unpatchify_video_tokens, + ) + + latent = torch.arange(2 * 3 * 2 * 4 * 6).reshape(2, 3, 2, 4, 6) + rows = minimax_h3_patchify_video_latent( + latent, + patch_size=(1, 2, 2), + ) + restored = minimax_h3_unpatchify_video_tokens( + rows, + latent_shape=(2, 2, 3, 3), + patch_size=(1, 2, 2), + ) + + torch.testing.assert_close(restored, latent) + + +def test_audio_pack_round_trip_preserves_channel_major_order(): + from vllm_omni.diffusion.models.minimax_h3.packed_tokens import ( + minimax_h3_pack_audio_latent, + minimax_h3_unpack_audio_tokens, + ) + + latent = torch.arange(2 * 4 * 5).reshape(2, 4, 5) + rows = minimax_h3_pack_audio_latent(latent) + restored = minimax_h3_unpack_audio_tokens( + rows, + audio_t=10, + audio_channel=2, + ) + + torch.testing.assert_close(restored, latent) + + +def test_t2va_and_fl2va_packing_keep_update_rows_separate(): + from vllm_omni.diffusion.models.minimax_h3.packed_sequence import ( + minimax_h3_packed_sequence, + ) + + common = dict( + text_len=4, + latent_t=2, + latent_h=4, + latent_w=6, + audio_t=3, + ) + t2va = minimax_h3_packed_sequence( + **common, + include_keyframe_cond=False, + ) + fl2va = minimax_h3_packed_sequence( + **common, + include_keyframe_cond=True, + keyframe_frame_indices=[0], + frame_count=5, + ) + + assert int(t2va["seq_len"]) == 64 + assert t2va["img_pos"].numel() == 12 + assert t2va["update_mask"].all() + assert fl2va["img_pos"].numel() == 18 + assert fl2va["update_mask"].sum().item() == 12 + assert (~fl2va["update_mask"]).sum().item() == 6 + assert t2va["audio_pos"].numel() == fl2va["audio_pos"].numel() == 6 + + +def test_ref2va_packing_tracks_video_and_audio_update_masks(): + from vllm_omni.diffusion.models.minimax_h3.packed_sequence import ( + minimax_h3_packed_sequence_ref2va_blocks, + ) + + packed = minimax_h3_packed_sequence_ref2va_blocks( + text_len=4, + latent_t=2, + latent_h=4, + latent_w=6, + audio_t=3, + ref_blocks=[ + {"kind": "image", "latent_h": 4, "latent_w": 4}, + {"kind": "audio", "ref_audio_t": 2}, + ], + ) + + assert int(packed["seq_len"]) == 64 + assert packed["img_pos"].numel() == 16 + assert packed["update_mask"].sum().item() == 12 + assert packed["audio_pos"].numel() == 10 + assert packed["audio_update_mask"].sum().item() == 6 + assert packed["cu_seqlens"].tolist() == [0, 30, 64] + + +def test_condition_noise_is_seeded_and_keeps_clean_anchor_at_timestep_one(): + from vllm_omni.diffusion.models.minimax_h3.condition_noise import ( + minimax_h3_audio_cond_noise_aug_rows, + minimax_h3_imgvid_cond_noise_aug_rows, + ) + + image_rows = torch.randn(4, 96) + audio_rows = torch.randn(6, 32) + + torch.testing.assert_close( + minimax_h3_imgvid_cond_noise_aug_rows( + image_rows, + condition_shapes=[(1, 4, 4)], + target_latent_t=2, + imgvid_cond_num_frames=1, + seed=42, + noise_aug=1.0, + ), + image_rows, + ) + torch.testing.assert_close( + minimax_h3_audio_cond_noise_aug_rows( + audio_rows, + condition_audio_t=[3], + seed=42, + noise_aug=1.0, + ), + audio_rows, + ) + + first = minimax_h3_audio_cond_noise_aug_rows( + audio_rows, + condition_audio_t=[3], + seed=42, + noise_aug=0.25, + ) + second = minimax_h3_audio_cond_noise_aug_rows( + audio_rows, + condition_audio_t=[3], + seed=42, + noise_aug=0.25, + ) + torch.testing.assert_close(first, second) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py new file mode 100644 index 000000000..d6f089d0a --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn as nn + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_grouped_qkv_checkpoint_reorder(): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + _reorder_grouped_qkv_to_qkv, + ) + + # Two groups with rows [q, k, v] become [q0, q1, k0, k1, v0, v1]. + grouped = torch.arange(6, dtype=torch.float32).reshape(6, 1) + reordered = _reorder_grouped_qkv_to_qkv( + grouped, + num_query_groups=2, + heads_per_group=1, + head_dim=1, + ) + + assert reordered[:, 0].tolist() == [0, 3, 1, 4, 2, 5] + + +def test_transformer_declares_cache_sp_layerwise_offload_and_hsdp(): + from cache_dit import ForwardPattern + + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3DiTModel, + ) + + assert MiniMaxH3DiTModel._repeated_blocks == ["MiniMaxH3DiTBlock"] + assert MiniMaxH3DiTModel._layerwise_offload_blocks_attrs == ["blocks"] + assert MiniMaxH3DiTModel._cache_dit_adapter_config.block_forward_patterns["blocks"] == ForwardPattern.Pattern_3 + assert not MiniMaxH3DiTModel._cache_dit_adapter_config.has_separate_cfg + assert set(MiniMaxH3DiTModel._sp_plan) == {"sp_prepare", "sp_gather"} + + model = object.__new__(MiniMaxH3DiTModel) + nn.Module.__init__(model) + model.blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)]) + model.token_refiner = nn.Module() + model.token_refiner.blocks = nn.ModuleList([nn.Linear(4, 4)]) + model.final_layer = nn.Linear(4, 4) + + matched = [ + name + for name, module in model.named_modules() + if any(condition(name, module) for condition in MiniMaxH3DiTModel._hsdp_shard_conditions) + ] + assert matched == ["blocks.0", "blocks.1"] + + +def test_packed_attention_is_a_regional_compile_boundary(): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3Attention, + ) + + assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False) + + +@pytest.mark.parametrize( + ("tp_size", "message"), + [ + (3, "num_attention_heads"), + (5, "num_attention_heads"), + ], +) +def test_tp_rejects_non_divisible_head_counts(tp_size, message): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3DiTArchConfig, + MiniMaxH3DiTModel, + ) + + model = object.__new__(MiniMaxH3DiTModel) + with pytest.raises(ValueError, match=message): + model._validate_tp_config( + arch=MiniMaxH3DiTArchConfig(), + tp_size=tp_size, + ) + + +def test_tp_accepts_checkpoint_supported_sizes(): + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3DiTArchConfig, + MiniMaxH3DiTModel, + ) + + model = object.__new__(MiniMaxH3DiTModel) + arch = MiniMaxH3DiTArchConfig() + for tp_size in (1, 2, 4, 7): + model._validate_tp_config(arch=arch, tp_size=tp_size) diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index ff69a2bf2..86757e56d 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -12,6 +12,7 @@ import os import threading import time from contextlib import asynccontextmanager +from pathlib import Path from types import SimpleNamespace import numpy as np @@ -62,11 +63,19 @@ class FakeAsyncOmni: self.stage_configs = [SimpleNamespace(stage_type="diffusion")] self.default_sampling_params_list = [OmniDiffusionSamplingParams()] self.captured_prompt = None + self.captured_reference_video_bytes = None self.captured_sampling_params_list = None async def generate(self, prompt, request_id, sampling_params_list): self.captured_prompt = prompt self.captured_sampling_params_list = sampling_params_list + reference_videos = prompt.get("multi_modal_data", {}).get("video") + if ( + isinstance(reference_videos, list) + and reference_videos + and all(isinstance(item, str) for item in reference_videos) + ): + self.captured_reference_video_bytes = [Path(item).read_bytes() for item in reference_videos] num_outputs = sampling_params_list[0].num_outputs_per_prompt videos = [object() for _ in range(num_outputs)] yield MockVideoResult(videos) @@ -501,6 +510,44 @@ def test_v2v_video_generation_with_video_reference_form(test_client, mocker: Moc assert input_video[0].size == (32, 24) +@pytest.mark.parametrize("endpoint", ["/v1/videos", "/v1/videos/sync"]) +def test_multi_video_generation_preserves_uploaded_files_until_generation( + endpoint, + test_client, + mocker: MockerFixture, +): + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", + ) + response = test_client.post( + endpoint, + data={ + "prompt": "Composite the subject from the first video into the second.", + "extra_params": json.dumps({"task": "ref2va", "duration": 15.0}), + }, + files=[ + ("input_references", ("subject.mp4", b"subject-video", "video/mp4")), + ("input_references", ("background.mov", b"background-video", "video/quicktime")), + ], + ) + + assert response.status_code == 200 + if endpoint.endswith("/sync"): + assert response.content == b"fake-video" + else: + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + + engine = test_client.app.state.openai_serving_video._engine_client + input_videos = engine.captured_prompt["multi_modal_data"]["video"] + assert engine.captured_reference_video_bytes == [b"subject-video", b"background-video"] + assert len(input_videos) == 2 + assert all(not Path(path).exists() for path in input_videos) + assert engine.captured_sampling_params_list[0].extra_args["task"] == "ref2va" + assert engine.captured_sampling_params_list[0].extra_args["duration"] == 15.0 + + def test_decode_video_bytes_can_keep_last_frames(): from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index e1795ddd0..e93abb74f 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -32,6 +32,7 @@ EXCLUDED_MODELS = [ "LTX2DistilledPipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", + "MiniMaxH3Pipeline", "StableAudioPipeline", "WanImageToVideoPipeline", "WanS2VPipeline", diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index c8fa51448..0207b0132 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -113,6 +113,7 @@ class _ParallelConfigEngineOverrides(TypedDict, total=False): ulysses_mode: str cfg_parallel_size: int vae_patch_parallel_size: int + text_encoder_tp_size: int use_hsdp: bool mask_sp_padding: bool hsdp_shard_size: int @@ -370,6 +371,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): ulysses_mode: str = "strict" cfg_parallel_size: int = Field(default=1, ge=1) vae_patch_parallel_size: int = Field(default=1, ge=1) + text_encoder_tp_size: int = Field(default=1, ge=1) vae_parallel_mode: str = "tile" use_hsdp: bool = False mask_sp_padding: bool = False diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 43944084c..649e5e4cf 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -377,6 +377,7 @@ class StageDeployConfig: cfg_parallel_size: int | None = None vae_patch_parallel_size: int | None = None vae_parallel_mode: str | None = None + text_encoder_tp_size: int | None = None use_hsdp: bool | None = None hsdp_shard_size: int | None = None hsdp_replicate_size: int | None = None diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index 43c6f525c..2e30877ca 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -79,6 +79,10 @@ class AttentionMetadata: # Well-known optional keys (convention, not required on all forwards): # "kv_cache_dtype": str | None — quantized KV dtype (e.g. "fp8"); backends # decide whether/how to apply. + # "cu_seqlens_q" / "cu_seqlens_k": int32 CUDA tensors describing packed + # variable-length query/key sequences for FlashAttention. + # "max_seqlen_q" / "max_seqlen_k": maximum sequence lengths paired with + # the packed cu_seqlens tensors. # Piecewise attention metadata (mixed causal/full masks). # full_attn_spans: per-sample [start, end) spans in global coordinates using full attention. diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index 82d61c25d..467a9ad16 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -141,6 +141,46 @@ class FlashAttentionImpl(AttentionImpl): out_unpad = self._unwrap_flash_output(out_unpad) return _pad_input(out_unpad, indices_q, query.size(0), query_length) + def _forward_varlen_packed( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + ) -> torch.Tensor: + """Run FlashAttention directly on an already packed sequence. + + Some diffusion transformers already maintain exact packed-document + boundaries. Reusing those boundaries avoids rebuilding boolean masks + and gathering/scattering Q/K/V in every attention layer. + """ + from vllm_omni.diffusion.attention.backends.utils.fa import ( + flash_attn_varlen_func, + ) + + if flash_attn_varlen_func is None: + raise ImportError("Packed variable-length attention requires flash_attn_varlen_func") + if query.shape[0] != 1 or key.shape[0] != 1 or value.shape[0] != 1: + raise ValueError("Packed variable-length attention currently requires batch size 1") + + out = flash_attn_varlen_func( + q=query.flatten(0, 1), + k=key.flatten(0, 1), + v=value.flatten(0, 1), + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + causal=self.causal, + softmax_scale=self.softmax_scale, + ) + out = self._unwrap_flash_output(out) + return out.reshape_as(query) + def _forward_varlen_dense( self, query: torch.Tensor, @@ -206,6 +246,7 @@ class FlashAttentionImpl(AttentionImpl): attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None + extra = attn_metadata.extra if attn_metadata is not None else {} # Try piecewise attention if full_attn_spans is not None: @@ -233,6 +274,22 @@ class FlashAttentionImpl(AttentionImpl): query_ranges=attn_metadata.query_ranges, ) + packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k") + present_packed_keys = [key for key in packed_keys if key in extra] + if present_packed_keys: + if len(present_packed_keys) != len(packed_keys): + missing = sorted(set(packed_keys) - set(present_packed_keys)) + raise ValueError(f"Incomplete packed FlashAttention metadata; missing {missing}") + return self._forward_varlen_packed( + query, + key, + value, + cu_seqlens_q=extra["cu_seqlens_q"], + cu_seqlens_k=extra["cu_seqlens_k"], + max_seqlen_q=extra["max_seqlen_q"], + max_seqlen_k=extra["max_seqlen_k"], + ) + if attention_mask is not None and torch.any(~attention_mask): return self._forward_varlen_masked( query, diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_globals.py b/vllm_omni/diffusion/attention/backends/ring/ring_globals.py index 4ac7d1456..25f9a41cc 100644 --- a/vllm_omni/diffusion/attention/backends/ring/ring_globals.py +++ b/vllm_omni/diffusion/attention/backends/ring/ring_globals.py @@ -23,6 +23,19 @@ HAS_FA3 = False fa3_fwd_func = None # Low-level forward function (_flash_attn_forward) fa3_attn_func = None # High-level attention function (flash_attn_func) +# FA4 detection. The CuTe API returns LSE directly, so it can participate in +# Ring Attention's numerically stable block-wise output accumulation. +HAS_FA4 = False +fa4_attn_func = None +try: + from flash_attn.cute import flash_attn_func as fa4_attn_func # noqa: F401 + + HAS_FA4 = True +except Exception: + # Optional CuTe/CUTLASS/Quack components can be importable but + # ABI-incompatible. Treat the whole optional backend as unavailable. + pass + # Try flash_attn_interface first (from flash-attention source build) try: from flash_attn_interface import _flash_attn_forward as fa3_fwd_func # noqa: F401 diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py b/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py index 51cf8bbc3..838726396 100644 --- a/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py +++ b/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py @@ -9,9 +9,11 @@ import torch from .ring_globals import ( HAS_AITER, HAS_FA3, + HAS_FA4, HAS_FLASH_ATTN, HAS_FLASHINFER, fa3_fwd_func, + fa4_attn_func, ) _scaled_dot_product_flash_attention = torch.ops.aten._scaled_dot_product_flash_attention @@ -190,6 +192,40 @@ def fa3_forward(q, k, v, dropout_p, softmax_scale, causal, window_size, softcap, flash_attn3_func_forward = fa3_forward +def flash_attn4_func_forward( + q, + k, + v, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), + softcap=None, + alibi_slopes=None, + return_softmax=False, +): + """FA4 forward pass returning the LSE required by Ring Attention.""" + del dropout_p, alibi_slopes, return_softmax + assert HAS_FA4, "FA4 is not available" + assert fa4_attn_func is not None, "FA4 high-level API is not available" + + left, right = window_size if window_size else (-1, -1) + out, softmax_lse = fa4_attn_func( + q, + k, + v, + softmax_scale=softmax_scale, + causal=causal, + window_size=( + None if left == -1 else left, + None if right == -1 else right, + ), + softcap=softcap or 0.0, + return_lse=True, + ) + return out, softmax_lse + + def flash_attn_forward_aiter( q, k, diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_selector.py b/vllm_omni/diffusion/attention/backends/ring/ring_selector.py index 77189c61b..05522c506 100644 --- a/vllm_omni/diffusion/attention/backends/ring/ring_selector.py +++ b/vllm_omni/diffusion/attention/backends/ring/ring_selector.py @@ -15,6 +15,7 @@ from .ring_globals import ( ) from .ring_kernels import ( flash_attn3_func_forward, + flash_attn4_func_forward, flash_attn_forward, flash_attn_forward_aiter, flashinfer_attn_forward, @@ -32,6 +33,7 @@ class AttnType(Enum): AITER = "aiter" FA = "fa" FA3 = "fa3" + FA4 = "fa4" FLASHINFER = "flashinfer" TORCH = "torch" SAGE_AUTO = "sage_auto" @@ -77,6 +79,9 @@ def select_flash_attn_impl( elif impl_type == AttnType.FA3: return flash_attn3_func_forward + elif impl_type == AttnType.FA4: + return flash_attn4_func_forward + elif impl_type == AttnType.FLASHINFER: return flashinfer_attn_forward diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py index e30d45ff2..d7a3a703f 100644 --- a/vllm_omni/diffusion/attention/backends/utils/fa.py +++ b/vllm_omni/diffusion/attention/backends/utils/fa.py @@ -22,6 +22,15 @@ from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) + +def _is_blackwell() -> bool: + try: + capability = current_omni_platform.get_device_capability() + except Exception: + return False + return capability is not None and capability.major >= 10 + + # Flash Attention function detection with fallback chain flash_attn_func = None flash_attn_varlen_func = None @@ -48,12 +57,27 @@ elif current_omni_platform.is_musa(): except (ImportError, ModuleNotFoundError): pass else: - # CUDA: try FA3 -> FA2 fallback chain + # CUDA: prefer the CuTe-based FA4 package when present. This matters on + # Blackwell, where an importable Hopper-only FA3 wheel may otherwise be + # selected and fail at launch with "no kernel image is available". + if _is_blackwell(): + try: + from flash_attn.cute import flash_attn_func, flash_attn_varlen_func # noqa: F401 + + logger.info("Using CuTe FlashAttention-4 on Blackwell") + except Exception as exc: + # Optional FA4 dependencies may be present but ABI-incompatible + # (for example, mismatched CUTLASS DSL and Quack builds). Treat + # that as unavailable so backend discovery can continue. + logger.debug("CuTe FlashAttention-4 is unavailable: %s", exc) + + # Try FA3 -> FA2 fallback chain. # Try FA3 from fa3-fwd PyPI package - try: - from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func # noqa: F401 - except (ImportError, ModuleNotFoundError): - pass + if flash_attn_func is None: + try: + from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func # noqa: F401 + except (ImportError, ModuleNotFoundError): + pass # Fallback: Try FA3 from flash-attention source build if flash_attn_func is None: @@ -100,7 +124,16 @@ def is_flash_attn_installed() -> bool: Shared by CUDA/ROCm/MUSA platforms. """ try: - # Check for any FA backend: FA3 (fa3_fwd_interface, flash_attn_interface) or FA2 (flash_attn) + # Check for any FA backend: FA4 (flash_attn.cute), FA3 + # (fa3_fwd_interface, flash_attn_interface), or FA2 (flash_attn). + if _is_blackwell(): + try: + from flash_attn.cute import flash_attn_varlen_func # noqa: F401 + + return True + except Exception as exc: + logger.debug("CuTe FlashAttention-4 is unavailable: %s", exc) + # Try FA3 from fa3-fwd PyPI package try: import fa3_fwd_interface # noqa: F401 diff --git a/vllm_omni/diffusion/attention/parallel/ring.py b/vllm_omni/diffusion/attention/parallel/ring.py index 745b25ccf..74c6fc3fb 100644 --- a/vllm_omni/diffusion/attention/parallel/ring.py +++ b/vllm_omni/diffusion/attention/parallel/ring.py @@ -10,7 +10,7 @@ import torch from vllm.logger import init_logger # import torch.distributed as dist # Not used directly here, but good practice if needed -from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_AITER, HAS_FA3, HAS_FLASH_ATTN +from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_AITER, HAS_FA3, HAS_FA4, HAS_FLASH_ATTN from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType from vllm_omni.diffusion.attention.parallel.base import ( ParallelAttentionContext, @@ -145,9 +145,12 @@ class RingParallelAttention: from vllm_omni.diffusion.attention.backends.ring_flash_attn import ring_flash_attn_func - # Prefer FA3 over FA2 for better performance (FA3 supports Ampere/Ada/Hopper) - # On ROCm, use AITER - if HAS_FA3: + # Prefer FA4 on Blackwell. An importable Hopper-only FA3 wheel can + # otherwise be selected and fail at launch with "no kernel image". + if HAS_FA4 and query.is_cuda and torch.cuda.get_device_capability(query.device)[0] >= 10: + attn_type = AttnType.FA4 + # Prefer FA3 over FA2 on Ampere/Ada/Hopper. On ROCm, use AITER. + elif HAS_FA3: attn_type = AttnType.FA3 elif HAS_AITER: attn_type = AttnType.AITER diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 622d3ce7d..493e33dd5 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -192,6 +192,16 @@ class DiffusionParallelConfig: vae_patch_parallel_size: int = 1 """Number of ranks used for VAE patch/tile parallelism (decode/encode).""" + text_encoder_tp_size: int = 1 + """Number of ranks used to tensor-parallel shard the diffusion text encoder. + + Ranks are the first ``text_encoder_tp_size`` DiT ranks. Defaults to 1, + which keeps the encoder fully resident on the DiT main rank (historical + behavior). Values > 1 shard the Qwen3-VL encoder across the encoder TP + ranks and run the encode with distributed collectives over the encoder + process group. + """ + vae_parallel_mode: str = "tile" """VAE parallel decode strategy. diff --git a/vllm_omni/diffusion/model_metadata.py b/vllm_omni/diffusion/model_metadata.py index 4952539b7..a6d6a53de 100644 --- a/vllm_omni/diffusion/model_metadata.py +++ b/vllm_omni/diffusion/model_metadata.py @@ -36,6 +36,10 @@ _DIFFUSION_MODEL_METADATA: dict[str, DiffusionModelMetadata] = { supports_multimodal_inputs=True, max_multimodal_image_inputs=BOOGU_IMAGE_MAX_INPUT_IMAGES, ), + "MiniMaxH3Pipeline": DiffusionModelMetadata( + supports_multimodal_inputs=True, + max_multimodal_image_inputs=1, + ), "WanPipeline": DiffusionModelMetadata(attention_mask_free=True), "WanImageToVideoPipeline": DiffusionModelMetadata(attention_mask_free=True), "WanVACEPipeline": DiffusionModelMetadata(attention_mask_free=True), diff --git a/vllm_omni/diffusion/models/minimax_h3/__init__.py b/vllm_omni/diffusion/models/minimax_h3/__init__.py new file mode 100644 index 000000000..53052ad42 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 audio-video diffusion support.""" + +from .pipeline_minimax_h3 import ( + MiniMaxH3Pipeline, + get_minimax_h3_post_process_func, +) + +__all__ = [ + "MiniMaxH3Pipeline", + "get_minimax_h3_post_process_func", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/condition_noise.py b/vllm_omni/diffusion/models/minimax_h3/condition_noise.py new file mode 100644 index 000000000..6534b884e --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/condition_noise.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 visual/audio condition-noise augmentation. + +The request's condition timestep is applied to both the tensor value and the +DiT timestep. Tokenizer artifacts remain clean +and reusable; this module materializes the fixed noised anchors immediately +before the denoise loop. +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import torch + +from .packed_tokens import minimax_h3_patchify_video_latent + +# Channel-major packed audio rows always carry a stereo layout. +MINIMAX_H3_AUDIO_COND_CHANNELS = 2 + + +def minimax_h3_imgvid_cond_noise_aug_rows( + clean_rows: torch.Tensor, + *, + condition_shapes: Sequence[Sequence[int]], + target_latent_t: int, + imgvid_cond_num_frames: int, + seed: int, + noise_aug: float, +) -> torch.Tensor: + """Apply the imgvid-condition RF noise recipe to packed clean rows. + + ``condition_shapes`` contains ``(latent_t, latent_h, latent_w)`` in packed + visual-condition order. A new CPU generator with the same row seed is + created for every condition. Under the dependent-noise policy, each draw + uses the target temporal length plus the template's imgvid-condition frame + count, then slices the prefix matching the current condition. + """ + + noise_aug = float(noise_aug) + if not 0.0 <= noise_aug <= 1.0: + raise ValueError(f"noise_aug must be in [0, 1], got {noise_aug}") + if noise_aug == 1.0: + return clean_rows + if clean_rows.ndim != 2 or int(clean_rows.shape[1]) != 96: + raise ValueError(f"clean imgvid condition rows must have shape [n, 96], got {list(clean_rows.shape)}") + + target_latent_t = int(target_latent_t) + imgvid_cond_num_frames = int(imgvid_cond_num_frames) + if target_latent_t <= 0: + raise ValueError(f"target_latent_t must be positive, got {target_latent_t}") + if imgvid_cond_num_frames <= 0: + raise ValueError( + f"imgvid_cond_num_frames must be positive when condition rows exist, got {imgvid_cond_num_frames}" + ) + + parsed_shapes: list[tuple[int, int, int]] = [] + expected_rows = 0 + for raw_shape in condition_shapes: + if len(raw_shape) != 3: + raise ValueError( + f"each imgvid condition shape must be (latent_t, latent_h, latent_w), got {list(raw_shape)}" + ) + latent_t, latent_h, latent_w = (int(value) for value in raw_shape) + if latent_t <= 0 or latent_h <= 0 or latent_w <= 0: + raise ValueError(f"imgvid condition shape must be positive, got {list(raw_shape)}") + if latent_h % 2 or latent_w % 2: + raise ValueError( + f"imgvid condition spatial dimensions must be divisible by 2, got {(latent_t, latent_h, latent_w)}" + ) + parsed_shapes.append((latent_t, latent_h, latent_w)) + expected_rows += latent_t * (latent_h // 2) * (latent_w // 2) + if not parsed_shapes: + raise ValueError("condition_shapes must not be empty") + if int(clean_rows.shape[0]) != expected_rows: + raise ValueError( + f"clean imgvid condition rows {int(clean_rows.shape[0])} != shape-derived rows {expected_rows}" + ) + + output_device = clean_rows.device + clean_rows = clean_rows.detach().to(device="cpu", dtype=torch.float32) + out: list[torch.Tensor] = [] + row_offset = 0 + timestep = torch.tensor(noise_aug, dtype=torch.float32, device="cpu") + for latent_t, latent_h, latent_w in parsed_shapes: + full_t = target_latent_t + imgvid_cond_num_frames + if full_t < latent_t: + raise ValueError(f"condition latent_t {latent_t} exceeds the noise draw length {full_t}") + generator = torch.Generator(device="cpu").manual_seed(int(seed)) + noise = torch.randn( + 1, + 24, + full_t, + latent_h, + latent_w, + generator=generator, + dtype=torch.float32, + device="cpu", + )[:, :, :latent_t] + noise_rows = minimax_h3_patchify_video_latent(noise, patch_size=[1, 2, 2]).to(dtype=torch.float32) + row_count = int(noise_rows.shape[0]) + clean_part = clean_rows[row_offset : row_offset + row_count].to(torch.float32) + out.append(timestep * clean_part + (1.0 - timestep) * noise_rows) + row_offset += row_count + return torch.cat(out, dim=0).to(device=output_device, dtype=torch.float32).contiguous() + + +def minimax_h3_audio_cond_noise_aug_rows( + clean_rows: torch.Tensor, + *, + condition_audio_t: Sequence[int], + seed: int, + noise_aug: float, +) -> torch.Tensor: + """Apply the audio-condition RF noise recipe to packed clean rows. + + ``condition_audio_t`` contains the latent T of each audio-bearing + condition in canonical request order. Noise is drawn per condition + element, with a fresh CPU generator seeded with ``seed + 1`` for every + element. Consequently each condition restarts the + same RNG stream; concatenating the rows and drawing once would be + numerically different for ordered multi-reference requests. + + The mix is intentionally evaluated on CPU in fp32 before the packed rows + are transferred to the DiT device. + """ + + noise_aug = float(noise_aug) + if not 0.0 <= noise_aug <= 1.0: + raise ValueError(f"noise_aug must be in [0, 1], got {noise_aug}") + if noise_aug == 1.0: + return clean_rows + if clean_rows.ndim != 2 or int(clean_rows.shape[1]) != 32: + raise ValueError(f"clean audio condition rows must have shape [n, 32], got {list(clean_rows.shape)}") + + audio_channels = MINIMAX_H3_AUDIO_COND_CHANNELS + parsed_audio_t = [int(value) for value in condition_audio_t] + if not parsed_audio_t: + raise ValueError("condition_audio_t must not be empty") + if any(value <= 0 for value in parsed_audio_t): + raise ValueError(f"condition audio latent lengths must be positive, got {parsed_audio_t}") + expected_rows = audio_channels * sum(parsed_audio_t) + if int(clean_rows.shape[0]) != expected_rows: + raise ValueError(f"clean audio condition rows {int(clean_rows.shape[0])} != shape-derived rows {expected_rows}") + + out: list[torch.Tensor] = [] + row_offset = 0 + timestep = torch.tensor(noise_aug, dtype=torch.float32, device="cpu") + for audio_t in parsed_audio_t: + row_count = audio_channels * audio_t + clean_part = clean_rows[row_offset : row_offset + row_count].detach().to(device="cpu", dtype=torch.float32) + generator = torch.Generator(device="cpu").manual_seed(int(seed) + 1) + noise = torch.randn( + clean_part.shape, + generator=generator, + dtype=torch.float32, + device="cpu", + ) + out.append(timestep * clean_part + (1.0 - timestep) * noise) + row_offset += row_count + return torch.cat(out, dim=0).to(device=clean_rows.device, dtype=torch.float32).contiguous() + + +__all__ = [ + "minimax_h3_audio_cond_noise_aug_rows", + "minimax_h3_imgvid_cond_noise_aug_rows", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py new file mode 100644 index 000000000..b051951a7 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py @@ -0,0 +1,249 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 cfg-distilled full denoise loop. + +Per step, the positive presentation is forwarded exactly once. Video and audio +target rows chain through the Euler-eta0 update while visual and audio condition +rows stay pinned to their noised step-0 anchors. +""" + +from __future__ import annotations + +from collections.abc import Callable +from contextlib import AbstractContextManager, nullcontext +from typing import Any + +import torch + +from .scheduling_minimax_h3_euler_ancestral import ( + minimax_h3_euler_eta0_step, + minimax_h3_rf_v_to_x0, +) + +MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999 +# ref2va audio reference anchor timestep +MINIMAX_H3_AUDIO_REF_COND_TIMESTEP = 1.0 +# Packed row widths: video rows are [1,2,2]-patchified 24-channel latents +# (24 * 1 * 2 * 2 = 96); audio rows carry the 32-dim audio latent. +MINIMAX_H3_VIDEO_ROW_WIDTH = 96 +MINIMAX_H3_AUDIO_ROW_WIDTH = 32 + + +class MiniMaxH3DenoiseBranch: + """Static per-branch state: packed layout + fixed forward kwargs. + + `packed` is a minimax_h3_packed_sequence(...) result (or equivalent layout + dict); `text_embeddings` is the branch's [text_len, 5120] hidden states; + `token_tags` must already carry any fl2va vision-span overrides. + """ + + def __init__( + self, + *, + packed: dict[str, torch.Tensor], + text_embeddings: torch.Tensor, + token_tags: torch.Tensor, + device: torch.device, + ) -> None: + seq_len = int(packed["seq_len"]) + self.seq_len = seq_len + self.img_pos = packed["img_pos"].view(-1).to(torch.long) + self.audio_pos = packed["audio_pos"].view(-1).to(torch.long) + self.update_mask = packed["update_mask"].view(-1).to(torch.bool) + # ref2va: audio_pos may include reference-audio anchor rows + # (audio_update_mask False); absent means all rows are targets. + if "audio_update_mask" in packed: + self.audio_update_mask = packed["audio_update_mask"].view(-1).to(torch.bool) + else: + self.audio_update_mask = torch.ones(self.audio_pos.shape[0], dtype=torch.bool) + text_len = int(packed["text_pos"].view(-1).shape[0]) + if list(text_embeddings.shape)[0] != text_len: + raise ValueError(f"text_embeddings rows {list(text_embeddings.shape)} != packed text_len {text_len}") + if int(token_tags.view(-1).shape[0]) != seq_len: + raise ValueError(f"token_tags length {int(token_tags.view(-1).shape[0])} != seq_len {seq_len}") + cu = packed["cu_seqlens"].to(torch.int32) + self.img_pos_dev = self.img_pos.to(device) + self.audio_pos_dev = self.audio_pos.to(device) + self.update_mask_dev = self.update_mask.to(device) + self.audio_update_mask_dev = self.audio_update_mask.to(device) + self.x_base = torch.zeros(1, seq_len, MINIMAX_H3_VIDEO_ROW_WIDTH, dtype=torch.float32, device=device) + self.audio_x_base = torch.zeros(1, seq_len, MINIMAX_H3_AUDIO_ROW_WIDTH, dtype=torch.float32, device=device) + text_pos_dev = packed["text_pos"].view(-1).to(torch.long).to(device) + self.static_kwargs: dict[str, Any] = { + "img_position_ids": packed["img_position_ids"][None].to(device), + "update_mask": self.update_mask_dev, + "token_tags": token_tags.view(-1).to(torch.long).to(device), + "skip_mask_out_condition": False, + "prompt_embeds": text_embeddings.to(device), + "img_pos_info": {"position_ids": self.img_pos_dev}, + "audio_pos_info": {"position_ids": self.audio_pos_dev}, + "text_pos_info": {"position_ids": text_pos_dev}, + "img_pos_for_infer_output_info": {"position_ids": self.img_pos_dev}, + "packed_seq_params": { + "cu_seqlens_q": cu.to(device), + "max_seqlen_q": int(cu[1]), + }, + "refiner_packed_seq_params": { + "cu_seqlens_q": torch.tensor([0, text_len, text_len], dtype=torch.int32, device=device), + "max_seqlen_q": text_len, + }, + } + + def forward_kwargs( + self, + *, + video_rows: torch.Tensor, + audio_rows: torch.Tensor, + t_video: float, + t_audio: float, + imgvid_cond_timestep: float, + audio_ref_cond_timestep: float, + ) -> dict[str, Any]: + x = self.x_base.clone() + x[0].index_copy_(0, self.img_pos_dev, video_rows) + audio_x = self.audio_x_base.clone() + audio_x[0].index_copy_(0, self.audio_pos_dev, audio_rows) + # Packed-sequence timestep semantics: non-media rows (text and + # padding) inherit the current video timestep. Later steps must reuse + # the previous step's updated rows; re-initializing from zeros is only + # valid at step 0. + timesteps = torch.full( + (self.seq_len,), + float(t_video), + dtype=torch.float32, + device=x.device, + ) + timesteps[self.img_pos_dev[self.update_mask_dev]] = t_video + timesteps[self.img_pos_dev[~self.update_mask_dev]] = imgvid_cond_timestep + timesteps[self.audio_pos_dev[self.audio_update_mask_dev]] = t_audio + timesteps[self.audio_pos_dev[~self.audio_update_mask_dev]] = audio_ref_cond_timestep + unique_timesteps, inverse_indices = torch.unique(timesteps, sorted=True, return_inverse=True) + return { + **self.static_kwargs, + "x": x, + "audio_x": audio_x, + "unique_timesteps": unique_timesteps, + "inverse_indices": inverse_indices, + } + + +def minimax_h3_denoise_loop( + *, + model: Any, + positive: MiniMaxH3DenoiseBranch, + initial_video_rows: torch.Tensor, + initial_audio_rows: torch.Tensor, + keyframe_cond_rows: torch.Tensor | None, + audio_ref_rows: torch.Tensor | None = None, + sigmas_video: list[float], + sigmas_audio: list[float], + device: torch.device, + imgvid_cond_noise_aug_for_inference: float = MINIMAX_H3_IMGVID_COND_TIMESTEP, + audio_cond_noise_aug_for_inference: float = MINIMAX_H3_AUDIO_REF_COND_TIMESTEP, + on_step: Callable[[int, torch.Tensor, torch.Tensor], None] | None = None, + step_profiler: Callable[[int], AbstractContextManager] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the full denoise loop; returns final (video_rows, audio_rows). + + ``initial_video_rows`` covers all image rows of the positive layout. For a + conditional task, pass ``keyframe_cond_rows`` and/or ``audio_ref_rows`` to + pin those rows across every step. The model's raw positive velocity is the + update signal; MiniMax H3 only supports cfg-distilled checkpoints. + """ + if len(sigmas_video) != len(sigmas_audio): + raise ValueError("video/audio sigma schedules must have equal length") + if len(sigmas_video) < 2: + raise ValueError("sigma schedules need at least 2 entries") + n_cond = int((~positive.update_mask).sum()) + if keyframe_cond_rows is None: + if n_cond != 0: + raise ValueError(f"layout has {n_cond} cond rows but keyframe_cond_rows is None") + else: + if int(keyframe_cond_rows.shape[0]) != n_cond: + raise ValueError(f"keyframe_cond_rows {int(keyframe_cond_rows.shape[0])} != layout cond rows {n_cond}") + video_rows = initial_video_rows.to(device=device, dtype=torch.float32).clone() + audio_rows = initial_audio_rows.to(device=device, dtype=torch.float32).clone() + update = positive.update_mask_dev + audio_update = positive.audio_update_mask_dev + if int(video_rows.shape[0]) != int(positive.img_pos.shape[0]): + raise ValueError( + f"initial video rows {int(video_rows.shape[0])} != positive layout rows {int(positive.img_pos.shape[0])}" + ) + if int(audio_rows.shape[0]) != int(positive.audio_pos.shape[0]): + raise ValueError( + f"initial audio rows {int(audio_rows.shape[0])} != positive layout rows {int(positive.audio_pos.shape[0])}" + ) + n_audio_ref = int((~positive.audio_update_mask).sum()) + if audio_ref_rows is None: + if n_audio_ref != 0: + raise ValueError(f"layout has {n_audio_ref} audio ref rows but audio_ref_rows is None") + audio_anchor = None + else: + if int(audio_ref_rows.shape[0]) != n_audio_ref: + raise ValueError(f"audio_ref_rows {int(audio_ref_rows.shape[0])} != layout audio ref rows {n_audio_ref}") + audio_anchor = audio_ref_rows.to(device=device, dtype=torch.float32) + cond_anchor = keyframe_cond_rows.to(device=device, dtype=torch.float32) if keyframe_cond_rows is not None else None + if cond_anchor is not None: + video_rows[~update] = cond_anchor + if audio_anchor is not None: + audio_rows[~audio_update] = audio_anchor + + num_steps = len(sigmas_video) - 1 + for step in range(num_steps): + step_cm = step_profiler(step) if step_profiler is not None else nullcontext() + with step_cm: + s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1] + s_a, s_a_next = sigmas_audio[step], sigmas_audio[step + 1] + t_v, t_a = 1.0 - s_v, 1.0 - s_a + imgvid_cond_t = max(t_v, float(imgvid_cond_noise_aug_for_inference)) + audio_ref_cond_t = max(t_a, float(audio_cond_noise_aug_for_inference)) + + fk = positive.forward_kwargs( + video_rows=video_rows, + audio_rows=audio_rows, + t_video=t_v, + t_audio=t_a, + imgvid_cond_timestep=imgvid_cond_t, + audio_ref_cond_timestep=audio_ref_cond_t, + ) + with torch.inference_mode(): + v_video, v_audio = model(**fk) + mv_video_t = v_video.float()[update] + mv_audio_t = v_audio.float()[audio_update] + + x0_video = minimax_h3_rf_v_to_x0( + video_rows[update], + mv_video_t, + torch.tensor(t_v, dtype=torch.float32, device=device), + ) + new_target = minimax_h3_euler_eta0_step(video_rows[update], x0_video, sigma_curr=s_v, sigma_next=s_v_next) + video_rows = video_rows.clone() + video_rows[update] = new_target + if cond_anchor is not None: + video_rows[~update] = cond_anchor # per-step imgvid cond reset + + x0_audio = minimax_h3_rf_v_to_x0( + audio_rows[audio_update], + mv_audio_t, + torch.tensor(t_a, dtype=torch.float32, device=device), + ) + new_audio = minimax_h3_euler_eta0_step( + audio_rows[audio_update], x0_audio, sigma_curr=s_a, sigma_next=s_a_next + ) + audio_rows = audio_rows.clone() + audio_rows[audio_update] = new_audio + if audio_anchor is not None: + audio_rows[~audio_update] = audio_anchor # per-step audio ref reset + if on_step is not None: + on_step(step, video_rows, audio_rows) + + return video_rows, audio_rows + + +__all__ = [ + "MINIMAX_H3_AUDIO_REF_COND_TIMESTEP", + "MINIMAX_H3_AUDIO_ROW_WIDTH", + "MINIMAX_H3_IMGVID_COND_TIMESTEP", + "MINIMAX_H3_VIDEO_ROW_WIDTH", + "MiniMaxH3DenoiseBranch", + "minimax_h3_denoise_loop", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/encoder.py b/vllm_omni/diffusion/models/minimax_h3/encoder.py new file mode 100644 index 000000000..438ff80a2 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/encoder.py @@ -0,0 +1,1214 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 Qwen3-VL layer-50 text/vision encoder. + +The encoder is reimplemented on top of vLLM-style tensor-parallel building +blocks (column/row/vocab-parallel linears bound to the MiniMax H3 encoder +process group) instead of the single-GPU ``transformers`` backbone. This +makes the large Qwen3-VL text model TP-shardable across ``text_encoder_tp_size`` +ranks, which removes the rank-0 memory hotspot that dominates no-offload peak +memory (the retained 50-layer encoder is ~51.5 GB in BF16). + +The computation contract is unchanged from the HF reference path: + +- plain BF16 weights, only the first ``MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER`` + decoder layers are kept (the checkpoint consumes the unnormalized hidden + state after decoder layer 49); +- cuDNN SDP is enabled during encode; +- the multimodal backbone runs with an all-ones attention mask and + ``mm_token_type_ids`` derived from the image/video token ids; +- DeepStack visual features are injected at the first + ``len(deepstack_visual_indexes)`` decoder layers; +- the output is ``hidden_states[50]`` of shape ``[seq, 5120]``. + +Tensor parallelism follows vLLM's TP semantics: vocab-parallel embeddings and +row-parallel projections all-reduce so the hidden state stays fully replicated +on every encoder rank, while column-parallel projections (QKV / MLP up) keep a +local shard. After the final layer every encoder rank therefore holds the +complete ``[seq, 5120]`` hidden state. +""" + +from __future__ import annotations + +import itertools +import json +import re +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.logger import init_logger + +MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER = 50 +MINIMAX_H3_QWEN3VL_HIDDEN_DIM = 5120 + +logger = init_logger(__name__) + + +def _default_weight_loader( + param: nn.Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: str | None = None, +) -> None: + del loaded_shard_id + param.data.copy_(loaded_weight) + + +def _tp_range(group: Any) -> tuple[int, int]: + return group.rank_in_group, group.world_size + + +# --------------------------------------------------------------------------- +# Tensor-parallel primitives bound to the encoder process group. +# +# These mirror the math of vLLM's ``VocabParallelEmbedding``, +# ``ColumnParallelLinear``, ``MergedColumnParallelLinear``, ``QKVParallelLinear`` +# and ``RowParallelLinear`` but route their collectives through the encoder +# group instead of the (DiT) tensor-model-parallel group, so the text encoder +# can be sharded independently of the DiT's TP configuration. +# --------------------------------------------------------------------------- + + +class MiniMaxH3Qwen3VLVocabParallelEmbedding(nn.Module): + def __init__(self, group: Any, num_embeddings: int, embedding_dim: int, dtype: torch.dtype) -> None: + super().__init__() + self.group = group + self.num_embeddings = num_embeddings + self.embedding_dim = embedding_dim + tp_rank, tp_size = _tp_range(group) + assert num_embeddings % tp_size == 0, ( + f"vocab_size {num_embeddings} must be divisible by text_encoder_tp_size {tp_size}" + ) + self.num_embeddings_per_partition = num_embeddings // tp_size + self.weight = nn.Parameter(torch.empty(self.num_embeddings_per_partition, embedding_dim, dtype=dtype)) + self.weight.weight_loader = self.weight_loader # type: ignore[attr-defined] + self._tp_rank = tp_rank + self._tp_size = tp_size + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + if self._tp_size == 1: + return F.embedding(input_, self.weight) + # Mirror vLLM's `VocabParallelEmbedding`: token ids that fall outside + # this rank's vocab partition are masked to 0 and in-partition ids are + # shifted to local indices before the local gather. The out-of-range + # rows are zeroed so the subsequent all-reduce reconstructs the full + # embedding for every token. + start_idx = self._tp_rank * self.num_embeddings_per_partition + end_idx = start_idx + self.num_embeddings_per_partition + input_mask = (input_ < start_idx) | (input_ >= end_idx) + masked_input = input_.clone() - start_idx + masked_input[input_mask] = 0 + output = F.embedding(masked_input, self.weight) + output[input_mask, :] = 0.0 + self.group.all_reduce(output) + return output + + def weight_loader( + self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None + ) -> None: + del loaded_shard_id + shard_size = self.num_embeddings_per_partition + start_idx = self._tp_rank * shard_size + param.data.copy_(loaded_weight.narrow(0, start_idx, shard_size)) + + +class MiniMaxH3Qwen3VLMergedColumnParallelLinear(nn.Module): + """Packed gate/up projection sharded along the output dimension.""" + + def __init__( + self, + group: Any, + input_size: int, + intermediate_size: int, + dtype: torch.dtype, + ) -> None: + super().__init__() + self.group = group + self.input_size = input_size + self.intermediate_size = intermediate_size + tp_rank, tp_size = _tp_range(group) + assert intermediate_size % tp_size == 0, ( + f"intermediate_size {intermediate_size} must be divisible by text_encoder_tp_size {tp_size}" + ) + self.intermediate_size_per_partition = intermediate_size // tp_size + self.weight = nn.Parameter(torch.empty(2 * self.intermediate_size_per_partition, input_size, dtype=dtype)) + self.weight.weight_loader = self.weight_loader # type: ignore[attr-defined] + self._tp_rank = tp_rank + self._tp_size = tp_size + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + return F.linear(input_, self.weight) + + def weight_loader( + self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None + ) -> None: + shard_size = self.intermediate_size_per_partition + start_idx = self._tp_rank * shard_size + if loaded_shard_id == 1: # up_proj + param.data[shard_size : 2 * shard_size].copy_(loaded_weight.narrow(0, start_idx, shard_size)) + else: # gate_proj + param.data[0:shard_size].copy_(loaded_weight.narrow(0, start_idx, shard_size)) + + +class MiniMaxH3Qwen3VLQKVParallelLinear(nn.Module): + """QKV projection with GQA head sharding along the output dimension.""" + + def __init__( + self, + group: Any, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + dtype: torch.dtype, + ) -> None: + super().__init__() + self.group = group + self.hidden_size = hidden_size + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads + self.head_dim = head_dim + tp_rank, tp_size = _tp_range(group) + assert num_heads % tp_size == 0, ( + f"num_attention_heads {num_heads} must be divisible by text_encoder_tp_size {tp_size}" + ) + assert num_kv_heads % tp_size == 0, ( + f"num_key_value_heads {num_kv_heads} must be divisible by text_encoder_tp_size {tp_size}" + ) + self.local_num_heads = num_heads // tp_size + self.local_num_kv_heads = num_kv_heads // tp_size + q_local = self.local_num_heads * head_dim + kv_local = self.local_num_kv_heads * head_dim + self.weight = nn.Parameter(torch.empty(q_local + 2 * kv_local, hidden_size, dtype=dtype)) + self.weight.weight_loader = self.weight_loader # type: ignore[attr-defined] + self._tp_rank = tp_rank + self._tp_size = tp_size + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + return F.linear(input_, self.weight) + + def weight_loader( + self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None + ) -> None: + head_dim = self.head_dim + q_local = self.local_num_heads * head_dim + kv_local = self.local_num_kv_heads * head_dim + if loaded_shard_id == "q": + start_idx = self._tp_rank * q_local + param.data[0:q_local].copy_(loaded_weight.narrow(0, start_idx, q_local)) + elif loaded_shard_id == "k": + start_idx = self._tp_rank * kv_local + param.data[q_local : q_local + kv_local].copy_(loaded_weight.narrow(0, start_idx, kv_local)) + elif loaded_shard_id == "v": + start_idx = self._tp_rank * kv_local + param.data[q_local + kv_local : q_local + 2 * kv_local].copy_(loaded_weight.narrow(0, start_idx, kv_local)) + else: + raise ValueError(f"unexpected QKV shard id {loaded_shard_id!r}") + + +class MiniMaxH3Qwen3VLRowParallelLinear(nn.Module): + def __init__( + self, + group: Any, + input_size: int, + output_size: int, + dtype: torch.dtype, + input_is_parallel: bool = True, + ) -> None: + super().__init__() + self.group = group + self.input_size = input_size + self.output_size = output_size + self.input_is_parallel = input_is_parallel + tp_rank, tp_size = _tp_range(group) + assert input_size % tp_size == 0, f"input_size {input_size} must be divisible by text_encoder_tp_size {tp_size}" + self.input_size_per_partition = input_size // tp_size + self.weight = nn.Parameter(torch.empty(output_size, self.input_size_per_partition, dtype=dtype)) + self.weight.weight_loader = self.weight_loader # type: ignore[attr-defined] + self._tp_rank = tp_rank + self._tp_size = tp_size + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + if self.input_is_parallel: + input_parallel = input_ + else: + split_input = input_.split(self.input_size_per_partition, dim=-1) + input_parallel = split_input[self._tp_rank].contiguous() + output_parallel = F.linear(input_parallel, self.weight) + if self._tp_size > 1: + # Reduce in fp32: the reference path accumulates the full (K=8192) + # dot product inside a single cuBLAS GEMM before rounding to bf16. + # A bf16 all-reduce of the per-rank partial GEMMs would round twice + # and amplify error on this model's large-magnitude activations. + output_parallel = output_parallel.float() + self.group.all_reduce(output_parallel) + output_parallel = output_parallel.to(self.weight.dtype) + return output_parallel + + def weight_loader( + self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str | None = None + ) -> None: + del loaded_shard_id + shard_size = self.input_size_per_partition + start_idx = self._tp_rank * shard_size + param.data.copy_(loaded_weight.narrow(1, start_idx, shard_size)) + + +# --------------------------------------------------------------------------- +# Shared math ports (identical to the HF reference path). +# --------------------------------------------------------------------------- + + +class MiniMaxH3Qwen3VLRMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def _apply_rotary_pos_emb( + q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + q_embed = (q * cos) + (_rotate_half(q) * sin) + k_embed = (k * cos) + (_rotate_half(k) * sin) + return q_embed, k_embed + + +def _apply_interleaved_mrope(freqs: torch.Tensor, mrope_section: list[int]) -> torch.Tensor: + """Reorganize chunked [TTT...HHH...WWW] into interleaved [THWTHW...] layout.""" + freqs_t = freqs[0] + for dim, offset in enumerate((1, 2), start=1): + length = mrope_section[dim] * 3 + idx = slice(offset, length, 3) + freqs_t[..., idx] = freqs[dim, ..., idx] + return freqs_t + + +class MiniMaxH3Qwen3VLTextRotaryEmbedding(nn.Module): + """Qwen3-VL text M-RoPE (interleaved sections). Same math as the reference.""" + + def __init__(self, config: Any) -> None: + super().__init__() + self.config = config + # ``dim`` and ``base`` are plain attributes so the inverse frequencies + # are always recomputed in float32 (module-wide dtype casts must not + # degrade the RoPE cache precision). + self.dim = int(getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads) + self.base = float(config.rope_parameters["rope_theta"]) + self.mrope_section = config.rope_parameters.get("mrope_section", [24, 20, 20]) + + def _get_inv_freq(self, device: torch.device) -> torch.Tensor: + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim)) + return inv_freq.to(device) + + def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # position_ids: (3, 1, seq) + inv_freq = self._get_inv_freq(position_ids.device) + inv_freq_expanded = inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1) + position_ids_expanded = position_ids[:, :, None, :].float() + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) + freqs = _apply_interleaved_mrope(freqs, self.mrope_section) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +# --------------------------------------------------------------------------- +# Qwen3-VL vision tower (replicated on every encoder rank; ~1.2 GB BF16). +# --------------------------------------------------------------------------- + + +class MiniMaxH3Qwen3VLVisionPatchEmbed(nn.Module): + def __init__(self, config: Any) -> None: + super().__init__() + self.patch_size = config.patch_size + self.temporal_patch_size = config.temporal_patch_size + self.in_channels = config.in_channels + self.embed_dim = config.hidden_size + kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size] + self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + target_dtype = self.proj.weight.dtype + hidden_states = hidden_states.view( + -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size + ) + hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim) + return hidden_states + + +class MiniMaxH3Qwen3VLVisionRotaryEmbedding(nn.Module): + def __init__(self, dim: int, theta: float = 10000.0) -> None: + super().__init__() + self.dim = dim + self.theta = theta + # Plain attribute (not a registered buffer): the reference keeps this + # table in float32 and module-wide dtype casts must not touch it. + self._inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) + + def forward(self, seqlen: int) -> torch.Tensor: + inv_freq = self._inv_freq.to(device=self._inv_freq.device) + seq = torch.arange(seqlen, device=inv_freq.device, dtype=torch.float) + freqs = torch.outer(seq, inv_freq) + return freqs + + +class MiniMaxH3Qwen3VLVisionPatchMerger(nn.Module): + def __init__(self, config: Any, use_postshuffle_norm: bool = False) -> None: + super().__init__() + self.hidden_size = config.hidden_size * (config.spatial_merge_size**2) + self.use_postshuffle_norm = use_postshuffle_norm + self.norm = nn.LayerNorm(self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6) + self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size) + self.act_fn = nn.GELU() + self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.norm(x.view(-1, self.hidden_size) if self.use_postshuffle_norm else x).view(-1, self.hidden_size) + x = self.linear_fc2(self.act_fn(self.linear_fc1(x))) + return x + + +class MiniMaxH3Qwen3VLVisionMLP(nn.Module): + def __init__(self, config: Any) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.linear_fc1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=True) + self.linear_fc2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=True) + # The reference vision MLP uses the tanh GELU approximation + # (``gelu_pytorch_tanh``), which differs from the exact erf GELU. + self.act_fn = nn.GELU(approximate="tanh") + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state))) + + +def _apply_rotary_pos_emb_vision( + q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + orig_q_dtype = q.dtype + orig_k_dtype = k.dtype + q, k = q.float(), k.float() + cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float() + q_embed = (q * cos) + (_rotate_half(q) * sin) + k_embed = (k * cos) + (_rotate_half(k) * sin) + q_embed = q_embed.to(orig_q_dtype) + k_embed = k_embed.to(orig_k_dtype) + return q_embed, k_embed + + +class MiniMaxH3Qwen3VLVisionAttention(nn.Module): + def __init__(self, config: Any) -> None: + super().__init__() + self.dim = config.hidden_size + self.num_heads = config.num_heads + self.head_dim = self.dim // self.num_heads + self.num_key_value_groups = 1 + self.qkv = nn.Linear(self.dim, self.dim * 3, bias=True) + self.proj = nn.Linear(self.dim, self.dim) + self.scaling = self.head_dim**-0.5 + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + seq_length = hidden_states.shape[0] + query_states, key_states, value_states = ( + self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0) + ) + cos, sin = position_embeddings + query_states, key_states = _apply_rotary_pos_emb_vision(query_states, key_states, cos, sin) + + query_states = query_states.transpose(0, 1).unsqueeze(0) + key_states = key_states.transpose(0, 1).unsqueeze(0) + value_states = value_states.transpose(0, 1).unsqueeze(0) + + lengths = cu_seqlens[1:] - cu_seqlens[:-1] + splits = [torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)] + attn_outputs = [ + F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False) for q, k, v in zip(*splits) + ] + attn_output = torch.cat(attn_outputs, dim=2) + # The SDPA kernel may return a non-contiguous (1, num_heads, seq_len, head_dim) + # tensor whose memory layout is seq-major; transpose before flattening exactly + # like the reference `sdpa_attention_forward` to avoid scrambling heads. + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(seq_length, -1).contiguous() + attn_output = self.proj(attn_output) + return attn_output + + +class MiniMaxH3Qwen3VLVisionBlock(nn.Module): + def __init__(self, config: Any) -> None: + super().__init__() + self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6) + self.attn = MiniMaxH3Qwen3VLVisionAttention(config=config) + self.mlp = MiniMaxH3Qwen3VLVisionMLP(config=config) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + hidden_states = hidden_states + self.attn( + self.norm1(hidden_states), + cu_seqlens=cu_seqlens, + position_embeddings=position_embeddings, + ) + hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) + return hidden_states + + +class MiniMaxH3Qwen3VLVisionModel(nn.Module): + """Qwen3-VL vision tower (patch embed + blocks + merger + DeepStack).""" + + def __init__(self, config: Any) -> None: + super().__init__() + self.config = config + self.spatial_merge_size = config.spatial_merge_size + self.patch_size = config.patch_size + self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size + + self.patch_embed = MiniMaxH3Qwen3VLVisionPatchEmbed(config=config) + self.pos_embed = nn.Embedding(config.num_position_embeddings, config.hidden_size) + self.num_grid_per_side = int(config.num_position_embeddings**0.5) + + head_dim = config.hidden_size // config.num_heads + self.rotary_pos_emb = MiniMaxH3Qwen3VLVisionRotaryEmbedding(head_dim // 2) + + self.blocks = nn.ModuleList([MiniMaxH3Qwen3VLVisionBlock(config=config) for _ in range(config.depth)]) + self.merger = MiniMaxH3Qwen3VLVisionPatchMerger(config=config, use_postshuffle_norm=False) + self.deepstack_visual_indexes = config.deepstack_visual_indexes + self.deepstack_merger_list = nn.ModuleList( + [ + MiniMaxH3Qwen3VLVisionPatchMerger(config=config, use_postshuffle_norm=True) + for _ in range(len(config.deepstack_visual_indexes)) + ] + ) + + def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor: + merge_size = self.spatial_merge_size + grid_thw_list = grid_thw.tolist() + max_hw = max(max(h, w) for _, h, w in grid_thw_list) + freq_table = self.rotary_pos_emb(max_hw).to(device=grid_thw.device) + device = freq_table.device + total_tokens = sum(t * h * w for t, h, w in grid_thw_list) + pos_ids = torch.empty((total_tokens, 2), dtype=torch.long, device=device) + offset = 0 + for num_frames, height, width in grid_thw_list: + merged_h, merged_w = height // merge_size, width // merge_size + block_rows = torch.arange(merged_h, device=device) + block_cols = torch.arange(merged_w, device=device) + intra_row = torch.arange(merge_size, device=device) + intra_col = torch.arange(merge_size, device=device) + row_idx = block_rows[:, None, None, None] * merge_size + intra_row[None, None, :, None] + col_idx = block_cols[None, :, None, None] * merge_size + intra_col[None, None, None, :] + row_idx = row_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1) + col_idx = col_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1) + coords = torch.stack((row_idx, col_idx), dim=-1) + if num_frames > 1: + coords = coords.repeat(num_frames, 1) + num_tokens = coords.shape[0] + pos_ids[offset : offset + num_tokens] = coords + offset += num_tokens + embeddings = freq_table[pos_ids] + embeddings = embeddings.flatten(1) + return embeddings + + def fast_pos_embed_interpolate(self, grid_thw: torch.Tensor) -> torch.Tensor: + grid_thw_list = grid_thw.tolist() + grid_ts = [row[0] for row in grid_thw_list] + grid_hs = [row[1] for row in grid_thw_list] + grid_ws = [row[2] for row in grid_thw_list] + device = self.pos_embed.weight.device + idx_list = [[] for _ in range(4)] + weight_list = [[] for _ in range(4)] + for _, h, w in grid_thw_list: + h_idxs = torch.linspace(0, self.num_grid_per_side - 1, h) + w_idxs = torch.linspace(0, self.num_grid_per_side - 1, w) + h_idxs_floor = h_idxs.int() + w_idxs_floor = w_idxs.int() + h_idxs_ceil = (h_idxs.int() + 1).clip(max=self.num_grid_per_side - 1) + w_idxs_ceil = (w_idxs.int() + 1).clip(max=self.num_grid_per_side - 1) + dh = h_idxs - h_idxs_floor + dw = w_idxs - w_idxs_floor + base_h = h_idxs_floor * self.num_grid_per_side + base_h_ceil = h_idxs_ceil * self.num_grid_per_side + indices = [ + (base_h[None].T + w_idxs_floor[None]).flatten(), + (base_h[None].T + w_idxs_ceil[None]).flatten(), + (base_h_ceil[None].T + w_idxs_floor[None]).flatten(), + (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(), + ] + weights = [ + ((1 - dh)[None].T * (1 - dw)[None]).flatten(), + ((1 - dh)[None].T * dw[None]).flatten(), + (dh[None].T * (1 - dw)[None]).flatten(), + (dh[None].T * dw[None]).flatten(), + ] + for i in range(4): + idx_list[i].extend(indices[i].tolist()) + weight_list[i].extend(weights[i].tolist()) + idx_tensor = torch.tensor(idx_list, dtype=torch.long, device=device) + weight_tensor = torch.tensor(weight_list, dtype=self.pos_embed.weight.dtype, device=device) + pos_embeds = self.pos_embed(idx_tensor).to(device) * weight_tensor[:, :, None] + patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3] + patch_pos_embeds = patch_pos_embeds.split([h * w for h, w in zip(grid_hs, grid_ws)]) + patch_pos_embeds_permute = [] + merge_size = self.spatial_merge_size + for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws): + pos_embed = pos_embed.repeat(t, 1) + pos_embed = ( + pos_embed.view(t, h // merge_size, merge_size, w // merge_size, merge_size, -1) + .permute(0, 1, 3, 2, 4, 5) + .flatten(0, 4) + ) + patch_pos_embeds_permute.append(pos_embed) + return torch.cat(patch_pos_embeds_permute) + + def forward( + self, + hidden_states: torch.Tensor, + grid_thw: torch.Tensor, + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + hidden_states = self.patch_embed(hidden_states) + pos_embeds = self.fast_pos_embed_interpolate(grid_thw) + hidden_states = hidden_states + pos_embeds + rotary_pos_emb = self.rot_pos_emb(grid_thw) + seq_len, _ = hidden_states.size() + hidden_states = hidden_states.reshape(seq_len, -1) + rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1) + emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + position_embeddings = (emb.cos(), emb.sin()) + cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( + dim=0, dtype=torch.int32 + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + deepstack_feature_lists = [] + for layer_num, blk in enumerate(self.blocks): + hidden_states = blk( + hidden_states, + cu_seqlens=cu_seqlens, + position_embeddings=position_embeddings, + ) + if layer_num in self.deepstack_visual_indexes: + deepstack_feature = self.deepstack_merger_list[self.deepstack_visual_indexes.index(layer_num)]( + hidden_states + ) + deepstack_feature_lists.append(deepstack_feature) + merged_hidden_states = self.merger(hidden_states) + return merged_hidden_states, deepstack_feature_lists + + +# --------------------------------------------------------------------------- +# Qwen3-VL text decoder layers (TP-sharded). +# --------------------------------------------------------------------------- + + +class MiniMaxH3Qwen3VLTextAttention(nn.Module): + def __init__(self, group: Any, config: Any, dtype: torch.dtype) -> None: + super().__init__() + self.group = group + self.config = config + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.scaling = self.head_dim**-0.5 + self.is_causal = True + + self.qkv_proj = MiniMaxH3Qwen3VLQKVParallelLinear( + group, + hidden_size=self.hidden_size, + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + dtype=dtype, + ) + self.o_proj = MiniMaxH3Qwen3VLRowParallelLinear( + group, + input_size=self.num_heads * self.head_dim, + output_size=self.hidden_size, + dtype=dtype, + input_is_parallel=True, + ) + self.q_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = MiniMaxH3Qwen3VLRMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + q_local = self.qkv_proj.local_num_heads * self.head_dim + kv_local = self.qkv_proj.local_num_kv_heads * self.head_dim + q_weight = self.qkv_proj.weight[0:q_local] + k_weight = self.qkv_proj.weight[q_local : q_local + kv_local] + v_weight = self.qkv_proj.weight[q_local + kv_local : q_local + 2 * kv_local] + # Three separate GEMMs keep the numerics identical to the reference + # path (packed QKV changes cuBLAS tiling and shifts bf16 rounding). + query_states = self.q_norm(F.linear(hidden_states, q_weight).view(hidden_shape)).transpose(1, 2) + key_states = self.k_norm(F.linear(hidden_states, k_weight).view(hidden_shape)).transpose(1, 2) + value_states = F.linear(hidden_states, v_weight).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = _apply_rotary_pos_emb(query_states, key_states, cos, sin) + + # Expand GQA keys/values to the local query-head count so SDPA works + # uniformly across backends (mirrors the reference eager path). + num_key_value_groups = self.num_heads // self.num_kv_heads + key_states = key_states.repeat_interleave(num_key_value_groups, dim=1) + value_states = value_states.repeat_interleave(num_key_value_groups, dim=1) + + attn_output = F.scaled_dot_product_attention( + query_states, + key_states, + value_states, + dropout_p=0.0, + is_causal=True, + ) + attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output + + +class MiniMaxH3Qwen3VLTextMLP(nn.Module): + def __init__(self, group: Any, config: Any, dtype: torch.dtype) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_up_proj = MiniMaxH3Qwen3VLMergedColumnParallelLinear( + group, + input_size=self.hidden_size, + intermediate_size=self.intermediate_size, + dtype=dtype, + ) + self.down_proj = MiniMaxH3Qwen3VLRowParallelLinear( + group, + input_size=self.intermediate_size, + output_size=self.hidden_size, + dtype=dtype, + input_is_parallel=True, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + ip_local = self.gate_up_proj.intermediate_size_per_partition + gate = F.linear(x, self.gate_up_proj.weight[0:ip_local]) + up = F.linear(x, self.gate_up_proj.weight[ip_local : 2 * ip_local]) + x = F.silu(gate) * up + return self.down_proj(x) + + +class MiniMaxH3Qwen3VLTextDecoderLayer(nn.Module): + def __init__(self, group: Any, config: Any, dtype: torch.dtype) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = MiniMaxH3Qwen3VLTextAttention(group, config, dtype) + self.mlp = MiniMaxH3Qwen3VLTextMLP(group, config, dtype) + self.input_layernorm = MiniMaxH3Qwen3VLRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = MiniMaxH3Qwen3VLRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, position_embeddings=position_embeddings) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +class MiniMaxH3Qwen3VLTextModel(nn.Module): + """TP-sharded Qwen3-VL text model returning unnormalized layer-50 states.""" + + def __init__(self, group: Any, config: Any, selected_layer: int, dtype: torch.dtype) -> None: + super().__init__() + self.group = group + self.config = config + self.selected_layer = selected_layer + self.vocab_size = config.vocab_size + self.hidden_size = config.hidden_size + self.num_layers = min(int(getattr(config, "num_hidden_layers", 0)), selected_layer) + self.embed_tokens = MiniMaxH3Qwen3VLVocabParallelEmbedding( + group, num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, dtype=dtype + ) + self.rotary_emb = MiniMaxH3Qwen3VLTextRotaryEmbedding(config) + self.layers = nn.ModuleList( + [MiniMaxH3Qwen3VLTextDecoderLayer(group, config, dtype) for _ in range(self.num_layers)] + ) + + def _deepstack_process( + self, + hidden_states: torch.Tensor, + visual_pos_masks: torch.Tensor, + visual_embeds: torch.Tensor, + ) -> torch.Tensor: + visual_pos_masks = visual_pos_masks.to(hidden_states.device) + visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype) + hidden_states = hidden_states.clone() + hidden_states[visual_pos_masks, :] = hidden_states[visual_pos_masks, :] + visual_embeds + return hidden_states + + def forward( + self, + inputs_embeds: torch.Tensor, + positions: torch.Tensor, + *, + visual_pos_masks: torch.Tensor | None = None, + deepstack_visual_embeds: list[torch.Tensor] | None = None, + ) -> torch.Tensor: + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, positions) + for layer_idx, decoder_layer in enumerate(self.layers): + hidden_states = decoder_layer(hidden_states, position_embeddings=position_embeddings) + if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)): + hidden_states = self._deepstack_process( + hidden_states, + visual_pos_masks, + deepstack_visual_embeds[layer_idx], + ) + return hidden_states + + +# --------------------------------------------------------------------------- +# Multimodal position / placeholder-mask helpers (ported from the reference). +# --------------------------------------------------------------------------- + + +def _get_vision_position_ids( + start_position: int, + grid_thw: list[int] | torch.Tensor, + temp_merge_size: int, + spatial_merge_size: int, + device: str | torch.device, +) -> torch.Tensor: + llm_grid_t, llm_grid_h, llm_grid_w = ( + grid_thw[0].item() // temp_merge_size, + grid_thw[1].item() // spatial_merge_size, + grid_thw[2].item() // spatial_merge_size, + ) + position_temporal = torch.arange(llm_grid_t, device=device) + position_width = torch.arange(llm_grid_w, device=device) + start_position + position_height = torch.arange(llm_grid_h, device=device) + start_position + position_width = position_width.repeat(llm_grid_h * llm_grid_t) + position_height = position_height.repeat_interleave(llm_grid_w).repeat(llm_grid_t) + position_temporal = position_temporal.repeat_interleave(llm_grid_h * llm_grid_w) + start_position + return torch.stack([position_temporal, position_height, position_width], dim=0) + + +def _get_rope_index( + input_ids: torch.Tensor, + mm_token_type_ids: torch.Tensor, + image_grid_thw: torch.Tensor | None, + video_grid_thw: torch.Tensor | None, + attention_mask: torch.Tensor | None, + spatial_merge_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if video_grid_thw is not None: + video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0) + video_grid_thw[:, 0] = 1 + + mrope_position_deltas = [] + position_ids = torch.zeros( + 3, + input_ids.shape[0], + input_ids.shape[1], + dtype=input_ids.dtype, + device=input_ids.device, + ) + grid_iters = { + 1: iter(image_grid_thw) if image_grid_thw is not None else None, + 2: iter(video_grid_thw) if video_grid_thw is not None else None, + } + + for batch_idx, current_input_ids in enumerate(input_ids): + input_token_type = mm_token_type_ids[batch_idx] + if attention_mask is not None: + current_input_ids = current_input_ids[attention_mask[batch_idx].bool()] + input_token_type = input_token_type[attention_mask[batch_idx].bool()] + + input_type_group = [] + for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]): + group = list(group) + start_index = group[0][0] + end_index = group[-1][0] + 1 + input_type_group.append((key, start_index, end_index)) + + current_pos = 0 + llm_pos_ids_list = [] + for modality_type, start_idx, end_idx in input_type_group: + if modality_type == 0: + text_len = end_idx - start_idx + llm_pos_ids_list.append( + torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos + ) + current_pos += text_len + else: + grid_thw = next(grid_iters[modality_type]) + vision_position_ids = _get_vision_position_ids( + current_pos, grid_thw, 1, spatial_merge_size, device=input_ids.device + ) + llm_pos_ids_list.append(vision_position_ids) + current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + if attention_mask is not None: + position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device) + else: + position_ids[:, batch_idx] = llm_positions.to(position_ids.device) + mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids)) + mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) + return position_ids, mrope_position_deltas + + +# --------------------------------------------------------------------------- +# Encoder entry point. +# --------------------------------------------------------------------------- + + +class MiniMaxH3Qwen3VLEncoder(nn.Module): + """Frozen, TP-capable Qwen3-VL backbone returning layer-50 states. + + ``encoder_group`` is a ``GroupCoordinator`` over the encoder tensor-parallel + ranks (by default the first ``text_encoder_tp_size`` DiT ranks). Only + ranks inside the group load weights; other ranks construct a parameter-free + stub that never runs the forward. + """ + + def __init__( + self, + model_path: str, + *, + device: torch.device, + load_model: bool, + encoder_group: Any | None = None, + ) -> None: + super().__init__() + self.device_target = device + self.encoder_group = encoder_group + self.image_token_id = 151655 + self.video_token_id = 151656 + self._tp_size = 1 + if not load_model: + return + + from transformers import Qwen3VLConfig + + config = Qwen3VLConfig.from_pretrained(model_path, trust_remote_code=False) + self.image_token_id = int(config.image_token_id) + self.video_token_id = int(config.video_token_id) + self._tp_size = int(encoder_group.world_size) if encoder_group is not None else 1 + dtype = torch.bfloat16 + self.vision = MiniMaxH3Qwen3VLVisionModel(config.vision_config) + self.vision.to(dtype=dtype) + self.text_model = MiniMaxH3Qwen3VLTextModel( + encoder_group, + config.text_config, + MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER, + dtype, + ) + self.text_model.to(dtype=dtype) + self._load_weights(model_path) + logger.info( + "MiniMax H3 Qwen3-VL encoder: %d retained decoder layers, text_encoder_tp_size=%d, vision replicated", + MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER, + self._tp_size, + ) + + @property + def is_loaded(self) -> bool: + return hasattr(self, "text_model") + + @property + def tp_size(self) -> int: + return self._tp_size + + def _map_weight_name(self, name: str) -> tuple[str, str | None] | None: + if name == "lm_head.weight" or name == "model.language_model.norm.weight": + return None + if name.startswith("model.visual."): + return ("vision." + name[len("model.visual.") :], None) + if name.startswith("model.language_model."): + rest = name[len("model.language_model.") :] + if rest.startswith("embed_tokens."): + return ("text_model." + rest, None) + match = re.match(r"layers\.(\d+)\.", rest) + if match is not None and int(match.group(1)) >= MINIMAX_H3_QWEN3VL_SELECTED_LM_LAYER: + return None + if ".self_attn.q_proj.weight" in rest: + return ( + "text_model." + rest.replace(".self_attn.q_proj.weight", ".self_attn.qkv_proj.weight"), + "q", + ) + if ".self_attn.k_proj.weight" in rest: + return ( + "text_model." + rest.replace(".self_attn.k_proj.weight", ".self_attn.qkv_proj.weight"), + "k", + ) + if ".self_attn.v_proj.weight" in rest: + return ( + "text_model." + rest.replace(".self_attn.v_proj.weight", ".self_attn.qkv_proj.weight"), + "v", + ) + if ".mlp.gate_proj.weight" in rest: + return ( + "text_model." + rest.replace(".mlp.gate_proj.weight", ".mlp.gate_up_proj.weight"), + 0, + ) + if ".mlp.up_proj.weight" in rest: + return ( + "text_model." + rest.replace(".mlp.up_proj.weight", ".mlp.gate_up_proj.weight"), + 1, + ) + return ("text_model." + rest, None) + return None + + def _load_weights(self, model_path: str) -> None: + import os + + import safetensors.torch + + index_path = os.path.join(model_path, "model.safetensors.index.json") + with open(index_path, encoding="utf-8") as handle: + index = json.load(handle) + weight_map: dict[str, str] = index["weight_map"] + files = sorted(set(weight_map.values())) + params = dict(self.named_parameters()) + loaded: set[str] = set() + for shard in files: + tensors = safetensors.torch.load_file(os.path.join(model_path, shard), device="cpu") + for name, tensor in tensors.items(): + mapped = self._map_weight_name(name) + if mapped is None: + continue + param_name, shard_id = mapped + param = params.get(param_name) + if param is None: + logger.warning("MiniMax H3 text encoder weight %s has no target parameter", name) + continue + weight_loader = getattr(param, "weight_loader", _default_weight_loader) + weight_loader(param, tensor, shard_id) + loaded.add(param_name) + missing = sorted(set(params) - loaded) + if missing: + logger.warning( + "MiniMax H3 text encoder weights not loaded (retained-layer subset): %d params, e.g. %s", + len(missing), + missing[:5], + ) + + def load_to_device(self) -> None: + if not self.is_loaded: + return + self.vision.to(self.device_target) + self.text_model.to(self.device_target) + + def offload_to_cpu(self) -> None: + if not self.is_loaded: + return + self.vision.to("cpu") + self.text_model.to("cpu") + torch.accelerator.empty_cache() + + def _encode( + self, + input_ids: torch.Tensor, + *, + pixel_values: torch.Tensor | None = None, + image_grid_thw: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + video_grid_thw: torch.Tensor | None = None, + ) -> torch.Tensor: + device = self.device_target + ids = input_ids.to(device, torch.long)[None] + attention_mask = torch.ones_like(ids) + mm_types = torch.zeros_like(ids, dtype=torch.int32) + mm_types[ids == self.image_token_id] = 1 + mm_types[ids == self.video_token_id] = 2 + + inputs_embeds = self.text_model.embed_tokens(ids) + image_mask = None + video_mask = None + deepstack_image_embeds: list[torch.Tensor] | None = None + deepstack_video_embeds: list[torch.Tensor] | None = None + + if pixel_values is not None: + image_embeds, deepstack_image_embeds = self.vision( + pixel_values.to(device, torch.bfloat16), + image_grid_thw.to(device, torch.long), + ) + image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + image_mask, _ = self._get_placeholder_mask(ids, inputs_embeds=inputs_embeds, image_features=image_embeds) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + if pixel_values_videos is not None: + video_embeds, deepstack_video_embeds = self.vision( + pixel_values_videos.to(device, torch.bfloat16), + video_grid_thw.to(device, torch.long), + ) + video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + _, video_mask = self._get_placeholder_mask(ids, inputs_embeds=inputs_embeds, video_features=video_embeds) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + visual_pos_masks, deepstack_visual_embeds = self._aggregate_deepstack( + image_mask, + video_mask, + deepstack_image_embeds, + deepstack_video_embeds, + ) + + positions, _ = _get_rope_index( + ids, + mm_token_type_ids=mm_types, + image_grid_thw=image_grid_thw.to(device, torch.long) if image_grid_thw is not None else None, + video_grid_thw=video_grid_thw.to(device, torch.long) if video_grid_thw is not None else None, + attention_mask=attention_mask, + spatial_merge_size=int(self.vision.config.spatial_merge_size), + ) + hidden = self.text_model( + inputs_embeds, + positions, + visual_pos_masks=visual_pos_masks, + deepstack_visual_embeds=deepstack_visual_embeds, + )[0] + expected = (int(ids.shape[1]), MINIMAX_H3_QWEN3VL_HIDDEN_DIM) + if tuple(hidden.shape) != expected: + raise ValueError(f"unexpected Qwen3-VL hidden shape {tuple(hidden.shape)}, expected {expected}") + return hidden.to(torch.bfloat16) + + def _get_placeholder_mask( + self, + input_ids: torch.Tensor, + *, + inputs_embeds: torch.Tensor, + image_features: torch.Tensor | None = None, + video_features: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + special_image_mask = input_ids == self.image_token_id + special_video_mask = input_ids == self.video_token_id + special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + if image_features is not None and special_image_mask.sum().item() != image_features.numel(): + raise ValueError( + "Image features and image tokens do not match: " + f"tokens: {special_image_mask.sum().item() // inputs_embeds.shape[-1]}, " + f"features: {image_features.shape[0]}" + ) + if video_features is not None and special_video_mask.sum().item() != video_features.numel(): + raise ValueError( + "Video features and video tokens do not match: " + f"tokens: {special_video_mask.sum().item() // inputs_embeds.shape[-1]}, " + f"features: {video_features.shape[0]}" + ) + return special_image_mask, special_video_mask + + def _aggregate_deepstack( + self, + image_mask: torch.Tensor | None, + video_mask: torch.Tensor | None, + deepstack_image_embeds: list[torch.Tensor] | None, + deepstack_video_embeds: list[torch.Tensor] | None, + ) -> tuple[torch.Tensor | None, list[torch.Tensor] | None]: + if image_mask is not None and video_mask is not None: + image_mask = image_mask[..., 0] + video_mask = video_mask[..., 0] + visual_pos_masks = image_mask | video_mask + deepstack_visual_embeds = [] + image_mask_joint = image_mask[visual_pos_masks] + video_mask_joint = video_mask[visual_pos_masks] + for img_embed, vid_embed in zip(deepstack_image_embeds, deepstack_video_embeds): + embed_joint = img_embed.new_zeros(visual_pos_masks.sum(), img_embed.shape[-1]).to(img_embed.device) + embed_joint[image_mask_joint, :] = img_embed + embed_joint[video_mask_joint, :] = vid_embed + deepstack_visual_embeds.append(embed_joint) + return visual_pos_masks, deepstack_visual_embeds + if image_mask is not None: + image_mask = image_mask[..., 0] + return image_mask, deepstack_image_embeds + if video_mask is not None: + video_mask = video_mask[..., 0] + return video_mask, deepstack_video_embeds + return None, None + + @torch.inference_mode() + def encode_ids( + self, + input_ids: torch.Tensor, + *, + pixel_values: torch.Tensor | None = None, + image_grid_thw: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + video_grid_thw: torch.Tensor | None = None, + ) -> torch.Tensor: + if not self.is_loaded: + raise RuntimeError("MiniMax H3 text encoder is only loaded on the encoder TP ranks") + if input_ids.ndim != 1: + raise ValueError(f"input_ids must be 1-D, got {list(input_ids.shape)}") + if (pixel_values is None) != (image_grid_thw is None): + raise ValueError("pixel_values and image_grid_thw must be provided together") + if (pixel_values_videos is None) != (video_grid_thw is None): + raise ValueError("pixel_values_videos and video_grid_thw must be provided together") + if next(self.parameters()).device.type != self.device_target.type: + raise RuntimeError("call load_to_device() before encode_ids()") + + torch.backends.cuda.enable_cudnn_sdp(True) + hidden = self._encode( + input_ids, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + pixel_values_videos=pixel_values_videos, + video_grid_thw=video_grid_thw, + ) + return hidden.cpu() + + def forward( + self, + input_ids: torch.Tensor, + *, + pixel_values: torch.Tensor | None = None, + image_grid_thw: torch.Tensor | None = None, + pixel_values_videos: torch.Tensor | None = None, + video_grid_thw: torch.Tensor | None = None, + ) -> torch.Tensor: + """Hook-compatible entry point for model-level CPU offloading.""" + kwargs = { + "pixel_values": pixel_values, + "image_grid_thw": image_grid_thw, + } + if pixel_values_videos is not None or video_grid_thw is not None: + kwargs.update( + pixel_values_videos=pixel_values_videos, + video_grid_thw=video_grid_thw, + ) + return self.encode_ids(input_ids, **kwargs) + + +__all__ = ["MiniMaxH3Qwen3VLEncoder"] diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py new file mode 100644 index 000000000..522ac2e8c --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -0,0 +1,1112 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 packed-token audio/video DiT for vLLM-Omni. + +vLLM tensor parallel linears and the unified attention layer provide TP and +Ulysses/Ring sequence parallel execution without changing the checkpoint +layout. +""" + +from __future__ import annotations + +import math +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn as nn +from cache_dit import ForwardPattern +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, + SequenceParallelOutput, +) + +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + ) + + from vllm_omni.diffusion.data import OmniDiffusionConfig + +logger = init_logger(__name__) + + +@dataclass +class MiniMaxH3DiTArchConfig: + num_layers: int = 50 + token_refiner_num_layers: int = 2 + hidden_size: int = 5376 + num_attention_heads: int = 56 + attention_head_dim: int = 128 + ffn_hidden_size: int = 14336 + latents_dim: int = 24 + audio_latents_dim: int = 32 + patch_size: tuple[int, int, int] = (1, 2, 2) + text_dim: int = 5120 + timestep_input_dim: int = 256 + time_embed_hidden_size: int = 5376 + time_embed_dim: int = 2688 + adaln_out_features: int = 18 * 5376 + final_adaln_out_features: int = 2 * 5376 + rope_inv_freq_len: int = 16 + norm_eps: float = 1e-5 + qk_norm_eps: float = 1e-5 + final_norm_eps: float = 1e-5 + + @classmethod + def from_mapping(cls, config: Mapping[str, Any]) -> MiniMaxH3DiTArchConfig: + fields = cls.__dataclass_fields__ + values = {name: config[name] for name in fields if name in config} + if "patch_size" in values: + values["patch_size"] = tuple(values["patch_size"]) + arch = cls(**values) + if len(arch.patch_size) != 3: + raise ValueError(f"patch_size must contain three values, got {arch.patch_size!r}") + return arch + + +_ARCH_DEFAULTS = MiniMaxH3DiTArchConfig() +_BF16_DTYPE = torch.bfloat16 +_FP32_DTYPE = torch.float32 + +MINIMAX_H3_FP32_PARAM_NAMES = frozenset( + { + "video_patch_proj.weight", + "video_patch_proj.bias", + "audio_patch_proj.weight", + "audio_patch_proj.bias", + "time_embedder.proj_in.weight", + "time_embedder.proj_in.bias", + "time_embedder.proj_out.weight", + "time_embedder.proj_out.bias", + "final_layer.video_out.weight", + "final_layer.video_out.bias", + "final_layer.audio_out.weight", + "final_layer.audio_out.bias", + } +) +MINIMAX_H3_FP32_BUFFER_NAMES = frozenset({"rope.inv_freq"}) + +# AdaLN modality count: token tags carry -1 for padding and 0/1/2 for +# video/text/audio tokens (padding is clamped to 0 before the embedding +# lookup and masked out afterwards). +MINIMAX_H3_ADALN_MODALITY_NUM = 3 + + +def _required_kwarg(kwargs: dict[str, Any], key: str) -> Any: + if key not in kwargs or kwargs[key] is None: + raise ValueError(f"MiniMaxH3DiTModel.forward requires kwarg {key!r}") + return kwargs[key] + + +# The exhaustive keyword contract of MiniMaxH3DiTModel.forward. Anything not +# listed here is rejected with a TypeError before any tensor work starts. +_FORWARD_SUPPORTED_KWARGS = frozenset( + { + "x", + "audio_x", + "img_position_ids", + "unique_timesteps", + "inverse_indices", + "update_mask", + "update_audio_mask", + "token_tags", + "skip_mask_out_condition", + "prompt_embeds", + "img_pos_info", + "audio_pos_info", + "text_pos_info", + "img_pos_for_infer_output_info", + "packed_seq_params", + "refiner_packed_seq_params", + } +) + + +def _reorder_grouped_qkv_to_qkv( + weight: torch.Tensor, + *, + num_query_groups: int, + heads_per_group: int, + head_dim: int, +) -> torch.Tensor: + per_group = (heads_per_group + 2) * head_dim + expected_out = num_query_groups * per_group + if weight.shape[0] != expected_out: + raise ValueError( + "qkv weight has incompatible output dim for grouped checkpoint layout: " + f"got {tuple(weight.shape)}, expected first dim {expected_out}." + ) + + rest_shape = weight.shape[1:] + grouped = weight.reshape(num_query_groups, per_group, *rest_shape) + q, k, v = torch.split( + grouped, + [heads_per_group * head_dim, head_dim, head_dim], + dim=1, + ) + return torch.cat( + [ + q.reshape(num_query_groups * heads_per_group * head_dim, *rest_shape), + k.reshape(num_query_groups * head_dim, *rest_shape), + v.reshape(num_query_groups * head_dim, *rest_shape), + ], + dim=0, + ) + + +def _norm(size: int, *, eps: float, dtype: torch.dtype = _BF16_DTYPE) -> nn.RMSNorm: + # RMSNorm uses fp32 accumulation with bf16 inputs and outputs. + # torch.nn.RMSNorm upcasts reduced-precision inputs for the variance + # reduction, matching that accumulation semantic. + return nn.RMSNorm(size, eps=eps, dtype=dtype) + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = torch.chunk(x, 2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +def _modulate_scale_shift( + x: torch.Tensor, + shift: torch.Tensor, + scale: torch.Tensor, + indices: torch.Tensor, + *, + dtype: torch.dtype, +) -> torch.Tensor: + # Apply per-index affine modulation: x * (1 + scale[idx]) + shift[idx]. + return (x * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(dtype) + + +def _modulate_gate( + x: torch.Tensor, + gate: torch.Tensor, + other: torch.Tensor, + indices: torch.Tensor, + *, + dtype: torch.dtype, +) -> torch.Tensor: + # Apply the per-index gated residual: x + gate[idx] * other. + return (x + gate.index_select(0, indices) * other).to(dtype) + + +class MiniMaxH3Rope(nn.Module): + """3D rope over (t, h, w); rotates 96 of 128 head dims (rotary_percent 0.75). + + Frequency layout concatenates temporal, height, and width embeddings twice, + with 16 frequencies per axis (inv_freq = base^-(arange(0,32,2)/32)). + """ + + def __init__(self, inv_freq_len: int) -> None: + super().__init__() + self.register_buffer( + "inv_freq", + torch.empty(inv_freq_len, dtype=_FP32_DTYPE), + persistent=True, + ) + + def forward(self, img_position_ids: torch.Tensor) -> torch.Tensor: + """img_position_ids: [1, S, 3] (t, h, w) -> freqs [S, rot_dim=96].""" + if img_position_ids.dim() != 3 or img_position_ids.shape[0] != 1: + raise ValueError(f"img_position_ids must be [1, S, 3], got {list(img_position_ids.shape)}") + pos = img_position_ids[0].to(_FP32_DTYPE) # [S, 3] + per_axis = pos.unsqueeze(-1) * self.inv_freq.view(1, 1, -1) # [S, 3, 16] + t_f, h_f, w_f = per_axis.unbind(dim=1) # each [S, 16] + half = torch.cat((t_f, h_f, w_f), dim=-1) # [S, 48] + return torch.cat((half, half), dim=-1) # [S, 96] + + +def _apply_rope(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + """Rotate the first rot_dim head dims; pass the rest through. + + x: [T, heads, head_dim]; freqs: [T, rot_dim]. In the unfused path, cos/sin + are cast to the activation dtype before the elementwise math. + """ + rot_dim = freqs.shape[-1] + x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:] + cos = torch.cos(freqs).to(x.dtype).unsqueeze(1) # [T, 1, rot_dim] + sin = torch.sin(freqs).to(x.dtype).unsqueeze(1) + x_rot = (x_rot * cos) + (_rotate_half(x_rot) * sin) + return torch.cat((x_rot, x_pass), dim=-1) + + +class MiniMaxH3TimeEmbedder(nn.Module): + def __init__( + self, + arch: MiniMaxH3DiTArchConfig, + quant_config: QuantizationConfig | None, + ) -> None: + super().__init__() + self.frequency_embedding_size = arch.timestep_input_dim + self.proj_in = ColumnParallelLinear( + arch.timestep_input_dim, + arch.time_embed_hidden_size, + bias=True, + gather_output=True, + params_dtype=_FP32_DTYPE, + quant_config=quant_config, + ) + self.proj_out = RowParallelLinear( + arch.time_embed_hidden_size, + arch.time_embed_dim, + bias=True, + input_is_parallel=False, + params_dtype=_FP32_DTYPE, + quant_config=quant_config, + ) + + def forward(self, t: torch.Tensor) -> torch.Tensor: + """t: [M] -> [M, time_embed_dim] fp32. + + The sinusoidal embedding stays fp32 throughout and concatenates cosine + values before sine values. + """ + half = self.frequency_embedding_size // 2 + freqs = torch.exp(-math.log(10000.0) * torch.arange(half, dtype=_FP32_DTYPE, device=t.device) / half) + args = t.to(_FP32_DTYPE)[:, None] * freqs[None] + t_freq = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + hidden, _ = self.proj_in(t_freq) + hidden = nn.functional.silu(hidden) + out, _ = self.proj_out(hidden) + return out + + +def _sdpa_varlen_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + """Segment-wise SDPA equivalent of the non-causal varlen FA call. + + Mirrors the generic attention layer's semantics: FA is the fast path, + SDPA is the correctness fallback when the platform resolves another + backend. Segments are delimited by ``cu_seqlens`` exactly like the + varlen kernel, so attention never crosses packed-document boundaries. + """ + out = torch.empty_like(q) + bounds = cu_seqlens.tolist() + for start, stop in zip(bounds[:-1], bounds[1:]): + if stop == start: + continue + seg_q = q[start:stop].transpose(0, 1).unsqueeze(0) + seg_k = k[start:stop].transpose(0, 1).unsqueeze(0) + seg_v = v[start:stop].transpose(0, 1).unsqueeze(0) + seg_out = torch.nn.functional.scaled_dot_product_attention( + seg_q, + seg_k, + seg_v, + scale=softmax_scale, + ) + out[start:stop] = seg_out.squeeze(0).transpose(0, 1) + return out + + +class MiniMaxH3Attention(nn.Module): + def __init__( + self, + arch: MiniMaxH3DiTArchConfig, + quant_config: QuantizationConfig | None, + *, + skip_sequence_parallel: bool = False, + ) -> None: + super().__init__() + self.total_num_heads = arch.num_attention_heads + self.head_dim = arch.attention_head_dim + inner_dim = self.total_num_heads * self.head_dim + self.softmax_scale = self.head_dim**-0.5 + self.qkv_proj = QKVParallelLinear( + hidden_size=arch.hidden_size, + head_size=self.head_dim, + total_num_heads=self.total_num_heads, + total_num_kv_heads=self.total_num_heads, + bias=False, + params_dtype=_BF16_DTYPE, + quant_config=quant_config, + return_bias=True, + ) + self.num_heads = self.qkv_proj.num_heads + self.num_kv_heads = self.qkv_proj.num_kv_heads + self._install_qkv_weight_loader(arch) + self.q_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps) + self.k_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps) + self.out_proj = RowParallelLinear( + inner_dim, + arch.hidden_size, + bias=False, + input_is_parallel=True, + params_dtype=_BF16_DTYPE, + quant_config=quant_config, + ) + self.attention = Attention( + num_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_size=self.head_dim, + softmax_scale=self.softmax_scale, + causal=False, + skip_sequence_parallel=skip_sequence_parallel, + ) + + def _install_qkv_weight_loader(self, arch: MiniMaxH3DiTArchConfig) -> None: + base_loader = self.qkv_proj.weight.weight_loader + + def _weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + # The grouped checkpoint layout is + # [num_query_groups, q_per_group + k + v] before splitting. + # MiniMax H3 uses MHA, so checkpoint rows are per-head [q, k, v], + # while qkv_proj expects [q_all, k_all, v_all]. + reordered = _reorder_grouped_qkv_to_qkv( + loaded_weight, + num_query_groups=arch.num_attention_heads, + heads_per_group=1, + head_dim=arch.attention_head_dim, + ) + base_loader(param, reordered) + + self.qkv_proj.weight.weight_loader = _weight_loader + + @torch.compiler.disable + def _run_packed_attention( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + cu_seqlens: torch.Tensor, + max_seqlen: int, + ) -> torch.Tensor: + """Run packed attention as a small eager island. + + The scalar packed-layout metadata and the CuTe FlashAttention-4 DSL + are intentionally opaque to Dynamo. Keeping this boundary narrow lets + regional compile fuse projections, norms, RoPE, and the surrounding + DiT block without repeatedly graph-breaking inside the FA4 compiler. + """ + used = int(cu_seqlens[1].item()) + packed_total = int(cu_seqlens[-1].item()) + attn_mask = None + if used < packed_total: + attn_mask = torch.arange(packed_total, device=q.device)[None] < used + metadata = AttentionMetadata( + attn_mask=attn_mask, + extra={ + "cu_seqlens_q": cu_seqlens, + "cu_seqlens_k": cu_seqlens, + "max_seqlen_q": max_seqlen, + "max_seqlen_k": max_seqlen, + }, + ) + return self.attention( + q.unsqueeze(0), + k.unsqueeze(0), + v.unsqueeze(0), + metadata, + ).squeeze(0) + + def forward( + self, + x: torch.Tensor, + *, + rope_freqs: torch.Tensor | None, + cu_seqlens: torch.Tensor, + max_seqlen: int, + sp_seq_lens: list[int] | None = None, + ) -> torch.Tensor: + """x: [T, hidden] packed thd rows -> [T, hidden]. + + Operation order: fused qkv projection -> per-head q/k RMSNorm -> RoPE + on q/k -> variable-length non-causal flash attention -> output projection. + + With Ulysses sequence parallelism, x holds this rank's row shard; + qkv/norm/RoPE run locally, an all-to-all trades sequence for heads. + Each rank attends the full sequence with heads/world_size local heads, + so cu_seqlens retains global packed-document semantics. The inverse + all-to-all restores the row shard before the output projection. + """ + total = x.shape[0] + qkv, _ = self.qkv_proj(x) + q_size = self.num_heads * self.head_dim + kv_size = self.num_kv_heads * self.head_dim + q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1) + q = q.view(total, self.num_heads, self.head_dim) + k = k.view(total, self.num_kv_heads, self.head_dim) + v = v.view(total, self.num_kv_heads, self.head_dim) + q = self.q_norm(q) + k = self.k_norm(k) + if rope_freqs is not None: + q = _apply_rope(q, rope_freqs) + k = _apply_rope(k, rope_freqs) + + # The packed layout uses a second document for alignment padding. + # Local/Ulysses backends unpad it, while Ring keeps aligned rows for + # fixed-size P2P buffers. + out = self._run_packed_attention( + q, + k, + v, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + out = out.reshape(total, self.num_heads * self.head_dim) + out, _ = self.out_proj(out) + return out + + +class MiniMaxH3MLP(nn.Module): + def __init__( + self, + arch: MiniMaxH3DiTArchConfig, + quant_config: QuantizationConfig | None, + ) -> None: + super().__init__() + self.fc1 = MergedColumnParallelLinear( + arch.hidden_size, + [arch.ffn_hidden_size, arch.ffn_hidden_size], + bias=False, + gather_output=False, + params_dtype=_BF16_DTYPE, + quant_config=quant_config, + ) + self._install_fc1_weight_loader() + # Chunk the fused fc1 output as [gate, up], then compute + # silu(gate) * up. + self.fc2 = RowParallelLinear( + arch.ffn_hidden_size, + arch.hidden_size, + bias=False, + input_is_parallel=True, + params_dtype=_BF16_DTYPE, + quant_config=quant_config, + ) + + def _install_fc1_weight_loader(self) -> None: + base_loader = self.fc1.weight.weight_loader + + def _weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + if loaded_weight.shape[0] % 2: + raise ValueError( + "MiniMax H3 fc1 checkpoint rows must split evenly into " + f"gate/up matrices, got {tuple(loaded_weight.shape)}" + ) + gate, up = loaded_weight.chunk(2, dim=0) + base_loader(param, gate, 0) + base_loader(param, up, 1) + + self.fc1.weight.weight_loader = _weight_loader + + def forward(self, x: torch.Tensor) -> torch.Tensor: + hidden, _ = self.fc1(x) + gate, up = hidden.chunk(2, dim=-1) + hidden = nn.functional.silu(gate) * up + out, _ = self.fc2(hidden) + return out + + +class MiniMaxH3AdalnProj(nn.Module): + """SiLU + zero-init linear over unique condition embeddings. + + Per block, three modalities each produce six H-wide vectors: + [M, t_dim] -> [M, 3*6H] -> view(M*3, 6H) -> chunk(6). + The final layer uses one modality and produces two H-wide vectors: + [M, t_dim] -> [M, 2H] -> chunk(2). + """ + + def __init__( + self, + arch: MiniMaxH3DiTArchConfig, + out_features: int, + quant_config: QuantizationConfig | None, + *, + expand_ratio: int, + modality_num: int, + ) -> None: + super().__init__() + if out_features != expand_ratio * arch.hidden_size * modality_num: + raise ValueError( + f"adaln out_features mismatch: {out_features} != {expand_ratio}*{arch.hidden_size}*{modality_num}" + ) + self.expand_ratio = expand_ratio + self.modality_num = modality_num + self.hidden_size = arch.hidden_size + self.linear = ColumnParallelLinear( + arch.time_embed_dim, + out_features, + bias=True, + gather_output=True, + params_dtype=_BF16_DTYPE, + quant_config=quant_config, + ) + + def forward(self, t_emb: torch.Tensor) -> tuple[torch.Tensor, ...]: + """t_emb: [M, t_dim] -> expand_ratio tensors of [M*modality_num, H].""" + x = nn.functional.silu(t_emb) + x, _ = self.linear(x.to(self.linear.weight.dtype)) + m = x.shape[0] + x = x.view(m * self.modality_num, self.expand_ratio * self.hidden_size) + return tuple(x.chunk(self.expand_ratio, dim=-1)) + + +class MiniMaxH3TokenRefinerBlock(nn.Module): + """Standard pre-norm transformer block without AdaLN or RoPE.""" + + def __init__( + self, + arch: MiniMaxH3DiTArchConfig, + quant_config: QuantizationConfig | None, + ) -> None: + super().__init__() + self.norm1 = _norm(arch.hidden_size, eps=arch.norm_eps) + self.norm2 = _norm(arch.hidden_size, eps=arch.norm_eps) + # Text refinement runs on replicated rows before ``sp_prepare``. + # Applying Ulysses here would all-to-all an unsharded sequence while + # retaining the original packed ``cu_seqlens`` metadata. + self.attn = MiniMaxH3Attention( + arch, + quant_config, + skip_sequence_parallel=True, + ) + self.mlp = MiniMaxH3MLP(arch, quant_config) + + def forward( + self, + x: torch.Tensor, + *, + cu_seqlens: torch.Tensor, + max_seqlen: int, + ) -> torch.Tensor: + x = x + self.attn( + self.norm1(x), + rope_freqs=None, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + x = x + self.mlp(self.norm2(x)) + return x + + +class MiniMaxH3TokenRefiner(nn.Module): + def __init__( + self, + arch: MiniMaxH3DiTArchConfig, + quant_config: QuantizationConfig | None, + ) -> None: + super().__init__() + self.blocks = nn.ModuleList( + [MiniMaxH3TokenRefinerBlock(arch, quant_config) for _ in range(arch.token_refiner_num_layers)] + ) + self.final_norm = _norm(arch.hidden_size, eps=arch.final_norm_eps) + + def forward( + self, + x: torch.Tensor, + *, + cu_seqlens: torch.Tensor, + max_seqlen: int, + ) -> torch.Tensor: + for block in self.blocks: + x = block(x, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen) + return self.final_norm(x) + + +class MiniMaxH3DiTBlock(nn.Module): + def __init__( + self, + arch: MiniMaxH3DiTArchConfig, + quant_config: QuantizationConfig | None, + ) -> None: + super().__init__() + self.norm1 = _norm(arch.hidden_size, eps=arch.norm_eps) + self.norm2 = _norm(arch.hidden_size, eps=arch.norm_eps) + self.attn = MiniMaxH3Attention(arch, quant_config) + self.mlp = MiniMaxH3MLP(arch, quant_config) + self.adaln_proj = MiniMaxH3AdalnProj( + arch, + arch.adaln_out_features, + quant_config, + expand_ratio=6, + modality_num=MINIMAX_H3_ADALN_MODALITY_NUM, + ) + + def forward( + self, + x: torch.Tensor, + *, + t_emb: torch.Tensor, + combined_indices: torch.Tensor, + rope_freqs: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + sp_seq_lens: list[int] | None = None, + ) -> torch.Tensor: + """x: [T, H]; t_emb: [M, t_dim]; combined_indices: [T] + (= inverse_indices * modality_num + token_tags.clamp(min=0)). + + Each block computes AdaLN parameters once, then applies + norm1 -> scale/shift -> attention -> gated residual, followed by + norm2 -> scale/shift -> MLP -> gated residual. + """ + ( + shift_msa, + scale_msa, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + ) = self.adaln_proj(t_emb) + + residual = x + h = self.norm1(x) + h = _modulate_scale_shift(h, shift_msa, scale_msa, combined_indices, dtype=_BF16_DTYPE) + h = self.attn( + h, + rope_freqs=rope_freqs, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + sp_seq_lens=sp_seq_lens, + ) + x = _modulate_gate(residual, gate_msa, h, combined_indices, dtype=_BF16_DTYPE) + + residual = x + h = self.norm2(x) + h = _modulate_scale_shift(h, shift_mlp, scale_mlp, combined_indices, dtype=_BF16_DTYPE) + h = self.mlp(h) + return _modulate_gate(residual, gate_mlp, h, combined_indices, dtype=_BF16_DTYPE) + + +class MiniMaxH3FinalLayer(nn.Module): + def __init__( + self, + arch: MiniMaxH3DiTArchConfig, + quant_config: QuantizationConfig | None, + ) -> None: + super().__init__() + video_patch_dim = arch.latents_dim * arch.patch_size[0] * arch.patch_size[1] * arch.patch_size[2] + self.norm = _norm(arch.hidden_size, eps=arch.final_norm_eps) + self.adaln_proj = MiniMaxH3AdalnProj( + arch, + arch.final_adaln_out_features, + quant_config, + expand_ratio=2, + modality_num=1, + ) + self.video_out = ColumnParallelLinear( + arch.hidden_size, + video_patch_dim, + bias=True, + gather_output=True, + params_dtype=_FP32_DTYPE, + quant_config=quant_config, + ) + self.audio_out = ColumnParallelLinear( + arch.hidden_size, + arch.audio_latents_dim, + bias=True, + gather_output=True, + params_dtype=_FP32_DTYPE, + quant_config=quant_config, + ) + + def forward( + self, + x: torch.Tensor, + *, + t_emb: torch.Tensor, + inverse_indices: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """x: [T, H] -> (video_logits [T, 96] fp32, audio_logits [T, 32] fp32). + + Apply single-modality shift/scale AdaLN to the final normalized + activations, cast to fp32, then apply both output heads to all rows. + """ + shift, scale = self.adaln_proj(t_emb) + h = self.norm(x) + h = _modulate_scale_shift(h, shift, scale, inverse_indices, dtype=_BF16_DTYPE) + # Preserve full precision through both final output projections. + h = h.to(_FP32_DTYPE) + video, _ = self.video_out(h) + audio, _ = self.audio_out(h) + return video, audio + + +class MiniMaxH3SPPrepare(nn.Module): + """Explicit boundary for sharding packed rows and their metadata together.""" + + def forward( + self, + hidden_states: torch.Tensor, + rope_freqs: torch.Tensor, + combined_indices: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return hidden_states, rope_freqs, combined_indices + + +class MiniMaxH3SPGather(nn.Module): + """Explicit boundary for restoring packed rows after the block stack.""" + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + + +class MiniMaxH3DiTModel(nn.Module): + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={"blocks": ForwardPattern.Pattern_3}, + # H3 is CFG-distilled and performs one transformer forward per step. + has_separate_cfg=False, + check_forward_pattern=False, + ) + _repeated_blocks = ["MiniMaxH3DiTBlock"] + _layerwise_offload_blocks_attrs = ["blocks"] + + @staticmethod + def _is_transformer_block(name: str, module: nn.Module) -> bool: + del module + parts = name.split(".") + return len(parts) == 2 and parts[0] == "blocks" and parts[1].isdigit() + + _hsdp_shard_conditions = [_is_transformer_block] + _hsdp_ignored_modules = [ + "video_patch_proj", + "audio_patch_proj", + "time_embedder", + "final_layer", + ] + _sp_plan = { + "sp_prepare": { + 0: SequenceParallelInput( + split_dim=0, + expected_dims=2, + split_output=True, + ), + 1: SequenceParallelInput( + split_dim=0, + expected_dims=2, + split_output=True, + ), + 2: SequenceParallelInput( + split_dim=0, + expected_dims=1, + split_output=True, + ), + }, + "sp_gather": SequenceParallelOutput(gather_dim=0, expected_dims=2), + } + packed_modules_mapping = {} + + def _validate_tp_config(self, *, arch: MiniMaxH3DiTArchConfig, tp_size: int) -> None: + if tp_size < 1: + raise ValueError(f"tensor_parallel_size must be positive, got {tp_size}") + if arch.num_attention_heads % tp_size: + raise ValueError( + "num_attention_heads must be divisible by tensor_parallel_size: " + f"{arch.num_attention_heads} % {tp_size} != 0" + ) + if arch.ffn_hidden_size % tp_size: + raise ValueError( + f"ffn_hidden_size must be divisible by tensor_parallel_size: {arch.ffn_hidden_size} % {tp_size} != 0" + ) + if arch.num_attention_heads <= 0: + raise ValueError("num_attention_heads must be positive.") + if arch.hidden_size <= 0: + raise ValueError("hidden_size must be positive.") + if arch.attention_head_dim <= 0: + raise ValueError("attention_head_dim must be positive.") + if arch.ffn_hidden_size <= 0: + raise ValueError("ffn_hidden_size must be positive.") + + def __init__( + self, + od_config: OmniDiffusionConfig, + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__() + tf_config = od_config.tf_model_config + config_mapping = tf_config.to_dict() if hasattr(tf_config, "to_dict") else dict(tf_config) + arch = MiniMaxH3DiTArchConfig.from_mapping(config_mapping) + self.arch = arch + self.od_config = od_config + self.parallel_config = od_config.parallel_config + self.hidden_size = arch.hidden_size + self.num_attention_heads = arch.num_attention_heads + self.num_channels_latents = arch.latents_dim + self._validate_tp_config( + arch=arch, + tp_size=get_tensor_model_parallel_world_size(), + ) + local_heads = arch.num_attention_heads // get_tensor_model_parallel_world_size() + ulysses_degree = int(self.parallel_config.ulysses_degree) + if local_heads % ulysses_degree: + raise ValueError( + "MiniMax H3 local attention heads must be divisible by " + "ulysses_degree: " + f"({arch.num_attention_heads} / " + f"{get_tensor_model_parallel_world_size()}) % " + f"{ulysses_degree} != 0" + ) + + self.video_patch_proj = ColumnParallelLinear( + arch.latents_dim * arch.patch_size[0] * arch.patch_size[1] * arch.patch_size[2], + arch.hidden_size, + bias=True, + gather_output=True, + params_dtype=_FP32_DTYPE, + quant_config=quant_config, + ) + self.audio_patch_proj = ColumnParallelLinear( + arch.audio_latents_dim, + arch.hidden_size, + bias=True, + gather_output=True, + params_dtype=_FP32_DTYPE, + quant_config=quant_config, + ) + self.condition_proj = ColumnParallelLinear( + arch.text_dim, + arch.hidden_size, + bias=True, + gather_output=True, + params_dtype=_BF16_DTYPE, + quant_config=quant_config, + ) + self.time_embedder = MiniMaxH3TimeEmbedder(arch, quant_config) + self.rope = MiniMaxH3Rope(arch.rope_inv_freq_len) + self.token_refiner = MiniMaxH3TokenRefiner(arch, quant_config) + self.blocks = nn.ModuleList([MiniMaxH3DiTBlock(arch, quant_config) for _ in range(arch.num_layers)]) + self.sp_prepare = MiniMaxH3SPPrepare() + self.sp_gather = MiniMaxH3SPGather() + self.final_layer = MiniMaxH3FinalLayer(arch, quant_config) + self._mark_missing_params_required() + + def _mark_missing_params_required(self) -> None: + for _, param in self.named_parameters(): + param.missing_param_init = "error" + + def post_load_weights(self) -> None: + for name, param in self.named_parameters(): + if name in MINIMAX_H3_FP32_PARAM_NAMES and param.dtype != _FP32_DTYPE: + raise ValueError(f"{name} must stay fp32 after load, got {param.dtype}.") + for name, buffer in self.named_buffers(): + if name in MINIMAX_H3_FP32_BUFFER_NAMES and buffer.dtype != _FP32_DTYPE: + raise ValueError(f"{name} must stay fp32 after load, got {buffer.dtype}.") + + def load_weights( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> set[str]: + """Load exact H3 checkpoint names with logical TP-aware loaders.""" + params = dict(self.named_parameters()) + params.update(dict(self.named_buffers())) + loaded: set[str] = set() + for name, loaded_weight in weights: + param = params.get(name) + if param is None: + logger.warning("Skipping MiniMax H3 weight not present in model: %s", name) + continue + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded.add(name) + return loaded + + @staticmethod + def _pos_ids(pos_info: Any, key: str) -> torch.Tensor: + if isinstance(pos_info, dict): + ids = pos_info.get("position_ids") + else: + ids = getattr(pos_info, "position_ids", None) + if ids is None: + raise ValueError(f"{key}.position_ids is required") + return ids.view(-1).to(torch.long) + + @staticmethod + def _psp_field(psp: Any, key: str, field: str) -> Any: + if isinstance(psp, dict): + value = psp.get(field) + else: + value = getattr(psp, field, None) + if value is None: + raise ValueError(f"{key}.{field} is required") + return value + + def _embed( + self, + *, + x: torch.Tensor, + audio_x: torch.Tensor, + text_embeddings_selected: torch.Tensor, + unique_timesteps: torch.Tensor, + img_pos: torch.Tensor, + audio_pos: torch.Tensor, + text_pos: torch.Tensor, + refiner_cu_seqlens: torch.Tensor, + refiner_max_seqlen: int, + seq_len: int, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Build packed multimodal embeddings for the TP=1/SP=1 inference path. + + Returns (decoder_input [S, H] bf16, t_emb [M, t_dim] fp32). + """ + # Latent embedders stay fp32 in and out; their outputs are cast to the + # bf16 sequence dtype only during indexed scattering. + x_rows = x.view(-1, x.shape[-1]).index_select(0, img_pos).to(_FP32_DTYPE) + video_embed, _ = self.video_patch_proj(x_rows) + audio_rows = audio_x.view(-1, audio_x.shape[-1]).index_select(0, audio_pos).to(_FP32_DTYPE) + audio_embed, _ = self.audio_patch_proj(audio_rows) + + text_rows = text_embeddings_selected.to(device=device, dtype=_BF16_DTYPE) + text_embed, _ = self.condition_proj(text_rows) + text_embed = self.token_refiner( + text_embed, + cu_seqlens=refiner_cu_seqlens, + max_seqlen=refiner_max_seqlen, + ) + + embeddings = torch.zeros((seq_len, self.hidden_size), device=device, dtype=_BF16_DTYPE) + embeddings.index_add_(0, text_pos, text_embed.to(_BF16_DTYPE)[: text_pos.shape[0]]) + embeddings.index_add_(0, img_pos, video_embed.to(_BF16_DTYPE)[: img_pos.shape[0]]) + embeddings.index_add_(0, audio_pos, audio_embed.to(_BF16_DTYPE)[: audio_pos.shape[0]]) + + t_emb = self.time_embedder(unique_timesteps) + return embeddings, t_emb + + def forward(self, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]: + """Packed inference forward. + + Keyword names follow the checkpoint's serving contract. + Returns `(video_logits, audio_logits)` from rows selected by + `img_pos_for_infer_output_info` and `audio_pos_info`, with condition + rows zeroed by update masks. + """ + # Strict keyword contract: refuse any kwarg forward does not consume. + unexpected = sorted(set(kwargs) - _FORWARD_SUPPORTED_KWARGS) + if unexpected: + raise TypeError( + "MiniMaxH3DiTModel.forward received unexpected kwargs: " + f"{unexpected}; supported kwargs: " + f"{sorted(_FORWARD_SUPPORTED_KWARGS)}" + ) + + x = _required_kwarg(kwargs, "x") + audio_x = _required_kwarg(kwargs, "audio_x") + img_position_ids = _required_kwarg(kwargs, "img_position_ids") + unique_timesteps = _required_kwarg(kwargs, "unique_timesteps") + inverse_indices = _required_kwarg(kwargs, "inverse_indices").view(-1).to(torch.long) + update_mask = _required_kwarg(kwargs, "update_mask") + token_tags = _required_kwarg(kwargs, "token_tags").view(-1).to(torch.long) + skip_mask_out_condition = bool(kwargs.get("skip_mask_out_condition", False)) + + text_selected = _required_kwarg(kwargs, "prompt_embeds") + + img_pos = self._pos_ids(_required_kwarg(kwargs, "img_pos_info"), "img_pos_info") + audio_pos = self._pos_ids(_required_kwarg(kwargs, "audio_pos_info"), "audio_pos_info") + text_pos = self._pos_ids( + _required_kwarg(kwargs, "text_pos_info"), + "text_pos_info", + ) + infer_out_pos = self._pos_ids( + _required_kwarg(kwargs, "img_pos_for_infer_output_info"), + "img_pos_for_infer_output_info", + ) + + psp = _required_kwarg(kwargs, "packed_seq_params") + cu_seqlens = self._psp_field(psp, "packed_seq_params", "cu_seqlens_q").to(torch.int32) + max_seqlen = int(self._psp_field(psp, "packed_seq_params", "max_seqlen_q")) + refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params") + refiner_cu = self._psp_field(refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q").to(torch.int32) + refiner_max = int(self._psp_field(refiner_psp, "refiner_packed_seq_params", "max_seqlen_q")) + + if x.dim() != 3 or x.shape[0] != 1: + raise ValueError(f"x must be [1, S, C], got {list(x.shape)}") + seq_len = int(x.shape[1]) + if token_tags.shape[0] != seq_len: + raise ValueError(f"token_tags must cover the full packed sequence ({seq_len}), got {token_tags.shape[0]}.") + if inverse_indices.shape[0] != seq_len: + raise ValueError(f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}") + device = x.device + # Compute RoPE frequencies over the full packed sequence. + rope_freqs = self.rope(img_position_ids).to(device) + + decoder_input, t_emb = self._embed( + x=x, + audio_x=audio_x, + text_embeddings_selected=text_selected, + unique_timesteps=unique_timesteps.view(-1).to(device), + img_pos=img_pos.to(device), + audio_pos=audio_pos.to(device), + text_pos=text_pos.to(device), + refiner_cu_seqlens=refiner_cu.to(device), + refiner_max_seqlen=refiner_max, + seq_len=seq_len, + device=device, + ) + + combined_indices = (inverse_indices * MINIMAX_H3_ADALN_MODALITY_NUM + token_tags.clamp(min=0)).to(device) + inverse_indices = inverse_indices.to(device) + + hidden = decoder_input + cu_seqlens = cu_seqlens.to(device) + block_rope = rope_freqs + block_combined = combined_indices + + hidden, block_rope, block_combined = self.sp_prepare( + hidden, + block_rope, + block_combined, + ) + for block in self.blocks: + hidden = block( + hidden, + t_emb=t_emb, + combined_indices=block_combined, + rope_freqs=block_rope, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ) + hidden = self.sp_gather(hidden) + + video_logits, audio_logits = self.final_layer( + hidden, + t_emb=t_emb, + inverse_indices=inverse_indices, + ) + + # Select target and condition rows at inference-output positions, then + # zero the condition rows. + video_logits = video_logits.index_select(0, infer_out_pos.to(device)) + audio_logits = audio_logits.index_select(0, audio_pos.to(device)) + if not skip_mask_out_condition: + update_mask = update_mask.view(-1).to(device) + if update_mask.shape[0] != video_logits.shape[0]: + raise ValueError(f"update_mask length mismatch: {update_mask.shape[0]} != {video_logits.shape[0]}") + video_logits = video_logits * update_mask.unsqueeze(-1) + # Audio has no condition rows in the supported tasks, so its + # derived update mask is all ones. Honor an explicit mask when + # provided. + update_audio_mask = kwargs.get("update_audio_mask") + if update_audio_mask is not None: + audio_logits = audio_logits * update_audio_mask.view(-1).unsqueeze(-1) + return video_logits, audio_logits + + +EntryClass = MiniMaxH3DiTModel + +__all__ = [ + "MINIMAX_H3_FP32_BUFFER_NAMES", + "MINIMAX_H3_FP32_PARAM_NAMES", + "MiniMaxH3DiTModel", + "_reorder_grouped_qkv_to_qkv", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py b/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py new file mode 100644 index 000000000..37353b459 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py @@ -0,0 +1,572 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 packed-sequence materialization for FL2VA and Ref2VA tasks. + +Layout: [text L | imgvid_cond C | audio A(=t*2ch) | video_target V | pad P]. +Builder rules: +- placeholder ids; block-derived pos infos/masks/cu_seqlens/document_id +- img_position_ids fp64 grid: text rows (row_idx,0,0); video/cond t counter + continues text_len with temporal interp spans (frame_rescale 5/3 x + frame_per_token (1,4,4,4,4)); each spatial sqrt_area axis uses evenly spaced + coordinates excluding the right endpoint, then scales them by INTERP; + audio channel-major blocks pinned to the w-grid extremes. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence + +import numpy as np +import torch + +MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES: tuple[tuple[int, ...], ...] = ( + (0,), + (-1,), + (0, -1), +) + +MINIMAX_H3_TEXT_ID = -5 +MINIMAX_H3_IMGVID_COND_ID = -11 +MINIMAX_H3_AUDIO_REF_COND_ID = -17 +MINIMAX_H3_AUDIO_FIRST_ID = -15 +MINIMAX_H3_AUDIO_ID = -14 +MINIMAX_H3_VIDEO_FIRST_ID = -3 +MINIMAX_H3_VIDEO_ID = -2 +MINIMAX_H3_VIDEO_LAST_ID = -4 +MINIMAX_H3_PAD_ID = -1 + +_INTERP = 32 +_T_GROUP = 5 +_FRAME_PER_TOKEN = (1, 4, 4, 4, 4) +_FRAME_RESCALE = 5.0 / 3.0 +_SEQ_ALIGN = 64 +# MiniMax H3 packs video latents with a fixed [1, 2, 2] (t, h, w) patch. +_PATCH_T = 1 +_PATCH_H = 2 +_PATCH_W = 2 + + +def _keyframe_cond_frame_indices( + *, + include_keyframe_cond: bool, + keyframe_frame_indices: list[int] | tuple[int, ...] | None, +) -> list[int]: + if not include_keyframe_cond: + if keyframe_frame_indices is not None: + raise ValueError("keyframe_frame_indices must be omitted when keyframe cond is not included") + return [] + if keyframe_frame_indices is None: + raise ValueError("strict fl2va packed layout requires keyframe_frame_indices") + if any(isinstance(value, bool) or not isinstance(value, int) for value in keyframe_frame_indices): + raise ValueError("strict fl2va packed layout requires integer keyframe_frame_indices") + out = list(keyframe_frame_indices) + if tuple(out) not in MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES: + raise ValueError( + "strict fl2va packed layout requires keyframe_frame_indices in " + f"{MINIMAX_H3_FL2VA_KEYFRAME_SIGNATURES!r}, got {out!r}" + ) + return out + + +def _resolve_keyframe_frame_indices( + frame_indices: Sequence[int], + *, + frame_count: int | None, +) -> list[int]: + if frame_indices and frame_count is None: + raise ValueError("frame_count is required when keyframe_frame_indices are provided") + if frame_count is None: + return [] + if frame_count <= 0: + raise ValueError("frame_count must be positive") + seen: dict[int, int] = {} + resolved: list[int] = [] + for block_index, semantic_index in enumerate(frame_indices): + if semantic_index == -1: + resolved_index = frame_count - 1 + elif 0 <= semantic_index < frame_count: + resolved_index = semantic_index + else: + raise ValueError(f"keyframe frame index {semantic_index} must be -1 or in [0, {frame_count})") + previous = seen.get(resolved_index) + if previous is not None: + raise ValueError( + f"keyframe frame index at block {block_index} resolves to " + f"{resolved_index}, already bound by block {previous}" + ) + seen[resolved_index] = block_index + resolved.append(resolved_index) + return resolved + + +def _temporal_position_span(temporal_length: int) -> float: + """Temporal position span for patch_t=1, in fp64. + + NOTE: intentionally NOT merged with ``_video_t_span``. This variant sums + via numpy (pairwise summation), matching the fl2va anchor computation, + while ``_video_t_span`` sums sequentially, matching the ref2va + t-origin accumulation. The two orders diverge in the last ulp + from n=16 onward, so each path must keep its own summation order. + """ + spans = np.ones(int(temporal_length), dtype=np.float64) * _FRAME_RESCALE + for token_index in range(_T_GROUP): + spans[token_index::_T_GROUP] *= _FRAME_PER_TOKEN[token_index] + return float(spans.sum()) + + +def minimax_h3_packed_sequence( + *, + text_len: int, + latent_t: int, + latent_h: int, + latent_w: int, + audio_t: int, + audio_channel: int = 2, + include_keyframe_cond: bool, + keyframe_frame_indices: list[int] | tuple[int, ...] | None = None, + frame_count: int | None = None, +) -> dict[str, torch.Tensor]: + """Build the packed-sequence structural fields for one CFG branch. + + The used length is padded up to a multiple of 64. + """ + ph, pw = latent_h // _PATCH_H, latent_w // _PATCH_W + frame_rows = ph * pw + cond_frame_indices = _keyframe_cond_frame_indices( + include_keyframe_cond=include_keyframe_cond, + keyframe_frame_indices=keyframe_frame_indices, + ) + resolved_cond_frame_indices = _resolve_keyframe_frame_indices( + cond_frame_indices, + frame_count=frame_count, + ) + cond_rows = len(cond_frame_indices) * frame_rows + video_rows = latent_t * frame_rows + audio_rows = audio_t * audio_channel + used = text_len + cond_rows + audio_rows + video_rows + seq_len = ((used + _SEQ_ALIGN - 1) // _SEQ_ALIGN) * _SEQ_ALIGN + + text_sl = slice(0, text_len) + cond_sl = slice(text_len, text_len + cond_rows) + audio_sl = slice(cond_sl.stop, cond_sl.stop + audio_rows) + video_sl = slice(audio_sl.stop, audio_sl.stop + video_rows) + pad_sl = slice(video_sl.stop, seq_len) + + input_ids = torch.full((seq_len,), MINIMAX_H3_PAD_ID, dtype=torch.int64) + input_ids[text_sl] = MINIMAX_H3_TEXT_ID + if cond_rows: + input_ids[cond_sl] = MINIMAX_H3_IMGVID_COND_ID + input_ids[audio_sl] = MINIMAX_H3_AUDIO_ID + input_ids[audio_sl.start] = MINIMAX_H3_AUDIO_FIRST_ID + input_ids[video_sl] = MINIMAX_H3_VIDEO_ID + input_ids[video_sl.start] = MINIMAX_H3_VIDEO_FIRST_ID + input_ids[video_sl.stop - 1] = MINIMAX_H3_VIDEO_LAST_ID + + image_mask = torch.zeros(seq_len, dtype=torch.bool) + image_mask[cond_sl] = True + image_mask[video_sl] = True + audio_mask = torch.zeros(seq_len, dtype=torch.bool) + audio_mask[audio_sl] = True + + img_pos = torch.cat( + [ + torch.arange(cond_sl.start, cond_sl.stop), + torch.arange(video_sl.start, video_sl.stop), + ] + ) + update_mask = torch.zeros(img_pos.shape[0], dtype=torch.bool) + update_mask[cond_rows:] = True + audio_pos = torch.arange(audio_sl.start, audio_sl.stop) + text_pos = torch.arange(0, text_len) + + g = torch.zeros(seq_len, 3, dtype=torch.float64) + g[text_sl, 0] = torch.arange(text_len, dtype=torch.float64) + + t_grid = _video_t_grid(latent_t, float(text_len)) + sqrt_area = np.sqrt(latent_h * latent_w) + h_grid = _axis_from_sqrt_area(latent_h, _PATCH_H, sqrt_area) + w_grid = _axis_from_sqrt_area(latent_w, _PATCH_W, sqrt_area) + hh, ww = torch.meshgrid(h_grid, w_grid, indexing="ij") + frame = torch.stack([hh.reshape(-1), ww.reshape(-1)], dim=-1) + video_g = torch.empty(latent_t, frame_rows, 3, dtype=torch.float64) + video_g[:, :, 0] = t_grid[:, None] + video_g[:, :, 1:] = frame[None] + for block_index, pixel_index in enumerate(resolved_cond_frame_indices): + sl = slice( + cond_sl.start + block_index * frame_rows, + cond_sl.start + (block_index + 1) * frame_rows, + ) + if pixel_index == 0: + cond_t = float(text_len) + elif frame_count is not None and pixel_index == frame_count - 1: + cond_t = float(text_len) + _temporal_position_span(latent_t) - _FRAME_RESCALE + else: + raise ValueError( + f"fl2va packed layout only supports first/last keyframe anchors, got resolved frame index {pixel_index}" + ) + g[sl, 0] = cond_t + g[sl, 1:] = frame + g[video_sl] = video_g.reshape(-1, 3) + audio_t_grid = float(text_len) + torch.arange(audio_t, dtype=torch.float64) + g[audio_sl, 0] = audio_t_grid.repeat(audio_channel) + g[audio_sl, 2] = torch.cat( + [ + torch.full((audio_t,), float(w_grid[0]), dtype=torch.float64), + torch.full((audio_rows - audio_t,), float(w_grid[-1]), dtype=torch.float64), + ] + ) + + token_tags = torch.full((seq_len,), -1, dtype=torch.long) # PADDING + token_tags[text_sl] = 1 # TEXT (fl2va image-segment override happens upstream) + token_tags[audio_sl] = 2 # AUDIO + token_tags[img_pos] = 0 # VIDEO + + cu = torch.tensor([0, used, seq_len], dtype=torch.int32) + doc = torch.zeros(seq_len, dtype=torch.int32) + doc[pad_sl] = 1 + return { + "seq_len": torch.tensor(seq_len), + "input_ids": input_ids, + "image_mask": image_mask, + "audio_mask": audio_mask, + "img_pos": img_pos, + "audio_pos": audio_pos, + "text_pos": text_pos, + "update_mask": update_mask, + "img_position_ids": g, + "token_tags": token_tags, + "cu_seqlens": cu, + "document_id": doc, + } + + +def _positive_int( + block: Mapping[str, object], + key: str, + path: str, + *, + allow_zero: bool = False, +) -> int: + value = block.get(key) + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{path}.{key} must be an integer") + if value < 0 or (value == 0 and not allow_zero): + predicate = "non-negative" if allow_zero else "positive" + raise ValueError(f"{path}.{key} must be {predicate}") + return int(value) + + +def _axis_from_sqrt_area(dim: int, patch: int, sqrt_area: float) -> torch.Tensor: + ratio = dim / sqrt_area + left = (1.0 - ratio) * 1.0 / 2.0 + right = left + ratio * 1.0 + grid = np.linspace(left, right, dim // patch, endpoint=False) * _INTERP + return torch.from_numpy(grid).to(torch.float64) + + +def _video_t_grid(n: int, origin: float) -> torch.Tensor: + spans = torch.tensor( + [_FRAME_RESCALE * _FRAME_PER_TOKEN[k % _T_GROUP] for k in range(n)], + dtype=torch.float64, + ) + return origin + torch.cat([torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)]) + + +def _video_t_span(n: int) -> float: + # Sequential fp64 summation on purpose — see _temporal_position_span for + # why the two span implementations must not be unified. + return sum(_FRAME_RESCALE * _FRAME_PER_TOKEN[k % _T_GROUP] for k in range(n)) + + +def _range_for_slice(sl: slice) -> torch.Tensor: + return torch.arange(sl.start, sl.stop, dtype=torch.long) + + +def _cat_ranges(parts: list[torch.Tensor]) -> torch.Tensor: + if parts: + return torch.cat(parts) + return torch.empty(0, dtype=torch.long) + + +def minimax_h3_packed_sequence_ref2va_blocks( + *, + text_len: int, + latent_t: int, + latent_h: int, + latent_w: int, + audio_t: int, + ref_blocks: Sequence[Mapping[str, object]], + audio_channel: int = 2, + seq_len: int | None = None, +) -> dict[str, torch.Tensor]: + """General ref2va-family packed layout. + + ``ref_blocks`` are consumed in request/plan order: + - ``{"kind": "image", "latent_h": H, "latent_w": W}`` + - ``{"kind": "audio", "ref_audio_t": T}`` + - ``{"kind": "video"|"video_audio", "ref_audio_t": T, + "latent_t": RT, "latent_h": RH, "latent_w": RW}`` + + Video-bearing blocks pack their audio rows immediately before their video + rows; both share the same temporal origin and advance by the longer of the + audio and video spans. Standalone audio advances the target origin by its + own T, and image blocks advance it by one integer slot. + """ + if not isinstance(ref_blocks, Sequence) or isinstance(ref_blocks, (str, bytes)): + raise ValueError("ref_blocks must be a sequence") + + parsed: list[dict[str, object]] = [] + ref_visual_rows = 0 + ref_audio_rows = 0 + for index, raw in enumerate(ref_blocks): + path = f"ref_blocks[{index}]" + if not isinstance(raw, Mapping): + raise ValueError(f"{path} must be an object") + kind = raw.get("kind", raw.get("type")) + if not isinstance(kind, str) or not kind: + raise ValueError(f"{path}.kind must be a non-empty string") + if kind == "image": + rh = _positive_int(raw, "latent_h", path) + rw = _positive_int(raw, "latent_w", path) + rows = (rh // _PATCH_H) * (rw // _PATCH_W) + item = {"kind": kind, "latent_h": rh, "latent_w": rw, "rows": rows} + ref_visual_rows += rows + elif kind == "audio": + rt = _positive_int(raw, "ref_audio_t", path, allow_zero=True) + rows = rt * audio_channel + item = {"kind": kind, "ref_audio_t": rt, "audio_rows": rows} + ref_audio_rows += rows + elif kind in ("video", "video_audio"): + rt = _positive_int(raw, "ref_audio_t", path, allow_zero=True) + vt = _positive_int(raw, "latent_t", path) + vh = _positive_int(raw, "latent_h", path) + vw = _positive_int(raw, "latent_w", path) + frame_rows = (vh // _PATCH_H) * (vw // _PATCH_W) + audio_rows = rt * audio_channel + video_rows = vt * frame_rows + item = { + "kind": kind, + "ref_audio_t": rt, + "latent_t": vt, + "latent_h": vh, + "latent_w": vw, + "frame_rows": frame_rows, + "audio_rows": audio_rows, + "video_rows": video_rows, + } + ref_audio_rows += audio_rows + ref_visual_rows += video_rows + else: + raise ValueError(f"{path}.kind unsupported for ref2va: {kind!r}") + parsed.append(item) + + ph, pw = latent_h // _PATCH_H, latent_w // _PATCH_W + frame_rows = ph * pw + video_rows = latent_t * frame_rows + audio_rows = audio_t * audio_channel + ref_rows = ref_visual_rows + ref_audio_rows + used = text_len + ref_rows + audio_rows + video_rows + if seq_len is None: + seq_len = ((used + _SEQ_ALIGN - 1) // _SEQ_ALIGN) * _SEQ_ALIGN + if seq_len < used: + raise ValueError(f"seq_len {seq_len} < used rows {used}") + + text_sl = slice(0, text_len) + cursor = text_len + block_slices: list[dict[str, object]] = [] + for item in parsed: + kind = str(item["kind"]) + if kind == "image": + rows = int(item["rows"]) + visual_sl = slice(cursor, cursor + rows) + cursor = visual_sl.stop + block_slices.append({**item, "visual_sl": visual_sl}) + elif kind == "audio": + rows = int(item["audio_rows"]) + audio_sl = slice(cursor, cursor + rows) + cursor = audio_sl.stop + block_slices.append({**item, "audio_sl": audio_sl}) + else: + a_rows = int(item["audio_rows"]) + v_rows = int(item["video_rows"]) + audio_sl = slice(cursor, cursor + a_rows) + visual_sl = slice(audio_sl.stop, audio_sl.stop + v_rows) + cursor = visual_sl.stop + block_slices.append({**item, "audio_sl": audio_sl, "visual_sl": visual_sl}) + + audio_sl = slice(cursor, cursor + audio_rows) + video_sl = slice(audio_sl.stop, audio_sl.stop + video_rows) + pad_sl = slice(video_sl.stop, seq_len) + + input_ids = torch.full((seq_len,), MINIMAX_H3_PAD_ID, dtype=torch.int64) + input_ids[text_sl] = MINIMAX_H3_TEXT_ID + image_mask = torch.zeros(seq_len, dtype=torch.bool) + audio_mask = torch.zeros(seq_len, dtype=torch.bool) + + ref_img_pos_parts: list[torch.Tensor] = [] + ref_audio_pos_parts: list[torch.Tensor] = [] + g = torch.zeros(seq_len, 3, dtype=torch.float64) + g[text_sl, 0] = torch.arange(text_len, dtype=torch.float64) + + target_area = np.sqrt(latent_h * latent_w) + h_grid = _axis_from_sqrt_area(latent_h, _PATCH_H, target_area) + w_grid = _axis_from_sqrt_area(latent_w, _PATCH_W, target_area) + hh, ww = torch.meshgrid(h_grid, w_grid, indexing="ij") + target_frame = torch.stack([hh.reshape(-1), ww.reshape(-1)], dim=-1) + + t_cursor = float(text_len) + for item in block_slices: + kind = str(item["kind"]) + if kind == "image": + visual_sl = item["visual_sl"] + assert isinstance(visual_sl, slice) + input_ids[visual_sl] = MINIMAX_H3_IMGVID_COND_ID + image_mask[visual_sl] = True + ref_img_pos_parts.append(_range_for_slice(visual_sl)) + rh = int(item["latent_h"]) + rw = int(item["latent_w"]) + area = np.sqrt(rh * rw) + ref_hh, ref_ww = torch.meshgrid( + _axis_from_sqrt_area(rh, _PATCH_H, area), + _axis_from_sqrt_area(rw, _PATCH_W, area), + indexing="ij", + ) + g[visual_sl, 0] = t_cursor + g[visual_sl, 1] = ref_hh.reshape(-1) + g[visual_sl, 2] = ref_ww.reshape(-1) + t_cursor += 1.0 + elif kind == "audio": + audio_ref_sl = item["audio_sl"] + assert isinstance(audio_ref_sl, slice) + ref_t = int(item["ref_audio_t"]) + input_ids[audio_ref_sl] = MINIMAX_H3_AUDIO_REF_COND_ID + audio_mask[audio_ref_sl] = True + ref_audio_pos_parts.append(_range_for_slice(audio_ref_sl)) + ref_t_grid = t_cursor + torch.arange(ref_t, dtype=torch.float64) + g[audio_ref_sl, 0] = ref_t_grid.repeat(audio_channel) + if ref_t: + g[audio_ref_sl, 2] = torch.cat( + [ + torch.full((ref_t,), float(w_grid[0]), dtype=torch.float64), + torch.full( + (int(item["audio_rows"]) - ref_t,), + float(w_grid[-1]), + dtype=torch.float64, + ), + ] + ) + t_cursor += float(ref_t) + else: + audio_ref_sl = item["audio_sl"] + visual_sl = item["visual_sl"] + assert isinstance(audio_ref_sl, slice) + assert isinstance(visual_sl, slice) + ref_t = int(item["ref_audio_t"]) + vt = int(item["latent_t"]) + vh = int(item["latent_h"]) + vw = int(item["latent_w"]) + input_ids[audio_ref_sl] = MINIMAX_H3_AUDIO_REF_COND_ID + input_ids[visual_sl] = MINIMAX_H3_IMGVID_COND_ID + audio_mask[audio_ref_sl] = True + image_mask[visual_sl] = True + ref_audio_pos_parts.append(_range_for_slice(audio_ref_sl)) + ref_img_pos_parts.append(_range_for_slice(visual_sl)) + + ref_area = np.sqrt(vh * vw) + rv_h_grid = _axis_from_sqrt_area(vh, _PATCH_H, ref_area) + rv_w_grid = _axis_from_sqrt_area(vw, _PATCH_W, ref_area) + rv_hh, rv_ww = torch.meshgrid(rv_h_grid, rv_w_grid, indexing="ij") + + ref_t_grid = t_cursor + torch.arange(ref_t, dtype=torch.float64) + g[audio_ref_sl, 0] = ref_t_grid.repeat(audio_channel) + if ref_t: + g[audio_ref_sl, 2] = torch.cat( + [ + torch.full((ref_t,), float(rv_w_grid[0]), dtype=torch.float64), + torch.full( + (int(item["audio_rows"]) - ref_t,), + float(rv_w_grid[-1]), + dtype=torch.float64, + ), + ] + ) + + rv_frame = torch.stack([rv_hh.reshape(-1), rv_ww.reshape(-1)], dim=-1) + rv_g = torch.empty(vt, int(item["frame_rows"]), 3, dtype=torch.float64) + rv_g[:, :, 0] = _video_t_grid(vt, t_cursor)[:, None] + rv_g[:, :, 1:] = rv_frame[None] + g[visual_sl] = rv_g.reshape(-1, 3) + t_cursor += max(float(ref_t), _video_t_span(vt)) + + input_ids[audio_sl] = MINIMAX_H3_AUDIO_ID + input_ids[audio_sl.start] = MINIMAX_H3_AUDIO_FIRST_ID + input_ids[video_sl] = MINIMAX_H3_VIDEO_ID + input_ids[video_sl.start] = MINIMAX_H3_VIDEO_FIRST_ID + input_ids[video_sl.stop - 1] = MINIMAX_H3_VIDEO_LAST_ID + audio_mask[audio_sl] = True + image_mask[video_sl] = True + + audio_t_grid = t_cursor + torch.arange(audio_t, dtype=torch.float64) + g[audio_sl, 0] = audio_t_grid.repeat(audio_channel) + g[audio_sl, 2] = torch.cat( + [ + torch.full((audio_t,), float(w_grid[0]), dtype=torch.float64), + torch.full((audio_rows - audio_t,), float(w_grid[-1]), dtype=torch.float64), + ] + ) + + video_g = torch.empty(latent_t, frame_rows, 3, dtype=torch.float64) + video_g[:, :, 0] = _video_t_grid(latent_t, t_cursor)[:, None] + video_g[:, :, 1:] = target_frame[None] + g[video_sl] = video_g.reshape(-1, 3) + + target_img_pos = _range_for_slice(video_sl) + target_audio_pos = _range_for_slice(audio_sl) + img_pos = _cat_ranges(ref_img_pos_parts + [target_img_pos]) + audio_pos = _cat_ranges(ref_audio_pos_parts + [target_audio_pos]) + + update_mask = torch.zeros(img_pos.shape[0], dtype=torch.bool) + update_mask[ref_visual_rows:] = True + audio_update_mask = torch.zeros(audio_pos.shape[0], dtype=torch.bool) + audio_update_mask[ref_audio_rows:] = True + text_pos = torch.arange(0, text_len) + + token_tags = torch.full((seq_len,), -1, dtype=torch.long) # PADDING + token_tags[text_sl] = 1 # TEXT + if ref_audio_rows: + token_tags[_cat_ranges(ref_audio_pos_parts)] = 2 # AUDIO (reference) + token_tags[audio_sl] = 2 # AUDIO (target) + token_tags[img_pos] = 0 # VIDEO (refs + target) + + cu = torch.tensor([0, used, seq_len], dtype=torch.int32) + doc = torch.zeros(seq_len, dtype=torch.int32) + doc[pad_sl] = 1 + return { + "seq_len": torch.tensor(seq_len), + "input_ids": input_ids, + "image_mask": image_mask, + "audio_mask": audio_mask, + "img_pos": img_pos, + "audio_pos": audio_pos, + "text_pos": text_pos, + "update_mask": update_mask, + "audio_update_mask": audio_update_mask, + "img_position_ids": g, + "token_tags": token_tags, + "cu_seqlens": cu, + "document_id": doc, + } + + +__all__ = [ + "MINIMAX_H3_AUDIO_FIRST_ID", + "MINIMAX_H3_AUDIO_ID", + "MINIMAX_H3_AUDIO_REF_COND_ID", + "MINIMAX_H3_IMGVID_COND_ID", + "MINIMAX_H3_PAD_ID", + "MINIMAX_H3_TEXT_ID", + "MINIMAX_H3_VIDEO_FIRST_ID", + "MINIMAX_H3_VIDEO_ID", + "MINIMAX_H3_VIDEO_LAST_ID", + "minimax_h3_packed_sequence", + "minimax_h3_packed_sequence_ref2va_blocks", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/packed_tokens.py b/vllm_omni/diffusion/models/minimax_h3/packed_tokens.py new file mode 100644 index 000000000..e37fe75d8 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/packed_tokens.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from collections.abc import Sequence + +import torch + + +def _int_tuple(value: Sequence[int], name: str, length: int) -> tuple[int, ...]: + if len(value) != length: + raise ValueError(f"{name} must have length {length}, got {list(value)!r}") + out = tuple(int(item) for item in value) + if any(item <= 0 for item in out): + raise ValueError(f"{name} values must be positive, got {list(value)!r}") + return out + + +def _rank(tensor: torch.Tensor, name: str, rank: int) -> None: + if tensor.ndim != rank: + raise ValueError(f"{name} must be rank {rank}, got shape={list(tensor.shape)}") + + +def minimax_h3_patchify_video_latent( + latent: torch.Tensor, + *, + patch_size: Sequence[int], +) -> torch.Tensor: + """Pack video latent [B,C,T,H,W] into DiT token rows.""" + + _rank(latent, "video latent", 5) + pt, ph, pw = _int_tuple(patch_size, "patch_size", 3) + batch, channel, full_t, full_h, full_w = (int(dim) for dim in latent.shape) + if full_t % pt or full_h % ph or full_w % pw: + raise ValueError( + "video latent spatial/time dims must be divisible by patch_size: " + f"shape={list(latent.shape)}, patch_size={[pt, ph, pw]}" + ) + t, h, w = full_t // pt, full_h // ph, full_w // pw + packed = latent.reshape(batch, channel, t, pt, h, ph, w, pw) + packed = torch.einsum("nctrhpwq->nthwcrpq", packed) + return packed.reshape(batch * t * h * w, channel * pt * ph * pw).contiguous() + + +def minimax_h3_unpatchify_video_tokens( + rows: torch.Tensor, + *, + latent_shape: Sequence[int], + patch_size: Sequence[int], +) -> torch.Tensor: + """Unpack DiT video token rows into latent [B,C,T,H,W].""" + + _rank(rows, "video token rows", 2) + t, h, w, channel = _int_tuple(latent_shape, "latent_shape", 4) + pt, ph, pw = _int_tuple(patch_size, "patch_size", 3) + expected_dim = pt * ph * pw * channel + if int(rows.shape[-1]) != expected_dim: + raise ValueError( + f"video token dim {int(rows.shape[-1])} != patch volume * channel " + f"{expected_dim} for latent_shape={list(latent_shape)}, " + f"patch_size={[pt, ph, pw]}" + ) + rows_per_sample = t * h * w + if int(rows.shape[0]) % rows_per_sample: + raise ValueError( + f"video rows {int(rows.shape[0])} must be divisible by t*h*w " + f"{rows_per_sample} for latent_shape={list(latent_shape)}" + ) + packed = rows.reshape(-1, t, h, w, channel, pt, ph, pw) + latent = torch.einsum("nthwcrpq->nctrhpwq", packed) + return latent.reshape(-1, channel, t * pt, h * ph, w * pw).contiguous() + + +def minimax_h3_pack_audio_latent(latent: torch.Tensor) -> torch.Tensor: + """Pack audio VAE latent [audio_channel, latent_dim, T] into rows.""" + + _rank(latent, "audio latent", 3) + audio_channel, latent_dim, steps = (int(dim) for dim in latent.shape) + return ( + latent.permute(0, 2, 1) + .reshape( + audio_channel * steps, + latent_dim, + ) + .contiguous() + ) + + +def minimax_h3_unpack_audio_tokens( + rows: torch.Tensor, + *, + audio_t: int, + audio_channel: int, +) -> torch.Tensor: + """Unpack DiT audio token rows into audio VAE latent [C,latent_dim,T].""" + + _rank(rows, "audio token rows", 2) + audio_t = int(audio_t) + audio_channel = int(audio_channel) + if audio_t <= 0 or audio_channel <= 0: + raise ValueError(f"audio_t and audio_channel must be positive, got {audio_t=} {audio_channel=}") + if int(rows.shape[0]) != audio_t: + raise ValueError(f"audio rows {int(rows.shape[0])} != audio_t {audio_t}") + if audio_t % audio_channel: + raise ValueError(f"audio_t must be divisible by audio_channel, got {audio_t=} {audio_channel=}") + native = rows.reshape(audio_channel, audio_t // audio_channel, int(rows.shape[-1])) + return native.permute(0, 2, 1).contiguous() + + +__all__ = [ + "minimax_h3_pack_audio_latent", + "minimax_h3_patchify_video_latent", + "minimax_h3_unpack_audio_tokens", + "minimax_h3_unpatchify_video_tokens", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py new file mode 100644 index 000000000..4bef10bcc --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -0,0 +1,1197 @@ +# SPDX-License-Identifier: Apache-2.0 +"""vLLM-Omni pipeline for MiniMax H3 FL2VA and Ref2VA partitions.""" + +from __future__ import annotations + +import json +import os +import tempfile +from collections.abc import Iterable +from pathlib import Path +from typing import Any, ClassVar + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn as nn +from PIL import Image +from transformers import Qwen2TokenizerFast, Qwen3VLProcessor +from vllm.logger import init_logger + +from vllm_omni.diffusion import envs +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.parallel_state import ( + get_dit_group, + init_world_group, +) +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import ( + DiffusersPipelineLoader, +) +from vllm_omni.diffusion.models.interface import ( + SupportAudioInput, + SupportAudioOutput, + SupportImageInput, + SupportsComponentDiscovery, +) +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( + DiffusionPipelineProfilerMixin, +) +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.platforms import current_omni_platform + +from .condition_noise import ( + minimax_h3_audio_cond_noise_aug_rows, + minimax_h3_imgvid_cond_noise_aug_rows, +) +from .denoise_loop import MiniMaxH3DenoiseBranch, minimax_h3_denoise_loop +from .encoder import MiniMaxH3Qwen3VLEncoder +from .minimax_h3_transformer import MiniMaxH3DiTModel +from .packed_sequence import ( + minimax_h3_packed_sequence, + minimax_h3_packed_sequence_ref2va_blocks, +) +from .packed_tokens import ( + minimax_h3_patchify_video_latent, + minimax_h3_unpack_audio_tokens, + minimax_h3_unpatchify_video_tokens, +) +from .presentation import ( + minimax_h3_multi_image_presentation_ids, + minimax_h3_multi_image_presentation_token_tags, + minimax_h3_ref2va_presentation, + minimax_h3_ref2va_video_presentation, + minimax_h3_text_only_ids, +) +from .reference_video import ( + load_video_audio, + load_video_frames, + prepare_reference_videos, + sample_reference_video_frames, +) +from .time_request import ( + MINIMAX_H3_SHAPE_PLANNER, + minimax_h3_align_frame_count, + minimax_h3_time_shift_sigmas, +) +from .vae import MiniMaxH3AudioVAE, MiniMaxH3VideoVAE + +logger = init_logger(__name__) + +MINIMAX_H3_FPS = 24 +MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000 +MINIMAX_H3_IMGVID_COND_TIMESTEP = 0.999 +MINIMAX_H3_AUDIO_REF_COND_TIMESTEP = 1.0 +MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE = 2048 +MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE = 32 + + +def _minimax_h3_post_process(output, output_type: str = "np"): + """Convert the joint video/audio output without capturing worker state. + + The callable crosses the multiprocessing result queue, so it must remain a + module-level function that the standard pickle module can resolve. + """ + if not isinstance(output, tuple) or len(output) != 2: + return output + video, audio = output + if output_type == "latent": + return output + if output_type == "np": + video = video.detach().float().cpu().permute(0, 2, 3, 4, 1).clamp(0, 1).numpy() + audio = audio.detach().float().cpu().numpy() + video = [sample for sample in video] + return { + "video": video, + "audio": audio, + "audio_sample_rate": MINIMAX_H3_AUDIO_SAMPLE_RATE, + "fps": MINIMAX_H3_FPS, + } + + +def get_minimax_h3_post_process_func( + od_config: OmniDiffusionConfig, +): + del od_config + return _minimax_h3_post_process + + +def _align_multiple(value: float, multiple: int = 32) -> int: + return max(multiple, int(round(float(value) / multiple)) * multiple) + + +def _load_image(value: Any) -> Image.Image: + if isinstance(value, list): + if len(value) != 1: + raise ValueError("MiniMax H3 currently supports exactly one image") + value = value[0] + if isinstance(value, (str, os.PathLike)): + return Image.open(value).convert("RGB") + if isinstance(value, Image.Image): + return value.convert("RGB") + if isinstance(value, torch.Tensor): + tensor = value.detach().float().cpu() + if tensor.ndim == 4 and tensor.shape[0] == 1: + tensor = tensor[0] + if tensor.ndim != 3: + raise ValueError(f"image tensor must be [C,H,W], got {tuple(tensor.shape)}") + if tensor.shape[0] in (1, 3, 4): + tensor = tensor.permute(1, 2, 0) + array = tensor.numpy() + if array.max(initial=0) <= 1.0: + array = array * 255.0 + return Image.fromarray(array.clip(0, 255).astype(np.uint8)).convert("RGB") + raise TypeError(f"unsupported MiniMax H3 image input {type(value)!r}") + + +def _load_audio(value: Any) -> tuple[torch.Tensor, int]: + import torchaudio + + if isinstance(value, list): + if len(value) != 1: + raise ValueError("MiniMax H3 currently supports exactly one audio") + value = value[0] + if isinstance(value, (str, os.PathLike)): + return torchaudio.load(str(value)) + if isinstance(value, tuple) and len(value) == 2: + waveform, sample_rate = value + waveform = torch.as_tensor(waveform).float() + return waveform, int(sample_rate) + if isinstance(value, dict): + waveform = value.get("waveform", value.get("array")) + sample_rate = value.get("sample_rate", value.get("sampling_rate")) + if waveform is not None and sample_rate is not None: + return torch.as_tensor(waveform).float(), int(sample_rate) + raise TypeError("MiniMax H3 audio input must be a path, (waveform, sample_rate), or a waveform mapping") + + +def _dit_rank_world() -> tuple[Any, int, int]: + if not dist.is_initialized(): + return None, 0, 1 + group = get_dit_group() + return group, dist.get_rank(group), dist.get_world_size(group) + + +def _broadcast_tensor( + tensor: torch.Tensor | None, + *, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + group, rank, world_size = _dit_rank_world() + if world_size == 1: + if tensor is None: + raise ValueError("source tensor is required for single-rank execution") + return tensor.to(device=device, dtype=dtype) + + shape = torch.zeros(5, dtype=torch.long, device=device) + if rank == 0: + if tensor is None: + raise ValueError("rank 0 must provide a tensor to broadcast") + shape[0] = tensor.ndim + shape[1 : tensor.ndim + 1] = torch.tensor( + tensor.shape, + device=device, + ) + dist.broadcast(shape, src=0, group=group) + ndim = int(shape[0].item()) + tensor_shape = tuple(int(v) for v in shape[1 : ndim + 1].tolist()) + if rank == 0: + output = tensor.to(device=device, dtype=dtype).contiguous() + else: + output = torch.empty(tensor_shape, device=device, dtype=dtype) + dist.broadcast(output, src=0, group=group) + return output + + +def _reference_image_shape(image: Image.Image) -> tuple[int, int]: + width, height = image.size + if width > 4 * height or height > 4 * width: + raise ValueError(f"reference image aspect ratio must be in [1:4, 4:1], got {width}x{height}") + scale = MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE / min(width, height) + return ( + _align_multiple( + width * scale, + MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE, + ), + _align_multiple( + height * scale, + MINIMAX_H3_REFERENCE_IMAGE_MULTIPLE, + ), + ) + + +class _SingleRankEncoderGroup: + """Lightweight encoder group for ``text_encoder_tp_size == 1``. + + Avoids creating a distributed ``GroupCoordinator`` with a single-member + rank set, which would assert on every other DiT rank that is not part of + the group. The pipeline and encoder only use the attributes below, and + all ``world_size == 1`` code paths short-circuit before any collective. + """ + + world_size: int = 1 + ranks: list[int] = [0] + + def __init__(self, rank: int) -> None: + self.rank_in_group = 0 if rank == 0 else -1 + self.device_group = None + + +class MiniMaxH3Pipeline( + nn.Module, + ProgressBarMixin, + DiffusionPipelineProfilerMixin, + SupportImageInput, + SupportAudioInput, + SupportAudioOutput, + SupportsComponentDiscovery, +): + """CFG-distilled joint video/audio generation for MiniMax H3.""" + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["video_vae", "audio_vae"] + _PROFILER_TARGETS: ClassVar[list[str]] = [ + "_prepare_reference_videos", + "encode_prompt", + "_encode_video_conditions", + "_encode_video_audio_conditions", + "diffuse", + "decode", + ] + dummy_run_num_frames: ClassVar[int] = 0 + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ) -> None: + del prefix + super().__init__() + self.od_config = od_config + self.parallel_config = od_config.parallel_config + if int(self.parallel_config.cfg_parallel_size) != 1: + raise ValueError("MiniMax-H3 is CFG-distilled and has no negative branch; cfg_parallel_size must be 1") + self.device = get_local_device() + model_path = str(od_config.model) + model_index = json.loads((Path(model_path) / "model_index.json").read_text(encoding="utf-8")) + release = model_index.get("_minimax_h3") or {} + self.partition = str(release.get("partition", "")) + self.supported_tasks = frozenset(release.get("tasks") or ()) + shifts = release.get("sigma_shift_scales") or {} + self.default_video_shift = float(shifts.get("video", 12.0)) + self.default_audio_shift = float(shifts.get("audio", 3.0)) + + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=model_path, + subfolder="transformer", + revision=od_config.revision, + prefix="transformer.", + fall_back_to_pt=False, + ) + ] + self.transformer = MiniMaxH3DiTModel( + od_config, + quant_config=od_config.quantization_config, + ) + + self.tokenizer = Qwen2TokenizerFast.from_pretrained( + model_path, + subfolder="tokenizer", + local_files_only=os.path.isdir(model_path), + ) + self.processor = Qwen3VLProcessor.from_pretrained( + model_path, + subfolder="processor", + local_files_only=os.path.isdir(model_path), + ) + + _, rank, dit_world = _dit_rank_world() + self._dit_rank = rank + text_encoder_tp_size = int(getattr(self.parallel_config, "text_encoder_tp_size", 1)) + if text_encoder_tp_size < 1: + raise ValueError(f"text_encoder_tp_size must be >= 1, got {text_encoder_tp_size}") + if text_encoder_tp_size > dit_world: + raise ValueError( + f"text_encoder_tp_size must not exceed the DiT group size ({dit_world}), got {text_encoder_tp_size}" + ) + # The Qwen3-VL text model uses 64 attention heads / 8 KV heads; the + # encoder shards them across the encoder TP ranks. + if 64 % text_encoder_tp_size or 8 % text_encoder_tp_size: + raise ValueError( + "text_encoder_tp_size must divide both Qwen3-VL " + f"num_attention_heads (64) and num_key_value_heads (8), " + f"got {text_encoder_tp_size}" + ) + self.text_encoder_tp_size = text_encoder_tp_size + self.text_encoder_group = self._build_text_encoder_group(text_encoder_tp_size) + self.text_encoder = MiniMaxH3Qwen3VLEncoder( + os.path.join(model_path, "text_encoder"), + device=self.device, + load_model=rank < text_encoder_tp_size, + encoder_group=self.text_encoder_group, + ) + self.video_vae = MiniMaxH3VideoVAE( + os.path.join(model_path, "video_vae"), + device=self.device, + ) + self.audio_vae = MiniMaxH3AudioVAE( + os.path.join(model_path, "audio_vae"), + device=self.device, + ) + # Registry-side VAE patch-parallel discovery uses ``pipeline.vae``. + self.vae = self.video_vae + + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=(od_config.enable_diffusion_pipeline_profiler) + ) + + def load_weights( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> set[str]: + prefix = "transformer." + + def transformer_weights(): + for name, tensor in weights: + if name.startswith(prefix): + yield name[len(prefix) :], tensor + + loaded = self.transformer.load_weights(transformer_weights()) + self.transformer.post_load_weights() + loaded_with_prefix = {prefix + name for name in loaded} + # The text encoder and both VAEs load eagerly in ``__init__`` rather + # than through ``weights_sources``. Record them for the runner's strict + # missing-parameter check. + for component_name in ("text_encoder", "video_vae", "audio_vae"): + component = getattr(self, component_name) + loaded_with_prefix.update(f"{component_name}.{name}" for name, _ in component.named_parameters()) + return loaded_with_prefix + + def _resolve_task( + self, + requested: str | None, + multi_modal_data: dict[str, Any], + ) -> str: + if requested is None: + if self.partition == "ref2va": + requested = "ref2va" + elif multi_modal_data.get("image") is not None: + requested = "fl2va" + else: + requested = "t2va" + task = str(requested).lower() + if task not in self.supported_tasks: + raise ValueError( + f"checkpoint partition {self.partition!r} supports {sorted(self.supported_tasks)}, got task={task!r}" + ) + return task + + def _resolve_shape( + self, + task: str, + sampling: Any, + image: Image.Image | None, + ) -> tuple[int, int, int, int, int]: + fps = int(sampling.fps or MINIMAX_H3_FPS) + if fps != MINIMAX_H3_FPS: + raise ValueError(f"MiniMax H3 output fps is fixed at {MINIMAX_H3_FPS}") + extra = sampling.extra_args or {} + duration = extra.get("duration") + if duration is not None: + requested_frames = int(round(float(duration) * fps)) + elif int(sampling.num_frames or 1) > 1: + requested_frames = int(sampling.num_frames) + else: + requested_frames = 124 if task == "ref2va" else 209 + num_frames = minimax_h3_align_frame_count(requested_frames) + + height = sampling.height + width = sampling.width + if height is None or width is None: + if task == "fl2va" and image is not None: + ratio = image.width / image.height + if ratio >= 1: + height = 768 + width = _align_multiple(768 * ratio) + else: + width = 768 + height = _align_multiple(768 / ratio) + else: + height, width = 768, 1344 + height = int(height) // 32 * 32 + width = int(width) // 32 * 32 + if min(height, width) <= 0: + raise ValueError(f"invalid MiniMax H3 canvas {width}x{height}") + if width > 4 * height or height > 4 * width: + raise ValueError("MiniMax H3 canvas aspect ratio must be in [1:4, 4:1]") + + latent_t = MINIMAX_H3_SHAPE_PLANNER.video_latent_t(num_frames) + audio_t = MINIMAX_H3_SHAPE_PLANNER.audio_latent_t(num_frames / fps) + return height, width, num_frames, latent_t, audio_t + + def encode_prompt( + self, + *, + task: str, + prompt: str, + image: Image.Image | None, + prepared_videos: list[dict[str, Any]] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + _, rank, _ = _dit_rank_world() + hidden = None + tags = None + ids = None + vision_kwargs: dict[str, torch.Tensor] = {} + if rank == 0: + if task == "t2va": + ids = minimax_h3_text_only_ids(self.tokenizer, prompt) + tags = torch.ones(ids.shape[0], dtype=torch.long) + vision_kwargs = {} + elif prepared_videos: + videos = [] + sampled_videos = [] + for index, item in enumerate(prepared_videos): + sampled = sample_reference_video_frames( + item["prepared_path"], + workdir=str(Path(item["prepared_path"]).parent / f"qwen_frames_{index}"), + ) + videos.append(np.stack(sampled["frames"])) + sampled_videos.append(sampled) + vision = self.processor.video_processor( + videos=videos, + do_sample_frames=False, + return_tensors="pt", + ) + video_grid = vision["video_grid_thw"] + merge = int(self.processor.image_processor.merge_size) ** 2 + block_counts = [] + block_timestamps = [] + for index, sampled in enumerate(sampled_videos): + blocks = int(video_grid[index, 0]) + per_block = int(video_grid[index, 1]) * int(video_grid[index, 2]) // merge + timestamps = sampled["block_timestamps"] + if len(timestamps) != blocks: + raise ValueError( + f"video block count mismatch: processor={blocks}, timestamps={len(timestamps)}" + ) + block_counts.append([per_block] * blocks) + block_timestamps.append(timestamps) + condition_labels: list[tuple[str, int]] = [] + audio_index = 0 + for video_index, item in enumerate(prepared_videos, start=1): + if item["input_has_audio"]: + audio_index += 1 + condition_labels.append(("audio", audio_index)) + condition_labels.append(("video", video_index)) + ids, tags = minimax_h3_ref2va_video_presentation( + self.tokenizer, + prompt=prompt, + condition_labels=condition_labels, + image_token_count=None, + video_block_token_counts=block_counts, + video_block_timestamps=block_timestamps, + ) + vision_kwargs = { + "pixel_values_videos": vision["pixel_values_videos"], + "video_grid_thw": video_grid, + } + else: + if image is None: + raise ValueError(f"{task} requires one image") + vision = self.processor.image_processor( + images=[image], + return_tensors="pt", + ) + image_grid = vision["image_grid_thw"] + merge = int(self.processor.image_processor.merge_size) ** 2 + image_tokens = int(image_grid[0].prod().item()) // merge + if task == "fl2va": + ids = minimax_h3_multi_image_presentation_ids( + self.tokenizer, + prompt=prompt, + image_token_counts=[image_tokens], + ) + tags = minimax_h3_multi_image_presentation_token_tags( + self.tokenizer, + prompt=prompt, + image_token_counts=[image_tokens], + ) + else: + ids, tags = minimax_h3_ref2va_presentation( + self.tokenizer, + prompt=prompt, + condition_labels=[("image", 1), ("audio", 1)], + image_token_count=image_tokens, + ) + vision_kwargs = { + "pixel_values": vision["pixel_values"], + "image_grid_thw": image_grid, + } + + logger.info( + "MiniMax H3 %s Qwen presentation: %d tokens%s", + task, + int(ids.shape[0]), + (f", {len(prepared_videos)} reference videos" if prepared_videos else ""), + ) + + if rank < self.text_encoder_tp_size: + # Distribute the encode inputs from the DiT main rank to the other + # encoder TP ranks, then run the distributed encode on all of them. + ids = self._distribute_encode_inputs(ids, vision_kwargs) + hidden = self._encode_text_hidden(ids, vision_kwargs) + + hidden = _broadcast_tensor( + hidden, + dtype=torch.bfloat16, + device=self.device, + ) + tags = _broadcast_tensor( + tags, + dtype=torch.long, + device=self.device, + ) + return hidden, tags + + def _build_text_encoder_group(self, text_encoder_tp_size: int) -> Any: + """Create the encoder tensor-parallel process group. + + The encoder group covers the first ``text_encoder_tp_size`` DiT ranks + (the DiT group is always global ranks ``[0, dit_world)``). Every rank + participates in ``new_group`` so the collective completes; ranks + outside the group never run encoder collectives. For a single-rank + encoder we return a lightweight placeholder so non-encoder ranks do + not need to join a ``GroupCoordinator`` that would assert on ranks + outside the group. + """ + if text_encoder_tp_size == 1: + return _SingleRankEncoderGroup(rank=self._dit_rank) + ranks = list(range(text_encoder_tp_size)) + return init_world_group( + ranks=ranks, + local_rank=envs.LOCAL_RANK, + backend=current_omni_platform.dist_backend, + ) + + def _encoder_group_broadcast_tensor( + self, + tensor: torch.Tensor | None, + *, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + """Broadcast a tensor from encoder rank 0 over the encoder TP group.""" + group = self.text_encoder_group + if group.world_size == 1: + if tensor is None: + raise ValueError("source tensor is required for single-rank execution") + return tensor.to(device=device, dtype=dtype) + + shape = torch.zeros(8, dtype=torch.long, device=device) + if group.rank_in_group == 0: + if tensor is None: + raise ValueError("encoder rank 0 must provide a tensor to broadcast") + shape[0] = tensor.ndim + shape[1 : tensor.ndim + 1] = torch.tensor(tensor.shape, device=device) + torch.distributed.broadcast(shape, src=group.ranks[0], group=group.device_group) + ndim = int(shape[0].item()) + tensor_shape = tuple(int(value) for value in shape[1 : ndim + 1].tolist()) + if group.rank_in_group == 0: + output = tensor.to(device=device, dtype=dtype).contiguous() + else: + output = torch.empty(tensor_shape, device=device, dtype=dtype) + torch.distributed.broadcast(output, src=group.ranks[0], group=group.device_group) + return output + + def _distribute_encode_inputs( + self, + ids: torch.Tensor | None, + vision_kwargs: dict[str, torch.Tensor], + ) -> torch.Tensor: + """Fan out encode inputs from encoder rank 0 to the encoder TP ranks. + + Mutates ``vision_kwargs`` in place so every encoder rank ends up with + the same vision tensors, and returns the broadcast ``input_ids``. + """ + keys = ("pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw") + key_dtypes = { + "pixel_values": torch.bfloat16, + "pixel_values_videos": torch.bfloat16, + "image_grid_thw": torch.long, + "video_grid_thw": torch.long, + } + group = self.text_encoder_group + device = self.device + if group.world_size == 1: + if ids is None: + raise ValueError("encoder rank 0 must produce input ids") + return ids.to(device=device, dtype=torch.long) + + mask = torch.zeros(len(keys), dtype=torch.long, device=device) + if group.rank_in_group == 0: + for index, key in enumerate(keys): + mask[index] = 1 if key in vision_kwargs else 0 + torch.distributed.broadcast(mask, src=group.ranks[0], group=group.device_group) + + if group.rank_in_group == 0: + ids = self._encoder_group_broadcast_tensor(ids, dtype=torch.long, device=device) + else: + ids = self._encoder_group_broadcast_tensor(None, dtype=torch.long, device=device) + for index, key in enumerate(keys): + if mask[index].item() == 0: + continue + source = vision_kwargs.get(key) if group.rank_in_group == 0 else None + vision_kwargs[key] = self._encoder_group_broadcast_tensor( + source, + dtype=key_dtypes[key], + device=device, + ) + return ids + + def _prepare_reference_videos( + self, + values: Any, + *, + target_frame_count: int, + workdir: str, + ) -> list[dict[str, Any]] | None: + _, rank, _ = _dit_rank_world() + if rank != 0: + return None + return prepare_reference_videos( + values, + target_frame_count=target_frame_count, + workdir=workdir, + ) + + def _encode_text_hidden( + self, + input_ids: torch.Tensor, + vision_kwargs: dict[str, torch.Tensor], + ) -> torch.Tensor: + if self.od_config.enable_cpu_offload: + # Invoke nn.Module.__call__ so the generic model-level offloader + # swaps the resident DiT and encoder. + return self.text_encoder(input_ids, **vision_kwargs) + + if self.od_config.enable_layerwise_offload: + # Layerwise DiT offload already provides the low-residency encoder + # phase used by the checkpoint reference. + self.text_encoder.load_to_device() + try: + return self.text_encoder.encode_ids(input_ids, **vision_kwargs) + finally: + self.text_encoder.offload_to_cpu() + + # Keep both Qwen and DiT resident across requests. Moving either model + # here makes encoder latency include a tens-of-gigabytes PCIe transfer, + # which defeats the no-offload contract. + self.text_encoder.load_to_device() + return self.text_encoder.encode_ids(input_ids, **vision_kwargs) + + def _encode_visual_condition( + self, + image: Image.Image, + ) -> torch.Tensor: + _, rank, _ = _dit_rank_world() + rows = self.video_vae.encode_image(image) if rank == 0 else None + return _broadcast_tensor( + rows, + dtype=torch.float32, + device=self.device, + ) + + def _encode_audio_condition( + self, + audio: tuple[torch.Tensor, int], + ) -> tuple[torch.Tensor, int]: + _, rank, _ = _dit_rank_world() + rows = None + audio_t = 0 + if rank == 0: + rows, audio_t = self.audio_vae.encode_waveform(*audio) + audio_t_tensor = torch.tensor( + [audio_t], + dtype=torch.long, + device=self.device, + ) + group, _, world_size = _dit_rank_world() + if world_size > 1: + dist.broadcast(audio_t_tensor, src=0, group=group) + rows = _broadcast_tensor( + rows, + dtype=torch.float32, + device=self.device, + ) + return rows, int(audio_t_tensor.item()) + + def _encode_video_conditions( + self, + prepared_videos: list[dict[str, Any]] | None, + *, + count: int, + ) -> tuple[torch.Tensor, list[tuple[int, int, int]]]: + group, rank, world_size = _dit_rank_world() + distributed_encode = self.video_vae.is_distributed_enabled() + if distributed_encode: + # Native tiled encode uses collectives, so every VPP rank must + # enter each reference encode in the same input order. + prepared_videos_list = [prepared_videos] + dist.broadcast_object_list( + prepared_videos_list, + src=0, + group=group, + device=self.device, + ) + prepared_videos = prepared_videos_list[0] + + rows = None + shapes = torch.zeros((count, 3), dtype=torch.long, device=self.device) + if rank == 0 or distributed_encode: + if prepared_videos is None or len(prepared_videos) != count: + raise ValueError("reference-video preparation is incomplete") + encoded = [ + self.video_vae.encode_video(load_video_frames(item["prepared_path"])) for item in prepared_videos + ] + rows = torch.cat([item[0] for item in encoded]) + shapes = torch.tensor( + [item[1] for item in encoded], + dtype=torch.long, + device=self.device, + ) + if distributed_encode: + return ( + rows.to(device=self.device, dtype=torch.float32), + [tuple(int(value) for value in item) for item in shapes.tolist()], + ) + + if world_size > 1: + dist.broadcast(shapes, src=0, group=group) + return ( + _broadcast_tensor(rows, dtype=torch.float32, device=self.device), + [tuple(int(value) for value in item) for item in shapes.tolist()], + ) + + def _encode_video_audio_conditions( + self, + prepared_videos: list[dict[str, Any]] | None, + *, + has_audio: list[bool], + ) -> tuple[torch.Tensor | None, list[int]]: + _, rank, _ = _dit_rank_world() + count = sum(has_audio) + if count == 0: + return None, [] + rows = None + lengths = torch.zeros(count, dtype=torch.long, device=self.device) + if rank == 0: + if prepared_videos is None: + raise ValueError("rank 0 reference-video preparation is incomplete") + encoded = [ + self.audio_vae.encode_waveform(*load_video_audio(item["original_path"])) + for item in prepared_videos + if item["input_has_audio"] + ] + rows = torch.cat([item[0] for item in encoded]) + lengths = torch.tensor( + [item[1] for item in encoded], + dtype=torch.long, + device=self.device, + ) + group, _, world_size = _dit_rank_world() + if world_size > 1: + dist.broadcast(lengths, src=0, group=group) + return ( + _broadcast_tensor(rows, dtype=torch.float32, device=self.device), + [int(value) for value in lengths.tolist()], + ) + + def _initial_noise( + self, + *, + seed: int, + latent_t: int, + latent_h: int, + latent_w: int, + audio_t: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + video_generator = torch.Generator(device="cpu").manual_seed(seed) + video = torch.randn( + 1, + 24, + latent_t, + latent_h, + latent_w, + generator=video_generator, + dtype=torch.float32, + ) + video_rows = minimax_h3_patchify_video_latent( + video, + patch_size=(1, 2, 2), + ) + audio_generator = torch.Generator(device="cpu").manual_seed(seed) + audio_rows = torch.randn( + audio_t * 2, + 32, + generator=audio_generator, + dtype=torch.float32, + ) + return video_rows, audio_rows + + def diffuse( + self, + *, + task: str, + text_embeddings: torch.Tensor, + text_tags: torch.Tensor, + seed: int, + latent_t: int, + latent_h: int, + latent_w: int, + audio_t: int, + num_frames: int, + num_steps: int, + video_shift: float, + audio_shift: float, + visual_condition: torch.Tensor | None, + visual_condition_shape: tuple[int, int, int] | None, + audio_condition: torch.Tensor | None, + ref_audio_t: int | None, + ref_blocks: list[dict[str, Any]] | None = None, + visual_condition_shapes: list[tuple[int, int, int]] | None = None, + audio_condition_lengths: list[int] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + initial_video, initial_audio = self._initial_noise( + seed=seed, + latent_t=latent_t, + latent_h=latent_h, + latent_w=latent_w, + audio_t=audio_t, + ) + if task == "ref2va": + if ref_blocks is None: + if visual_condition_shape is None or ref_audio_t is None: + raise ValueError("ref2va condition metadata is missing") + _, ref_h, ref_w = visual_condition_shape + ref_blocks = [ + {"kind": "image", "latent_h": ref_h, "latent_w": ref_w}, + {"kind": "audio", "ref_audio_t": ref_audio_t}, + ] + packed = minimax_h3_packed_sequence_ref2va_blocks( + text_len=int(text_embeddings.shape[0]), + latent_t=latent_t, + latent_h=latent_h, + latent_w=latent_w, + audio_t=audio_t, + ref_blocks=ref_blocks, + ) + else: + packed = minimax_h3_packed_sequence( + text_len=int(text_embeddings.shape[0]), + latent_t=latent_t, + latent_h=latent_h, + latent_w=latent_w, + audio_t=audio_t, + include_keyframe_cond=task == "fl2va", + keyframe_frame_indices=[0] if task == "fl2va" else None, + frame_count=num_frames if task == "fl2va" else None, + ) + + tags = packed["token_tags"].clone() + tags[packed["text_pos"]] = text_tags.cpu() + branch = MiniMaxH3DenoiseBranch( + packed=packed, + text_embeddings=text_embeddings, + token_tags=tags, + device=self.device, + ) + + visual_anchor = visual_condition + if visual_anchor is not None: + condition_shapes = visual_condition_shapes + if condition_shapes is None and visual_condition_shape is not None: + condition_shapes = [visual_condition_shape] + if not condition_shapes: + raise ValueError("visual condition shape is missing") + visual_anchor = minimax_h3_imgvid_cond_noise_aug_rows( + visual_anchor, + condition_shapes=condition_shapes, + target_latent_t=latent_t, + imgvid_cond_num_frames=len(condition_shapes), + seed=seed, + noise_aug=MINIMAX_H3_IMGVID_COND_TIMESTEP, + ) + full_video = torch.zeros( + branch.img_pos.shape[0], + 96, + dtype=torch.float32, + ) + full_video[branch.update_mask] = initial_video + initial_video = full_video + + audio_anchor = audio_condition + if audio_anchor is not None: + condition_audio_t = audio_condition_lengths + if condition_audio_t is None and ref_audio_t is not None: + condition_audio_t = [ref_audio_t] + if not condition_audio_t: + raise ValueError("reference audio length is missing") + audio_anchor = minimax_h3_audio_cond_noise_aug_rows( + audio_anchor, + condition_audio_t=condition_audio_t, + seed=seed, + noise_aug=MINIMAX_H3_AUDIO_REF_COND_TIMESTEP, + ) + full_audio = torch.zeros( + branch.audio_pos.shape[0], + 32, + dtype=torch.float32, + ) + full_audio[branch.audio_update_mask] = initial_audio + initial_audio = full_audio + + video_sigmas = minimax_h3_time_shift_sigmas( + num_steps=num_steps, + shift_scale=video_shift, + ) + audio_sigmas = minimax_h3_time_shift_sigmas( + num_steps=num_steps, + shift_scale=audio_shift, + ) + with self.progress_bar(total=len(video_sigmas) - 1) as progress: + video_rows, audio_rows = minimax_h3_denoise_loop( + model=self.transformer, + positive=branch, + initial_video_rows=initial_video, + initial_audio_rows=initial_audio, + keyframe_cond_rows=visual_anchor, + audio_ref_rows=audio_anchor, + sigmas_video=video_sigmas, + sigmas_audio=audio_sigmas, + device=self.device, + imgvid_cond_noise_aug_for_inference=(MINIMAX_H3_IMGVID_COND_TIMESTEP), + audio_cond_noise_aug_for_inference=(MINIMAX_H3_AUDIO_REF_COND_TIMESTEP), + on_step=lambda step, video, audio: progress.update(), + ) + + target_video = video_rows[branch.update_mask_dev] + video_latent = minimax_h3_unpatchify_video_tokens( + target_video, + latent_shape=( + latent_t, + latent_h // 2, + latent_w // 2, + 24, + ), + patch_size=(1, 2, 2), + ) + target_audio = audio_rows[branch.audio_update_mask_dev] + audio_latent = minimax_h3_unpack_audio_tokens( + target_audio, + audio_t=audio_t * 2, + audio_channel=2, + ) + return video_latent, audio_latent + + def decode( + self, + video_latent: torch.Tensor, + audio_latent: torch.Tensor, + *, + height: int, + width: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + with current_omni_platform.create_autocast_context( + device_type=self.device.type, + dtype=torch.float16, + enabled=True, + ): + video = self.video_vae.decode_latent(video_latent) + video = video[..., :height, :width].contiguous() + audio = self.audio_vae.decode_latent(audio_latent) + return video, audio + + @torch.no_grad() + def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput: + if len(request.prompts) != 1: + raise ValueError("MiniMax H3 supports one request at a time") + raw_prompt = request.prompts[0] + if isinstance(raw_prompt, str): + prompt = raw_prompt + multi_modal_data: dict[str, Any] = {} + else: + prompt = str(raw_prompt.get("prompt") or "") + multi_modal_data = raw_prompt.get("multi_modal_data") or {} + if not prompt: + raise ValueError("MiniMax H3 requires a non-empty prompt") + + sampling = request.sampling_params + extra = sampling.extra_args or {} + task = self._resolve_task(extra.get("task"), multi_modal_data) + + raw_image = multi_modal_data.get("image") + raw_videos = multi_modal_data.get("video") + image = _load_image(raw_image) if raw_image is not None else None + if task == "fl2va" and image is None: + raise ValueError(f"{task} requires multi_modal_data.image") + if task == "ref2va" and image is None and raw_videos is None: + raise ValueError("ref2va requires multi_modal_data.image or multi_modal_data.video") + if task == "ref2va" and image is not None and raw_videos is not None: + raise ValueError("ref2va currently accepts image+audio or one or more videos, not both") + if task != "ref2va" and raw_videos is not None: + raise ValueError(f"{task} does not accept a video condition") + if task == "ref2va" and raw_videos is not None and multi_modal_data.get("audio") is not None: + raise ValueError( + "video Ref2VA uses the reference-video soundtracks and does not accept a separate audio condition" + ) + if task == "t2va" and image is not None: + raise ValueError("t2va does not accept an image condition") + + height, width, num_frames, latent_t, audio_t = self._resolve_shape(task, sampling, image) + prepared_image = image + if task == "fl2va" and image is not None: + prepared_image = image.resize( + (width, height), + Image.Resampling.LANCZOS, + ) + elif task == "ref2va" and image is not None: + ref_width, ref_height = _reference_image_shape(image) + prepared_image = image.resize( + (ref_width, ref_height), + Image.Resampling.LANCZOS, + ) + + visual_condition = None + visual_shape = None + visual_shapes = None + audio_condition = None + ref_audio_t = None + audio_lengths = None + ref_blocks = None + with tempfile.TemporaryDirectory(prefix="minimax_h3_ref2va_") as workdir: + prepared_videos = None + has_audio: list[bool] = [] + video_count = 0 + if raw_videos is not None: + video_count = len(raw_videos) if isinstance(raw_videos, (list, tuple)) else 1 + prepared_videos = self._prepare_reference_videos( + raw_videos, + target_frame_count=num_frames, + workdir=workdir, + ) + has_audio_tensor = torch.zeros( + video_count, + dtype=torch.long, + device=self.device, + ) + _, rank, world_size = _dit_rank_world() + if rank == 0: + has_audio_tensor = torch.tensor( + [int(item["input_has_audio"]) for item in prepared_videos or []], + dtype=torch.long, + device=self.device, + ) + if world_size > 1: + dist.broadcast( + has_audio_tensor, + src=0, + group=get_dit_group(), + ) + has_audio = [bool(value) for value in has_audio_tensor.tolist()] + + text_embeddings, text_tags = self.encode_prompt( + task=task, + prompt=prompt, + image=prepared_image, + prepared_videos=prepared_videos, + ) + + if prepared_videos is not None or raw_videos is not None: + visual_condition, visual_shapes = self._encode_video_conditions( + prepared_videos, + count=video_count, + ) + audio_condition, audio_lengths = self._encode_video_audio_conditions( + prepared_videos, + has_audio=has_audio, + ) + audio_iterator = iter(audio_lengths) + ref_blocks = [] + for shape, contributes_audio in zip( + visual_shapes, + has_audio, + strict=True, + ): + ref_audio = next(audio_iterator) if contributes_audio else 0 + ref_blocks.append( + { + "kind": "video", + "ref_audio_t": ref_audio, + "latent_t": shape[0], + "latent_h": shape[1], + "latent_w": shape[2], + } + ) + elif prepared_image is not None: + visual_condition = self._encode_visual_condition(prepared_image) + visual_shape = ( + 1, + prepared_image.height // 16, + prepared_image.width // 16, + ) + + if task == "ref2va" and raw_videos is None: + raw_audio = multi_modal_data.get("audio") + if raw_audio is None: + raise ValueError("image Ref2VA requires multi_modal_data.audio") + audio_condition, ref_audio_t = self._encode_audio_condition(_load_audio(raw_audio)) + elif task != "ref2va" and multi_modal_data.get("audio") is not None: + raise ValueError(f"{task} does not accept an audio condition") + + seed = int(sampling.seed if sampling.seed is not None else 42) + num_steps = int(sampling.num_inference_steps or 50) + video_shift = float(extra.get("flow_shift", self.default_video_shift)) + audio_shift = float(extra.get("audio_flow_shift", self.default_audio_shift)) + video_latent, audio_latent = self.diffuse( + task=task, + text_embeddings=text_embeddings, + text_tags=text_tags, + seed=seed, + latent_t=latent_t, + latent_h=height // 16, + latent_w=width // 16, + audio_t=audio_t, + num_frames=num_frames, + num_steps=num_steps, + video_shift=video_shift, + audio_shift=audio_shift, + visual_condition=visual_condition, + visual_condition_shape=visual_shape, + audio_condition=audio_condition, + ref_audio_t=ref_audio_t, + ref_blocks=ref_blocks, + visual_condition_shapes=visual_shapes, + audio_condition_lengths=audio_lengths, + ) + video, audio = self.decode( + video_latent, + audio_latent, + height=height, + width=width, + ) + return DiffusionOutput( + output=(video, audio), + post_process_func=get_minimax_h3_post_process_func(self.od_config), + stage_durations=(self.stage_durations if hasattr(self, "_stage_durations") else {}), + ) + + +__all__ = [ + "MiniMaxH3Pipeline", + "get_minimax_h3_post_process_func", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/presentation.py b/vllm_omni/diffusion/models/minimax_h3/presentation.py new file mode 100644 index 000000000..a8e142189 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/presentation.py @@ -0,0 +1,297 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax H3 Qwen presentation building. + +Builds the positive presentation token stream: +- fl2va: ': ' label + vision block (<|vision_start|> + + N*<|image_pad|> + <|vision_end|>) + prompt text. +- t2va: prompt text only (no vision block). +Prompt text passes through verbatim (no stripping or rewriting). + +All presentation variants are emitted through the shared ``_Presentation`` +accumulator so ids and AdaLN token tags cannot drift apart. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import torch + +VISION_START = "<|vision_start|>" +VISION_END = "<|vision_end|>" +IMAGE_PAD = "<|image_pad|>" +VIDEO_PAD = "<|video_pad|>" + +_TEXT_TAG = 1 +_VIDEO_TAG = 0 + + +def _text_ids(tokenizer: Any, text: str) -> list[int]: + return list(tokenizer(text, add_special_tokens=False)["input_ids"]) + + +def _vision_block_ids(tokenizer: Any, pad_token: str, count: int) -> list[int]: + return ( + [tokenizer.convert_tokens_to_ids(VISION_START)] + + [tokenizer.convert_tokens_to_ids(pad_token)] * int(count) + + [tokenizer.convert_tokens_to_ids(VISION_END)] + ) + + +class _Presentation: + """Accumulates aligned (ids, token_tags) presentation segments.""" + + def __init__(self) -> None: + self.ids: list[int] = [] + self.tags: list[int] = [] + + def text(self, token_ids: list[int]) -> None: + self.ids += token_ids + self.tags += [_TEXT_TAG] * len(token_ids) + + def vision(self, token_ids: list[int]) -> None: + self.ids += token_ids + self.tags += [_VIDEO_TAG] * len(token_ids) + + def build(self) -> tuple[torch.Tensor, torch.Tensor]: + return ( + torch.tensor(self.ids, dtype=torch.long), + torch.tensor(self.tags, dtype=torch.long), + ) + + +def _timestamped_video_blocks( + presentation: _Presentation, + tokenizer: Any, + *, + counts: Sequence[int], + timestamps: Sequence[float], + context: str, +) -> None: + """Emit per-temporal-block ``<{t:.1f} seconds>`` text + VIDEO vision.""" + + counts = [int(value) for value in counts] + timestamps = [float(value) for value in timestamps] + if not counts or len(counts) != len(timestamps): + raise ValueError(f"{context}video block token counts and timestamps must align") + for count, timestamp in zip(counts, timestamps): + if count <= 0: + raise ValueError(f"{context}video block token count must be positive") + presentation.text(_text_ids(tokenizer, f"<{timestamp:.1f} seconds>")) + presentation.vision(_vision_block_ids(tokenizer, VIDEO_PAD, count)) + + +def minimax_h3_text_only_ids(tokenizer: Any, prompt: str) -> torch.Tensor: + """t2va presentation: verbatim prompt, no special tokens.""" + if not prompt: + raise ValueError("prompt must be non-empty") + return torch.tensor(_text_ids(tokenizer, prompt), dtype=torch.long) + + +def _multi_image_presentation( + tokenizer: Any, + *, + prompt: str, + image_token_counts: list[int], +) -> tuple[torch.Tensor, torch.Tensor]: + if not image_token_counts: + raise ValueError("image_token_counts must be non-empty") + presentation = _Presentation() + for index, count in enumerate(image_token_counts, start=1): + if int(count) <= 0: + raise ValueError("image_token_count must be positive") + presentation.text(_text_ids(tokenizer, f": ")) + presentation.vision(_vision_block_ids(tokenizer, IMAGE_PAD, count)) + presentation.text(_text_ids(tokenizer, prompt)) + return presentation.build() + + +def minimax_h3_multi_image_presentation_ids( + tokenizer: Any, + *, + prompt: str, + image_token_counts: list[int], +) -> torch.Tensor: + """Positive fl2va presentation with one or more labeled image slots.""" + ids, _ = _multi_image_presentation(tokenizer, prompt=prompt, image_token_counts=image_token_counts) + return ids + + +def minimax_h3_multi_image_presentation_token_tags( + tokenizer: Any, + *, + prompt: str, + image_token_counts: list[int], +) -> torch.Tensor: + """AdaLN token tags aligned with minimax_h3_multi_image_presentation_ids.""" + _, tags = _multi_image_presentation(tokenizer, prompt=prompt, image_token_counts=image_token_counts) + return tags + + +def minimax_h3_ref2va_presentation( + tokenizer: Any, + *, + prompt: str, + condition_labels: list[tuple[str, int]], + image_token_count: int | list[int] | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """ref2va positive presentation: + + per condition in request order — image i: ``: `` label followed + by the vision block; audio j: ``