diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md index 9dd1b760e97..1eb1fa866f1 100644 --- a/docs/user_guide/diffusion/attention_backends.md +++ b/docs/user_guide/diffusion/attention_backends.md @@ -36,7 +36,7 @@ For the internal selector, registry, and platform contract, see | `FLASH_ATTN_HUB` | Hub kernel | FlashAttention 2 from Hugging Face `kernels` | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) | | `FLASH_ATTN_3_HUB` | Hub kernel | FlashAttention 3 from Hugging Face `kernels` on Hopper or newer | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) | | `RAINFUSION_ATTN` | Block sparse | MindIE-SD RainFusion video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) | -| `FASTVIDEO_VSA` | Block sparse | FastVideo variable sparse self-attention for FastWan2.2-TI2V-5B and FastH3 MiniMax-H3 on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) | +| `FASTVIDEO_VSA` | Block sparse | VSA routing for FastWan2.2-TI2V-5B and FastH3 MiniMax-H3; compute-provider scope is documented separately | [FastVideo VSA](attention_backends/fastvideo_vsa.md) | ## Configuration diff --git a/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md b/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md index c03b9ffdeaf..db21a7c4b00 100644 --- a/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md +++ b/docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md @@ -20,6 +20,46 @@ dtypes, sequence-parallel execution, or kernel failures fall back to sequence-parallel context is one of those fallbacks; the H3 route supports pure Ulysses and rejects ring or all-gather sequence parallelism at startup. +## SM120: FlashInfer compute with VSA + +Keep `FASTVIDEO_VSA` as the attention backend when following the accelerated +H3 recipe. VSA selects sparse blocks; FlashInfer supplies the attention kernel. +Selecting `FLASHINFER_ATTN`, `SAGE_ATTN` or `SAGE_ATTN_3` instead does not +reproduce that configuration. + +| Configuration | Computation | Numerical behavior | +| --- | --- | --- | +| Standard H3 VSA | FastVideo block-sparse kernel | FP16/BF16 attention | +| H3 VSA with FlashInfer BF16 | FlashInfer block-sparse kernel | BF16 attention; validate output when switching providers | +| Accelerated H3 VSA on SM120 | FlashInfer/CAKE Sage block-sparse kernel | INT8 Q/K, FP8 V, BF16 output; not bitwise equivalent to BF16 attention | + +The accelerated configuration currently belongs to the +[H3 draft integration](https://github.com/vllm-project/vllm-omni/pull/7519). +The [ownership refactor](https://github.com/vllm-project/vllm-omni/pull/7535) +alone does not add FlashInfer/Sage support. FastH3 supports T2VA only, loading +the original model's FL2VA weight partition. + +### Installation status + +A single qualified installation entrypoint for this accelerated configuration +is not yet available. The current draft still requires `fastvideo-kernel` +because of its backend availability check, as well as a compatible FlashInfer +build for the selected compute provider. The final recipe must supply one +validated dependency set; an arbitrary latest FlashInfer wheel or a successful +import does not establish compatibility or the reported performance. + +The Sage path relies on the descriptor-lifetime fix in +[FlashInfer #5127](https://github.com/flashinfer-ai/flashinfer/pull/5127), merged +September 12, 2026. The exact released package or pinned build containing the +required behavior still needs qualification with this integration. + +The current SM120 Sage path requires head dimension 128 and a compatible +64-row sparse layout. Other models and layouts need their own correctness and +quality checks. Choosing a FlashInfer attention kernel does not enable RDMA; +communication setup is a separate choice. The current public draft has no new +full E2E qualification, and historical H3 timings are not a cross-model or +cross-machine speed guarantee. + ## Enable the backend For online serving, select the backend with the existing attention backend @@ -27,7 +67,7 @@ flag. Use `--fastvideo-vsa-topk` to set the number of key/value blocks retained for every query block: ```bash -vllm-omni serve \ +vllm serve --omni \ --diffusion-attention-backend FASTVIDEO_VSA \ --fastvideo-vsa-topk 64 ``` @@ -37,14 +77,14 @@ default `topk=64`: ```bash export DIFFUSION_ATTENTION_BACKEND=FASTVIDEO_VSA -vllm-omni serve +vllm serve --omni ``` To tune top-k, pass the CLI backend and top-k flags together as shown above, or use the equivalent structured configuration: ```bash -vllm-omni serve \ +vllm serve --omni \ --diffusion-attention-config \ '{"default":{"backend":"FASTVIDEO_VSA","fastvideo_vsa_topk":64}}' ``` @@ -75,6 +115,10 @@ stages: ## Choose top-k +The limits below describe the Wan route. H3 retains dense prefix attention and +checkpoint compensation, and its top-k selects video blocks; use the H3 recipe +for that route rather than assuming the Wan all-block fallback rules apply. + At runtime the backend logs the sequence shape and derived block count: ```text diff --git a/recipes/MiniMaxAI/MiniMax-H3-ultra.md b/recipes/MiniMaxAI/MiniMax-H3-ultra.md new file mode 100644 index 00000000000..62f7d70320c --- /dev/null +++ b/recipes/MiniMaxAI/MiniMax-H3-ultra.md @@ -0,0 +1,129 @@ +# MiniMax-H3 acceleration + +The [MiniMax-H3 recipe](MiniMax-H3.md) covers checkpoints and standard serving. +This recipe configures the optional FastH3 execution path: cached AdaLN, +Sage block-sparse attention, MXFP8 DiT and VAE projections, and overlapped +attention/output and video/audio decoding. + +## Requirements + +- Use a vLLM version compatible with the vLLM-Omni checkout. +- MXFP8 requires CUDA Blackwell hardware and PyTorch exposing + `torch.nn.functional.scaled_mm`, `ScalingType.BlockWise1x32`, and + `SwizzleType.SWIZZLE_32_4_4`. +- Sage requires FlashInfer's SM120 block-sparse Sage provider with the + by-value descriptor behavior from + [FlashInfer #5127](https://github.com/flashinfer-ai/flashinfer/pull/5127). + The reproducible software version remains to be pinned and qualified. +- The current `FASTVIDEO_VSA` selector also requires `fastvideo-kernel`, even + for the FlashInfer compute path. +- The optional RDMA adapter requires `flashinfer.comm.UlyssesCommunicator` + with the PCIe backend and registered input/output buffers, supplied by + [FlashInfer #4876](https://github.com/flashinfer-ai/flashinfer/pull/4876). + This dependency is not yet available in a released FlashInfer wheel. +- VAE quantization uses `comfy-kitchen==0.2.33`. MP4 output uses PyAV/libx264. + +The optional compute providers are imported on their execution paths; the +selector still checks `fastvideo-kernel` availability before provider selection. +The complete preset currently requires TP1, Ulysses8, ring1, eight SM120 GPUs, +the VSA/Data-Free adapter, four denoiser evaluations, and top-k 162. Its +communication schedule requires 11,992 local rows, prefix segments `(558, 1206)`, +and a target token grid of `(107, 22, 40)`. It rejects other layouts. +The paired VAE schedule covers 21 temporal windows and 28 spatial tiles per +window. Standard serving remains available for other request geometries. + +## Attention provider on SM120 + +Keep `FASTVIDEO_VSA` in `--diffusion-attention-config`: it selects the VSA +routing contract. This draft's complete preset chooses **FlashInfer/CAKE Sage** +for the fine attention kernel, with INT8 Q/K, FP8 V and BF16 output. It does not +use the default FastVideo BF16 kernel for that computation. Changing the public +backend to `FLASHINFER_ATTN`, `SAGE_ATTN` or `SAGE_ATTN_3` does not reproduce this +configuration. + +The preset selects this compute path automatically. The FlashInfer BF16 +provider is a separate numerical configuration. Neither provider choice implies +FlashInfer Ulysses/RDMA is enabled. Users should not need to set per-kernel +implementation flags to reproduce the final recipe. + +The current public integration has not completed full E2E qualification with a +clean, pinned dependency installation. Historical sub-15-second results do not +qualify this PR revision or a different model, layout, provider or NIC topology. + +## Prepare the model + +This FastH3 adapter supports **T2VA only**. It loads the original H3 `FL2VA` +weight partition; the partition name does not enable first/last-frame tasks +with this adapter. In the server command, `--task-type fl2va` selects those +weights. Requests must set `extra_params.task` to `t2va` and +`num_inference_steps` to `4`. + +Follow the base recipe to download H3 and the FastH3 adapter. Select +`vsa-datafree/adapter_model.safetensors`; the base recipe's `dense-datafree` +example is a different adapter. Set `MODEL_DIR` to the local H3 repository +root and `FASTH3_LORA` to the VSA/Data-Free adapter file. Build the exact, +adapter-bound AdaLN cache: + +```bash +export ADALN_CACHE="$PWD/minimax-h3-t2va-adaln.safetensors" +python tools/minimax_h3/build_adaln_cache.py \ + --transformer-path "${MODEL_DIR}/FL2VA/transformer" \ + --output "${ADALN_CACHE}" \ + --model-variant fl2va --mode t2va \ + --num-inference-steps 4 --flow-shift 12 --audio-flow-shift 3 \ + --fasth3-adapter "${FASTH3_LORA}" \ + --base-schedule 0.999 0.749 0.5 0.25 0.0 +``` + +The cache filename is arbitrary; it does not select an optimization mode. +The sidecar is bound to the adapter fingerprint, task, sigma schedule, and +modality shifts. Startup checks that identity before omitting the cached +AdaLN projections. MXFP8 conversion follows student fusion; it does not +quantize the teacher and then apply a student delta. + +## Configure serving + +The current draft integration still requires `VLLM_OMNI_H3_ULTRA=1` to select +the complete schedule. The split PRs will replace this preset with automatic +selection of validated exact optimizations and explicit backend, precision, +and transport settings. Until that implementation lands, omitting the switch +does not run the same accelerated configuration. Conflicting explicit settings +are rejected. Backend selection, parallelism, and the cache use the existing +Omni interfaces: + +```bash +export VLLM_OMNI_H3_ULTRA=1 +export OMP_NUM_THREADS=28 +export PYTORCH_ALLOC_CONF=expandable_segments:False +export VLLM_WORKER_MULTIPROC_METHOD=spawn + +vllm serve "${MODEL_DIR}/FL2VA" --omni --trust-remote-code \ + --host 127.0.0.1 --port 8093 \ + --task-type fl2va --lora-path "${FASTH3_LORA}" \ + --num-gpus 8 --tensor-parallel-size 1 --usp 8 --ring 1 \ + --ulysses-mode strict --ulysses-a2a-permute \ + --text-encoder-tp-size 8 \ + --vae-patch-parallel-size 8 --vae-parallel-mode tile --vae-use-tiling \ + --cache-backend none \ + --cache-config "{\"minimax_h3_adaln_cache_path\":\"${ADALN_CACHE}\"}" \ + --diffusion-attention-config '{"default":{"backend":"FASTVIDEO_VSA","fastvideo_vsa_topk":162},"per_role":{"minimax_h3.token_refiner":{"backend":"CUDNN_ATTN"}}}' +``` + +Configure GPU/NIC affinity for the host topology. Warm up a resident server +before measuring request latency. This integration has no E2E latency claim +until the combined dependency build and PR commit pass model qualification. + +FastH3/VSA, Sage and MXFP8 change numerical precision or model computation; +the full configuration is not bitwise equivalent to dense BF16 H3. + +## Implementation boundaries + +Device detection, memory queries and synchronization use the existing platform +interfaces. Shared layers provide indexed modulation and MXFP8 producers; +shared attention ops provide Sage quantization and sparse attention. Ulysses +owns the reusable communication and buffer lifecycle. + +`models/minimax_h3/attention/` owns H3 layout metadata, QKV/gate scheduling, +and output projection overlap. `models/minimax_h3/ops/attention/` owns its tile +maps and chunk kernels. VAE schedules, adapter/cache identity and quantized +layer selection also remain in the model directory. diff --git a/tests/diffusion/attention/test_acceleration_ownership.py b/tests/diffusion/attention/test_acceleration_ownership.py new file mode 100644 index 00000000000..60142062141 --- /dev/null +++ b/tests/diffusion/attention/test_acceleration_ownership.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import ast +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +@pytest.mark.parametrize( + "module", + [ + "attention/layer.py", + "attention/parallel/ulysses.py", + "attention/backends/fastvideo_vsa.py", + "attention/ops/sage_block_sparse_attention.py", + "attention/ops/sage_quantization.py", + "distributed/flashinfer_ulysses.py", + "layers/fused_qk_norm_rope.py", + "layers/indexed_modulation.py", + "layers/mxfp8.py", + ], +) +def test_shared_acceleration_does_not_depend_on_model_policy(module): + root = Path(__file__).resolve().parents[3] / "vllm_omni/diffusion" + source = (root / module).read_text() + for node in ast.walk(ast.parse(source)): + if isinstance(node, ast.ImportFrom): + assert not (node.module or "").startswith("vllm_omni.diffusion.models.") + elif isinstance(node, ast.Import): + assert all(not alias.name.startswith("vllm_omni.diffusion.models.") for alias in node.names) + elif isinstance(node, ast.Constant) and isinstance(node.value, str): + assert "VLLM_OMNI_H3_" not in node.value + assert "VLLM_OMNI_MINIMAX_H3_" not in node.value diff --git a/tests/diffusion/attention/test_fastvideo_vsa.py b/tests/diffusion/attention/test_fastvideo_vsa.py index 1bd7192ad7d..7a840133fd9 100644 --- a/tests/diffusion/attention/test_fastvideo_vsa.py +++ b/tests/diffusion/attention/test_fastvideo_vsa.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -import math import sys import types from typing import Any @@ -11,15 +10,10 @@ from vllm_omni.diffusion.attention.backends.abstract import ( AttentionMetadata, - PackedPaddingMetadata, - VideoTokenLayout, - VideoTokenSpan, ) from vllm_omni.diffusion.attention.backends.fastvideo_vsa import ( FastVideoVSABackend, FastVideoVSAImpl, - _build_h3_block_map, - _get_h3_tile_metadata, ) from vllm_omni.diffusion.attention.backends.registry import ( DiffusionAttentionBackendEnum, @@ -216,190 +210,3 @@ def test_fastvideo_vsa_allows_topk_equal_to_num_blocks(): # CPU/float32 is rejected later, but k=N itself must not trigger fallback. assert all_blocks._fallback_reason(query, query, query, metadata) == "dtype torch.float32 is not supported" assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2" - - -def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None): - extra: dict[str, Any] = { - "vsa_h3_prefix_segments": prefix_segments, - } - if gate is not None: - extra["gate_compress"] = gate - return AttentionMetadata( - packed_padding=packed_padding, - video_layout=VideoTokenLayout( - used_len=sum(prefix_segments) + math.prod(video_shape), - video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),), - ), - extra=extra, - ) - - -def _fake_block_sparse_kernel(monkeypatch, calls): - """Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA.""" - - def block_sparse_attn(q, k, v, block_map, variable_block_sizes): - calls["block_map"] = block_map.clone() - calls["variable_block_sizes"] = variable_block_sizes.tolist() - calls["q_shape"] = tuple(q.shape) - blocks = block_map.shape[-1] - assert blocks * 64 == q.shape[2] - # Expand the per-block map to a per-token mask so the stand-in honours - # the sparsity contract the real kernel implements, and drop the rows - # each partial block pads with, which the real kernel skips through - # variable_block_sizes. - token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2) - within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None] - token_mask = token_mask & within.reshape(1, 1, 1, -1) - out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask) - return out.to(q.dtype), None - - module = types.ModuleType("fastvideo_kernel.block_sparse_attn") - setattr(module, "block_sparse_attn", block_sparse_attn) - monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel")) - monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module) - - -def _h3_impl(**backend_kwargs): - kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False} - kwargs.update(backend_kwargs) - return FastVideoVSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs) - - -def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch): - calls: dict[str, Any] = {} - _fake_block_sparse_kernel(monkeypatch, calls) - impl = _h3_impl(topk=64) - prefix_segments, video_shape = (5, 70), (2, 4, 4) - seq_len = sum(prefix_segments) + 2 * 4 * 4 - torch.manual_seed(0) - query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16) - - output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) - - assert output.shape == query.shape - # topk covers every video block, so the block map is dense and the tiled - # kernel must reproduce plain attention over the same rows. - reference = torch.nn.functional.scaled_dot_product_attention( - *(x.float().transpose(1, 2) for x in (query, query, query)) - ).transpose(1, 2) - torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2) - # 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile. - assert calls["variable_block_sizes"] == [5, 64, 6, 32] - - -def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch): - calls: dict[str, Any] = {} - _fake_block_sparse_kernel(monkeypatch, calls) - impl = _h3_impl(topk=1) - prefix_segments, video_shape = (5,), (8, 4, 4) - seq_len = 5 + 8 * 4 * 4 - torch.manual_seed(0) - query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16) - - impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) - - # 1 prefix block + 2 video blocks. The odd count is padded to an even one - # for the paired-CTA kernel contract, and the Triton route slices that - # transport-only partner off again before launching. - block_map = calls["block_map"] - assert block_map.shape[-1] == 3 - assert calls["variable_block_sizes"] == [5, 64, 64] - assert block_map[:, :, :1, :].all(), "prefix queries stay dense" - assert block_map[..., :1].all(), "prefix keys are exempt from selection" - assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query" - - # The same odd geometry must still untile back onto the original rows. - dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) - reference = torch.nn.functional.scaled_dot_product_attention( - *(x.float().transpose(1, 2) for x in (query, query, query)) - ).transpose(1, 2) - torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2) - - -def test_h3_forward_applies_the_learned_compression_gate(monkeypatch): - _fake_block_sparse_kernel(monkeypatch, {}) - impl = _h3_impl(topk=64) - prefix_segments, video_shape = (5,), (2, 4, 4) - seq_len = 5 + 2 * 4 * 4 - torch.manual_seed(0) - query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16) - - without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) - zero_gate = torch.zeros_like(query) - with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate)) - live_gate = torch.full_like(query, 0.5) - with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate)) - - torch.testing.assert_close(with_zero_gate.float(), without_gate.float()) - assert not torch.allclose(with_live_gate.float(), without_gate.float()) - - -def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch): - _fake_block_sparse_kernel(monkeypatch, {}) - impl = _h3_impl(topk=64) - prefix_segments, video_shape = (5,), (2, 4, 4) - valid = 5 + 2 * 4 * 4 - torch.manual_seed(0) - query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16) - padding = PackedPaddingMetadata( - q_length=valid, - kv_length=valid, - cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32), - cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32), - ) - - output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding)) - - assert output.shape == query.shape - assert torch.count_nonzero(output[:, valid:]) == 0 - - -def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch): - # The backend advertises supports_packed_mask_free, so MiniMax-H3 skips - # building the padding mask. Every fallback out of the VSA path must honour - # that contract itself; SDPA reads attn_mask and nothing else. - impl = FastVideoVSAImpl( - num_heads=2, - head_size=8, - softmax_scale=8**-0.5, - backend_kwargs={"topk": 1, "min_seq_len": 4096}, - ) - valid = 40 - torch.manual_seed(0) - query = torch.randn(1, 64, 2, 8) - padding = PackedPaddingMetadata( - q_length=valid, - kv_length=valid, - cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32), - cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32), - ) - - baseline = impl.forward_cuda(query, query, query, AttentionMetadata(packed_padding=padding, extra={})) - perturbed_input = query.clone() - perturbed_input[:, valid:] += 100.0 - perturbed = impl.forward_cuda( - perturbed_input, perturbed_input, perturbed_input, AttentionMetadata(packed_padding=padding, extra={}) - ) - - assert baseline.shape == query.shape - torch.testing.assert_close(baseline[:, :valid], perturbed[:, :valid]) - assert torch.count_nonzero(baseline[:, valid:]) == 0 - - -def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d(): - partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata( - (5, 70, 9), (5, 6, 6), torch.device("cpu") - ) - assert sizes[:4].tolist() == [5, 64, 6, 9] - assert prefix_blocks == 4 - assert video_blocks == 8 - assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6 - assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum()) - - -def test_h3_block_map_makes_prefix_queries_dense_and_prefix_keys_exempt(): - scores = torch.arange(1 * 2 * 5 * 5, dtype=torch.float32).reshape(1, 2, 5, 5) - block_map = _build_h3_block_map(scores, num_prefix_blocks=2, num_video_blocks=3, topk=1) - assert block_map[:, :, :2].all() - assert block_map[..., :2].all() - assert (block_map[:, :, 2:, 2:].sum(dim=-1) == 1).all() diff --git a/tests/diffusion/attention/test_impl_overrides.py b/tests/diffusion/attention/test_impl_overrides.py new file mode 100644 index 00000000000..d1447529ee9 --- /dev/null +++ b/tests/diffusion/attention/test_impl_overrides.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest + +from vllm_omni.diffusion.attention import layer +from vllm_omni.diffusion.attention.backends.fastvideo_vsa import FastVideoVSABackend, FastVideoVSAImpl +from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend +from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention +from vllm_omni.diffusion.data import AttentionSpec + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +class _SpecializedImpl(FastVideoVSAImpl): + pass + + +@pytest.fixture +def local_attention(monkeypatch): + monkeypatch.setattr(layer, "get_current_diffusion_config_or_none", lambda: None) + monkeypatch.setattr(layer, "build_parallel_attention_strategy", lambda **kwargs: NoParallelAttention()) + + +@pytest.mark.parametrize("explicit", [False, True]) +@pytest.mark.parametrize("override", [False, True]) +def test_specialization_preserves_selected_backend_and_options(mocker, local_attention, explicit, override): + spec = AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=7) if explicit else None + select = mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(FastVideoVSABackend, spec)) + attention = layer.Attention( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + causal=False, + role="video.self", + role_category="self", + qkv_layout="BSND", + impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl} if override else None, + ) + + assert attention.attn_backend is FastVideoVSABackend + assert type(attention.attention) is (_SpecializedImpl if override else FastVideoVSAImpl) + assert attention.attention.topk == (7 if explicit else 64) + assert attention.attention.qkv_layout == "BSND" + assert attention.backend_pref == "FASTVIDEO_VSA" + assert attention.backend_explicit is explicit + assert attention.attn_spec is spec + assert isinstance(attention.parallel_strategy, NoParallelAttention) + select.assert_called_once_with( + role="video.self", head_size=8, attention_config=None, role_category="self", allow_trtllm_default=False + ) + + +def test_unselected_specialization_does_not_replace_dense_backend(mocker, local_attention): + mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(SDPABackend, None)) + attention = layer.Attention( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + causal=False, + impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl}, + ) + assert attention.attn_backend is SDPABackend + assert type(attention.attention) is SDPABackend.get_impl_cls() + + +def test_specialization_does_not_bypass_selection_errors(mocker, local_attention): + mocker.patch.object(layer, "get_attn_backend_for_role", side_effect=ImportError("optional kernel unavailable")) + with pytest.raises(ImportError, match="optional kernel unavailable"): + layer.Attention( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + causal=False, + impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl}, + ) + + +def test_specialization_must_extend_selected_implementation(mocker, local_attention): + mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(FastVideoVSABackend, None)) + with pytest.raises(TypeError, match="must subclass the selected implementation"): + layer.Attention( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + causal=False, + impl_overrides={"FASTVIDEO_VSA": SDPABackend.get_impl_cls()}, + ) + + +def test_platform_backend_capabilities_survive_model_specialization(mocker, local_attention): + from vllm_omni.diffusion.models.minimax_h3.attention.vsa import MiniMaxH3VSAImpl + + class PlatformBackend(FastVideoVSABackend): + @classmethod + def supports_packed_mask_free(cls): + return False + + @staticmethod + def get_supported_head_sizes(): + return [128] + + # A registry/platform may refine capabilities without changing the + # implementation. The model must not replace those capability decisions. + mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(PlatformBackend, None)) + attention = layer.Attention( + num_heads=2, + head_size=128, + softmax_scale=128**-0.5, + causal=False, + impl_overrides={"FASTVIDEO_VSA": MiniMaxH3VSAImpl}, + ) + assert attention.attn_backend is PlatformBackend + assert not attention.attn_backend.supports_packed_mask_free() + assert attention.attn_backend.get_supported_head_sizes() == [128] + assert type(attention.attention) is MiniMaxH3VSAImpl + + +def test_incompatible_platform_implementation_is_not_discarded(mocker, local_attention): + class PlatformImpl(FastVideoVSAImpl): + pass + + class PlatformBackend(FastVideoVSABackend): + @staticmethod + def get_impl_cls(): + return PlatformImpl + + mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(PlatformBackend, None)) + with pytest.raises(TypeError, match="selected implementation .*PlatformImpl"): + layer.Attention( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + causal=False, + impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl}, + ) diff --git a/tests/diffusion/attention/test_sage_block_sparse_interface.py b/tests/diffusion/attention/test_sage_block_sparse_interface.py new file mode 100644 index 00000000000..f8e6d9a7d3d --- /dev/null +++ b/tests/diffusion/attention/test_sage_block_sparse_interface.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Structural checks stay usable without loading FlashInfer or a CUDA runtime.""" + +import pytest +import torch + +from vllm_omni.diffusion.attention.ops.sage_block_sparse_attention import ( + sage_block_sparse_attention, + validate_sparse_inputs, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def _inputs(sq=129, sk=191, batch=2, heads=3): + q = torch.empty(batch, sq, heads, 128, dtype=torch.bfloat16) + k = torch.empty(batch, sk, heads, 128, dtype=torch.bfloat16) + indices = torch.zeros(batch, heads, (sq + 63) // 64, 2, dtype=torch.int32) + counts = torch.ones(indices.shape[:3], dtype=torch.int32) + sizes = torch.full(((sk + 63) // 64,), 64, dtype=torch.int32) + return q, k, torch.empty_like(k), indices, counts, sizes + + +@pytest.mark.parametrize("metadata_rank", [0, 1, 2, 3]) +def test_rectangular_ragged_attention_accepts_public_block_size_layouts(metadata_rank): + args = list(_inputs()) + sizes = args[-1] + args[-1] = { + 0: None, + 1: sizes, + 2: sizes.repeat(2, 1), + 3: sizes.repeat(2, 3, 1), + }[metadata_rank] + validate_sparse_inputs(*args) + + +@pytest.mark.parametrize("operand", [0, 1, 2]) +def test_rejects_incorrect_operand_dtype(operand): + args = list(_inputs()) + args[operand] = args[operand].float() + with pytest.raises(TypeError, match="BF16"): + validate_sparse_inputs(*args) + + +@pytest.mark.parametrize("metadata", [3, 4, 5]) +def test_rejects_non_int32_metadata(metadata): + args = list(_inputs()) + args[metadata] = args[metadata].long() + with pytest.raises(TypeError, match="INT32"): + validate_sparse_inputs(*args) + + +def test_rejects_metadata_for_wrong_query_length(): + args = list(_inputs()) + args[3] = args[3][:, :, :-1].contiguous() + with pytest.raises(ValueError, match="Sparse indices"): + validate_sparse_inputs(*args) + + +def test_rejects_wrong_key_block_sizes(): + args = list(_inputs(sq=65, sk=191)) + args[-1] = torch.full((2,), 64, dtype=torch.int32) + with pytest.raises(ValueError, match="Block sizes"): + validate_sparse_inputs(*args) + + +def test_rejects_cross_device_metadata_before_launch(): + args = list(_inputs()) + args[4] = args[4].to("meta") + with pytest.raises(ValueError, match="share one device"): + validate_sparse_inputs(*args) + + +@pytest.mark.parametrize("scale", [0.0, -1.0, float("inf"), float("nan")]) +def test_invalid_scale_fails_before_cuda_import(scale): + with pytest.raises(ValueError, match="finite and positive"): + sage_block_sparse_attention(*_inputs(), scale) + + +def test_prepared_query_cannot_reuse_scales_from_another_length(): + q_int8 = torch.empty(2, 3, 129, 128, dtype=torch.int8) + wrong_scales = torch.empty(2, 3, 4, dtype=torch.float32) + with pytest.raises(ValueError, match="prepared_q must match"): + sage_block_sparse_attention(*_inputs(), 0.125, prepared_q=(q_int8, wrong_scales)) diff --git a/tests/diffusion/attention/test_ulysses_transport_interface.py b/tests/diffusion/attention/test_ulysses_transport_interface.py new file mode 100644 index 00000000000..b75d4fce882 --- /dev/null +++ b/tests/diffusion/attention/test_ulysses_transport_interface.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _strategy(): + group = SimpleNamespace(ulysses_group=object(), ulysses_world_size=2, ulysses_rank=0, ring_world_size=1) + strategy = UlyssesParallelAttention(group, 2, 1, False) + strategy._ulysses_a2a_permute = True + return strategy + + +def test_exchange_keeps_gate_with_query_layout(monkeypatch): + strategy = _strategy() + calls = [] + + def scatter(tensor, slot): + calls.append(slot) + return tensor.repeat_interleave(2, dim=1)[:, :, :2].contiguous() + + monkeypatch.setattr(strategy, "_scatter_heads", scatter) + q = torch.randn(1, 3, 4, 8) + metadata = AttentionMetadata(extra={"gate_compress": q.clone()}) + query, key, value, metadata, ctx = strategy.pre_attention(q, q, q, metadata) + assert calls == ["q", "k", "v", "g"] + assert query.shape == key.shape == value.shape == metadata.extra["gate_compress"].shape == (1, 6, 2, 8) + assert not ctx.use_uaa + + +def test_backend_selection_keeps_optional_dependency_lazy(monkeypatch): + from vllm_omni.diffusion.distributed import flashinfer_ulysses + + def forbidden(): + pytest.fail("default Ulysses must not inspect optional FlashInfer dependencies") + + monkeypatch.setattr(flashinfer_ulysses, "ensure_flashinfer_pcie_available", forbidden) + monkeypatch.setenv("VLLM_OMNI_ULYSSES_A2A_BACKEND", "flashinfer-pcie") + assert not _strategy().supports_qk_input_landing + + +def test_strict_validation_precedes_exchange(monkeypatch): + strategy = _strategy() + monkeypatch.setattr(strategy, "_exchange_qkv", lambda *args: pytest.fail("invalid heads reached transport")) + q = torch.zeros(1, 3, 3, 8) + with pytest.raises(ValueError, match="head_cnt divisible"): + strategy.pre_attention(q, q, q, None) diff --git a/tests/diffusion/cache/test_teacache_extractors.py b/tests/diffusion/cache/test_teacache_extractors.py index d821541d006..fb76b518423 100644 --- a/tests/diffusion/cache/test_teacache_extractors.py +++ b/tests/diffusion/cache/test_teacache_extractors.py @@ -68,12 +68,12 @@ def get_extractor(self): pass @abstractmethod - def get_module(self): + def get_module(self, module, /): """Return model module instance.""" pass @abstractmethod - def get_sample_inputs(self): + def get_sample_inputs(self, sample_inputs): """Return sample inputs for model.""" pass @@ -737,10 +737,10 @@ def test_teacache_strict_sp_keeps_cache_state_rank_local( monkeypatch, ): """Strict SP must not mix local TeaCache state with gathered block rows.""" - from vllm_omni.diffusion.attention.ops import minimax_h3_modulation from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook from vllm_omni.diffusion.distributed import parallel_state + from vllm_omni.diffusion.layers import indexed_modulation as minimax_h3_modulation seq_len = sample_inputs["x"].shape[1] local_len = seq_len // 2 diff --git a/tests/diffusion/models/minimax_h3/test_attention_schedule_interface.py b/tests/diffusion/models/minimax_h3/test_attention_schedule_interface.py new file mode 100644 index 00000000000..3454497aa89 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_attention_schedule_interface.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +from vllm_omni.diffusion.models.minimax_h3.attention import overlap, parallel, qkv_overlap + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def test_split_v_is_replaced_before_transport(monkeypatch): + strategy = object.__new__(parallel.H3UlyssesAttention) + strategy._ulysses_pg = object() + order = [] + query, key, placeholder, produced_v = (torch.tensor(i) for i in range(4)) + + def scatter(tensor, slot): + order.append(slot) + if slot == "v": + assert tensor is produced_v + return tensor + + monkeypatch.setattr(strategy, "_scatter_heads", scatter) + monkeypatch.setattr(overlap, "after_q", lambda *args: order.append("prepare_q")) + monkeypatch.setattr(qkv_overlap, "after_q", lambda: order.append("submit_v")) + + def before_v(q, k, v, metadata, group): + assert q is query and k is key and v is placeholder + order.append("join_v") + return produced_v + + monkeypatch.setattr(overlap, "before_v", before_v) + token = overlap.ACTIVE.set({}) + try: + _, _, result = strategy._exchange_qkv(query, key, placeholder, None) + finally: + overlap.ACTIVE.reset(token) + assert result is produced_v + assert order == ["q", "prepare_q", "submit_v", "k", "join_v", "v"] + + +def test_default_path_uses_generic_exchange(monkeypatch): + strategy = object.__new__(parallel.H3UlyssesAttention) + monkeypatch.setattr(strategy, "_scatter_heads", lambda tensor, slot: tensor) + monkeypatch.setattr(overlap, "after_q", lambda *args: pytest.fail("inactive schedule ran")) + q, k, v = (torch.tensor(i) for i in range(3)) + assert strategy._exchange_qkv(q, k, v, None) == (q, k, v) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py index e10fbe56a9a..7fde01354be 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py @@ -6,7 +6,7 @@ from vllm.triton_utils import HAS_TRITON from tests.helpers.mark import hardware_marks -from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import ( +from vllm_omni.diffusion.layers.indexed_modulation import ( _MAX_1D_GRID_SIZE, _iter_row_chunks, _launch_row_chunks, diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_ultra_schedule.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_ultra_schedule.py new file mode 100644 index 00000000000..58d6b75e404 --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_ultra_schedule.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from collections import Counter + +import pytest + +from vllm_omni.diffusion.models.minimax_h3.vae_batching import jobs + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +def test_paired_vae_has_no_missing_or_repeated_window_tiles(): + rounds = jobs() + observed = Counter(job for assignments in rounds for rank in assignments for job in rank) + assert observed == Counter({(window, tile): 1 for window in range(21) for tile in range(28)}) + assert [sum(len(assignments[rank]) for assignments in rounds) for rank in range(8)] == [74] * 4 + [73] * 4 + # Collective rounds must not expose windows out of temporal order. + assert [sorted({window for rank in assignments for window, _ in rank}) for assignments in rounds] == [ + list(range(first, min(first + 2, 21))) for first in range(0, 21, 2) + ] diff --git a/tests/diffusion/models/minimax_h3/test_vsa_layout.py b/tests/diffusion/models/minimax_h3/test_vsa_layout.py new file mode 100644 index 00000000000..358b060106a --- /dev/null +++ b/tests/diffusion/models/minimax_h3/test_vsa_layout.py @@ -0,0 +1,211 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import math +import sys +import types +from typing import Any + +import pytest +import torch + +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionMetadata, + PackedPaddingMetadata, + VideoTokenLayout, + VideoTokenSpan, +) +from vllm_omni.diffusion.models.minimax_h3.attention.vsa import ( + MiniMaxH3VSAImpl, + _build_h3_block_map, + _get_h3_tile_metadata, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None): + extra: dict[str, Any] = { + "vsa_h3_prefix_segments": prefix_segments, + } + if gate is not None: + extra["gate_compress"] = gate + return AttentionMetadata( + packed_padding=packed_padding, + video_layout=VideoTokenLayout( + used_len=sum(prefix_segments) + math.prod(video_shape), + video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),), + ), + extra=extra, + ) + + +def _fake_block_sparse_kernel(monkeypatch, calls): + """Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA.""" + + def block_sparse_attn(q, k, v, block_map, variable_block_sizes): + calls["block_map"] = block_map.clone() + calls["variable_block_sizes"] = variable_block_sizes.tolist() + calls["q_shape"] = tuple(q.shape) + blocks = block_map.shape[-1] + assert blocks * 64 == q.shape[2] + # Expand the per-block map to a per-token mask so the stand-in honours + # the sparsity contract the real kernel implements, and drop the rows + # each partial block pads with, which the real kernel skips through + # variable_block_sizes. + token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2) + within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None] + token_mask = token_mask & within.reshape(1, 1, 1, -1) + out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask) + return out.to(q.dtype), None + + module = types.ModuleType("fastvideo_kernel.block_sparse_attn") + setattr(module, "block_sparse_attn", block_sparse_attn) + monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel")) + monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module) + + +def _h3_impl(**backend_kwargs): + kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False} + kwargs.update(backend_kwargs) + return MiniMaxH3VSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs) + + +def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch): + calls: dict[str, Any] = {} + _fake_block_sparse_kernel(monkeypatch, calls) + impl = _h3_impl(topk=64) + prefix_segments, video_shape = (5, 70), (2, 4, 4) + seq_len = sum(prefix_segments) + 2 * 4 * 4 + torch.manual_seed(0) + query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16) + + output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) + + assert output.shape == query.shape + # topk covers every video block, so the block map is dense and the tiled + # kernel must reproduce plain attention over the same rows. + reference = torch.nn.functional.scaled_dot_product_attention( + *(x.float().transpose(1, 2) for x in (query, query, query)) + ).transpose(1, 2) + torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2) + # 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile. + assert calls["variable_block_sizes"] == [5, 64, 6, 32] + + +def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch): + calls: dict[str, Any] = {} + _fake_block_sparse_kernel(monkeypatch, calls) + impl = _h3_impl(topk=1) + prefix_segments, video_shape = (5,), (8, 4, 4) + seq_len = 5 + 8 * 4 * 4 + torch.manual_seed(0) + query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16) + + impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) + + # 1 prefix block + 2 video blocks. The odd count is padded to an even one + # for the paired-CTA kernel contract, and the Triton route slices that + # transport-only partner off again before launching. + block_map = calls["block_map"] + assert block_map.shape[-1] == 3 + assert calls["variable_block_sizes"] == [5, 64, 64] + assert block_map[:, :, :1, :].all(), "prefix queries stay dense" + assert block_map[..., :1].all(), "prefix keys are exempt from selection" + assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query" + + # The same odd geometry must still untile back onto the original rows. + dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) + reference = torch.nn.functional.scaled_dot_product_attention( + *(x.float().transpose(1, 2) for x in (query, query, query)) + ).transpose(1, 2) + torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2) + + +def test_h3_forward_applies_the_learned_compression_gate(monkeypatch): + _fake_block_sparse_kernel(monkeypatch, {}) + impl = _h3_impl(topk=64) + prefix_segments, video_shape = (5,), (2, 4, 4) + seq_len = 5 + 2 * 4 * 4 + torch.manual_seed(0) + query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16) + + without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape)) + zero_gate = torch.zeros_like(query) + with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate)) + live_gate = torch.full_like(query, 0.5) + with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate)) + + torch.testing.assert_close(with_zero_gate.float(), without_gate.float()) + assert not torch.allclose(with_live_gate.float(), without_gate.float()) + + +def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch): + _fake_block_sparse_kernel(monkeypatch, {}) + impl = _h3_impl(topk=64) + prefix_segments, video_shape = (5,), (2, 4, 4) + valid = 5 + 2 * 4 * 4 + torch.manual_seed(0) + query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16) + padding = PackedPaddingMetadata( + q_length=valid, + kv_length=valid, + cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32), + cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32), + ) + + output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding)) + + assert output.shape == query.shape + assert torch.count_nonzero(output[:, valid:]) == 0 + + +def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch): + # The backend advertises supports_packed_mask_free, so MiniMax-H3 skips + # building the padding mask. Every fallback out of the VSA path must honour + # that contract itself; SDPA reads attn_mask and nothing else. + impl = MiniMaxH3VSAImpl( + num_heads=2, + head_size=8, + softmax_scale=8**-0.5, + backend_kwargs={"topk": 1, "min_seq_len": 4096}, + ) + monkeypatch.setattr(impl.sdpa_fallback, "forward", impl.sdpa_fallback.forward_cuda) + valid = 40 + torch.manual_seed(0) + query = torch.randn(1, 64, 2, 8) + padding = PackedPaddingMetadata( + q_length=valid, + kv_length=valid, + cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32), + cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32), + ) + + baseline = impl.forward_cuda(query, query, query, AttentionMetadata(packed_padding=padding, extra={})) + perturbed_input = query.clone() + perturbed_input[:, valid:] += 100.0 + perturbed = impl.forward_cuda( + perturbed_input, perturbed_input, perturbed_input, AttentionMetadata(packed_padding=padding, extra={}) + ) + + assert baseline.shape == query.shape + torch.testing.assert_close(baseline[:, :valid], perturbed[:, :valid]) + assert torch.count_nonzero(baseline[:, valid:]) == 0 + + +def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d(): + partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata( + (5, 70, 9), (5, 6, 6), torch.device("cpu") + ) + assert sizes[:4].tolist() == [5, 64, 6, 9] + assert prefix_blocks == 4 + assert video_blocks == 8 + assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6 + assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum()) + + +def test_h3_block_map_makes_prefix_queries_dense_and_prefix_keys_exempt(): + scores = torch.arange(1 * 2 * 5 * 5, dtype=torch.float32).reshape(1, 2, 5, 5) + block_map = _build_h3_block_map(scores, num_prefix_blocks=2, num_video_blocks=3, topk=1) + assert block_map[:, :, :2].all() + assert block_map[..., :2].all() + assert (block_map[:, :, 2:, 2:].sum(dim=-1) == 1).all() diff --git a/tests/diffusion/quantization/test_mxfp8_producer_interface.py b/tests/diffusion/quantization/test_mxfp8_producer_interface.py new file mode 100644 index 00000000000..89d78ca51fc --- /dev/null +++ b/tests/diffusion/quantization/test_mxfp8_producer_interface.py @@ -0,0 +1,36 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +from vllm_omni.diffusion.layers.mxfp8 import mxfp8_linear, mxfp8_quantize_project, silu_mxfp8_linear + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@pytest.mark.parametrize("rows,hidden,out_features", [(1, 32, 128), (137, 96, 256), (256, 128, 384)]) +def test_export_keeps_quantized_input_and_scale_as_explicit_outputs(rows, hidden, out_features): + class Projection(torch.nn.Module): + def forward(self, x, weight, scale): + return mxfp8_quantize_project(x, weight, scale) + + x = torch.empty(rows, hidden, dtype=torch.bfloat16, device="meta") + weight = torch.empty(out_features, hidden, dtype=torch.float8_e4m3fn, device="meta") + scale = torch.empty(4096, dtype=torch.float8_e8m0fnu, device="meta") + exported = torch.export.export(Projection(), (x, weight, scale)) + output, payload, activation_scale = exported.module()(x, weight, scale) + assert output.shape == (rows, out_features) and output.dtype == torch.bfloat16 + assert payload.shape == x.shape and payload.dtype == torch.float8_e4m3fn + assert activation_scale.shape == (((rows + 127) // 128) * 128 * ((hidden // 32 + 3) // 4) * 4,) + assert activation_scale.dtype == torch.uint8 + assert len(exported.graph_signature.output_specs) == 3 + + +@pytest.mark.parametrize("operation", [mxfp8_linear, silu_mxfp8_linear]) +def test_projection_fake_preserves_leading_dimensions(operation): + x = torch.empty(2, 3, 64, dtype=torch.bfloat16, device="meta") + weight = torch.empty(128, 64 if operation is mxfp8_linear else 32, dtype=torch.float8_e4m3fn, device="meta") + scale = torch.empty(512, dtype=torch.float8_e8m0fnu, device="meta") + output = operation(x, weight, scale) + assert output.shape == (2, 3, 128) and output.dtype == torch.bfloat16 diff --git a/tests/diffusion/quantization/test_mxfp8_producers_cuda.py b/tests/diffusion/quantization/test_mxfp8_producers_cuda.py new file mode 100644 index 00000000000..e803cf16324 --- /dev/null +++ b/tests/diffusion/quantization/test_mxfp8_producers_cuda.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +from vllm_omni.diffusion.layers.mxfp8 import ( + _scale_numel, + mxfp8_quantize_project, + mxfp8_quantize_swizzled, + mxfp8_scaled_mm, + silu_mul_mxfp8_cuda, +) + +pytestmark = [pytest.mark.local_model, pytest.mark.cuda, pytest.mark.diffusion] + + +@pytest.fixture(autouse=True) +def sm120_only(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (12, 0): + pytest.skip("requires SM120") + if not hasattr(torch.nn.functional, "scaled_mm"): + pytest.skip("requires the public PyTorch block-scaled GEMM API") + + +@pytest.mark.parametrize("rows,hidden", [(1, 32), (127, 96), (137, 128), (256, 512), (11992, 14336)]) +def test_swiglu_quantization_preserves_cuda_activation_bytes(rows, hidden): + from vllm.model_executor.layers.activation import SiluAndMul + + generator = torch.Generator(device="cuda").manual_seed(1101) + x = torch.randn(rows, 2 * hidden, dtype=torch.bfloat16, device="cuda", generator=generator) + reference = SiluAndMul().forward_cuda(x) + expected_q, expected_scale = mxfp8_quantize_swizzled(reference) + q, scale, activation = silu_mul_mxfp8_cuda(x, return_bf16=True) + assert torch.equal(activation.view(torch.uint8), reference.view(torch.uint8)) + assert torch.equal(q.view(torch.uint8), expected_q.view(torch.uint8)) + assert torch.equal(scale, expected_scale) + + +@pytest.mark.parametrize( + "rows,hidden,width,split", [(137, 128, 384, 256), (256, 512, 384, 256), (11992, 5376, 21504, 14336)] +) +def test_weight_slice_reuses_one_quantized_activation(rows, hidden, width, split): + generator = torch.Generator(device="cuda").manual_seed(1101) + x = torch.randn(rows, hidden, dtype=torch.bfloat16, device="cuda", generator=generator) + weight = torch.randn(width, hidden, dtype=x.dtype, device=x.device, generator=generator) + wq, ws = mxfp8_quantize_swizzled(weight) + cut = _scale_numel(split, hidden) + first, q, scale = mxfp8_quantize_project(x, wq[:split], ws[:cut]) + second = mxfp8_scaled_mm(q, wq[split:], scale, ws[cut:]) + full = mxfp8_scaled_mm(q, wq, scale, ws) + assert torch.equal(first.view(torch.uint8), full[:, :split].contiguous().view(torch.uint8)) + assert torch.equal(second.view(torch.uint8), full[:, split:].contiguous().view(torch.uint8)) diff --git a/tests/diffusion/test_indexed_gate.py b/tests/diffusion/test_indexed_gate.py new file mode 100644 index 00000000000..5d84c4bc2ba --- /dev/null +++ b/tests/diffusion/test_indexed_gate.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +from vllm_omni.diffusion.layers.indexed_modulation import bf16_indexed_gate_add_ +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda] + + +@pytest.mark.parametrize("batch,rows,heads,dim,tail", [(2, 17, 3, 32, 5), (1, 2998, 56, 128, 270)]) +def test_indexed_gate_preserves_bf16_rounding_and_coarse_tail(batch, rows, heads, dim, tail): + if not current_omni_platform.is_cuda(): + pytest.skip("CUDA BF16 arithmetic required") + device = current_omni_platform.get_torch_device() + generator = torch.Generator(device=device).manual_seed(27) + bundle = torch.randn(batch, rows + tail, heads, dim, device=device, dtype=torch.bfloat16, generator=generator) + gate = torch.randn(batch, rows, heads, dim, device=device, dtype=torch.bfloat16, generator=generator) + indices = torch.arange(rows, device=device, dtype=torch.int32) % tail + indices[::7] = -1 + before = bundle.clone() + valid = indices >= 0 + expected = before[:, :rows].clone() + expected[:, valid] += before[:, rows:].index_select(1, indices[valid].long()) * gate[:, valid] + actual = bf16_indexed_gate_add_(bundle, gate, indices) + assert torch.equal(actual, expected) + assert torch.equal(bundle[:, rows:], before[:, rows:]) diff --git a/tools/minimax_h3/build_adaln_cache.py b/tools/minimax_h3/build_adaln_cache.py new file mode 100644 index 00000000000..4ca4e72975b --- /dev/null +++ b/tools/minimax_h3/build_adaln_cache.py @@ -0,0 +1,410 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Build a fixed-schedule MiniMax H3 AdaLN output sidecar. + +Base H3 emits the SGLang-compatible v2 format. Supplying a FastH3 adapter emits +v3: the adapter's timestep and AdaLN edits are fused before projection, and the +result is bound to the exact adapter and four-step sampling contract. +""" + +from __future__ import annotations + +import argparse +import json +import math +from contextlib import ExitStack +from pathlib import Path +from typing import Any + +import torch +import torch.nn.functional as F +from safetensors.torch import safe_open, save_file + +from vllm_omni.diffusion.models.minimax_h3.adaln_cache import ( + MiniMaxH3AdalnCacheBinding, + fingerprint_minimax_h3_fasth3_adapter, +) +from vllm_omni.diffusion.models.minimax_h3.denoise_loop import ( + MINIMAX_H3_AUDIO_REF_COND_TIMESTEP, + MINIMAX_H3_IMGVID_COND_TIMESTEP, +) +from vllm_omni.diffusion.models.minimax_h3.time_request import ( + minimax_h3_time_shift_sigmas, +) +from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule + +_HIDDEN_SIZE = 5376 +_TIMESTEP_INPUT_DIM = 256 +_NUM_BLOCKS = 50 +_NUM_REFINER_BLOCKS = 2 +_ATTENTION_HEAD_DIM = 128 +_BLOCK_PARAM_WIDTH = 18 * _HIDDEN_SIZE +_FINAL_PARAM_WIDTH = 2 * _HIDDEN_SIZE +_CACHE_MODES = { + "t2va": ("video", "audio"), + "fl2va": ("video", "audio", "image"), + "ref2va-image": ("video", "audio", "image"), + "ref2va-audio": ("video", "audio", "audio_ref"), + "ref2va-mixed": ("video", "audio", "image", "audio_ref"), +} +_MODE_VARIANTS = { + "t2va": "fl2va", + "fl2va": "fl2va", + "ref2va-image": "ref2va", + "ref2va-audio": "ref2va", + "ref2va-mixed": "ref2va", +} + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=("Build a fixed-schedule MiniMax H3 AdaLN cache from local native transformer weights.") + ) + parser.add_argument("--transformer-path", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--model-variant", choices=("fl2va", "ref2va"), required=True) + parser.add_argument("--mode", choices=tuple(_CACHE_MODES), default="t2va") + parser.add_argument("--num-inference-steps", type=int, default=50) + parser.add_argument("--flow-shift", type=float, default=12.0) + parser.add_argument("--audio-flow-shift", type=float, default=3.0) + parser.add_argument( + "--imgvid-cond-noise-aug", + type=float, + default=MINIMAX_H3_IMGVID_COND_TIMESTEP, + ) + parser.add_argument( + "--audio-cond-noise-aug", + type=float, + default=MINIMAX_H3_AUDIO_REF_COND_TIMESTEP, + ) + parser.add_argument( + "--timesteps", + type=float, + nargs="+", + help="Override the scheduler-derived timestep plan with one exact plan.", + ) + parser.add_argument( + "--fasth3-adapter", + type=Path, + help=( + "FastVideo FastH3 adapter file or variant directory. Its edits to " + "the timestep and AdaLN projections are fused before building a " + "format-v3 sidecar." + ), + ) + parser.add_argument( + "--base-schedule", + type=float, + nargs="+", + help=( + "Exact pre-shift rectified-flow positions. Required with " + "--fasth3-adapter; FastH3 Preview v1 uses " + "0.999 0.749 0.5 0.25 0.0." + ), + ) + parser.add_argument("--device", default="cuda") + return parser.parse_args() + + +def _cache_timestep_plans(args: argparse.Namespace) -> list[torch.Tensor]: + if args.timesteps is not None: + return [torch.tensor(args.timesteps, dtype=torch.float32).unique(sorted=True)] + + if args.base_schedule is None: + video_sigmas = minimax_h3_time_shift_sigmas( + num_steps=args.num_inference_steps, + shift_scale=args.flow_shift, + ) + audio_sigmas = minimax_h3_time_shift_sigmas( + num_steps=args.num_inference_steps, + shift_scale=args.audio_flow_shift, + ) + else: + # An explicit distilled schedule owns its cardinality. In particular, + # FastH3's five positions mean four transformer evaluations regardless + # of the legacy base-H3 --num-inference-steps default (50). + schedule = DMD2SigmaSchedule.from_positions(args.base_schedule) + video_sigmas = schedule.shifted_sigmas(args.flow_shift) + audio_sigmas = schedule.shifted_sigmas(args.audio_flow_shift) + fields = _CACHE_MODES[args.mode] + plans = [] + for video_sigma, audio_sigma in zip(video_sigmas[:-1], audio_sigmas[:-1]): + video_timestep = 1.0 - video_sigma + audio_timestep = 1.0 - audio_sigma + candidates = { + "video": video_timestep, + "audio": audio_timestep, + "image": max(video_timestep, args.imgvid_cond_noise_aug), + "audio_ref": max(audio_timestep, args.audio_cond_noise_aug), + } + plans.append(torch.tensor([candidates[field] for field in fields], dtype=torch.float32).unique(sorted=True)) + + deduplicated = [] + seen = set() + for plan in plans: + key = tuple(plan.tolist()) + if key not in seen: + seen.add(key) + deduplicated.append(plan) + return deduplicated + + +def _time_embed( + timesteps: torch.Tensor, + *, + proj_in_weight: torch.Tensor, + proj_in_bias: torch.Tensor, + proj_out_weight: torch.Tensor, + proj_out_bias: torch.Tensor, +) -> torch.Tensor: + half = _TIMESTEP_INPUT_DIM // 2 + freqs = torch.exp(-math.log(10000.0) * torch.arange(half, dtype=torch.float32, device=timesteps.device) / half) + args = timesteps[:, None] * freqs[None] + t_freq = torch.cat((torch.cos(args), torch.sin(args)), dim=-1) + return F.linear( + F.silu(F.linear(t_freq, proj_in_weight, proj_in_bias)), + proj_out_weight, + proj_out_bias, + ) + + +def _load_tensor( + name: str, + *, + weight_map: dict[str, str], + files: dict[str, Any], + device: torch.device, +) -> torch.Tensor: + tensor_file = files[weight_map[name]] + return tensor_file.get_tensor(name).to(device) + + +def _validate_build_contract(args: argparse.Namespace) -> DMD2SigmaSchedule | None: + """Validate the mutually exclusive base-H3 and FastH3 build contracts.""" + if args.num_inference_steps < 2 and args.timesteps is None: + raise ValueError("--num-inference-steps must be at least 2") + mode_variant = _MODE_VARIANTS[args.mode] + if args.model_variant != mode_variant: + raise ValueError(f"--mode {args.mode} requires {mode_variant}") + if not math.isfinite(args.flow_shift) or args.flow_shift <= 0: + raise ValueError("--flow-shift must be finite and positive") + if not math.isfinite(args.audio_flow_shift) or args.audio_flow_shift <= 0: + raise ValueError("--audio-flow-shift must be finite and positive") + + has_adapter = args.fasth3_adapter is not None + has_schedule = args.base_schedule is not None + if has_adapter != has_schedule: + raise ValueError("--fasth3-adapter and --base-schedule must be supplied together") + if not has_adapter: + return None + if args.timesteps is not None: + raise ValueError("--timesteps cannot override an adapter-bound FastH3 base schedule") + if args.model_variant != "fl2va" or args.mode != "t2va": + raise ValueError("FastH3 Preview v1 AdaLN sidecars require --model-variant fl2va and --mode t2va") + assert args.base_schedule is not None + return DMD2SigmaSchedule.from_positions(args.base_schedule) + + +def _load_fasth3_fusion( + args: argparse.Namespace, + *, + schedule: DMD2SigmaSchedule | None, +): + """Load and validate the canonical FastH3 fusion only when requested.""" + if args.fasth3_adapter is None: + return None, None + assert schedule is not None + from vllm_omni.diffusion.models.minimax_h3.fasth3 import FastH3WeightFusion + + fusion = FastH3WeightFusion.from_path( + args.fasth3_adapter, + head_dim=_ATTENTION_HEAD_DIM, + num_blocks=_NUM_BLOCKS, + num_refiner_blocks=_NUM_REFINER_BLOCKS, + ) + if fusion is None: + raise ValueError(f"--fasth3-adapter is not a recognized FastH3 release: {args.fasth3_adapter}") + if schedule.base_schedule != fusion.base_schedule: + raise ValueError( + "--base-schedule does not match the FastH3 release ladder: " + f"{schedule.base_schedule!r} != {fusion.base_schedule!r}" + ) + identity = fingerprint_minimax_h3_fasth3_adapter(fusion.source) + return fusion, identity + + +def _fuse_required_tensor( + name: str, + tensor: torch.Tensor, + *, + fusion, +) -> torch.Tensor: + """Fuse a FastH3 patch and reject a sidecar built from an unedited weight.""" + if fusion is None: + return tensor + fused = fusion.fuse(name, tensor) + if fused is tensor: + raise ValueError(f"FastH3 adapter does not patch required AdaLN input {name!r}") + return fused + + +def main() -> None: + args = _parse_args() + schedule = _validate_build_contract(args) + device = torch.device(args.device) + if device.type != "cuda" or not torch.cuda.is_available(): + raise ValueError("MiniMax H3 AdaLN cache must be built on CUDA") + + index_path = args.transformer_path / "model.safetensors.index.json" + with index_path.open(encoding="utf-8") as index_file: + weight_map = json.load(index_file)["weight_map"] + + fasth3_fusion, fasth3_identity = _load_fasth3_fusion( + args, + schedule=schedule, + ) + + plans = _cache_timestep_plans(args) + if not plans or any(plan.numel() == 0 for plan in plans): + raise ValueError("AdaLN cache must cover at least one timestep plan") + max_plan_length = max(plan.numel() for plan in plans) + plan_timesteps = torch.zeros((len(plans), max_plan_length), dtype=torch.float32) + plan_lengths = torch.tensor([plan.numel() for plan in plans], dtype=torch.int64) + block_params = torch.empty( + (len(plans), max_plan_length, _NUM_BLOCKS, _BLOCK_PARAM_WIDTH), + dtype=torch.bfloat16, + ) + final_params = torch.empty( + (len(plans), max_plan_length, _FINAL_PARAM_WIDTH), + dtype=torch.bfloat16, + ) + + print( + f"Building {len(plans)} AdaLN plans (width <= {max_plan_length}); " + f"sidecar tensors use " + f"{(block_params.numel() + final_params.numel()) * 2 / 2**30:.2f} GiB" + ) + with ExitStack() as stack: + files = { + filename: stack.enter_context( + safe_open( + str(args.transformer_path / filename), + framework="pt", + device="cpu", + ) + ) + for filename in set(weight_map.values()) + } + time_kwargs = {} + for module, name in ( + ("proj_in", "weight"), + ("proj_in", "bias"), + ("proj_out", "weight"), + ("proj_out", "bias"), + ): + parameter_name = f"time_embedder.{module}.{name}" + time_kwargs[f"{module}_{name}"] = _fuse_required_tensor( + parameter_name, + _load_tensor( + parameter_name, + weight_map=weight_map, + files=files, + device=device, + ), + fusion=fasth3_fusion, + ) + adaln_inputs = [] + for plan_index, plan in enumerate(plans): + plan_length = plan.numel() + plan_timesteps[plan_index, :plan_length].copy_(plan) + adaln_inputs.append(F.silu(_time_embed(plan.to(device), **time_kwargs)).to(torch.bfloat16)) + + for index in range(_NUM_BLOCKS): + prefix = f"blocks.{index}.adaln_proj.linear" + weight_name = f"{prefix}.weight" + bias_name = f"{prefix}.bias" + weight = _fuse_required_tensor( + weight_name, + _load_tensor( + weight_name, + weight_map=weight_map, + files=files, + device=device, + ), + fusion=fasth3_fusion, + ) + bias = _fuse_required_tensor( + bias_name, + _load_tensor( + bias_name, + weight_map=weight_map, + files=files, + device=device, + ), + fusion=fasth3_fusion, + ) + for plan_index, adaln_input in enumerate(adaln_inputs): + plan_length = adaln_input.shape[0] + block_params[plan_index, :plan_length, index].copy_(F.linear(adaln_input, weight, bias).cpu()) + del weight, bias + if index == 0 or (index + 1) % 5 == 0: + print(f"Projected AdaLN block {index + 1}/{_NUM_BLOCKS}") + + prefix = "final_layer.adaln_proj.linear" + weight_name = f"{prefix}.weight" + bias_name = f"{prefix}.bias" + weight = _fuse_required_tensor( + weight_name, + _load_tensor( + weight_name, + weight_map=weight_map, + files=files, + device=device, + ), + fusion=fasth3_fusion, + ) + bias = _fuse_required_tensor( + bias_name, + _load_tensor( + bias_name, + weight_map=weight_map, + files=files, + device=device, + ), + fusion=fasth3_fusion, + ) + for plan_index, adaln_input in enumerate(adaln_inputs): + plan_length = adaln_input.shape[0] + final_params[plan_index, :plan_length].copy_(F.linear(adaln_input, weight, bias).cpu()) + + args.output.parent.mkdir(parents=True, exist_ok=True) + if fasth3_identity is None: + binding = MiniMaxH3AdalnCacheBinding( + format_version="2", + model_variant=args.model_variant, + ) + else: + assert schedule is not None + binding = MiniMaxH3AdalnCacheBinding.for_fasth3( + adapter=fasth3_identity, + model_variant=args.model_variant, + mode=args.mode, + base_schedule=schedule.base_schedule, + flow_shift=args.flow_shift, + audio_flow_shift=args.audio_flow_shift, + ) + save_file( + { + "plan_timesteps": plan_timesteps, + "plan_lengths": plan_lengths, + "block_params": block_params, + "final_params": final_params, + }, + str(args.output), + metadata=binding.to_metadata(), + ) + print(f"Wrote MiniMax H3 AdaLN cache v{binding.format_version}: {args.output}") + + +if __name__ == "__main__": + main() diff --git a/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py b/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py index 465912d3942..0d0cdc8f7c2 100644 --- a/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py +++ b/vllm_omni/diffusion/attention/backends/fastvideo_vsa.py @@ -6,7 +6,6 @@ import functools import importlib.util import math -import os from collections.abc import Mapping from typing import Any @@ -88,77 +87,6 @@ def _( _fastvideo_vsa_bshd_op = torch.ops.vllm_omni.fastvideo_vsa_bshd -# H3 needs an explicit per-query block map: prefix queries are dense, while -# video queries select prefix + top-k video tiles. The generic -# video_sparse_attn() entry point cannot express that contract. -if not hasattr(torch.ops.vllm_omni, "fastvideo_h3_vsa_bhsd"): - - @torch.library.custom_op("vllm_omni::fastvideo_h3_vsa_bhsd", mutates_args=()) - def _fastvideo_h3_vsa_bhsd_op( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - block_map: torch.Tensor, - variable_block_sizes: torch.Tensor, - logical_blocks: int, - ) -> torch.Tensor: - q = query.transpose(1, 2).contiguous() - k = key.transpose(1, 2).contiguous() - v = value.transpose(1, 2).contiguous() - - # Official FastVideo's measured FastH3 route opts into this native - # Blackwell forward. Wheels without the extension (including SM103 - # builds today) retain the corrected explicit-mask Triton route. - if os.environ.get("FASTVIDEO_VSA_SM100A", "0") == "1": - try: - from fastvideo_kernel import block_sparse_attn_sm100a - from fastvideo_kernel.triton_kernels.index import map_to_index - - if block_sparse_attn_sm100a.is_supported(q, variable_block_sizes): - q2k_idx, q2k_num = map_to_index(block_map) - out, _ = block_sparse_attn_sm100a.block_sparse_attn_sm100a( - q, - k, - v, - q2k_idx.to(torch.int32).contiguous(), - q2k_num.to(torch.int32).contiguous(), - variable_block_sizes.to(torch.int32).contiguous(), - need_lse=False, - ) - return out.transpose(1, 2).contiguous() - except (ImportError, RuntimeError) as exc: - # Opting in explicitly and then silently getting a different - # numeric path is worse than the slower route it lands on. - logger.warning_once( - "FASTVIDEO_VSA_SM100A=1 requested but the native Blackwell forward is " - "unavailable (%s); using the Triton block-sparse route instead.", - exc, - ) - - from fastvideo_kernel.block_sparse_attn import block_sparse_attn - - logical_len = logical_blocks * 64 - out, _ = block_sparse_attn( - q[:, :, :logical_len].contiguous(), - k[:, :, :logical_len].contiguous(), - v[:, :, :logical_len].contiguous(), - block_map[..., :logical_blocks, :logical_blocks].contiguous(), - variable_block_sizes[:logical_blocks].to(torch.int32).contiguous(), - ) - out = out.transpose(1, 2).contiguous() - if out.shape[1] != query.shape[1]: - out = torch.nn.functional.pad(out, (0, 0, 0, 0, 0, query.shape[1] - out.shape[1])) - return out - - @_fastvideo_h3_vsa_bhsd_op.register_fake - def _(query, key, value, block_map, variable_block_sizes, logical_blocks): - del key, value, block_map, variable_block_sizes, logical_blocks - return torch.empty_like(query) - - -_fastvideo_h3_vsa_bhsd_op = torch.ops.vllm_omni.fastvideo_h3_vsa_bhsd - - @functools.lru_cache(maxsize=32) def _get_tile_partition_indices( dit_seq_shape: tuple[int, int, int], @@ -226,81 +154,6 @@ def _get_tile_metadata( return tile_partition_indices, variable_block_sizes, non_pad_index, untile_combined_index -@functools.lru_cache(maxsize=32) -def _get_h3_tile_metadata( - prefix_segments: tuple[int, ...], - video_shape: tuple[int, int, int], - device: torch.device, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]: - """Official FastVideo H3 geometry: pure prefix chunks + 3-D video tiles.""" - block_size = (4, 4, 4) - block_elements = 64 - prefix_len = sum(prefix_segments) - prefix_sizes: list[int] = [] - for segment in prefix_segments: - full, remainder = divmod(segment, block_elements) - prefix_sizes.extend([block_elements] * full) - if remainder: - prefix_sizes.append(remainder) - - video_indices = _get_tile_partition_indices(video_shape, block_size, device) + prefix_len - video_sizes = _construct_variable_block_sizes(video_shape, block_size, device) - partition = torch.cat([torch.arange(prefix_len, device=device, dtype=torch.long), video_indices]) - sizes = torch.cat([torch.tensor(prefix_sizes, device=device, dtype=torch.int32), video_sizes.to(torch.int32)]) - non_pad = _get_non_pad_index(sizes, block_elements) - untile = non_pad[torch.argsort(partition)] - total = prefix_len + math.prod(video_shape) - if int(sizes.sum()) != total or untile.numel() != total: - raise ValueError( - f"invalid H3 VSA geometry: prefix={prefix_segments}, video={video_shape}, " - f"sizes_sum={int(sizes.sum())}, total={total}" - ) - return partition, sizes, non_pad, untile, len(prefix_sizes), int(video_sizes.numel()) - - -def _get_h3_layout( - attn_metadata: AttentionMetadata | None, -) -> tuple[tuple[int, ...], tuple[int, int, int], int] | None: - if attn_metadata is None or attn_metadata.video_layout is None: - return None - prefix = attn_metadata.extra.get("vsa_h3_prefix_segments") - if not isinstance(prefix, (tuple, list)): - return None - target = next( - (span for span in reversed(attn_metadata.video_layout.video_spans) if span.role == "target"), - None, - ) - if target is None: - return None - return tuple(int(x) for x in prefix if int(x) > 0), target.latent_grid, target.start - - -def _pool_h3_tiles(x: torch.Tensor, sizes: torch.Tensor) -> torch.Tensor: - batch, seq_len, heads, dim = x.shape - blocks = seq_len // 64 - pooled = x.view(batch, blocks, 64, heads, dim).sum(dim=2, dtype=torch.float32) - pooled = pooled / sizes.view(1, -1, 1, 1).clamp_min(1) - return pooled.permute(0, 2, 1, 3) - - -def _build_h3_block_map( - scores: torch.Tensor, - num_prefix_blocks: int, - num_video_blocks: int, - topk: int, -) -> torch.Tensor: - """Prefix K/V are exempt and prefix queries stay dense, as in FastVideo.""" - keep_video = min(topk, num_video_blocks) - if keep_video == num_video_blocks: - return torch.ones_like(scores, dtype=torch.bool) - block_map = torch.zeros_like(scores, dtype=torch.bool) - indices = scores[..., num_prefix_blocks:].topk(keep_video, dim=-1).indices + num_prefix_blocks - block_map.scatter_(-1, indices, True) - block_map[..., :num_prefix_blocks] = True - block_map[:, :, :num_prefix_blocks, :] = True - return block_map - - def _get_vsa_dit_seq_shape(attn_metadata: AttentionMetadata | None) -> tuple[int, int, int] | None: if attn_metadata is None: return None @@ -504,93 +357,6 @@ def _fallback_reason( return "sequence parallel context is active" return None - def _forward_h3( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_metadata: AttentionMetadata, - ) -> torch.Tensor: - layout = _get_h3_layout(attn_metadata) - if layout is None: - raise ValueError("incomplete VSA-H3 layout metadata") - prefix_segments, video_shape, target_start = layout - if sum(prefix_segments) != target_start: - raise ValueError(f"VSA-H3 prefix segments sum to {sum(prefix_segments)}, target starts at {target_start}") - expected = target_start + math.prod(video_shape) - if query.shape[1] != expected: - raise ValueError(f"VSA-H3 layout has {expected} rows but attention received {query.shape[1]}") - if query.dtype not in (torch.float16, torch.bfloat16): - raise ValueError(f"VSA-H3 requires fp16/bf16 tensors, got {query.dtype}") - gate = _get_gate_compress(attn_metadata) - if gate is not None: - # H3 pads the packed document to 64 rows, while VSA operates on - # the valid prefix. Validate/slice before launching attention so a - # metadata error can never trigger an unsafe dense fallback after - # an asynchronous custom kernel. - if gate.shape[0] != query.shape[0] or gate.shape[2:] != query.shape[2:] or gate.shape[1] < query.shape[1]: - raise ValueError(f"gate_compress shape {gate.shape} cannot cover query shape {query.shape}") - gate = gate[:, : query.shape[1]] - - partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata( - prefix_segments, video_shape, query.device - ) - logical_blocks = int(sizes.numel()) - # The native sm100a kernel assigns pairs of query blocks to CTAs. Its - # contract requires an even block count; the synthetic partner is - # transport-only and is removed before returning. - pair_pad = logical_blocks % 2 - kernel_blocks = logical_blocks + pair_pad - target_shape = (query.shape[0], kernel_blocks * 64, query.shape[2], query.shape[3]) - q_tiled = torch.zeros(target_shape, device=query.device, dtype=query.dtype) - k_tiled = torch.zeros_like(q_tiled) - v_tiled = torch.zeros_like(q_tiled) - q_tiled[:, non_pad] = query[:, partition] - k_tiled[:, non_pad] = key[:, partition] - v_tiled[:, non_pad] = value[:, partition] - - q_pool = _pool_h3_tiles(q_tiled[:, : logical_blocks * 64], sizes) - k_pool = _pool_h3_tiles(k_tiled[:, : logical_blocks * 64], sizes) - scores = torch.matmul(q_pool, k_pool.transpose(-2, -1)) * self.softmax_scale - block_map = _build_h3_block_map(scores, prefix_blocks, video_blocks, self.topk) - kernel_sizes = sizes - if pair_pad: - block_map = torch.nn.functional.pad(block_map, (0, 1, 0, 1), value=False) - kernel_sizes = torch.nn.functional.pad(sizes, (0, 1), value=0) - - logger.info_once( - "FASTVIDEO_VSA H3 routing: seq_len=%d, prefix_segments=%s, video_shape=%s, " - "prefix_blocks=%d, video_blocks=%d, topk=%d, kernel_blocks=%d", - query.shape[1], - prefix_segments, - video_shape, - prefix_blocks, - video_blocks, - min(self.topk, video_blocks), - kernel_blocks, - ) - output = _fastvideo_h3_vsa_bhsd_op( - q_tiled.contiguous(), - k_tiled.contiguous(), - v_tiled.contiguous(), - block_map.contiguous(), - kernel_sizes.contiguous(), - logical_blocks, - )[:, : logical_blocks * 64] - - if gate is not None: - gate_tiled = torch.zeros_like(q_tiled[:, : logical_blocks * 64]) - gate_tiled[:, non_pad] = gate[:, partition] - v_pool = _pool_h3_tiles(v_tiled[:, : logical_blocks * 64], sizes) - compressed = torch.matmul(torch.softmax(scores, dim=-1), v_pool) - compressed = compressed.permute(0, 2, 1, 3).to(output.dtype) - output = ( - output.view(output.shape[0], logical_blocks, 64, output.shape[2], output.shape[3]) - + compressed.unsqueeze(2) - * gate_tiled.view(gate_tiled.shape[0], logical_blocks, 64, gate_tiled.shape[2], gate_tiled.shape[3]) - ).view_as(output) - return output[:, untile].contiguous() - def forward_cuda( self, query: torch.Tensor, @@ -611,26 +377,6 @@ def forward_cuda( key = key[:, :valid_seq_len] value = value[:, :valid_seq_len] - if attn_metadata is not None and _get_h3_layout(attn_metadata) is not None: - try: - output = self._forward_h3(query, key, value, attn_metadata) - if valid_seq_len == original_seq_len: - return output - restored = torch.zeros_like(original_query) - restored[:, :valid_seq_len] = output - return restored - except Exception as exc: - # A CUDA fault poisons the process context; attempting SDPA - # afterwards obscures the original kernel failure and cannot - # recover the request. - if isinstance(exc, torch.AcceleratorError): - raise - if not self.fallback_on_error: - raise - return self._fallback( - original_query, original_key, original_value, attn_metadata, f"VSA-H3 kernel failed: {exc}" - ) - reason = self._fallback_reason(query, key, value, attn_metadata) if reason is not None: return self._fallback(original_query, original_key, original_value, attn_metadata, reason) diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 4c588af0cfb..76a3a250076 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -7,7 +7,9 @@ # https://github.com/feifeibear/long-context-attention/blob/main/yunchang/attention/layer.py +from collections.abc import Mapping from dataclasses import replace +from typing import cast import torch import torch.nn as nn @@ -16,7 +18,7 @@ from vllm.model_executor.models.utils import extract_layer_index from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec -from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend, AttentionImpl, AttentionMetadata from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention @@ -83,6 +85,10 @@ def __init__( custom_attention: nn.Module | None = None, # Preserve dense FP32 inference for models opting into CUDA auto fallback. allow_fp32_fallback: bool = False, + # Model-owned implementation subclasses, keyed by backend name. + # Keep the platform-selected backend and its capabilities unchanged. + # Each override must preserve the selected implementation's contract. + impl_overrides: Mapping[str, type[AttentionImpl]] | None = None, ): super().__init__() @@ -181,8 +187,18 @@ def __init__( self.backend_pref = attn_backend_cls.get_name() logger.debug("Attention(role=%s) → platform default (%s)", role, self.backend_pref) - self.attn_backend = attn_backend_cls + self.attn_backend: type[AttentionBackend] | None = attn_backend_cls self.attn_impl_cls = self.attn_backend.get_impl_cls() + if impl_overrides is not None: + override = impl_overrides.get(attn_backend_cls.get_name()) + if override is not None: + if not issubclass(override, self.attn_impl_cls): + raise TypeError( + f"Attention implementation override {override.__qualname__} must subclass " + f"the selected implementation {self.attn_impl_cls.__qualname__} " + f"for backend {attn_backend_cls.__qualname__}" + ) + self.attn_impl_cls = override self.attention = self.attn_impl_cls( num_heads=num_heads, head_size=head_size, @@ -197,7 +213,7 @@ def __init__( ) # Compatibility kernels run inside shared dispatch, between the # parallel strategy's input preparation and output restoration. - self.sdpa_fallback = SDPABackend.get_impl_cls()( + self.sdpa_fallback: AttentionImpl | None = SDPABackend.get_impl_cls()( num_heads=num_heads, head_size=head_size, softmax_scale=softmax_scale, @@ -296,6 +312,29 @@ def _get_active_parallel_strategy(self): return self._no_parallel_strategy return self.parallel_strategy + @property + def supports_qk_input_landing(self) -> bool: + """Expose the static producer-direct capability to model layers.""" + if self.skip_sequence_parallel: + return False + return bool(getattr(self.parallel_strategy, "supports_qk_input_landing", False)) + + @torch.compiler.disable + def prepare_qk_input_landings( + self, + query: torch.Tensor, + key: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + """Acquire Q/K landing buffers only from the existing eager island.""" + strategy = self._get_active_parallel_strategy() + prepare = getattr(strategy, "prepare_qk_input_landings", None) + if prepare is None: + return None + return prepare( + query, + key, + ) + def _init_kv_cache_quantization(self, config) -> None: if config is None or self._has_custom_attention: return @@ -305,6 +344,7 @@ def _init_kv_cache_quantization(self, config) -> None: parallel_config = getattr(config, "parallel_config", None) ring_degree = getattr(parallel_config, "ring_degree", 1) if dtype: + assert self.attn_backend is not None if ring_degree > 1: raise ValueError( "KV quantization is not compatible with ring attention " @@ -400,6 +440,7 @@ def _forward_impl( ) use_paged_attention = paged_adapter is not None and self.paged_kv_cache_role is not None if use_paged_attention and not getattr(self.attn_backend, "supports_paged_kv", False): + assert self.attn_backend is not None raise NotImplementedError( f"Diffusion paged KV requires an Omni backend with paged support; " f"selected {self.attn_backend.get_name()}" @@ -496,7 +537,7 @@ def is_paged_kv_active(self) -> bool: def _run_local_attention(self, query, key, value, attn_metadata): if self._has_custom_attention: - return self.attention(query, key, value, attn_metadata) + return cast(nn.Module, self.attention)(query, key, value, attn_metadata) self._assert_metadata_compatible(attn_metadata) @@ -506,7 +547,7 @@ def _run_local_attention(self, query, key, value, attn_metadata): and query.dtype == torch.float32 and query.ndim == 4 and not self.backend_explicit - and self.attn_backend.get_name() == "FLASH_ATTN" + and cast(type[AttentionBackend], self.attn_backend).get_name() == "FLASH_ATTN" and ( attn_metadata is None or ( @@ -519,7 +560,7 @@ def _run_local_attention(self, query, key, value, attn_metadata): ) ): logger.warning_once("Using SDPA for this layer's FP32 input with automatic CUDA FlashAttention selection.") - return self.sdpa_fallback.forward(query, key, value, attn_metadata) + return cast(AttentionImpl, self.sdpa_fallback).forward(query, key, value, attn_metadata) in_kv_memory_profile = is_forward_context_available() and get_forward_context().in_diffusion_kv_memory_profile # The startup KV-capacity profile needs tensor shapes, not a paged @@ -531,14 +572,14 @@ def _run_local_attention(self, query, key, value, attn_metadata): self._scheduler_paged_kv and self.paged_kv_cache_role is not None and in_kv_memory_profile - and self.attn_backend.get_name() == "FLASH_ATTN" + and cast(type[AttentionBackend], self.attn_backend).get_name() == "FLASH_ATTN" and not current_omni_platform.supports_diffusion_dense_flash_attention() ): logger.warning_once( "The startup KV memory profile is using SDPA because dense FLASH_ATTN is unavailable. " "Formal paged requests still use the platform-native paged attention backend." ) - return self.sdpa_fallback.forward(query, key, value, attn_metadata) + return cast(AttentionImpl, self.sdpa_fallback).forward(query, key, value, attn_metadata) return self.attention.forward(query, key, value, attn_metadata) diff --git a/vllm_omni/diffusion/attention/ops/sage_block_sparse_attention.py b/vllm_omni/diffusion/attention/ops/sage_block_sparse_attention.py new file mode 100644 index 00000000000..07243f9191f --- /dev/null +++ b/vllm_omni/diffusion/attention/ops/sage_block_sparse_attention.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Sage fine attention with an explicit sparse-layout and prepared-Q contract. + +The scheduler owns stream dependencies and sparse selection. This adapter owns +operand quantization, the CAKE kernel invocation, and output layout. Tensor-map +descriptors are owned by FlashInfer's by-value launch ABI, not by a global cache. +""" + +import math +from functools import cache + +import torch + + +def validate_sparse_inputs( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + indices: torch.Tensor, + counts: torch.Tensor, + block_sizes: torch.Tensor | None, +) -> None: + """Check structural metadata without reading device values on the host.""" + for tensor in (query, key, value): + if tensor.ndim != 4 or tensor.shape[-1] != 128 or min(tensor.shape) < 1: + raise ValueError("Sage expects nonempty BSHD inputs with head dimension 128") + if tensor.stride(-1) != 1: + raise ValueError("Sage requires a contiguous head dimension") + if key.shape != value.shape or query.shape[::2] != key.shape[::2]: + raise ValueError("Sage requires matching batch/head dimensions and matching K/V shapes") + if any(t.dtype != torch.bfloat16 for t in (query, key, value)): + raise TypeError("Sage inputs must use BF16") + batch, rows, heads, _ = query.shape + blocks = (rows + 63) // 64 + if indices.ndim != 4 or indices.shape[:3] != (batch, heads, blocks): + raise ValueError("Sparse indices must have shape [B,H,ceil(Sq/64),capacity]") + if indices.shape[-1] < 1 or counts.shape != indices.shape[:3]: + raise ValueError("Sparse counts must match the query blocks and capacity must be positive") + metadata: tuple[torch.Tensor, ...] = (indices, counts) + if block_sizes is not None: + key_blocks = (key.shape[1] + 63) // 64 + if block_sizes.shape not in ((key_blocks,), (batch, key_blocks), (batch, heads, key_blocks)): + raise ValueError("Block sizes must have shape [Kblocks], [B,Kblocks], or [B,H,Kblocks]") + metadata += (block_sizes,) + for tensor in metadata: + if tensor.dtype != torch.int32 or not tensor.is_contiguous(): + raise TypeError("Sparse metadata must be contiguous INT32") + if any(t.device != query.device for t in (key, value, *metadata)): + raise ValueError("Sage inputs and sparse metadata must share one device") + + +@cache +def _empty_descriptor_workspace(device: torch.device) -> torch.Tensor: + # The public API still accepts this argument; the by-value ABI never uses + # its storage. There is no descriptor state shared by concurrent launches. + return torch.empty(0, dtype=torch.uint8, device=device) + + +def sage_block_sparse_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + indices: torch.Tensor, + counts: torch.Tensor, + block_sizes: torch.Tensor | None, + softmax_scale: float, + *, + prepared_q: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> torch.Tensor: + """Compute Sage fine attention and return contiguous BF16 BSHD output. + + A prepared Q may be supplied after the caller joins its producer stream. + Otherwise all three operands are quantized here. Ordered top-k indices + are not contiguous blocks, and prefix rows have different counts: neither + FlashInfer's uniform-count nor contiguous-index specialization is enabled. + """ + validate_sparse_inputs(query, key, value, indices, counts, block_sizes) + if isinstance(softmax_scale, bool) or not isinstance(softmax_scale, (int, float)): + raise TypeError("softmax_scale must be a real scalar") + if not math.isfinite(softmax_scale) or softmax_scale <= 0: + raise ValueError("softmax_scale must be finite and positive") + if prepared_q is not None: + batch, rows, heads, dim = query.shape + expected = (((batch, heads, rows, dim), torch.int8), ((batch, heads, ((rows + 127) // 128) * 4), torch.float32)) + if len(prepared_q) != 2: + raise ValueError("prepared_q must contain an INT8 tensor and FP32 scales") + for tensor, (shape, dtype) in zip(prepared_q, expected, strict=True): + if tensor.shape != shape or tensor.dtype != dtype or tensor.device != query.device: + raise ValueError("prepared_q must match the query shape, scale layout, and device") + if not tensor.is_contiguous(): + raise ValueError("prepared_q must be contiguous") + if query.device.type != "cuda": + raise ValueError("Sage execution requires CUDA") + from flashinfer.cute_dsl.sparse.bsa_attn_sm120 import bsa_attn_sm120_blk64_sage_fwd + + from .sage_quantization import quantize_sage_kv_sm120, quantize_sage_qkv_sm120 + + q_bhsd, k_bhsd, v_bhsd = (t.transpose(1, 2) for t in (query, key, value)) + if prepared_q is None: + quantized = quantize_sage_qkv_sm120(q_bhsd, k_bhsd, v_bhsd) + else: + k_int8, v_fp8, k_scale, v_scale = quantize_sage_kv_sm120(k_bhsd, v_bhsd) + quantized = (prepared_q[0], k_int8, v_fp8, prepared_q[1], k_scale, v_scale) + output = torch.empty_like(q_bhsd, memory_format=torch.contiguous_format) + result = bsa_attn_sm120_blk64_sage_fwd( + *quantized, + indices, + int(indices.shape[-1]), + block_sizes=block_sizes, + q2k_block_nums=counts, + softmax_scale=softmax_scale, + out=output, + tma_descriptor_workspace=_empty_descriptor_workspace(query.device), + uniform_block_count=False, + contiguous_block_indices=False, + backend="cake", + ) + return result.transpose(1, 2).contiguous() diff --git a/vllm_omni/diffusion/attention/ops/sage_quantization.py b/vllm_omni/diffusion/attention/ops/sage_quantization.py new file mode 100644 index 00000000000..22112176235 --- /dev/null +++ b/vllm_omni/diffusion/attention/ops/sage_quantization.py @@ -0,0 +1,564 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# Copyright (c) 2025 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Ported from Block-Sparse-Attention's csrc/fwd/sm120_blk64/bsa_quant_sm120_sage.py +# and bsa_sage_quant.py: Triton quantization kernels that produce the exact +# QK-INT8 / PV-FP8 operand layout consumed by +# flash_fwd_sm120_sage.BlockSparseAttnForwardSageSm120Blk64. +# +# Phase-1 scope note: unlike upstream, this module has no AOT-compiled quant +# runtime fallback and no external-workspace reuse parameter -- both are +# optional performance paths, not required for correctness. Only the plain +# Triton JIT path is ported here. + +from collections.abc import Sequence + +import torch +from vllm.triton_utils import tl, triton +from vllm.triton_utils import tldevice as libdevice + +from vllm_omni.platforms import current_omni_platform + +SAGE_Q_GROUP_SIZE = 32 +SAGE_Q_BLOCK_SIZE = 128 +SAGE_K_BLOCK_SIZE = 64 +SAGE_KV_STATS_CHUNK = 256 +SAGE_HEAD_DIM = 128 +SAGE_V_SCALE_MAX = 2.25 + + +# ============================================================================= +# Triton kernels +# ============================================================================= + + +@triton.jit +def _quantize_sage_q_kernel( + q_ptr, + q8_ptr, + scale_ptr, + q_stride_b, + q_stride_h, + q_stride_s, + q8_stride_b, + q8_stride_h, + q8_stride_s, + scale_stride_b, + scale_stride_h, + seqlen_q, + batch_size, + HEAD_DIM: tl.constexpr, # noqa: N803 + GROUP_SIZE: tl.constexpr, # noqa: N803 +): + group_idx = tl.program_id(0) + head_idx = tl.program_id(1) + batch_idx = tl.program_id(2) + + row_idx = group_idx * GROUP_SIZE + tl.arange(0, GROUP_SIZE) + dim_idx = tl.arange(0, HEAD_DIM) + valid = row_idx < seqlen_q + q = tl.load( + q_ptr + batch_idx * q_stride_b + head_idx * q_stride_h + row_idx[:, None] * q_stride_s + dim_idx[None, :], + mask=valid[:, None], + other=0.0, + ).to(tl.float32) + + row_max = tl.max(tl.abs(q), axis=1) + amax = tl.maximum(tl.max(row_max, axis=0), 1.0e-7) + scale = amax / 127.0 + q_quant = tl.maximum( + tl.minimum(libdevice.rint(q / scale), 127.0), + -127.0, + ) + tl.store( + q8_ptr + batch_idx * q8_stride_b + head_idx * q8_stride_h + row_idx[:, None] * q8_stride_s + dim_idx[None, :], + q_quant.to(tl.int8), + mask=valid[:, None], + ) + tl.store( + scale_ptr + batch_idx * scale_stride_b + head_idx * scale_stride_h + group_idx, + scale, + ) + + +@triton.jit +def _sage_kv_stats_partial_kernel( + k_ptr, + v_ptr, + k_partial_ptr, + v_partial_ptr, + k_stride_b, + k_stride_h, + k_stride_s, + v_stride_b, + v_stride_h, + v_stride_s, + partial_stride_b, + partial_stride_h, + partial_stride_c, + seqlen_k, + heads, + batch_size, + HEAD_DIM: tl.constexpr, # noqa: N803 + CHUNK_SIZE: tl.constexpr, # noqa: N803 + DIM_TILE: tl.constexpr, # noqa: N803 +): + chunk_idx = tl.program_id(0) + dim_tile_idx = tl.program_id(1) + batch_head_idx = tl.program_id(2) + batch_idx = batch_head_idx // heads + head_idx = batch_head_idx - batch_idx * heads + + row_idx = chunk_idx * CHUNK_SIZE + tl.arange(0, CHUNK_SIZE) + dim_idx = dim_tile_idx * DIM_TILE + tl.arange(0, DIM_TILE) + valid = row_idx < seqlen_k + k = tl.load( + k_ptr + batch_idx * k_stride_b + head_idx * k_stride_h + row_idx[:, None] * k_stride_s + dim_idx[None, :], + mask=valid[:, None], + other=0.0, + ).to(tl.float32) + v = tl.load( + v_ptr + batch_idx * v_stride_b + head_idx * v_stride_h + row_idx[:, None] * v_stride_s + dim_idx[None, :], + mask=valid[:, None], + other=0.0, + ).to(tl.float32) + + partial_base = batch_idx * partial_stride_b + head_idx * partial_stride_h + chunk_idx * partial_stride_c + dim_idx + tl.store(partial_base + k_partial_ptr, tl.sum(k, axis=0)) + tl.store(partial_base + v_partial_ptr, tl.max(tl.abs(v), axis=0)) + + +@triton.jit +def _sage_kv_stats_finalize_kernel( + k_partial_ptr, + v_partial_ptr, + k_mean_ptr, + v_scale_ptr, + partial_stride_b, + partial_stride_h, + partial_stride_c, + mean_stride_b, + mean_stride_h, + scale_stride_b, + scale_stride_h, + num_chunks, + seqlen_k, + heads, + batch_size, + DIM_TILE: tl.constexpr, # noqa: N803 + REDUCE_TILE: tl.constexpr, # noqa: N803 + SCALE_MAX: tl.constexpr, # noqa: N803 +): + dim_tile_idx = tl.program_id(0) + batch_head_idx = tl.program_id(1) + batch_idx = batch_head_idx // heads + head_idx = batch_head_idx - batch_idx * heads + dim_idx = dim_tile_idx * DIM_TILE + tl.arange(0, DIM_TILE) + + sum_acc = tl.zeros((DIM_TILE,), tl.float32) + max_acc = tl.zeros((DIM_TILE,), tl.float32) + chunk_base = 0 + while chunk_base < num_chunks: + chunk_idx = chunk_base + tl.arange(0, REDUCE_TILE) + mask = chunk_idx < num_chunks + partial_offset = ( + batch_idx * partial_stride_b + + head_idx * partial_stride_h + + chunk_idx[:, None] * partial_stride_c + + dim_idx[None, :] + ) + partial_sum = tl.load( + k_partial_ptr + partial_offset, + mask=mask[:, None], + other=0.0, + ) + partial_max = tl.load( + v_partial_ptr + partial_offset, + mask=mask[:, None], + other=0.0, + ) + sum_acc += tl.sum(partial_sum, axis=0) + max_acc = tl.maximum(max_acc, tl.max(partial_max, axis=0)) + chunk_base += REDUCE_TILE + + mean = sum_acc / seqlen_k + scale = tl.maximum(max_acc, 1.0e-7) / SCALE_MAX + tl.store( + k_mean_ptr + batch_idx * mean_stride_b + head_idx * mean_stride_h + dim_idx, + mean, + ) + tl.store( + v_scale_ptr + batch_idx * scale_stride_b + head_idx * scale_stride_h + dim_idx, + scale, + ) + + +@triton.jit +def _quantize_sage_kv_kernel( + k_ptr, + v_ptr, + k_mean_ptr, + v_scale_ptr, + k8_ptr, + v8_ptr, + k_scale_ptr, + k_stride_b, + k_stride_h, + k_stride_s, + v_stride_b, + v_stride_h, + v_stride_s, + mean_stride_b, + mean_stride_h, + v_scale_stride_b, + v_scale_stride_h, + k8_stride_b, + k8_stride_h, + k8_stride_s, + v8_stride_b, + v8_stride_h, + v8_stride_d, + k_scale_stride_b, + k_scale_stride_h, + seqlen_k, + batch_size, + HEAD_DIM: tl.constexpr, # noqa: N803 + BLOCK_SIZE: tl.constexpr, # noqa: N803 +): + block_idx = tl.program_id(0) + head_idx = tl.program_id(1) + batch_idx = tl.program_id(2) + local_row = tl.arange(0, BLOCK_SIZE) + row_idx = block_idx * BLOCK_SIZE + local_row + dim_idx = tl.arange(0, HEAD_DIM) + valid = row_idx < seqlen_k + + mean = tl.load(k_mean_ptr + batch_idx * mean_stride_b + head_idx * mean_stride_h + dim_idx).to(tl.float32) + v_scale = tl.load(v_scale_ptr + batch_idx * v_scale_stride_b + head_idx * v_scale_stride_h + dim_idx).to(tl.float32) + k = tl.load( + k_ptr + batch_idx * k_stride_b + head_idx * k_stride_h + row_idx[:, None] * k_stride_s + dim_idx[None, :], + mask=valid[:, None], + other=0.0, + ).to(tl.float32) + v = tl.load( + v_ptr + batch_idx * v_stride_b + head_idx * v_stride_h + row_idx[:, None] * v_stride_s + dim_idx[None, :], + mask=valid[:, None], + other=0.0, + ).to(tl.float32) + + k_centered = tl.where(valid[:, None], k - mean[None, :], 0.0) + row_max = tl.max(tl.abs(k_centered), axis=1) + k_amax = tl.maximum(tl.max(row_max, axis=0), 1.0e-7) + k_scale = k_amax / 127.0 + k_quant = tl.maximum( + tl.minimum(libdevice.rint(k_centered / k_scale), 127.0), + -127.0, + ) + tl.store( + k8_ptr + batch_idx * k8_stride_b + head_idx * k8_stride_h + row_idx[:, None] * k8_stride_s + dim_idx[None, :], + k_quant.to(tl.int8), + mask=valid[:, None], + ) + tl.store( + k_scale_ptr + batch_idx * k_scale_stride_b + head_idx * k_scale_stride_h + block_idx, + k_scale, + ) + + # This 16-token physical permutation must stay bit-for-bit consistent + # with the P-fragment byte-lane shuffle in + # flash_fwd_sm120_sage._make_acc_into_fp8_op -- the two together make the + # FP8 PV MMA's A-operand mapping lane-local without a runtime transpose. + row_mod = local_row % 16 + physical_row = ( + block_idx * BLOCK_SIZE + (local_row // 16) * 16 + (row_mod // 8) * 2 + ((row_mod // 2) % 4) * 4 + row_mod % 2 + ) + v_quant = tl.where(valid[:, None], v / v_scale[None, :], 0.0) + tl.store( + v8_ptr + + batch_idx * v8_stride_b + + head_idx * v8_stride_h + + dim_idx[:, None] * v8_stride_d + + physical_row[None, :], + tl.trans(v_quant).to(v8_ptr.dtype.element_ty), + ) + + +# ============================================================================= +# Public API +# ============================================================================= + + +def _require_sm120_bf16_bhsd(name: str, tensor: torch.Tensor) -> None: + if tensor.ndim != 4: + raise ValueError(f"{name} must be a rank-4 BHSD tensor") + if tensor.dtype != torch.bfloat16: + raise TypeError(f"{name} must use torch.bfloat16") + if not tensor.is_cuda: + raise ValueError(f"{name} must be a CUDA tensor") + if tensor.stride(-1) != 1: + raise ValueError(f"{name} must have a contiguous head-dim (stride(-1) == 1)") + if tensor.shape[-1] != SAGE_HEAD_DIM: + raise ValueError(f"{name} must have head dimension 128") + if tensor.shape[0] < 1 or tensor.shape[1] < 1 or tensor.shape[2] < 1: + raise ValueError(f"{name} requires positive B, H, and sequence length") + if current_omni_platform.get_device_capability( + tensor.device.index if tensor.device.index is not None else torch.accelerator.current_device_index() + ) != (12, 0): + raise RuntimeError("SM120 Sage quantization requires compute capability 12.0") + + +def _require_output( + name: str, + tensor: torch.Tensor, + shape: tuple, + dtype: torch.dtype, + device: torch.device, +) -> None: + if tensor.shape != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}") + if tensor.dtype != dtype: + raise TypeError(f"{name} must use {dtype}") + if tensor.device != device: + raise ValueError(f"{name} must be on {device}") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + + +def _normalize_out( + out: Sequence[torch.Tensor] | None, + specs: tuple, + device: torch.device, +) -> tuple: + if out is None: + return tuple(torch.empty(shape, dtype=dtype, device=device) for _, shape, dtype in specs) + if len(out) != len(specs): + raise ValueError(f"out must contain {len(specs)} tensors") + result = tuple(out) + for tensor, (name, shape, dtype) in zip(result, specs, strict=True): + _require_output(name, tensor, shape, dtype, device) + return result + + +def quantize_sage_q_sm120( + q: torch.Tensor, + *, + out: Sequence[torch.Tensor] | None = None, +): + """Quantize BF16 BHSD Q to Sage per-32-token-group INT8 on SM120. + + Args: + q: Query tensor, contiguous BHSD bfloat16, head_dim=128. + out: Optional pre-allocated ``(q_int8, q_scale)`` output tensors. + + Returns: + ``(q_int8, q_scale)``: ``q_int8`` is int8 with the same BHSD shape as + ``q``. ``q_scale`` is float32 with shape + ``(B, H, ceil(Sq/128) * 4)`` -- one scale per 32-token group, padded + to four groups per 128-query tile. + """ + _require_sm120_bf16_bhsd("Q", q) + batch, heads, seqlen_q, head_dim = q.shape + num_groups = triton.cdiv(seqlen_q, SAGE_Q_BLOCK_SIZE) * (SAGE_Q_BLOCK_SIZE // SAGE_Q_GROUP_SIZE) + q_int8, q_scale = _normalize_out( + out, + ( + ("q_int8", tuple(q.shape), torch.int8), + ("q_scale", (batch, heads, num_groups), torch.float32), + ), + q.device, + ) + _quantize_sage_q_kernel[(num_groups, heads, batch)]( + q, + q_int8, + q_scale, + q.stride(0), + q.stride(1), + q.stride(2), + q_int8.stride(0), + q_int8.stride(1), + q_int8.stride(2), + q_scale.stride(0), + q_scale.stride(1), + seqlen_q, + batch, + HEAD_DIM=head_dim, + GROUP_SIZE=SAGE_Q_GROUP_SIZE, + num_warps=8, + ) + return q_int8, q_scale + + +def quantize_sage_kv_sm120( + k: torch.Tensor, + v: torch.Tensor, + *, + out: Sequence[torch.Tensor] | None = None, +): + """Quantize BF16 BHSD K/V to Sage K64 INT8 and per-channel FP8 on SM120. + + K is channel-mean-centered before INT8 quantization (one scale per K64 + tile). V is FP8 E4M3 with a per-``[B,H,D]``-channel scale, stored in + Sage's ``[B, H, D, round_up(Sk, 64)]`` layout with the 16-token physical + permutation baked in by the kernel used to feed the FP8 PV MMA directly + (see flash_fwd_sm120_sage._make_acc_into_fp8_op). + + Args: + k: Key tensor, contiguous BHSD bfloat16, head_dim=128. + v: Value tensor, same shape/dtype/layout as ``k``. + out: Optional pre-allocated ``(k_int8, v_fp8, k_scale, v_scale)`` + output tensors. + + Returns: + ``(k_int8, v_fp8, k_scale, v_scale)``. + """ + _require_sm120_bf16_bhsd("K", k) + _require_sm120_bf16_bhsd("V", v) + if v.shape != k.shape: + raise ValueError("V must have the same shape as K") + if v.device != k.device: + raise ValueError("K and V must be on the same CUDA device") + + batch, heads, seqlen_k, head_dim = k.shape + padded_len = triton.cdiv(seqlen_k, SAGE_K_BLOCK_SIZE) * SAGE_K_BLOCK_SIZE + num_blocks = padded_len // SAGE_K_BLOCK_SIZE + specs = ( + ("k_int8", tuple(k.shape), torch.int8), + ("v_fp8", (batch, heads, head_dim, padded_len), torch.float8_e4m3fn), + ("k_scale", (batch, heads, num_blocks), torch.float32), + ("v_scale", (batch, heads, head_dim), torch.float32), + ) + k_int8, v_fp8, k_scale, v_scale = _normalize_out(out, specs, k.device) + + num_chunks = triton.cdiv(seqlen_k, SAGE_KV_STATS_CHUNK) + dim_tile = 32 + k_partial = torch.empty((batch, heads, num_chunks, head_dim), dtype=torch.float32, device=k.device) + v_partial = torch.empty_like(k_partial) + k_mean = torch.empty((batch, heads, head_dim), dtype=torch.bfloat16, device=k.device) + + _sage_kv_stats_partial_kernel[(num_chunks, head_dim // dim_tile, batch * heads)]( + k, + v, + k_partial, + v_partial, + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + k_partial.stride(0), + k_partial.stride(1), + k_partial.stride(2), + seqlen_k, + heads, + batch, + HEAD_DIM=head_dim, + CHUNK_SIZE=SAGE_KV_STATS_CHUNK, + DIM_TILE=dim_tile, + num_warps=8, + ) + _sage_kv_stats_finalize_kernel[(head_dim // dim_tile, batch * heads)]( + k_partial, + v_partial, + k_mean, + v_scale, + k_partial.stride(0), + k_partial.stride(1), + k_partial.stride(2), + k_mean.stride(0), + k_mean.stride(1), + v_scale.stride(0), + v_scale.stride(1), + num_chunks, + seqlen_k, + heads, + batch, + DIM_TILE=dim_tile, + REDUCE_TILE=16, + SCALE_MAX=SAGE_V_SCALE_MAX, + num_warps=4, + ) + _quantize_sage_kv_kernel[(num_blocks, heads, batch)]( + k, + v, + k_mean, + v_scale, + k_int8, + v_fp8, + k_scale, + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + k_mean.stride(0), + k_mean.stride(1), + v_scale.stride(0), + v_scale.stride(1), + k_int8.stride(0), + k_int8.stride(1), + k_int8.stride(2), + v_fp8.stride(0), + v_fp8.stride(1), + v_fp8.stride(2), + k_scale.stride(0), + k_scale.stride(1), + seqlen_k, + batch, + HEAD_DIM=head_dim, + BLOCK_SIZE=SAGE_K_BLOCK_SIZE, + num_warps=8, + ) + return k_int8, v_fp8, k_scale, v_scale + + +def quantize_sage_qkv_sm120( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + out: Sequence[torch.Tensor] | None = None, +): + """Quantize BF16 BHSD Q/K/V using the native Sage SM120 contract. + + Convenience wrapper around :func:`quantize_sage_q_sm120` and + :func:`quantize_sage_kv_sm120`. + + Returns: + ``(q_int8, k_int8, v_fp8, q_scale, k_scale, v_scale)``. + """ + if q.shape[:2] != k.shape[:2] or k.shape != v.shape: + raise ValueError("Q, K, and V must use matching batch/head dimensions") + if out is None: + q_out = None + kv_out = None + else: + if len(out) != 6: + raise ValueError("out must contain six tensors") + q_out = (out[0], out[3]) + kv_out = (out[1], out[2], out[4], out[5]) + + q_int8, q_scale = quantize_sage_q_sm120(q, out=q_out) + k_int8, v_fp8, k_scale, v_scale = quantize_sage_kv_sm120(k, v, out=kv_out) + return q_int8, k_int8, v_fp8, q_scale, k_scale, v_scale + + +__all__ = [ + "quantize_sage_kv_sm120", + "quantize_sage_q_sm120", + "quantize_sage_qkv_sm120", +] diff --git a/vllm_omni/diffusion/attention/parallel/ulysses.py b/vllm_omni/diffusion/attention/parallel/ulysses.py index 28298c0a280..1a9e9f83e3d 100644 --- a/vllm_omni/diffusion/attention/parallel/ulysses.py +++ b/vllm_omni/diffusion/attention/parallel/ulysses.py @@ -254,15 +254,31 @@ def __init__( self._gather_idx = gather_idx self._use_sync = use_sync self._ulysses_a2a_permute = ulysses_a2a_permute + self._ulysses_a2a_backend = "symmetric" + self._qk_input_landing_enabled = False if _a2a_permute_enabled( ulysses_a2a_permute, scatter_idx, gather_idx, sp_group.ulysses_world_size, ): - from vllm_omni.diffusion.distributed.a2a_permute import ensure_a2a_permute_available + from vllm_omni.diffusion.distributed.flashinfer_ulysses import configured_backend - ensure_a2a_permute_available() + self._ulysses_a2a_backend = configured_backend() + if self._ulysses_a2a_backend == "flashinfer-pcie": + from vllm_omni.diffusion.distributed.flashinfer_ulysses import ( + ensure_flashinfer_pcie_available, + is_qk_producer_direct_enabled, + ) + + ensure_flashinfer_pcie_available() + self._qk_input_landing_enabled = is_qk_producer_direct_enabled() + elif self._ulysses_a2a_backend == "symmetric": + from vllm_omni.diffusion.distributed.a2a_permute import ensure_a2a_permute_available + + ensure_a2a_permute_available() + else: + raise ValueError("VLLM_OMNI_ULYSSES_A2A_BACKEND must be symmetric or flashinfer-pcie") @property def enabled(self) -> bool: @@ -272,6 +288,45 @@ def enabled(self) -> bool: def name(self) -> str: return "ulysses" + @property + def supports_qk_input_landing(self) -> bool: + return self._qk_input_landing_enabled + + def prepare_qk_input_landings(self, query: torch.Tensor, key: torch.Tensor): + if not self.supports_qk_input_landing or get_ulysses_mode(default="strict") != "strict": + return None + if self._sp_group.ring_world_size != 1: + return None + from vllm_omni.diffusion.distributed.flashinfer_ulysses import flashinfer_ulysses_qk_input_landings + + return flashinfer_ulysses_qk_input_landings(query, key, self._ulysses_pg, self._sp_group.ulysses_world_size) + + def _scatter_heads(self, tensor: torch.Tensor, slot: str) -> torch.Tensor: + group_name = self._ulysses_pg.group_name + world_size = self._sp_group.ulysses_world_size + if self._ulysses_a2a_backend == "flashinfer-pcie": + from vllm_omni.diffusion.distributed.flashinfer_ulysses import flashinfer_ulysses_qkv_fwd + + return flashinfer_ulysses_qkv_fwd(tensor, group_name, world_size, slot, self._use_sync) + from vllm_omni.diffusion.distributed.a2a_permute import ulysses_qkv_fwd + + return ulysses_qkv_fwd(tensor, group_name, world_size) + + def _exchange_qkv(self, query, key, value, attn_metadata): + """Exchange independent operands; subclasses may schedule producers between exchanges.""" + return self._scatter_heads(query, "q"), self._scatter_heads(key, "k"), self._scatter_heads(value, "v") + + def _gather_heads(self, tensor: torch.Tensor, ctx: _UlyssesCtx) -> torch.Tensor: + group_name = ctx.ulysses_pg.group_name + world_size = dist.get_world_size(ctx.ulysses_pg) + if self._ulysses_a2a_backend == "flashinfer-pcie": + from vllm_omni.diffusion.distributed.flashinfer_ulysses import flashinfer_ulysses_o_rev + + return flashinfer_ulysses_o_rev(tensor, group_name, world_size, ctx.use_sync) + from vllm_omni.diffusion.distributed.a2a_permute import ulysses_o_rev + + return ulysses_o_rev(tensor, group_name, world_size) + def pre_attention( self, query: torch.Tensor, @@ -526,14 +581,9 @@ def pre_attention( self._gather_idx, ulysses_world_size, ): - from vllm_omni.diffusion.distributed.a2a_permute import ulysses_qkv_fwd - - group_name = self._ulysses_pg.group_name - query = ulysses_qkv_fwd(query, group_name, ulysses_world_size) - key = ulysses_qkv_fwd(key, group_name, ulysses_world_size) - value = ulysses_qkv_fwd(value, group_name, ulysses_world_size) + query, key, value = self._exchange_qkv(query, key, value, attn_metadata) if gate_compress is not None: - gate_compress = ulysses_qkv_fwd(gate_compress, group_name, ulysses_world_size) + gate_compress = self._scatter_heads(gate_compress, "g") else: query = SeqAllToAll4D.apply( self._ulysses_pg, query, self._scatter_idx, self._gather_idx, self._use_sync @@ -666,13 +716,7 @@ def post_attention(self, attn_output: torch.Tensor, ctx: ParallelAttentionContex ctx.gather_idx, dist.get_world_size(ctx.ulysses_pg), ): - from vllm_omni.diffusion.distributed.a2a_permute import ulysses_o_rev - - output_img = ulysses_o_rev( - output_img, - ctx.ulysses_pg.group_name, - dist.get_world_size(ctx.ulysses_pg), - ) + output_img = self._gather_heads(output_img, ctx) else: output_img = SeqAllToAll4D.apply( ctx.ulysses_pg, output_img, ctx.gather_idx, ctx.scatter_idx, ctx.use_sync @@ -711,11 +755,5 @@ def post_attention(self, attn_output: torch.Tensor, ctx: ParallelAttentionContex ctx.gather_idx, dist.get_world_size(ctx.ulysses_pg), ): - from vllm_omni.diffusion.distributed.a2a_permute import ulysses_o_rev - - return ulysses_o_rev( - attn_output, - ctx.ulysses_pg.group_name, - dist.get_world_size(ctx.ulysses_pg), - ) + return self._gather_heads(attn_output, ctx) return SeqAllToAll4D.apply(ctx.ulysses_pg, attn_output, ctx.gather_idx, ctx.scatter_idx, ctx.use_sync) diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 7f0242d57ab..2cc321050ef 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -199,7 +199,7 @@ def extract_qwen_context( hidden_states, vid_freqs, txt_freqs = module.image_rope_prepare(hidden_states, img_shapes, txt_seq_lens) image_rotary_emb = (vid_freqs, txt_freqs) - timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype) + timestep = torch.as_tensor(timestep, device=hidden_states.device, dtype=hidden_states.dtype) # Call modulate_index_prepare instead of handling timestep directly. # For zero_cond_t=False: timestep unchanged, modulate_index=None. @@ -618,6 +618,7 @@ def extract_flux2_klein_context( # ============================================================================ dtype = hidden_states.dtype + assert encoder_hidden_states is not None num_txt_tokens = encoder_hidden_states.shape[1] timestep = timestep.to(dtype=dtype) * 1000 @@ -969,6 +970,7 @@ def extract_flux2_context( # ============================================================================ # PREPROCESSING (Flux2-specific) # ============================================================================ + assert encoder_hidden_states is not None num_txt_tokens = encoder_hidden_states.shape[1] timestep = timestep.to(hidden_states.dtype) * 1000 @@ -1260,7 +1262,7 @@ def extract_minimax_h3_context( preprocessing via ``_embed``, cacheable ``blocks`` loop, and ``final_layer`` postprocessing with row selection and update masks. """ - from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import indexed_scale_shift_ + from vllm_omni.diffusion.layers.indexed_modulation import indexed_scale_shift_ from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( _FORWARD_SUPPORTED_KWARGS, MINIMAX_H3_ADALN_MODALITY_NUM, diff --git a/vllm_omni/diffusion/distributed/flashinfer_ulysses.py b/vllm_omni/diffusion/distributed/flashinfer_ulysses.py new file mode 100644 index 00000000000..9f6b3f7e699 --- /dev/null +++ b/vllm_omni/diffusion/distributed/flashinfer_ulysses.py @@ -0,0 +1,1056 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Opt-in FlashInfer PCIe/RDMA transport for strict Ulysses attention. + +This adapter intentionally sits behind the existing ``--ulysses-a2a-permute`` +switch plus ``VLLM_OMNI_ULYSSES_A2A_BACKEND=flashinfer-pcie``. The default +continues to be vLLM-Omni's symmetric-memory kernel. + +Q/K producer-direct landing is a second, independent opt-in through +``VLLM_OMNI_FLASHINFER_ULYSSES_QK_PRODUCER_DIRECT=1``. This keeps the existing +FlashInfer adapter path unchanged for isolated A/B measurements. + +Attention kernels may independently write its fused tile-to-aligned output into the +registered reverse-Ulysses source buffer through +``VLLM_OMNI_FLASHINFER_ULYSSES_O_PRODUCER_DIRECT=1``. The communicator still +performs the identical gather-heads exchange; only its source staging copy is +removed. + +FlashInfer's PCIe communicator owns registered output and RDMA landing buffers. +One communicator is shared by all attention layers in a process-group. Separate +registered slots keep Q/K/V and the inverse-attention result live at the same time. Non-contiguous projection views +are copied directly into the registered landing buffer, avoiding an intermediate +``contiguous()`` allocation and the transport's otherwise-required second device +copy. +""" + +from __future__ import annotations + +import os +import threading +from dataclasses import dataclass, field +from typing import Any + +import torch +import torch.distributed as dist +from torch.distributed.distributed_c10d import _resolve_process_group +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_BACKEND_ENV = "VLLM_OMNI_ULYSSES_A2A_BACKEND" +_MAX_BYTES_ENV = "VLLM_OMNI_FLASHINFER_ULYSSES_MAX_BYTES" +_REQUIRE_RDMA_ENV = "VLLM_OMNI_FLASHINFER_ULYSSES_REQUIRE_RDMA" +_QK_PRODUCER_DIRECT_ENV = "VLLM_OMNI_FLASHINFER_ULYSSES_QK_PRODUCER_DIRECT" +_O_PRODUCER_DIRECT_ENV = "VLLM_OMNI_FLASHINFER_ULYSSES_O_PRODUCER_DIRECT" +_RELEASE_AFTER_DENOISE_ENV = "VLLM_OMNI_FLASHINFER_ULYSSES_RELEASE_AFTER_DENOISE" +_EXPECTED_BACKEND = "flashinfer-pcie" + +_LOCK = threading.RLock() + +_TensorIdentity = tuple[ + int, + tuple[int, ...], + tuple[int, ...], + int, + torch.dtype, + torch.device, +] + + +def _tensor_identity(x: torch.Tensor) -> _TensorIdentity: + """Describe the exact tensor view handed to the communicator.""" + return ( + x.data_ptr(), + tuple(x.shape), + tuple(x.stride()), + x.storage_offset(), + x.dtype, + x.device, + ) + + +def configured_backend() -> str: + """Return the explicitly selected optimized Ulysses implementation.""" + return os.getenv(_BACKEND_ENV, "symmetric").strip().lower() + + +def is_flashinfer_pcie_enabled() -> bool: + return configured_backend() == _EXPECTED_BACKEND + + +def is_qk_producer_direct_enabled() -> bool: + """Return whether a producer may write Q/K into registered source storage.""" + value = os.getenv(_QK_PRODUCER_DIRECT_ENV, "0").strip().lower() + if value not in ("0", "1", "false", "true"): + raise ValueError(f"{_QK_PRODUCER_DIRECT_ENV} must be 0/1/false/true, got {value!r}") + return is_flashinfer_pcie_enabled() and value in ("1", "true") + + +def is_o_producer_direct_enabled() -> bool: + """Return whether Attention kernels may target the registered reverse-O source.""" + value = os.getenv(_O_PRODUCER_DIRECT_ENV, "0").strip().lower() + if value not in ("0", "1", "false", "true"): + raise ValueError(f"{_O_PRODUCER_DIRECT_ENV} must be 0/1/false/true, got {value!r}") + return is_flashinfer_pcie_enabled() and value in ("1", "true") + + +def release_after_denoise_enabled() -> bool: + """Whether registered PCIe/RDMA windows are scoped to one denoise. + + Re-registering the windows adds cold work to the next request, so the + memory-saving lifecycle remains explicitly disabled by default. + """ + value = os.getenv(_RELEASE_AFTER_DENOISE_ENV, "0").strip().lower() + if value not in ("0", "1", "false", "true"): + raise ValueError(f"{_RELEASE_AFTER_DENOISE_ENV} must be 0/1/false/true, got {value!r}") + return is_flashinfer_pcie_enabled() and value in ("1", "true") + + +def _require_rdma() -> bool: + value = os.getenv(_REQUIRE_RDMA_ENV, "1").strip().lower() + if value not in ("0", "1", "false", "true"): + raise ValueError(f"{_REQUIRE_RDMA_ENV} must be 0/1/false/true, got {value!r}") + return value in ("1", "true") + + +def ensure_flashinfer_pcie_available() -> None: + """Validate the PR #4876 Python/API surface without compiling a kernel.""" + try: + from flashinfer.comm import UlyssesCommunicator + from flashinfer.comm.ulysses import missing_ulysses_pcie_dependencies + except (ImportError, AttributeError) as exc: + raise RuntimeError( + "flashinfer-pcie requires FlashInfer PR #4876 (UlyssesCommunicator " + "with backend='pcie'); prepend its pinned source checkout to PYTHONPATH" + ) from exc + + missing = missing_ulysses_pcie_dependencies() + if missing: + raise RuntimeError("flashinfer-pcie JIT dependencies are missing: " + ", ".join(missing)) + methods = ["allocate_output", "input_buffer", "scatter_heads", "gather_heads", "close"] + for method in methods: + if not hasattr(UlyssesCommunicator, method): + raise RuntimeError(f"FlashInfer UlyssesCommunicator lacks required method {method!r}") + + +def _configured_capacity(first_nbytes: int) -> int: + raw = os.getenv(_MAX_BYTES_ENV) + if raw is None or not raw.strip(): + return first_nbytes + try: + capacity = int(raw) + except ValueError as exc: + raise ValueError(f"{_MAX_BYTES_ENV} must be a positive integer, got {raw!r}") from exc + if capacity <= 0: + raise ValueError(f"{_MAX_BYTES_ENV} must be positive, got {capacity}") + return capacity + + +def _group_name(group: dist.ProcessGroup) -> str: + name = getattr(group, "group_name", None) + return str(name) if name is not None else f"process-group-{id(group)}" + + +@dataclass +class _PcieState: + group: dist.ProcessGroup + world_size: int + device: torch.device + capacity_bytes: int + communicator: Any | None = None + disabled_reason: str | None = None + # ``release_after_denoise`` returns the registered raw-CUDA allocations + # before VAE decode, but the VAE then leaves its much larger scratch arena + # in PyTorch's caching allocator. FlashInfer's next registered windows use + # native cudaMalloc and cannot reuse those cached segments. Remember that + # lifecycle transition so only the next cold re-arm returns inactive VAE + # cache to CUDA; the first request and every hot layer call stay untouched. + rearm_pending: bool = False + producer_direct_logged: bool = False + producer_direct_generation: int = 0 + producer_direct_expected: dict[str, _TensorIdentity] = field( + default_factory=dict, + repr=False, + ) + producer_direct_exchanged: set[str] = field(default_factory=set, repr=False) + producer_direct_lock: threading.Lock = field( + default_factory=threading.Lock, + repr=False, + compare=False, + ) + o_producer_direct_logged: bool = False + o_producer_direct_expected: _TensorIdentity | None = field( + default=None, + repr=False, + ) + o_producer_direct_lock: threading.Lock = field( + default_factory=threading.Lock, + repr=False, + compare=False, + ) + # One registered allocation per logical Q/K/V/O slot. FlashInfer sizes + # each allocation (and its RDMA landing buffer) at ``capacity_bytes`` even + # though allocate_output returns a view for the first call's geometry. + # Reuse that storage for smaller/later geometries instead of registering + # another capacity-sized pair for every token-refiner/main-DiT shape. + outputs: dict[tuple[str, str, torch.dtype], torch.Tensor] = field(default_factory=dict) + input_landing_identities: dict[tuple[str, str, torch.dtype], _TensorIdentity] = field( + default_factory=dict, repr=False + ) + + def initialize(self, dtype: torch.dtype) -> bool: + """Collectively construct once; a clean cold-start failure falls back.""" + if self.disabled_reason is not None: + return False + if self.communicator is not None: + return True + + if self.rearm_pending: + # Every rank reaches initialize through the same strict-Ulysses + # operation. Drain its local work and release only *inactive* + # allocator segments before FlashInfer asks CUDA for raw registered + # storage. Live model, compiled-region and current-layer tensors + # are unaffected by empty_cache(). Vote before entering the + # communicator constructor so a rank-local allocator failure cannot + # strand peers in its collective setup protocol. + reclaim_error: BaseException | None = None + try: + torch.accelerator.synchronize(self.device) + torch.accelerator.empty_cache() + except BaseException as exc: + reclaim_error = exc + reclaim_failures = self._collective_sum(int(reclaim_error is not None)) + if reclaim_failures: + raise RuntimeError( + "FlashInfer PCIe Ulysses re-arm could not reclaim the " + f"post-VAE allocator cache on {reclaim_failures}/{self.world_size} ranks" + ) from reclaim_error + logger.info( + "FlashInfer PCIe Ulysses re-arm reclaimed post-VAE allocator cache: rank=%d group=%s device=%s", + dist.get_rank(self.group), + _group_name(self.group), + self.device, + ) + + from flashinfer.comm import UlyssesCommunicator + + try: + communicator = UlyssesCommunicator( + self.group, + max_bytes=self.capacity_bytes, + dtype=dtype, + backend="pcie", + device=self.device, + ) + except Exception as exc: # FlashInfer coordinates init cleanup across ranks. + if _require_rdma(): + raise RuntimeError("required FlashInfer all-RDMA Ulysses initialization failed") from exc + self.disabled_reason = f"cold initialization failed ({type(exc).__name__}: {exc})" + logger.warning("FlashInfer PCIe Ulysses unavailable; using NCCL: %s", self.disabled_reason) + return False + + if _require_rdma() and communicator.transport != "rdma": + observed = communicator.transport + communicator.close() + raise RuntimeError( + f"required FlashInfer all-RDMA Ulysses route was not selected; observed transport={observed!r}" + ) + + self.communicator = communicator + self.rearm_pending = False + logger.info( + "FlashInfer PCIe Ulysses armed: group=%s world=%d device=%s transport=%s capacity=%d", + _group_name(self.group), + self.world_size, + self.device, + communicator.transport, + self.capacity_bytes, + ) + return True + + def output(self, x: torch.Tensor, *, op: str, slot: str) -> torch.Tensor: + assert self.communicator is not None + key = (slot, op, x.dtype) + out = self.outputs.get(key) + if out is None: + # allocate_output validates contiguity but never reads the values. + prototype = x if x.is_contiguous() else torch.empty(x.shape, dtype=x.dtype, device=x.device) + # PR #4876 permits a per-call element type on the PCIe backend. + # Always state it explicitly: the communicator default is fixed by + # the first collective, while callers may + # use E4M3 Q/K/V and retain BF16 for the reverse-O exchange. + out = self.communicator.allocate_output( + prototype, + op=op, + dtype=x.dtype, + ) + self.outputs[key] = out + return out + + batch, seq, heads, head_dim = x.shape + if op == "scatter_heads": + shape = (batch, seq * self.world_size, heads // self.world_size, head_dim) + elif op == "gather_heads": + shape = (batch, seq // self.world_size, heads * self.world_size, head_dim) + else: # This adapter deliberately exposes only the two Ulysses transforms. + raise ValueError(f"unsupported cached FlashInfer Ulysses op: {op!r}") + + # ``out`` may be a narrow view of a capacity-sized native allocation. + # as_strided checks against the backing storage, so it can safely form + # a different contiguous geometry without allocating or changing the + # registered base pointer. + strides = [1] * len(shape) + for index in range(len(shape) - 2, -1, -1): + strides[index] = strides[index + 1] * shape[index + 1] + return out.as_strided(shape, tuple(strides)) + + def registered_shape_prototype( + self, + shape: tuple[int, int, int, int], + *, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor: + """Return a shape-only view over an existing capacity-sized slot. + + FlashInfer's PCIe ``allocate_output`` validates the input geometry and + element type but never reads its values. After Q/K/V have armed the + communicator, one of their registered native outputs therefore + provides a safe untyped backing store for a reverse-O *prototype* + view. The view may intentionally use a different dtype (FP8 QKV -> + BF16 O); this avoids allocating a transient full-size CUDA tensor + solely to register the Attention kernels O or O-bundle slot. + """ + if len(shape) != 4 or any(int(dim) <= 0 for dim in shape): + raise ValueError(f"registered shape prototype requires positive 4-D shape, got {shape}") + required_bytes = int(torch.Size(shape).numel()) * dtype.itemsize + if required_bytes > self.capacity_bytes: + raise RuntimeError( + "registered shape prototype exceeds communicator capacity: " + f"operand={required_bytes}, capacity={self.capacity_bytes}" + ) + strides = [1] * len(shape) + for index in range(len(shape) - 2, -1, -1): + strides[index] = strides[index + 1] * shape[index + 1] + # Prefer an exact typed view for the ordinary BF16-only path. If the + # forward collective used FP8, reinterpret only the metadata over the + # same untyped registered storage. ``allocate_output`` never reads + # this prototype, and the newly allocated O slot remains separately + # typed and registered by FlashInfer. + candidates = sorted( + self.outputs.values(), + key=lambda candidate: candidate.dtype != dtype, + ) + for candidate in candidates: + if candidate.device != device: + continue + storage = candidate.untyped_storage() + if candidate.storage_offset() != 0 or candidate.data_ptr() != storage.data_ptr(): + continue + if storage.nbytes() < required_bytes: + continue + if candidate.dtype == dtype: + return candidate.as_strided(shape, tuple(strides)) + prototype = torch.empty(0, dtype=dtype, device=device) + return prototype.set_(storage, 0, shape, tuple(strides)) + raise RuntimeError( + "no existing registered Ulysses output has capacity for the " + f"shape-only prototype {shape}; Q/K/V must be armed first" + ) + + def scatter(self, x: torch.Tensor, *, slot: str) -> torch.Tensor: + assert self.communicator is not None + out = self.output(x, op="scatter_heads", slot=slot) + if x.is_contiguous(): + source = x + else: + # Directly materialize a strided Q/K/V projection view in the + # registered source buffer. Passing this exact pointer suppresses + # FlashInfer's ordinary source -> landing staging copy. + source = self.communicator.input_buffer(out, x.shape) + source.copy_(x) + result = self.communicator.scatter_heads( + source, + out=out, + dtype=source.dtype, + ) + if slot == "g": + logger.info_once( + "FlashInfer PCIe Ulysses VSA gate exchange active: rank=%d group=%s shape=%s dtype=%s", + dist.get_rank(self.group), + _group_name(self.group), + tuple(source.shape), + source.dtype, + # The default vLLM "local" scope logs only local rank zero. + # Qualification needs evidence that every SP rank completed + # this separate registered-RDMA collective. + scope="process", + ) + # This branch becomes a single false boolean check after the first + # confirmed Q/K pair. Keep all bookkeeping out of the steady-state + # transport path while requiring the collective call to return first. + landing_key = (slot, "scatter_heads", source.dtype) + landed_directly = self.input_landing_identities.get(landing_key) == _tensor_identity(source) + if slot in ("q", "k", "v") and landed_directly and not self.producer_direct_logged: + self.mark_producer_direct_exchange(source, slot=slot) + return result + + def input_landing( + self, + x: torch.Tensor, + *, + slot: str, + op: str = "scatter_heads", + ) -> torch.Tensor: + """Return the registered source storage paired with ``slot`` output. + + An eager producer may write this view directly. A later + ``scatter(..., slot=slot)`` resolves the same output allocation, and + FlashInfer recognizes the exact landing pointer and skips staging. + """ + assert self.communicator is not None + out = self.output(x, op=op, slot=slot) + landing = self.communicator.input_buffer(out, x.shape) + if landing.shape != x.shape: + raise RuntimeError( + "FlashInfer PCIe input landing shape mismatch: " + f"expected={tuple(x.shape)}, observed={tuple(landing.shape)}" + ) + if landing.dtype != x.dtype or landing.device != x.device: + raise RuntimeError( + "FlashInfer PCIe input landing dtype/device mismatch: " + f"expected=({x.dtype}, {x.device}), " + f"observed=({landing.dtype}, {landing.device})" + ) + if not landing.is_contiguous(): + raise RuntimeError("FlashInfer PCIe input landing must be contiguous") + if landing.untyped_storage().data_ptr() == out.untyped_storage().data_ptr(): + raise RuntimeError("FlashInfer PCIe input and output buffers unexpectedly alias") + self.input_landing_identities[(slot, op, x.dtype)] = _tensor_identity(landing) + return landing + + def arm_o_producer_direct(self, landing: torch.Tensor) -> None: + """Bind the next reverse-O exchange to one exact registered view.""" + landing_key = ("o", "gather_heads", landing.dtype) + identity = _tensor_identity(landing) + if self.input_landing_identities.get(landing_key) != identity: + raise RuntimeError("FlashInfer reverse-O producer landing is not registered for the gather-heads slot") + with self.o_producer_direct_lock: + self.o_producer_direct_expected = identity + + def arm_qk_producer_direct( + self, + q_landing: torch.Tensor, + k_landing: torch.Tensor, + ) -> None: + """Atomically bind one Q/K producer generation to exact landing views.""" + if self.producer_direct_logged: + return + with self.producer_direct_lock: + if self.producer_direct_logged: + return + self.producer_direct_generation += 1 + self.producer_direct_expected = { + "q": _tensor_identity(q_landing), + "k": _tensor_identity(k_landing), + } + self.producer_direct_exchanged.clear() + + def mark_producer_direct_exchange( + self, + source: torch.Tensor, + *, + slot: str, + ) -> None: + """Record a direct scatter only after its communicator call succeeds.""" + if self.producer_direct_logged or slot not in self.producer_direct_expected: + return + with self.producer_direct_lock: + if self.producer_direct_logged: + return + if self.producer_direct_expected.get(slot) != _tensor_identity(source): + return + self.producer_direct_exchanged.add(slot) + expected_slots = set(self.producer_direct_expected) + if self.producer_direct_exchanged != expected_slots: + return + + logger.info( + "FlashInfer PCIe Ulysses %s producer-direct exchange active: rank=%d group=%s shape=%s dtype=%s", + "/".join(name.upper() for name in ("q", "k", "v") if name in expected_slots), + dist.get_rank(self.group), + _group_name(self.group), + tuple(source.shape), + source.dtype, + ) + self.producer_direct_logged = True + + def gather(self, x: torch.Tensor, *, slot: str) -> torch.Tensor: + assert self.communicator is not None + out = self.output(x, op="gather_heads", slot=slot) + input_identity = _tensor_identity(x) + with self.o_producer_direct_lock: + expected = self.o_producer_direct_expected if slot == "o" else None + producer_direct = expected is not None and expected == input_identity + if x.is_contiguous(): + source = x + else: + source = self.communicator.input_buffer(out, x.shape) + source.copy_(x) + result = self.communicator.gather_heads( + source, + out=out, + dtype=source.dtype, + ) + if expected is not None: + with self.o_producer_direct_lock: + # Clear only the generation consumed by this exchange. A + # failed communicator call leaves it armed for diagnostics or + # an explicit retry. + if self.o_producer_direct_expected == expected: + self.o_producer_direct_expected = None + if producer_direct and not self.o_producer_direct_logged: + logger.info( + "FlashInfer PCIe Ulysses O producer-direct exchange active: rank=%d group=%s shape=%s dtype=%s", + dist.get_rank(self.group), + _group_name(self.group), + tuple(x.shape), + x.dtype, + ) + self.o_producer_direct_logged = True + if not producer_direct: + logger.warning_once( + "FlashInfer PCIe Ulysses O producer-direct exchange fell back: " + "attention output did not match the armed registered landing" + ) + return result + + def close(self) -> None: + communicator = self.communicator + if communicator is not None: + # FlashInfer keeps a failed collective close in CLOSING and + # requires every rank to retry that same communicator. Retain the + # handle and every registered tensor until teardown succeeds. + communicator.close() + self.communicator = None + self.outputs.clear() + self.input_landing_identities.clear() + with self.producer_direct_lock: + self.producer_direct_logged = False + self.producer_direct_generation = 0 + self.producer_direct_expected.clear() + self.producer_direct_exchanged.clear() + with self.o_producer_direct_lock: + self.o_producer_direct_logged = False + self.o_producer_direct_expected = None + + def _collective_sum(self, value: int) -> int: + """Return a group-wide lifecycle vote on the existing CUDA group.""" + vote = torch.tensor([value], dtype=torch.int32, device=self.device) + dist.all_reduce(vote, op=dist.ReduceOp.SUM, group=self.group) + return int(vote.item()) + + def release_after_denoise(self) -> bool: + """Collectively release registered windows between execution stages. + + The state object remains cached, so the next strict-Ulysses operation + lazily re-enters :meth:`initialize`. Every rank votes before draining + CUDA work and before reporting close success, turning an accidental + asymmetric lifecycle into an explicit failure instead of silent reuse. + """ + if not release_after_denoise_enabled(): + return False + + armed_count = self._collective_sum(int(self.communicator is not None)) + if armed_count == 0: + return False + if armed_count != self.world_size: + raise RuntimeError( + f"FlashInfer PCIe Ulysses lifecycle is rank-asymmetric: {armed_count}/{self.world_size} ranks are armed" + ) + + sync_error: BaseException | None = None + try: + torch.accelerator.synchronize(self.device) + except BaseException as exc: # every rank must vote before failing + sync_error = exc + sync_failures = self._collective_sum(int(sync_error is not None)) + if sync_failures: + raise RuntimeError( + f"FlashInfer PCIe Ulysses release could not drain CUDA work on {sync_failures}/{self.world_size} ranks" + ) from sync_error + + close_error: BaseException | None = None + try: + self.close() + # Registered allocations are now ordinary cached PyTorch storage. + # Return it to the allocator before the much larger video VAE tail. + torch.accelerator.empty_cache() + except BaseException as exc: # report one collective verdict + close_error = exc + close_failures = self._collective_sum(int(close_error is not None)) + if close_failures: + raise RuntimeError( + f"FlashInfer PCIe Ulysses release failed on {close_failures}/{self.world_size} ranks" + ) from close_error + + # VAE decode follows immediately and repopulates PyTorch's cache. The + # next request must return those inactive segments before native RDMA + # output registration; doing it lazily avoids adding another sync/cache + # flush to this request's decode/output critical path. + self.rearm_pending = True + + logger.info( + "FlashInfer PCIe Ulysses released after denoise; next request will re-arm " + "(capacity %.2f GiB per registered slot)", + self.capacity_bytes / 2**30, + ) + return True + + +_STATES: dict[tuple[str, torch.device], _PcieState] = {} + + +def _state_for(x: torch.Tensor, group: dist.ProcessGroup, world_size: int) -> _PcieState: + actual_world_size = dist.get_world_size(group) + if actual_world_size != world_size: + raise RuntimeError( + f"FlashInfer PCIe Ulysses world-size mismatch: caller={world_size}, process-group={actual_world_size}" + ) + key = (_group_name(group), x.device) + with _LOCK: + state = _STATES.get(key) + if state is None: + state = _PcieState( + group=group, + world_size=world_size, + device=x.device, + capacity_bytes=_configured_capacity(x.nbytes), + ) + _STATES[key] = state + elif state.group is not group or state.world_size != world_size: + raise RuntimeError("FlashInfer PCIe Ulysses process-group identity changed for a live cache entry") + return state + + +def _nccl_scatter_heads( + x: torch.Tensor, + group: dist.ProcessGroup, + world_size: int, +) -> torch.Tensor: + batch, local_seq, heads, head_dim = x.shape + local_heads = heads // world_size + send = x.reshape(batch, local_seq, world_size, local_heads, head_dim).transpose(0, 2).contiguous() + recv = torch.empty_like(send) + dist.all_to_all_single(recv, send, group=group) + return ( + recv.reshape(local_seq * world_size, batch, local_heads, head_dim) + .transpose(0, 1) + .contiguous() + .reshape(batch, local_seq * world_size, local_heads, head_dim) + ) + + +def _nccl_gather_heads( + x: torch.Tensor, + group: dist.ProcessGroup, + world_size: int, +) -> torch.Tensor: + batch, global_seq, local_heads, head_dim = x.shape + local_seq = global_seq // world_size + heads = local_heads * world_size + send = ( + x.reshape(batch, world_size, local_seq, local_heads, head_dim) + .transpose(0, 3) + .transpose(0, 1) + .contiguous() + .reshape(world_size, local_heads, local_seq, batch, head_dim) + ) + recv = torch.empty_like(send) + dist.all_to_all_single(recv, send, group=group) + return ( + recv.reshape(heads, local_seq, batch, head_dim) + .transpose(0, 2) + .contiguous() + .reshape(batch, local_seq, heads, head_dim) + ) + + +def _scatter_impl( + x: torch.Tensor, + group_name: str, + world_size: int, + slot: str, + use_sync: bool, +) -> torch.Tensor: + group = _resolve_process_group(group_name) + state = _state_for(x, group, world_size) + if x.nbytes > state.capacity_bytes and _require_rdma(): + raise RuntimeError( + "required FlashInfer all-RDMA Ulysses operand exceeds configured " + f"capacity: operand={x.nbytes}, capacity={state.capacity_bytes}" + ) + if x.nbytes <= state.capacity_bytes and state.initialize(x.dtype): + return state.scatter(x, slot=slot) + if x.nbytes > state.capacity_bytes: + logger.warning_once( + "FlashInfer PCIe Ulysses operand %d exceeds capacity %d; using NCCL for this geometry", + x.nbytes, + state.capacity_bytes, + ) + out = _nccl_scatter_heads(x, group, world_size) + if use_sync: + torch.accelerator.synchronize(x.device) + return out + + +def flashinfer_ulysses_qk_input_landings( + q: torch.Tensor, + k: torch.Tensor, + group: dist.ProcessGroup, + world_size: int, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """Acquire Q/K registered source buffers from an eager graph island. + + ``None`` means the optional PCIe backend declined its cold start and the + caller should produce ordinary Q/K tensors for the NCCL fallback. With + ``VLLM_OMNI_FLASHINFER_ULYSSES_REQUIRE_RDMA=1`` every inability to use the + all-RDMA route remains a hard error. + """ + if torch.compiler.is_compiling(): + raise RuntimeError("FlashInfer input landing buffers must be acquired outside torch.compile") + if q.ndim != 4 or k.ndim != 4: + raise ValueError(f"q and k must be [batch, sequence, heads, head_dim], got {q.shape} and {k.shape}") + if q.shape != k.shape: + raise ValueError(f"producer-direct Q/K requires equal shapes, got {q.shape} and {k.shape}") + if q.dtype != k.dtype or q.device != k.device: + raise ValueError("producer-direct Q/K must have the same dtype and device") + if q.device.type != "cuda": + raise ValueError(f"FlashInfer PCIe input landing requires CUDA tensors, got {q.device}") + with torch.get_device_module().device(q.device): + capturing = torch.cuda.is_current_stream_capturing() + if capturing: + raise RuntimeError("FlashInfer PCIe/RDMA input landing does not support CUDA graph capture") + + state = _state_for(q, group, world_size) + if q.nbytes > state.capacity_bytes: + if _require_rdma(): + raise RuntimeError( + "required FlashInfer all-RDMA Ulysses Q/K exceeds configured " + f"capacity: operand={q.nbytes}, capacity={state.capacity_bytes}" + ) + return None + if not state.initialize(q.dtype): + return None + communicator = state.communicator + assert communicator is not None + if communicator.transport not in ("hybrid", "rdma"): + # All-P2P reads its caller tensor directly and therefore has no + # registered source landing to target. Producing ordinary contiguous + # Q/K is already the zero-staging path for that transport. + return None + q_op = "scatter_heads" + q_slot = "q" + landings = ( + state.input_landing(q, slot=q_slot, op=q_op), + state.input_landing(k, slot="k"), + ) + state.arm_qk_producer_direct(*landings) + return landings + + +def _typed_contiguous_shape_prototype( + source: torch.Tensor, + *, + dtype: torch.dtype, +) -> torch.Tensor: + """Build allocation metadata over source storage without a full temporary.""" + storage = source.untyped_storage() + byte_offset = source.storage_offset() * source.element_size() + if byte_offset < 0 or byte_offset % dtype.itemsize: + raise RuntimeError("source storage cannot back the requested typed prototype") + required_bytes = source.numel() * dtype.itemsize + if byte_offset + required_bytes > storage.nbytes(): + raise RuntimeError( + "source storage is too small for the requested typed prototype: " + f"offset={byte_offset}, required={required_bytes}, " + f"storage={storage.nbytes()}" + ) + shape = tuple(int(dim) for dim in source.shape) + strides = [1] * len(shape) + for index in range(len(shape) - 2, -1, -1): + strides[index] = strides[index + 1] * shape[index + 1] + prototype = torch.empty(0, dtype=dtype, device=source.device) + return prototype.set_( + storage, + byte_offset // dtype.itemsize, + shape, + tuple(strides), + ) + + +def flashinfer_ulysses_o_input_landing( + prototype: torch.Tensor, + group: dist.ProcessGroup, + world_size: int, + *, + input_shape: tuple[int, int, int, int] | None = None, + output_dtype: torch.dtype | None = None, +) -> torch.Tensor | None: + """Acquire the registered source for the next reverse-Ulysses exchange. + + The caller supplies the post-Ulysses attention-output geometry + ``[B, S_global, H_local, D]``. ``input_shape`` may expand only its sequence + dimension for an owner-piggyback payload. ``output_dtype`` may differ from + the query prototype for low-precision forward transport; an already + registered capacity-sized Q/K/V output then backs a typed, shape-only + prototype without a full CUDA allocation. A returned tensor is the exact storage that must be passed to + :func:`flashinfer_ulysses_o_rev`; otherwise FlashInfer stages the source + before the gather-heads collective. + """ + if not is_o_producer_direct_enabled(): + return None + if torch.compiler.is_compiling(): + raise RuntimeError("FlashInfer reverse-O input landing must be acquired outside torch.compile") + if prototype.ndim != 4: + raise ValueError( + "reverse-O producer-direct prototype must be " + f"[batch, sequence, heads, head_dim], got {tuple(prototype.shape)}" + ) + if prototype.device.type != "cuda": + raise ValueError(f"FlashInfer reverse-O input landing requires a CUDA prototype, got {prototype.device}") + requested_dtype = prototype.dtype if output_dtype is None else output_dtype + if requested_dtype not in (torch.float16, torch.bfloat16): + raise ValueError(f"FlashInfer reverse-O producer-direct supports fp16/bf16, got {requested_dtype}") + if prototype.requires_grad: + raise ValueError("FlashInfer reverse-O producer-direct is inference-only") + if not prototype.is_contiguous(): + raise ValueError("FlashInfer reverse-O prototype must be contiguous") + if world_size <= 1: + raise ValueError(f"FlashInfer reverse-O producer-direct requires world_size > 1, got {world_size}") + if any(int(dim) <= 0 for dim in prototype.shape): + raise ValueError(f"FlashInfer reverse-O prototype dimensions must be positive, got {tuple(prototype.shape)}") + if prototype.shape[1] % world_size: + raise ValueError( + "FlashInfer reverse-O global sequence must divide the Ulysses world: " + f"sequence={prototype.shape[1]}, world={world_size}" + ) + landing_shape = tuple(int(dim) for dim in (input_shape or tuple(prototype.shape))) + if len(landing_shape) != 4 or any(dim <= 0 for dim in landing_shape): + raise ValueError( + f"reverse-O producer-direct input_shape must contain four positive dimensions, got {landing_shape}" + ) + if ( + landing_shape[0] != prototype.shape[0] + or landing_shape[2] != prototype.shape[2] + or landing_shape[3] != prototype.shape[3] + ): + raise ValueError( + "reverse-O expanded landing may change only the sequence dimension: " + f"prototype={tuple(prototype.shape)}, requested={landing_shape}" + ) + if landing_shape[1] % world_size: + raise ValueError( + "FlashInfer reverse-O landing sequence must divide the Ulysses world: " + f"sequence={landing_shape[1]}, world={world_size}" + ) + with torch.get_device_module().device(prototype.device): + capturing = torch.cuda.is_current_stream_capturing() + if capturing: + raise RuntimeError("FlashInfer PCIe/RDMA reverse-O landing does not support CUDA graph capture") + + state = _state_for(prototype, group, world_size) + landing_nbytes = int(torch.Size(landing_shape).numel()) * requested_dtype.itemsize + if landing_nbytes > state.capacity_bytes: + if _require_rdma(): + raise RuntimeError( + "required FlashInfer all-RDMA Ulysses reverse-O exceeds " + f"configured capacity: operand={landing_nbytes}, " + f"capacity={state.capacity_bytes}" + ) + logger.warning_once( + "FlashInfer PCIe Ulysses O producer-direct exchange fell back: operand %d exceeds capacity %d", + landing_nbytes, + state.capacity_bytes, + ) + return None + if not state.initialize(requested_dtype): + logger.warning_once( + "FlashInfer PCIe Ulysses O producer-direct exchange fell back: communicator initialization declined" + ) + return None + communicator = state.communicator + assert communicator is not None + if communicator.transport not in ("hybrid", "rdma"): + # All-P2P reads the producer tensor directly and therefore has no + # staging copy for this optimization to remove. + logger.warning_once( + "FlashInfer PCIe Ulysses O producer-direct exchange fell back: " + "transport %s has no registered source landing", + communicator.transport, + ) + return None + + landing_prototype = prototype + if landing_shape != tuple(prototype.shape) or requested_dtype != prototype.dtype: + landing_prototype = state.registered_shape_prototype( + landing_shape, + dtype=requested_dtype, + device=prototype.device, + ) + landing = state.input_landing( + landing_prototype, + slot="o", + op="gather_heads", + ) + state.arm_o_producer_direct(landing) + return landing + + +def _gather_impl( + x: torch.Tensor, + group_name: str, + world_size: int, + use_sync: bool, +) -> torch.Tensor: + group = _resolve_process_group(group_name) + state = _state_for(x, group, world_size) + if x.nbytes > state.capacity_bytes and _require_rdma(): + raise RuntimeError( + "required FlashInfer all-RDMA Ulysses operand exceeds configured " + f"capacity: operand={x.nbytes}, capacity={state.capacity_bytes}" + ) + if x.nbytes <= state.capacity_bytes and state.initialize(x.dtype): + return state.gather(x, slot="o") + if x.nbytes > state.capacity_bytes: + logger.warning_once( + "FlashInfer PCIe Ulysses operand %d exceeds capacity %d; using NCCL for this geometry", + x.nbytes, + state.capacity_bytes, + ) + out = _nccl_gather_heads(x, group, world_size) + if use_sync: + torch.accelerator.synchronize(x.device) + return out + + +@torch.library.custom_op( + "vllm_omni_flashinfer_ulysses::scatter_heads", + mutates_args=(), + device_types="cuda", +) +def flashinfer_ulysses_qkv_fwd( + x: torch.Tensor, + group_name: str, + world_size: int, + slot: str, + use_sync: bool, +) -> torch.Tensor: + """Opaque compiled-graph node for PCIe Ulysses Q/K/V exchange.""" + return _scatter_impl(x, group_name, world_size, slot, use_sync) + + +@flashinfer_ulysses_qkv_fwd.register_fake +def _( + x: torch.Tensor, + group_name: str, + world_size: int, + slot: str, + use_sync: bool, +) -> torch.Tensor: + batch, local_seq, heads, head_dim = x.shape + return x.new_empty(batch, local_seq * world_size, heads // world_size, head_dim) + + +@torch.library.custom_op( + "vllm_omni_flashinfer_ulysses::gather_heads", + mutates_args=(), + device_types="cuda", +) +def flashinfer_ulysses_o_rev( + x: torch.Tensor, + group_name: str, + world_size: int, + use_sync: bool, +) -> torch.Tensor: + """Opaque compiled-graph node for inverse PCIe Ulysses output exchange.""" + return _gather_impl(x, group_name, world_size, use_sync) + + +@flashinfer_ulysses_o_rev.register_fake +def _( + x: torch.Tensor, + group_name: str, + world_size: int, + use_sync: bool, +) -> torch.Tensor: + batch, global_seq, local_heads, head_dim = x.shape + return x.new_empty( + batch, + global_seq // world_size, + local_heads * world_size, + head_dim, + ) + + +def clear_flashinfer_ulysses_communicators() -> None: + """Collectively close all process-local communicators before PG teardown.""" + with _LOCK: + states = list(_STATES.items()) + closed: list[tuple[tuple[str, torch.device], _PcieState]] = [] + errors: list[Exception] = [] + for key, state in states: + try: + state.close() + closed.append((key, state)) + except Exception as exc: # teardown errors must remain visible to shutdown logs + errors.append(exc) + with _LOCK: + for key, state in closed: + if _STATES.get(key) is state: + del _STATES[key] + if errors: + raise RuntimeError(f"failed to close {len(errors)} FlashInfer Ulysses communicator(s): {errors}") + + +def release_flashinfer_ulysses_after_denoise() -> bool: + """Release every process-local FlashInfer Ulysses state, if opted in. + + Call this on every rank at the request/step boundary. Keeping + the states in ``_STATES`` preserves their process-group and capacity + contracts while allowing the communicator itself to be initialized lazily + by the next request. + """ + if not release_after_denoise_enabled(): + return False + with _LOCK: + # Match collective order across ranks even if unrelated Python object + # construction happened in a different order in each worker. + states = [ + state + for _key, state in sorted( + _STATES.items(), + key=lambda item: (item[0][0], str(item[0][1])), + ) + ] + released = False + for state in states: + released = state.release_after_denoise() or released + return released + + +__all__ = [ + "clear_flashinfer_ulysses_communicators", + "configured_backend", + "ensure_flashinfer_pcie_available", + "flashinfer_ulysses_o_input_landing", + "flashinfer_ulysses_o_rev", + "flashinfer_ulysses_qk_input_landings", + "flashinfer_ulysses_qkv_fwd", + "is_flashinfer_pcie_enabled", + "is_o_producer_direct_enabled", + "is_qk_producer_direct_enabled", + "release_after_denoise_enabled", + "release_flashinfer_ulysses_after_denoise", +] diff --git a/vllm_omni/diffusion/layers/fused_qk_norm_rope.py b/vllm_omni/diffusion/layers/fused_qk_norm_rope.py index 9f66057262e..ba4e40ddcc0 100644 --- a/vllm_omni/diffusion/layers/fused_qk_norm_rope.py +++ b/vllm_omni/diffusion/layers/fused_qk_norm_rope.py @@ -11,15 +11,15 @@ ``[cos(theta), sin(theta)]`` with ``theta`` of width ``rotary_dim // 2``; its dtype is either the activation dtype or float32; * ``interleaved=False`` rotates half-split pairs ``(d, d + rotary_dim/2)`` - (MiniMax-H3); ``interleaved=True`` rotates adjacent pairs ``(2i, 2i + 1)`` + (three-axis); ``interleaved=True`` rotates adjacent pairs ``(2i, 2i + 1)`` with ``theta_i`` (Boogu-Image, matching its ``apply_rotary_emb``). The CUDA fast path fuses RMSNorm and RoPE without materializing normalized Q/K or rotary-product intermediates. The interleaved mode supports any even ``rotary_dim <= head_dim <= 256`` (``tl.arange`` padding to the next power of -two); the half-split mode keeps the pre-existing kernel and its MiniMax-H3 +two); the half-split mode keeps the pre-existing kernel and its three-axis geometry contract (``head_dim == 128``, ``rotary_dim == 96``). Ascend -composes its RMSNorm and rotary fused primitives on the MiniMax-H3 geometry; +composes its RMSNorm and rotary fused primitives on the three-axis rotary geometry; unsupported inputs use the eager reference. """ @@ -30,12 +30,12 @@ import torch import torch.nn.functional as F from torch.library import Library -from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, tl, triton from vllm.utils.torch_utils import direct_register_custom_op from vllm_omni.diffusion import envs from vllm_omni.platforms import current_omni_platform +from vllm_omni.platforms import current_omni_platform as current_platform _FUSED_HEAD_DIM = 128 _FUSED_ROTARY_DIM = 96 @@ -63,7 +63,7 @@ def _apply_rope_table( ``[cos(theta) | sin(theta)]`` with ``theta`` of width ``rotary_dim // 2``. ``interleaved`` selects the pairing: half-split ``(d, d + rotary_dim/2)`` - sharing ``theta_d`` (MiniMax-H3) or adjacent pairs ``(2i, 2i + 1)`` + sharing ``theta_d`` (three-axis) or adjacent pairs ``(2i, 2i + 1)`` sharing ``theta_i`` (Boogu-Image). ``dtype`` is the precision the rotation arithmetic runs in and ``output_dtype`` the precision the result is cast to; both default to ``x``'s dtype (the historical half-split @@ -166,135 +166,6 @@ def _rms_norm_rope_kernel( out_offsets = token * out_stride_t + heads[:, None] * out_stride_h + dims[None, :] * out_stride_d tl.store(out_ptr + out_offsets, output, mask=mask) - @triton.jit - def _qk_norm_rope_kernel( - q_ptr, - k_ptr, - q_weight_ptr, - k_weight_ptr, - rope_table_ptr, - out_q_ptr, - out_k_ptr, - q_stride_t, - q_stride_h, - q_stride_d, - k_stride_t, - k_stride_h, - k_stride_d, - rope_stride_t, - out_q_stride_t, - out_q_stride_h, - out_q_stride_d, - out_k_stride_t, - out_k_stride_h, - out_k_stride_d, - num_q_heads: tl.constexpr, - num_kv_heads: tl.constexpr, - head_dim: tl.constexpr, - padded_dim: tl.constexpr, # next power of 2 >= head_dim (tl.arange needs a power of 2) - rotary_half: tl.constexpr, - eps: tl.constexpr, - heads_per_program: tl.constexpr, - q_head_groups: tl.constexpr, - interleaved: tl.constexpr, - ): - """Fused per-head RMSNorm + RoPE for Q and K in one launch. - - Grid axis 1 assigns the first ``q_head_groups`` programs of a token - to Q and the rest to K, so the small K grid hides inside Q's waves - instead of paying its own latency-bound launch. ``interleaved`` - selects the pairing, mirroring ``_apply_rope_table``: - - * ``True`` — adjacent pairs ``(2i, 2i+1)`` sharing ``theta_i``. - * ``False`` — half-split pairs ``(d, d + rotary_half)`` sharing - ``theta_d``. Production traffic still routes half-split inputs - through ``_rms_norm_rope_kernel``; this mode is test-covered and - awaits the maintainers' call before taking over that path. - - The modes differ only in the pair/frequency/selector index - expressions; the arithmetic is shared. Beyond the rotary width - (partial rotary or lane padding) the output is the normalized value - unchanged. - """ - token = tl.program_id(0) - head_group = tl.program_id(1) - dims = tl.arange(0, padded_dim) - dim_mask = dims < head_dim - if head_group < q_head_groups: - heads = head_group * heads_per_program + tl.arange(0, heads_per_program) - # Valid-domain mask: rows beyond the head count (when it is not a - # multiple of heads_per_program) and padded lanes beyond head_dim. - mask = (heads[:, None] < num_q_heads) & dim_mask[None, :] - in_ptr = q_ptr - weight_ptr = q_weight_ptr - out_ptr = out_q_ptr - in_stride_t, in_stride_h, in_stride_d = q_stride_t, q_stride_h, q_stride_d - out_stride_t, out_stride_h, out_stride_d = ( - out_q_stride_t, - out_q_stride_h, - out_q_stride_d, - ) - else: - heads = (head_group - q_head_groups) * heads_per_program + tl.arange(0, heads_per_program) - mask = (heads[:, None] < num_kv_heads) & dim_mask[None, :] - in_ptr = k_ptr - weight_ptr = k_weight_ptr - out_ptr = out_k_ptr - in_stride_t, in_stride_h, in_stride_d = k_stride_t, k_stride_h, k_stride_d - out_stride_t, out_stride_h, out_stride_d = ( - out_k_stride_t, - out_k_stride_h, - out_k_stride_d, - ) - - offsets = token * in_stride_t + heads[:, None] * in_stride_h + dims[None, :] * in_stride_d - x = tl.load(in_ptr + offsets, mask=mask, other=0.0).to(tl.float32) - weight = tl.load(weight_ptr + dims, mask=dim_mask, other=0.0).to(tl.float32) - # Padding lanes load 0 and do not perturb the sum; divide by the true - # head_dim, not the padded width. - inv_rms = tl.rsqrt(tl.sum(x * x, axis=1) / head_dim + eps) - normalized = (x * inv_rms[:, None] * weight[None, :]).to(tl.bfloat16) - - # The two pairings differ only in three index expressions; the pair - # reload, the table gather, the rotation formulas and the write-back - # are shared (the same structure as _apply_rope_table). - rotary_dim = rotary_half * 2 - if interleaved: - # Adjacent pairs (2i, 2i+1) sharing theta_i. - pair_dims = tl.where(dims < rotary_dim, dims ^ 1, dims) - freq_dims = tl.where(dims < rotary_dim, dims // 2, 0) - is_first = (dims % 2) == 0 - else: - # Half-split pairs (d, d + rotary_half) sharing theta_d. - pair_dims = tl.where( - dims < rotary_half, - dims + rotary_half, - tl.where(dims < rotary_dim, dims - rotary_half, dims), - ) - freq_dims = tl.where( - dims < rotary_half, - dims, - tl.where(dims < rotary_dim, dims - rotary_half, 0), - ) - is_first = dims < rotary_half - pair_offsets = token * in_stride_t + heads[:, None] * in_stride_h + pair_dims[None, :] * in_stride_d - pair_x = tl.load(in_ptr + pair_offsets, mask=mask, other=0.0).to(tl.float32) - pair_weight = tl.load(weight_ptr + pair_dims, mask=pair_dims < head_dim, other=0.0).to(tl.float32) - pair_normalized = (pair_x * inv_rms[:, None] * pair_weight[None, :]).to(tl.bfloat16) - table_offsets = token * rope_stride_t + freq_dims - cos = tl.load(rope_table_ptr + table_offsets).to(tl.float32) - sin = tl.load(rope_table_ptr + table_offsets + rotary_half).to(tl.float32) - first = normalized.to(tl.float32) * cos - pair_normalized.to(tl.float32) * sin - second = normalized.to(tl.float32) * cos + pair_normalized.to(tl.float32) * sin - output = tl.where( - dims < rotary_dim, - tl.where(is_first, first, second), - normalized.to(tl.float32), - ) - - out_offsets = token * out_stride_t + heads[:, None] * out_stride_h + dims[None, :] * out_stride_d - tl.store(out_ptr + out_offsets, output, mask=mask) - def _eager_qk_norm_rope( q: torch.Tensor, @@ -310,7 +181,7 @@ def _eager_qk_norm_rope( q_norm = F.rms_norm(q, (head_dim,), q_weight, eps) k_norm = F.rms_norm(k, (head_dim,), k_weight, eps) # fp32 rotation for the interleaved (Boogu) semantics; the half-split - # (H3) reference keeps its historical x-dtype arithmetic via the defaults. + # half-split reference keeps its historical x-dtype arithmetic via the defaults. rope_dtype = torch.float32 if interleaved else None return ( _apply_rope_table(q_norm, rope_table, rotary_dim, interleaved=interleaved, dtype=rope_dtype), @@ -323,13 +194,13 @@ def _npu_apply_rope_table( rope_table: torch.Tensor, rotary_dim: int, ) -> torch.Tensor: - """Apply H3's rotary dimensions through Ascend's fused rotary kernel. + """Apply packed rotary dimensions through Ascend's fused rotary kernel. ``mindiesd.rotary_position_embedding`` takes a 4-D BSND tensor, whereas - MiniMax-H3 keeps packed activations as ``[tokens, heads, head_dim]``. + The packed input keeps packed activations as ``[tokens, heads, head_dim]``. The shared MindIE-SD wrapper normalizes this 3-D layout to BSND and restores it afterwards. It receives the half-width cos/sin values from - the packed table; the wrapper expands them to H3's non-interleaved 96-D + the packed table; the wrapper expands them to the non-interleaved 96-D rotary layout. The 32 non-rotary head dimensions bypass the kernel. Some CANN environments package ``torch_npu`` without MindIE-SD. Keep @@ -354,7 +225,7 @@ def _npu_apply_rope_table( half_head_dim=True, ) else: - # npu_rotary_mul uses BSND and full rotary-width cos/sin. H3 uses + # npu_rotary_mul uses BSND and full rotary-width cos/sin. The packed table uses # NeoX/rotated-half ordering, so duplicate rather than interleave the # half-width table along the final dimension. cos = cos.unsqueeze(0).unsqueeze(2).repeat(1, 1, 1, 2) @@ -373,7 +244,7 @@ def _npu_qk_norm_rope( eps: float, rotary_dim: int, ) -> tuple[torch.Tensor, torch.Tensor]: - """Use Ascend RMSNorm and RoPE fused primitives for MiniMax-H3 DiT.""" + """Use Ascend RMSNorm and RoPE fused primitives for packed diffusion attention.""" import torch_npu q_norm = torch_npu.npu_rms_norm(q, q_weight, epsilon=eps)[0] @@ -403,7 +274,7 @@ def _fused_cuda_supported( if interleaved: return rotary_dim % 2 == 0 and 2 <= rotary_dim <= head_dim <= _FUSED_MAX_HEAD_DIM # Half-split traffic keeps the pre-existing per-tensor kernel and its - # exact MiniMax-H3 geometry contract (that kernel is untouched by the + # exact three-axis rotary geometry contract (that kernel is untouched by the # interleaved extension). return head_dim == _FUSED_HEAD_DIM and rotary_dim == _FUSED_ROTARY_DIM @@ -415,7 +286,7 @@ def _fused_npu_supported( rotary_dim: int, interleaved: bool = False, ) -> bool: - """Return whether the MiniMax-H3 Ascend fused-op contract is satisfied.""" + """Return whether the packed Ascend fused-op contract is satisfied.""" return ( not interleaved and current_omni_platform.is_npu() @@ -433,10 +304,13 @@ def _launch_fused_rms_norm_rope( weight: torch.Tensor, rope_table: torch.Tensor, eps: float, + *, + out: torch.Tensor | None = None, ) -> torch.Tensor: tokens, heads, head_dim = x.shape rotary_half = rope_table.shape[-1] // 2 - out = torch.empty(x.shape, dtype=x.dtype, device=x.device) + if out is None: + out = torch.empty(x.shape, dtype=x.dtype, device=x.device) if tokens == 0: return out grid = (tokens, triton.cdiv(heads, _HEADS_PER_PROGRAM)) @@ -462,67 +336,56 @@ def _launch_fused_rms_norm_rope( return out -def _launch_fused_qk_norm_rope( +def _validate_qk_norm_rope_inputs( q: torch.Tensor, k: torch.Tensor, q_weight: torch.Tensor, k_weight: torch.Tensor, rope_table: torch.Tensor, - eps: float, - interleaved: bool, -) -> tuple[torch.Tensor, torch.Tensor]: - tokens, q_heads, head_dim = q.shape - kv_heads = k.shape[1] - rotary_half = rope_table.shape[-1] // 2 - out_q = torch.empty(q.shape, dtype=q.dtype, device=q.device) - out_k = torch.empty(k.shape, dtype=k.dtype, device=k.device) - if tokens == 0: - return out_q, out_k - if interleaved: - heads_per_program = _INTERLEAVED_HEADS_PER_PROGRAM - num_warps, num_stages = 1, 2 - else: - # Match the historical per-tensor launch shape (8 heads per program, - # 8 warps) so the half-split arithmetic keeps its reduction tree. - heads_per_program = _HEADS_PER_PROGRAM - num_warps, num_stages = 8, 3 - q_head_groups = triton.cdiv(q_heads, heads_per_program) - kv_head_groups = triton.cdiv(kv_heads, heads_per_program) - grid = (tokens, q_head_groups + kv_head_groups) - _qk_norm_rope_kernel[grid]( - q, - k, - q_weight, - k_weight, - rope_table, - out_q, - out_k, - q.stride(0), - q.stride(1), - q.stride(2), - k.stride(0), - k.stride(1), - k.stride(2), - rope_table.stride(0), - out_q.stride(0), - out_q.stride(1), - out_q.stride(2), - out_k.stride(0), - out_k.stride(1), - out_k.stride(2), - num_q_heads=q_heads, - num_kv_heads=kv_heads, - head_dim=head_dim, - padded_dim=triton.next_power_of_2(head_dim), - rotary_half=rotary_half, - eps=eps, - heads_per_program=heads_per_program, - q_head_groups=q_head_groups, - interleaved=interleaved, - num_warps=num_warps, - num_stages=num_stages, - ) - return out_q, out_k + head_dim: int | None, + rotary_dim: int | None, +) -> tuple[int, int]: + if q.ndim != 3 or k.ndim != 3: + raise ValueError(f"q and k must be [tokens, heads, head_dim], got {q.shape} and {k.shape}") + if q.shape[0] != k.shape[0] or q.shape[2] != k.shape[2]: + raise ValueError(f"q and k shapes are incompatible: {q.shape} and {k.shape}") + if q.dtype != k.dtype or q.device != k.device: + raise ValueError("q and k must have the same dtype and device") + if q.dtype not in (torch.bfloat16, torch.float16, torch.float32): + raise TypeError(f"Fused QK RMSNorm/RoPE requires floating inputs, got {q.dtype}") + + head_dim = q.shape[-1] if head_dim is None else head_dim + rotary_dim = rope_table.shape[-1] if rotary_dim is None else rotary_dim + if q.shape[-1] != head_dim: + raise ValueError(f"Expected q/k head_dim={head_dim}, got {q.shape[-1]}") + if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2: + raise ValueError(f"rotary_dim must be even and in [2, {head_dim}], got {rotary_dim}") + if q_weight.shape != (head_dim,) or k_weight.shape != (head_dim,): + raise ValueError(f"Expected norm weights [{head_dim}], got {tuple(q_weight.shape)} and {tuple(k_weight.shape)}") + if q_weight.device != q.device or k_weight.device != q.device: + raise ValueError("Q/K norm weights must be on the activation device") + if rope_table.device != q.device or rope_table.dtype != q.dtype: + raise ValueError("rope_table must have the same dtype and device as q/k") + if rope_table.shape != (q.shape[0], rotary_dim): + raise ValueError(f"Expected rope_table [{q.shape[0]}, {rotary_dim}], got {tuple(rope_table.shape)}") + return head_dim, rotary_dim + + +def _validate_qk_norm_rope_outputs( + q: torch.Tensor, + k: torch.Tensor, + q_out: torch.Tensor, + k_out: torch.Tensor, +) -> None: + for name, source, out in (("q", q, q_out), ("k", k, k_out)): + if out.shape != source.shape: + raise ValueError(f"{name}_out must have shape {tuple(source.shape)}, got {tuple(out.shape)}") + if out.dtype != source.dtype or out.device != source.device: + raise ValueError(f"{name}_out must have the same dtype and device as {name}") + if not out.is_contiguous(): + raise ValueError(f"{name}_out must be contiguous") + if q_out.data_ptr() == k_out.data_ptr() and q_out.numel() != 0: + raise ValueError("q_out and k_out must use distinct storage") def _fused_qk_norm_rope_impl( @@ -639,33 +502,16 @@ def fused_qk_norm_rope( rotary_dim: int | None = None, interleaved: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: - """Apply Q/K RMSNorm and packed RoPE (half-split or adjacent pairs).""" - if q.ndim != 3 or k.ndim != 3: - raise ValueError(f"q and k must be [tokens, heads, head_dim], got {q.shape} and {k.shape}") - if q.shape[0] != k.shape[0] or q.shape[2] != k.shape[2]: - raise ValueError(f"q and k shapes are incompatible: {q.shape} and {k.shape}") - if q.dtype != k.dtype or q.device != k.device: - raise ValueError("q and k must have the same dtype and device") - if q.dtype not in (torch.bfloat16, torch.float16, torch.float32): - raise TypeError(f"Fused QK RMSNorm/RoPE requires floating inputs, got {q.dtype}") - - head_dim = q.shape[-1] if head_dim is None else head_dim - rotary_dim = rope_table.shape[-1] if rotary_dim is None else rotary_dim - if q.shape[-1] != head_dim: - raise ValueError(f"Expected q/k head_dim={head_dim}, got {q.shape[-1]}") - if rotary_dim <= 0 or rotary_dim > head_dim or rotary_dim % 2: - raise ValueError(f"rotary_dim must be even and in [2, {head_dim}], got {rotary_dim}") - if q_weight.shape != (head_dim,) or k_weight.shape != (head_dim,): - raise ValueError(f"Expected norm weights [{head_dim}], got {tuple(q_weight.shape)} and {tuple(k_weight.shape)}") - if q_weight.device != q.device or k_weight.device != q.device: - raise ValueError("Q/K norm weights must be on the activation device") - if rope_table.device != q.device or rope_table.dtype not in ( - q.dtype, - torch.float32, - ): - raise ValueError("rope_table must be on q/k's device with their dtype or float32") - if rope_table.shape != (q.shape[0], rotary_dim): - raise ValueError(f"Expected rope_table [{q.shape[0]}, {rotary_dim}], got {tuple(rope_table.shape)}") + """Apply Q/K RMSNorm and packed non-interleaved RoPE.""" + head_dim, rotary_dim = _validate_qk_norm_rope_inputs( + q, + k, + q_weight, + k_weight, + rope_table, + head_dim, + rotary_dim, + ) q_weight = q_weight.contiguous() k_weight = k_weight.contiguous() @@ -705,4 +551,261 @@ def fused_qk_norm_rope( ) -__all__ = ["fused_qk_norm_rope", "fused_qk_norm_rope_min_tokens"] +@torch.compiler.disable +def fused_qk_norm_rope_out( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rope_table: torch.Tensor, + eps: float, + *, + q_out: torch.Tensor, + k_out: torch.Tensor, + head_dim: int | None = None, + rotary_dim: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the CUDA fused kernel into caller-owned contiguous storage. + + This eager-only entry point is for registered communication landing + buffers. It deliberately does not masquerade as a functional custom op: + the caller owns and reuses ``q_out``/``k_out``, so mutation must remain + visible to Dynamo and CUDA graph capture must stay disabled. + """ + head_dim, rotary_dim = _validate_qk_norm_rope_inputs( + q, + k, + q_weight, + k_weight, + rope_table, + head_dim, + rotary_dim, + ) + _validate_qk_norm_rope_outputs(q, k, q_out, k_out) + if not _fused_cuda_supported(q, k, head_dim, rotary_dim): + raise RuntimeError( + "caller-owned Q/K output storage requires the CUDA Triton " + f"fast path (bf16, head_dim={_FUSED_HEAD_DIM}, " + f"rotary_dim={_FUSED_ROTARY_DIM})" + ) + + q_weight = q_weight.contiguous() + k_weight = k_weight.contiguous() + rope_table = rope_table.contiguous() + return ( + _launch_fused_rms_norm_rope( + q, + q_weight, + rope_table, + eps, + out=q_out, + ), + _launch_fused_rms_norm_rope( + k, + k_weight, + rope_table, + eps, + out=k_out, + ), + ) + + +__all__ = [ + "fused_qk_norm_rope", + "fused_qk_norm_rope_out", +] + + +if HAS_TRITON: + + @triton.jit + def _qk_norm_rope_kernel( + q_ptr, + k_ptr, + q_weight_ptr, + k_weight_ptr, + rope_table_ptr, + out_q_ptr, + out_k_ptr, + q_stride_t, + q_stride_h, + q_stride_d, + k_stride_t, + k_stride_h, + k_stride_d, + rope_stride_t, + out_q_stride_t, + out_q_stride_h, + out_q_stride_d, + out_k_stride_t, + out_k_stride_h, + out_k_stride_d, + num_q_heads: tl.constexpr, + num_kv_heads: tl.constexpr, + head_dim: tl.constexpr, + padded_dim: tl.constexpr, # next power of 2 >= head_dim (tl.arange needs a power of 2) + rotary_half: tl.constexpr, + eps: tl.constexpr, + heads_per_program: tl.constexpr, + q_head_groups: tl.constexpr, + interleaved: tl.constexpr, + ): + """Fused per-head RMSNorm + RoPE for Q and K in one launch. + + Grid axis 1 assigns the first ``q_head_groups`` programs of a token + to Q and the rest to K, so the small K grid hides inside Q's waves + instead of paying its own latency-bound launch. ``interleaved`` + selects the pairing, mirroring ``_apply_rope_table``: + + * ``True`` — adjacent pairs ``(2i, 2i+1)`` sharing ``theta_i``. + * ``False`` — half-split pairs ``(d, d + rotary_half)`` sharing + ``theta_d``. Production traffic still routes half-split inputs + through ``_rms_norm_rope_kernel``; this mode is test-covered and + awaits the maintainers' call before taking over that path. + + The modes differ only in the pair/frequency/selector index + expressions; the arithmetic is shared. Beyond the rotary width + (partial rotary or lane padding) the output is the normalized value + unchanged. + """ + token = tl.program_id(0) + head_group = tl.program_id(1) + dims = tl.arange(0, padded_dim) + dim_mask = dims < head_dim + if head_group < q_head_groups: + heads = head_group * heads_per_program + tl.arange(0, heads_per_program) + # Valid-domain mask: rows beyond the head count (when it is not a + # multiple of heads_per_program) and padded lanes beyond head_dim. + mask = (heads[:, None] < num_q_heads) & dim_mask[None, :] + in_ptr = q_ptr + weight_ptr = q_weight_ptr + out_ptr = out_q_ptr + in_stride_t, in_stride_h, in_stride_d = q_stride_t, q_stride_h, q_stride_d + out_stride_t, out_stride_h, out_stride_d = ( + out_q_stride_t, + out_q_stride_h, + out_q_stride_d, + ) + else: + heads = (head_group - q_head_groups) * heads_per_program + tl.arange(0, heads_per_program) + mask = (heads[:, None] < num_kv_heads) & dim_mask[None, :] + in_ptr = k_ptr + weight_ptr = k_weight_ptr + out_ptr = out_k_ptr + in_stride_t, in_stride_h, in_stride_d = k_stride_t, k_stride_h, k_stride_d + out_stride_t, out_stride_h, out_stride_d = ( + out_k_stride_t, + out_k_stride_h, + out_k_stride_d, + ) + + offsets = token * in_stride_t + heads[:, None] * in_stride_h + dims[None, :] * in_stride_d + x = tl.load(in_ptr + offsets, mask=mask, other=0.0).to(tl.float32) + weight = tl.load(weight_ptr + dims, mask=dim_mask, other=0.0).to(tl.float32) + # Padding lanes load 0 and do not perturb the sum; divide by the true + # head_dim, not the padded width. + inv_rms = tl.rsqrt(tl.sum(x * x, axis=1) / head_dim + eps) + normalized = (x * inv_rms[:, None] * weight[None, :]).to(tl.bfloat16) + + # The two pairings differ only in three index expressions; the pair + # reload, the table gather, the rotation formulas and the write-back + # are shared (the same structure as _apply_rope_table). + rotary_dim = rotary_half * 2 + if interleaved: + # Adjacent pairs (2i, 2i+1) sharing theta_i. + pair_dims = tl.where(dims < rotary_dim, dims ^ 1, dims) + freq_dims = tl.where(dims < rotary_dim, dims // 2, 0) + is_first = (dims % 2) == 0 + else: + # Half-split pairs (d, d + rotary_half) sharing theta_d. + pair_dims = tl.where( + dims < rotary_half, + dims + rotary_half, + tl.where(dims < rotary_dim, dims - rotary_half, dims), + ) + freq_dims = tl.where( + dims < rotary_half, + dims, + tl.where(dims < rotary_dim, dims - rotary_half, 0), + ) + is_first = dims < rotary_half + pair_offsets = token * in_stride_t + heads[:, None] * in_stride_h + pair_dims[None, :] * in_stride_d + pair_x = tl.load(in_ptr + pair_offsets, mask=mask, other=0.0).to(tl.float32) + pair_weight = tl.load(weight_ptr + pair_dims, mask=pair_dims < head_dim, other=0.0).to(tl.float32) + pair_normalized = (pair_x * inv_rms[:, None] * pair_weight[None, :]).to(tl.bfloat16) + table_offsets = token * rope_stride_t + freq_dims + cos = tl.load(rope_table_ptr + table_offsets).to(tl.float32) + sin = tl.load(rope_table_ptr + table_offsets + rotary_half).to(tl.float32) + first = normalized.to(tl.float32) * cos - pair_normalized.to(tl.float32) * sin + second = normalized.to(tl.float32) * cos + pair_normalized.to(tl.float32) * sin + output = tl.where( + dims < rotary_dim, + tl.where(is_first, first, second), + normalized.to(tl.float32), + ) + + out_offsets = token * out_stride_t + heads[:, None] * out_stride_h + dims[None, :] * out_stride_d + tl.store(out_ptr + out_offsets, output, mask=mask) + + +def _launch_fused_qk_norm_rope( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + rope_table: torch.Tensor, + eps: float, + interleaved: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + tokens, q_heads, head_dim = q.shape + kv_heads = k.shape[1] + rotary_half = rope_table.shape[-1] // 2 + out_q = torch.empty(q.shape, dtype=q.dtype, device=q.device) + out_k = torch.empty(k.shape, dtype=k.dtype, device=k.device) + if tokens == 0: + return out_q, out_k + if interleaved: + heads_per_program = _INTERLEAVED_HEADS_PER_PROGRAM + num_warps, num_stages = 1, 2 + else: + # Match the historical per-tensor launch shape (8 heads per program, + # 8 warps) so the half-split arithmetic keeps its reduction tree. + heads_per_program = _HEADS_PER_PROGRAM + num_warps, num_stages = 8, 3 + q_head_groups = triton.cdiv(q_heads, heads_per_program) + kv_head_groups = triton.cdiv(kv_heads, heads_per_program) + grid = (tokens, q_head_groups + kv_head_groups) + _qk_norm_rope_kernel[grid]( + q, + k, + q_weight, + k_weight, + rope_table, + out_q, + out_k, + q.stride(0), + q.stride(1), + q.stride(2), + k.stride(0), + k.stride(1), + k.stride(2), + rope_table.stride(0), + out_q.stride(0), + out_q.stride(1), + out_q.stride(2), + out_k.stride(0), + out_k.stride(1), + out_k.stride(2), + num_q_heads=q_heads, + num_kv_heads=kv_heads, + head_dim=head_dim, + padded_dim=triton.next_power_of_2(head_dim), + rotary_half=rotary_half, + eps=eps, + heads_per_program=heads_per_program, + q_head_groups=q_head_groups, + interleaved=interleaved, + num_warps=num_warps, + num_stages=num_stages, + ) + return out_q, out_k diff --git a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py b/vllm_omni/diffusion/layers/indexed_modulation.py similarity index 69% rename from vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py rename to vllm_omni/diffusion/layers/indexed_modulation.py index 20327b7e945..9e7f044aa1e 100644 --- a/vllm_omni/diffusion/attention/ops/minimax_h3_modulation.py +++ b/vllm_omni/diffusion/layers/indexed_modulation.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project -"""Fused MiniMax H3 modulation with FP32 accumulation.""" +"""Fused tensor-indexed modulation with FP32 accumulation.""" from __future__ import annotations @@ -244,7 +244,7 @@ def rms_norm_indexed_scale_shift( indices: torch.Tensor, eps: float, ) -> torch.Tensor: - """Fuse H3 RMSNorm with its indexed AdaLN affine transform.""" + """Fuse RMSNorm with a row-indexed affine transform.""" if x.is_cpu: input_dtype = x.dtype normalized = x.float() @@ -336,3 +336,111 @@ def indexed_gate_rms_norm_scale_shift( "indexed_scale_shift_", "rms_norm_indexed_scale_shift", ] + + +@triton.jit +def _bf16_indexed_gate_add_kernel( + bundle_ptr, + gate_ptr, + local_row_to_tail_ptr, + bundle_stride_batch, + bundle_stride_seq, + bundle_stride_head, + bundle_stride_dim, + gate_stride_batch, + gate_stride_seq, + gate_stride_head, + gate_stride_dim, + local_rows: tl.constexpr, + head_dim: tl.constexpr, + row_width: tl.constexpr, + feature_block: tl.constexpr, +): + flat_local_row = tl.program_id(0) + feature_block_index = tl.program_id(1) + batch_index = flat_local_row // local_rows + local_row = flat_local_row - batch_index * local_rows + tail_slot = tl.load(local_row_to_tail_ptr + local_row) + valid_row = tail_slot >= 0 + + features = feature_block_index * feature_block + tl.arange(0, feature_block) + feature_mask = features < row_width + head_index = features // head_dim + dim_index = features - head_index * head_dim + fine_offsets = ( + batch_index * bundle_stride_batch + + local_row * bundle_stride_seq + + head_index * bundle_stride_head + + dim_index * bundle_stride_dim + ) + coarse_offsets = ( + batch_index * bundle_stride_batch + + (local_rows + tail_slot) * bundle_stride_seq + + head_index * bundle_stride_head + + dim_index * bundle_stride_dim + ) + gate_offsets = ( + batch_index * gate_stride_batch + + local_row * gate_stride_seq + + head_index * gate_stride_head + + dim_index * gate_stride_dim + ) + mask = valid_row & feature_mask + fine_values = tl.load(bundle_ptr + fine_offsets, mask=mask, other=0.0) + coarse_values = tl.load(bundle_ptr + coarse_offsets, mask=mask, other=0.0) + gate_values = tl.load(gate_ptr + gate_offsets, mask=mask, other=0.0) + values = tl.inline_asm_elementwise( + """ + { + .reg .b32 gated; + mul.rn.bf16x2 gated, $2, $3; + add.rn.bf16x2 $0, $1, gated; + } + """, + constraints="=r,r,r,r", + args=[fine_values, coarse_values, gate_values], + dtype=tl.bfloat16, + is_pure=True, + pack=2, + ) + tl.store(bundle_ptr + fine_offsets, values, mask=mask) + + +def bf16_indexed_gate_add_(bundle: torch.Tensor, gate: torch.Tensor, row_indices: torch.Tensor) -> torch.Tensor: + """Add indexed coarse rows to a fine prefix with separate BF16 mul/add rounding. + + ``bundle`` is [B, fine_rows + coarse_rows, H, D]; ``gate`` covers the fine + prefix. Negative indices leave a fine row unchanged. The caller validates + nonnegative indices against the coarse row count before device execution. + The output aliases the fine prefix and the coarse tail remains unchanged. + """ + if bundle.ndim != 4 or gate.ndim != 4 or bundle.shape[::2] != gate.shape[::2] or bundle.shape[-1] != gate.shape[-1]: + raise ValueError("bundle and gate must be BSHD tensors with matching batch/head dimensions") + rows = gate.shape[1] + if rows < 1 or bundle.shape[1] <= rows or bundle.shape[-1] % 2: + raise ValueError("expected a nonempty fine prefix, coarse tail and even head dimension") + if bundle.dtype != torch.bfloat16 or gate.dtype != torch.bfloat16: + raise TypeError("indexed gate addition requires BF16 tensors") + if row_indices.shape != (rows,) or row_indices.dtype != torch.int32: + raise ValueError("row_indices must be INT32 with one entry per fine row") + if not all( + t.is_cuda and t.device == bundle.device and t.is_contiguous() and not t.requires_grad + for t in (bundle, gate, row_indices) + ): + raise ValueError("indexed gate addition requires contiguous inference tensors on one CUDA device") + if bundle.untyped_storage().data_ptr() == gate.untyped_storage().data_ptr(): + raise ValueError("gate must not alias the mutable output bundle") + width = bundle.shape[2] * bundle.shape[3] + _bf16_indexed_gate_add_kernel[(bundle.shape[0] * rows, triton.cdiv(width, 1024))]( + bundle, + gate, + row_indices, + *bundle.stride(), + *gate.stride(), + local_rows=rows, + head_dim=bundle.shape[3], + row_width=width, + feature_block=1024, + num_warps=8, + ) + return bundle[:, :rows] diff --git a/vllm_omni/diffusion/layers/mxfp8.py b/vllm_omni/diffusion/layers/mxfp8.py new file mode 100644 index 00000000000..9ef5ee7cbfe --- /dev/null +++ b/vllm_omni/diffusion/layers/mxfp8.py @@ -0,0 +1,327 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 producers: e4m3 payload plus one E8M0 scale per 32 elements along K, +the scales in the cuBLASLt ``SWIZZLE_32_4_4`` layout that +``torch.nn.functional.scaled_mm(..., BlockWise1x32)`` consumes on Blackwell GPUs. + +Scale ``(r, c)`` of the ``[rows, K/32]`` scale matrix lives at byte +``((r // 128) * ceil(K/32 / 4) + c // 4) * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + c % 4``, +rows padded to 128 and scale columns to 4, padding zero. Exponent +``e = ceil(log2(amax / 448))`` exactly from the float bits; ``q = e4m3(x * 2**-e)``; +scale byte ``e + 127``. Every producer quantizes the bf16-rounded value the +unfused bf16 kernel stores, so each is byte-exact against that kernel followed +by ``mxfp8_quantize_swizzled`` (itself byte-exact vs ``flashinfer.mxfp8_quantize``). +""" + +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from __future__ import annotations + +import torch +from vllm.triton_utils import tl, triton +from vllm.triton_utils import tldevice as libdevice + +_E4M3 = torch.float8_e4m3fn + + +def _scale_numel(rows: int, k: int) -> int: + n_groups = k // 32 + return -(-rows // 128) * 128 * (-(-n_groups // 4) * 4) + + +@triton.jit +def _mx_e8m0_from_amax(amax): + bits = amax.to(tl.int32, bitcast=True) + e0 = ((bits >> 23) & 0xFF) - 135 + # amax / 2**e0 lies in [256, 512); bump e0 when it exceeds 448 = 1.75 * 2**8 + the = (((e0 + 135) << 23) | 0x600000).to(tl.float32, bitcast=True) + e = e0 + (amax > the).to(tl.int32) + e = tl.maximum(e, -127) + inv = ((127 - e) << 23).to(tl.float32, bitcast=True) + return e + 127, inv + + +@triton.jit +def _mx_scale_offsets(r, c, n_col_blocks): + tile = (r // 128) * n_col_blocks + (c // 4) + return tile * 512 + (r % 32) * 16 + ((r % 128) // 32) * 4 + (c % 4) + + +@triton.jit +def _mxfp8_quant_kernel( + x_ptr, + q_ptr, + s_ptr, + rows, + k, + n_groups, + n_col_blocks, + stride_x, + BLOCK_R: tl.constexpr, # noqa: N803 + G: tl.constexpr, # noqa: N803 +): + pid_r = tl.program_id(0) + pid_g = tl.program_id(1) + r = pid_r * BLOCK_R + tl.arange(0, BLOCK_R) + g = pid_g * G + tl.arange(0, G) + c = pid_g * (G * 32) + tl.arange(0, G * 32) + rmask = r < rows + mask = rmask[:, None] & (c < k)[None, :] + x = tl.load(x_ptr + r[:, None].to(tl.int64) * stride_x + c[None, :], mask=mask, other=0.0).to(tl.float32) + x3 = tl.reshape(x, [BLOCK_R, G, 32]) + amax = tl.max(tl.abs(x3), axis=2) + sbyte, inv = _mx_e8m0_from_amax(amax) + q = tl.reshape(x3 * inv[:, :, None], [BLOCK_R, G * 32]) + tl.store( + q_ptr + r[:, None].to(tl.int64) * k + c[None, :], + q.to(tl.float8e4nv), + mask=mask, + ) + smask = rmask[:, None] & (g < n_groups)[None, :] + tl.store( + s_ptr + _mx_scale_offsets(r[:, None], g[None, :], n_col_blocks), + sbyte.to(tl.uint8), + mask=smask, + ) + + +def can_use_mxfp8_swizzled(x: torch.Tensor) -> bool: + """Row-major bf16 CUDA 2D tensor with K % 32 == 0, outside torch.compile.""" + return ( + x.is_cuda + and x.ndim == 2 + and x.dtype == torch.bfloat16 + and x.stride(-1) == 1 + and x.shape[-1] % 32 == 0 + and not torch.compiler.is_compiling() + ) + + +def _alloc(rows: int, k: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]: + q = torch.empty(rows, k, dtype=_E4M3, device=device) + s = torch.zeros(_scale_numel(rows, k), dtype=torch.uint8, device=device) + return q, s + + +def mxfp8_quantize_swizzled(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """bf16 ``[rows, k]`` -> ``(fp8 [rows, k], swizzled e8m0 scale bytes)``.""" + if not can_use_mxfp8_swizzled(x): + raise ValueError("expected a row-major bf16 CUDA [rows, k] tensor, k % 32 == 0") + rows, k = x.shape + q, s = _alloc(rows, k, x.device) + if rows == 0: + return q, s + n_groups = k // 32 + n_col_blocks = -(-n_groups // 4) + block_r, g = 32, 8 + grid = (triton.cdiv(rows, block_r), triton.cdiv(n_groups, g)) + with torch.get_device_module().device(x.device): + _mxfp8_quant_kernel[grid]( + x, + q, + s, + rows, + k, + n_groups, + n_col_blocks, + x.stride(0), + BLOCK_R=block_r, + G=g, + num_warps=4, + ) + return q, s + + +@triton.jit +def _bf16_round(x): + # Explicit conversion pair preserves the CUDA kernel's two BF16 roundings. + # Unlike the integer rounding shortcut this also preserves NaN conversion. + return tl.inline_asm_elementwise( + "{ .reg .b16 h; cvt.rn.bf16.f32 h, $1; cvt.f32.bf16 $0, h; }", + constraints="=f,f", + args=[x], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _cuda_silu_product(gate, up): + # The installed vLLM SM120 SASS uses expf range reduction and corrected + # division gate / (1 + expf(-gate)), then BF16 round, FMUL, BF16 round. + # tl.sigmoid uses a different approximation: never assume byte identity. + act = _bf16_round(libdevice.div_rn(gate, 1.0 + libdevice.exp(-gate))) + # Current generic CUDA act_and_mul keeps its (up + beta), beta=+0, add. + # Preserve its signed-zero behavior too; an ordinary +0 can be optimized out. + up_zero = tl.inline_asm_elementwise( + "add.rn.f32 $0, $1, 0f00000000;", + constraints="=f,f", + args=[up], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + return _bf16_round(libdevice.mul_rn(act, up_zero)) + + +@triton.jit +def _silu_mxfp8_cuda_kernel( + X, # noqa: N803 + Q, # noqa: N803 + S, # noqa: N803 + Y, # noqa: N803 + rows, + hidden, + groups, + col_blocks, + stride_row, + STORE_BF16: tl.constexpr, # noqa: N803 + BLOCK_R: tl.constexpr, # noqa: N803 + G: tl.constexpr, # noqa: N803 +): + r = tl.program_id(0) * BLOCK_R + tl.arange(0, BLOCK_R) + g = tl.program_id(1) * G + tl.arange(0, G) + c = tl.program_id(1) * (G * 32) + tl.arange(0, G * 32) + mask = (r < rows)[:, None] & (c < hidden)[None, :] + base = X + r[:, None].to(tl.int64) * stride_row + gate = tl.load(base + c[None, :], mask, other=0).to(tl.float32) + up = tl.load(base + hidden + c[None, :], mask, other=0).to(tl.float32) + prod = _cuda_silu_product(gate, up) + if STORE_BF16: + tl.store(Y + r[:, None].to(tl.int64) * hidden + c[None, :], prod, mask) + p3 = tl.reshape(prod, [BLOCK_R, G, 32]) + amax = tl.max(tl.abs(p3), axis=2) + scale_byte, inv = _mx_e8m0_from_amax(amax) + quant = tl.reshape(p3 * inv[:, :, None], [BLOCK_R, G * 32]) + tl.store(Q + r[:, None].to(tl.int64) * hidden + c[None, :], quant.to(tl.float8e4nv), mask) + tl.store( + S + _mx_scale_offsets(r[:, None], g[None, :], col_blocks), + scale_byte.to(tl.uint8), + (r < rows)[:, None] & (g < groups)[None, :], + ) + + +def silu_mul_mxfp8_cuda(x: torch.Tensor, *, return_bf16: bool = False): + """BF16 [M, 2K] gate|up -> E4M3 [M,K], padded swizzled E8M0 bytes. + + return_bf16 additionally returns the rounded activation for validation. + """ + if ( + not x.is_cuda + or x.dtype != torch.bfloat16 + or x.ndim != 2 + or x.stride(1) != 1 + or x.shape[1] % 64 + or torch.compiler.is_compiling() + ): + raise ValueError("Expected resident row-major BF16 [M,2K], K multiple of 32, outside compile") + rows, twice = x.shape + hidden = twice // 2 + q, s = _alloc(rows, hidden, x.device) + y = torch.empty((rows, hidden), dtype=torch.bfloat16, device=x.device) if return_bf16 else q + if rows: + groups = hidden // 32 + with torch.get_device_module().device(x.device): + _silu_mxfp8_cuda_kernel[(triton.cdiv(rows, 16), triton.cdiv(groups, 8))]( + x, + q, + s, + y, + rows, + hidden, + groups, + triton.cdiv(groups, 4), + x.stride(0), + STORE_BF16=return_bf16, + BLOCK_R=16, + G=8, + num_warps=4, + enable_fp_fusion=True, + # vLLM's installed CUDA SASS uses non-FTZ exp/div arithmetic. + # Triton's libdevice reflection otherwise defaults to FTZ. + enable_reflect_ftz=False, + ) + return (q, s, y) if return_bf16 else (q, s) + + +def mxfp8_scaled_mm( + quantized: torch.Tensor, + weight: torch.Tensor, + activation_scale: torch.Tensor, + weight_scale: torch.Tensor, + *, + output_dtype: torch.dtype = torch.bfloat16, + bias: torch.Tensor | None = None, + use_fast_accum: bool = False, +) -> torch.Tensor: + """Multiply E4M3 operands with swizzled E8M0 scales, one scale per 32 values. + + Weight is stored as [N,K]. Keep this import lazy: the block-scaled public + PyTorch API is required only by callers explicitly selecting MXFP8. + """ + from torch.nn.functional import ScalingType, SwizzleType, scaled_mm + + return scaled_mm( + quantized, + weight.t(), + scale_a=activation_scale.view(torch.float8_e8m0fnu), + scale_b=weight_scale.view(torch.float8_e8m0fnu), + scale_recipe_a=ScalingType.BlockWise1x32, + scale_recipe_b=ScalingType.BlockWise1x32, + swizzle_a=SwizzleType.SWIZZLE_32_4_4, + swizzle_b=SwizzleType.SWIZZLE_32_4_4, + bias=bias, + output_dtype=output_dtype, + use_fast_accum=use_fast_accum, + ) + + +@torch.library.custom_op("vllm_omni::mxfp8_linear", mutates_args=()) +def mxfp8_linear(x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Opaque quantize/GEMM boundary for regional compilation.""" + q, activation_scale = mxfp8_quantize_swizzled(x.reshape(-1, x.shape[-1]).contiguous()) + out = mxfp8_scaled_mm(q, weight, activation_scale, scale) + return out.reshape(*x.shape[:-1], weight.shape[0]) + + +@mxfp8_linear.register_fake +def _mxfp8_linear_fake(x, weight, scale): + return x.new_empty((*x.shape[:-1], weight.shape[0])) + + +@torch.library.custom_op("vllm_omni::silu_mxfp8_linear", mutates_args=()) +def silu_mxfp8_linear(x: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Fuse CUDA-compatible BF16 SwiGLU rounding and MXFP8 quantization.""" + q, activation_scale = silu_mul_mxfp8_cuda(x.reshape(-1, x.shape[-1]).contiguous()) + out = mxfp8_scaled_mm(q, weight, activation_scale, scale) + return out.reshape(*x.shape[:-1], weight.shape[0]) + + +@silu_mxfp8_linear.register_fake +def _silu_mxfp8_linear_fake(x, weight, scale): + return x.new_empty((*x.shape[:-1], weight.shape[0])) + + +@torch.library.custom_op("vllm_omni::mxfp8_quantize_project", mutates_args=()) +def mxfp8_quantize_project( + x: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Project one weight slice and expose quantized X for another consumer. + + Returning both payload and scale makes the data dependency explicit across + compile boundaries. No tensor is retained in a process-global registry. + """ + q, activation_scale = mxfp8_quantize_swizzled(x) + out = mxfp8_scaled_mm(q, weight, activation_scale, scale) + return out, q, activation_scale + + +@mxfp8_quantize_project.register_fake +def _mxfp8_quantize_project_fake(x, weight, scale): + return ( + x.new_empty((x.shape[0], weight.shape[0])), + x.new_empty(x.shape, dtype=torch.float8_e4m3fn), + x.new_empty((_scale_numel(*x.shape),), dtype=torch.uint8), + ) diff --git a/vllm_omni/diffusion/models/minimax_h3/__init__.py b/vllm_omni/diffusion/models/minimax_h3/__init__.py index 02937c05881..723fa04b929 100644 --- a/vllm_omni/diffusion/models/minimax_h3/__init__.py +++ b/vllm_omni/diffusion/models/minimax_h3/__init__.py @@ -1,15 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniMax H3 audio-video diffusion support.""" -from vllm_omni.model_executor.models.minimax_h3 import reference_video +__all__ = ["MiniMaxH3Pipeline", "get_minimax_h3_post_process_func", "reference_video"] -from .pipeline_minimax_h3 import ( - MiniMaxH3Pipeline, - get_minimax_h3_post_process_func, -) -__all__ = [ - "MiniMaxH3Pipeline", - "get_minimax_h3_post_process_func", - "reference_video", -] +def __getattr__(name: str): + if name == "reference_video": + from vllm_omni.model_executor.models.minimax_h3 import reference_video + + return reference_video + if name in ("MiniMaxH3Pipeline", "get_minimax_h3_post_process_func"): + from . import pipeline_minimax_h3 + + return getattr(pipeline_minimax_h3, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm_omni/diffusion/models/minimax_h3/adaln_cache.py b/vllm_omni/diffusion/models/minimax_h3/adaln_cache.py new file mode 100644 index 00000000000..302e822c8f7 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/adaln_cache.py @@ -0,0 +1,710 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Precomputed AdaLN plan cache for the MiniMax H3 DiT. + +Version 2 is the SGLang-compatible base-H3 format. Version 3 is deliberately a +FastH3-only format: its outputs were computed after fusing one exact FastH3 +adapter, so the sidecar metadata binds that adapter and its sampling contract. + +The expensive adapter digest belongs in the runner/coordinator. Workers receive +the resulting :class:`MiniMaxH3AdalnCacheBinding`, bind it to the cache, and +validate only the small sidecar metadata. That avoids eight workers hashing the +same multi-GiB adapter while still failing closed before its AdaLN weights are +elided. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import os +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn as nn +from safetensors.torch import safe_open + +if TYPE_CHECKING: + from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import ( + MiniMaxH3DiTArchConfig, + ) + +_BF16_DTYPE = torch.bfloat16 +_FP32_DTYPE = torch.float32 +_ADALN_MODALITY_NUM = 3 +_SUPPORTED_FORMAT_VERSIONS = frozenset({"2", "3"}) +_V3_FORMAT_VERSION = "3" +_SHA256_HEX_LENGTH = 64 +_FASTH3_FORMAT = "fastvideo-lora-v2" +_FASTH3_BASE_MODEL = "MiniMaxAI/MiniMax-H3" +_FASTH3_ADAPTER_KEYS = { + "sha256": "fasth3_adapter_sha256", + "size_bytes": "fasth3_adapter_size_bytes", + "header_sha256": "fasth3_adapter_header_sha256", + "header_size_bytes": "fasth3_adapter_header_size_bytes", + "header_identity": "fasth3_adapter_header_identity", +} +_V3_CONTRACT_KEYS = { + "base_schedule": "base_schedule", + "mode": "mode", + "flow_shift": "flow_shift", + "audio_flow_shift": "audio_flow_shift", +} + + +def _validate_sha256(value: str, *, field: str) -> str: + normalized = str(value).lower() + if len(normalized) != _SHA256_HEX_LENGTH: + raise ValueError(f"{field} must be a 64-character SHA256 digest") + try: + int(normalized, 16) + except ValueError as exc: + raise ValueError(f"{field} must be a hexadecimal SHA256 digest") from exc + return normalized + + +def _canonical_json(value: Any) -> str: + return json.dumps( + value, + allow_nan=False, + ensure_ascii=True, + separators=(",", ":"), + sort_keys=True, + ) + + +def _parse_header_identity(value: str) -> dict[str, str]: + try: + parsed = json.loads(value) + except (TypeError, ValueError) as exc: + raise ValueError("fasth3_adapter_header_identity must be canonical JSON") from exc + if not isinstance(parsed, dict) or any( + not isinstance(key, str) or not isinstance(item, str) for key, item in parsed.items() + ): + raise ValueError("fasth3_adapter_header_identity must encode string metadata") + canonical = _canonical_json(parsed) + if canonical != value: + raise ValueError("fasth3_adapter_header_identity must use canonical JSON encoding") + return parsed + + +def _validate_fasth3_header_identity(value: str) -> str: + metadata = _parse_header_identity(value) + if metadata.get("format") != _FASTH3_FORMAT: + raise ValueError(f"FastH3 adapter header identity must declare format={_FASTH3_FORMAT!r}") + finetuned_model = metadata.get("finetuned_model", "").casefold() + if not finetuned_model.startswith("fastvideo/") or "fasth3" not in finetuned_model: + raise ValueError("FastH3 adapter header identity must name a FastVideo FastH3 release") + base_model = metadata.get("base_model") + if base_model is not None and base_model.casefold() != _FASTH3_BASE_MODEL.casefold(): + raise ValueError("FastH3 adapter header identity targets an unexpected base model") + return value + + +@dataclass(frozen=True) +class MiniMaxH3FastH3AdapterIdentity: + """Content and header identity of one FastH3 adapter artifact.""" + + sha256: str + size_bytes: int + header_sha256: str + header_size_bytes: int + header_identity: str + + def __post_init__(self) -> None: + object.__setattr__( + self, + "sha256", + _validate_sha256(self.sha256, field="fasth3_adapter_sha256"), + ) + object.__setattr__( + self, + "header_sha256", + _validate_sha256( + self.header_sha256, + field="fasth3_adapter_header_sha256", + ), + ) + if int(self.size_bytes) <= 8: + raise ValueError("fasth3_adapter_size_bytes must exceed the header prefix") + if int(self.header_size_bytes) <= 0: + raise ValueError("fasth3_adapter_header_size_bytes must be positive") + if int(self.header_size_bytes) + 8 > int(self.size_bytes): + raise ValueError("FastH3 adapter header extends beyond the artifact") + object.__setattr__(self, "size_bytes", int(self.size_bytes)) + object.__setattr__(self, "header_size_bytes", int(self.header_size_bytes)) + object.__setattr__( + self, + "header_identity", + _validate_fasth3_header_identity(self.header_identity), + ) + + def to_dict(self) -> dict[str, str | int]: + return { + "sha256": self.sha256, + "size_bytes": self.size_bytes, + "header_sha256": self.header_sha256, + "header_size_bytes": self.header_size_bytes, + "header_identity": self.header_identity, + } + + @classmethod + def from_dict( + cls, + value: Mapping[str, Any], + ) -> MiniMaxH3FastH3AdapterIdentity: + expected = set(_FASTH3_ADAPTER_KEYS) + unknown = set(value) - expected + missing = expected - set(value) + if missing or unknown: + raise ValueError( + f"invalid FastH3 adapter identity fields: missing={sorted(missing)}, unknown={sorted(unknown)}" + ) + return cls( + sha256=str(value["sha256"]), + size_bytes=int(value["size_bytes"]), + header_sha256=str(value["header_sha256"]), + header_size_bytes=int(value["header_size_bytes"]), + header_identity=str(value["header_identity"]), + ) + + +@dataclass(frozen=True) +class MiniMaxH3AdalnCacheBinding: + """Serializable identity and sampling contract for one AdaLN sidecar.""" + + format_version: str + model_variant: str | None + mode: str | None = None + base_schedule: tuple[float, ...] | None = None + flow_shift: float | None = None + audio_flow_shift: float | None = None + fasth3_adapter: MiniMaxH3FastH3AdapterIdentity | None = None + + def __post_init__(self) -> None: + version = str(self.format_version) + object.__setattr__(self, "format_version", version) + if version not in _SUPPORTED_FORMAT_VERSIONS: + raise ValueError(f"unsupported MiniMax H3 AdaLN cache format_version {version!r}") + if self.model_variant is not None and self.model_variant not in { + "fl2va", + "ref2va", + }: + raise ValueError(f"invalid MiniMax H3 AdaLN cache model_variant {self.model_variant!r}") + + v3_values = ( + self.mode, + self.base_schedule, + self.flow_shift, + self.audio_flow_shift, + self.fasth3_adapter, + ) + if version == "2": + if any(value is not None for value in v3_values): + raise ValueError("AdaLN cache format v2 cannot carry a FastH3 binding") + return + + if any(value is None for value in v3_values): + raise ValueError("AdaLN cache format v3 requires a complete FastH3 binding") + if self.model_variant != "fl2va": + raise ValueError("FastH3 AdaLN cache v3 requires model_variant='fl2va'") + if self.mode != "t2va": + raise ValueError("FastH3 AdaLN cache v3 requires mode='t2va'") + + assert self.base_schedule is not None + from vllm_omni.diffusion.sched.sigma_schedule import DMD2SigmaSchedule + + schedule = DMD2SigmaSchedule.from_positions(self.base_schedule) + object.__setattr__(self, "base_schedule", schedule.base_schedule) + for field, value in ( + ("flow_shift", self.flow_shift), + ("audio_flow_shift", self.audio_flow_shift), + ): + assert value is not None + normalized = float(value) + if not math.isfinite(normalized) or normalized <= 0: + raise ValueError(f"{field} must be finite and positive") + object.__setattr__(self, field, normalized) + + @classmethod + def for_fasth3( + cls, + *, + adapter: MiniMaxH3FastH3AdapterIdentity, + model_variant: str, + mode: str, + base_schedule: Sequence[float], + flow_shift: float, + audio_flow_shift: float, + ) -> MiniMaxH3AdalnCacheBinding: + return cls( + format_version=_V3_FORMAT_VERSION, + model_variant=model_variant, + mode=mode, + base_schedule=tuple(float(value) for value in base_schedule), + flow_shift=float(flow_shift), + audio_flow_shift=float(audio_flow_shift), + fasth3_adapter=adapter, + ) + + def to_dict(self) -> dict[str, Any]: + return { + "format_version": self.format_version, + "model_variant": self.model_variant, + "mode": self.mode, + "base_schedule": (list(self.base_schedule) if self.base_schedule is not None else None), + "flow_shift": self.flow_shift, + "audio_flow_shift": self.audio_flow_shift, + "fasth3_adapter": (self.fasth3_adapter.to_dict() if self.fasth3_adapter is not None else None), + } + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> MiniMaxH3AdalnCacheBinding: + expected = { + "format_version", + "model_variant", + "mode", + "base_schedule", + "flow_shift", + "audio_flow_shift", + "fasth3_adapter", + } + unknown = set(value) - expected + missing = expected - set(value) + if missing or unknown: + raise ValueError( + f"invalid MiniMax H3 AdaLN binding fields: missing={sorted(missing)}, unknown={sorted(unknown)}" + ) + raw_adapter = value["fasth3_adapter"] + adapter = MiniMaxH3FastH3AdapterIdentity.from_dict(raw_adapter) if isinstance(raw_adapter, Mapping) else None + raw_schedule = value["base_schedule"] + return cls( + format_version=str(value["format_version"]), + model_variant=(str(value["model_variant"]) if value["model_variant"] is not None else None), + mode=str(value["mode"]) if value["mode"] is not None else None, + base_schedule=(tuple(float(item) for item in raw_schedule) if raw_schedule is not None else None), + flow_shift=(float(value["flow_shift"]) if value["flow_shift"] is not None else None), + audio_flow_shift=(float(value["audio_flow_shift"]) if value["audio_flow_shift"] is not None else None), + fasth3_adapter=adapter, + ) + + def to_metadata(self) -> dict[str, str]: + metadata = { + "format_version": self.format_version, + "model_variant": self.model_variant or "", + } + if self.format_version == "2": + return metadata + assert self.base_schedule is not None + assert self.mode is not None + assert self.flow_shift is not None + assert self.audio_flow_shift is not None + assert self.fasth3_adapter is not None + metadata.update( + { + _V3_CONTRACT_KEYS["mode"]: self.mode, + _V3_CONTRACT_KEYS["base_schedule"]: _canonical_json(list(self.base_schedule)), + _V3_CONTRACT_KEYS["flow_shift"]: repr(self.flow_shift), + _V3_CONTRACT_KEYS["audio_flow_shift"]: repr(self.audio_flow_shift), + _FASTH3_ADAPTER_KEYS["sha256"]: self.fasth3_adapter.sha256, + _FASTH3_ADAPTER_KEYS["size_bytes"]: str(self.fasth3_adapter.size_bytes), + _FASTH3_ADAPTER_KEYS["header_sha256"]: self.fasth3_adapter.header_sha256, + _FASTH3_ADAPTER_KEYS["header_size_bytes"]: str(self.fasth3_adapter.header_size_bytes), + _FASTH3_ADAPTER_KEYS["header_identity"]: self.fasth3_adapter.header_identity, + } + ) + return metadata + + @classmethod + def from_metadata( + cls, + metadata: Mapping[str, str], + ) -> MiniMaxH3AdalnCacheBinding: + version = metadata.get("format_version", "") + model_variant = metadata.get("model_variant") + if version == "2": + return cls(format_version=version, model_variant=model_variant) + if version != _V3_FORMAT_VERSION: + raise ValueError("MiniMax H3 AdaLN cache has an unsupported or missing format_version") + required = set(_V3_CONTRACT_KEYS.values()) | set(_FASTH3_ADAPTER_KEYS.values()) + missing = sorted(key for key in required if key not in metadata) + if missing: + raise ValueError(f"MiniMax H3 AdaLN cache v3 metadata is incomplete: {missing}") + try: + raw_schedule = json.loads(metadata[_V3_CONTRACT_KEYS["base_schedule"]]) + if not isinstance(raw_schedule, list): + raise TypeError("base_schedule must be a JSON list") + adapter = MiniMaxH3FastH3AdapterIdentity( + sha256=metadata[_FASTH3_ADAPTER_KEYS["sha256"]], + size_bytes=int(metadata[_FASTH3_ADAPTER_KEYS["size_bytes"]]), + header_sha256=metadata[_FASTH3_ADAPTER_KEYS["header_sha256"]], + header_size_bytes=int(metadata[_FASTH3_ADAPTER_KEYS["header_size_bytes"]]), + header_identity=metadata[_FASTH3_ADAPTER_KEYS["header_identity"]], + ) + return cls.for_fasth3( + adapter=adapter, + model_variant=str(model_variant), + mode=metadata[_V3_CONTRACT_KEYS["mode"]], + base_schedule=raw_schedule, + flow_shift=float(metadata[_V3_CONTRACT_KEYS["flow_shift"]]), + audio_flow_shift=float(metadata[_V3_CONTRACT_KEYS["audio_flow_shift"]]), + ) + except (TypeError, ValueError) as exc: + raise ValueError(f"MiniMax H3 AdaLN cache v3 metadata is malformed: {exc}") from exc + + +def _resolve_fasth3_adapter_file(path: str | os.PathLike[str]) -> Path: + candidate = Path(path).expanduser() + if candidate.is_file(): + if candidate.suffix != ".safetensors": + raise ValueError(f"FastH3 adapter is not a safetensors file: {candidate}") + return candidate.resolve() + if not candidate.is_dir(): + raise ValueError(f"FastH3 adapter does not exist: {candidate}") + named = candidate / "adapter_model.safetensors" + if named.is_file(): + return named.resolve() + files = sorted(candidate.glob("*.safetensors")) + if len(files) != 1: + raise ValueError(f"FastH3 adapter directory must contain one safetensors file: {candidate}") + return files[0].resolve() + + +def fingerprint_minimax_h3_fasth3_adapter( + path: str | os.PathLike[str], +) -> MiniMaxH3FastH3AdapterIdentity: + """Hash one FastH3 adapter once, for runner-side prevalidation. + + This intentionally reads the full payload. Do not call it independently in + every diffusion worker; serialize the returned identity instead. + """ + adapter_path = _resolve_fasth3_adapter_file(path) + digest = hashlib.sha256() + with adapter_path.open("rb") as adapter_file: + before = os.fstat(adapter_file.fileno()) + prefix = adapter_file.read(8) + if len(prefix) != 8: + raise ValueError(f"FastH3 adapter has a truncated header: {adapter_path}") + header_size = int.from_bytes(prefix, byteorder="little", signed=False) + if header_size <= 0 or header_size + 8 > before.st_size: + raise ValueError(f"FastH3 adapter has an invalid header size: {adapter_path}") + header = adapter_file.read(header_size) + if len(header) != header_size: + raise ValueError(f"FastH3 adapter has a truncated header: {adapter_path}") + try: + parsed_header = json.loads(header) + except (TypeError, ValueError) as exc: + raise ValueError(f"FastH3 adapter has malformed safetensors JSON: {adapter_path}") from exc + if not isinstance(parsed_header, dict): + raise ValueError(f"FastH3 adapter safetensors header is not an object: {adapter_path}") + metadata = parsed_header.get("__metadata__") or {} + if not isinstance(metadata, dict) or any( + not isinstance(key, str) or not isinstance(value, str) for key, value in metadata.items() + ): + raise ValueError(f"FastH3 adapter has invalid safetensors metadata: {adapter_path}") + header_identity = _canonical_json(metadata) + _validate_fasth3_header_identity(header_identity) + + adapter_file.seek(0) + while chunk := adapter_file.read(16 * 1024 * 1024): + digest.update(chunk) + after = os.fstat(adapter_file.fileno()) + identity_before = ( + before.st_dev, + before.st_ino, + before.st_size, + before.st_mtime_ns, + ) + identity_after = ( + after.st_dev, + after.st_ino, + after.st_size, + after.st_mtime_ns, + ) + if identity_before != identity_after: + raise ValueError(f"FastH3 adapter changed while it was hashed: {adapter_path}") + return MiniMaxH3FastH3AdapterIdentity( + sha256=digest.hexdigest(), + size_bytes=before.st_size, + header_sha256=hashlib.sha256(header).hexdigest(), + header_size_bytes=header_size, + header_identity=header_identity, + ) + + +def read_minimax_h3_adaln_cache_binding( + path: str | os.PathLike[str], +) -> MiniMaxH3AdalnCacheBinding: + """Read and validate only the sidecar's metadata binding.""" + cache_path = os.path.abspath(os.path.expanduser(os.fspath(path))) + if not os.path.isfile(cache_path): + raise ValueError(f"MiniMax H3 AdaLN cache does not exist: {cache_path}") + with safe_open(cache_path, framework="pt", device="cpu") as cache_file: + return MiniMaxH3AdalnCacheBinding.from_metadata(cache_file.metadata() or {}) + + +def _coerce_binding( + binding: MiniMaxH3AdalnCacheBinding | Mapping[str, Any], +) -> MiniMaxH3AdalnCacheBinding: + if isinstance(binding, MiniMaxH3AdalnCacheBinding): + return binding + if isinstance(binding, Mapping): + return MiniMaxH3AdalnCacheBinding.from_dict(binding) + raise TypeError("MiniMax H3 AdaLN runtime binding must be a binding or mapping") + + +def validate_minimax_h3_adaln_sidecar_binding( + path: str | os.PathLike[str], + binding: MiniMaxH3AdalnCacheBinding | Mapping[str, Any], +) -> MiniMaxH3AdalnCacheBinding: + """Require ``path`` metadata to equal a runner-prevalidated binding.""" + expected = _coerce_binding(binding) + actual = read_minimax_h3_adaln_cache_binding(path) + if actual != expected: + raise ValueError( + f"MiniMax H3 AdaLN cache binding mismatch: expected={expected.to_dict()!r}, actual={actual.to_dict()!r}" + ) + return actual + + +def minimax_h3_adaln_sidecar_satisfies( + path: str | os.PathLike[str], + binding: MiniMaxH3AdalnCacheBinding | Mapping[str, Any], +) -> bool: + """Cheap metadata-only predicate for a pipeline/core sidecar selection.""" + try: + validate_minimax_h3_adaln_sidecar_binding(path, binding) + except (OSError, TypeError, ValueError): + return False + return True + + +def prevalidate_minimax_h3_fasth3_adaln_sidecar( + sidecar_path: str | os.PathLike[str], + *, + adapter_path: str | os.PathLike[str], + model_variant: str, + mode: str, + base_schedule: Sequence[float], + flow_shift: float, + audio_flow_shift: float, +) -> MiniMaxH3AdalnCacheBinding: + """Runner-only full adapter verification followed by sidecar validation.""" + adapter = fingerprint_minimax_h3_fasth3_adapter(adapter_path) + expected = MiniMaxH3AdalnCacheBinding.for_fasth3( + adapter=adapter, + model_variant=model_variant, + mode=mode, + base_schedule=base_schedule, + flow_shift=flow_shift, + audio_flow_shift=audio_flow_shift, + ) + return validate_minimax_h3_adaln_sidecar_binding(sidecar_path, expected) + + +class MiniMaxH3AdalnCache(nn.Module): + """Fixed-schedule AdaLN output sidecar. + + Base-H3 v2 files keep their original behavior. A FastH3 v3 file must be + bound before :meth:`load`; otherwise the runtime cannot know that the + fused student, schedule and shifts match the precomputed outputs. + """ + + plan_timesteps: torch.Tensor + plan_lengths: torch.Tensor + block_params: torch.Tensor + final_params: torch.Tensor + + def __init__( + self, + arch: MiniMaxH3DiTArchConfig, + *, + path: str, + model_variant: str | None = None, + runtime_binding: MiniMaxH3AdalnCacheBinding | Mapping[str, Any] | None = None, + ) -> None: + super().__init__() + if not path: + raise ValueError("MiniMax H3 AdaLN cache path must not be empty") + self.path = os.path.abspath(os.path.expanduser(path)) + self.model_variant = model_variant + self.num_layers = arch.num_layers + self.hidden_size = arch.hidden_size + self.block_width = 6 * _ADALN_MODALITY_NUM * arch.hidden_size + self.final_width = 2 * arch.hidden_size + self._loaded = False + self._runtime_binding: MiniMaxH3AdalnCacheBinding | None = None + self._sidecar_binding: MiniMaxH3AdalnCacheBinding | None = None + self._binding_validated = False + if runtime_binding is not None: + self.bind_runtime_contract(runtime_binding) + + def bind_runtime_contract( + self, + binding: MiniMaxH3AdalnCacheBinding | Mapping[str, Any], + ) -> None: + """Bind runner-prevalidated identity without rehashing the adapter. + + The runner calls :func:`prevalidate_minimax_h3_fasth3_adaln_sidecar` + once, serializes ``binding.to_dict()``, and gives the mapping to every + worker. This method reads only the sidecar header. + """ + if self._loaded: + raise RuntimeError("MiniMax H3 AdaLN runtime contract must be bound before load") + expected = _coerce_binding(binding) + if self.model_variant is not None and expected.model_variant != self.model_variant: + raise ValueError( + "MiniMax H3 AdaLN runtime binding model_variant does not match " + f"the loaded variant ({expected.model_variant!r} != " + f"{self.model_variant!r})" + ) + actual = validate_minimax_h3_adaln_sidecar_binding(self.path, expected) + if self._runtime_binding is not None and self._runtime_binding != expected: + raise ValueError("MiniMax H3 AdaLN cache is already bound to another runtime contract") + self._runtime_binding = expected + self._sidecar_binding = actual + self._binding_validated = True + + def sidecar_satisfied_parameter_names(self) -> frozenset[str]: + """Return FastH3 patches replaced by a verified v3 sidecar. + + These are the parameters a load-time FastH3 fusion may mark satisfied + instead of reconstructing. ``time_embedder`` is intentionally absent: + the current transformer still executes it in ``_embed`` and therefore + still needs its four FastH3 dense patches loaded. + """ + binding = self._sidecar_binding + if not self._binding_validated or binding is None or binding.format_version != _V3_FORMAT_VERSION: + raise RuntimeError("FastH3 AdaLN sidecar parameters require a validated v3 binding") + names = { + f"blocks.{index}.adaln_proj.linear.{suffix}" + for index in range(self.num_layers) + for suffix in ("weight", "bias") + } + names.update( + { + "final_layer.adaln_proj.linear.weight", + "final_layer.adaln_proj.linear.bias", + } + ) + return frozenset(names) + + def load(self, device: torch.device) -> None: + """Validate the sidecar on CPU, then install its tensors on ``device``.""" + if self._loaded: + if self.block_params.device != device: + raise ValueError( + "MiniMax H3 AdaLN cache is already loaded on " + f"{self.block_params.device}, cannot reload it on {device}" + ) + return + if not os.path.isfile(self.path): + raise ValueError(f"MiniMax H3 AdaLN cache does not exist: {self.path}") + + with safe_open(self.path, framework="pt", device="cpu") as cache_file: + metadata = cache_file.metadata() or {} + sidecar_binding = MiniMaxH3AdalnCacheBinding.from_metadata(metadata) + if self.model_variant is not None and sidecar_binding.model_variant != self.model_variant: + raise ValueError( + "MiniMax H3 AdaLN cache model_variant does not match the " + "loaded variant " + f"({sidecar_binding.model_variant!r} != {self.model_variant!r})" + ) + if sidecar_binding.format_version == _V3_FORMAT_VERSION: + if self._runtime_binding is None or not self._binding_validated: + raise ValueError("MiniMax H3 AdaLN cache v3 requires bind_runtime_contract() before load") + if sidecar_binding != self._runtime_binding: + raise ValueError("MiniMax H3 AdaLN cache binding changed after runtime prevalidation") + elif self._runtime_binding is not None: + if sidecar_binding != self._runtime_binding: + raise ValueError("MiniMax H3 AdaLN cache binding does not match the runtime contract") + plan_timesteps = cache_file.get_tensor("plan_timesteps") + plan_lengths = cache_file.get_tensor("plan_lengths") + block_params = cache_file.get_tensor("block_params") + final_params = cache_file.get_tensor("final_params") + + num_plans = plan_timesteps.shape[0] if plan_timesteps.ndim == 2 else -1 + max_width = plan_timesteps.shape[1] if plan_timesteps.ndim == 2 else -1 + if ( + plan_timesteps.dtype != _FP32_DTYPE + or plan_timesteps.ndim != 2 + or plan_lengths.dtype != torch.int64 + or plan_lengths.shape != (num_plans,) + or bool((plan_lengths < 1).any()) + or bool((plan_lengths > max_width).any()) + ): + raise ValueError("MiniMax H3 AdaLN cache has invalid timestep plans") + if block_params.dtype != _BF16_DTYPE or block_params.shape != ( + num_plans, + max_width, + self.num_layers, + self.block_width, + ): + raise ValueError("MiniMax H3 AdaLN cache has invalid block_params") + if final_params.dtype != _BF16_DTYPE or final_params.shape != ( + num_plans, + max_width, + self.final_width, + ): + raise ValueError("MiniMax H3 AdaLN cache has invalid final_params") + + # These tensors are derived data and must not become checkpoint state. + self.register_buffer("plan_timesteps", plan_timesteps.to(device), persistent=False) + self.register_buffer("plan_lengths", plan_lengths.to(device), persistent=False) + self.register_buffer("block_params", block_params.to(device), persistent=False) + self.register_buffer("final_params", final_params.to(device), persistent=False) + self._sidecar_binding = sidecar_binding + self._loaded = True + + def lookup(self, unique_timesteps: torch.Tensor) -> torch.Tensor: + """Return the exact matching plan slot or fail closed.""" + if not self._loaded: + raise RuntimeError("MiniMax H3 AdaLN cache has not been loaded") + unique_timesteps = unique_timesteps.to( + device=self.plan_timesteps.device, + dtype=_FP32_DTYPE, + ).view(-1) + num_timesteps = unique_timesteps.shape[0] + if num_timesteps > self.plan_timesteps.shape[1]: + raise ValueError("MiniMax H3 AdaLN cache does not cover the request timestep plan") + matches = self.plan_lengths.eq(num_timesteps) & self.plan_timesteps[:, :num_timesteps].eq(unique_timesteps).all( + dim=-1 + ) + if not bool(matches.any()): + raise ValueError("MiniMax H3 AdaLN cache does not cover the request timestep plan") + return matches.to(torch.int64).argmax() + + def block_all( + self, + *, + cache_plan_index: torch.Tensor, + num_timesteps: int, + ) -> tuple[tuple[torch.Tensor, ...], ...]: + """Return all blocks' six AdaLN tensors with one layer-major gather.""" + stacked = self.block_params.permute(2, 0, 1, 3)[:, cache_plan_index, :num_timesteps] + stacked = stacked.reshape(self.num_layers, -1, 6, self.hidden_size) + return tuple(tuple(layer.unbind(dim=1)) for layer in stacked) + + def final( + self, + cache_plan_index: torch.Tensor, + num_timesteps: int, + ) -> tuple[torch.Tensor, ...]: + params = self.final_params[cache_plan_index, :num_timesteps] + return tuple(params.reshape(-1, 2, self.hidden_size).unbind(dim=1)) + + +__all__ = [ + "MiniMaxH3AdalnCache", + "MiniMaxH3AdalnCacheBinding", + "MiniMaxH3FastH3AdapterIdentity", + "fingerprint_minimax_h3_fasth3_adapter", + "minimax_h3_adaln_sidecar_satisfies", + "prevalidate_minimax_h3_fasth3_adaln_sidecar", + "read_minimax_h3_adaln_cache_binding", + "validate_minimax_h3_adaln_sidecar_binding", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/attention/__init__.py b/vllm_omni/diffusion/models/minimax_h3/attention/__init__.py new file mode 100644 index 00000000000..4244ce4b4e1 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/attention/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project diff --git a/vllm_omni/diffusion/models/minimax_h3/attention/overlap.py b/vllm_omni/diffusion/models/minimax_h3/attention/overlap.py new file mode 100644 index 00000000000..557dea20391 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/attention/overlap.py @@ -0,0 +1,503 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""H3 VSA execution scope: Q preparation, coarse attention, and reverse-output overlap. + +All producer tickets share one per-invocation lifetime. Layout and numerical +contracts remain VSA-specific; the Ulysses transport is provided by shared code. +""" + +from __future__ import annotations + +import os +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from functools import lru_cache +from typing import Any + +import torch +import torch.distributed as dist +from vllm.logger import init_logger + +logger = init_logger(__name__) +_BUNDLE_LOGGED = False +_GATE_LOGGED = False + +ACTIVE: ContextVar[dict[str, Any] | None] = ContextVar("minimax_h3_overlap", default=None) + + +@lru_cache(None) +def streams(device): + return torch.get_device_module().Stream(device=device, priority=-1), torch.get_device_module().Stream(device=device) + + +def run( + attention, + query, + key, + value, + metadata, + projector, + *, + gate_input=None, + gate_projector=None, + vsplit_input=None, + vsplit_projection=None, + vsplit_layer_index=None, +): + assert ACTIVE.get() is None, "overlap forward cannot be nested" + impl = attention.attention + assert impl.h3_kernel_backend == "flashinfer" and impl.topk == 162 + assert query.shape == key.shape == value.shape == (1, 11992, 56, 128) + assert query.dtype == key.dtype == value.dtype == torch.bfloat16 + caller = torch.get_device_module().current_stream() + comm, side = streams(query.device) + comm.wait_stream(caller) + state = dict(comm=comm, side=side, projector=projector, impl=impl, prepared=None, reverse=None, projected=False) + token = ACTIVE.set(state) + try: + with torch.get_device_module().stream(comm), torch.cuda.nvtx.range("h3.overlap.attention"): + if vsplit_input is not None: + from vllm_omni.diffusion.models.minimax_h3.attention.qkv_overlap import begin + + begin(state, attention, vsplit_projection, vsplit_input, value, vsplit_layer_index, metadata) + if gate_input is not None: + assert gate_input.shape == (11992, 5376) and gate_input.dtype == torch.bfloat16 + assert not gate_input.requires_grad and gate_input.is_contiguous() + assert gate_projector is not None and "gate_compress" not in metadata.extra + assert "vsa_h3_prefix_segments" in metadata.extra + # Q/K norm+RoPE have already been submitted to caller. Gate + # uses exactly the original GEMM on the same side stream as + # QK preparation and reverse gate/O. It has no dependency on + # received Q/K; preprocessing only inspects its shape/dtype. + side.wait_stream(comm) + ticket = state.get("vsplit_ticket") + if ticket is not None: + assert ticket.launched and ticket.done is not None + with torch.cuda.nvtx.range("h3.vsplit.gate_after_v_ready"): + side.wait_event(ticket.done) + with torch.get_device_module().stream(side), torch.cuda.nvtx.range("h3.lossless.gate_gemm"): + gate_input.record_stream(side) + gate_result = gate_projector(gate_input) + gate = gate_result[0] if isinstance(gate_result, tuple) else gate_result + assert gate.shape == (11992, 7168) and gate.dtype == torch.bfloat16 + metadata.extra["gate_compress"] = gate.view(1, 11992, 56, 128) + global _GATE_LOGGED + if not _GATE_LOGGED: + logger.info( + "H3_LOSSLESS_GATE rank=%d M=11992 N=7168 K=5376 dtype=bf16 " + "shape_unchanged=1 stream=overlap_side", + dist.get_rank(), + ) + _GATE_LOGGED = True + out = attention(query, key, value, metadata) + assert state["projected"], "chunked reverse/projector did not run" + if vsplit_input is not None: + from vllm_omni.diffusion.models.minimax_h3.attention.qkv_overlap import finish + + finish(state) + caller.wait_stream(comm) + out.record_stream(caller) + return out + finally: + # Also cover exceptions before the normal preparation/reverse joins. + if vsplit_input is not None: + from vllm_omni.diffusion.models.minimax_h3.attention.qkv_overlap import cleanup + + cleanup(state, caller) + + drain_output_producer(state, caller) + caller.wait_stream(side) + caller.wait_stream(comm) + ACTIVE.reset(token) + + +def before_v(query, key, value, metadata, group): + active = ACTIVE.get() + if active is None: + return value + if active.get("vsplit_ticket") is not None: + from vllm_omni.diffusion.models.minimax_h3.attention.qkv_overlap import before_v as split_before_v + + return split_before_v(active, query, key, value, metadata, group) + from vllm_omni.diffusion.distributed.flashinfer_ulysses import _state_for + from vllm_omni.diffusion.models.minimax_h3.attention.vsa import ( + _build_h3_ordered_q2k_indices, + _get_h3_layout, + _get_h3_tile_metadata, + _get_h3_tiled_source_rows, + _pool_h3_tiles, + h3_vsa_tile_pack, + ) + + state = _state_for(value, group, 8) + assert state.communicator is not None + out = state.output(value, op="scatter_heads", slot="v") + landing = state.communicator.input_buffer(out, value.shape) + landing.copy_(value) + layout = _get_h3_layout(metadata) + assert layout == ((558, 1206), (107, 22, 40), 1764), layout + prefix, shape, _ = layout + active["side"].wait_stream(active["comm"]) + with torch.get_device_module().stream(active["side"]), torch.cuda.nvtx.range("h3.overlap.qk_prepare"): + _, sizes, _, _, prefix_blocks, video_blocks = _get_h3_tile_metadata(prefix, shape, query.device) + maps = _get_h3_tiled_source_rows(prefix, shape, 105472, query.device) + q_ready = prepared_q() + q = h3_vsa_tile_pack(query[:, :95924], maps) if q_ready is None else q_ready[0] + k = h3_vsa_tile_pack(key[:, :95924], maps) + qp = _pool_h3_tiles(q, sizes) if q_ready is None else q_ready[1] + kp = _pool_h3_tiles(k, sizes) + scores = torch.matmul(qp, kp.transpose(-2, -1)) * active["impl"].softmax_scale + idx, num = _build_h3_ordered_q2k_indices(scores, prefix_blocks, video_blocks, 162) + active["prepared"] = (q, k, qp, kp, scores, idx, num) + return landing + + +def prepared(): + active = ACTIVE.get() + if active is None: + return None + result = active["prepared"] + assert result is not None, "forward QK preparation hook did not run" + active["comm"].wait_stream(active["side"]) + for tensor in result: + tensor.record_stream(active["comm"]) + return result + + +def publish_reverse(fine, coarse, plan): + active = ACTIVE.get() + if active is None: + return False + + join_coarse(coarse) + assert active["reverse"] is None + active["reverse"] = (fine, coarse, plan) + return True + + +def finish_reverse(attn_output, ctx): + active = ACTIVE.get() + if active is None: + return None + from vllm_omni.diffusion.distributed.flashinfer_ulysses import _state_for + from vllm_omni.diffusion.models.minimax_h3.ops.attention.chunks import gate_chunk + from vllm_omni.diffusion.models.minimax_h3.ops.attention.owner_route import h3_vsa_owner_route_plan_is_trusted + + fine, coarse, plan = active["reverse"] + assert fine is attn_output + assert set(ctx.o_bundle_state) == {"plan"} and ctx.o_bundle_state["plan"] is plan + assert ctx.strict_a2a_backend == "flashinfer-pcie" and ctx.joint_len == 0 and not ctx.use_uaa + assert h3_vsa_owner_route_plan_is_trusted(plan) + assert plan.local_rows == 11992 and plan.kmax == 270 and plan.sp_world_size == 8 + gate = ctx.o_bundle_gate_local + assert gate is not None and gate.shape == (1, 11992, 56, 128) + rank = dist.get_rank(ctx.ulysses_pg) + state = _state_for(fine, ctx.ulysses_pg, 8) + comm = state.communicator + assert comm is not None + # Cache these maps on the trusted, geometry-bound route's CPU tensors. + tiles, row_map = route_maps(plan, rank, fine.device) + count = 2998 + result = torch.empty(1, 11992, 5376, device=fine.device, dtype=torch.bfloat16) + + prepare_output_chunks(active, fine, coarse, tiles, state) + for index in range(4): + start = index * count + out, landing = consume_output_chunk(active, index) + with torch.cuda.nvtx.range(f"h3.overlap.o_rdma.{index}"): + comm.gather_heads(landing, out=out) + active["side"].wait_stream(active["comm"]) + with ( + torch.get_device_module().stream(active["side"]), + torch.cuda.nvtx.range(f"h3.overlap.o_gate_project.{index}"), + ): + view = gate_chunk(out, gate[:, start : start + count], row_map[start : start + count], count) + projected = active["projector"](view) + result[:, start : start + count].copy_(projected) + active["comm"].wait_stream(active["side"]) + result.record_stream(active["side"]) + gate.record_stream(active["side"]) + finish_output_chunks(active) + active["projected"] = True + # Separate markers bind the validator to the chunked protocol and its + # actual volume; they do not impersonate the original whole-bundle path. + if not state.o_producer_direct_logged: + logger.info( + "H3_OVERLAP_EXCHANGE rank=%d group=%s chunks=4 rows=2998 tail=270 " + "remote_bytes=163975168 direct=1 completed=1", + rank, + ctx.ulysses_pg.group_name, + ) + state.o_producer_direct_logged = True + state.o_producer_direct_expected = None + global _BUNDLE_LOGGED + if not _BUNDLE_LOGGED: + logger.info( + "H3_OVERLAP_BUNDLE rank=%d fine_rows=%d local_rows=11992 coarse_rows=1648 " + "chunks=4 tail=270 gate=bf16 projector=active", + rank, + fine.shape[1], + ) + _BUNDLE_LOGGED = True + return result + + +_ROUTE_MAPS = {} + + +def route_maps(plan, rank, device): + key = (id(plan), rank, device) + if key not in _ROUTE_MAPS: + rows = plan.local_row_to_tail[rank] + _ROUTE_MAPS[key] = (plan, plan.owner_tile_ids.to(device), rows.to(device)) + saved, tiles, row_map = _ROUTE_MAPS[key] + assert saved is plan + return tiles, row_map + + +_LOGGED: set[tuple[int, str]] = set() +_PREFIX = "VLLM_OMNI_H3_LOSSLESS_" + + +def _enabled(feature: str) -> bool: + return os.environ.get(_PREFIX + feature, "0") == "1" + + +def _log(feature: str) -> None: + rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 + key = (rank, feature) + if key not in _LOGGED: + logger.info("H3_SCHEDULE_V3_FEATURE rank=%d feature=%s enabled=1", rank, feature) + _LOGGED.add(key) + + +def layout_view_enabled() -> bool: + enabled = _enabled("LAYOUT_VIEW") + if enabled: + _log("layout_view") + return enabled + + +def after_q(query: torch.Tensor, metadata: Any) -> None: + active = ACTIVE.get() + if not _enabled("EARLY_Q") or active is None: + return + # The V-first candidate submits Q-only and K-dependent preparation + # together after K exchange, leaving K transport to gate projection. + if active.get("vsplit_ticket") is not None and not active.get("vsplit_k_ready", False): + return + from vllm_omni.diffusion.attention.ops.sage_quantization import quantize_sage_q_sm120 + from vllm_omni.diffusion.models.minimax_h3.attention.vsa import ( + _get_h3_layout, + _get_h3_tile_metadata, + _get_h3_tiled_source_rows, + _pool_h3_tiles, + h3_vsa_tile_pack, + ) + + assert query.shape == (1, 95936, 7, 128) and query.dtype == torch.bfloat16 + assert query.is_cuda and not query.requires_grad + assert torch.get_device_module().current_stream(query.device) == active["comm"] + assert "v3_q_ready" not in active, "Q preparation must run once per layer" + layout = _get_h3_layout(metadata) + assert layout == ((558, 1206), (107, 22, 40), 1764), layout + prefix, shape, _ = layout + side = active["side"] + side.wait_stream(active["comm"]) + with torch.get_device_module().stream(side), torch.cuda.nvtx.range("h3.lossless.v3.early_q"): + query.record_stream(side) + _, sizes, _, _, _, _ = _get_h3_tile_metadata(prefix, shape, query.device) + maps = _get_h3_tiled_source_rows(prefix, shape, 105472, query.device) + q = h3_vsa_tile_pack(query[:, :95924], maps) + qp = _pool_h3_tiles(q, sizes) + q_int8, q_scale = quantize_sage_q_sm120(q.transpose(1, 2)) + active["v3_q_ready"] = dict(q=q, qp=qp, q_int8=q_int8, q_scale=q_scale) + _log("early_q") + + +def prepared_q(): + active = ACTIVE.get() + if active is None: + return None + ready = active.get("v3_q_ready") + if ready is None: + return None + return ready["q"], ready["qp"] + + +def quantized_q(query: torch.Tensor): + active = ACTIVE.get() + if active is None: + return None + ready = active.get("v3_q_ready") + if ready is None: + return None + q = ready["q"] + assert ( + query.data_ptr() == q.data_ptr() + and query.shape == q.shape + and query.stride() == q.stride() + and query.dtype == q.dtype + and query.device == q.device + ), "prepared Q identity changed before Sage quantization" + assert torch.get_device_module().current_stream(query.device) == active["comm"] + # prepared() already enqueued the producer-stream join. Keep the tensors + # in this invocation's state and protect their use on the consumer stream. + q_int8, q_scale = ready["q_int8"], ready["q_scale"] + q_int8.record_stream(active["comm"]) + q_scale.record_stream(active["comm"]) + return q_int8, q_scale + + +@dataclass +class _CoarseTicket: + active_id: int + ready: Any + inputs: tuple[torch.Tensor, torch.Tensor] + used: bool = False + + +def coarse_ready(v_tiled: torch.Tensor, scores: torch.Tensor): + active = ACTIVE.get() + if not _enabled("COARSE_OVERLAP") or active is None: + return None + # ACTIVE is scoped to the fixed geometry, owner-local gate/O bundle route. + prepared = active["prepared"] + assert prepared is not None and prepared[4] is scores + assert active["reverse"] is None + assert v_tiled.shape == (1, 105472, 7, 128) and v_tiled.dtype == torch.bfloat16 + assert scores.shape == (1, 7, 1648, 1648) and scores.dtype == torch.float32 + assert v_tiled.device == scores.device + assert torch.get_device_module().current_stream(v_tiled.device) == active["comm"] + assert "v3_coarse_ticket" not in active, "coarse preparation must run once per layer" + ready = torch.get_device_module().Event(enable_timing=False) + ready.record(active["comm"]) + v_tiled.record_stream(active["side"]) + scores.record_stream(active["side"]) + ticket = _CoarseTicket(id(active), ready, (v_tiled, scores)) + active["v3_coarse_ticket"] = ticket + _log("coarse_overlap") + return ticket + + +@contextmanager +def coarse_scope(ticket): + if ticket is None: + yield + return + active = ACTIVE.get() + assert isinstance(ticket, _CoarseTicket) and active is not None and ticket.active_id == id(active) + assert not ticket.used, "coarse ticket must be consumed once" + ticket.used = True + side = active["side"] + # The event was recorded before fine attention. Waiting on the current + # comm stream here would also wait for fine and eliminate the overlap. + side.wait_event(ticket.ready) + with torch.get_device_module().stream(side), torch.cuda.nvtx.range("h3.lossless.v3.coarse"): + try: + yield + finally: + done = torch.get_device_module().Event(enable_timing=False) + done.record(side) + active["v3_coarse_done"] = done + + +def join_coarse(compressed: torch.Tensor) -> None: + active = ACTIVE.get() + if active is None: + return + ticket = active.get("v3_coarse_ticket") + if ticket is None: + return + assert ticket.used and "v3_coarse_done" in active + assert torch.get_device_module().current_stream(compressed.device) == active["comm"] + active["comm"].wait_event(active["v3_coarse_done"]) + compressed.record_stream(active["comm"]) + + +OUTPUT_LOOKAHEAD_ENV = "VLLM_OMNI_H3_O_PRODUCER_LOOKAHEAD" + + +@lru_cache(None) +def output_producer_stream(device): + return torch.get_device_module().Stream(device=device, priority=-1) + + +def prepare_output_chunks(active, fine, coarse, tiles, state): + from vllm_omni.diffusion.models.minimax_h3.ops.attention.chunks import produce + + if os.environ.get(OUTPUT_LOOKAHEAD_ENV) != "1" or "o_lookahead" in active: + raise RuntimeError("This arm requires the fixed four-chunk O lookahead route") + assert torch.get_device_module().current_stream(fine.device) == active["comm"] + assert fine.dtype == coarse.dtype == torch.bfloat16 + assert tuple(fine.shape) == (1, 95936, 7, 128) + assert tuple(coarse.shape) == (1, 1648, 7, 128) + assert dist.get_world_size() == 8 and state.communicator is not None + count = 2998 + shape = (1, 8 * (count + 270), 7, 128) + entries = [] + for index in range(4): + prototype = state.registered_shape_prototype(shape, dtype=fine.dtype, device=fine.device) + out = state.output(prototype, op="gather_heads", slot=f"h3_o_chunk4_{index}") + landing = state.communicator.input_buffer(out, shape) + assert landing.is_contiguous() and out.is_contiguous() + entries.append(dict(out=out, landing=landing, ready=None)) + # Concurrent writes must use separate registered source slots and must not + # overlap any receive buffer consumed by gate/projection on the side stream. + ranges = sorted( + (t.data_ptr(), t.data_ptr() + t.numel() * t.element_size()) for e in entries for t in (e["landing"], e["out"]) + ) + assert all(left[1] <= right[0] for left, right in zip(ranges, ranges[1:])), "O registered ranges overlap" + stream = output_producer_stream(fine.device) + ticket = dict(stream=stream, entries=entries, consumed=0) + active["o_lookahead"] = ticket # Install before submitting anything: cleanup also covers exceptions. + stream.wait_stream(active["comm"]) # fine and the original coarse join precede this event. + for tensor in (fine, coarse, tiles): + tensor.record_stream(stream) + with torch.get_device_module().stream(stream): + for index, e in enumerate(entries): + e["landing"].record_stream(stream) + with torch.cuda.nvtx.range(f"h3.oproducer.lookahead.produce.{index}"): + produce(fine, coarse, tiles, e["landing"], index * count, count) + ready = torch.get_device_module().Event(enable_timing=False) + ready.record(stream) + e["ready"] = ready + return ticket + + +def consume_output_chunk(active, index): + ticket = active["o_lookahead"] + assert index == ticket["consumed"] and 0 <= index < 4 + entry = ticket["entries"][index] + assert entry["ready"] is not None + active["comm"].wait_event(entry["ready"]) + entry["landing"].record_stream(active["comm"]) + ticket["consumed"] += 1 + return entry["out"], entry["landing"] + + +def finish_output_chunks(active): + if active["o_lookahead"]["consumed"] != 4: + raise RuntimeError("O producer must be consumed once per output chunk") + + +def drain_output_producer(active, caller): + ticket = active.get("o_lookahead") + if ticket is not None: + caller.wait_stream(ticket["stream"]) + for entry in ticket["entries"]: + entry["landing"].record_stream(caller) + + +H3_VSA_O_BUNDLE_ENV = "VLLM_OMNI_FASTVIDEO_VSA_O_BUNDLE" +H3_VSA_O_BUNDLE_ACTIVE_KEY = "vsa_h3_o_bundle_active" +H3_VSA_O_BUNDLE_STATE_KEY = "vsa_h3_o_bundle_state" + + +def h3_vsa_o_bundle_enabled() -> bool: + """Return whether reverse-O compact-coarse piggybacking was requested.""" + raw = os.environ.get(H3_VSA_O_BUNDLE_ENV, "0").strip() + if raw not in {"0", "1"}: + raise ValueError(f"{H3_VSA_O_BUNDLE_ENV} must be exactly '0' or '1', got {raw!r}") + return raw == "1" diff --git a/vllm_omni/diffusion/models/minimax_h3/attention/parallel.py b/vllm_omni/diffusion/models/minimax_h3/attention/parallel.py new file mode 100644 index 00000000000..1d68c6d60cd --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/attention/parallel.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""H3 gate ownership and producer scheduling over the shared Ulysses transport.""" + +from dataclasses import dataclass + +import torch +import torch.distributed as dist + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.parallel.base import ParallelAttentionContext +from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention +from vllm_omni.diffusion.forward_context import get_ulysses_mode + +from . import overlap, qkv_overlap +from .overlap import H3_VSA_O_BUNDLE_ACTIVE_KEY, H3_VSA_O_BUNDLE_STATE_KEY, h3_vsa_o_bundle_enabled +from .vsa import _get_h3_layout, get_h3_vsa_owner_route_plan + + +@dataclass(frozen=True) +class H3OverlapContext(ParallelAttentionContext): + ulysses_pg: dist.ProcessGroup + o_bundle_gate_local: torch.Tensor + o_bundle_metadata: AttentionMetadata + o_bundle_state: dict + strict_a2a_backend: str = "flashinfer-pcie" + joint_len: int = 0 + use_uaa: bool = False + + +class H3UlyssesAttention(UlyssesParallelAttention): + """Reuse ordinary resharding; specialize only the active H3 overlap scope.""" + + def _exchange_qkv(self, query, key, value, attn_metadata): + if overlap.ACTIVE.get() is None: + return super()._exchange_qkv(query, key, value, attn_metadata) + query = self._scatter_heads(query, "q") + overlap.after_q(query, attn_metadata) + qkv_overlap.after_q() + key = self._scatter_heads(key, "k") + value = overlap.before_v(query, key, value, attn_metadata, self._ulysses_pg) + value = self._scatter_heads(value, "v") + return query, key, value + + def pre_attention(self, query, key, value, attn_metadata): + if overlap.ACTIVE.get() is None: + return super().pre_attention(query, key, value, attn_metadata) + if ( + get_ulysses_mode(default="strict") != "strict" + or self._sp_group.ring_world_size != 1 + or self._sp_group.ulysses_world_size != 8 + or self._ulysses_a2a_backend != "flashinfer-pcie" + or not self._ulysses_a2a_permute + or (self._scatter_idx, self._gather_idx) != (2, 1) + ): + raise RuntimeError("H3 overlap requires strict TP1/SP8 FlashInfer Ulysses without Ring") + if not h3_vsa_o_bundle_enabled() or attn_metadata is None: + raise RuntimeError("H3 overlap requires reverse-O bundling and model layout metadata") + if any(getattr(attn_metadata, key) is not None for key in ("joint_query", "joint_key", "joint_value")): + raise RuntimeError("H3 overlap does not support joint attention rows") + if any(key in attn_metadata.extra for key in (H3_VSA_O_BUNDLE_ACTIVE_KEY, H3_VSA_O_BUNDLE_STATE_KEY)): + raise RuntimeError("H3 overlap received stale per-layer output state") + gate = attn_metadata.extra.get("gate_compress") + if not isinstance(gate, torch.Tensor) or gate.shape != query.shape or gate.dtype != torch.bfloat16: + raise ValueError("H3 overlap requires a BF16 gate matching the local Q layout") + if gate.device != query.device or not gate.is_contiguous() or gate.requires_grad: + raise ValueError("H3 overlap gate must be contiguous, inference-only and on the Q device") + layout = _get_h3_layout(attn_metadata) + if layout is None: + raise ValueError("H3 overlap requires prefix and target-video layout metadata") + prefix, video_shape, target_start = layout + if sum(prefix) != target_start: + raise ValueError("H3 prefix/target boundary mismatch") + import math + + valid_rows = sum(prefix) + math.prod(video_shape) + aligned_rows = query.shape[1] * self._sp_group.ulysses_world_size + packed = attn_metadata.packed_padding + if packed is None: + if valid_rows != aligned_rows: + raise ValueError("H3 aligned rows require packed-padding metadata") + elif packed.q_length != valid_rows or packed.kv_length != valid_rows: + raise ValueError("H3 packed lengths must match the prefix/video geometry") + plan = get_h3_vsa_owner_route_plan(prefix, video_shape, aligned_rows, self._sp_group.ulysses_world_size) + state = {"plan": plan} + attn_metadata.extra.pop("gate_compress") + attn_metadata.extra[H3_VSA_O_BUNDLE_ACTIVE_KEY] = True + attn_metadata.extra[H3_VSA_O_BUNDLE_STATE_KEY] = state + try: + query, key, value, attn_metadata, _ = super().pre_attention(query, key, value, attn_metadata) + except BaseException: + attn_metadata.extra.pop(H3_VSA_O_BUNDLE_ACTIVE_KEY, None) + attn_metadata.extra.pop(H3_VSA_O_BUNDLE_STATE_KEY, None) + raise + return ( + query, + key, + value, + attn_metadata, + H3OverlapContext( + name=self.name, + ulysses_pg=self._ulysses_pg, + o_bundle_gate_local=gate, + o_bundle_metadata=attn_metadata, + o_bundle_state=state, + ), + ) + + def post_attention(self, attn_output, ctx): + if not isinstance(ctx, H3OverlapContext): + return super().post_attention(attn_output, ctx) + try: + result = overlap.finish_reverse(attn_output, ctx) + if result is None: + raise RuntimeError("H3 output producer scope ended before reverse exchange") + return result + finally: + ctx.o_bundle_metadata.extra.pop(H3_VSA_O_BUNDLE_ACTIVE_KEY, None) + ctx.o_bundle_metadata.extra.pop(H3_VSA_O_BUNDLE_STATE_KEY, None) + + +def configure_parallel_attention(attention) -> None: + strategy = attention.parallel_strategy + if isinstance(strategy, UlyssesParallelAttention) and strategy._ulysses_a2a_backend == "flashinfer-pcie": + attention.parallel_strategy = H3UlyssesAttention( + strategy._sp_group, + strategy._scatter_idx, + strategy._gather_idx, + strategy._use_sync, + strategy._ulysses_a2a_permute, + ) diff --git a/vllm_omni/diffusion/models/minimax_h3/attention/qkv_overlap.py b/vllm_omni/diffusion/models/minimax_h3/attention/qkv_overlap.py new file mode 100644 index 00000000000..b087615e937 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/attention/qkv_overlap.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""QK/V producer scheduling for the supported H3 overlap geometry. + +Only tensor arguments cross the regional-compile boundary. Tickets, streams, +events and producer state live entirely in the existing eager attention +island. The V placeholder aliases raw Q and must never be read as V: the +strictly guarded Ulysses path replaces it before its first V data access. +""" + +import json +import os +from dataclasses import dataclass +from functools import lru_cache +from typing import Any + +import torch +import torch.distributed as dist +from vllm.logger import init_logger + +logger = init_logger(__name__) +ROWS, HIDDEN, WIDTH = 11992, 5376, 7168 +_AUDIT_CALLS: dict[int, int] = {} +_GLOBAL_AUDIT_CALL = 0 +_ACTIVE_LOGGED = set() + + +def configured_mode(prefix: str) -> str: + if os.getenv("VLLM_OMNI_H3_VSPLIT", "0") != "1" or not prefix.startswith("blocks."): + return "off" + mode = os.getenv("VLLM_OMNI_H3_VSPLIT_MODE", "split") + if mode != "split" or os.environ.get("VLLM_OMNI_H3_MXFP8_VSPLIT") != "1": + raise ValueError(f"Invalid H3 VSPLIT mode: {mode!r}") + for flag in ( + "ATTENTION_OVERLAP", + "LOSSLESS_GATE", + "LOSSLESS_EARLY_Q", + "LOSSLESS_LAYOUT_VIEW", + "LOSSLESS_COARSE_OVERLAP", + ): + if os.getenv("VLLM_OMNI_H3_" + flag, "0") != "1": + raise RuntimeError(f"H3 VSPLIT requires the existing v3 flag H3_{flag}=1") + if os.getenv("VLLM_BATCH_INVARIANT", "0") != "0": + raise RuntimeError("H3 VSPLIT does not replace batch-invariant GEMM") + return mode + + +def validate_projection(projection, prepared) -> None: + from vllm.distributed import get_tensor_model_parallel_world_size + + from vllm_omni.diffusion.layers.mxfp8 import _scale_numel + from vllm_omni.diffusion.models.minimax_h3.quantization import validate_split_projection + + if not dist.is_initialized() or dist.get_world_size() != 8 or get_tensor_model_parallel_world_size() != 1: + raise RuntimeError("MXFP8 split requires initialized TP1/SP8") + if not isinstance(prepared, tuple) or len(prepared) != 3: + raise RuntimeError("MXFP8 split requires explicit quantized X, scale and QK tensors") + x, scale, qk = prepared + expected = ( + ((ROWS, HIDDEN), torch.float8_e4m3fn), + ((_scale_numel(ROWS, HIDDEN),), torch.uint8), + ((ROWS, 2 * WIDTH), torch.bfloat16), + ) + for tensor, (shape, dtype) in zip(prepared, expected): + if ( + tuple(tensor.shape) != shape + or tensor.dtype != dtype + or not tensor.is_cuda + or tensor.device != projection.weight.device + or not tensor.is_contiguous() + or tensor.requires_grad + or torch.is_grad_enabled() + ): + raise RuntimeError("MXFP8 split prepared tensor contract changed") + validate_split_projection(projection) + + +@lru_cache(None) +def _v_stream(device_index: int): + return torch.get_device_module().Stream(device=device_index) + + +@dataclass +class Ticket: + x: torch.Tensor + weight: torch.Tensor + placeholder: torch.Tensor + stream: Any + x_ready: Any + layer_index: int + x_scale: torch.Tensor + qk: torch.Tensor + projection: Any + result: Any = None + done: Any = None + launched: bool = False + consumed: bool = False + + +def begin(active, attention, projection, prepared, placeholder, layer_index: int, metadata): + validate_projection(projection, prepared) + x, x_scale, qk = prepared + validate_route(attention, metadata) + if not 0 <= layer_index < 50 or tuple(placeholder.shape) != (1, ROWS, 56, 128): + raise RuntimeError("H3 VSPLIT placeholder geometry changed") + ready = torch.get_device_module().Event(enable_timing=False) + ready.record(active["comm"]) + stream = _v_stream(x.device.index) + for tensor in (x, x_scale, qk, projection.weight, projection.mxfp8_weight_scale): + tensor.record_stream(stream) + ticket = Ticket(x, projection.weight, placeholder, stream, ready, layer_index, x_scale, qk, projection) + active["vsplit_ticket"] = ticket + # V depends only on the prepared input. Launch before Q transport, then + # hand the SMs to gate projection rather than running both large GEMMs + # concurrently. Gate/early-Q/QK preparation are queued on the side stream. + stream.wait_event(ready) + ticket.launched = True # Exception cleanup joins partial submissions. + with torch.get_device_module().stream(stream): + from vllm_omni.diffusion.models.minimax_h3.quantization import project_split_v + + ticket.result = project_split_v(projection, x, x_scale, qk).view(1, ROWS, 56, 128) + ticket.done = torch.get_device_module().Event(enable_timing=False) + ticket.done.record(stream) + key = (dist.get_rank(), layer_index) + if key not in _ACTIVE_LOGGED: + _ACTIVE_LOGGED.add(key) + logger.info( + "H3_VSPLIT_ACTIVE %s", + json.dumps( + dict( + schema=1, + rank=key[0], + layer_index=layer_index, + mode="split", + schedule="before_q_v_then_gate_qk_after_k", + v_stream="dedicated", + qk_shape=[ROWS, 2 * WIDTH], + v_shape=[ROWS, WIDTH], + ), + sort_keys=True, + ), + ) + + +def validate_route(attention, metadata) -> None: + from vllm_omni.diffusion.forward_context import get_ulysses_mode + + strategy = attention._get_active_parallel_strategy() + if ( + strategy is not attention.parallel_strategy + or strategy.name != "ulysses" + or strategy._sp_group.ulysses_world_size != 8 + or strategy._sp_group.ring_world_size != 1 + or get_ulysses_mode(default="strict") != "strict" + or strategy._ulysses_a2a_backend != "flashinfer-pcie" + or not strategy._ulysses_a2a_permute + or strategy._scatter_idx != 2 + or strategy._gather_idx != 1 + or not strategy._qk_input_landing_enabled + or attention.use_ring + or attention._active_paged_kv_adapter() is not None + or any(getattr(metadata, name, None) is not None for name in ("joint_query", "joint_key", "joint_value")) + ): + raise RuntimeError("H3 VSPLIT requires the exact non-joint BF16 Q/K/V Ulysses SP8 route") + + +def after_q() -> None: + from vllm_omni.diffusion.models.minimax_h3.attention.overlap import ACTIVE + + active = ACTIVE.get() + ticket = None if active is None else active.get("vsplit_ticket") + if ticket is None: + return + assert active is not None + if not ticket.launched or ticket.done is None or "v3_q_ready" in active: + raise RuntimeError("H3 VSPLIT requires V submitted before Q and Q preparation deferred until K") + + +def before_v(active, query, key, placeholder, metadata, group): + """Schedule K-dependent work before joining the independent V producer.""" + from vllm_omni.diffusion.distributed.flashinfer_ulysses import _state_for + from vllm_omni.diffusion.models.minimax_h3.attention.overlap import prepared_q + from vllm_omni.diffusion.models.minimax_h3.attention.vsa import ( + _build_h3_ordered_q2k_indices, + _get_h3_layout, + _get_h3_tile_metadata, + _get_h3_tiled_source_rows, + _pool_h3_tiles, + h3_vsa_tile_pack, + ) + + ticket = active["vsplit_ticket"] + if ( + not ticket.launched + or ticket.done is None + or ticket.consumed + or placeholder.data_ptr() != ticket.placeholder.data_ptr() + or placeholder.shape != ticket.placeholder.shape + or placeholder.stride() != ticket.placeholder.stride() + ): + raise RuntimeError("H3 VSPLIT V ticket missing, repeated, or rebound") + layout = _get_h3_layout(metadata) + if layout != ((558, 1206), (107, 22, 40), 1764): + raise RuntimeError("H3 VSPLIT layout changed") + prefix, shape, _ = layout + if active.get("vsplit_k_ready", False) or "v3_q_ready" in active: + raise RuntimeError("Deferred Q preparation must be submitted once after K") + active["vsplit_k_ready"] = True + from vllm_omni.diffusion.models.minimax_h3.attention.overlap import after_q + + after_q(query, metadata) + with torch.cuda.nvtx.range("h3.vsplit.qk_ready"): + active["side"].wait_stream(active["comm"]) + with torch.get_device_module().stream(active["side"]), torch.cuda.nvtx.range("h3.overlap.qk_prepare"): + _, sizes, _, _, prefix_blocks, video_blocks = _get_h3_tile_metadata(prefix, shape, query.device) + maps = _get_h3_tiled_source_rows(prefix, shape, 105472, query.device) + q_ready = prepared_q() + if q_ready is None: + raise RuntimeError("H3 VSPLIT requires prepared earlyQ") + q, qp = q_ready + k = h3_vsa_tile_pack(key[:, :95924], maps) + kp = _pool_h3_tiles(k, sizes) + scores = torch.matmul(qp, kp.transpose(-2, -1)) * active["impl"].softmax_scale + idx, num = _build_h3_ordered_q2k_indices(scores, prefix_blocks, video_blocks, 162) + active["prepared"] = (q, k, qp, kp, scores, idx, num) + # This is the first V data access. Waiting here cannot delay the already + # submitted side-stream QK preparation above. + with torch.cuda.nvtx.range("h3.vsplit.v_ready_wait"): + active["comm"].wait_event(ticket.done) + value = ticket.result + value.record_stream(active["comm"]) + state = _state_for(value, group, 8) + if state.communicator is None: + raise RuntimeError("H3 VSPLIT requires the active RDMA communicator") + out = state.output(value, op="scatter_heads", slot="v") + landing = state.communicator.input_buffer(out, value.shape) + landing.copy_(value) + ticket.consumed = True + return landing + + +def finish(active) -> None: + ticket = active.get("vsplit_ticket") + if ticket is not None and not ticket.consumed: + raise RuntimeError("H3 VSPLIT V hook was not consumed") + + +def cleanup(active, caller) -> None: + ticket = active.get("vsplit_ticket") + if ticket is not None: + caller.wait_stream(ticket.stream) + for tensor in ( + ticket.x, + ticket.x_scale, + ticket.qk, + ticket.weight, + ticket.projection.mxfp8_weight_scale, + ticket.result, + ): + if tensor is not None: + tensor.record_stream(caller) diff --git a/vllm_omni/diffusion/models/minimax_h3/attention/vsa.py b/vllm_omni/diffusion/models/minimax_h3/attention/vsa.py new file mode 100644 index 00000000000..721977c22ac --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/attention/vsa.py @@ -0,0 +1,918 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from __future__ import annotations + +import functools +import math +import os +from contextlib import contextmanager +from typing import Any + +import torch +from vllm.logger import init_logger + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.fastvideo_vsa import ( + FastVideoVSAImpl, + _construct_variable_block_sizes, + _get_gate_compress, + _get_non_pad_index, + _get_tile_partition_indices, +) +from vllm_omni.diffusion.models.minimax_h3.attention.overlap import ( + H3_VSA_O_BUNDLE_ACTIVE_KEY, + H3_VSA_O_BUNDLE_ENV, + H3_VSA_O_BUNDLE_STATE_KEY, + h3_vsa_o_bundle_enabled, +) +from vllm_omni.diffusion.models.minimax_h3.ops.attention.layout import ( + H3_VSA_FUSED_TILE_PACK_ENV, + H3_VSA_FUSED_UNTILE_ENV, + build_h3_aligned_untile_source_rows, + build_h3_tiled_source_rows, + h3_vsa_fused_tile_pack_enabled, + h3_vsa_fused_untile_enabled, + h3_vsa_tile_pack, + h3_vsa_tile_pack_cuda_supported, + h3_vsa_tile_untile, + h3_vsa_tile_untile_cuda_supported, + h3_vsa_tile_untile_out, +) +from vllm_omni.diffusion.models.minimax_h3.ops.attention.owner_route import ( + H3VSAOwnerRoutePlan, + build_h3_vsa_owner_route_plan, + h3_vsa_owner_route_plan_is_trusted, +) +from vllm_omni.platforms import current_omni_platform + +logger = init_logger(__name__) +H3_VSA_DIRECT_Q2K_ENV = "VLLM_OMNI_FASTVIDEO_VSA_DIRECT_Q2K" +H3_VSA_SKIP_SOFTMAX_THRESHOLD_ENV = "VLLM_OMNI_FASTVIDEO_VSA_SKIP_SOFTMAX_THRESHOLD_SCALE_FACTOR" +H3_VSA_NVTX_ENV = "VLLM_OMNI_MINIMAX_H3_VSA_NVTX" +H3_VSA_NVTX_DOMAIN = "vllm_omni.minimax_h3.vsa" +H3_VSA_EFFECTIVE_TILE_SHAPE = (4, 4, 4) +H3_VSA_EFFECTIVE_TILE_SIZE = math.prod(H3_VSA_EFFECTIVE_TILE_SHAPE) + + +def h3_vsa_direct_q2k_enabled() -> bool: + return os.environ.get(H3_VSA_DIRECT_Q2K_ENV, "0") == "1" + + +@contextmanager +def _h3_vsa_nvtx_stage(name: str): + """Emit nested VSA phase ranges only for an explicit diagnostic run.""" + if os.environ.get(H3_VSA_NVTX_ENV, "0") != "1": + yield + return + import nvtx + + nvtx.push_range(name, domain=H3_VSA_NVTX_DOMAIN) + try: + yield + finally: + nvtx.pop_range(domain=H3_VSA_NVTX_DOMAIN) + + +def h3_vsa_skip_softmax_threshold_scale_factor() -> float: + """Return the opt-in SM120 VSA skip-softmax threshold. + + Zero is the exact/default kernel specialization. A positive value enables + an approximate early-skip policy in FlashInfer's SM120 CuTeDSL kernel, so + malformed values must fail closed instead of silently selecting a numeric + path that was not requested. + """ + raw = os.environ.get(H3_VSA_SKIP_SOFTMAX_THRESHOLD_ENV, "0").strip() + try: + value = float(raw) + except ValueError as exc: + raise ValueError( + f"{H3_VSA_SKIP_SOFTMAX_THRESHOLD_ENV} must be a finite non-negative float, got {raw!r}" + ) from exc + if not math.isfinite(value) or value < 0: + raise ValueError(f"{H3_VSA_SKIP_SOFTMAX_THRESHOLD_ENV} must be a finite non-negative float, got {raw!r}") + return value + + +if not hasattr(torch.ops.vllm_omni, "fastvideo_h3_vsa_bhsd"): + + @torch.library.custom_op("vllm_omni::fastvideo_h3_vsa_bhsd", mutates_args=()) + def _fastvideo_h3_vsa_bhsd_op( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, + logical_blocks: int, + ) -> torch.Tensor: + q = query.transpose(1, 2).contiguous() + k = key.transpose(1, 2).contiguous() + v = value.transpose(1, 2).contiguous() + if os.environ.get("FASTVIDEO_VSA_SM100A", "0") == "1": + try: + from fastvideo_kernel import block_sparse_attn_sm100a + from fastvideo_kernel.triton_kernels.index import map_to_index + + if block_sparse_attn_sm100a.is_supported(q, variable_block_sizes): + (q2k_idx, q2k_num) = map_to_index(block_map) + (out, _) = block_sparse_attn_sm100a.block_sparse_attn_sm100a( + q, + k, + v, + q2k_idx.to(torch.int32).contiguous(), + q2k_num.to(torch.int32).contiguous(), + variable_block_sizes.to(torch.int32).contiguous(), + need_lse=False, + ) + return out.transpose(1, 2).contiguous() + except (ImportError, RuntimeError) as exc: + logger.warning_once( + "FASTVIDEO_VSA_SM100A=1 requested but the native Blackwell forward is " + "unavailable (%s); using the Triton block-sparse route instead.", + exc, + ) + from fastvideo_kernel.block_sparse_attn import block_sparse_attn + + logical_len = logical_blocks * H3_VSA_EFFECTIVE_TILE_SIZE + (out, _) = block_sparse_attn( + q[:, :, :logical_len].contiguous(), + k[:, :, :logical_len].contiguous(), + v[:, :, :logical_len].contiguous(), + block_map[..., :logical_blocks, :logical_blocks].contiguous(), + variable_block_sizes[:logical_blocks].to(torch.int32).contiguous(), + ) + out = out.transpose(1, 2).contiguous() + if out.shape[1] != query.shape[1]: + out = torch.nn.functional.pad(out, (0, 0, 0, 0, 0, query.shape[1] - out.shape[1])) + return out + + @_fastvideo_h3_vsa_bhsd_op.register_fake + def _(query, key, value, block_map, variable_block_sizes, logical_blocks): + del key, value, block_map, variable_block_sizes, logical_blocks + return torch.empty_like(query) + + +_fastvideo_h3_vsa_bhsd_op = torch.ops.vllm_omni.fastvideo_h3_vsa_bhsd + + +def _flashinfer_h3_vsa_q2k_bshd_impl( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + variable_block_sizes: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + """Dispatch H3's blk64 VSA contract to the native FlashInfer kernel.""" + skip_softmax_threshold = h3_vsa_skip_softmax_threshold_scale_factor() + if q2k_idx.dtype != torch.int32 or not q2k_idx.is_contiguous(): + raise ValueError("q2k_idx must be a contiguous int32 tensor") + if q2k_num.dtype != torch.int32 or not q2k_num.is_contiguous(): + raise ValueError("q2k_num must be a contiguous int32 tensor") + blocks = query.shape[1] // H3_VSA_EFFECTIVE_TILE_SIZE + expected_idx_shape = (query.shape[0], query.shape[2], blocks, blocks) + expected_num_shape = expected_idx_shape[:-1] + if tuple(q2k_idx.shape) != expected_idx_shape: + raise ValueError(f"q2k_idx shape {tuple(q2k_idx.shape)} must be {expected_idx_shape}") + if tuple(q2k_num.shape) != expected_num_shape: + raise ValueError(f"q2k_num shape {tuple(q2k_num.shape)} must be {expected_num_shape}") + if q2k_idx.device != query.device or q2k_num.device != query.device: + raise ValueError("q2k_idx and q2k_num must be on the same device as query") + block_sizes = variable_block_sizes.to(torch.int32).contiguous() + sage_mode = os.environ.get("VLLM_OMNI_H3_VSA_SAGE", "bf16") + if sage_mode not in ("bf16", "sage"): + raise ValueError("H3 VSA compute must be bf16 or sage") + if sage_mode != "bf16": + if skip_softmax_threshold: + raise ValueError("Sage requires the zero skip-softmax threshold") + from vllm_omni.diffusion.attention.ops.sage_block_sparse_attention import sage_block_sparse_attention + from vllm_omni.diffusion.models.minimax_h3.attention.overlap import quantized_q + + with _h3_vsa_nvtx_stage("vsa.sage"): + return sage_block_sparse_attention( + query, key, value, q2k_idx, q2k_num, block_sizes, softmax_scale, prepared_q=quantized_q(query) + ) + try: + (major, minor) = current_omni_platform.get_device_capability( + query.device.index if query.device.index is not None else torch.accelerator.current_device_index() + ) + except (AssertionError, RuntimeError, ValueError) as exc: + raise RuntimeError("FlashInfer H3 blk64 VSA requires a CUDA tensor on a supported GPU") from exc + if (major, minor) in ((10, 0), (10, 3)): + if skip_softmax_threshold: + raise RuntimeError( + f"{H3_VSA_SKIP_SOFTMAX_THRESHOLD_ENV} is supported only by FlashInfer's " + f"SM120/SM121 blk64 CuTeDSL kernel" + ) + if query.dtype != torch.bfloat16: + raise ValueError(f"FlashInfer SM100/SM103 blk64 VSA requires bfloat16, got {query.dtype}") + try: + from flashinfer.cute_dsl.sparse.bsa_attn_sm100_blk64 import bsa_attn_sm100_blk64_fwd + except (ImportError, AttributeError) as exc: + raise ImportError( + "FlashInfer H3 VSA on SM100/SM103 requires flashinfer.cute_dsl.sparse.bsa_attn_sm100_blk64" + ) from exc + kernel = bsa_attn_sm100_blk64_fwd + kernel_name = "vsa_sm100_blk64_cuda" + elif (major, minor) in ((12, 0), (12, 1)): + try: + from flashinfer.cute_dsl.sparse.bsa_attn_sm120 import bsa_attn_sm120_blk64_fwd + except (ImportError, AttributeError) as exc: + raise ImportError( + "FlashInfer H3 VSA on SM120/SM121 requires the #4944 blk64 provider " + "flashinfer.cute_dsl.sparse.bsa_attn_sm120" + ) from exc + kernel = bsa_attn_sm120_blk64_fwd + kernel_name = "vsa_sm120_blk64_cute_dsl" + else: + raise RuntimeError( + f"FlashInfer H3 blk64 VSA supports SM100/SM103 and SM120/SM121; current device is SM{major}{minor}" + ) + logger.info_once("FASTVIDEO_VSA H3 compute kernel: FlashInfer %s", kernel_name) + if skip_softmax_threshold: + logger.warning_once( + "FASTVIDEO_VSA H3 approximate skip-softmax enabled: %s=%g; this output is not part " + "of the exact VSA baseline", + H3_VSA_SKIP_SOFTMAX_THRESHOLD_ENV, + skip_softmax_threshold, + ) + kernel_kwargs: dict[str, Any] = {} + if (major, minor) in ((12, 0), (12, 1)) and skip_softmax_threshold > 0: + kernel_kwargs["skip_softmax_threshold_scale_factor"] = skip_softmax_threshold + (output, _) = kernel( + query, + key, + value, + q2k_idx, + int(q2k_idx.shape[-1]), + block_sizes=block_sizes, + q2k_block_nums=q2k_num, + softmax_scale=softmax_scale, + return_lse=False, + **kernel_kwargs, + ) + return output + + +def _flashinfer_h3_vsa_bshd_impl( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + """Legacy FlashInfer wrapper accepting FastVideo's dense bool map.""" + from fastvideo_kernel.triton_kernels.index import map_to_index + + (q2k_idx, q2k_num) = map_to_index(block_map) + return _flashinfer_h3_vsa_q2k_bshd_impl( + query, + key, + value, + q2k_idx.to(torch.int32).contiguous(), + q2k_num.to(torch.int32).contiguous(), + variable_block_sizes, + softmax_scale, + ) + + +if not hasattr(torch.ops.vllm_omni, "flashinfer_h3_vsa_bshd"): + + @torch.library.custom_op("vllm_omni::flashinfer_h3_vsa_bshd", mutates_args=()) + def _flashinfer_h3_vsa_bshd_op( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + block_map: torch.Tensor, + variable_block_sizes: torch.Tensor, + softmax_scale: float, + ) -> torch.Tensor: + return _flashinfer_h3_vsa_bshd_impl(query, key, value, block_map, variable_block_sizes, softmax_scale) + + @_flashinfer_h3_vsa_bshd_op.register_fake + def _(query, key, value, block_map, variable_block_sizes, softmax_scale): + del key, value, block_map, variable_block_sizes, softmax_scale + return torch.empty_like(query) + + +_flashinfer_h3_vsa_bshd_op = torch.ops.vllm_omni.flashinfer_h3_vsa_bshd +if not hasattr(torch.ops.vllm_omni, "flashinfer_h3_vsa_q2k_bshd"): + + @torch.library.custom_op("vllm_omni::flashinfer_h3_vsa_q2k_bshd", mutates_args=()) + def _flashinfer_h3_vsa_q2k_bshd_op( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + q2k_idx: torch.Tensor, + q2k_num: torch.Tensor, + variable_block_sizes: torch.Tensor, + softmax_scale: float, + ) -> torch.Tensor: + return _flashinfer_h3_vsa_q2k_bshd_impl( + query, key, value, q2k_idx, q2k_num, variable_block_sizes, softmax_scale + ) + + @_flashinfer_h3_vsa_q2k_bshd_op.register_fake + def _(query, key, value, q2k_idx, q2k_num, variable_block_sizes, softmax_scale): + del key, value, q2k_idx, q2k_num, variable_block_sizes, softmax_scale + return torch.empty_like(query) + + +_flashinfer_h3_vsa_q2k_bshd_op = torch.ops.vllm_omni.flashinfer_h3_vsa_q2k_bshd + + +@functools.lru_cache(maxsize=32) +def _get_h3_tile_metadata( + prefix_segments: tuple[int, ...], video_shape: tuple[int, int, int], device: torch.device +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]: + """Official FastVideo H3 geometry: pure prefix chunks + 3-D video tiles.""" + block_size = H3_VSA_EFFECTIVE_TILE_SHAPE + block_elements = H3_VSA_EFFECTIVE_TILE_SIZE + prefix_len = sum(prefix_segments) + prefix_sizes: list[int] = [] + for segment in prefix_segments: + (full, remainder) = divmod(segment, block_elements) + prefix_sizes.extend([block_elements] * full) + if remainder: + prefix_sizes.append(remainder) + video_indices = _get_tile_partition_indices(video_shape, block_size, device) + prefix_len + video_sizes = _construct_variable_block_sizes(video_shape, block_size, device) + partition = torch.cat([torch.arange(prefix_len, device=device, dtype=torch.long), video_indices]) + sizes = torch.cat([torch.tensor(prefix_sizes, device=device, dtype=torch.int32), video_sizes.to(torch.int32)]) + non_pad = _get_non_pad_index(sizes, block_elements) + untile = non_pad[torch.argsort(partition)] + total = prefix_len + math.prod(video_shape) + if int(sizes.sum()) != total or untile.numel() != total: + raise ValueError( + f"invalid H3 VSA geometry: prefix={prefix_segments}, video={video_shape}, " + f"sizes_sum={int(sizes.sum())}, total={total}" + ) + return (partition, sizes, non_pad, untile, len(prefix_sizes), int(video_sizes.numel())) + + +@functools.lru_cache(maxsize=32) +def _get_h3_tiled_source_rows( + prefix_segments: tuple[int, ...], video_shape: tuple[int, int, int], padded_rows: int, device: torch.device +) -> torch.Tensor: + """Cache the validated destination-to-source map for fused H3 tiling.""" + (partition, _sizes, non_pad, _untile, _prefix_blocks, _video_blocks) = _get_h3_tile_metadata( + prefix_segments, video_shape, device + ) + return build_h3_tiled_source_rows(partition, non_pad, padded_rows) + + +@functools.lru_cache(maxsize=32) +def _get_h3_aligned_untile_source_rows( + prefix_segments: tuple[int, ...], + video_shape: tuple[int, int, int], + aligned_rows: int, + tiled_rows: int, + device: torch.device, +) -> torch.Tensor: + """Cache the validated tile-64 source map for aligned H3 output.""" + (_partition, sizes, _non_pad, untile, _prefix_blocks, _video_blocks) = _get_h3_tile_metadata( + prefix_segments, video_shape, device + ) + expected_tiled_rows = int(sizes.numel()) * H3_VSA_EFFECTIVE_TILE_SIZE + if tiled_rows != expected_tiled_rows: + raise ValueError(f"H3 VSA output has {tiled_rows} logical tiled rows, expected {expected_tiled_rows}") + return build_h3_aligned_untile_source_rows(untile, aligned_rows=aligned_rows, tiled_rows=tiled_rows) + + +@functools.lru_cache(maxsize=32) +@torch.compiler.disable +def get_h3_vsa_owner_route_plan( + prefix_segments: tuple[int, ...], video_shape: tuple[int, int, int], aligned_rows: int, sp_world_size: int +) -> H3VSAOwnerRoutePlan: + """Build and cache the exact aligned H3 row-owner route on CPU. + + Ulysses needs ``Kmax`` before the attention backend runs so it can acquire + an expanded reverse-O landing. Deriving the plan from the same H3 untile + metadata used below keeps the producer and consumer on one row-order + contract without synchronizing CUDA metadata back to the host. + """ + if aligned_rows <= 0 or sp_world_size <= 1: + raise ValueError( + f"H3 VSA owner routing requires positive aligned_rows and SP world_size > 1, got " + f"aligned_rows={aligned_rows}, sp_world_size={sp_world_size}" + ) + if aligned_rows % sp_world_size: + raise ValueError(f"H3 VSA aligned rows must divide the SP world: rows={aligned_rows}, world={sp_world_size}") + cpu = torch.device("cpu") + (_partition, sizes, _non_pad, untile, _prefix_blocks, _video_blocks) = _get_h3_tile_metadata( + prefix_segments, video_shape, cpu + ) + logical_blocks = int(sizes.numel()) + aligned_source_rows = build_h3_aligned_untile_source_rows( + untile, aligned_rows=aligned_rows, tiled_rows=logical_blocks * H3_VSA_EFFECTIVE_TILE_SIZE + ) + return build_h3_vsa_owner_route_plan( + aligned_source_rows, + sp_world_size=sp_world_size, + global_rows=aligned_rows, + local_rows=aligned_rows // sp_world_size, + logical_block_count=logical_blocks, + ) + + +def _get_h3_layout(attn_metadata: AttentionMetadata | None) -> tuple[tuple[int, ...], tuple[int, int, int], int] | None: + if attn_metadata is None or attn_metadata.video_layout is None: + return None + prefix = attn_metadata.extra.get("vsa_h3_prefix_segments") + if not isinstance(prefix, (tuple, list)): + return None + target = next((span for span in reversed(attn_metadata.video_layout.video_spans) if span.role == "target"), None) + if target is None: + return None + return (tuple(int(x) for x in prefix if int(x) > 0), target.latent_grid, target.start) + + +def _pool_h3_tiles(x: torch.Tensor, sizes: torch.Tensor) -> torch.Tensor: + (batch, seq_len, heads, dim) = x.shape + blocks = seq_len // H3_VSA_EFFECTIVE_TILE_SIZE + pooled = x.view(batch, blocks, H3_VSA_EFFECTIVE_TILE_SIZE, heads, dim).sum(dim=2, dtype=torch.float32) + pooled = pooled / sizes.view(1, -1, 1, 1).clamp_min(1) + return pooled.permute(0, 2, 1, 3) + + +def _build_h3_block_map(scores: torch.Tensor, num_prefix_blocks: int, num_video_blocks: int, topk: int) -> torch.Tensor: + """Prefix K/V are exempt and prefix queries stay dense, as in FastVideo.""" + keep_video = min(topk, num_video_blocks) + if keep_video == num_video_blocks: + return torch.ones_like(scores, dtype=torch.bool) + block_map = torch.zeros_like(scores, dtype=torch.bool) + indices = scores[..., num_prefix_blocks:].topk(keep_video, dim=-1).indices + num_prefix_blocks + block_map.scatter_(-1, indices, True) + block_map[..., :num_prefix_blocks] = True + block_map[:, :, :num_prefix_blocks, :] = True + return block_map + + +def _build_h3_ordered_q2k_indices( + scores: torch.Tensor, num_prefix_blocks: int, num_video_blocks: int, topk: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Build FlashInfer's ordered sparse metadata without a dense bool map. + + ``map_to_index`` scans each bool-map row from left to right. Sorting the + selected video block IDs and prepending the exempt prefix therefore emits + exactly the same index order and ``-1`` tail, while avoiding both the bool + map allocation and its full-width scan. + """ + if scores.ndim != 4: + raise ValueError(f"scores must be [B, H, Nq, Nkv], got {tuple(scores.shape)}") + if num_prefix_blocks < 0 or num_video_blocks < 0: + raise ValueError("num_prefix_blocks and num_video_blocks must be non-negative") + total_blocks = num_prefix_blocks + num_video_blocks + if scores.shape[-2:] != (total_blocks, total_blocks): + raise ValueError( + f"scores must match the H3 prefix+video block geometry: got " + f"{tuple(scores.shape[-2:])}, expected {(total_blocks, total_blocks)}" + ) + if topk < 0: + raise ValueError(f"topk must be non-negative, got {topk}") + keep_video = min(topk, num_video_blocks) + selected: torch.Tensor | None = None + if keep_video < num_video_blocks and keep_video: + selected = scores[..., num_prefix_blocks:, num_prefix_blocks:].topk(keep_video, dim=-1).indices.to(torch.int32) + selected = selected.sort(dim=-1).values + output_shape = tuple(scores.shape) + q2k_idx = torch.full(output_shape, -1, dtype=torch.int32, device=scores.device) + q2k_num = torch.full(output_shape[:-1], num_prefix_blocks + keep_video, dtype=torch.int32, device=scores.device) + all_indices = torch.arange(total_blocks, dtype=torch.int32, device=scores.device) + if keep_video == num_video_blocks: + q2k_idx.copy_(all_indices) + q2k_num.fill_(total_blocks) + return (q2k_idx, q2k_num) + if num_prefix_blocks: + q2k_idx[..., num_prefix_blocks:, :num_prefix_blocks] = all_indices[:num_prefix_blocks] + if selected is not None: + q2k_idx[..., num_prefix_blocks:, num_prefix_blocks : num_prefix_blocks + keep_video] = ( + selected + num_prefix_blocks + ) + if num_prefix_blocks: + q2k_idx[..., :num_prefix_blocks, :] = all_indices + q2k_num[..., :num_prefix_blocks] = total_blocks + return (q2k_idx, q2k_num) + + +class MiniMaxH3VSAImpl(FastVideoVSAImpl): + def __init__( + self, + num_heads, + head_size, + softmax_scale, + causal=False, + num_kv_heads=None, + prefix="", + qkv_layout=None, + backend_kwargs=None, + **extra_impl_args, + ): + super().__init__( + num_heads, + head_size, + softmax_scale, + causal, + num_kv_heads, + prefix, + qkv_layout, + backend_kwargs, + **extra_impl_args, + ) + self.h3_kernel_backend = os.getenv("VLLM_OMNI_H3_VSA_KERNEL", "fastvideo") + if self.h3_kernel_backend not in {"fastvideo", "flashinfer"}: + raise ValueError("VLLM_OMNI_H3_VSA_KERNEL must be fastvideo or flashinfer") + + h3_effective_tile_shape = H3_VSA_EFFECTIVE_TILE_SHAPE + h3_effective_tile_size = H3_VSA_EFFECTIVE_TILE_SIZE + + def _forward_h3( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata, + aligned_rows: int, + ) -> torch.Tensor: + skip_softmax_threshold = h3_vsa_skip_softmax_threshold_scale_factor() + if skip_softmax_threshold and self.h3_kernel_backend != "flashinfer": + raise RuntimeError( + f"{H3_VSA_SKIP_SOFTMAX_THRESHOLD_ENV}>0 requires FASTVIDEO_VSA h3_kernel_backend='flashinfer'" + ) + layout = _get_h3_layout(attn_metadata) + if layout is None: + raise ValueError("incomplete VSA-H3 layout metadata") + (prefix_segments, video_shape, target_start) = layout + if sum(prefix_segments) != target_start: + raise ValueError(f"VSA-H3 prefix segments sum to {sum(prefix_segments)}, target starts at {target_start}") + expected = target_start + math.prod(video_shape) + if query.shape[1] != expected: + raise ValueError(f"VSA-H3 layout has {expected} rows but attention received {query.shape[1]}") + if query.dtype not in (torch.float16, torch.bfloat16): + raise ValueError(f"VSA-H3 requires fp16/bf16 tensors, got {query.dtype}") + o_bundle = attn_metadata.extra.get(H3_VSA_O_BUNDLE_ACTIVE_KEY, False) + if not isinstance(o_bundle, bool): + raise TypeError(f"{H3_VSA_O_BUNDLE_ACTIVE_KEY} must be a bool") + o_bundle_state = attn_metadata.extra.get(H3_VSA_O_BUNDLE_STATE_KEY) + owner_route: H3VSAOwnerRoutePlan | None = None + if o_bundle: + if not h3_vsa_o_bundle_enabled(): + raise RuntimeError(f"{H3_VSA_O_BUNDLE_ACTIVE_KEY}=True requires {H3_VSA_O_BUNDLE_ENV}=1") + if query.dtype != torch.bfloat16: + raise TypeError("H3 VSA reverse-O bundling currently requires BF16") + if not isinstance(o_bundle_state, dict): + raise RuntimeError("H3 VSA reverse-O bundle mutable state is absent") + if set(o_bundle_state) != {"plan"}: + raise RuntimeError(f"stale H3 VSA reverse-O bundle mutable state: keys={sorted(o_bundle_state)}") + candidate_plan = o_bundle_state["plan"] + if not isinstance(candidate_plan, H3VSAOwnerRoutePlan) or not h3_vsa_owner_route_plan_is_trusted( + candidate_plan + ): + raise RuntimeError("H3 VSA reverse-O bundle route plan is absent or untrusted") + owner_route = candidate_plan + if owner_route.sp_world_size * owner_route.local_rows != aligned_rows: + raise ValueError( + f"H3 VSA reverse-O bundle route does not match aligned attention rows: " + f"route={owner_route.sp_world_size}*{owner_route.local_rows}, attention={aligned_rows}" + ) + gate = _get_gate_compress(attn_metadata) + if o_bundle and gate is not None: + raise RuntimeError("deferred/bundled H3 VSA must retain gate_compress on the original SP shard") + if gate is not None: + if gate.shape[0] != query.shape[0] or gate.shape[2:] != query.shape[2:] or gate.shape[1] < query.shape[1]: + raise ValueError(f"gate_compress shape {gate.shape} cannot cover query shape {query.shape}") + gate = gate[:, : query.shape[1]] + (partition, sizes, non_pad, untile, prefix_blocks, video_blocks) = _get_h3_tile_metadata( + prefix_segments, video_shape, query.device + ) + logical_blocks = int(sizes.numel()) + if owner_route is not None and owner_route.logical_block_count != logical_blocks: + raise ValueError( + f"H3 VSA reverse-O bundle route block count does not match the backend layout: " + f"route={owner_route.logical_block_count}, backend={logical_blocks}" + ) + pair_pad = logical_blocks % 2 if self.h3_kernel_backend == "fastvideo" else 0 + kernel_blocks = logical_blocks + pair_pad + target_shape = (query.shape[0], kernel_blocks * H3_VSA_EFFECTIVE_TILE_SIZE, query.shape[2], query.shape[3]) + use_fused_tile_pack = h3_vsa_fused_tile_pack_enabled() + use_fused_untile = h3_vsa_fused_untile_enabled() + qkv_source_rows: torch.Tensor | None = None + gate_source_rows: torch.Tensor | None = None + if use_fused_tile_pack: + qkv_source_rows = _get_h3_tiled_source_rows(prefix_segments, video_shape, target_shape[1], query.device) + if gate is not None and pair_pad: + gate_source_rows = _get_h3_tiled_source_rows( + prefix_segments, video_shape, logical_blocks * H3_VSA_EFFECTIVE_TILE_SIZE, query.device + ) + else: + gate_source_rows = qkv_source_rows + pack_gate = gate is not None + packed_tensor_count = 3 + int(pack_gate) + fused_inputs = (query, key, value, gate) if pack_gate else (query, key, value) + fused_maps = ( + (qkv_source_rows, qkv_source_rows, qkv_source_rows, gate_source_rows) + if pack_gate + else (qkv_source_rows,) * 3 + ) + use_fused_tile_pack = all( + source_rows is not None and h3_vsa_tile_pack_cuda_supported(tensor, source_rows) + for (tensor, source_rows) in zip(fused_inputs, fused_maps, strict=True) + ) + if use_fused_tile_pack: + logger.info_once( + "FASTVIDEO_VSA H3 tile pack: fused Triton compact->tile64 enabled " + "(compact_rows=%d, padded_rows=%d, tensors=%d)", + query.shape[1], + target_shape[1], + packed_tensor_count, + scope="process", + ) + else: + logger.warning_once( + "%s=1 requested but the H3 input shape/platform is unsupported; using the reference tile pack", + H3_VSA_FUSED_TILE_PACK_ENV, + scope="process", + ) + from vllm_omni.diffusion.models.minimax_h3.attention.overlap import prepared + + overlap_prepared = prepared() + if overlap_prepared is not None: + assert use_fused_tile_pack and gate is None + (q_tiled, k_tiled, q_pool, k_pool, scores, q2k_idx, q2k_num) = overlap_prepared + with _h3_vsa_nvtx_stage("vsa.layout.tile_pack_v"): + v_tiled = h3_vsa_tile_pack(value, qkv_source_rows) + else: + with _h3_vsa_nvtx_stage("vsa.layout.tile_pack"): + if use_fused_tile_pack: + assert qkv_source_rows is not None + q_tiled = h3_vsa_tile_pack(query, qkv_source_rows) + k_tiled = h3_vsa_tile_pack(key, qkv_source_rows) + v_tiled = h3_vsa_tile_pack(value, qkv_source_rows) + else: + assert True + q_tiled = torch.zeros(target_shape, device=query.device, dtype=query.dtype) + k_tiled = torch.zeros_like(q_tiled) + v_tiled = torch.zeros_like(q_tiled) + q_tiled[:, non_pad] = query[:, partition] + k_tiled[:, non_pad] = key[:, partition] + v_tiled[:, non_pad] = value[:, partition] + with _h3_vsa_nvtx_stage("vsa.coarse.pool_qk"): + q_pool = _pool_h3_tiles(q_tiled[:, : logical_blocks * H3_VSA_EFFECTIVE_TILE_SIZE], sizes) + k_pool = _pool_h3_tiles(k_tiled[:, : logical_blocks * H3_VSA_EFFECTIVE_TILE_SIZE], sizes) + with _h3_vsa_nvtx_stage("vsa.coarse.qk_scores"): + scores = torch.matmul(q_pool, k_pool.transpose(-2, -1)) * self.softmax_scale + kernel_sizes = sizes + from vllm_omni.diffusion.models.minimax_h3.attention.overlap import coarse_ready, coarse_scope + + coarse_ticket = coarse_ready(v_tiled, scores) if gate is not None or o_bundle else None + + def compute_coarse(): + with coarse_scope(coarse_ticket): + with _h3_vsa_nvtx_stage("vsa.coarse.pool_v"): + v_pool = _pool_h3_tiles(v_tiled[:, : logical_blocks * H3_VSA_EFFECTIVE_TILE_SIZE], sizes) + with _h3_vsa_nvtx_stage("vsa.coarse.softmax"): + coarse_weights = torch.softmax(scores, dim=-1) + with _h3_vsa_nvtx_stage("vsa.coarse.pv"): + coarse = torch.matmul(coarse_weights, v_pool) + return coarse.permute(0, 2, 1, 3).to(q_tiled.dtype) + + compressed = None + logger.info_once( + "FASTVIDEO_VSA H3 routing: seq_len=%d, prefix_segments=%s, video_shape=%s, " + "prefix_blocks=%d, video_blocks=%d, topk=%d, kernel_blocks=%d, compute_backend=%s", + query.shape[1], + prefix_segments, + video_shape, + prefix_blocks, + video_blocks, + min(self.topk, video_blocks), + kernel_blocks, + self.h3_kernel_backend, + ) + if self.h3_kernel_backend == "flashinfer": + if h3_vsa_direct_q2k_enabled(): + logger.info_once( + "FASTVIDEO_VSA H3 metadata: direct ordered q2k enabled by %s=1", + H3_VSA_DIRECT_Q2K_ENV, + scope="process", + ) + with _h3_vsa_nvtx_stage("vsa.route.topk_q2k"): + if overlap_prepared is None: + (q2k_idx, q2k_num) = _build_h3_ordered_q2k_indices( + scores, prefix_blocks, video_blocks, self.topk + ) + with _h3_vsa_nvtx_stage("vsa.fine.block_sparse_attention"): + fine_op = _flashinfer_h3_vsa_q2k_bshd_op + output = fine_op( + q_tiled.contiguous(), + k_tiled.contiguous(), + v_tiled.contiguous(), + q2k_idx, + q2k_num, + kernel_sizes.contiguous(), + self.softmax_scale, + ) + else: + logger.info_once( + "FASTVIDEO_VSA H3 metadata: legacy bool-map conversion active; set %s=1 to " + "qualify direct ordered q2k", + H3_VSA_DIRECT_Q2K_ENV, + scope="process", + ) + with _h3_vsa_nvtx_stage("vsa.route.topk_q2k"): + block_map = _build_h3_block_map(scores, prefix_blocks, video_blocks, self.topk) + with _h3_vsa_nvtx_stage("vsa.fine.block_sparse_attention"): + output = _flashinfer_h3_vsa_bshd_op( + q_tiled.contiguous(), + k_tiled.contiguous(), + v_tiled.contiguous(), + block_map.contiguous(), + kernel_sizes.contiguous(), + self.softmax_scale, + ) + else: + with _h3_vsa_nvtx_stage("vsa.route.topk_q2k"): + block_map = _build_h3_block_map(scores, prefix_blocks, video_blocks, self.topk) + if pair_pad: + block_map = torch.nn.functional.pad(block_map, (0, 1, 0, 1), value=False) + kernel_sizes = torch.nn.functional.pad(sizes, (0, 1), value=0) + with _h3_vsa_nvtx_stage("vsa.fine.block_sparse_attention"): + output = _fastvideo_h3_vsa_bhsd_op( + q_tiled.contiguous(), + k_tiled.contiguous(), + v_tiled.contiguous(), + block_map.contiguous(), + kernel_sizes.contiguous(), + logical_blocks, + ) + output = output[:, : logical_blocks * H3_VSA_EFFECTIVE_TILE_SIZE] + if gate is not None or o_bundle: + compressed = compute_coarse() + aligned_untile_source_rows: torch.Tensor | None = None + o_input_landing = attn_metadata.extra.get("ulysses_o_input_landing") + if o_input_landing is not None and (not isinstance(o_input_landing, torch.Tensor)): + raise TypeError(f"ulysses_o_input_landing must be a torch.Tensor, got {type(o_input_landing)!r}") + if o_bundle: + assert compressed is not None + assert owner_route is not None + assert isinstance(o_bundle_state, dict) + with _h3_vsa_nvtx_stage("vsa.output.untile"): + aligned_untile_source_rows = _get_h3_aligned_untile_source_rows( + prefix_segments, video_shape, aligned_rows, output.shape[1], output.device + ) + if use_fused_untile: + if not h3_vsa_tile_untile_cuda_supported(output, aligned_untile_source_rows): + raise RuntimeError( + "H3 VSA reverse-O bundling requested with fused untile, but the " + "tile64 output geometry is unsupported" + ) + fine_aligned = h3_vsa_tile_untile(output, aligned_untile_source_rows) + logger.info_once( + "FASTVIDEO_VSA H3 output untile: fused Triton tile64->aligned enabled " + "(valid_rows=%d, aligned_rows=%d, tiled_rows=%d, destination=o_bundle_staging)", + untile.numel(), + aligned_rows, + output.shape[1], + scope="process", + ) + else: + fine_aligned = torch.zeros( + (output.shape[0], aligned_rows, output.shape[2], output.shape[3]), + dtype=output.dtype, + device=output.device, + ) + fine_aligned[:, : untile.numel()].copy_(output[:, untile]) + from vllm_omni.diffusion.models.minimax_h3.attention.overlap import publish_reverse + + if publish_reverse(fine_aligned, compressed, owner_route): + return fine_aligned + raise RuntimeError("H3 reverse-O bundling requires the active model overlap scope") + if use_fused_untile: + aligned_untile_source_rows = _get_h3_aligned_untile_source_rows( + prefix_segments, video_shape, aligned_rows, output.shape[1], output.device + ) + if gate is not None: + assert compressed is not None + if use_fused_tile_pack: + if gate_source_rows is None: + gate_source_rows = _get_h3_tiled_source_rows( + prefix_segments, video_shape, logical_blocks * H3_VSA_EFFECTIVE_TILE_SIZE, query.device + ) + if h3_vsa_tile_pack_cuda_supported(gate, gate_source_rows): + gate_tiled = h3_vsa_tile_pack(gate, gate_source_rows) + else: + gate_tiled = torch.zeros_like(q_tiled[:, : logical_blocks * H3_VSA_EFFECTIVE_TILE_SIZE]) + gate_tiled[:, non_pad] = gate[:, partition] + else: + gate_tiled = torch.zeros_like(q_tiled[:, : logical_blocks * H3_VSA_EFFECTIVE_TILE_SIZE]) + gate_tiled[:, non_pad] = gate[:, partition] + output = ( + output.view( + output.shape[0], logical_blocks, H3_VSA_EFFECTIVE_TILE_SIZE, output.shape[2], output.shape[3] + ) + + compressed.unsqueeze(2) + * gate_tiled.view( + gate_tiled.shape[0], + logical_blocks, + H3_VSA_EFFECTIVE_TILE_SIZE, + gate_tiled.shape[2], + gate_tiled.shape[3], + ) + ).view_as(output) + if use_fused_untile: + assert aligned_untile_source_rows is not None + if h3_vsa_tile_untile_cuda_supported(output, aligned_untile_source_rows): + logger.info_once( + "FASTVIDEO_VSA H3 output untile: fused Triton tile64->aligned enabled " + "(valid_rows=%d, aligned_rows=%d, tiled_rows=%d, destination=%s)", + untile.numel(), + aligned_rows, + output.shape[1], + "registered_reverse_o" if o_input_landing is not None else "functional", + scope="process", + ) + if o_input_landing is not None: + return h3_vsa_tile_untile_out(output, aligned_untile_source_rows, out=o_input_landing) + return h3_vsa_tile_untile(output, aligned_untile_source_rows) + logger.warning_once( + "FASTVIDEO_VSA H3 output untile fallback: %s=1 requested but the shape/platform " + "is unsupported; using the reference tile64->compact->aligned path " + "(valid_rows=%d, aligned_rows=%d, tiled_rows=%d)", + H3_VSA_FUSED_UNTILE_ENV, + untile.numel(), + aligned_rows, + output.shape[1], + scope="process", + ) + return output[:, untile].contiguous() + + def forward_cuda( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + (original_query, original_key, original_value) = (query, key, value) + original_seq_len = query.shape[1] + valid_seq_len = original_seq_len + if attn_metadata is not None and attn_metadata.packed_padding is not None: + valid_seq_len = attn_metadata.packed_padding.q_length + if attn_metadata.packed_padding.kv_length != valid_seq_len: + return self._fallback( + original_query, original_key, original_value, attn_metadata, "packed Q/KV lengths must match" + ) + query = query[:, :valid_seq_len] + key = key[:, :valid_seq_len] + value = value[:, :valid_seq_len] + h3_layout = _get_h3_layout(attn_metadata) + if attn_metadata is not None and h3_layout is not None: + try: + output = self._forward_h3(query, key, value, attn_metadata, original_seq_len) + o_bundle = attn_metadata.extra.get(H3_VSA_O_BUNDLE_ACTIVE_KEY, False) + if o_bundle is True: + state = attn_metadata.extra.get(H3_VSA_O_BUNDLE_STATE_KEY) + plan = state.get("plan") if isinstance(state, dict) else None + if not isinstance(plan, H3VSAOwnerRoutePlan): + raise RuntimeError("H3 VSA reverse-O bundle plan disappeared after producer execution") + from vllm_omni.diffusion.models.minimax_h3.attention.overlap import ACTIVE + + overlap_state = ACTIVE.get() + if overlap_state is not None: + (fine, _coarse, owner_plan) = overlap_state["reverse"] + expected_fine_rows = original_seq_len + if output is not fine or owner_plan is not plan or output.shape[1] != expected_fine_rows: + raise RuntimeError("H3 chunked reverse-O producer lost its fine/route identity") + return output + expected_bundle_rows = plan.sp_world_size * (plan.local_rows + plan.kmax) + if output.shape[1] != expected_bundle_rows: + raise ValueError( + f"H3 VSA reverse-O producer returned the wrong bundle rows: " + f"got={output.shape[1]}, expected={expected_bundle_rows}" + ) + return output + if output.shape[1] == original_seq_len: + return output + if output.shape[1] != valid_seq_len: + raise ValueError( + f"VSA-H3 output must contain either the valid or aligned sequence, got " + f"{output.shape[1]} rows for valid={valid_seq_len}, aligned={original_seq_len}" + ) + restored = torch.zeros(original_query.shape, dtype=output.dtype, device=output.device) + restored[:, :valid_seq_len] = output + return restored + except Exception as exc: + o_bundle = attn_metadata.extra.get(H3_VSA_O_BUNDLE_ACTIVE_KEY, False) + if ( + o_bundle is True + or self.h3_kernel_backend == "flashinfer" + or isinstance(exc, torch.AcceleratorError) + ): + raise + if not self.fallback_on_error: + raise + return self._fallback( + original_query, original_key, original_value, attn_metadata, f"VSA-H3 kernel failed: {exc}" + ) + return super().forward_cuda(original_query, original_key, original_value, attn_metadata) diff --git a/vllm_omni/diffusion/models/minimax_h3/chunked_cpu_output.py b/vllm_omni/diffusion/models/minimax_h3/chunked_cpu_output.py new file mode 100644 index 00000000000..5488b3f3696 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/chunked_cpu_output.py @@ -0,0 +1,688 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Bounded finalized-chunk output pipeline for MiniMax H3. + +The remote H3 video VAE can expose temporally finalized decoder chunks. This +module converts each chunk to the exact uint8 representation used by the +ordinary response path, copies it through a small pinned-host ring, and feeds +one persistent PyAV/libx264 MP4 session. Consequently, rank zero never has to +assemble or transport a full-resolution floating-point video. + +The feature is deliberately opt-in. ``submit_decoded`` also defers failures +until ``finish`` so a rank-zero media failure cannot make it leave PP8 VAE +collectives while peer ranks are still decoding. +""" + +from __future__ import annotations + +import io +import os +import queue +import struct +import threading +import time +from contextlib import suppress +from dataclasses import dataclass +from fractions import Fraction +from typing import Any, cast + +import numpy as np +import torch + +from vllm_omni.platforms import current_omni_platform + +MINIMAX_H3_CHUNKED_CPU_MP4_OUTPUT_ENV = "VLLM_OMNI_MINIMAX_H3_CHUNKED_CPU_MP4_OUTPUT" +MINIMAX_H3_CHUNKED_CPU_MP4_SLOTS_ENV = "VLLM_OMNI_MINIMAX_H3_CHUNKED_CPU_MP4_SLOTS" +MINIMAX_H3_CHUNKED_CPU_MP4_ROUTE = "finalized_chunk_gpu_u8_pinned_d2h_pyav_libx264" +MINIMAX_H3_CHUNKED_CPU_MP4_DEFAULT_CODEC_OPTIONS = { + "preset": "ultrafast", + "threads": "0", +} +_MINIMAX_H3_AAC_GLOBAL_HEADER_FLAG = 1 << 22 +_MINIMAX_H3_AAC_PAYLOAD_MAGIC = b"H3AACV1\0" +_MINIMAX_H3_AAC_PAYLOAD_HEADER = struct.Struct("<8sIIH") +_MINIMAX_H3_AAC_PACKET_HEADER = struct.Struct(" int: + """Return the bounded pinned-ring size (only one or two slots are legal).""" + + if raw is None: + raw = os.environ.get(MINIMAX_H3_CHUNKED_CPU_MP4_SLOTS_ENV, "2") + try: + slots = int(raw) + except (TypeError, ValueError) as exc: + raise ValueError(f"{MINIMAX_H3_CHUNKED_CPU_MP4_SLOTS_ENV} must be 1 or 2, got {raw!r}") from exc + if slots not in (1, 2): + raise ValueError(f"{MINIMAX_H3_CHUNKED_CPU_MP4_SLOTS_ENV} must be 1 or 2, got {slots}") + return slots + + +def validate_chunked_cpu_mp4_runtime() -> None: + """Fail at pipeline startup if PyAV has no usable H.264 encoder.""" + + try: + import av + + av.CodecContext.create("h264", "w") + except BaseException as exc: + raise RuntimeError( + f"{MINIMAX_H3_CHUNKED_CPU_MP4_OUTPUT_ENV}=1 requires PyAV with an H.264/libx264 encoder" + ) from exc + resolve_chunked_cpu_mp4_slots() + + +def resolve_chunked_cpu_mp4_codec_options( + options: dict[str, str] | None, +) -> dict[str, str]: + """Mirror the synchronous video API's default-or-replacement contract.""" + + if options is None: + return dict(MINIMAX_H3_CHUNKED_CPU_MP4_DEFAULT_CODEC_OPTIONS) + if not isinstance(options, dict) or any( + not isinstance(key, str) or not isinstance(value, str) for key, value in options.items() + ): + raise TypeError("video_codec_options must be a string-to-string mapping") + return dict(options) + + +def exact_rgb_u8_bthwc( + frames: torch.Tensor, + *, + height: int, + width: int, +) -> torch.Tensor: + """Crop BCTHW RGB and reproduce the baseline FP32/NumPy uint8 rounding. + + NumPy ``rint`` and ``torch.round`` both use round-to-nearest-even. The + explicit FP32 conversion is important: doing the multiply in FP16 changes + pixels at half-way boundaries. + """ + + if frames.ndim != 5 or int(frames.shape[0]) != 1 or int(frames.shape[1]) != 3: + raise ValueError(f"expected one normalized BCTHW RGB chunk, got shape={tuple(frames.shape)}") + if not frames.dtype.is_floating_point: + raise TypeError(f"expected floating-point RGB chunk, got {frames.dtype}") + if int(frames.shape[-2]) < height or int(frames.shape[-1]) < width: + raise ValueError( + f"decoded RGB chunk is smaller than the requested crop: shape={tuple(frames.shape)}, crop={height}x{width}" + ) + rgb = frames[0, :, :, :height, :width].permute(1, 2, 3, 0) + work = rgb.float() + if not work.is_contiguous(): + work = work.contiguous() + work.clamp_(0.0, 1.0).mul_(255.0).round_() + return work.to(torch.uint8) + + +def _audio_to_numpy(audio: Any | None) -> np.ndarray: + if audio is None: + raise ValueError("MiniMax H3 chunked MP4 output requires stereo audio") + if isinstance(audio, torch.Tensor): + samples = audio.detach().float().cpu().numpy() + else: + samples = np.asarray(audio, dtype=np.float32) + samples = np.squeeze(samples).astype(np.float32, copy=False) + if samples.ndim == 2 and samples.shape[0] > samples.shape[1]: + samples = samples.T + if samples.ndim != 2 or int(samples.shape[0]) != 2: + raise ValueError(f"MiniMax H3 audio must be stereo channel-first PCM, got {samples.shape}") + return np.ascontiguousarray(samples) + + +@dataclass(frozen=True) +class MiniMaxH3EncodedAudioPacket: + data: bytes + pts: int + dts: int + duration: int + time_base: Fraction + + +@dataclass(frozen=True) +class MiniMaxH3EncodedAudio: + """AAC packets encoded independently but muxed in the original order.""" + + sample_rate: int + extradata: bytes + packets: tuple[MiniMaxH3EncodedAudioPacket, ...] + + +def encode_minimax_h3_aac( + audio: Any, + audio_sample_rate: int, +) -> MiniMaxH3EncodedAudio: + """Encode the fixed H3 stereo AAC stream without opening an MP4 container.""" + + import av + + sample_rate = int(audio_sample_rate) + if sample_rate <= 0: + raise ValueError(f"audio_sample_rate must be positive, got {sample_rate}") + samples = _audio_to_numpy(audio) + codec = av.CodecContext.create("aac", "w") + codec.sample_rate = sample_rate + codec.layout = "stereo" + codec.format = "fltp" + codec.flags |= _MINIMAX_H3_AAC_GLOBAL_HEADER_FLAG + codec.open() + + frame = av.AudioFrame.from_ndarray(samples, format="fltp", layout="stereo") + frame.sample_rate = sample_rate + frame.pts = 0 + frame.time_base = Fraction(1, sample_rate) + raw_packets = list(codec.encode(frame)) + raw_packets.extend(codec.encode()) + packets: list[MiniMaxH3EncodedAudioPacket] = [] + for packet in raw_packets: + if packet.pts is None or packet.dts is None or packet.duration is None or packet.time_base is None: + raise RuntimeError("MiniMax H3 AAC encoder returned incomplete packet timing metadata") + packets.append( + MiniMaxH3EncodedAudioPacket( + data=bytes(packet), + pts=int(packet.pts), + dts=int(packet.dts), + duration=int(packet.duration), + time_base=Fraction(packet.time_base), + ) + ) + if not packets: + raise RuntimeError("MiniMax H3 AAC encoder returned no packets") + return MiniMaxH3EncodedAudio( + sample_rate=sample_rate, + extradata=bytes(codec.extradata or b""), + packets=tuple(packets), + ) + + +def serialize_minimax_h3_encoded_audio(audio: MiniMaxH3EncodedAudio) -> bytes: + """Pack encoded AAC and timing metadata into one broadcast payload.""" + + if len(audio.extradata) > 0xFFFF: + raise ValueError("MiniMax H3 AAC extradata exceeds the payload format") + output = io.BytesIO() + output.write( + _MINIMAX_H3_AAC_PAYLOAD_HEADER.pack( + _MINIMAX_H3_AAC_PAYLOAD_MAGIC, + int(audio.sample_rate), + len(audio.packets), + len(audio.extradata), + ) + ) + output.write(audio.extradata) + for packet in audio.packets: + output.write( + _MINIMAX_H3_AAC_PACKET_HEADER.pack( + int(packet.pts), + int(packet.dts), + int(packet.duration), + int(packet.time_base.numerator), + int(packet.time_base.denominator), + len(packet.data), + ) + ) + output.write(packet.data) + return output.getvalue() + + +def deserialize_minimax_h3_encoded_audio(payload: bytes) -> MiniMaxH3EncodedAudio: + """Validate and unpack one rank-local AAC broadcast payload.""" + + view = memoryview(payload) + header_size = _MINIMAX_H3_AAC_PAYLOAD_HEADER.size + if len(view) < header_size: + raise ValueError("MiniMax H3 AAC payload is truncated before its header") + magic, sample_rate, packet_count, extradata_size = _MINIMAX_H3_AAC_PAYLOAD_HEADER.unpack_from(view) + if magic != _MINIMAX_H3_AAC_PAYLOAD_MAGIC or sample_rate <= 0 or packet_count <= 0: + raise ValueError("MiniMax H3 AAC payload header is invalid") + offset = header_size + if offset + extradata_size > len(view): + raise ValueError("MiniMax H3 AAC payload is truncated in codec extradata") + extradata = bytes(view[offset : offset + extradata_size]) + offset += extradata_size + packets: list[MiniMaxH3EncodedAudioPacket] = [] + for _ in range(packet_count): + packet_header_size = _MINIMAX_H3_AAC_PACKET_HEADER.size + if offset + packet_header_size > len(view): + raise ValueError("MiniMax H3 AAC payload is truncated in packet metadata") + pts, dts, duration, numerator, denominator, data_size = _MINIMAX_H3_AAC_PACKET_HEADER.unpack_from(view, offset) + offset += packet_header_size + if duration < 0 or denominator == 0 or offset + data_size > len(view): + raise ValueError("MiniMax H3 AAC packet metadata is invalid") + packets.append( + MiniMaxH3EncodedAudioPacket( + data=bytes(view[offset : offset + data_size]), + pts=pts, + dts=dts, + duration=duration, + time_base=Fraction(numerator, denominator), + ) + ) + offset += data_size + if offset != len(view): + raise ValueError("MiniMax H3 AAC payload contains trailing bytes") + return MiniMaxH3EncodedAudio( + sample_rate=sample_rate, + extradata=extradata, + packets=tuple(packets), + ) + + +@dataclass(frozen=True) +class _PinnedSlot: + tensor: torch.Tensor + index: int + + +@dataclass +class _D2HItem: + slot: _PinnedSlot + frame_count: int + first_frame: int + ready: torch.cuda.Event + gpu_owner: torch.Tensor + + +class MiniMaxH3ChunkedCpuMp4Output: + """Finalized VAE chunk -> exact GPU uint8 -> pinned D2H -> MP4 sink.""" + + def __init__( + self, + *, + width: int, + height: int, + fps: int, + audio_sample_rate: int, + max_chunk_frames: int, + device: torch.device, + video_codec_options: dict[str, str] | None = None, + queue_slots: int | None = None, + ) -> None: + if device.type != "cuda": + raise ValueError(f"MiniMax H3 chunked CPU MP4 output requires a CUDA device, got {device}") + if device.index is None: + device = torch.device("cuda", torch.accelerator.current_device_index()) + if ( + min( + int(width), + int(height), + int(fps), + int(audio_sample_rate), + int(max_chunk_frames), + ) + <= 0 + ): + raise ValueError("width, height, fps, audio_sample_rate, and max_chunk_frames must be positive") + self.width = int(width) + self.height = int(height) + self.fps = int(fps) + self.audio_sample_rate = int(audio_sample_rate) + self.device = device + self._queue_slots = ( + resolve_chunked_cpu_mp4_slots() if queue_slots is None else resolve_chunked_cpu_mp4_slots(str(queue_slots)) + ) + self._video_codec_options = resolve_chunked_cpu_mp4_codec_options(video_codec_options) + self._copy_stream = torch.get_device_module().Stream(device=device) + self._free: queue.Queue[_PinnedSlot] = queue.Queue(maxsize=self._queue_slots) + self._work: queue.Queue[_D2HItem | None] = queue.Queue(maxsize=self._queue_slots) + self._slot_frame_capacity = int(max_chunk_frames) + for index in range(self._queue_slots): + tensor = torch.empty( + (self._slot_frame_capacity, self.height, self.width, 3), + dtype=torch.uint8, + device="cpu", + pin_memory=True, + ) + self._free.put(_PinnedSlot(tensor=tensor, index=index)) + + self._error: BaseException | None = None + self._error_lock = threading.Lock() + self._audio: np.ndarray | MiniMaxH3EncodedAudio | None = None + self._result: bytes | None = None + self._closed = False + self._submitted_frames = 0 + self._expected_frames: int | None = None + self._next_frame = 0 + self._raw_rgb_bytes = 0 + self._max_inflight = 0 + self._inflight = 0 + self._pipeline_start: float | None = None + self._producer_backpressure_s = 0.0 + self._gpu_prepare_enqueue_s = 0.0 + self._d2h_wait_s = 0.0 + self._video_encode_mux_s = 0.0 + self._audio_mux_finalize_s = 0.0 + + self._thread = threading.Thread( + target=self._worker, + name="minimax-h3-pinned-pyav", + daemon=True, + ) + self._thread.start() + + def _set_error(self, error: BaseException) -> None: + with self._error_lock: + if self._error is None: + self._error = error + + def _get_error(self) -> BaseException | None: + with self._error_lock: + return self._error + + def _ensure_slots(self, frame_count: int) -> None: + if frame_count < 1: + raise ValueError("finalized VAE chunks must contain at least one frame") + if frame_count > self._slot_frame_capacity: + raise ValueError( + "finalized VAE chunk exceeds the checkpoint-derived pinned-ring " + f"capacity: frames={frame_count}, capacity={self._slot_frame_capacity}" + ) + + def submit_decoded( + self, + decoded: torch.Tensor, + first_frame: int, + total_frames: int, + processor: Any, + ) -> None: + """Queue one chunk without ever escaping an error into VAE collectives.""" + + if self._get_error() is not None: + return + slot: _PinnedSlot | None = None + ready: torch.cuda.Event | None = None + queued = False + inflight_accounted = False + try: + if self._closed: + raise RuntimeError("cannot submit after the chunked MP4 sink is closed") + if self._pipeline_start is None: + self._pipeline_start = time.perf_counter() + first_frame = int(first_frame) + total_frames = int(total_frames) + if first_frame != self._next_frame: + raise RuntimeError( + f"non-contiguous finalized VAE chunk: first_frame={first_frame}, expected={self._next_frame}" + ) + if total_frames < 1 or first_frame >= total_frames: + raise ValueError(f"invalid finalized VAE range {first_frame}/{total_frames}") + if self._expected_frames not in (None, total_frames): + raise RuntimeError( + f"finalized VAE total changed between chunks: {self._expected_frames} -> {total_frames}" + ) + self._expected_frames = total_frames + + submit_start = time.perf_counter() + if decoded.ndim != 5: + raise ValueError(f"expected a rank-5 finalized decoder chunk, got shape={tuple(decoded.shape)}") + if decoded.device.type != "cuda" or decoded.device.index != self.device.index: + raise ValueError( + "finalized decoder chunks must remain on the sink CUDA device: " + f"chunk={decoded.device}, sink={self.device}" + ) + frame_count = int(decoded.shape[2]) + self._ensure_slots(frame_count) + wait_start = time.perf_counter() + slot = self._free.get() + self._producer_backpressure_s += time.perf_counter() - wait_start + + # Record the producer before switching streams. The callback owns + # ``decoded`` only until this method returns. + producer_ready = torch.get_device_module().Event() + producer_ready.record(torch.get_device_module().current_stream(device=decoded.device)) + with torch.get_device_module().stream(self._copy_stream): + self._copy_stream.wait_event(producer_ready) + decoded.record_stream(self._copy_stream) + frames = processor.revert_tensor(decoded) + if frames.ndim == 4: + frames = frames.unsqueeze(0).transpose(1, 2) + if frames.device != decoded.device: + raise ValueError( + "MiniMax H3 processor moved a finalized chunk off its CUDA " + f"device: {decoded.device} -> {frames.device}" + ) + rgb_u8 = exact_rgb_u8_bthwc( + frames, + height=self.height, + width=self.width, + ) + if int(rgb_u8.shape[0]) != frame_count: + raise RuntimeError( + f"MiniMax H3 processor changed finalized chunk length: {frame_count} -> {int(rgb_u8.shape[0])}" + ) + if first_frame + frame_count > total_frames: + raise ValueError( + f"finalized VAE chunk exceeds advertised total: {first_frame}+{frame_count}>{total_frames}" + ) + + ready = torch.get_device_module().Event() + with torch.get_device_module().stream(self._copy_stream): + slot.tensor[:frame_count].copy_(rgb_u8, non_blocking=True) + frames.record_stream(self._copy_stream) + rgb_u8.record_stream(self._copy_stream) + ready.record(self._copy_stream) + self._inflight += 1 + inflight_accounted = True + self._max_inflight = max(self._max_inflight, self._inflight) + self._work.put( + _D2HItem( + slot=slot, + frame_count=frame_count, + first_frame=first_frame, + ready=ready, + gpu_owner=rgb_u8, + ) + ) + queued = True + self._submitted_frames += frame_count + self._next_frame += frame_count + self._raw_rgb_bytes += frame_count * self.height * self.width * 3 + self._gpu_prepare_enqueue_s += time.perf_counter() - submit_start + except BaseException as exc: + if slot is not None and not queued: + if ready is not None: + with suppress(BaseException): + ready.synchronize() + if inflight_accounted: + self._inflight -= 1 + self._free.put(slot) + self._set_error(exc) + + def _worker(self) -> None: + stop_seen = False + try: + import av + + current_omni_platform.set_device(self.device) + buffer = io.BytesIO() + with cast(Any, av.open(buffer, mode="w", format="mp4")) as container: + video_stream = cast( + Any, + container.add_stream( + "h264", + rate=Fraction(self.fps).limit_denominator(10_000), + ), + ) + video_stream.width = self.width + video_stream.height = self.height + video_stream.pix_fmt = "yuv420p" + options: dict[str, object] = {"crf": "18"} + options.update(self._video_codec_options) + video_stream.options = options + + audio_stream = cast( + Any, + container.add_stream("aac", rate=self.audio_sample_rate), + ) + # H3 always emits stereo audio. Set the layout before video + # packets can cause PyAV/FFmpeg to open the codec context; + # recent PyAV versions reject layout changes after that point. + audio_stream.layout = "stereo" + + while True: + item = self._work.get() + try: + if item is None: + stop_seen = True + break + wait_start = time.perf_counter() + item.ready.synchronize() + self._d2h_wait_s += time.perf_counter() - wait_start + array = item.slot.tensor[: item.frame_count].numpy() + encode_start = time.perf_counter() + for frame_data in array: + frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24") + for packet in video_stream.encode(frame): + container.mux(packet) + self._video_encode_mux_s += time.perf_counter() - encode_start + finally: + if item is not None: + self._inflight -= 1 + self._free.put(item.slot) + self._work.task_done() + + finalize_start = time.perf_counter() + for packet in video_stream.encode(): + container.mux(packet) + audio = self._audio + if isinstance(audio, MiniMaxH3EncodedAudio): + if audio.sample_rate != self.audio_sample_rate: + raise ValueError( + "MiniMax H3 encoded audio sample rate changed during output: " + f"{audio.sample_rate}/{self.audio_sample_rate}" + ) + if bytes(audio_stream.codec_context.extradata or b"") != audio.extradata: + raise RuntimeError("MiniMax H3 independently encoded AAC extradata does not match MP4 stream") + for encoded_packet in audio.packets: + packet = av.Packet(encoded_packet.data) + packet.pts = encoded_packet.pts + packet.dts = encoded_packet.dts + packet.duration = encoded_packet.duration + packet.time_base = encoded_packet.time_base + packet.stream = audio_stream + container.mux(packet) + elif audio is not None: + audio_frame = av.AudioFrame.from_ndarray( + audio, + format="fltp", + layout="stereo", + ) + audio_frame.sample_rate = self.audio_sample_rate + audio_frame.pts = 0 + audio_frame.time_base = Fraction(1, self.audio_sample_rate) + for packet in audio_stream.encode(audio_frame): + container.mux(packet) + for packet in audio_stream.encode(): + container.mux(packet) + self._audio_mux_finalize_s += time.perf_counter() - finalize_start + self._result = buffer.getvalue() + except BaseException as exc: + self._set_error(exc) + # Return every queued slot so a producer blocked by bounded + # backpressure cannot strand PP8 peers in a later callback. Once + # the sentinel was consumed, however, finalize/close failures + # must return immediately: waiting for a second sentinel would + # deadlock ``finish`` in its thread join. + if not stop_seen: + while True: + item = self._work.get() + self._work.task_done() + if item is None: + break + # A media-thread failure can race with a D2H submitted by + # the callback. Keep the pinned destination alive until + # that copy is done before returning its slot. + with suppress(BaseException): + item.ready.synchronize() + self._inflight -= 1 + self._free.put(item.slot) + + def finish( + self, + audio: Any | None, + audio_sample_rate: int, + ) -> tuple[bytes, dict[str, float]]: + """Flush the persistent codec session and return only final MP4 bytes.""" + + if self._closed: + raise RuntimeError("MiniMax H3 chunked CPU MP4 sink is already closed") + if int(audio_sample_rate) != self.audio_sample_rate: + self._closed = True + self._work.put(None) + self._thread.join() + raise ValueError( + f"MiniMax H3 audio sample rate changed during output: {audio_sample_rate}/{self.audio_sample_rate}" + ) + self._closed = True + audio_start = time.perf_counter() + if isinstance(audio, MiniMaxH3EncodedAudio): + self._audio = audio + else: + try: + self._audio = _audio_to_numpy(audio) + except BaseException: + self._work.put(None) + self._thread.join() + raise + audio_prepare_s = time.perf_counter() - audio_start + self._work.put(None) + self._thread.join() + + error = self._get_error() + if error is not None: + raise RuntimeError("MiniMax H3 chunked CPU MP4 pipeline failed") from error + if self._result is None: + raise RuntimeError("MiniMax H3 chunked CPU MP4 pipeline produced no MP4") + if self._expected_frames is None or self._submitted_frames != self._expected_frames: + raise RuntimeError( + "MiniMax H3 chunked CPU MP4 frame count mismatch: " + f"submitted={self._submitted_frames}, expected={self._expected_frames}" + ) + pipeline_s = 0.0 if self._pipeline_start is None else time.perf_counter() - self._pipeline_start + return self._result, { + "chunked_cpu_mp4_frames": float(self._submitted_frames), + "chunked_cpu_mp4_rgb_u8_bytes": float(self._raw_rgb_bytes), + "chunked_cpu_mp4_bytes": float(len(self._result)), + "chunked_cpu_mp4_queue_slots": float(self._queue_slots), + "chunked_cpu_mp4_max_inflight": float(self._max_inflight), + "chunked_cpu_mp4_producer_backpressure_s": self._producer_backpressure_s, + "chunked_cpu_mp4_gpu_prepare_enqueue_s": self._gpu_prepare_enqueue_s, + "chunked_cpu_mp4_d2h_wait_s": self._d2h_wait_s, + "chunked_cpu_mp4_video_encode_mux_s": self._video_encode_mux_s, + "chunked_cpu_mp4_audio_prepare_s": audio_prepare_s, + "chunked_cpu_mp4_audio_mux_finalize_s": self._audio_mux_finalize_s, + "chunked_cpu_mp4_audio_preencoded": float(isinstance(self._audio, MiniMaxH3EncodedAudio)), + "chunked_cpu_mp4_audio_packet_bytes": float( + sum(len(packet.data) for packet in self._audio.packets) + if isinstance(self._audio, MiniMaxH3EncodedAudio) + else 0 + ), + "chunked_cpu_mp4_pipeline_s": pipeline_s, + } + + def abort(self) -> None: + """Drain and stop the worker without masking an upstream exception.""" + + if self._closed: + return + self._closed = True + self._work.put(None) + self._thread.join() + + +__all__ = [ + "MINIMAX_H3_CHUNKED_CPU_MP4_OUTPUT_ENV", + "MINIMAX_H3_CHUNKED_CPU_MP4_DEFAULT_CODEC_OPTIONS", + "MINIMAX_H3_CHUNKED_CPU_MP4_ROUTE", + "MINIMAX_H3_CHUNKED_CPU_MP4_SLOTS_ENV", + "MiniMaxH3ChunkedCpuMp4Output", + "MiniMaxH3EncodedAudio", + "MiniMaxH3EncodedAudioPacket", + "deserialize_minimax_h3_encoded_audio", + "encode_minimax_h3_aac", + "exact_rgb_u8_bthwc", + "resolve_chunked_cpu_mp4_slots", + "resolve_chunked_cpu_mp4_codec_options", + "serialize_minimax_h3_encoded_audio", + "validate_chunked_cpu_mp4_runtime", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/fasth3.py b/vllm_omni/diffusion/models/minimax_h3/fasth3.py index c4860c57ec4..548179d4c71 100644 --- a/vllm_omni/diffusion/models/minimax_h3/fasth3.py +++ b/vllm_omni/diffusion/models/minimax_h3/fasth3.py @@ -342,6 +342,8 @@ def __init__( self._injections = dict(injections or {}) self._injected: set[str] = set() self._applied: set[str] = set() + self._sidecar_satisfied: set[str] = set() + self._stream_started = False self._device: torch.device | None = None @property @@ -506,6 +508,8 @@ def fuse(self, name: str, weight: torch.Tensor) -> torch.Tensor: patch = self._patches.get(name) if patch is None: return weight + if name in self._sidecar_satisfied: + return weight self._applied.add(name) device = self._compute_device(weight) @@ -545,11 +549,12 @@ def fuse(self, name: str, weight: torch.Tensor) -> torch.Tensor: def apply(self, weights: Iterable[tuple[str, torch.Tensor]]) -> Iterator[tuple[str, torch.Tensor]]: """Fuse every streamed checkpoint tensor on its way into the model.""" - if self._applied: + if self._stream_started or (self._applied - self._sidecar_satisfied): # ``validate_fully_applied`` released the deltas, so a second stream # would fuse nothing and then pass its own completeness check: the # server would serve base H3 weights on the student's ladder. raise FastH3AdapterError(f"{self._source} has already been fused into this checkpoint") + self._stream_started = True for name, weight in weights: if name in self._injections: raise FastH3AdapterError( @@ -563,6 +568,26 @@ def apply(self, weights: Iterable[tuple[str, torch.Tensor]]) -> Iterator[tuple[s self._injected.add(name) yield name, weight + def mark_sidecar_satisfied(self, names: Iterable[str]) -> None: + """Mark adapter edits supplied by a verified AdaLN sidecar. + + The cache owns only block/final AdaLN projections. Its runtime binding + verifies that it was built from this adapter and exact four-step + schedule before exposing the names. Keep this method strict as the last + boundary: a sidecar may satisfy only edits this adapter really carries, + and it must be bound before checkpoint streaming starts. + """ + if self._stream_started: + raise FastH3AdapterError("FastH3 AdaLN sidecar must be bound before checkpoint weights start streaming") + satisfied = frozenset(str(name) for name in names) + unknown = sorted(satisfied - set(self._patches)) + if unknown: + raise FastH3AdapterError( + f"FastH3 AdaLN sidecar claims parameters this adapter does not edit: {unknown[:5]}" + ) + self._sidecar_satisfied.update(satisfied) + self._applied.update(satisfied) + def validate_fully_applied(self, loaded: Iterable[str] | None = None) -> None: """Close the fusion: every edit must have met its parameter. diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index 3abc9047639..685c9a6d101 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -10,6 +10,7 @@ from __future__ import annotations import math +import os from collections.abc import Iterable, Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -33,22 +34,28 @@ VideoTokenLayout, ) from vllm_omni.diffusion.attention.layer import Attention -from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import ( - indexed_gate, - indexed_gate_rms_norm_scale_shift, - indexed_scale_shift_, - rms_norm_indexed_scale_shift, -) from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, SequenceParallelOutput, ) from vllm_omni.diffusion.layers.activation import SiluAndMul -from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope +from vllm_omni.diffusion.layers.fused_qk_norm_rope import ( + fused_qk_norm_rope, + fused_qk_norm_rope_out, +) +from vllm_omni.diffusion.layers.indexed_modulation import ( + indexed_gate, + indexed_gate_rms_norm_scale_shift, + indexed_scale_shift_, + rms_norm_indexed_scale_shift, +) from vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract +from vllm_omni.diffusion.models.minimax_h3.ops.attention.layout import ( + H3_VSA_ATTENTION_ACTIVE_KEY, +) from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: @@ -61,6 +68,33 @@ logger = init_logger(__name__) +def _validate_adaln_cache_quant_config( + quant_config: QuantizationConfig | None, +) -> None: + """Allow only the validated BF16-to-FP8 online quantization path. + + Cached AdaLN projections are absent from the module tree. Online + ``Fp8Config`` quantizes the remaining BF16 linears after their (possibly + FastH3-fused) weights are loaded, so it does not require serialized AdaLN + weights or scales. Every other quantization format remains fail-closed. + """ + if quant_config is None: + return + + # Keep this import local: constructing an H3 model without an AdaLN cache + # should not eagerly import vLLM's full FP8 implementation. + from vllm.model_executor.layers.quantization.fp8 import Fp8Config + + if type(quant_config) is Fp8Config and quant_config.is_checkpoint_fp8_serialized is False: + return + + raise ValueError( + "MiniMax H3 AdaLN cache supports only unquantized weights or upstream " + "online FP8 from BF16 via " + "Fp8Config(is_checkpoint_fp8_serialized=False)" + ) + + # Packed multi-request forwards require the attention backend to actually # consume cu_seqlens as a block-diagonal plan (not a padding-mask rebuild that # spans the full packed row). The pipeline gates on this capability before @@ -422,13 +456,24 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.out_proj", ) - # VSA compression gate. A FastH3 VSA artifact assigns this projection - # with ``.set_weight``; the dense path never builds it, so the module is - # created only once the loader knows a VSA artifact is coming. + # FastH3 VSA compression gate. Dense H3 does not own this parameter; + # the module is materialized only after the adapter contract identifies + # a sparse student and before checkpoint loading begins. self.to_gate_compress: ColumnParallelLinear | None = None self._gate_hidden_size = arch.hidden_size + self._h3_ulysses_overlap = os.getenv("VLLM_OMNI_H3_ATTENTION_OVERLAP", "0") == "1" and prefix.startswith( + "blocks." + ) + self._h3_lossless_gate = self._h3_ulysses_overlap and os.getenv("VLLM_OMNI_H3_LOSSLESS_GATE", "0") == "1" + from vllm_omni.diffusion.models.minimax_h3.attention.qkv_overlap import configured_mode + + self._h3_vsplit_mode = configured_mode(prefix) + self._h3_vsplit_layer_index = int(prefix.split(".")[1]) if self._h3_vsplit_mode != "off" else -1 self._gate_quant_config = quant_config self._gate_prefix = f"{prefix}.to_gate_compress" + from .attention.parallel import configure_parallel_attention + from .attention.vsa import MiniMaxH3VSAImpl + self.attention = Attention( num_heads=self.num_heads, num_kv_heads=self.num_kv_heads, @@ -441,7 +486,20 @@ def __init__( role_category=role_category, skip_sequence_parallel=skip_sequence_parallel, prefix=prefix, + impl_overrides={"FASTVIDEO_VSA": MiniMaxH3VSAImpl}, ) + configure_parallel_attention(self.attention) + # Static and strictly opt-in: only FlashInfer PCIe Ulysses advertises + # registered producer output. Keeping this as a plain bool lets the + # compiled forward specialize without inspecting process-group state. + self._use_qk_input_landing = self.attention.supports_qk_input_landing + + def _project_attention_chunk(self, out: torch.Tensor) -> torch.Tensor: + """Apply the unchanged H3 output projection to one gathered chunk.""" + # Project this TP rank's head shard using the model's existing linear. + out = out.reshape(-1, self.num_heads * self.head_dim) + out, _ = self.out_proj(out) + return out def enable_vsa_gate(self) -> None: """Build the VSA compression gate this attention would otherwise lack. @@ -490,14 +548,77 @@ def _run_packed_attention( video_layout: VideoTokenLayout | None = None, vsa_prefix_segments: tuple[int, ...] = (), gate_compress: torch.Tensor | None = None, + lossless_gate_input: torch.Tensor | None = None, + qk_norm_rope: tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + float, + ] + | None = None, + vsplit_input: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> torch.Tensor: """Run packed attention as a small eager island. The scalar packed-layout metadata and backend-specific attention kernels are intentionally opaque to Dynamo. Keeping this boundary - narrow lets regional compile fuse projections, norms, RoPE, and the - surrounding DiT block without repeated graph breaks. + narrow lets regional compile fuse the surrounding DiT block without + repeated graph breaks. The opt-in PCIe path also performs Q/K norm and + RoPE here so their Triton producer can target registered input storage. """ + backend = getattr(self.attention, "attn_backend", None) + get_backend_name = getattr(backend, "get_name", None) + h3_vsa_attention_active = bool( + callable(get_backend_name) + and get_backend_name() == "FASTVIDEO_VSA" + and video_layout is not None + and any(span.role == "target" for span in video_layout.video_spans) + ) + if self._h3_vsplit_mode != "off": + if ( + vsplit_input is None + or not h3_vsa_attention_active + or (not self._h3_ulysses_overlap) + or (not self._use_qk_input_landing) + or (qk_norm_rope is None) + or (lossless_gate_input is None) + or (gate_compress is not None) + or (num_requests != 1) + ): + raise RuntimeError("H3 VSPLIT requires the exact BF16 v3 producer-direct inference path") + gate_batched = gate_compress.unsqueeze(0) if gate_compress is not None else None + if qk_norm_rope is not None: + q_weight, k_weight, rope_table, eps = qk_norm_rope + landings = self.attention.prepare_qk_input_landings( + q.unsqueeze(0), + k.unsqueeze(0), + ) + if landings is None: + if self._h3_vsplit_mode != "off": + raise RuntimeError("H3 VSPLIT cannot fall back from the qualified QK landings") + # Optional PCIe cold-start failure with REQUIRE_RDMA=0 keeps + # correctness by producing ordinary tensors for NCCL. + q, k = fused_qk_norm_rope( + q, + k, + q_weight, + k_weight, + rope_table, + eps, + ) + else: + q_landing, k_landing = landings + q, k = fused_qk_norm_rope_out( + q, + k, + q_weight, + k_weight, + rope_table, + eps, + q_out=q_landing.squeeze(0), + k_out=k_landing.squeeze(0), + ) + # max_seqlen is already the longest packed document length. Do not read # the CUDA cu_seqlens scalars here: this function runs once per layer # and .item() would serialize every attention launch. ``num_requests`` @@ -568,25 +689,43 @@ def _run_packed_attention( # (see MINIMAX_H3_LASER_INPUT_SCALE). Ignored by every other # backend/path. "laser_input_scale": MINIMAX_H3_LASER_INPUT_SCALE, - # Present only for a VSA artifact; the VSA backend reads it as - # the learned compression gate and every other backend ignores it. - **({"gate_compress": gate_compress.unsqueeze(0)} if gate_compress is not None else {}), - # FastH3 uses segment-pure prefix chunks. The target video and - # its true 3-D shape remain in the shared typed video layout. + # Only FastH3 VSA consumes these entries. The gate has the + # same local BSHD layout as Q here and follows Q/K/V through + # the Ulysses sequence-to-head exchange. + **({H3_VSA_ATTENTION_ACTIVE_KEY: True} if h3_vsa_attention_active else {}), + **({"gate_compress": gate_batched} if gate_batched is not None else {}), **( {"vsa_h3_prefix_segments": vsa_prefix_segments} - if gate_compress is not None and video_layout is not None and video_layout.video_spans + if (gate_batched is not None or lossless_gate_input is not None) + and video_layout is not None + and video_layout.video_spans else {} ), }, video_layout=video_layout, ) - return self.attention( - q.unsqueeze(0), - k.unsqueeze(0), - v.unsqueeze(0), - metadata, - ).squeeze(0) + q_batched = q.unsqueeze(0) + k_batched = k.unsqueeze(0) + v_batched = v.unsqueeze(0) + if self._h3_ulysses_overlap: + if not h3_vsa_attention_active: + raise RuntimeError("H3 overlap requires main-block BF16-wire VSA without query chunking") + from vllm_omni.diffusion.models.minimax_h3.attention.overlap import run + + return run( + self.attention, + q_batched, + k_batched, + v_batched, + metadata, + self._project_attention_chunk, + gate_input=lossless_gate_input, + gate_projector=self.to_gate_compress, + vsplit_input=vsplit_input if self._h3_vsplit_mode == "split" else None, + vsplit_projection=self.qkv_proj if self._h3_vsplit_mode == "split" else None, + vsplit_layer_index=self._h3_vsplit_layer_index, + ).squeeze(0) + return self.attention(q_batched, k_batched, v_batched, metadata).squeeze(0) def forward( self, @@ -613,16 +752,44 @@ def forward( 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) + if self._h3_vsplit_mode == "split": + from .quantization import project_split_qk + + qk, quantized_x, activation_scale = project_split_qk(self.qkv_proj, x) + q, k = qk.split([q_size, kv_size], dim=-1) + # Metadata only. The guarded Ulysses before_v hook replaces this + # alias with the independently projected V before any V data read. + v = q + else: + qkv, _ = self.qkv_proj(x) + 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) + gate_compress = None + lossless_gate_input = None + if self.to_gate_compress is not None: + if self._h3_lossless_gate: + lossless_gate_input = x + else: + gate_result = self.to_gate_compress(x) + gate_compress = gate_result[0] if isinstance(gate_result, tuple) else gate_result + gate_compress = gate_compress.view(total, self.num_heads, self.head_dim) if rope_table is None: q = self.q_norm(q) k = self.k_norm(k) + qk_norm_rope = None + elif self._use_qk_input_landing: + # Defer fused norm+RoPE into the existing eager attention island, + # where it may write directly into registered RDMA source storage. + qk_norm_rope = ( + self.q_norm.weight, + self.k_norm.weight, + rope_table, + self.q_norm.variance_epsilon, + ) else: q, k = fused_qk_norm_rope( q, @@ -632,15 +799,7 @@ def forward( rope_table, self.q_norm.variance_epsilon, ) - - # The gate is projected from the same local rows as Q. Pure Ulysses - # reshards it alongside Q/K/V in UlyssesParallelAttention so each VSA - # rank receives the full sequence for its local head shard. - gate_compress = None - if self.to_gate_compress is not None: - gate_result = self.to_gate_compress(x) - gate_compress = gate_result[0] if isinstance(gate_result, tuple) else gate_result - gate_compress = gate_compress.view(total, self.num_heads, self.head_dim) + qk_norm_rope = None # Each request contributes a document for its rows plus one for any # nonempty alignment padding. Local/Ulysses backends unpad it, while @@ -659,7 +818,12 @@ def forward( video_layout=video_layout, vsa_prefix_segments=vsa_prefix_segments, gate_compress=gate_compress, + lossless_gate_input=lossless_gate_input, + qk_norm_rope=qk_norm_rope, + vsplit_input=(quantized_x, activation_scale, qk) if self._h3_vsplit_mode == "split" else None, ) + if self._h3_ulysses_overlap: + return out out = out.reshape(total, self.num_heads * self.head_dim) out, _ = self.out_proj(out) return out @@ -698,6 +862,10 @@ def __init__( def forward(self, x: torch.Tensor) -> torch.Tensor: hidden, _ = self.fc1(x) + if getattr(self, "_h3_swiglu_mxfp8_fused", False): + from .quantization import fused_swiglu_fc2 + + return fused_swiglu_fc2(self.fc2, hidden) hidden = self.act_fn(hidden) out, _ = self.fc2(hidden) return out @@ -839,6 +1007,7 @@ def __init__( quant_config: QuantizationConfig | None, *, prefix: str, + use_adaln_cache: bool = False, ) -> None: super().__init__() self.norm1 = _norm(arch.hidden_size, eps=arch.norm_eps) @@ -855,13 +1024,17 @@ def __init__( quant_config, prefix=f"{prefix}.mlp", ) - self.adaln_proj = MiniMaxH3AdalnProj( - arch, - arch.adaln_out_features, - quant_config, - expand_ratio=6, - modality_num=MINIMAX_H3_ADALN_MODALITY_NUM, - prefix=f"{prefix}.adaln_proj", + self.adaln_proj = ( + None + if use_adaln_cache + else MiniMaxH3AdalnProj( + arch, + arch.adaln_out_features, + quant_config, + expand_ratio=6, + modality_num=MINIMAX_H3_ADALN_MODALITY_NUM, + prefix=f"{prefix}.adaln_proj", + ) ) def forward( @@ -878,6 +1051,7 @@ def forward( sp_seq_lens: list[int] | None = None, video_layout: VideoTokenLayout | None = None, vsa_prefix_segments: tuple[int, ...] = (), + adaln_params: tuple[torch.Tensor, ...] | None = None, ) -> torch.Tensor: """x: [T, H]; t_emb: [M, t_dim]; combined_indices: [T] (= inverse_indices * modality_num + token_tags.clamp(min=0)). @@ -886,6 +1060,10 @@ def forward( norm1 -> scale/shift -> attention -> gated residual, followed by norm2 -> scale/shift -> MLP -> gated residual. """ + if adaln_params is None: + if self.adaln_proj is None: + raise RuntimeError("MiniMax H3 AdaLN cache parameters are required") + adaln_params = self.adaln_proj(t_emb) ( shift_msa, scale_msa, @@ -893,7 +1071,7 @@ def forward( shift_mlp, scale_mlp, gate_mlp, - ) = self.adaln_proj(t_emb) + ) = adaln_params residual = x h = rms_norm_indexed_scale_shift( @@ -937,17 +1115,22 @@ def __init__( quant_config: QuantizationConfig | None, *, prefix: str, + use_adaln_cache: bool = False, ) -> 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, - prefix=f"{prefix}.adaln_proj", + self.adaln_proj = ( + None + if use_adaln_cache + else MiniMaxH3AdalnProj( + arch, + arch.final_adaln_out_features, + quant_config, + expand_ratio=2, + modality_num=1, + prefix=f"{prefix}.adaln_proj", + ) ) self.video_out = ColumnParallelLinear( arch.hidden_size, @@ -974,13 +1157,18 @@ def forward( *, t_emb: torch.Tensor, inverse_indices: torch.Tensor, + adaln_params: tuple[torch.Tensor, ...] | None = None, ) -> 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) + if adaln_params is None: + if self.adaln_proj is None: + raise RuntimeError("MiniMax H3 AdaLN cache parameters are required") + adaln_params = self.adaln_proj(t_emb) + shift, scale = adaln_params h = self.norm(x) h = indexed_scale_shift_(h, shift, scale, inverse_indices) # Preserve full precision through both final output projections. @@ -1107,6 +1295,9 @@ def __init__( self, od_config: OmniDiffusionConfig, quant_config: QuantizationConfig | None = None, + *, + adaln_cache_path: str | None = None, + adaln_cache_model_variant: str | None = None, ) -> None: super().__init__() tf_config = od_config.tf_model_config @@ -1114,10 +1305,14 @@ def __init__( arch = MiniMaxH3DiTArchConfig.from_mapping(config_mapping) self.arch = arch self.od_config = od_config + self.quant_config = quant_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 + if adaln_cache_path is not None: + _validate_adaln_cache_quant_config(quant_config) + self._adaln_precomputed = adaln_cache_path is not None self._validate_tp_config( arch=arch, tp_size=get_tensor_model_parallel_world_size(), @@ -1176,6 +1371,7 @@ def __init__( arch, quant_config, prefix=f"blocks.{i}", + use_adaln_cache=self._adaln_precomputed, ) for i in range(arch.num_layers) ] @@ -1188,17 +1384,26 @@ def __init__( arch, quant_config, prefix="final_layer", + use_adaln_cache=self._adaln_precomputed, ) + if adaln_cache_path is None: + self.adaln_cache = None + else: + # Local import avoids a module import cycle: the cache shares this + # file's architecture dataclass and modality constant. + from vllm_omni.diffusion.models.minimax_h3.adaln_cache import ( + MiniMaxH3AdalnCache, + ) + + self.adaln_cache = MiniMaxH3AdalnCache( + arch, + path=adaln_cache_path, + model_variant=adaln_cache_model_variant, + ) self._mark_missing_params_required() def enable_vsa_gates(self) -> None: - """Give every DiT block's attention a VSA compression gate. - - A FastH3 VSA artifact assigns these projections rather than adding to - them, so they have to exist before the weight stream reaches them. The - token refiner is left alone: the artifact carries gates for the 50 DiT - blocks only. - """ + """Materialize the 50 DiT VSA gates before adapter weight loading.""" if self.vsa_gates_enabled: return for block in self.blocks: @@ -1264,6 +1469,13 @@ def post_load_weights(self) -> None: 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}.") + if self.adaln_cache is not None: + self.adaln_cache.load(self.video_patch_proj.weight.device) + + if os.getenv("VLLM_OMNI_H3_DIT_MXFP8", "0") == "1": + from .quantization import install_dit_mxfp8 + + install_dit_mxfp8(self) def validate_restored_host_weights(self) -> None: """Validate mixed-precision invariants after lease-backed restore.""" @@ -1277,9 +1489,15 @@ def load_weights( params = dict(self.named_parameters()) params.update(dict(self.named_buffers())) loaded: set[str] = set() + adaln_cache = getattr(self, "adaln_cache", None) for name, loaded_weight in weights: param = params.get(name) if param is None: + if adaln_cache is not None and ( + name.startswith("final_layer.adaln_proj.") + or (name.startswith("blocks.") and ".adaln_proj." in name) + ): + continue logger.warning("Skipping MiniMax H3 weight not present in model: %s", name) continue weight_loader = getattr(param, "weight_loader", default_weight_loader) @@ -1526,6 +1744,16 @@ def forward(self, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]: block_rope = rope_table block_combined = combined_indices + adaln_cache = getattr(self, "adaln_cache", None) + adaln_cache_plan_index = None + block_adaln_params = None + if adaln_cache is not None: + adaln_cache_plan_index = adaln_cache.lookup(unique_timesteps.view(-1).to(device)) + block_adaln_params = adaln_cache.block_all( + cache_plan_index=adaln_cache_plan_index, + num_timesteps=t_emb.shape[0], + ) + if local_len == seq_len: hidden, block_rope, block_combined = self.sp_prepare( hidden, @@ -1538,37 +1766,49 @@ def forward(self, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]: block_rope, block_combined, ) - for block in self.blocks: - hidden = block( - hidden, - t_emb=t_emb, - combined_indices=block_combined, - rope_table=block_rope, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, - packed_total=seq_len, - num_requests=num_requests, - video_layout=video_layout, - vsa_prefix_segments=vsa_prefix_segments, - ) + for block_index, block in enumerate(self.blocks): + block_kwargs: dict[str, Any] = { + "t_emb": t_emb, + "combined_indices": block_combined, + "rope_table": block_rope, + "cu_seqlens": cu_seqlens, + "max_seqlen": max_seqlen, + "packed_total": seq_len, + "num_requests": num_requests, + "video_layout": video_layout, + "vsa_prefix_segments": vsa_prefix_segments, + } + if block_adaln_params is not None: + block_kwargs["adaln_params"] = block_adaln_params[block_index] + hidden = block(hidden, **block_kwargs) if local_len == seq_len: hidden = self.sp_gather(hidden) - video_logits, audio_logits = self.final_layer( - hidden, - t_emb=t_emb, - inverse_indices=inverse_indices, - ) + final_kwargs: dict[str, Any] = { + "t_emb": t_emb, + "inverse_indices": inverse_indices, + } + if adaln_cache_plan_index is not None: + final_kwargs["adaln_params"] = adaln_cache.final( + adaln_cache_plan_index, + t_emb.shape[0], + ) + video_logits, audio_logits = self.final_layer(hidden, **final_kwargs) else: local_inverse_indices = inverse_indices.narrow( 0, local_start, local_len, ) - video_logits, audio_logits = self.final_layer( - hidden, - t_emb=t_emb, - inverse_indices=local_inverse_indices, - ) + final_kwargs = { + "t_emb": t_emb, + "inverse_indices": local_inverse_indices, + } + if adaln_cache_plan_index is not None: + final_kwargs["adaln_params"] = adaln_cache.final( + adaln_cache_plan_index, + t_emb.shape[0], + ) + video_logits, audio_logits = self.final_layer(hidden, **final_kwargs) compact_logits = torch.cat((video_logits, audio_logits), dim=-1) compact_logits = self.sp_gather(compact_logits) video_width = self.arch.latents_dim * math.prod(self.arch.patch_size) diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/README.md b/vllm_omni/diffusion/models/minimax_h3/ops/README.md index e882ee78e28..c8449f1351a 100644 --- a/vllm_omni/diffusion/models/minimax_h3/ops/README.md +++ b/vllm_omni/diffusion/models/minimax_h3/ops/README.md @@ -27,10 +27,15 @@ dtype already selected by decode autocast. Unsupported model contracts, tensor inputs, execution modes, and devices retain the original implementation. Hardware dispatch is an explicit allowlist. SM90, SM100, and SM103 are enabled; -other capabilities fall back to the reference path. Bit-exact full-decode and -operator evidence has been collected on SM90, with independent full-decode and -stress validation on SM103. Enabling a target and claiming it as validated are -kept separate so the evidence remains clear. +other capabilities fall back to the reference path. SM120 has an experimental, +default-off entry gated by +`VLLM_OMNI_MINIMAX_H3_VAE_EXACT_OPS_SM120=1`; this is a process-startup setting, +not a request-time toggle. Bit-exact full-decode and operator evidence has been +collected on SM90, with independent full-decode and stress validation on SM103. +On SM120, production-shape operator checks and one 7-token, 704x1280, tile-256 +real-weight decode were bit-exact; full PP8 and complete-video qualification is +still required before default enablement. Enabling a target and claiming it as +validated are kept separate so the evidence remains clear. ## Extending platform support diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/__init__.py b/vllm_omni/diffusion/models/minimax_h3/ops/__init__.py index 6164781288b..7a66d902af1 100644 --- a/vllm_omni/diffusion/models/minimax_h3/ops/__init__.py +++ b/vllm_omni/diffusion/models/minimax_h3/ops/__init__.py @@ -2,6 +2,14 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniMax H3 model-specific optimized operators.""" -from .vae import install_h3_vae_optimizations +from .vae import ( + H3VAEExactOpStatsSnapshot, + install_h3_vae_optimizations, + snapshot_h3_vae_exact_op_stats, +) -__all__ = ["install_h3_vae_optimizations"] +__all__ = [ + "H3VAEExactOpStatsSnapshot", + "install_h3_vae_optimizations", + "snapshot_h3_vae_exact_op_stats", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/attention/__init__.py b/vllm_omni/diffusion/models/minimax_h3/ops/attention/__init__.py new file mode 100644 index 00000000000..4244ce4b4e1 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ops/attention/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/attention/chunks.py b/vllm_omni/diffusion/models/minimax_h3/ops/attention/chunks.py new file mode 100644 index 00000000000..615764b4b61 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ops/attention/chunks.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""H3 row-chunk producer using the shared BF16 indexed gate operator. + +Each chunk duplicates the owner's existing compact coarse tail. This keeps +the current native collective and its completion/abort protocol unchanged. +The additional bytes are measured and reported by the benchmark. +""" + +import torch +from vllm.triton_utils import tl, triton + +from vllm_omni.diffusion.layers.indexed_modulation import bf16_indexed_gate_add_ + + +@triton.jit +def _produce( + fine, + coarse, + owners, + out, + coarse_row_stride, + coarse_head_stride, + ROWS: tl.constexpr, # noqa: N803 + OFFSET: tl.constexpr, # noqa: N803 + COUNT: tl.constexpr, # noqa: N803 + TAIL: tl.constexpr, # noqa: N803 + WIDTH: tl.constexpr, # noqa: N803 + BLOCK: tl.constexpr, # noqa: N803 +): + row = tl.program_id(0) + owner = row // (COUNT + TAIL) + within = row % (COUNT + TAIL) + col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + fine_row = owner * ROWS + OFFSET + within + fine_value = tl.load(fine + fine_row * WIDTH + col, (within < COUNT) & (col < WIDTH), other=0) + tile = tl.load(owners + owner * TAIL + (within - COUNT), within >= COUNT, other=-1) + coarse_offset = tile * coarse_row_stride + (col // 128) * coarse_head_stride + col % 128 + coarse_value = tl.load(coarse + coarse_offset, (within >= COUNT) & (tile >= 0) & (col < WIDTH), other=0) + value = tl.where(within < COUNT, fine_value, coarse_value) + tl.store(out + row * WIDTH + col, value, col < WIDTH) + + +def produce(fine, coarse, owner_tiles, landing, start, count): + assert fine.shape == (1, 95936, 7, 128) and fine.is_contiguous() + assert coarse.shape == (1, 1648, 7, 128) and coarse.stride(-1) == 1 + assert 0 <= start < start + count <= 11992 + assert landing.shape == (1, 8 * (count + 270), 7, 128) and landing.is_contiguous() + _produce[(8 * (count + 270), triton.cdiv(896, 1024))]( + fine, + coarse, + owner_tiles, + landing, + coarse.stride(1), + coarse.stride(2), + ROWS=11992, + OFFSET=start, + COUNT=count, + TAIL=270, + WIDTH=896, + BLOCK=1024, + num_warps=8, + ) + + +def gate_chunk(bundle, gate, row_map, count): + """The source kernel's unchanged BF16 multiply then BF16 add, on a row slice.""" + assert bundle.shape == (1, count + 270, 56, 128) and bundle.is_contiguous() + assert gate.shape == (1, count, 56, 128) and gate.is_contiguous() + assert row_map.shape == (count,) and row_map.dtype == torch.int32 + return bf16_indexed_gate_add_(bundle, gate, row_map) diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/attention/layout.py b/vllm_omni/diffusion/models/minimax_h3/ops/attention/layout.py new file mode 100644 index 00000000000..8bbe238655e --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ops/attention/layout.py @@ -0,0 +1,592 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Inference-only layout helpers for MiniMax-H3 tile-64 VSA. + +The reference H3 path first zeroes a padded tile tensor and then evaluates an +advanced gather followed by indexed assignment:: + + tiled = torch.zeros(...) + tiled[:, non_pad] = compact[:, partition] + +For long H3 documents the RHS is another full compact tensor. The optional +CUDA paths below write every destination row exactly once and use cached +``destination row -> source row`` maps. Negative map entries produce exact +zeros, so edge-tile, optional partner-tile, and packed-sequence suffix padding +retain the reference contract. + +These are deliberately narrow, feature-gated primitives. Routing, sparse +block selection, and attention arithmetic are unchanged. +""" + +from __future__ import annotations + +import os +import weakref + +import torch +from vllm.triton_utils import HAS_TRITON, tl, triton + +H3_VSA_FUSED_TILE_PACK_ENV = "VLLM_OMNI_FASTVIDEO_VSA_FUSED_TILE_PACK" +H3_VSA_FUSED_UNTILE_ENV = "VLLM_OMNI_FASTVIDEO_VSA_FUSED_UNTILE" +# Producer-owned marker that identifies an H3 target-video call whose resolved +# attention backend is FASTVIDEO_VSA. It deliberately does not depend on the +# learned gate tensor being present: opt-in communication experiments use it +# to fail before their first collective if gate production regresses. +H3_VSA_ATTENTION_ACTIVE_KEY = "vsa_h3_attention_active" +_MAX_FUSED_ROW_WIDTH = 4096 + +# Avoid a device synchronization on every one of H3's 200 attention layers: +# CUDA kernels accept only maps that were range/uniqueness checked by the +# builders below. The weak reference prevents allocator pointer reuse from +# inheriting stale trust after a map is freed. +_VALIDATED_SOURCE_ROWS: dict[ + tuple[str, int | None, int], + tuple[weakref.ReferenceType[torch.Tensor], int, str, int, int], +] = {} + + +def _source_rows_key(source_rows: torch.Tensor) -> tuple[str, int | None, int]: + return ( + source_rows.device.type, + source_rows.device.index, + source_rows.data_ptr(), + ) + + +def _source_rows_version(source_rows: torch.Tensor) -> int | None: + """Return the mutation version used by the trusted-map registry. + + Tensors created while ``torch.inference_mode`` is active intentionally do + not have a version counter. Such a tensor cannot participate safely in a + cache whose trust contract depends on detecting in-place mutation, so it + must fail closed instead of making ``._version`` raise during model setup. + """ + if source_rows.is_inference(): + return None + return source_rows._version + + +def _register_source_rows( + source_rows: torch.Tensor, + *, + kind: str, + source_bound: int, + valid_rows: int, +) -> torch.Tensor: + key = _source_rows_key(source_rows) + version = _source_rows_version(source_rows) + if version is None: + # Do not let allocator pointer reuse inherit a previous tensor's trust + # if an inference tensor is ever passed here outside the builders. + _VALIDATED_SOURCE_ROWS.pop(key, None) + return source_rows + + def cleanup(reference: weakref.ReferenceType[torch.Tensor]) -> None: + current = _VALIDATED_SOURCE_ROWS.get(key) + if current is not None and current[0] is reference: + _VALIDATED_SOURCE_ROWS.pop(key, None) + + reference = weakref.ref(source_rows, cleanup) + _VALIDATED_SOURCE_ROWS[key] = ( + reference, + version, + kind, + source_bound, + valid_rows, + ) + return source_rows + + +def _source_rows_are_registered( + source_rows: torch.Tensor, + *, + kind: str, + source_bound: int, + valid_rows: int | None = None, +) -> bool: + entry = _VALIDATED_SOURCE_ROWS.get(_source_rows_key(source_rows)) + if entry is None: + return False + reference, version, registered_kind, registered_bound, registered_valid = entry + current_version = _source_rows_version(source_rows) + return ( + reference() is source_rows + and current_version is not None + and version == current_version + and registered_kind == kind + and registered_bound == source_bound + and (valid_rows is None or registered_valid == valid_rows) + ) + + +def h3_vsa_fused_tile_pack_enabled() -> bool: + """Return whether the opt-in H3 compact-to-tile CUDA path is requested.""" + raw = os.environ.get(H3_VSA_FUSED_TILE_PACK_ENV, "0").strip().lower() + return raw in {"1", "true", "yes", "on"} + + +def h3_vsa_fused_untile_enabled() -> bool: + """Return whether the opt-in H3 tile-to-aligned CUDA path is requested.""" + raw = os.environ.get(H3_VSA_FUSED_UNTILE_ENV, "0").strip().lower() + return raw in {"1", "true", "yes", "on"} + + +def build_h3_tiled_source_rows( + partition: torch.Tensor, + non_pad: torch.Tensor, + padded_rows: int, +) -> torch.Tensor: + """Build the inverse map consumed by :func:`h3_vsa_tile_pack`. + + ``partition`` and ``non_pad`` encode the existing reference assignment + ``dst[:, non_pad] = src[:, partition]``. The returned int32 vector has + length ``padded_rows``; ``-1`` denotes a destination padding row. + + Geometry is constructed once per cached H3 layout. Strict validation is + intentional: a malformed map must never reach an asynchronous CUDA + kernel, where an out-of-bounds source row would poison the process. + """ + if partition.ndim != 1 or non_pad.ndim != 1: + raise ValueError("partition and non_pad must be one-dimensional") + if partition.device != non_pad.device: + raise ValueError("partition and non_pad must be on the same device") + if partition.numel() != non_pad.numel(): + raise ValueError(f"partition and non_pad must have equal length, got {partition.numel()} and {non_pad.numel()}") + if partition.dtype not in (torch.int32, torch.int64) or non_pad.dtype not in (torch.int32, torch.int64): + raise TypeError("partition and non_pad must use int32 or int64 indices") + if padded_rows < partition.numel(): + raise ValueError(f"padded_rows={padded_rows} cannot hold {partition.numel()} compact rows") + if padded_rows > torch.iinfo(torch.int32).max: + raise ValueError("padded_rows exceeds the int32 layout-map contract") + + compact_rows = partition.numel() + if compact_rows: + partition64 = partition.to(torch.int64) + non_pad64 = non_pad.to(torch.int64) + if int(partition64.min()) < 0 or int(partition64.max()) >= compact_rows: + raise ValueError("partition contains an out-of-bounds compact source row") + if int(non_pad64.min()) < 0 or int(non_pad64.max()) >= padded_rows: + raise ValueError("non_pad contains an out-of-bounds padded destination row") + if torch.unique(partition64).numel() != compact_rows: + raise ValueError("partition must be a permutation of all compact source rows") + if torch.unique(non_pad64).numel() != compact_rows: + raise ValueError("non_pad destination rows must be unique") + + # Layout maps must retain a version counter even when model construction + # runs under the outer inference-mode context. The trusted-map registry + # uses that counter to reject maps modified after geometry validation. + with torch.inference_mode(False): + source_rows = torch.full( + (padded_rows,), + -1, + dtype=torch.int32, + device=partition.device, + ) + if compact_rows: + source_rows.index_copy_( + 0, + non_pad.to(torch.int64), + partition.to(torch.int32), + ) + return _register_source_rows( + source_rows, + kind="tile_pack", + source_bound=compact_rows, + valid_rows=compact_rows, + ) + + +def build_h3_aligned_untile_source_rows( + untile: torch.Tensor, + *, + aligned_rows: int, + tiled_rows: int, +) -> torch.Tensor: + """Build the map for one-pass tile-64 output un-tiling and pad restore. + + ``untile[compact_row]`` is the source row in the VSA kernel's tile-64 + output. The returned int32 map has ``aligned_rows`` entries: the valid + prefix preserves ``untile`` exactly and every structural suffix row is + ``-1``, which the CUDA kernel writes as an exact zero. + """ + if untile.ndim != 1: + raise ValueError("untile must be one-dimensional") + if untile.dtype not in (torch.int32, torch.int64): + raise TypeError("untile must use int32 or int64 indices") + if aligned_rows < untile.numel(): + raise ValueError(f"aligned_rows={aligned_rows} cannot hold {untile.numel()} valid rows") + if aligned_rows > torch.iinfo(torch.int32).max: + raise ValueError("aligned_rows exceeds the int32 layout-map contract") + if tiled_rows <= 0: + raise ValueError(f"tiled_rows must be positive, got {tiled_rows}") + + if untile.numel(): + untile64 = untile.to(torch.int64) + if int(untile64.min()) < 0 or int(untile64.max()) >= tiled_rows: + raise ValueError("untile contains an out-of-bounds tiled source row") + if torch.unique(untile64).numel() != untile.numel(): + raise ValueError("untile source rows must be unique") + + # See build_h3_tiled_source_rows: mutation tracking is part of the CUDA + # trust contract, so the map itself must not be an inference tensor. + with torch.inference_mode(False): + source_rows = torch.full( + (aligned_rows,), + -1, + dtype=torch.int32, + device=untile.device, + ) + if untile.numel(): + source_rows[: untile.numel()] = untile.to(torch.int32) + return _register_source_rows( + source_rows, + kind="aligned_untile", + source_bound=tiled_rows, + valid_rows=untile.numel(), + ) + + +def _validate_h3_tile_pack_inputs(x: torch.Tensor, source_rows: torch.Tensor) -> None: + if x.ndim != 4: + raise ValueError(f"H3 VSA tile pack expects [B,S,H,D], got {tuple(x.shape)}") + if source_rows.ndim != 1 or source_rows.dtype != torch.int32: + raise ValueError("source_rows must be a one-dimensional int32 tensor") + if source_rows.device != x.device: + raise ValueError("source_rows must be on the activation device") + if x.shape[1] > source_rows.numel(): + raise ValueError(f"padded destination has {source_rows.numel()} rows but compact input has {x.shape[1]} rows") + if source_rows.numel() == 0 and x.shape[1] != 0: + raise ValueError("a non-empty compact input requires a non-empty destination map") + + +def h3_vsa_tile_pack_cuda_supported(x: torch.Tensor, source_rows: torch.Tensor) -> bool: + """Whether the inference-only Triton implementation can safely run.""" + row_width = x.shape[2] * x.shape[3] if x.ndim == 4 else 0 + return ( + HAS_TRITON + and x.is_cuda + and x.ndim == 4 + and x.dtype in (torch.float16, torch.bfloat16) + and not x.requires_grad + and source_rows.is_cuda + and source_rows.device == x.device + and source_rows.ndim == 1 + and source_rows.dtype == torch.int32 + and source_rows.is_contiguous() + and x.shape[0] > 0 + and x.shape[1] > 0 + and x.shape[2] > 0 + and x.shape[3] > 0 + and x.shape[1] <= source_rows.numel() + and 0 < row_width <= _MAX_FUSED_ROW_WIDTH + and _source_rows_are_registered( + source_rows, + kind="tile_pack", + source_bound=x.shape[1], + valid_rows=x.shape[1], + ) + ) + + +def h3_vsa_tile_pack_reference(x: torch.Tensor, source_rows: torch.Tensor) -> torch.Tensor: + """Pure-Torch reference used by CPU tests and correctness qualification.""" + _validate_h3_tile_pack_inputs(x, source_rows) + output = torch.zeros( + (x.shape[0], source_rows.numel(), x.shape[2], x.shape[3]), + dtype=x.dtype, + device=x.device, + ) + valid_destinations = torch.nonzero(source_rows >= 0, as_tuple=False).flatten() + if valid_destinations.numel(): + compact_sources = source_rows.index_select(0, valid_destinations).to(torch.int64) + if int(compact_sources.min()) < 0 or int(compact_sources.max()) >= x.shape[1]: + raise ValueError("source_rows contains an out-of-bounds compact source row") + if torch.unique(compact_sources).numel() != x.shape[1] or compact_sources.numel() != x.shape[1]: + raise ValueError("source_rows must reference every compact source row exactly once") + output.index_copy_(1, valid_destinations, x.index_select(1, compact_sources)) + elif x.shape[1]: + raise ValueError("source_rows does not reference the non-empty compact input") + return output + + +def _validate_h3_tile_untile_inputs(x: torch.Tensor, source_rows: torch.Tensor) -> None: + if x.ndim != 4: + raise ValueError(f"H3 VSA tile untile expects [B,S,H,D], got {tuple(x.shape)}") + if source_rows.ndim != 1 or source_rows.dtype != torch.int32: + raise ValueError("source_rows must be a one-dimensional int32 tensor") + if source_rows.device != x.device: + raise ValueError("source_rows must be on the activation device") + if source_rows.numel() == 0: + raise ValueError("H3 VSA tile untile requires a non-empty destination map") + + +def _validate_h3_tile_untile_output( + x: torch.Tensor, + source_rows: torch.Tensor, + out: torch.Tensor, +) -> None: + expected_shape = ( + x.shape[0], + source_rows.numel(), + x.shape[2], + x.shape[3], + ) + if tuple(out.shape) != expected_shape: + raise ValueError(f"H3 untile output shape mismatch: expected {expected_shape}, got {tuple(out.shape)}") + if out.dtype != x.dtype or out.device != x.device: + raise ValueError( + "H3 VSA tile untile output dtype/device mismatch: " + f"expected=({x.dtype}, {x.device}), observed=({out.dtype}, {out.device})" + ) + if not out.is_contiguous(): + raise ValueError("H3 VSA tile untile output must be contiguous") + if out.requires_grad: + raise ValueError("H3 VSA tile untile output must not require gradients") + if out.untyped_storage().data_ptr() == x.untyped_storage().data_ptr(): + raise ValueError("H3 VSA tile untile input and output must not share storage") + + +def h3_vsa_tile_untile_cuda_supported( + x: torch.Tensor, + source_rows: torch.Tensor, +) -> bool: + """Whether the one-pass tile-to-aligned CUDA implementation can run.""" + row_width = x.shape[2] * x.shape[3] if x.ndim == 4 else 0 + return ( + HAS_TRITON + and x.is_cuda + and x.ndim == 4 + and x.is_contiguous() + and x.dtype in (torch.float16, torch.bfloat16) + and not x.requires_grad + and source_rows.is_cuda + and source_rows.device == x.device + and source_rows.ndim == 1 + and source_rows.dtype == torch.int32 + and source_rows.is_contiguous() + and x.shape[0] > 0 + and x.shape[1] > 0 + and x.shape[2] > 0 + and x.shape[3] > 0 + and source_rows.numel() > 0 + and 0 < row_width <= _MAX_FUSED_ROW_WIDTH + and _source_rows_are_registered( + source_rows, + kind="aligned_untile", + source_bound=x.shape[1], + ) + ) + + +def h3_vsa_tile_untile_reference( + x: torch.Tensor, + source_rows: torch.Tensor, +) -> torch.Tensor: + """Pure-Torch tile-to-aligned reference, including exact suffix zeros.""" + _validate_h3_tile_untile_inputs(x, source_rows) + output = torch.zeros( + (x.shape[0], source_rows.numel(), x.shape[2], x.shape[3]), + dtype=x.dtype, + device=x.device, + ) + valid_destinations = torch.nonzero(source_rows >= 0, as_tuple=False).flatten() + if valid_destinations.numel(): + tiled_sources = source_rows.index_select(0, valid_destinations).to(torch.int64) + if int(tiled_sources.min()) < 0 or int(tiled_sources.max()) >= x.shape[1]: + raise ValueError("source_rows contains an out-of-bounds tiled source row") + if torch.unique(tiled_sources).numel() != tiled_sources.numel(): + raise ValueError("source_rows tiled source rows must be unique") + output.index_copy_(1, valid_destinations, x.index_select(1, tiled_sources)) + return output + + +if HAS_TRITON: + + @triton.jit + def _h3_vsa_tile_pack_kernel( + source_ptr, + source_rows_ptr, + output_ptr, + source_stride_batch, + source_stride_seq, + source_stride_head, + source_stride_dim, + padded_rows: tl.constexpr, + head_dim: tl.constexpr, + row_width: tl.constexpr, + block_size: tl.constexpr, + ): + flat_destination_row = tl.program_id(0) + batch_index = flat_destination_row // padded_rows + destination_row = flat_destination_row - batch_index * padded_rows + source_row = tl.load(source_rows_ptr + destination_row) + + feature_offsets = tl.arange(0, block_size) + feature_mask = feature_offsets < row_width + head_index = feature_offsets // head_dim + dim_index = feature_offsets - head_index * head_dim + valid_source = source_row >= 0 + source_offsets = ( + batch_index * source_stride_batch + + source_row * source_stride_seq + + head_index * source_stride_head + + dim_index * source_stride_dim + ) + values = tl.load(source_ptr + source_offsets, mask=feature_mask & valid_source, other=0.0) + output_offsets = flat_destination_row * row_width + feature_offsets + tl.store(output_ptr + output_offsets, values, mask=feature_mask) + + +def _h3_vsa_tile_pack_cuda_impl(x: torch.Tensor, source_rows: torch.Tensor) -> torch.Tensor: + _validate_h3_tile_pack_inputs(x, source_rows) + if not h3_vsa_tile_pack_cuda_supported(x, source_rows): + raise RuntimeError("H3 VSA fused tile pack received an unsupported CUDA input") + output = torch.empty( + (x.shape[0], source_rows.numel(), x.shape[2], x.shape[3]), + dtype=x.dtype, + device=x.device, + ) + row_width = x.shape[2] * x.shape[3] + block_size = triton.next_power_of_2(row_width) + _h3_vsa_tile_pack_kernel[(x.shape[0] * source_rows.numel(),)]( + x, + source_rows, + output, + x.stride(0), + x.stride(1), + x.stride(2), + x.stride(3), + padded_rows=source_rows.numel(), + head_dim=x.shape[3], + row_width=row_width, + block_size=block_size, + num_warps=4 if block_size <= 256 else 8, + ) + return output + + +def _h3_vsa_tile_untile_cuda_impl( + x: torch.Tensor, + source_rows: torch.Tensor, + out: torch.Tensor | None = None, +) -> torch.Tensor: + _validate_h3_tile_untile_inputs(x, source_rows) + if not h3_vsa_tile_untile_cuda_supported(x, source_rows): + raise RuntimeError("H3 VSA fused tile untile received an unsupported CUDA input") + if out is None: + output = torch.empty( + (x.shape[0], source_rows.numel(), x.shape[2], x.shape[3]), + dtype=x.dtype, + device=x.device, + ) + else: + _validate_h3_tile_untile_output(x, source_rows, out) + output = out + row_width = x.shape[2] * x.shape[3] + block_size = triton.next_power_of_2(row_width) + _h3_vsa_tile_pack_kernel[(x.shape[0] * source_rows.numel(),)]( + x, + source_rows, + output, + x.stride(0), + x.stride(1), + x.stride(2), + x.stride(3), + padded_rows=source_rows.numel(), + head_dim=x.shape[3], + row_width=row_width, + block_size=block_size, + num_warps=4 if block_size <= 256 else 8, + ) + return output + + +if not hasattr(torch.ops.vllm_omni, "h3_vsa_tile_pack"): + + @torch.library.custom_op("vllm_omni::h3_vsa_tile_pack", mutates_args=()) + def _h3_vsa_tile_pack_op(x: torch.Tensor, source_rows: torch.Tensor) -> torch.Tensor: + return _h3_vsa_tile_pack_cuda_impl(x, source_rows) + + @_h3_vsa_tile_pack_op.register_fake + def _(x: torch.Tensor, source_rows: torch.Tensor) -> torch.Tensor: + return torch.empty( + (x.shape[0], source_rows.shape[0], x.shape[2], x.shape[3]), + dtype=x.dtype, + device=x.device, + ) + + +if not hasattr(torch.ops.vllm_omni, "h3_vsa_tile_untile"): + + @torch.library.custom_op("vllm_omni::h3_vsa_tile_untile", mutates_args=()) + def _h3_vsa_tile_untile_op( + x: torch.Tensor, + source_rows: torch.Tensor, + ) -> torch.Tensor: + return _h3_vsa_tile_untile_cuda_impl(x, source_rows) + + @_h3_vsa_tile_untile_op.register_fake + def _(x: torch.Tensor, source_rows: torch.Tensor) -> torch.Tensor: + return torch.empty( + (x.shape[0], source_rows.shape[0], x.shape[2], x.shape[3]), + dtype=x.dtype, + device=x.device, + ) + + +def h3_vsa_tile_pack(x: torch.Tensor, source_rows: torch.Tensor) -> torch.Tensor: + """Pack compact H3 rows into tile-64 order through the fused CUDA op.""" + _validate_h3_tile_pack_inputs(x, source_rows) + if not h3_vsa_tile_pack_cuda_supported(x, source_rows): + raise RuntimeError("H3 VSA fused tile pack is unavailable for this input") + return torch.ops.vllm_omni.h3_vsa_tile_pack(x, source_rows) + + +def h3_vsa_tile_untile( + x: torch.Tensor, + source_rows: torch.Tensor, +) -> torch.Tensor: + """Un-tile H3 output into aligned packed order through the fused CUDA op.""" + _validate_h3_tile_untile_inputs(x, source_rows) + if not h3_vsa_tile_untile_cuda_supported(x, source_rows): + raise RuntimeError("H3 VSA fused tile untile is unavailable for this input") + return torch.ops.vllm_omni.h3_vsa_tile_untile(x, source_rows) + + +@torch.compiler.disable +def h3_vsa_tile_untile_out( + x: torch.Tensor, + source_rows: torch.Tensor, + *, + out: torch.Tensor, +) -> torch.Tensor: + """Un-tile directly into caller-owned eager output storage. + + This mutating form intentionally bypasses the functional custom op. H3's + packed-attention island is already eager, and the caller-owned storage is + a registered FlashInfer communication buffer that cannot be represented + as a functional compiler allocation. + """ + _validate_h3_tile_untile_inputs(x, source_rows) + _validate_h3_tile_untile_output(x, source_rows, out) + if not h3_vsa_tile_untile_cuda_supported(x, source_rows): + raise RuntimeError("H3 VSA fused tile untile-out is unavailable for this input") + return _h3_vsa_tile_untile_cuda_impl(x, source_rows, out=out) + + +__all__ = [ + "H3_VSA_ATTENTION_ACTIVE_KEY", + "H3_VSA_FUSED_TILE_PACK_ENV", + "H3_VSA_FUSED_UNTILE_ENV", + "build_h3_aligned_untile_source_rows", + "build_h3_tiled_source_rows", + "h3_vsa_fused_tile_pack_enabled", + "h3_vsa_fused_untile_enabled", + "h3_vsa_tile_pack", + "h3_vsa_tile_pack_cuda_supported", + "h3_vsa_tile_pack_reference", + "h3_vsa_tile_untile", + "h3_vsa_tile_untile_out", + "h3_vsa_tile_untile_cuda_supported", + "h3_vsa_tile_untile_reference", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/attention/owner_route.py b/vllm_omni/diffusion/models/minimax_h3/ops/attention/owner_route.py new file mode 100644 index 00000000000..eed8d4bdc76 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ops/attention/owner_route.py @@ -0,0 +1,361 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU route planning for MiniMax-H3 VSA coarse-output piggybacking. + +The H3 VSA output is restored from tile-major order through an aligned +``destination row -> tiled source row`` map. Under Ulysses, aligned output +rows are split into equal contiguous chunks. A receiver therefore needs only +the coarse tiles referenced by its own chunk, rather than the full coarse +gate/output tensor. + +This module builds that static owner route. It deliberately contains no CUDA +or distributed communication code: the plan is constructed once on CPU and +can be consumed by a later producer-direct reverse-O implementation. +""" + +from __future__ import annotations + +import weakref +from dataclasses import dataclass + +import torch + +H3_VSA_TILE_ROWS = 64 +_INT32_MAX = torch.iinfo(torch.int32).max + +_TRUSTED_PLANS: dict[ + int, + tuple[ + weakref.ReferenceType[H3VSAOwnerRoutePlan], + tuple[weakref.ReferenceType[torch.Tensor], ...], + tuple[int, ...], + tuple[int, int, int, int], + tuple[int, ...], + ], +] = {} + + +@dataclass(frozen=True) +class H3VSAOwnerRoutePlan: + """Static owner route for one aligned H3 VSA layout. + + ``owner_tile_ids[owner, :owner_tile_counts[owner]]`` lists the sorted, + unique logical coarse tiles needed by that owner's sequence rows. The + remaining columns are exactly ``-1``. + + ``local_row_to_tail[owner, local_row]`` maps an effective local row to a + relative slot in that owner's compact tail. Structural aligned padding is + represented by ``-1``. + """ + + owner_tile_ids: torch.Tensor + owner_tile_counts: torch.Tensor + local_row_to_tail: torch.Tensor + logical_block_count: int + + @property + def sp_world_size(self) -> int: + return self.owner_tile_ids.shape[0] + + @property + def local_rows(self) -> int: + return self.local_row_to_tail.shape[1] + + @property + def kmax(self) -> int: + return self.owner_tile_ids.shape[1] + + @property + def local_row_to_receive_row(self) -> torch.Tensor: + """Map local rows to absolute rows in ``[local O; compact tail]``. + + The returned tensor has the same shape as ``local_row_to_tail``. + Valid entries are ``local_rows + relative_tail_slot``; aligned padding + remains ``-1``. + """ + return torch.where( + self.local_row_to_tail >= 0, + self.local_row_to_tail + self.local_rows, + self.local_row_to_tail, + ) + + @property + def owner_tile_entries(self) -> int: + """Total tile payload rows across all owners, before Kmax padding.""" + return int(self.owner_tile_counts.to(torch.int64).sum()) + + @property + def duplicate_owner_tile_entries(self) -> int: + """Extra owner/tile entries caused by tiles spanning SP owners.""" + return self.owner_tile_entries - self.logical_block_count + + @property + def cross_owner_tile_count(self) -> int: + """Number of logical tiles referenced by more than one SP owner.""" + valid_tiles = self.owner_tile_ids[self.owner_tile_ids >= 0].to(torch.int64) + frequencies = torch.bincount(valid_tiles, minlength=self.logical_block_count) + return int(torch.count_nonzero(frequencies > 1)) + + +def _register_owner_route_plan(plan: H3VSAOwnerRoutePlan) -> H3VSAOwnerRoutePlan: + key = id(plan) + + def cleanup(reference: weakref.ReferenceType[H3VSAOwnerRoutePlan]) -> None: + current = _TRUSTED_PLANS.get(key) + if current is not None and current[0] is reference: + _TRUSTED_PLANS.pop(key, None) + + tensors = ( + plan.owner_tile_ids, + plan.owner_tile_counts, + plan.local_row_to_tail, + ) + plan_reference = weakref.ref(plan, cleanup) + tensor_references = tuple(weakref.ref(tensor) for tensor in tensors) + versions = tuple(tensor._version for tensor in tensors) + geometry = ( + plan.sp_world_size, + plan.local_rows, + plan.kmax, + plan.logical_block_count, + ) + data_ptrs = tuple(tensor.data_ptr() for tensor in tensors) + _TRUSTED_PLANS[key] = ( + plan_reference, + tensor_references, + versions, + geometry, + data_ptrs, + ) + return plan + + +def h3_vsa_owner_route_plan_is_trusted(plan: H3VSAOwnerRoutePlan) -> bool: + """Whether ``plan`` is an unmodified result of the validated builder.""" + if not isinstance(plan, H3VSAOwnerRoutePlan): + return False + entry = _TRUSTED_PLANS.get(id(plan)) + if entry is None: + return False + plan_reference, tensor_references, versions, geometry, data_ptrs = entry + tensors = ( + plan.owner_tile_ids, + plan.owner_tile_counts, + plan.local_row_to_tail, + ) + return ( + plan_reference() is plan + and all(reference() is tensor for reference, tensor in zip(tensor_references, tensors)) + and versions == tuple(tensor._version for tensor in tensors) + and data_ptrs == tuple(tensor.data_ptr() for tensor in tensors) + and geometry + == ( + plan.sp_world_size, + plan.local_rows, + plan.kmax, + plan.logical_block_count, + ) + ) + + +def _validate_positive_int(name: str, value: int) -> None: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + + +def _validate_aligned_source_rows( + aligned_untile_source_rows: torch.Tensor, + *, + global_rows: int, + logical_block_count: int, +) -> None: + if aligned_untile_source_rows.device.type != "cpu": + raise ValueError("aligned_untile_source_rows must be a CPU tensor") + if aligned_untile_source_rows.ndim != 1: + raise ValueError("aligned_untile_source_rows must be one-dimensional") + if aligned_untile_source_rows.dtype != torch.int32: + raise TypeError("aligned_untile_source_rows must use int32 indices") + if not aligned_untile_source_rows.is_contiguous(): + raise ValueError("aligned_untile_source_rows must be contiguous") + if aligned_untile_source_rows.numel() != global_rows: + raise ValueError( + "aligned_untile_source_rows length must equal global_rows, got " + f"{aligned_untile_source_rows.numel()} and {global_rows}" + ) + + source_rows64 = aligned_untile_source_rows.to(torch.int64) + if torch.any(source_rows64 < -1): + raise ValueError("aligned_untile_source_rows may use only -1 for padding") + + valid_mask = source_rows64 >= 0 + valid_count = int(valid_mask.sum()) + if valid_count == 0: + raise ValueError("aligned_untile_source_rows must contain valid rows") + if not torch.all(valid_mask[:valid_count]) or torch.any(valid_mask[valid_count:]): + raise ValueError("aligned padding must be a contiguous -1 suffix") + + valid_rows = source_rows64[:valid_count] + tiled_rows = logical_block_count * H3_VSA_TILE_ROWS + if int(valid_rows.max()) >= tiled_rows: + raise ValueError("aligned_untile_source_rows contains an out-of-bounds tiled source row") + if torch.unique(valid_rows).numel() != valid_count: + raise ValueError("valid aligned_untile_source_rows must be globally unique") + + covered_tiles = torch.unique( + torch.div(valid_rows, H3_VSA_TILE_ROWS, rounding_mode="floor"), + sorted=True, + ) + expected_tiles = torch.arange(logical_block_count, dtype=torch.int64) + if not torch.equal(covered_tiles, expected_tiles): + raise ValueError("aligned_untile_source_rows must cover every logical tile") + + +def _validate_owner_route_plan( + plan: H3VSAOwnerRoutePlan, + aligned_untile_source_rows: torch.Tensor, +) -> None: + """Fail closed if a constructed route violates its source-map contract.""" + sp_world_size = plan.sp_world_size + local_rows = plan.local_rows + kmax = plan.kmax + + if plan.owner_tile_ids.dtype != torch.int32: + raise RuntimeError("owner_tile_ids must use int32 indices") + if plan.owner_tile_counts.dtype != torch.int32: + raise RuntimeError("owner_tile_counts must use int32 indices") + if plan.local_row_to_tail.dtype != torch.int32: + raise RuntimeError("local_row_to_tail must use int32 indices") + if plan.owner_tile_counts.shape != (sp_world_size,): + raise RuntimeError("owner_tile_counts shape does not match owner_tile_ids") + if plan.local_row_to_tail.shape != (sp_world_size, local_rows): + raise RuntimeError("local_row_to_tail has an invalid shape") + + local_sources = aligned_untile_source_rows.view(sp_world_size, local_rows).to(torch.int64) + for owner in range(sp_world_size): + count = int(plan.owner_tile_counts[owner]) + if not 0 <= count <= kmax: + raise RuntimeError("owner_tile_counts contains an out-of-range count") + + tile_ids = plan.owner_tile_ids[owner, :count].to(torch.int64) + padding = plan.owner_tile_ids[owner, count:] + if count and ( + int(tile_ids.min()) < 0 + or int(tile_ids.max()) >= plan.logical_block_count + or torch.unique(tile_ids).numel() != count + or (count > 1 and not torch.all(tile_ids[1:] > tile_ids[:-1])) + ): + raise RuntimeError("an owner's active tile IDs must be sorted, unique, and in range") + if torch.any(padding != -1): + raise RuntimeError("owner_tile_ids padding must be exactly -1") + + source = local_sources[owner] + valid_mask = source >= 0 + slots = plan.local_row_to_tail[owner].to(torch.int64) + if torch.any(slots[~valid_mask] != -1): + raise RuntimeError("aligned padding rows must map to -1 tail slots") + if torch.any(slots[valid_mask] < 0) or torch.any(slots[valid_mask] >= count): + raise RuntimeError("a valid local row maps outside its owner's active tail") + if valid_mask.any(): + expected_tiles = torch.div(source[valid_mask], H3_VSA_TILE_ROWS, rounding_mode="floor") + routed_tiles = tile_ids.index_select(0, slots[valid_mask]) + if not torch.equal(routed_tiles, expected_tiles): + raise RuntimeError("local_row_to_tail does not route to the row's logical tile") + + +def build_h3_vsa_owner_route_plan( + aligned_untile_source_rows: torch.Tensor, + *, + sp_world_size: int, + global_rows: int, + local_rows: int, + logical_block_count: int, +) -> H3VSAOwnerRoutePlan: + """Build a compact per-owner coarse-tile route for reverse-O piggybacking. + + Args: + aligned_untile_source_rows: Validated aligned output-row to tile-major + source-row map. It must be contiguous CPU int32 with only a + trailing ``-1`` padding suffix. + sp_world_size: Number of contiguous Ulysses sequence owners. + global_rows: Aligned global sequence length. + local_rows: Aligned rows received by each owner. + logical_block_count: Number of tile-64 coarse blocks. + """ + _validate_positive_int("sp_world_size", sp_world_size) + _validate_positive_int("global_rows", global_rows) + _validate_positive_int("local_rows", local_rows) + _validate_positive_int("logical_block_count", logical_block_count) + if global_rows != sp_world_size * local_rows: + raise ValueError( + f"global_rows must equal sp_world_size * local_rows, got {global_rows} != {sp_world_size} * {local_rows}" + ) + if global_rows > _INT32_MAX or logical_block_count * H3_VSA_TILE_ROWS > _INT32_MAX: + raise ValueError("route geometry exceeds the int32 indexing contract") + + _validate_aligned_source_rows( + aligned_untile_source_rows, + global_rows=global_rows, + logical_block_count=logical_block_count, + ) + + # These three persistent maps participate in the trusted-plan registry, + # whose mutation check relies on their tensor version counters. Model + # execution constructs this plan under torch.inference_mode(), where new + # tensors ordinarily have no version counter. Build the maps with + # inference mode locally disabled, matching the trusted H3 layout-map + # builders, so production and test execution retain the same fail-closed + # mutation contract. + with torch.inference_mode(False): + local_sources = aligned_untile_source_rows.view(sp_world_size, local_rows).to(torch.int64).clone() + owner_tiles: list[torch.Tensor] = [] + counts: list[int] = [] + for owner in range(sp_world_size): + valid_sources = local_sources[owner][local_sources[owner] >= 0] + tile_ids = torch.unique( + torch.div(valid_sources, H3_VSA_TILE_ROWS, rounding_mode="floor"), + sorted=True, + ) + owner_tiles.append(tile_ids) + counts.append(tile_ids.numel()) + + kmax = max(counts) + owner_tile_ids = torch.full( + (sp_world_size, kmax), + -1, + dtype=torch.int32, + ) + owner_tile_counts = torch.tensor(counts, dtype=torch.int32) + local_row_to_tail = torch.full( + (sp_world_size, local_rows), + -1, + dtype=torch.int32, + ) + + for owner, tile_ids in enumerate(owner_tiles): + count = tile_ids.numel() + owner_tile_ids[owner, :count] = tile_ids.to(torch.int32) + + source = local_sources[owner] + valid_mask = source >= 0 + local_tiles = torch.div(source[valid_mask], H3_VSA_TILE_ROWS, rounding_mode="floor") + local_row_to_tail[owner, valid_mask] = torch.searchsorted( + tile_ids, + local_tiles, + ).to(torch.int32) + + plan = H3VSAOwnerRoutePlan( + owner_tile_ids=owner_tile_ids, + owner_tile_counts=owner_tile_counts, + local_row_to_tail=local_row_to_tail, + logical_block_count=logical_block_count, + ) + _validate_owner_route_plan(plan, aligned_untile_source_rows) + return _register_owner_route_plan(plan) + + +__all__ = [ + "H3_VSA_TILE_ROWS", + "H3VSAOwnerRoutePlan", + "build_h3_vsa_owner_route_plan", + "h3_vsa_owner_route_plan_is_trusted", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py index 247e072c82b..2eba6a0a6b0 100644 --- a/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py +++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/__init__.py @@ -4,6 +4,7 @@ from __future__ import annotations +from dataclasses import dataclass from types import MethodType from typing import Any @@ -15,6 +16,55 @@ from .dispatch import resolve_h3_vae_operators +@dataclass(frozen=True) +class H3VAEExactOpStatsSnapshot: + """Immutable per-decoder exact-op dispatch counters.""" + + transformer_fast: int = 0 + transformer_fallback: int = 0 + qk_norm_rope_fast: int = 0 + qk_norm_rope_fallback: int = 0 + swiglu_fast: int = 0 + swiglu_fallback: int = 0 + scaled_residual_fast: int = 0 + scaled_residual_fallback: int = 0 + + def delta(self, previous: H3VAEExactOpStatsSnapshot) -> H3VAEExactOpStatsSnapshot: + """Return counters accumulated after ``previous``.""" + + return H3VAEExactOpStatsSnapshot( + **{name: int(getattr(self, name)) - int(getattr(previous, name)) for name in self.__dataclass_fields__} + ) + + +@dataclass +class _H3VAEExactOpStats: + """Mutable counters shared by all optimized blocks in one decoder.""" + + transformer_fast: int = 0 + transformer_fallback: int = 0 + qk_norm_rope_fast: int = 0 + qk_norm_rope_fallback: int = 0 + swiglu_fast: int = 0 + swiglu_fallback: int = 0 + scaled_residual_fast: int = 0 + scaled_residual_fallback: int = 0 + + def snapshot(self) -> H3VAEExactOpStatsSnapshot: + return H3VAEExactOpStatsSnapshot( + **{name: int(getattr(self, name)) for name in H3VAEExactOpStatsSnapshot.__dataclass_fields__} + ) + + +def snapshot_h3_vae_exact_op_stats(decoder: nn.Module) -> H3VAEExactOpStatsSnapshot | None: + """Snapshot an installed decoder, or return ``None`` when not installed.""" + + stats = getattr(decoder, "_omni_h3_vae_exact_op_stats", None) + if not isinstance(stats, _H3VAEExactOpStats): + return None + return stats.snapshot() + + def _is_boolean_flag(value: Any) -> bool: """Accept bools and the JSON-derived 0/1 flags used by remote VAE code.""" @@ -44,10 +94,13 @@ def _optimized_feed_forward(self: nn.Module, hidden_states: torch.Tensor) -> tor if torch.compiler.is_compiling(): return type(self).forward(self, hidden_states) + stats = self._omni_h3_vae_exact_op_stats hidden_states = self.w1(hidden_states) if hidden_states.is_cuda and hidden_states.dtype == torch.float16: + stats.swiglu_fast += 1 hidden_states = self._omni_silu_and_mul(hidden_states) else: + stats.swiglu_fallback += 1 gate, hidden_states = hidden_states.chunk(2, dim=-1) hidden_states = self.act_fn(gate) * hidden_states return self.w2(hidden_states) @@ -60,7 +113,10 @@ def _optimized_attention( pack_info: dict[str, Any] | None = None, ) -> torch.Tensor: pack_info = {} if pack_info is None else pack_info + stats = self._omni_h3_vae_exact_op_stats if torch.compiler.is_compiling() or self.spatial_parallel or rotary_pos_emb is None: + if not torch.compiler.is_compiling(): + stats.qk_norm_rope_fallback += 1 return type(self).forward(self, hidden_states, rotary_pos_emb, pack_info) batch_size, sequence, _ = hidden_states.shape @@ -78,8 +134,10 @@ def _optimized_attention( float(self.norm_q.eps), ) if optimized_qk is None: + stats.qk_norm_rope_fallback += 1 return type(self).forward(self, hidden_states, rotary_pos_emb, pack_info) + stats.qk_norm_rope_fast += 1 query, key = optimized_qk hidden_states = self.perform_attention(query, key, value, pack_info) hidden_states = hidden_states.reshape(batch_size, sequence, -1) @@ -93,9 +151,13 @@ def _optimized_transformer_block( pack_info: dict[str, Any] | None = None, ) -> torch.Tensor: pack_info = {} if pack_info is None else pack_info + stats = self._omni_h3_vae_exact_op_stats if torch.compiler.is_compiling() or hidden_states.dtype != torch.float32: + if not torch.compiler.is_compiling(): + stats.transformer_fallback += 1 return type(self).forward(self, hidden_states, rotary_pos_emb, pack_info) + stats.transformer_fast += 1 normalized = self.norm1(hidden_states.float()).to(hidden_states.dtype) attention_output = self.attn(normalized, rotary_pos_emb, pack_info) updated = self._omni_scaled_residual( @@ -103,6 +165,10 @@ def _optimized_transformer_block( attention_output, self.scale1, ) + if updated is None: + stats.scaled_residual_fallback += 1 + else: + stats.scaled_residual_fast += 1 hidden_states = hidden_states + attention_output * self.scale1 if updated is None else updated normalized = self.norm2(hidden_states.float()).to(hidden_states.dtype) @@ -112,6 +178,10 @@ def _optimized_transformer_block( feed_forward_output, self.scale2, ) + if updated is None: + stats.scaled_residual_fallback += 1 + else: + stats.scaled_residual_fast += 1 return hidden_states + feed_forward_output * self.scale2 if updated is None else updated @@ -184,6 +254,7 @@ def install_h3_vae_optimizations( decoder: nn.Module, *, device: torch.device, + persist_fp16_weights: bool = True, ) -> bool: """Install the operators selected for ``device`` once.""" @@ -199,19 +270,29 @@ def install_h3_vae_optimizations( # The H3 decode path always uses FP16 CUDA autocast. Persisting these # rounded decoder-block weights avoids rebuilding the same casts per tile. - for linear in linears: - linear.to(dtype=torch.float16) + if persist_fp16_weights: + for linear in linears: + linear.to(dtype=torch.float16) + stats = _H3VAEExactOpStats() for block in decoder.transformer_blocks: + block._omni_h3_vae_exact_op_stats = stats block.ff._omni_silu_and_mul = SiluAndMul() + block.ff._omni_h3_vae_exact_op_stats = stats block.ff.forward = MethodType(_optimized_feed_forward, block.ff) block.attn._omni_qk_norm_rope = operators.qk_norm_rope + block.attn._omni_h3_vae_exact_op_stats = stats block.attn.forward = MethodType(_optimized_attention, block.attn) block._omni_scaled_residual = operators.scaled_residual block.forward = MethodType(_optimized_transformer_block, block) + decoder._omni_h3_vae_exact_op_stats = stats decoder._omni_h3_vae_optimizations_installed = True return True -__all__ = ["install_h3_vae_optimizations"] +__all__ = [ + "H3VAEExactOpStatsSnapshot", + "install_h3_vae_optimizations", + "snapshot_h3_vae_exact_op_stats", +] diff --git a/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py b/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py index d413ef85008..3028c19ecd6 100644 --- a/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py +++ b/vllm_omni/diffusion/models/minimax_h3/ops/vae/dispatch.py @@ -4,6 +4,7 @@ from __future__ import annotations +import os from collections.abc import Callable from dataclasses import dataclass @@ -15,6 +16,8 @@ from .qk_norm_rope import try_qk_norm_rope_exact from .scaled_residual import try_scaled_residual_exact +MINIMAX_H3_VAE_EXACT_OPS_SM120_ENV = "VLLM_OMNI_MINIMAX_H3_VAE_EXACT_OPS_SM120" + QKNormRopeOp = Callable[ [ torch.Tensor, @@ -65,6 +68,25 @@ def _supports_cuda_sm103(device: torch.device) -> bool: return _supports_cuda_capability(device, 103) +def resolve_minimax_h3_vae_exact_ops_sm120(raw: str | None = None) -> bool: + """Resolve the experimental SM120 exact-op gate without silent fallback.""" + + if raw is None: + raw = os.environ.get(MINIMAX_H3_VAE_EXACT_OPS_SM120_ENV, "0") + normalized = str(raw).strip().lower() + if normalized in {"1", "true", "yes", "on", "enable", "enabled"}: + return True + if normalized in {"0", "false", "no", "off", "disable", "disabled", ""}: + return False + raise ValueError(f"{MINIMAX_H3_VAE_EXACT_OPS_SM120_ENV} must be 0 or 1, got {raw!r}") + + +def _supports_cuda_sm120(device: torch.device) -> bool: + # Parse the opt-in only on the exact experimental target. An invalid value + # must fail that worker at startup, but must not affect other architectures. + return _supports_cuda_capability(device, 120) and resolve_minimax_h3_vae_exact_ops_sm120() + + # Keep hardware selection flat: adding a backend means adding one operator set, # without changing the installer or the model execution path. H3_VAE_OPERATOR_TABLE: tuple[H3VAEOperatorSet, ...] = ( @@ -83,6 +105,11 @@ def _supports_cuda_sm103(device: torch.device) -> bool: qk_norm_rope=try_qk_norm_rope_exact, scaled_residual=try_scaled_residual_exact, ), + H3VAEOperatorSet( + supports=_supports_cuda_sm120, + qk_norm_rope=try_qk_norm_rope_exact, + scaled_residual=try_scaled_residual_exact, + ), ) @@ -94,7 +121,9 @@ def resolve_h3_vae_operators(device: torch.device) -> H3VAEOperatorSet | None: __all__ = [ + "MINIMAX_H3_VAE_EXACT_OPS_SM120_ENV", "H3VAEOperatorSet", "H3_VAE_OPERATOR_TABLE", + "resolve_minimax_h3_vae_exact_ops_sm120", "resolve_h3_vae_operators", ] diff --git a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py index f80513d4eb6..ae0b33c82f1 100644 --- a/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py +++ b/vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py @@ -8,8 +8,10 @@ import math import os import tempfile +import time from collections.abc import Iterable, Mapping, Sequence from contextlib import contextmanager, nullcontext +from dataclasses import dataclass from itertools import groupby from pathlib import Path from typing import TYPE_CHECKING, Any, ClassVar @@ -114,6 +116,11 @@ ) from .batched_packing import minimax_h3_batched_forward_kwargs +from .chunked_cpu_output import ( + MINIMAX_H3_CHUNKED_CPU_MP4_OUTPUT_ENV, + MINIMAX_H3_CHUNKED_CPU_MP4_ROUTE, + validate_chunked_cpu_mp4_runtime, +) from .condition_noise import ( minimax_h3_audio_cond_noise_aug_rows, minimax_h3_imgvid_cond_noise_aug_rows, @@ -221,6 +228,141 @@ ], } +# TAEH3 handles only decode, so the FL2VA transformer, audio VAE and optional +# conditioner still come from the base checkpoint. Deliberately omit the +# multi-gigabyte full video VAE from remote snapshot downloads in this mode. + + +def _minimax_h3_env_enabled(name: str, default: bool = False) -> bool: + raw = os.environ.get(name) + if raw is None: + return default + return raw.strip().lower() not in {"0", "false", "no", "off", "disable", "disabled"} + + +_MINIMAX_H3_GPU_UINT8_OUTPUT_ENV = "VLLM_OMNI_MINIMAX_H3_GPU_UINT8_OUTPUT" +_MINIMAX_H3_UINT8_BCTHW_MARKER = "__minimax_h3_uint8_bcthw_v1__" +_MINIMAX_H3_NVTX_DENOISE_STEP_ENV = "VLLM_OMNI_MINIMAX_H3_NVTX_DENOISE_STEP" +_MINIMAX_H3_NVTX_DECODE_REQUEST_ENV = "VLLM_OMNI_MINIMAX_H3_NVTX_DECODE_REQUEST" +_MINIMAX_H3_NVTX_DOMAIN = "vllm_omni.minimax_h3" +_MINIMAX_H3_QKV_E4M3_TRANSPORT_ENV = "VLLM_OMNI_FLASHINFER_ULYSSES_QKV_E4M3_TRANSPORT" +_MINIMAX_H3_VIDEO_PROFILER_KEY = "MiniMaxH3Pipeline.video_vae.decode_latent_to_chunked_cpu_mp4" +_MINIMAX_H3_AUDIO_PROFILER_KEY = "MiniMaxH3Pipeline.audio_vae.decode_latent" + + +@dataclass(frozen=True) +class _MiniMaxH3AudioDecodeTicket: + stream: Any + ready_event: Any + start_event: Any | None + done_event: Any + audio: torch.Tensor + + +def _minimax_h3_profiler_targets(*, audio_decode_overlap: bool) -> list[str]: + targets = [ + "encode_prompt", + "text_encoder.forward", + "diffuse", + "decode", + "denoise_step", + "video_vae.decode_latent", + "video_vae.decode_latent_to_chunked_cpu_mp4", + "audio_vae.decode_latent", + ] + if audio_decode_overlap: + # The generic profiler synchronizes the whole device around every + # target. That would serialize the two decoder streams. The overlap + # path records equivalent stage keys from route-local CUDA events. + targets.remove("video_vae.decode_latent_to_chunked_cpu_mp4") + targets.remove("audio_vae.decode_latent") + return targets + + +def _minimax_h3_runtime_gpu_architecture(device: torch.device) -> str: + if device.type != "cuda": + return device.type + capability = current_omni_platform.get_device_capability( + device.index if device.index is not None else torch.accelerator.current_device_index() + ) + if capability is None: + raise RuntimeError("The active platform did not report a GPU architecture") + major, minor = capability + return f"sm{major}{minor}" + + +def _minimax_h3_video_codec_options(sampling: Any) -> dict[str, str] | None: + """Validate the API codec override consumed inside the worker route. + + ``None`` deliberately means the normal API default + ``preset=ultrafast,threads=0``; the sink applies that exact default. An + explicit request mapping replaces the default, matching ``serving_video``. + """ + + raw = (getattr(sampling, "extra_args", None) or {}).get("video_codec_options") + if raw is None: + return None + if not isinstance(raw, Mapping) or any( + not isinstance(key, str) or not isinstance(value, str) for key, value in raw.items() + ): + raise OmniClientError("MiniMax H3 video_codec_options must be a string-to-string mapping") + return dict(raw) + + +def _parse_positive_nvtx_selector(raw: str | None, *, env_name: str) -> int | None: + if raw is None: + return None + value = raw.strip().lower() + if value in {"", "0", "false", "no", "off", "disable", "disabled"}: + return None + try: + selected = int(value) + except ValueError as exc: + raise ValueError(f"{env_name} must be a positive one-based integer or 0/off, got {raw!r}") from exc + if selected < 1: + raise ValueError(f"{env_name} must be a positive one-based integer or 0/off, got {raw!r}") + return selected + + +def _parse_minimax_h3_nvtx_denoise_step(raw: str | None) -> int | None: + """Parse the one-based denoise step selected for an NVTX capture.""" + return _parse_positive_nvtx_selector(raw, env_name=_MINIMAX_H3_NVTX_DENOISE_STEP_ENV) + + +def _parse_minimax_h3_nvtx_decode_request(raw: str | None) -> int | None: + """Parse the one-based decode invocation selected for an NVTX capture.""" + return _parse_positive_nvtx_selector(raw, env_name=_MINIMAX_H3_NVTX_DECODE_REQUEST_ENV) + + +def _push_minimax_h3_nvtx_range(name: str) -> None: + # Nsight capture-range triggers ignore dynamically supplied NVTX messages + # by default. The NVIDIA `nvtx` package caches messages as registered + # strings, unlike torch.cuda.nvtx.range_push(str). Keep the import lazy so + # the disabled production path has no new import or runtime dependency. + import nvtx + + nvtx.push_range(name, domain=_MINIMAX_H3_NVTX_DOMAIN) + + +def _pop_minimax_h3_nvtx_range() -> None: + import nvtx + + nvtx.pop_range(domain=_MINIMAX_H3_NVTX_DOMAIN) + + +@contextmanager +def _minimax_h3_nvtx_stage(enabled: bool, name: str): + """Publish a nested decode stage only for a selected diagnostic request.""" + + if not enabled: + yield + return + _push_minimax_h3_nvtx_range(name) + try: + yield + finally: + _pop_minimax_h3_nvtx_range() + def _resolve_minimax_h3_text_encoder_quant_config( quant_config: QuantizationConfig | None, @@ -332,6 +474,93 @@ def _read_base_schedule(release: Mapping[str, Any]) -> DMD2SigmaSchedule | None: return DMD2SigmaSchedule.from_metadata(release) +def _minimax_h3_uint8_bcthw_output(video: torch.Tensor) -> dict[str, Any]: + """Consume a decoded BCTHW video and return its compact transport form. + + The reference media path first converts the decoded tensor to FP32, clips + it to ``[0, 1]``, multiplies by 255, and applies NumPy ``rint`` before the + uint8 cast. ``torch.round`` has the same round-to-nearest-even contract. + Keeping the arithmetic in-place avoids another full-resolution FP32 + allocation. Keeping the compact tensor in BCTHW also avoids a GPU relayout; + after SHM, postprocess exposes a zero-copy BTHWC view whose individual + channel planes are contiguous for the direct-planar encoder. + + This helper intentionally consumes ``video``. The caller invokes it only + after the VAE outputs have been cropped/concatenated and will not use the + floating-point tensor again. + """ + if video.ndim != 5 or int(video.shape[1]) != 3: + raise ValueError(f"MiniMax H3 expected decoded BCTHW RGB video, got {tuple(video.shape)}") + if not video.dtype.is_floating_point: + raise TypeError(f"MiniMax H3 expected floating-point decoded video, got {video.dtype}") + + work = video.detach() + if work.dtype != torch.float32: + work = work.float() + if not work.is_contiguous(): + work = work.contiguous() + work.clamp_(0.0, 1.0).mul_(255.0).round_() + + compact = work.to(torch.uint8) + + reference_fp32_nbytes = int(compact.numel()) * torch.empty((), dtype=torch.float32).element_size() + payload_nbytes = int(compact.numel()) * compact.element_size() + return { + _MINIMAX_H3_UINT8_BCTHW_MARKER: True, + "layout": "BCTHW", + "frames": compact, + "reference_fp32_nbytes": reference_fp32_nbytes, + "payload_nbytes": payload_nbytes, + "saved_nbytes": reference_fp32_nbytes - payload_nbytes, + } + + +def _minimax_h3_prepare_video_transport( + video: Any, + *, + enabled: bool, + output_type: str | None, +) -> Any: + """Optionally compact an ordinary decoded video before worker IPC.""" + if not enabled or not isinstance(video, torch.Tensor): + return video + if output_type not in (None, "np"): + return video + packed = _minimax_h3_uint8_bcthw_output(video) + logger.info( + "MINIMAX_H3_GPU_UINT8_OUTPUT marker=v1 layout=BCTHW reference_fp32_bytes=%d payload_bytes=%d saved_bytes=%d", + packed["reference_fp32_nbytes"], + packed["payload_nbytes"], + packed["saved_nbytes"], + ) + return packed + + +def _minimax_h3_uint8_bcthw_frames(video: Any) -> torch.Tensor | None: + """Validate and unwrap a compact MiniMax H3 transport marker.""" + if not isinstance(video, dict) or video.get(_MINIMAX_H3_UINT8_BCTHW_MARKER) is not True: + return None + frames = video.get("frames") + if not isinstance(frames, torch.Tensor): + raise TypeError("MiniMax H3 uint8 BCTHW marker is missing its tensor payload") + if frames.dtype != torch.uint8 or frames.ndim != 5 or int(frames.shape[1]) != 3: + raise ValueError( + "MiniMax H3 uint8 BCTHW marker requires a rank-5 uint8 RGB tensor, " + f"got shape={tuple(frames.shape)} dtype={frames.dtype}" + ) + if video.get("layout") != "BCTHW": + raise ValueError(f"MiniMax H3 uint8 transport has unsupported layout {video.get('layout')!r}") + payload_nbytes = int(frames.numel()) * frames.element_size() + reference_fp32_nbytes = int(frames.numel()) * torch.empty((), dtype=torch.float32).element_size() + if int(video.get("payload_nbytes", -1)) != payload_nbytes: + raise ValueError("MiniMax H3 uint8 transport payload byte counter does not match its tensor") + if int(video.get("reference_fp32_nbytes", -1)) != reference_fp32_nbytes: + raise ValueError("MiniMax H3 uint8 transport FP32 byte counter does not match its tensor") + if int(video.get("saved_nbytes", -1)) != reference_fp32_nbytes - payload_nbytes: + raise ValueError("MiniMax H3 uint8 transport saved-byte counter is inconsistent") + return frames + + def resolve_minimax_h3_diffusion_model_path( model: str, revision: str | None, @@ -389,8 +618,19 @@ def _minimax_h3_post_process(output, output_type: str = "np"): ) if output_type == "latent": return output + if isinstance(video, (bytes, bytearray, memoryview)): + return { + "video": bytes(video), + "audio": None, + "audio_sample_rate": MINIMAX_H3_AUDIO_SAMPLE_RATE, + "fps": MINIMAX_H3_FPS, + } if output_type == "np": - video = video.detach().cpu().numpy() + compact_frames = _minimax_h3_uint8_bcthw_frames(video) + if compact_frames is None: + video = video.detach().float().cpu().permute(0, 2, 3, 4, 1).clamp(0, 1).numpy() + else: + video = compact_frames.detach().cpu().permute(0, 2, 3, 4, 1).numpy() audio = audio.detach().float().cpu().numpy() video = [sample for sample in video] return { @@ -622,30 +862,33 @@ def _broadcast_tensor( *, dtype: torch.dtype, device: torch.device, + source_rank: int = 0, ) -> torch.Tensor: group, rank, world_size = _dit_rank_world() + if not 0 <= source_rank < world_size: + raise ValueError(f"source rank {source_rank} is outside DiT world size {world_size}") 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 rank == source_rank: if tensor is None: - raise ValueError("rank 0 must provide a tensor to broadcast") + raise ValueError(f"rank {source_rank} 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) + dist.broadcast(shape, src=source_rank, group=group) ndim = int(shape[0].item()) tensor_shape = tuple(int(v) for v in shape[1 : ndim + 1].tolist()) - if rank == 0: + if rank == source_rank: 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) + dist.broadcast(output, src=source_rank, group=group) return output @@ -695,7 +938,9 @@ class MiniMaxH3Pipeline( "encode_prompt", "_encode_visual_conditions", "_encode_reference_audio_conditions", + "_build_denoise_inputs", "diffuse", + "_unpack_denoised_rows", "decode", "video_vae.decode_latent", "audio_vae.decode_latent", @@ -709,6 +954,8 @@ class MiniMaxH3Pipeline( _base_schedule_by_partition: ClassVar[Mapping[str, DMD2SigmaSchedule | None]] = {} # Set from --lora-path during construction; absent means no FastH3 adapter. _fasth3: FastH3WeightFusion | None = None + # The benchmark-only E4M3 transport artifact is bound during construction + # and installed only after FastH3's streamed-weight proof has closed. def _load_diffusion_lora_adapter( self, @@ -905,11 +1152,37 @@ def __init__( ) -> None: del prefix super().__init__() + from .ultra import configure_ultra + + configure_ultra() 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() + self._chunked_cpu_mp4_output_enabled = _minimax_h3_env_enabled( + MINIMAX_H3_CHUNKED_CPU_MP4_OUTPUT_ENV, + ) + self._gpu_uint8_output_enabled = _minimax_h3_env_enabled( + _MINIMAX_H3_GPU_UINT8_OUTPUT_ENV, + ) + enabled_output_routes = sum( + ( + self._chunked_cpu_mp4_output_enabled, + self._gpu_uint8_output_enabled, + ) + ) + if enabled_output_routes > 1: + raise ValueError( + "MiniMax H3 output routes are mutually exclusive; enable only one of " + f"{MINIMAX_H3_CHUNKED_CPU_MP4_OUTPUT_ENV}, or " + f"{_MINIMAX_H3_GPU_UINT8_OUTPUT_ENV}" + ) + self._nvtx_denoise_step = _parse_minimax_h3_nvtx_denoise_step(os.environ.get(_MINIMAX_H3_NVTX_DENOISE_STEP_ENV)) + self._nvtx_decode_request = _parse_minimax_h3_nvtx_decode_request( + os.environ.get(_MINIMAX_H3_NVTX_DECODE_REQUEST_ENV) + ) + self._decode_invocations = 0 self.load_text_encoder = od_config.model_loaded["text_encoder"] self.partition = _minimax_h3_partition_for_task( getattr(od_config, "task_type", None), @@ -984,14 +1257,29 @@ def __init__( od_config.quantization_config, "transformer", ) + cache_config = getattr(od_config, "cache_config", None) + adaln_cache_path = getattr( + cache_config, + "minimax_h3_adaln_cache_path", + None, + ) self.transformer = MiniMaxH3DiTModel( od_config, quant_config=transformer_quant_config, + adaln_cache_path=adaln_cache_path, + adaln_cache_model_variant=expected_partition, ) if ref2va_model_path is not None: + ref_adaln_cache_path = getattr( + cache_config, + "minimax_h3_ref_adaln_cache_path", + None, + ) self.transformers_ref = MiniMaxH3DiTModel( od_config, quant_config=transformer_quant_config, + adaln_cache_path=ref_adaln_cache_path, + adaln_cache_model_variant="ref2va", ) self._fasth3 = resolve_fasth3_fusion(od_config, self.transformer) @@ -1010,6 +1298,8 @@ def __init__( audio_shift=self.default_audio_shift, ) + self._bind_fasth3_adaln_cache_contract(cache_config) + if self.load_text_encoder: self.tokenizer = Qwen2TokenizerFast.from_pretrained( str(model_path), @@ -1027,6 +1317,20 @@ def __init__( _, rank, dit_world = _dit_rank_world() self._dit_rank = rank + self._full_vae_audio_overlap_enabled = os.environ.get("VLLM_OMNI_H3_FULL_VAE_AUDIO_OVERLAP") == "1" + if self._full_vae_audio_overlap_enabled and ( + not self._chunked_cpu_mp4_output_enabled + or self.device.type != "cuda" + or dit_world != 8 + or any( + bool(getattr(od_config, k, False)) + for k in ("enable_cpu_offload", "enable_layerwise_offload", "enable_distributed_layerwise_offload") + ) + ): + raise RuntimeError("Full-VAE/audio overlap requires the resident CUDA SP8 chunked-MP4 route") + self._audio_decode_stream = None + if self._full_vae_audio_overlap_enabled and rank == 0: + self._audio_decode_stream = torch.get_device_module().Stream(device=self.device) if self.load_text_encoder: text_encoder_tp_size = int(getattr(self.parallel_config, "text_encoder_tp_size", 1)) if text_encoder_tp_size < 1: @@ -1069,15 +1373,13 @@ def __init__( legacy_manual_components = getattr(od_config, "diffusion_offload_config", None) is None and bool( od_config.enable_layerwise_offload or getattr(od_config, "enable_distributed_layerwise_offload", False) ) - # Preserve the legacy MiniMax-H3 low-residency path. The compact API - # deliberately limits explicit component selection to dit/text_encoder, - # so VAEs stay resident for new configurations. component_load_device = torch.device("cpu") if legacy_manual_components else self.device self.video_vae = MiniMaxH3VideoVAE( os.path.join(model_path, "video_vae"), device=self.device, load_device=component_load_device, ) + self._vae_modules = ["video_vae", "audio_vae"] self.audio_vae = MiniMaxH3AudioVAE( os.path.join(model_path, "audio_vae"), device=self.device, @@ -1085,6 +1387,26 @@ def __init__( ) # Registry-side VAE patch-parallel discovery uses ``pipeline.vae``. self.vae = self.video_vae + if self._chunked_cpu_mp4_output_enabled: + if self.video_vae is not None: + try: + self.video_vae.validate_chunked_output() + except RuntimeError as exc: + raise RuntimeError( + f"{MINIMAX_H3_CHUNKED_CPU_MP4_OUTPUT_ENV}=1 has an " + "incompatible MiniMax H3 video decoder output contract" + ) from exc + validate_chunked_cpu_mp4_runtime() + logger.info( + "MiniMax H3 output route enabled: route=%s gate=%s worker_transport=final_mp4_bytes", + MINIMAX_H3_CHUNKED_CPU_MP4_ROUTE, + MINIMAX_H3_CHUNKED_CPU_MP4_OUTPUT_ENV, + ) + if self._gpu_uint8_output_enabled: + logger.info( + "MiniMax H3 exact GPU uint8 BCTHW worker transport enabled via %s", + _MINIMAX_H3_GPU_UINT8_OUTPUT_ENV, + ) self._dlo_component_cache = None offloads_text_encoder = should_offload_component(od_config, TEXT_ENCODER_COMPONENT) @@ -1105,7 +1427,10 @@ def __init__( self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self) self.setup_diffusion_pipeline_profiler( - enable_diffusion_pipeline_profiler=(od_config.enable_diffusion_pipeline_profiler) + profiler_targets=_minimax_h3_profiler_targets( + audio_decode_overlap=False, + ), + enable_diffusion_pipeline_profiler=(od_config.enable_diffusion_pipeline_profiler), ) def load_weights( @@ -1135,6 +1460,16 @@ def source_prefix(item: tuple[str, torch.Tensor]) -> str: loaded = component.load_weights(stream) if prefix == "transformer.": transformer_loaded = set(loaded) + from .quantization import dit_mxfp8_mode + + if dit_mxfp8_mode() is not None: + if self._fasth3 is None: + raise ValueError("H3 MXFP8 requires the complete FastH3 student") + required_patches = set(self._fasth3._patches) - self._fasth3._sidecar_satisfied + if not required_patches.issubset(transformer_loaded): + raise ValueError("FastH3 fused patches were not consumed by the native loader") + self._fasth3.validate_fully_applied(transformer_loaded) + component._h3_student_fusion_complete = True if prefix != "text_encoder.": component.post_load_weights() loaded_with_prefix.update(prefix + name for name in loaded) @@ -1148,8 +1483,8 @@ def source_prefix(item: tuple[str, torch.Tensor]) -> str: continue loaded_with_prefix.update(f"{component_name}.{name}" for name, _ in component.named_parameters()) if self._fasth3 is not None: - # load_weights only warns on a parameter the model does not have, so - # close the adapter against what the DiT actually consumed. + # Gate tensors are assignments into dynamically-created modules; + # close the contract against names the model actually consumed. self._fasth3.validate_fully_applied(transformer_loaded) return loaded_with_prefix @@ -1158,6 +1493,27 @@ def lora_is_fused(self) -> bool: """True when --lora-path was consumed as a load-time weight fusion.""" return self._fasth3 is not None + def _bind_fasth3_adaln_cache_contract(self, cache_config: Any) -> None: + """Accept only an AdaLN sidecar verified against this FastH3 adapter.""" + if self._fasth3 is None: + return + cache = getattr(self.transformer, "adaln_cache", None) + if cache is None: + return + from .adaln_cache import prevalidate_minimax_h3_fasth3_adaln_sidecar + + binding = prevalidate_minimax_h3_fasth3_adaln_sidecar( + cache.path, + adapter_path=self._fasth3.source, + model_variant="fl2va", + mode="t2va", + base_schedule=self._fasth3.base_schedule, + flow_shift=self.default_video_shift, + audio_flow_shift=self.default_audio_shift, + ) + cache.bind_runtime_contract(binding) + self._fasth3.mark_sidecar_satisfied(cache.sidecar_satisfied_parameter_names()) + def _transformer_for_task(self, task: str) -> MiniMaxH3DiTModel: if task == "ref2va" and hasattr(self, "transformers_ref"): return self.transformers_ref @@ -1908,6 +2264,52 @@ def _resident_dit_layers_on_device(self, *, enabled: bool = True): if controller is not None and enabled: controller.offload_resident_layers() + @contextmanager + def _profile_denoise_iteration(self, step: int, num_steps: int): + """Profile a DiT update and optionally mark one complete step with NVTX.""" + profile_enabled = getattr(self, "enable_diffusion_pipeline_profiler", False) + one_based_step = step + 1 + nvtx_enabled = ( + getattr(self, "_dit_rank", 0) == 0 and getattr(self, "_nvtx_denoise_step", None) == one_based_step + ) + if not profile_enabled and not nvtx_enabled: + yield + return + + current_omni_platform.synchronize() + start_time = time.perf_counter() if profile_enabled else None + if nvtx_enabled: + range_name = f"minimax_h3.denoise.step_{one_based_step:02d}" + logger.info( + "[MiniMaxH3NVTX] capturing step=%d/%d range=%s", + one_based_step, + num_steps, + range_name, + ) + _push_minimax_h3_nvtx_range(range_name) + try: + yield + finally: + try: + # Keep every GPU launch belonging to this update inside the + # selected NVTX range before Nsight ends and terminates the run. + current_omni_platform.synchronize() + finally: + if nvtx_enabled: + _pop_minimax_h3_nvtx_range() + + if profile_enabled: + assert start_time is not None + duration = time.perf_counter() - start_time + logger.info( + "[MiniMaxH3StepProfiler] step=%d/%d took %.6fs", + one_based_step, + num_steps, + duration, + ) + with self._profiler_lock: + self._stage_durations[f"MiniMaxH3Pipeline.diffuse.step_{one_based_step:02d}"] = duration + def _build_denoise_inputs( self, *, @@ -2149,6 +2551,7 @@ def diffuse( 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(), + step_profiler=lambda step: self._profile_denoise_iteration(step, len(inputs["sigmas_video"]) - 1), ) return self._unpack_denoised_rows( @@ -2161,99 +2564,289 @@ def diffuse( audio_t=audio_t, ) - def decode_to_mp4( + def _start_audio_decode_overlap( + self, + audio_latent: torch.Tensor, + ) -> _MiniMaxH3AudioDecodeTicket: + stream = self._audio_decode_stream + producer_stream = torch.get_device_module().current_stream(device=audio_latent.device) + profile_enabled = bool(getattr(self, "enable_diffusion_pipeline_profiler", False)) + ready_event = torch.get_device_module().Event(enable_timing=False) + # The completion event is both the consumer dependency and, when the + # profiler is active, the elapsed-time endpoint. CUDA requires both + # endpoints passed to elapsed_time() to be timing-enabled. + done_event = torch.get_device_module().Event(enable_timing=profile_enabled) + start_event = torch.get_device_module().Event(enable_timing=True) if profile_enabled else None + ready_event.record(producer_stream) + try: + with torch.get_device_module().stream(stream): + stream.wait_event(ready_event) + if start_event is not None: + start_event.record(stream) + audio_latent.record_stream(stream) + with self._component_on_device(self.audio_vae): + audio = self.audio_vae.decode_latent(audio_latent) + if not isinstance(audio, torch.Tensor): + raise TypeError("MiniMax H3 audio VAE must return a CUDA tensor for audio/video overlap") + if audio.device.type != "cuda": + raise RuntimeError("MiniMax H3 audio VAE returned a non-CUDA tensor during audio/video overlap") + audio.record_stream(stream) + done_event.record(stream) + except BaseException: + try: + stream.synchronize() + except BaseException: + logger.exception("Failed to drain MiniMax H3 audio decode stream after launch failure") + raise + return _MiniMaxH3AudioDecodeTicket( + stream=stream, + ready_event=ready_event, + start_event=start_event, + done_event=done_event, + audio=audio, + ) + + def _finish_audio_decode_overlap( + self, + ticket: _MiniMaxH3AudioDecodeTicket, + ) -> torch.Tensor: + consumer_stream = torch.get_device_module().current_stream(device=ticket.audio.device) + consumer_stream.wait_event(ticket.done_event) + ticket.audio.record_stream(consumer_stream) + return ticket.audio + + @staticmethod + def _drain_audio_decode_overlap( + ticket: _MiniMaxH3AudioDecodeTicket | None, + ) -> None: + if ticket is None: + return + try: + ticket.done_event.synchronize() + except BaseException: + try: + ticket.stream.synchronize() + except BaseException: + logger.exception("Failed to drain MiniMax H3 audio decode stream after request failure") + + def _record_audio_decode_overlap_profile( + self, + *, + ticket: _MiniMaxH3AudioDecodeTicket, + video_start_event: Any | None, + video_done_event: Any | None, + ) -> None: + if not getattr(self, "enable_diffusion_pipeline_profiler", False): + return + if ticket.start_event is None or video_start_event is None or video_done_event is None: + raise RuntimeError("MiniMax H3 audio/video overlap profiler events are incomplete") + # Synchronizing these route-local completion events preserves overlap; + # unlike the generic profiler it never fences before either decoder. + video_done_event.synchronize() + ticket.done_event.synchronize() + durations = { + _MINIMAX_H3_VIDEO_PROFILER_KEY: float(video_start_event.elapsed_time(video_done_event)) / 1000.0, + _MINIMAX_H3_AUDIO_PROFILER_KEY: float(ticket.start_event.elapsed_time(ticket.done_event)) / 1000.0, + } + for name, duration in durations.items(): + logger.info("[DiffusionPipelineProfiler] %s took %.6fs", name, duration) + with self._profiler_lock: + for name, duration in durations.items(): + self._stage_durations[name] = self._stage_durations.get(name, 0.0) + duration + + def decode( self, video_latent: torch.Tensor, audio_latent: torch.Tensor, *, height: int, width: int, - max_pending: int = 2, - batch_frames: int = 17, + fps: int = MINIMAX_H3_FPS, + use_chunked_cpu_mp4: bool = False, video_codec_options: dict[str, str] | None = None, - ) -> bytes: - """Decode and encode one output on the worker without full-video materialization. + ) -> tuple[Any, Any]: + """Decode one request, optionally exposing a clean Nsight range. - Audio is decoded first so the incremental mux session can attach its audio - stream before temporal video chunks arrive. The callback receives committed - float32 ``BCTHW`` frames, performs the requested-size crop and uint8 - conversion in the worker, then applies bounded backpressure to the encoder. + The selector is an invocation number rather than a boolean so a runner + can warm the model once and capture only the measured request. The + default-disabled path does not synchronize or import NVTX. """ - from vllm_omni.diffusion.utils.media_utils import ChunkedMP4Encoder - - if batch_frames <= 0: - raise ValueError("batch_frames must be positive") - - with self._component_on_device(self.audio_vae): - audio = self.audio_vae.decode_latent(audio_latent) - audio_np = audio.detach().float().cpu().numpy() - if audio_np.ndim == 3 and audio_np.shape[0] == 1: - audio_np = audio_np[0] - encoder = ChunkedMP4Encoder( - width=width, - height=height, - fps=MINIMAX_H3_FPS, - audio_waveform=audio_np, - audio_sample_rate=MINIMAX_H3_AUDIO_SAMPLE_RATE, - max_pending=max_pending, - video_codec_options=video_codec_options, - ) - - pending_chunks: list[torch.Tensor] = [] - pending_frames = 0 - - def flush_pending() -> None: - nonlocal pending_frames - if not pending_chunks: - return - batched = torch.cat(pending_chunks, dim=1) - encoder.push(batched[0].cpu().numpy()) - pending_chunks.clear() - pending_frames = 0 - - def on_chunk(frames: torch.Tensor) -> None: - nonlocal pending_frames - prepared = _prepare_minimax_h3_video_output(frames[..., :height, :width]) - if prepared.shape[0] != 1: - raise ValueError("MiniMax H3 chunked MP4 encoding currently expects one output per decoder") - pending_chunks.append(prepared) - pending_frames += int(prepared.shape[1]) - if pending_frames >= batch_frames: - flush_pending() + self._decode_invocations = getattr(self, "_decode_invocations", 0) + 1 + selected_request = getattr(self, "_nvtx_decode_request", None) + nvtx_enabled = getattr(self, "_dit_rank", 0) == 0 and selected_request == self._decode_invocations + if nvtx_enabled: + range_name = f"minimax_h3.decode.request_{self._decode_invocations:02d}" + current_omni_platform.synchronize() + logger.info( + "[MiniMaxH3NVTX] capturing decode request=%d range=%s", + self._decode_invocations, + range_name, + ) + _push_minimax_h3_nvtx_range(range_name) try: - with self._component_on_device(self.video_vae): - with current_omni_platform.create_autocast_context( - device_type=self.device.type, - dtype=torch.float16, - enabled=True, - ): - self.video_vae.decode_with_chunks(video_latent, on_chunk=on_chunk) - flush_pending() - return encoder.finish() - except BaseException: - encoder.abort() - raise + return self._decode_impl( + video_latent, + audio_latent, + height=height, + width=width, + fps=fps, + use_chunked_cpu_mp4=use_chunked_cpu_mp4, + video_codec_options=video_codec_options, + ) + finally: + if nvtx_enabled: + current_omni_platform.synchronize() + _pop_minimax_h3_nvtx_range() - def decode( + def _decode_impl( self, video_latent: torch.Tensor, audio_latent: torch.Tensor, *, height: int, width: int, - ) -> tuple[torch.Tensor, torch.Tensor]: - with self._component_on_device(self.video_vae): - 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() - with self._component_on_device(self.audio_vae): - audio = self.audio_vae.decode_latent(audio_latent) - return video, audio + fps: int = MINIMAX_H3_FPS, + use_chunked_cpu_mp4: bool = False, + video_codec_options: dict[str, str] | None = None, + ) -> tuple[Any, Any]: + # Both request execution and step execution enter VAE decode through + # this method on every rank. Keep the call itself rank-unconditional: + # FlashInfer communicator teardown is collective. The helper remains a + # default-off no-op unless its request-scoped lifecycle was selected. + from vllm_omni.diffusion.distributed.flashinfer_ulysses import ( + release_flashinfer_ulysses_after_denoise, + ) + + release_flashinfer_ulysses_after_denoise() + chunked_cpu_output = None + audio_overlap_ticket = None + video_profile_start_event = None + video_profile_done_event = None + # The multiprocess executor consumes rank 0 for this DiT replica. The + # native video VAE still has to run through ``decode_base`` on every + # rank because its tile path contains collectives, but only this rank + # needs a full RGB result or an MP4 sink. + return_output = getattr(self, "_dit_rank", 0) == 0 + nvtx_stages = getattr(self, "_nvtx_decode_request", None) == getattr(self, "_decode_invocations", 0) + try: + if self._full_vae_audio_overlap_enabled and return_output: + if ( + self._uses_manual_component_offload(self.audio_vae) + or getattr(self, "_model_cpu_offload_modules", None) + or audio_latent.device != video_latent.device + or (audio_latent.device.type != "cuda") + or torch.cuda.is_current_stream_capturing() + or getattr(self, "enable_diffusion_pipeline_profiler", False) + or (self._audio_decode_stream is None) + ): + raise RuntimeError("Full-VAE/audio overlap runtime contract changed") + audio_overlap_ticket = self._start_audio_decode_overlap(audio_latent) + logger.info( + "H3_FULL_VAE_AUDIO_START %s", + json.dumps( + dict(rank=0, request=self._decode_invocations, auxiliary_stream=True, before_video_decode=True) + ), + ) + video_vae = self.video_vae + if video_vae is None: + raise RuntimeError("MiniMax H3 full video VAE is not loaded") + with self._component_on_device(video_vae): + with current_omni_platform.create_autocast_context( + device_type=self.device.type, + dtype=torch.float16, + enabled=True, + ): + if use_chunked_cpu_mp4: + chunked_cpu_output = video_vae.decode_latent_to_chunked_cpu_mp4( + video_latent, + height=height, + width=width, + fps=fps, + audio_sample_rate=MINIMAX_H3_AUDIO_SAMPLE_RATE, + video_codec_options=video_codec_options, + return_output=return_output, + ) + video = None + else: + video = video_vae.decode_latent( + video_latent, + return_output=return_output, + ) + # ``decode_latent`` gates only after all native video-VAE + # collectives have completed. The audio VAE is rank-local (it + # contains no distributed collectives), so non-output ranks can + # now finish with a lightweight but valid worker result. + if not return_output: + if chunked_cpu_output is not None: + chunked_cpu_output.abort() + return None, None + audio = None + if video is not None: + video = video[..., :height, :width].contiguous() + + # Encode finalized video chunks on a host thread during audio decode. + # The video VAE can be offloaded after submission. + if audio_overlap_ticket is None: + with self._component_on_device(self.audio_vae): + audio = self.audio_vae.decode_latent(audio_latent) + else: + audio = self._finish_audio_decode_overlap(audio_overlap_ticket) + if self._full_vae_audio_overlap_enabled: + logger.info( + "H3_FULL_VAE_AUDIO_JOIN %s", + json.dumps(dict(rank=0, request=self._decode_invocations, consumer_wait_event=True)), + ) + self._record_audio_decode_overlap_profile( + ticket=audio_overlap_ticket, + video_start_event=video_profile_start_event, + video_done_event=video_profile_done_event, + ) + + if use_chunked_cpu_mp4: + if chunked_cpu_output is None: + raise RuntimeError("MiniMax H3 chunked CPU MP4 output rank has no output sink") + with _minimax_h3_nvtx_stage(nvtx_stages, "minimax_h3.decode.mp4_finish"): + video, output_stats = chunked_cpu_output.finish( + audio, + MINIMAX_H3_AUDIO_SAMPLE_RATE, + ) + logger.info( + "MiniMax H3 output complete: route=%s frames=%d " + "rgb_u8=%.2fMiB mp4=%.2fMiB slots=%d max_inflight=%d " + "backpressure=%.3fs d2h_wait=%.3fs video_encode_mux=%.3fs " + "audio_finalize=%.3fs audio_preencoded=%s " + "audio_packets=%.2fKiB pipeline=%.3fs", + MINIMAX_H3_CHUNKED_CPU_MP4_ROUTE, + int(output_stats["chunked_cpu_mp4_frames"]), + output_stats["chunked_cpu_mp4_rgb_u8_bytes"] / (1024 * 1024), + output_stats["chunked_cpu_mp4_bytes"] / (1024 * 1024), + int(output_stats["chunked_cpu_mp4_queue_slots"]), + int(output_stats["chunked_cpu_mp4_max_inflight"]), + output_stats["chunked_cpu_mp4_producer_backpressure_s"], + output_stats["chunked_cpu_mp4_d2h_wait_s"], + output_stats["chunked_cpu_mp4_video_encode_mux_s"], + output_stats["chunked_cpu_mp4_audio_mux_finalize_s"], + bool(output_stats["chunked_cpu_mp4_audio_preencoded"]), + output_stats["chunked_cpu_mp4_audio_packet_bytes"] / 1024, + output_stats["chunked_cpu_mp4_pipeline_s"], + ) + if getattr(self, "enable_diffusion_pipeline_profiler", False): + with self._profiler_lock: + self._stage_durations.update( + { + f"MiniMaxH3Pipeline.{name}": value + for name, value in output_stats.items() + if name.endswith("_s") + } + ) + return video, None + return video, audio + except BaseException: + self._drain_audio_decode_overlap(audio_overlap_ticket) + if chunked_cpu_output is not None: + chunked_cpu_output.abort() + raise @staticmethod def _extract_prompt(raw_prompt: Any) -> tuple[str, dict[str, Any]]: @@ -2378,6 +2971,10 @@ def _prepare_request_inputs( image = images[0] if images else None height, width, num_frames, latent_t, audio_t = self._resolve_shape(task, sampling, image) + # ``_resolve_shape`` validates that MiniMax H3 uses its fixed output + # frame rate. Keep the resolved value in the request context for the + # decode/output path instead of relying on an undefined local. + fps = int(sampling.fps or MINIMAX_H3_FPS) if task == "fl2va": for item in images: _validate_reference_image(item) @@ -2572,6 +3169,7 @@ def _prepare_request_inputs( "task": task, "height": height, "width": width, + "fps": fps, "num_frames": num_frames, "latent_t": latent_t, "latent_h": height // 16, @@ -2620,37 +3218,66 @@ def forward(self, request: DiffusionRequestBatch) -> DiffusionOutput: ) denoise_kwargs = self._denoise_kwargs(context) num_outputs = context["num_outputs"] + use_chunked_cpu_mp4 = getattr( + self, + "_chunked_cpu_mp4_output_enabled", + False, + ) + if use_chunked_cpu_mp4 and num_outputs != 1: + raise OmniClientError( + f"{MINIMAX_H3_CHUNKED_CPU_MP4_OUTPUT_ENV}=1 supports exactly one output, got {num_outputs}" + ) + if use_chunked_cpu_mp4 and request.sampling_params.output_type not in ( + None, + "np", + ): + raise OmniClientError( + f"{MINIMAX_H3_CHUNKED_CPU_MP4_OUTPUT_ENV}=1 returns a finalized " + "MP4 and requires output_type=np (or unset)" + ) + video_codec_options = _minimax_h3_video_codec_options(request.sampling_params) if use_chunked_cpu_mp4 else None videos = [] audios = [] for output_seed in _minimax_h3_output_seeds(context["seed"], num_outputs): video_latent, audio_latent = self.diffuse(**{**denoise_kwargs, "seed": output_seed}) - if context["preencode_mp4"]: - videos.append( - self.decode_to_mp4( - video_latent, - audio_latent, - height=context["height"], - width=context["width"], - video_codec_options=context["video_codec_options"], - batch_frames=context["preencode_batch_frames"], - ) - ) - audios.append(None) - else: - video, audio = self.decode( - video_latent, - audio_latent, - height=context["height"], - width=context["width"], - ) - videos.append(_prepare_minimax_h3_video_output(video)) - audios.append(audio) - if videos and isinstance(videos[0], bytes): - video = videos[0] if len(videos) == 1 else videos + video, audio = self.decode( + video_latent, + audio_latent, + height=context["height"], + width=context["width"], + fps=context["fps"], + use_chunked_cpu_mp4=use_chunked_cpu_mp4, + video_codec_options=video_codec_options, + ) + videos.append(video) + audios.append(audio) + if getattr(self, "_dit_rank", 0) != 0: + # Non-primary ranks deliberately return lightweight placeholders. + # They already completed every collective in ``decode``; their + # result is not selected by the multiprocess executor. + video = None + audio = None + elif use_chunked_cpu_mp4: + video = videos[0] audio = None else: - video = videos[0] if len(videos) == 1 else torch.cat(videos, dim=0) - audio = audios[0] if len(audios) == 1 else torch.cat(audios, dim=0) + # Avoid copying an entire decoded clip for the overwhelmingly + # common single-output request. Besides the redundant device + # bandwidth, a 15-second H3 FP32 result is several GiB and the + # duplicate can consume the last available HBM immediately before + # compact output transport. Preserve concatenation for the true + # multi-output case. + if num_outputs == 1: + video = videos[0] + audio = audios[0] + else: + video = torch.cat(videos, dim=0) + audio = torch.cat(audios, dim=0) + video = _minimax_h3_prepare_video_transport( + video, + enabled=getattr(self, "_gpu_uint8_output_enabled", False), + output_type=request.sampling_params.output_type, + ) return DiffusionOutput( output=(video, audio), post_process_func=get_minimax_h3_post_process_func(self.od_config), @@ -2697,6 +3324,13 @@ def prepare_encode(self, state: StepRequestState, **kwargs: Any) -> StepRequestS raise OmniClientError( f"MiniMax H3 step execution produces one output per request, got num_outputs_per_prompt={num_outputs}" ) + if getattr(self, "_chunked_cpu_mp4_output_enabled", False): + if state.sampling.output_type not in (None, "np"): + raise OmniClientError( + f"{MINIMAX_H3_CHUNKED_CPU_MP4_OUTPUT_ENV}=1 returns a " + "finalized MP4 and requires output_type=np (or unset)" + ) + _minimax_h3_video_codec_options(state.sampling) if getattr(self, "_dlo_residency_controller", None) is not None: raise ValueError( "MiniMax H3 step execution is not compatible with distributed layerwise offload; " @@ -2771,6 +3405,7 @@ def prepare_encode(self, state: StepRequestState, **kwargs: Any) -> StepRequestS _STEP_SHAPE: { "height": context["height"], "width": context["width"], + "fps": context["fps"], "latent_t": context["latent_t"], "latent_h": context["latent_h"], "latent_w": context["latent_w"], @@ -2949,24 +3584,28 @@ def post_decode(self, state: StepRequestState, **kwargs: Any) -> DiffusionOutput latent_w=shape["latent_w"], audio_t=shape["audio_t"], ) - if shape.get("preencode_mp4", False): - video = self.decode_to_mp4( - video_latent, - audio_latent, - height=shape["height"], - width=shape["width"], - video_codec_options=shape.get("video_codec_options"), - batch_frames=shape.get("preencode_batch_frames", 17), - ) - audio = None - else: - video, audio = self.decode( - video_latent, - audio_latent, - height=shape["height"], - width=shape["width"], - ) - video = _prepare_minimax_h3_video_output(video) + video, audio = self.decode( + video_latent, + audio_latent, + height=shape["height"], + width=shape["width"], + fps=shape["fps"], + use_chunked_cpu_mp4=getattr( + self, + "_chunked_cpu_mp4_output_enabled", + False, + ), + video_codec_options=( + _minimax_h3_video_codec_options(state.sampling) + if getattr(self, "_chunked_cpu_mp4_output_enabled", False) + else None + ), + ) + video = _minimax_h3_prepare_video_transport( + video, + enabled=getattr(self, "_gpu_uint8_output_enabled", False), + output_type=state.sampling.output_type, + ) return DiffusionOutput( output=(video, audio), post_process_func=get_minimax_h3_post_process_func(self.od_config), diff --git a/vllm_omni/diffusion/models/minimax_h3/quantization.py b/vllm_omni/diffusion/models/minimax_h3/quantization.py new file mode 100644 index 00000000000..917263135db --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/quantization.py @@ -0,0 +1,263 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""H3 DiT and VAE weight-conversion policies over shared MXFP8 arithmetic. + +DiT conversion follows student fusion; VAE conversion preserves its original +FP32 weight source and FP16 activation boundary. The policies stay distinct. +""" + +from __future__ import annotations + +import collections +import json +import os + +import torch +from torch import nn +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + +from vllm_omni.diffusion.layers.mxfp8 import ( + _scale_numel, + mxfp8_linear, + mxfp8_quantize_project, + mxfp8_quantize_swizzled, + mxfp8_scaled_mm, + silu_mxfp8_linear, +) +from vllm_omni.platforms import current_omni_platform + +DIT_MXFP8_ENV = "VLLM_OMNI_H3_DIT_MXFP8" +SUFFIXES = ("attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2") +SHAPES = dict(zip(SUFFIXES, ((21504, 5376), (5376, 7168), (28672, 5376), (5376, 14336)))) +_QK_WIDTH = 14336 + + +def dit_mxfp8_mode() -> str | None: + value = os.getenv(DIT_MXFP8_ENV, "0") + if value not in ("0", "1"): + raise ValueError(f"{DIT_MXFP8_ENV} must be 0 or 1") + return "mxfp8" if value == "1" else None + + +class H3MXFP8Method(UnquantizedLinearMethod): + def apply(self, layer, x, bias=None): + if bias is not None: + raise ValueError("H3 MXFP8 main projections must be bias-free") + return mxfp8_linear(x, layer.weight, layer.mxfp8_weight_scale) + + +def quantize_layer(layer): + """Convert one BF16 projection after the student adapter has been fused.""" + weight = layer.weight + if ( + weight.dtype != torch.bfloat16 + or not weight.is_cuda + or weight.ndim != 2 + or weight.shape[0] % 16 + or weight.shape[1] % 32 + or layer.bias is not None + or type(layer.quant_method) is not UnquantizedLinearMethod + ): + raise ValueError("MXFP8 source must be a loaded CUDA BF16 unquantized projection") + if hasattr(layer, "mxfp8_weight_scale"): + raise ValueError("Duplicate MXFP8 weight quantization") + q, scale = mxfp8_quantize_swizzled(weight.detach().contiguous()) + layer.weight = torch.nn.Parameter(q, requires_grad=False) + layer.register_buffer("mxfp8_weight_scale", scale.view(torch.float8_e8m0fnu)) + layer.quant_method = H3MXFP8Method() + return layer + + +def install_dit_mxfp8(model) -> None: + if dit_mxfp8_mode() is None or getattr(model, "_h3_mxfp8_installed", False): + return + if not getattr(model, "_h3_student_fusion_complete", False): + raise ValueError("MXFP8 conversion requires completed FastH3 student fusion") + if getattr(model, "adaln_cache", None) is None or not model.adaln_cache._loaded: + raise ValueError("H3 Ultra requires the loaded exact AdaLN cache") + from vllm.distributed import get_tensor_model_parallel_world_size + + from vllm_omni.diffusion.distributed.parallel_state import get_sp_group + + if ( + len(model.blocks) != 50 + or current_omni_platform.get_device_capability(torch.accelerator.current_device_index()) != (12, 0) + or get_tensor_model_parallel_world_size() != 1 + or get_sp_group().world_size != 8 + ): + raise ValueError("H3 Ultra MXFP8 is qualified for the 50-block SM120 TP1/SP8 model") + names = {f"blocks.{i}.{suffix}" for i in range(50) for suffix in SUFFIXES} + modules = dict(model.named_modules()) + if not names <= modules.keys(): + raise ValueError("Incomplete H3 main projection inventory") + # Validate the entire inventory before making any in-place conversion. + for name, layer in modules.items(): + if not isinstance(layer, LinearBase): + continue + if type(layer.quant_method) is not UnquantizedLinearMethod: + raise ValueError("H3 Ultra requires BF16 loading before MXFP8 conversion: " + name) + if name in names and ( + tuple(layer.weight.shape) != SHAPES[name.split(".", 2)[2]] or layer.bias is not None or layer.tp_size != 1 + ): + raise ValueError("Unsupported H3 projection shape or parallelism: " + name) + for block in model.blocks: + if block.attn.to_gate_compress is None or block.attn.to_gate_compress.weight.dtype != torch.bfloat16: + raise ValueError("The VSA compression gate must remain BF16") + if block.adaln_proj is not None: + raise ValueError("H3 Ultra requires precomputed AdaLN") + fusion = os.getenv("VLLM_OMNI_H3_SWIGLU_MXFP8_FUSION", "0") + if fusion not in ("0", "1"): + raise ValueError("VLLM_OMNI_H3_SWIGLU_MXFP8_FUSION must be 0 or 1") + for name in sorted(names): + quantize_layer(modules[name]) + for block in model.blocks: + block.mlp._h3_swiglu_mxfp8_fused = fusion == "1" + model._h3_mxfp8_installed = True + + +def fused_swiglu_fc2(layer, hidden): + return silu_mxfp8_linear(hidden, layer.weight, layer.mxfp8_weight_scale) + + +def validate_split_projection(projection) -> None: + if ( + type(projection.quant_method) is not H3MXFP8Method + or projection.bias is not None + or projection.tp_size != 1 + or tuple(projection.weight.shape) != (21504, 5376) + ): + raise ValueError("Split QK/V requires the TP1 H3 MXFP8 QKV projection") + + +def project_split_qk(projection, x): + cut = _scale_numel(_QK_WIDTH, 5376) + return mxfp8_quantize_project(x, projection.weight[:_QK_WIDTH], projection.mxfp8_weight_scale[:cut]) + + +def project_split_v(projection, q, activation_scale, qk): + validate_split_projection(projection) + cut = _scale_numel(_QK_WIDTH, 5376) + return mxfp8_scaled_mm(q, projection.weight[_QK_WIDTH:], activation_scale, projection.mxfp8_weight_scale[cut:]) + + +VAE_MXFP8_ENV = "VLLM_OMNI_H3_VAE_MXFP8" + + +def vae_mxfp8_enabled(): + value = os.environ.get(VAE_MXFP8_ENV, "0") + if value not in ("0", "1"): + raise ValueError(f"{VAE_MXFP8_ENV} must be 0 or 1") + if value == "1" and os.environ.get("VLLM_OMNI_H3_VAE_INT8_CONVROT", "0") != "0": + raise ValueError("Choose exactly one VAE quantization method") + return value == "1" + + +class OnlineMXFP8Linear(nn.Module): + def __init__(self, linear, name): + super().__init__() + from comfy_kitchen.backends.cuda import quantize_mxfp8 + from comfy_kitchen.backends.eager.quantization import quantize_mxfp8 as quantize_fp32_weight + + if linear.weight.dtype != torch.float32 or not linear.weight.is_cuda: + raise RuntimeError("Online MXFP8 requires original FP32 weights on CUDA") + self.in_features, self.out_features = linear.in_features, linear.out_features + if self.in_features % 32 or self.out_features % 32: + raise RuntimeError("Unexpected H3 VAE weight dimensions") + self.name, self.calls = name, 0 + self._quantize = quantize_mxfp8 + with torch.no_grad(), torch.autocast("cuda", enabled=False): + # The wheel's CUDA MX producer actually accepts only FP16/BF16. + # Use its FP32-capable implementation once at load time so weights + # are not silently rounded to FP16 before the intended quantizer. + q, scale = quantize_fp32_weight(linear.weight.detach().contiguous(), pad_32x=False) + assert q.dtype == torch.float8_e4m3fn and scale.dtype == torch.float8_e8m0fnu + assert q.shape == linear.weight.shape + self.register_buffer("weight", q) + self.register_buffer("weight_scale", scale) + self.register_buffer("bias", None if linear.bias is None else linear.bias.detach().to(torch.float16)) + + def forward(self, x): + if not x.is_cuda or x.shape[-1] != self.in_features: + raise RuntimeError("Unsupported VAE MXFP8 input") + if not torch.is_autocast_enabled("cuda") or torch.get_autocast_dtype("cuda") != torch.float16: + raise RuntimeError("H3 VAE MXFP8 is qualified only under FP16 decode autocast") + shape = x.shape + with torch.autocast("cuda", enabled=False): + x2d = x.to(torch.float16).reshape(-1, self.in_features).contiguous() + rows = x2d.shape[0] + q, scale = self._quantize(x2d, pad_32x=True) + if q.dtype != torch.float8_e4m3fn or scale.dtype != torch.float8_e8m0fnu: + raise RuntimeError("MXFP8 activation producer returned the wrong format") + # Call the same native cuBLASLt block-scaled API as the DiT. + # Unsupported dtype/shape/device raises; no dequantized fallback. + out = mxfp8_scaled_mm( + q, + self.weight, + scale, + self.weight_scale, + bias=self.bias, + output_dtype=torch.float16, + use_fast_accum=False, + ) + if out.dtype != torch.float16 or out.shape != (q.shape[0], self.out_features): + raise RuntimeError("MXFP8 GEMM returned the wrong output contract") + self.calls += 1 + return out[:rows].reshape(*shape[:-1], self.out_features) + + +def install_vae_mxfp8(decoder): + from importlib.metadata import version + + assert version("comfy-kitchen") == "0.2.33" + assert len(decoder.transformer_blocks) == 36 + modules = [] + for i, block in enumerate(decoder.transformer_blocks): + for owner, key, kind in ( + (block.attn, "to_qkv", "qkv"), + (block.attn, "to_out", "out"), + (block.ff, "w1", "fc1"), + (block.ff, "w2", "fc2"), + ): + old = getattr(owner, key) + assert isinstance(old, nn.Linear) + name = f"transformer_blocks.{i}.{'attn' if kind in ('qkv', 'out') else 'ff'}.{key}" + module = OnlineMXFP8Linear(old, name) + module.kind = kind + setattr(owner, key, module) + modules.append(module) + decoder._h3_online_mxfp8_modules = tuple(modules) + return dict( + rank=torch.distributed.get_rank() if torch.distributed.is_initialized() else 0, + layers=144, + group_size=32, + weight_source="original_fp32", + activation_dtype="float8_e4m3fn", + scale_dtype="float8_e8m0fnu", + output_dtype="float16", + backend="torch_scaled_mm_cublaslt", + weight_quantizer="comfy_kitchen_eager_fp32", + activation_quantizer="comfy_kitchen_cuda", + version="0.2.33", + approximate=True, + convrot=False, + ) + + +def audit_vae_mxfp8_calls(decoder, rank, logger): + modules = getattr(decoder, "_h3_online_mxfp8_modules", ()) + if not modules: + return + counts = collections.Counter() + for module in modules: + expected = (74 if rank < 4 else 73) if module.kind == "out" else 21 + if module.calls != expected: + raise RuntimeError(f"MXFP8 dispatch count mismatch: {module.name}: {module.calls} != {expected}") + counts[module.kind] += module.calls + logger.info( + "H3_VAE_MXFP8_COMPLETE %s", + json.dumps( + dict(rank=rank, layers=144, calls=dict(counts), group_size=32, fallback=0, status="PASS"), sort_keys=True + ), + ) + for module in modules: + module.calls = 0 diff --git a/vllm_omni/diffusion/models/minimax_h3/ultra.py b/vllm_omni/diffusion/models/minimax_h3/ultra.py new file mode 100644 index 00000000000..b53989e9e42 --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/ultra.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""One opt-in for the qualified H3 schedule; geometry remains model-owned.""" + +import os +from collections.abc import MutableMapping + +# Explicitly keep nonselected precision/scheduling routes disabled. Enabling +# the preset must not silently accept a conflicting process environment. +ULTRA_ENV = { + "VLLM_OMNI_FASTVIDEO_VSA_DIRECT_Q2K": "1", + "VLLM_OMNI_H3_VSA_KERNEL": "flashinfer", + "VLLM_OMNI_H3_ATTENTION_OVERLAP": "1", + "VLLM_OMNI_H3_DIT_MXFP8": "1", + "VLLM_OMNI_H3_SWIGLU_MXFP8_FUSION": "1", + "VLLM_OMNI_H3_VSA_SAGE": "sage", + "VLLM_OMNI_H3_VSPLIT": "1", + "VLLM_OMNI_H3_VSPLIT_MODE": "split", + "VLLM_OMNI_H3_MXFP8_VSPLIT": "1", + "VLLM_OMNI_H3_LOSSLESS_GATE": "1", + "VLLM_OMNI_H3_LOSSLESS_EARLY_Q": "1", + "VLLM_OMNI_H3_LOSSLESS_LAYOUT_VIEW": "1", + "VLLM_OMNI_H3_LOSSLESS_COARSE_OVERLAP": "1", + "VLLM_OMNI_H3_O_PRODUCER_LOOKAHEAD": "1", + "VLLM_OMNI_ULYSSES_A2A_BACKEND": "flashinfer-pcie", + "VLLM_OMNI_FLASHINFER_ULYSSES_REQUIRE_RDMA": "1", + "VLLM_OMNI_FLASHINFER_ULYSSES_MAX_BYTES": "175788032", + "VLLM_OMNI_FLASHINFER_ULYSSES_QK_PRODUCER_DIRECT": "1", + "VLLM_OMNI_FLASHINFER_ULYSSES_O_PRODUCER_DIRECT": "1", + "VLLM_OMNI_FLASHINFER_ULYSSES_RELEASE_AFTER_DENOISE": "0", + "VLLM_OMNI_FASTVIDEO_VSA_FUSED_TILE_PACK": "1", + "VLLM_OMNI_FASTVIDEO_VSA_FUSED_UNTILE": "1", + "VLLM_OMNI_FASTVIDEO_VSA_O_BUNDLE": "1", + "VLLM_OMNI_FASTVIDEO_VSA_SKIP_SOFTMAX_THRESHOLD_SCALE_FACTOR": "0", + "VLLM_OMNI_MINIMAX_H3_CHUNKED_CPU_MP4_OUTPUT": "1", + "VLLM_OMNI_H3_VAE_PAIR_PIPELINE": "1", + "VLLM_OMNI_H3_VAE_MIXED_BATCH": "1", + "VLLM_OMNI_H3_VAE_BATCH4": "1", + "VLLM_OMNI_H3_VAE_GATHER_OVERLAP": "1", + "VLLM_OMNI_H3_VAE_MXFP8": "1", + "VLLM_OMNI_H3_FULL_VAE_AUDIO_OVERLAP": "1", + "VLLM_OMNI_MINIMAX_H3_VAE_EXACT_OPS_SM120": "1", +} + + +def configure_ultra(env: MutableMapping[str, str] | None = None) -> bool: + env = os.environ if env is None else env + value = env.get("VLLM_OMNI_H3_ULTRA", "0") + if value not in ("0", "1"): + raise ValueError("VLLM_OMNI_H3_ULTRA must be 0 or 1") + if value == "0": + return False + conflicts = [name for name, expected in ULTRA_ENV.items() if name in env and env[name] != expected] + if conflicts: + raise ValueError("H3 Ultra conflicts with explicit settings: " + ", ".join(sorted(conflicts))) + env.update(ULTRA_ENV) + return True diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index ca8359904ea..664893cd391 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -5,11 +5,14 @@ from __future__ import annotations import importlib +import inspect import json +import os from collections.abc import Iterator -from contextlib import AbstractContextManager, contextmanager, nullcontext +from contextlib import AbstractContextManager, ExitStack, contextmanager, nullcontext +from dataclasses import dataclass from pathlib import Path -from typing import Any, ClassVar +from typing import Any import torch import torch.distributed as dist @@ -29,36 +32,91 @@ ) from .chunked_decode import decode_h3_chunks -from .ops import install_h3_vae_optimizations +from .ops import ( + H3VAEExactOpStatsSnapshot, + install_h3_vae_optimizations, + snapshot_h3_vae_exact_op_stats, +) from .packed_tokens import minimax_h3_patchify_video_latent +from .vae_parallel import _agree_on_failure MINIMAX_H3_KEYFRAME_ENCODE_SEED = 42 MINIMAX_H3_AUDIO_SAMPLE_RATE = 32000 MINIMAX_H3_AUDIO_CHANNELS = 2 +MINIMAX_H3_VAE_DECODER_TILE_SIZE_ENV = "VLLM_OMNI_MINIMAX_H3_VAE_DECODER_TILE_SIZE" logger = init_logger(__name__) -@contextmanager -def _minimax_h3_keyframe_encode_context( - device: torch.device, -) -> Iterator[None]: - if device.type != "cuda": - yield +def resolve_minimax_h3_vae_decoder_tile_size(raw: str | None = None) -> int: + """Resolve the experimental decoder tile size without silent fallback. + + Zero preserves the checkpoint configuration. A nonzero override must use + the 16-pixel alignment required by H3's spatial VAE geometry. + """ + + if raw is None: + raw = os.environ.get(MINIMAX_H3_VAE_DECODER_TILE_SIZE_ENV, "0") + normalized = str(raw).strip() + try: + value = int(normalized, 10) + except (TypeError, ValueError) as exc: + raise ValueError( + f"{MINIMAX_H3_VAE_DECODER_TILE_SIZE_ENV} must be 0 or a positive multiple of 16, got {raw!r}" + ) from exc + if value == 0: + return 0 + if value < 0 or value % 16 != 0: + raise ValueError(f"{MINIMAX_H3_VAE_DECODER_TILE_SIZE_ENV} must be 0 or a positive multiple of 16, got {raw!r}") + return value + + +def _apply_minimax_h3_vae_decoder_tile_size_override(remote: nn.Module, tile_size: int) -> None: + """Apply an already validated tile-size override to a loaded component.""" + + if tile_size == 0: return + model = getattr(remote, "model", None) + if model is None: + raise RuntimeError("MiniMax H3 video VAE tile-size override requires checkpoint attribute model") + required = ("decoder_tile_size", "decoder_tile_overlap_min") + missing = [name for name in required if not hasattr(model, name)] + if missing: + raise RuntimeError( + "MiniMax H3 video VAE tile-size override requires checkpoint attribute(s): " + ", ".join(missing) + ) + try: + old_tile_size = int(model.decoder_tile_size) + overlap = int(model.decoder_tile_overlap_min) + except (TypeError, ValueError) as exc: + raise RuntimeError("MiniMax H3 checkpoint decoder tile geometry must be integer-valued") from exc + if old_tile_size <= 0 or overlap < 0 or overlap >= tile_size: + raise RuntimeError( + "MiniMax H3 video VAE tile-size override is incompatible with checkpoint geometry: " + f"old_tile_size={old_tile_size}, new_tile_size={tile_size}, overlap={overlap}" + ) + model.decoder_tile_size = tile_size + logger.warning( + "MiniMax H3 video VAE decoder tile-size override enabled: old=%d new=%d " + "decoder_tile_overlap_min=%d (unchanged); output is numerical/non-bit-exact", + old_tile_size, + tile_size, + overlap, + ) - # The official keyframe latent uses cuDNN's TF32 convolution path. The - # default non-deterministic algorithm can select numerically different - # reductions on H100s, and the difference is amplified by the denoiser. - # Pin both the algorithm and math mode for this sensitive encode only. - with torch.backends.cudnn.flags( - enabled=True, - benchmark=False, - deterministic=True, - allow_tf32=True, - ): - yield + +@dataclass(frozen=True) +class _TemporalDecodePlan: + """Exact released-H3 temporal layout used by the communication shortcut.""" + + num_windows: int + decoded_frames_per_window: int + gathered_frames_per_window: int + main_frames: int + overlap_frames: int + output_frames: int + keep_ranges: tuple[tuple[int, int], ...] def _load_component_config(component_path: str) -> dict[str, Any]: @@ -152,6 +210,10 @@ def __init__( component_path, self.config_dict, ) + _apply_minimax_h3_vae_decoder_tile_size_override( + self.remote, + resolve_minimax_h3_vae_decoder_tile_size(), + ) # Match the reference loader contract before installing inference-only # decoder fast paths. Keyframe encoding remains FP32; decoder Linear # weights may be materialized in FP16 because reference decode casts @@ -159,11 +221,33 @@ def __init__( initial_device = load_device or device self.remote.eval().to(device=initial_device, dtype=torch.float32) decoder = getattr(self.remote.model, "decoder", None) + from . import quantization + + mxfp8_enabled = quantization.vae_mxfp8_enabled() + if mxfp8_enabled and (decoder is None or initial_device.type != "cuda"): + raise RuntimeError("MXFP8 arm requires the resident CUDA H3 video decoder") if decoder is not None: - install_h3_vae_optimizations( + exact_ops_installed = install_h3_vae_optimizations( decoder, device=device, + persist_fp16_weights=not mxfp8_enabled, ) + if mxfp8_enabled: + if not exact_ops_installed: + raise RuntimeError("MXFP8 requires the current-best VAE operators") + recipe = quantization.install_vae_mxfp8(decoder) + logger.info("H3_VAE_MXFP8_READY %s", json.dumps(recipe, sort_keys=True)) + if exact_ops_installed: + rank = dist.get_rank() if dist.is_initialized() else 0 + blocks = getattr(decoder, "transformer_blocks", ()) + logger.info( + "MiniMax H3 VAE exact operators installed: rank=%d device=%s " + "blocks=%d linear_dtype=%s counters=enabled", + rank, + device, + len(blocks), + "mxfp8" if mxfp8_enabled else "float16", + ) self._stager = None if initial_device.type == "cpu" and device.type not in ("cpu", "meta"): self._stager = PinnedModuleStager( @@ -379,6 +463,15 @@ def encode_video( ).float() return rows, shape + @torch.inference_mode() + def decode_latent(self, latent: torch.Tensor, *, return_output: bool = True) -> torch.Tensor | None: + with self._decode_tiling_context(latent): + decoded = self.model.decode_base(self._denormalize_latent(latent)) + # All decoder collectives finish before non-output ranks leave. + if not return_output: + return None + return self._normalize_decoded_frames(self.model.processor.revert_tensor(decoded)) + def _decode_tiling_context(self, latent: torch.Tensor) -> AbstractContextManager: """Pick the tiling mode a decode of ``latent`` can safely use. @@ -401,12 +494,6 @@ def _decode_tiling_context(self, latent: torch.Tensor) -> AbstractContextManager return self._rank_local_tiling() return nullcontext() - @torch.inference_mode() - def decode_latent(self, latent: torch.Tensor) -> torch.Tensor: - with self._decode_tiling_context(latent): - decoded = self.model.decode_base(self._denormalize_latent(latent)) - return self._normalize_decoded_frames(self.model.processor.revert_tensor(decoded)) - def _denormalize_latent(self, latent: torch.Tensor) -> torch.Tensor: channels = int(self.config_dict["latent_channels"]) mean = torch.tensor(self.config_dict["latents_mean"], device=latent.device, dtype=latent.dtype) @@ -424,10 +511,6 @@ def _normalize_decoded_frames(decoded: torch.Tensor) -> torch.Tensor: raise ValueError(f"unexpected decoded video shape {tuple(frames.shape)}") return frames.float() - # Chunks are reverted through the checkpoint's processor, which - # denormalizes and clamps into the unit interval. - chunk_value_range: ClassVar[tuple[float, float]] = (0.0, 1.0) - @torch.inference_mode() def decode_with_chunks(self, z: torch.Tensor, *, on_chunk: DecodedChunkConsumer) -> None: """Decode temporal clips and synchronously publish frames-only chunks. @@ -456,6 +539,496 @@ def decode_with_chunks(self, z: torch.Tensor, *, on_chunk: DecodedChunkConsumer) with self._decode_tiling_context(z): decode_h3_chunks(self, z, on_chunk, group=group) + def supports_chunked_output(self) -> bool: + """Whether this checkpoint exposes the finalized-chunk callback hook.""" + + try: + parameters = inspect.signature(self.model.decode_base).parameters + except (TypeError, ValueError): + return False + return "output_callback" in parameters + + def validate_chunked_output(self) -> None: + """Validate the checkpoint and environment before PP8 decode starts.""" + + if not self.supports_chunked_output(): + raise RuntimeError("MiniMax H3 checkpoint does not expose chunked VAE decode output") + raw = ( + os.environ.get( + "MINIMAX_H3_VAE_DECODER_STREAM_TEMPORAL_CAT", + "1", + ) + .strip() + .lower() + ) + if raw in {"0", "false", "no", "off", "disable", "disabled"}: + raise RuntimeError( + "MiniMax H3 finalized-chunk output requires MINIMAX_H3_VAE_DECODER_STREAM_TEMPORAL_CAT=1" + ) + self._finalized_chunk_frame_capacity() + + def _finalized_chunk_frame_capacity(self) -> int: + """Upper bound emitted by the checkpoint's temporal callback. + + ``decode_temporal`` slices each decoded clip at + ``tokens_chunk_size * vae_ratio_t`` before removing padding. Isolated + head/tail callbacks contain one frame, so the unsliced size is a safe + bound for every callback and avoids sizing the ring from a tiny first + chunk. + """ + + try: + capacity = int(self.model.tokens_chunk_size) * int(self.model.vae_ratio_t) + except (AttributeError, TypeError, ValueError) as exc: + raise RuntimeError("MiniMax H3 checkpoint does not expose its temporal chunk geometry") from exc + if capacity < 1: + raise RuntimeError(f"invalid MiniMax H3 temporal chunk capacity {capacity}") + return capacity + + def _log_exact_op_chunked_decode( + self, + latent: torch.Tensor, + *, + num_tiles: int, + route: str, + before: H3VAEExactOpStatsSnapshot | None, + ) -> None: + """Emit one fail-closed-friendly dispatch summary per VAE rank. + + This is deliberately a delta around one decode request. A cumulative + process counter cannot distinguish the warmup from the measured arm, + and an installation marker alone cannot prove that every PP8 rank hit + the optimized kernels without a silent eager fallback. + """ + + decoder = getattr(self.model, "decoder", None) + after = snapshot_h3_vae_exact_op_stats(decoder) if decoder is not None else None + if before is None or after is None: + return + delta = after.delta(before) + from . import quantization + + quantization.audit_vae_mxfp8_calls(decoder, dist.get_rank() if dist.is_initialized() else 0, logger) + + model = self.model + latent_tokens = int(latent.shape[2]) + tokens_per_window = int(model.tokens_chunk_size) + isolated_tokens = int(bool(model.isolated_first_frame) and int(model.frame_pre_padding) == 0) + int( + bool(model.isolated_last_frame) + ) + pseudo_tokens = latent_tokens - isolated_tokens + int(model.token_drop) + remainder = pseudo_tokens % tokens_per_window + if remainder: + pseudo_tokens += tokens_per_window - remainder + temporal_windows = pseudo_tokens // tokens_per_window - int(int(model.token_drop) > 0) + + state = self._native_parallel_state() + sp_rank = int(state.get("sp_rank", 0)) + sp_size = int(state.get("sp_size", 1)) + if route == "temporal_window_pp": + owned_windows = len(range(sp_rank, temporal_windows, sp_size)) + local_tiles = num_tiles + else: + owned_windows = temporal_windows + if self.parallel_size > 1 and num_tiles >= self.parallel_size: + local_tiles = len(range(sp_rank, num_tiles, sp_size)) + else: + # ``_rank_local_tiling`` makes every rank own the full grid + # when there are fewer spatial tiles than PP ranks. + local_tiles = num_tiles + + stack_tiling = bool(getattr(model, "stack_tiling", False)) + decoder_calls_per_window = int(local_tiles > 0) if stack_tiling else local_tiles + block_count = len(getattr(decoder, "transformer_blocks", ())) + expected_block_calls = owned_windows * decoder_calls_per_window * block_count + if route in ("temporal_spatial_pair_pp", "temporal_spatial_mixed_pp"): + from .vae_batching import jobs + + assert (temporal_windows, num_tiles, sp_size) == (21, 28, 8) and not stack_tiling + # Per-window ownership varies. -1 is an explicit sentinel, while + # the exact total comes from the fixed, coverage-checked schedule. + local_tiles = -1 + expected_block_calls = sum(len(row[sp_rank]) for row in jobs()) * block_count + if route == "temporal_spatial_mixed_pp": + from .vae_batching import batch_plan + + expected_block_calls = len(batch_plan(sp_rank)) * block_count + fallback_total = ( + delta.transformer_fallback + + delta.qk_norm_rope_fallback + + delta.swiglu_fallback + + delta.scaled_residual_fallback + ) + counts_ok = ( + delta.transformer_fast == expected_block_calls + and delta.qk_norm_rope_fast == expected_block_calls + and delta.swiglu_fast == expected_block_calls + and delta.scaled_residual_fast == 2 * expected_block_calls + and fallback_total == 0 + ) + global_rank = dist.get_rank() if dist.is_initialized() else 0 + logger.info( + "MiniMax H3 VAE exact operators complete: rank=%d sp_rank=%d " + "route=%s latent_tokens=%d temporal_windows=%d tiles=%d " + "local_tiles=%d owned_windows=%d blocks=%d expected_block_calls=%d " + "transformer_fast=%d transformer_fallback=%d " + "qk_norm_rope_fast=%d qk_norm_rope_fallback=%d " + "swiglu_fast=%d swiglu_fallback=%d " + "scaled_residual_fast=%d scaled_residual_fallback=%d status=%s", + global_rank, + sp_rank, + route, + latent_tokens, + temporal_windows, + num_tiles, + local_tiles, + owned_windows, + block_count, + expected_block_calls, + delta.transformer_fast, + delta.transformer_fallback, + delta.qk_norm_rope_fast, + delta.qk_norm_rope_fallback, + delta.swiglu_fast, + delta.swiglu_fallback, + delta.scaled_residual_fast, + delta.scaled_residual_fallback, + "ok" if counts_ok else "failed", + ) + + @staticmethod + def _require_remote_signature(function: Any, name: str, parameters: tuple[str, ...]) -> None: + try: + actual = tuple(inspect.signature(function).parameters) + except (TypeError, ValueError) as exc: + raise RuntimeError(f"cannot inspect MiniMax H3 checkpoint method {name}") from exc + if actual != parameters: + raise RuntimeError( + f"MiniMax H3 checkpoint method {name} has unsupported signature {actual}; expected {parameters}" + ) + + def _temporal_decode_plan( + self, + latent: torch.Tensor, + *, + num_tiles: int, + require_spatial_rank_share: bool = True, + ) -> _TemporalDecodePlan: + """Validate and describe the one released temporal layout we can prune. + + The optimization deliberately has no heuristic path. The released H3 + decoder produces 28 frames from every seven-token temporal window, then + drops frames 0:3 and 20:23 before any cross-window blend. Spatial tile + blending is frame-local, so gathering only the other 22 frames is exact. + Any checkpoint/configuration drift is rejected before PP ranks enter the + first tile collective. + """ + + if latent.ndim != 5 or int(latent.shape[0]) != 1: + raise RuntimeError( + "MiniMax H3 temporal pre-gather prune requires one rank-5 video latent, " + f"got shape={tuple(latent.shape)}" + ) + if self.parallel_size <= 1: + raise RuntimeError("MiniMax H3 temporal pre-gather prune requires parallel VAE tiling") + if require_spatial_rank_share and num_tiles < self.parallel_size: + raise RuntimeError( + "MiniMax H3 temporal pre-gather prune requires at least one decoder tile per rank: " + f"tiles={num_tiles}, ranks={self.parallel_size}" + ) + + model = self.model + expected_attributes: dict[str, Any] = { + "decoder_tiling": True, + "parallel_tiling": True, + "use_3d_conv": True, + "tokens_chunk_size": 5, + "token_overlap": 2, + "vae_ratio_t": 4, + "frame_pre_padding": 3, + "frame_overlap": 5, + "token_drop": 3, + "isolated_first_frame": False, + "isolated_last_frame": False, + } + mismatches = [] + for name, expected in expected_attributes.items(): + actual = getattr(model, name, None) + if actual != expected: + mismatches.append(f"{name}={actual!r} (expected {expected!r})") + if bool(getattr(model, "training", False)): + mismatches.append("training=True (expected inference mode)") + if mismatches: + raise RuntimeError( + "MiniMax H3 temporal pre-gather prune does not match this checkpoint plan: " + ", ".join(mismatches) + ) + + gather = getattr(model, "_all_gather_tiled_results", None) + adaptive_decode = getattr(model, "_adaptive_decode", None) + blend = getattr(model, "blend", None) + output_plan = getattr(model, "_decode_temporal_output_frame_plan", None) + temporal_dtype = getattr( + importlib.import_module(model.__class__.__module__), "_resolve_temporal_cat_dtype", None + ) + for function, name, parameters in ( + (gather, "_all_gather_tiled_results", ("tasks", "num_tiles")), + (adaptive_decode, "_adaptive_decode", ("z",)), + (blend, "blend", ("a", "b", "blend_extent", "dim")), + ( + output_plan, + "_decode_temporal_output_frame_plan", + ("z", "z_head", "z_tail", "num_chunks", "pad_tokens"), + ), + (temporal_dtype, "_resolve_temporal_cat_dtype", ()), + ): + if not callable(function): + raise RuntimeError(f"MiniMax H3 checkpoint is missing required method {name}") + self._require_remote_signature(function, name, parameters) + + tokens_chunk_size = int(model.tokens_chunk_size) + token_overlap = int(model.token_overlap) + vae_ratio_t = int(model.vae_ratio_t) + frame_pre_padding = int(model.frame_pre_padding) + pseudo_total_tokens = int(latent.shape[2]) + int(model.token_drop) + if pseudo_total_tokens % tokens_chunk_size: + raise RuntimeError( + "MiniMax H3 temporal pre-gather prune does not support a padded final token window: " + f"latent_tokens={int(latent.shape[2])}, token_drop={int(model.token_drop)}" + ) + num_windows = pseudo_total_tokens // tokens_chunk_size - 1 + if num_windows <= 0: + raise RuntimeError(f"MiniMax H3 temporal pre-gather prune planned {num_windows} temporal windows") + + decoded_frames = (tokens_chunk_size + token_overlap) * vae_ratio_t + chunk_dec = tokens_chunk_size * vae_ratio_t + keep_ranges = tuple( + (start + frame_pre_padding, min(start + chunk_dec, decoded_frames)) + for start in range(0, decoded_frames, chunk_dec) + ) + if keep_ranges != ((3, 20), (23, 28)): + raise RuntimeError( + "MiniMax H3 temporal pre-gather prune derived an unsupported frame plan: " + f"decoded_frames={decoded_frames}, keep_ranges={keep_ranges}" + ) + main_frames = keep_ranges[0][1] - keep_ranges[0][0] + overlap_frames = keep_ranges[1][1] - keep_ranges[1][0] + gathered_frames = main_frames + overlap_frames + output_frames = num_windows * main_frames + overlap_frames + + total, pad, checkpoint_output = output_plan(latent, None, None, num_windows, 0) + if (int(total), int(pad), int(checkpoint_output)) != (output_frames, 0, output_frames): + raise RuntimeError( + "MiniMax H3 checkpoint temporal plan disagrees with pre-gather prune: " + f"checkpoint={(int(total), int(pad), int(checkpoint_output))}, " + f"expected={(output_frames, 0, output_frames)}" + ) + return _TemporalDecodePlan( + num_windows=num_windows, + decoded_frames_per_window=decoded_frames, + gathered_frames_per_window=gathered_frames, + main_frames=main_frames, + overlap_frames=overlap_frames, + output_frames=output_frames, + keep_ranges=keep_ranges, + ) + + def _paired_decode_plan( + self, + latent: torch.Tensor, + *, + num_tiles: int, + ) -> _TemporalDecodePlan: + """Validate the one canonical 107-token temporal-PP serving shape.""" + + plan = self._temporal_decode_plan( + latent, + num_tiles=num_tiles, + require_spatial_rank_share=False, + ) + observed = ( + int(latent.shape[2]), + plan.num_windows, + plan.main_frames, + plan.overlap_frames, + plan.output_frames, + ) + expected = (107, 21, 17, 5, 362) + if observed != expected: + raise RuntimeError( + "MiniMax H3 temporal chunk parallel supports only the canonical " + "107-token/21-window/17+5 plan: " + f"observed={observed}, expected={expected}" + ) + return plan + + @torch.inference_mode() + def decode_latent_to_chunked_cpu_mp4( + self, + latent: torch.Tensor, + *, + height: int, + width: int, + fps: int, + audio_sample_rate: int, + video_codec_options: dict[str, str] | None, + return_output: bool, + ) -> Any | None: + """Stream finalized decoder chunks into rank zero's CPU MP4 sink. + + The default/pruned routes enter the checkpoint-native spatial PP8 + decode on every rank. The explicit temporal route instead assigns + complete windows round-robin and runs spatial tiles rank-locally. + Only the designated output rank constructs a sink. Sink-construction + and route-validation status is reduced before decode so a local error + cannot strand peers in the first VAE data collective. + """ + from .chunked_cpu_output import MiniMaxH3ChunkedCpuMp4Output + + sink = None + paired_plan: _TemporalDecodePlan | None = None + pair_stats = None + gather_stream = None + num_tiles = 0 + world = get_world_group() if dist.is_initialized() else None + creation_error: BaseException | None = None + try: + self.validate_chunked_output() + num_tiles = self._decoder_tile_count(latent) + pair_flag = os.environ.get("VLLM_OMNI_H3_VAE_PAIR_PIPELINE", "0") + if pair_flag not in ("0", "1"): + raise ValueError("VLLM_OMNI_H3_VAE_PAIR_PIPELINE must be 0 or 1") + pair_enabled = pair_flag == "1" + mixed_flag = os.environ.get("VLLM_OMNI_H3_VAE_MIXED_BATCH", "0") + if mixed_flag not in ("0", "1"): + raise ValueError("VLLM_OMNI_H3_VAE_MIXED_BATCH must be 0 or 1") + mixed_enabled = mixed_flag == "1" + if mixed_enabled and (not pair_enabled): + raise ValueError("Mixed VAE batching requires paired VAE scheduling") + from .vae_parallel import prepare_gather_stream + + gather_stream = prepare_gather_stream(latent.device, paired=pair_enabled, mixed=mixed_enabled, owner=self) + if pair_enabled: + if self.parallel_size <= 1 or world is None: + raise RuntimeError("MiniMax H3 temporal chunk parallel requires distributed VAE parallel_size > 1") + if int(world.world_size) != int(self.parallel_size): + raise RuntimeError( + f"MiniMax H3 temporal chunk group size must match VAE parallel_size: " + f"group={world.world_size}, VAE={self.parallel_size}" + ) + expected_return_output = int(world.rank_in_group) == 0 + if bool(return_output) != expected_return_output: + raise RuntimeError( + f"MiniMax H3 temporal chunk output rank must be group rank zero: " + f"rank={world.rank_in_group}, return_output={return_output}" + ) + paired_plan = self._paired_decode_plan(latent, num_tiles=num_tiles) + if return_output: + sink = MiniMaxH3ChunkedCpuMp4Output( + width=width, + height=height, + fps=fps, + audio_sample_rate=audio_sample_rate, + max_chunk_frames=self._finalized_chunk_frame_capacity(), + device=latent.device, + video_codec_options=video_codec_options, + ) + except BaseException as exc: + creation_error = exc + group = world.device_group if world is not None else None + if self.parallel_size > 1 and group is not None: + failed = torch.tensor([creation_error is not None], dtype=torch.int32, device=latent.device) + dist.all_reduce(failed, op=dist.ReduceOp.MAX, group=group) + if bool(failed.item()): + if sink is not None: + sink.abort() + if creation_error is not None: + raise RuntimeError("failed to construct chunked CPU MP4 output sink") from creation_error + raise RuntimeError("chunked CPU MP4 output sink failed on the output rank") + elif creation_error is not None: + raise RuntimeError("failed to construct chunked CPU MP4 output sink") from creation_error + channels = int(self.config_dict["latent_channels"]) + mean = torch.tensor(self.config_dict["latents_mean"], device=latent.device, dtype=latent.dtype).view( + 1, channels, 1, 1, 1 + ) + std = torch.tensor(self.config_dict["latents_std"], device=latent.device, dtype=latent.dtype).view( + 1, channels, 1, 1, 1 + ) + if paired_plan is not None: + tiling_context = self._rank_local_tiling() + elif self.parallel_size > 1 and num_tiles < self.parallel_size: + logger.warning_once( + "MiniMax-H3 VAE decode splits into %d tile(s) but the tile group has %d ranks; " + "decoding rank-locally for this shape instead.", + num_tiles, + self.parallel_size, + ) + tiling_context: AbstractContextManager = self._rank_local_tiling() + else: + tiling_context = nullcontext() + + def consume(decoded: torch.Tensor, first_frame: int, total_frames: int) -> None: + if sink is not None: + sink.submit_decoded(decoded, first_frame, total_frames, self.model.processor) + + decoder = getattr(self.model, "decoder", None) + exact_stats_before = snapshot_h3_vae_exact_op_stats(decoder) if decoder is not None else None + exact_route = ( + ("temporal_spatial_mixed_pp" if mixed_enabled else "temporal_spatial_pair_pp") + if pair_enabled + else "temporal_window_pp" + if paired_plan is not None + else "spatial_pp_chunked" + ) + try: + with tiling_context: + denormalized_latent = latent * std + mean + if paired_plan is not None: + assert world is not None + temporal_cat_dtype = importlib.import_module( + self.model.__class__.__module__ + )._resolve_temporal_cat_dtype() + from .vae_batching import decode_pairs + + with ExitStack() as mixed_stack: + mixed_stats = None + if mixed_enabled: + setup_error = None + try: + from .vae_batching import mixed_decode + + mixed_stats = mixed_stack.enter_context( + mixed_decode(self.model, denormalized_latent, world.rank_in_group) + ) + except BaseException as error: + setup_error = error + if _agree_on_failure(world, denormalized_latent.device, setup_error is not None): + raise RuntimeError("Mixed VAE adapter setup failed") from setup_error + pair_stats = decode_pairs( + self.model, + denormalized_latent, + world, + consume if return_output else None, + temporal_cat_dtype=temporal_cat_dtype, + gather_stream=gather_stream, + ) + if mixed_stats is not None: + pair_stats["mixed"] = mixed_stats + decoded = None + else: + decoded = self.model.decode_base(denormalized_latent, output_callback=consume) + except BaseException: + if sink is not None: + sink.abort() + raise + if decoded is not None: + if sink is not None: + sink.abort() + raise RuntimeError("chunked MiniMax H3 VAE decode unexpectedly returned a full tensor") + self._log_exact_op_chunked_decode(latent, num_tiles=num_tiles, route=exact_route, before=exact_stats_before) + if pair_stats is not None: + logger.info("H3_VAE_PAIR_COMPLETE %s", json.dumps(pair_stats, sort_keys=True)) + return sink + class MiniMaxH3AudioVAE(nn.Module): def __init__( @@ -577,3 +1150,24 @@ def decode_latent(self, latent: torch.Tensor) -> torch.Tensor: __all__ = ["MiniMaxH3AudioVAE", "MiniMaxH3VideoVAE"] + + +@contextmanager +def _minimax_h3_keyframe_encode_context( + device: torch.device, +) -> Iterator[None]: + if device.type != "cuda": + yield + return + + # The official keyframe latent uses cuDNN's TF32 convolution path. The + # default non-deterministic algorithm can select numerically different + # reductions on H100s, and the difference is amplified by the denoiser. + # Pin both the algorithm and math mode for this sensitive encode only. + with torch.backends.cudnn.flags( + enabled=True, + benchmark=False, + deterministic=True, + allow_tf32=True, + ): + yield diff --git a/vllm_omni/diffusion/models/minimax_h3/vae_batching.py b/vllm_omni/diffusion/models/minimax_h3/vae_batching.py new file mode 100644 index 00000000000..17ad39f5c5a --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/vae_batching.py @@ -0,0 +1,260 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""H3 temporal/spatial tile scheduling, decoder batching, and ordered assembly.""" + +from __future__ import annotations + +import types +from contextlib import contextmanager +from functools import lru_cache + +import torch + + +@lru_cache(None) +def jobs(windows=21, tiles=28, world=8): + assert (windows, tiles, world) == (21, 28, 8) + rounds = [] + for first in range(0, windows, 2): + count = min(2, windows - first) * tiles + rounds.append( + tuple(tuple((first + i // tiles, i % tiles) for i in range(r, count, world)) for r in range(world)) + ) + counts = tuple(sum(len(row[r]) for row in rounds) for r in range(world)) + assert counts == (74, 74, 74, 74, 73, 73, 73, 73) + flat = [job for row in rounds for rank in row for job in rank] + assert len(flat) == len(set(flat)) == windows * tiles + assert set(flat) == {(w, t) for w in range(windows) for t in range(tiles)} + return tuple(rounds) + + +def geometry(model, latent): + assert tuple(latent.shape) == (1, 24, 107, 44, 80) and latent.dtype == torch.float32 + assert model.vae_ratio == 16 and model.vae_ratio_t == 4 + assert model.tokens_chunk_size == 5 and model.token_overlap == 2 + assert model.frame_pre_padding == 3 and model.frame_overlap == 5 and model.token_drop == 3 + assert not model.training and not model.stack_tiling and not model.decoder_parallel + assert model.decoder_tiling and model.decoder_tile_size == 256 + yi, yl, yo = model.split_tiles(704, True) + xi, xl, xo = model.split_tiles(1280, True) + assert len(yi) == 4 and len(xi) == 7 and set(yl + xl) == {256} + slices = tuple( + (slice(y // 16, (y + length) // 16), slice(x // 16, (x + k) // 16)) + for y, length in zip(yi, yl) + for x, k in zip(xi, xl) + ) + return slices, tuple(yo), tuple(xo) + + +def assemble_spatial(model, tiles, y_overlap, x_overlap): + # Same sequence as checkpoint AutoencoderKL.tiled_decode after gathering: + # vertical blend, horizontal blend, trims, horizontal cats, vertical cat. + rows = [list(tiles[i * 7 : (i + 1) * 7]) for i in range(4)] + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + if i > 0: + tile = model.blend(rows[i - 1][j], tile, y_overlap[i - 1], dim=-2) + if j > 0: + tile = model.blend(row[j - 1], tile, x_overlap[j - 1], dim=-1) + if i < len(rows) - 1: + tile = tile[..., : -y_overlap[i], :] + if j < len(row) - 1: + tile = tile[..., :, : -x_overlap[j]] + result_row.append(tile) + result_rows.append(torch.cat(result_row, dim=-1)) + return torch.cat(result_rows, dim=-2) + + +def decode_pairs(model, latent, group, callback, *, temporal_cat_dtype=None, gather_stream=None): + """Group wrapper requires world_size/rank_in_group/ranks/device_group/all_reduce.""" + from vllm_omni.diffusion.models.minimax_h3.vae_parallel import ( + _agree_on_failure, + _gather_stack_to_rank_zero, + _LeaderAssembler, + ) + + assert group.world_size == 8 and (callback is not None) == (group.rank_in_group == 0) + spatial, y_overlap, x_overlap = geometry(model, latent) + schedule = jobs() + rank = group.rank_in_group + assembler = ( + _LeaderAssembler(model, callback, main_frames=17, overlap_frames=5, output_frames=362, device=latent.device) + if callback is not None + else None + ) + local_error = None + tile_shape = (1, 3, 22, 256, 256) + tile_dtype = torch.float32 + tile_count = 0 + try: + for round_id, round_jobs in enumerate(schedule): + width = max(map(len, round_jobs)) + # Uniform collective shape, including four padding tiles in the + # final round. Padding is never decoded or consumed by assembly. + packed = torch.empty((width, *tile_shape), device=latent.device, dtype=tile_dtype) + for slot, (window, tile) in enumerate(round_jobs[rank]): + if local_error is not None: + packed[slot].zero_() + continue + try: + sy, sx = spatial[tile] + clip = latent[:, :, window * 5 : window * 5 + 7, sy, sx] + with torch.cuda.nvtx.range(f"h3.vae.pair.tile.{window}.{tile}"): + decoded = model.decode(clip) + assert decoded.shape == (1, 3, 28, 256, 256) and decoded.dtype == tile_dtype + # Blends are per frame. Gather only the original live + # ranges, preserving FP32 until spatial blending finishes. + packed[slot, :, :, :17].copy_(decoded[:, :, 3:20]) + packed[slot, :, :, 17:].copy_(decoded[:, :, 23:28]) + tile_count += 1 + except BaseException as e: + e.__traceback__ = None + local_error = e + packed[slot].zero_() + if len(round_jobs[rank]) < width: + packed[len(round_jobs[rank]) :].zero_() + if assembler is not None: + assembler.drain_previous_round() + with torch.cuda.nvtx.range(f"h3.vae.pair.gather.{round_id}"): + if gather_stream is None: + gathered = _gather_stack_to_rank_zero(packed, group) + else: + from .vae_parallel import gather_on_stream + + gathered = gather_on_stream(packed, group, gather_stream) + if assembler is not None and gathered is not None: + # Assemble on the dedicated output stream, allowing the next + # pair's decoder to run concurrently. The original temporal + # assembler preserves callback order and the five-frame tail. + stream = assembler._stream + stream.wait_stream( + gather_stream + if gather_stream is not None + else torch.get_device_module().current_stream(latent.device) + ) + gathered.record_stream(stream) + try: + with torch.get_device_module().stream(stream): + for window in range(round_id * 2, min(round_id * 2 + 2, 21)): + parts = [None] * 28 + for owner, owner_jobs in enumerate(round_jobs): + for slot, (w, t) in enumerate(owner_jobs): + if w == window: + parts[t] = gathered[owner, slot] + assert all(t is not None for t in parts) + segment = assemble_spatial(model, parts, y_overlap, x_overlap) + if temporal_cat_dtype is not None: + segment = segment.to(temporal_cat_dtype) + assembler.push(segment) + assembler.finish_round() + except BaseException as e: + assembler._remember_error(e) + del gathered, packed + if assembler is not None: + assembler.finalize() + finally: + # Include non-output ranks and exceptional exits. No gather may still + # access a send/receive allocation when this request is released or + # the all-rank failure agreement enters the original caller stream. + try: + if gather_stream is not None: + gather_stream.synchronize() + finally: + if assembler is not None: + assembler.synchronize() + callback_error = assembler.error if assembler is not None else None + if _agree_on_failure(group, latent.device, local_error is not None or callback_error is not None): + raise RuntimeError("Paired VAE failed after draining collectives") from (local_error or callback_error) + assert tile_count == sum(len(row[rank]) for row in schedule) + if gather_stream is not None: + from vllm.logger import init_logger + + init_logger(__name__).info( + "H3_VAE_GATHER_OVERLAP rank=%d gathers=11 separate_stream=1 " + "send_record_stream=1 decoder_calls=42 tiles=%d full_vae=1 " + "mixed=1 gather_drained=1 frames=362 complete=1", + rank, + tile_count, + ) + return dict(rank=rank, tile_calls=tile_count, gathers=11, windows=21, padding_tiles=4, frames=362) + + +@lru_cache(None) +def batch_plan(rank): + return tuple( + tuple(round_jobs[rank][i : i + 4]) for round_jobs in jobs() for i in range(0, len(round_jobs[rank]), 4) + ) + + +@contextmanager +def mixed_decode(model, latent, rank): + """Yield measured calls, without changing existing exact-op counters. + + Each original paired decode call still consumes its original tile. At the + first call in a B4 group, compute the group together and retain its outputs + until their original consumers run. Groups never cross a gather boundary. The only linear + override executes attention to_out at its original B1 shape and stride. + """ + spatial, _, _ = geometry(model, latent) + plan = batch_plan(rank) + ordered = tuple(job for batch in plan for job in batch) + original_decode = model.decode + originals = [] + stats = dict(tile_outputs=0, decoder_calls=0, batch_sizes=[], to_out_calls=0) + pending = [] + batch_index = 0 + + def tile_view(job): + window, tile = job + sy, sx = spatial[tile] + return latent[:, :, window * 5 : window * 5 + 7, sy, sx] + + def decode(self, clip): + nonlocal batch_index + expected = tile_view(ordered[stats["tile_outputs"]]) + assert (clip.data_ptr(), clip.shape, clip.stride(), clip.dtype) == ( + expected.data_ptr(), + expected.shape, + expected.stride(), + expected.dtype, + ) + if not pending: + batch = plan[batch_index] + batch_index += 1 + x = clip if len(batch) == 1 else torch.cat([tile_view(job) for job in batch], dim=0) + decoded = original_decode(x) + assert decoded.shape == (len(batch), 3, 28, 256, 256) + assert decoded.dtype == torch.float32 + pending.extend(decoded[i : i + 1] for i in range(len(batch))) + stats["decoder_calls"] += 1 + stats["batch_sizes"].append(len(batch)) + stats["tile_outputs"] += 1 + return pending.pop(0) + + try: + for block in model.decoder.transformer_blocks: + module = block.attn.to_out + original = module.forward + originals.append((module, original)) + + def forward(self, x, _original=original): + assert x.ndim == 3 and x.shape[1:] == (1797, 2048) + assert x.dtype == torch.float16 and x.stride() == (1797 * 2048, 2048, 1) + stats["to_out_calls"] += x.shape[0] + if x.shape[0] == 1: + return _original(x) + return torch.cat([_original(x[i : i + 1]) for i in range(x.shape[0])], dim=0) + + module.forward = types.MethodType(forward, module) + model.decode = types.MethodType(decode, model) + yield stats + assert not pending and batch_index == len(plan) + assert stats["tile_outputs"] == len(ordered) + assert stats["decoder_calls"] == 21 + assert stats["to_out_calls"] == 36 * len(ordered) + finally: + model.decode = original_decode + for module, original in originals: + module.forward = original diff --git a/vllm_omni/diffusion/models/minimax_h3/vae_parallel.py b/vllm_omni/diffusion/models/minimax_h3/vae_parallel.py new file mode 100644 index 00000000000..8a2c0cb485e --- /dev/null +++ b/vllm_omni/diffusion/models/minimax_h3/vae_parallel.py @@ -0,0 +1,198 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""H3 VAE gather, stream lifetime, failure agreement, and temporal output assembly.""" + +from __future__ import annotations + +import os +from collections.abc import Callable +from typing import Any + +import torch +import torch.distributed as dist + +DecodedChunkCallback = Callable[[torch.Tensor, int, int], None] + + +def _gather_stack_to_rank_zero(segment: torch.Tensor, group: Any) -> torch.Tensor | None: + """Destination-only gather without a second full-size ``torch.cat``. + + ``dist.gather`` writes directly into contiguous views of one leader-owned + ``[world_size, *segment.shape]`` allocation. Non-leaders allocate no + receive tensor. + """ + segment = segment.contiguous() + output = segment.new_empty((group.world_size, *segment.shape)) if group.rank_in_group == 0 else None + gather_list = list(output.unbind(0)) if output is not None else None + dist.gather(segment, gather_list=gather_list, dst=group.ranks[0], group=group.device_group) + return output + + +class _LeaderAssembler: + """Serially join gathered segments while later windows decode.""" + + def __init__( + self, + model: Any, + callback: DecodedChunkCallback, + *, + main_frames: int, + overlap_frames: int, + output_frames: int, + device: torch.device, + ) -> None: + self._model = model + self._callback = callback + self._main_frames = main_frames + self._overlap_frames = overlap_frames + self._output_frames = output_frames + self._previous_overlap: torch.Tensor | None = None + self._write_pos = 0 + self._error: BaseException | None = None + self._stream = torch.get_device_module().Stream(device=device) if device.type == "cuda" else None + self._round_done: torch.cuda.Event | None = None + + def _remember_error(self, error: BaseException) -> None: + if self._error is None: + error.__traceback__ = None + self._error = error + + def push(self, segment: torch.Tensor) -> None: + if self._error is not None: + return + try: + if self._stream is None: + self._push(segment) + return + self._stream.wait_stream(torch.get_device_module().current_stream(segment.device)) + segment.record_stream(self._stream) + with torch.get_device_module().stream(self._stream): + self._push(segment) + except BaseException as error: + self._remember_error(error) + + def _push(self, segment: torch.Tensor) -> None: + main = segment[:, :, : self._main_frames] + overlap = segment[:, :, self._main_frames :] + if self._previous_overlap is not None: + main = self._model.blend(self._previous_overlap, main, self._overlap_frames, dim=-3) + frame_count = min(int(main.shape[2]), self._output_frames - self._write_pos) + if frame_count > 0: + self._callback(main[:, :, :frame_count], self._write_pos, self._output_frames) + self._write_pos += frame_count + self._previous_overlap = overlap.clone() + + def finish_round(self) -> None: + if self._stream is None or self._error is not None: + return + try: + self._round_done = torch.get_device_module().Event() + self._round_done.record(self._stream) + except BaseException as error: + self._remember_error(error) + + def drain_previous_round(self) -> None: + if self._round_done is not None: + try: + self._round_done.synchronize() + except BaseException as error: + self._remember_error(error) + finally: + self._round_done = None + + def finalize(self) -> None: + if self._error is not None: + return + try: + if self._previous_overlap is None: + raise RuntimeError("MiniMax H3 temporal chunk parallel produced no final overlap") + remaining = self._output_frames - self._write_pos + tail = self._previous_overlap[:, :, :remaining] + if self._stream is None: + self._emit_tail(tail) + else: + with torch.get_device_module().stream(self._stream): + self._emit_tail(tail) + except BaseException as error: + self._remember_error(error) + + def _emit_tail(self, tail: torch.Tensor) -> None: + if int(tail.shape[2]) > 0: + self._callback(tail, self._write_pos, self._output_frames) + self._write_pos += int(tail.shape[2]) + if self._write_pos != self._output_frames: + raise RuntimeError( + f"MiniMax H3 temporal chunk parallel output length mismatch: " + f"frames={self._write_pos}/{self._output_frames}" + ) + + def synchronize(self) -> None: + if self._stream is not None: + try: + self._stream.synchronize() + except BaseException as error: + self._remember_error(error) + + @property + def error(self) -> BaseException | None: + return self._error + + +def _agree_on_failure(group: Any, device: torch.device, failed: bool) -> bool: + flag = torch.tensor([int(failed)], dtype=torch.int32, device=device) + result = group.all_reduce(flag, op=dist.ReduceOp.MAX) + return bool(result.item()) + + +GATHER_OVERLAP_ENV = "VLLM_OMNI_H3_VAE_GATHER_OVERLAP" + + +def gather_overlap_enabled(*, paired: bool, mixed: bool) -> bool: + value = os.environ.get(GATHER_OVERLAP_ENV, "0") + if value not in ("0", "1"): + raise ValueError(f"{GATHER_OVERLAP_ENV} must be 0 or 1") + if value == "1" and not (paired and mixed): + raise ValueError(f"{GATHER_OVERLAP_ENV}=1 requires paired full-H3 VAE and mixed B2/B1") + return value == "1" + + +def prepare_gather_stream(device, *, paired: bool, mixed: bool, owner): + """Run during existing all-rank sink/setup admission, before any gather.""" + if not gather_overlap_enabled(paired=paired, mixed=mixed): + return None + if device.type != "cuda": + raise ValueError(f"{GATHER_OVERLAP_ENV}=1 requires CUDA") + for name in ("TORCH_NCCL_BLOCKING_WAIT", "NCCL_BLOCKING_WAIT"): + if os.environ.get(name, "0") != "0": + raise ValueError(f"{GATHER_OVERLAP_ENV}=1 requires {name} unset or 0") + # This VAE already requires serial decode requests (mixed_decode temporarily + # wraps model.decode). Reuse one stream per instance/device after the prior + # request's unconditional drain, so allocator blocks retain a stable owner. + index = device.index if device.index is not None else torch.accelerator.current_device_index() + streams = getattr(owner, "_h3_vae_gather_streams", None) + if streams is None: + streams = {} + owner._h3_vae_gather_streams = streams + if index not in streams: + streams[index] = torch.get_device_module().Stream(device=device) + return streams[index] + + +def gather_on_stream(packed, group, stream): + """Keep the NCCL completion dependency off the decoder caller stream. + + The existing synchronous Python gather waits for NCCL on its current CUDA + stream. Here that current stream is dedicated to gather, so the subsequent + decoder round may be submitted on the original caller stream. The leader + must wait on this gather stream before consuming the returned tensor. + """ + + caller = torch.get_device_module().current_stream(packed.device) + if stream == caller: + raise RuntimeError("VAE gather overlap requires a distinct stream") + stream.wait_stream(caller) + # `packed` was allocated/written on caller. Its Python lifetime ends after + # this round is submitted, potentially before the gather stream finishes. + packed.record_stream(stream) + with torch.get_device_module().stream(stream): + return _gather_stack_to_rank_zero(packed, group)