From 93457f29ff88b34b9e1851f207f989c3b1bb6e0b Mon Sep 17 00:00:00 2001 From: Zijun Huang Date: Thu, 6 Aug 2026 14:15:21 +0800 Subject: [PATCH] [NPU][Quantization] Add RainFusion attention and INT8 online quantization for MiniMax H3 (#5706) Signed-off-by: Zijun Huang Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com> --- docs/design/module/dit_module.md | 2 +- .../diffusion/attention_backends.md | 77 +++- recipes/MiniMaxAI/MiniMax-H3-NPU.md | 32 +- .../attention/test_attention_config.py | 29 ++ .../attention/test_rainfusion_plan.py | 157 ++++++++ .../quantization/test_int8_config.py | 124 ++++++ .../diffusion/attention/backends/abstract.py | 36 ++ .../attention/backends/rainfusion_attn.py | 366 ++++++++++++++++++ .../diffusion/attention/backends/registry.py | 1 + vllm_omni/diffusion/data.py | 46 ++- .../model_loader/diffusers_loader.py | 37 ++ .../models/minimax_h3/denoise_loop.py | 14 + .../minimax_h3/minimax_h3_transformer.py | 21 +- .../models/minimax_h3/packed_sequence.py | 6 + vllm_omni/diffusion/models/minimax_h3/vae.py | 9 +- vllm_omni/platforms/cuda/platform.py | 1 + vllm_omni/platforms/interface.py | 23 ++ vllm_omni/platforms/musa/platform.py | 1 + vllm_omni/platforms/npu/platform.py | 1 + vllm_omni/platforms/rocm/platform.py | 1 + vllm_omni/platforms/xpu/platform.py | 1 + vllm_omni/quantization/int8_config.py | 78 ++++ 22 files changed, 1055 insertions(+), 8 deletions(-) create mode 100644 tests/diffusion/attention/test_rainfusion_plan.py create mode 100644 vllm_omni/diffusion/attention/backends/rainfusion_attn.py diff --git a/docs/design/module/dit_module.md b/docs/design/module/dit_module.md index f6883f12c..36fd16877 100644 --- a/docs/design/module/dit_module.md +++ b/docs/design/module/dit_module.md @@ -539,7 +539,7 @@ def get_attn_backend_for_role(role, head_size, attention_config=None, role_categ 4. **Platform default** — `current_omni_platform.get_diffusion_attn_backend_cls(...)` picks the best-available kernel for the hardware. -A backend that needs configuration exposes it as a typed field on `AttentionSpec`; `AttentionSpec.backend_kwargs()` serializes it into the `backend_kwargs` the impl consumes. Today the only such field is `skip_softmax` (TRTLLM_ATTN), so e.g. `target_sparsity` can be set without changing model code. +A backend that needs configuration exposes it as a typed field on `AttentionSpec`; `AttentionSpec.backend_kwargs()` serializes it into the `backend_kwargs` the impl consumes. Today those fields are `skip_softmax` (TRTLLM_ATTN) and `block_sparse` (any block-sparse backend, currently RAINFUSION_ATTN), so e.g. `target_sparsity` or `sparsity` can be set without changing model code. For the user-facing CLI surface, see [Diffusion Attention Backends](../../user_guide/diffusion/attention_backends.md). For the role declaration contract on the model side, see [Adding a Diffusion Model](../../contributing/model/adding_diffusion_model.md). diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md index 92a13aa74..54e5ae8e3 100644 --- a/docs/user_guide/diffusion/attention_backends.md +++ b/docs/user_guide/diffusion/attention_backends.md @@ -23,6 +23,7 @@ The full set of backends and their platform defaults is in the **Backend Options | `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. | | `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. | | `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. | +| `RAINFUSION_ATTN` | MindIE-SD **RainFusion** block-sparse video attention — see [below](#rainfusion_attn-backend-and-block-sparse-video-attention). Ascend NPU only; requires `mindiesd`. Delegates to `FLASH_ATTN` for anything that is not a packed video sequence. | ## Configuration @@ -69,8 +70,9 @@ vllm-omni serve \ --diffusion-attention-config '{"default":{"backend":"FLASH_ATTN"},"per_role":{"cross":{"backend":"TORCH_SDPA"}}}' ``` -A backend that needs configuration exposes it as a typed field on the spec. Today the only -one is `TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)): +A backend that needs configuration exposes it as a typed field on the spec: +`TRTLLM_ATTN`'s Skip-Softmax (see [below](#trtllm_attn-backend-and-skip-softmax)) and +`RAINFUSION_ATTN`'s RainFusion block (see [below](#rainfusion_attn-backend-and-block-sparse-video-attention)): ```bash --diffusion-attention-config.default.backend TRTLLM_ATTN \ @@ -204,11 +206,80 @@ AttentionConfig( ) ``` - **Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103** (B300); `int8` QK kernels are compiled for **SM100 only**. +## RAINFUSION_ATTN Backend and Block-Sparse Video Attention + +`RAINFUSION_ATTN` runs MindIE-SD's RainFusion (`rf_v2`) kernel on Ascend NPU. It pools each +128-token block of keys, ranks them per query block, and computes attention against only the +top-scoring ones. Video tokens are rearranged into `(t, h, w)` order first, so the kept blocks are +spatiotemporal neighbours rather than arbitrary rows of the packed sequence. + +Only the video segment is sparsified. The prefix rows (text, visual conditions, audio) and the +first-frame blocks are always kept dense, which is why the realized sparsity is lower than the +nominal `sparsity` you configure. Anything the kernel cannot handle — a warmup denoise step, an +exempt layer, a sequence with no published video geometry, a video segment under 32 blocks — +delegates to `FLASH_ATTN`, so a model can select this backend unconditionally. + +Enable it through the typed `block_sparse` block on the attention spec, shared by every +block-sparse backend: + +| Key | Valid values | Meaning | +|---|---|---| +| `sparsity` | finite, `[0, 1]` | Nominal fraction of key blocks dropped per query block. `0` disables sparsity. Defaults to `0.8`. | +| `start_step` | `≥ 0` | Keep the first N denoise steps dense. Layout is decided early, so these steps dominate structural fidelity. | +| `skip_layers` | index selector, e.g. `"0-3,38"` | DiT blocks that always stay dense. | + +```bash +vllm-omni serve MiniMaxAI/MiniMax-H3 \ + --diffusion-attention-config '{"default": {"backend": "RAINFUSION_ATTN", + "block_sparse": {"sparsity": 0.8, "start_step": 0}}}' +``` + +Programmatically the same block is a typed `BlockSparseSpec` (values validated at construction): + +```python +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, BlockSparseSpec + +AttentionConfig( + default=AttentionSpec( + backend="RAINFUSION_ATTN", + block_sparse=BlockSparseSpec(sparsity=0.8, start_step=0, skip_layers="0-1"), + ), +) +``` + +**Tune in the order `start_step` → `sparsity` → `skip_layers`.** Raise `start_step` first: it is the +cheapest way to recover structure, because the early high-noise steps decide global layout while the +later steps only refine texture. Then walk `sparsity` up to your quality bar. Reach for +`skip_layers` last, once an A/B against dense at the same seed points at specific blocks. + +Requires Ascend NPU with `mindiesd`; selecting it on any other platform raises. It is also +incompatible with ring sequence parallelism, since `rf_v2` needs the whole key sequence to rank +blocks — use Ulysses SP (`ring_degree=1`). + +### Which geometries run sparse + +Sparsity is only applied when the video segment is a **multiple of 128 rows**, where the row count +is `latent_t × (height / 32) × (width / 32)`. Otherwise `rf_v2`'s block mask and the kernel's own +tiling disagree on the block count, and the block straddling the seam mixes video and prefix rows +that selection may then drop. Any resolution still runs — an unaligned geometry falls back to dense +attention and logs `RAINFUSION_ATTN staying dense` with the row count it computed — but it gets no +speedup, so pick an aligned geometry when you want one. + +Alignment is necessary, not sufficient, for good quality. `rf_v2` groups video positions into 8x8 +spatial tiles (two tiles fill one 128-row block), which is what makes a selected block a compact +patch of the frame. When the latent `h` or `w` is not a multiple of 8, the leftover rows or columns +are peeled off and appended as a flat run instead of being tiled, so they get pooled with spatially +distant positions and selection can no longer rank them meaningfully. This is silent, and invisible +to a `sparsity=0` check, because a fully populated mask does not care how blocks are grouped. + +"Multiple of 8" on the latent grid means **width and height that are multiples of 256**. Aligned +resolutions off that grid still run sparse; they just lose more fidelity at the same `sparsity`, +which you can buy back with `start_step`. + ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell) Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`. diff --git a/recipes/MiniMaxAI/MiniMax-H3-NPU.md b/recipes/MiniMaxAI/MiniMax-H3-NPU.md index e6120e916..afc3ffcc1 100755 --- a/recipes/MiniMaxAI/MiniMax-H3-NPU.md +++ b/recipes/MiniMaxAI/MiniMax-H3-NPU.md @@ -50,7 +50,8 @@ uv pip install -e . ``` Install the **mindie-sd** operator library to enable Ascend-optimized fused -operators (`adalayernorm`, etc.): +operators (`adalayernorm`, etc.) and the RainFusion `rf_v2` block-sparse +attention kernel: ```bash git clone https://gitcode.com/Ascend/MindIE-SD.git && cd MindIE-SD @@ -111,6 +112,28 @@ Do not add `--enforce-eager`. The first request includes regional compilation; warm the server once before measuring steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1. +### Optional optimizations + +Two independent optimizations may be enabled on top of the configuration +above. Both are validated for T2VA only. + +**RainFusion block-sparse attention** — switch the attention backend, keeping +every other flag unchanged: + +```bash + --diffusion-attention-backend RAINFUSION_ATTN +``` + +**INT8 online quantization** — add one flag to either command above: + +```bash + --quantization int8 +``` + +Keep `--ring 1` when using RainFusion: the `rf_v2` kernel ranks key blocks +over the whole sequence, so ring parallelism would split away the keys it +needs. Scale with `--usp` instead. + To serve Ref2VA, stop the FL2VA server and restart with: ```bash @@ -154,9 +177,16 @@ throughput guarantee. the same HTTP request contract. - VAE patch parallelism requires size 1 or the full DiT group size and supports the H3 native `tile` mode only. +- RainFusion block-sparse attention and INT8 quantization are validated for + T2VA only; use the BF16 dense configuration for FL2VA and Ref2VA. +- Online quantization cannot be combined with distributed layerwise offload + while AllGather is enabled; pass `--dlo-no-use-allgather` in that case. ## Additional resources - [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide +- [Attention backends § RAINFUSION_ATTN](../../docs/user_guide/diffusion/attention_backends.md#rainfusion_attn-backend-and-block-sparse-video-attention) + — RainFusion knobs and tuning +- [Int8 quantization](../../docs/user_guide/quantization/int8.md) - [Supported models](../../docs/models/supported_models.md) - [Video API](../../docs/serving/videos_api.md) diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index 72b409de5..4a2c65cbb 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/tests/diffusion/attention/test_attention_config.py @@ -90,6 +90,35 @@ class TestAttentionSpec: with pytest.raises(ValueError, match=match): AttentionSpec(**spec) + def test_block_sparse_defaults_applied_when_backend_selected(self): + spec = AttentionSpec(backend="RAINFUSION_ATTN") + assert spec.block_sparse.sparsity == 0.8 + assert spec.backend_kwargs() == {"sparsity": 0.8, "start_step": 0} + + def test_block_sparse_skip_layers_selector_expanded(self): + spec = AttentionSpec( + backend="RAINFUSION_ATTN", + block_sparse={"sparsity": 0.9, "start_step": 12, "skip_layers": "0-2,38"}, + ) + assert spec.block_sparse.skip_layer_indices == {0, 1, 2, 38} + assert spec.backend_kwargs() == { + "sparsity": 0.9, + "start_step": 12, + "skip_layers": [0, 1, 2, 38], + } + + def test_block_sparse_rejected_on_dense_backend(self): + with pytest.raises(ValueError, match="block_sparse is only supported by"): + AttentionSpec(backend="FLASH_ATTN", block_sparse={"sparsity": 0.8}) + + @pytest.mark.parametrize( + "block_sparse", + [{"sparsity": 1.5}, {"start_step": -1}], + ) + def test_block_sparse_invalid_values(self, block_sparse): + with pytest.raises(ValueError): + AttentionSpec(backend="RAINFUSION_ATTN", block_sparse=block_sparse) + class TestAttentionConfig: def test_empty_config(self): diff --git a/tests/diffusion/attention/test_rainfusion_plan.py b/tests/diffusion/attention/test_rainfusion_plan.py new file mode 100644 index 000000000..e88a973a9 --- /dev/null +++ b/tests/diffusion/attention/test_rainfusion_plan.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Geometry contract RAINFUSION_ATTN enforces before handing a forward to rf_v2. + +The grids here are the ones MiniMax-H3 FL2VA actually produces for an 8.7s clip: +1280x768 gives a 62x24x40 latent grid whose 59520 video rows land on the 128-row +kernel block, and 1344x768 gives 62x24x42 whose 62496 rows straddle it. Only the +first can run sparse: rf_v2 tiles the sequence as a whole but pools its block +mask from the video and prefix segments separately, so the two only agree on the +block count when the video segment ends on a boundary. +""" + +import dataclasses + +import pytest +import torch + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout +from vllm_omni.diffusion.attention.backends.rainfusion_attn import ( + _BLOCK_SIZE, + RainFusionAttentionImpl, + RainFusionPlan, +) +from vllm_omni.platforms import current_omni_platform + +PREFIX_ROWS = 710 # 14 text rows + 696 audio rows +ALIGNED_GRID = (62, 24, 40) # 1280x768 -> 59520 video rows, 465 blocks +MISALIGNED_GRID = (62, 24, 42) # 1344x768 -> 62496 video rows, 488 blocks + 32 rows + + +def make_impl(**backend_kwargs): + return RainFusionAttentionImpl( + num_heads=8, + head_size=128, + softmax_scale=128**-0.5, + prefix="transformer_blocks.0.attn", + qkv_layout="BSND", + backend_kwargs={"sparsity": 0.8, **backend_kwargs}, + ) + + +def make_metadata(grid, prefix_len=PREFIX_ROWS, max_seqlen_q=None): + video_rows = grid[0] * grid[1] * grid[2] + return AttentionMetadata( + extra={"max_seqlen_q": prefix_len + video_rows if max_seqlen_q is None else max_seqlen_q}, + video_layout=VideoTokenLayout(prefix_len=prefix_len, latent_grid=grid), + ) + + +def test_block_aligned_video_segment_runs_sparse(): + plan = make_impl()._resolve_plan(make_metadata(ALIGNED_GRID)) + + assert plan is not None + assert plan.prefix_len == PREFIX_ROWS + assert plan.used_len == PREFIX_ROWS + 59520 + assert plan.latent_shape == list(ALIGNED_GRID) + + +def test_misaligned_video_segment_stays_dense(): + # Realigning the sequence would mean padding it, and rf_v2 ignores attn_mask: + # pad keys would take a share of every real query's softmax denominator. + assert make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID)) is None + + +@pytest.mark.parametrize("prefix_len", [1, 127, _BLOCK_SIZE, 710, 900]) +def test_prefix_length_does_not_affect_alignment(prefix_len): + # Only the video segment has to land on a block boundary; rf_v2 pools the + # prefix separately and keeps every one of its blocks. + plan = make_impl()._resolve_plan(make_metadata(ALIGNED_GRID, prefix_len=prefix_len)) + + assert plan is not None + assert plan.prefix_len == prefix_len + + +def test_sparsity_zero_never_resolves_a_plan(): + assert make_impl(sparsity=0.0)._resolve_plan(make_metadata(ALIGNED_GRID)) is None + + +def test_missing_video_layout_falls_back_to_dense(): + assert make_impl()._resolve_plan(AttentionMetadata(extra={"max_seqlen_q": 60230})) is None + + +def test_missing_max_seqlen_falls_back_to_dense(): + metadata = AttentionMetadata( + video_layout=VideoTokenLayout(prefix_len=PREFIX_ROWS, latent_grid=ALIGNED_GRID), + ) + + assert make_impl()._resolve_plan(metadata) is None + + +def test_video_segment_must_be_the_tail_of_packed_document_zero(): + metadata = make_metadata(ALIGNED_GRID, max_seqlen_q=PREFIX_ROWS + 59520 + 128) + + assert make_impl()._resolve_plan(metadata) is None + + +@pytest.mark.parametrize("grid", [(4, 24, 40), (1, 24, 40)]) +def test_short_video_stays_dense(grid): + assert make_impl()._resolve_plan(make_metadata(grid)) is None + + +@pytest.mark.parametrize("qkv_layout", ["BNSD", "BSH"]) +def test_explicitly_wrong_layout_is_rejected(qkv_layout): + with pytest.raises(ValueError, match="BSND"): + RainFusionAttentionImpl( + num_heads=8, + head_size=128, + softmax_scale=128**-0.5, + qkv_layout=qkv_layout, + backend_kwargs={"sparsity": 0.8}, + ) + + +def test_undeclared_layout_stays_dense(): + # An undeclared layout is not an error -- the layer keeps working, just + # densely, and the fallback sees the same absent layout plain FLASH_ATTN would. + impl = RainFusionAttentionImpl( + num_heads=8, + head_size=128, + softmax_scale=128**-0.5, + prefix="transformer_blocks.0.attn", + backend_kwargs={"sparsity": 0.8}, + ) + + assert impl._resolve_plan(make_metadata(ALIGNED_GRID)) is None + assert impl.dense_fallback.qkv_layout is None + + +@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="rf_v2 runs on Ascend NPU only.") +def test_fully_populated_mask_reproduces_dense_attention(): + """At sparsity=0 every key block is kept, so rf_v2 must match dense attention. + + This is the guard on the geometry the backend hands the kernel: a wrong + prefix length or latent grid still produces plausible output under sparse + selection, but shows up here as a mismatch against SDPA. + """ + torch.manual_seed(0) + prefix_len, grid = 384, (4, 16, 64) # 4096 video rows = 32 blocks + video_len = grid[0] * grid[1] * grid[2] + used = prefix_len + video_len + heads, head_dim = 4, 128 + q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3)) + + impl = make_impl() + impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0) + plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid)) + out = impl._forward_sparse_npu(q, k, v, plan) + + reference = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), + k.transpose(1, 2), + v.transpose(1, 2), + scale=impl.softmax_scale, + ).transpose(1, 2) + error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean() + assert error < 2e-3, f"mean relative error {error:.4%} against dense attention" diff --git a/tests/diffusion/quantization/test_int8_config.py b/tests/diffusion/quantization/test_int8_config.py index 875277ece..f706cc68b 100644 --- a/tests/diffusion/quantization/test_int8_config.py +++ b/tests/diffusion/quantization/test_int8_config.py @@ -147,6 +147,130 @@ def test_get_npu_quant_method(mocker: MockerFixture, monkeypatch: pytest.MonkeyP assert isinstance(method, UnquantizedLinearMethod) +def test_fused_layer_can_be_ignored_by_its_prefix(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture): + """A layer opts out of quantization under the prefix its module was built with. + + MiniMax H3 stores qkv as one checkpoint tensor, so the fused name is the + only handle a user has on it. + """ + from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod + + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False) + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True) + config = build_quant_config("int8", ignored_layers=["blocks.0.attn.qkv_proj"]) + layer = mocker.Mock(spec=LinearBase) + + assert isinstance(config.get_quant_method(layer, "blocks.0.attn.qkv_proj"), UnquantizedLinearMethod) + assert isinstance(config.get_quant_method(layer, "blocks.1.attn.qkv_proj"), NPUInt8OnlineLinearMethod) + assert isinstance(config.get_quant_method(layer, "blocks.0.mlp.fc1"), NPUInt8OnlineLinearMethod) + + +class TestNPUQuantMatmulShapeFallback: + """npu_quant_matmul rejects wide outputs, so those layers stay unquantized. + + The limit applies to the per-rank shard, which is what the kernel sees. The + same layer can therefore be quantized at a higher TP degree. + """ + + @pytest.fixture(autouse=True) + def _mock_tp(self, mocker): + # The fallback delegates to UnquantizedLinearMethod, which reads the TP + # group; stand one in so the shape gating can be tested without a + # distributed init. + mock_group = mocker.Mock() + mock_group.rank_in_group = 0 + mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1) + mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0) + mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group) + + @pytest.fixture + def quant_config(self): + from vllm_omni.quantization.int8_config import DiffusionInt8Config + + return DiffusionInt8Config(is_checkpoint_int8_serialized=False, activation_scheme="dynamic") + + def create_weights(self, method, layer, output_partition_sizes): + # Narrow input dim: the fallback path allocates for real, and only the + # output dimension is under test. + method.create_weights( + layer, + input_size_per_partition=8, + output_partition_sizes=output_partition_sizes, + input_size=8, + output_size=sum(output_partition_sizes), + params_dtype=torch.bfloat16, + weight_loader=lambda param, loaded_weight, *args, **kwargs: None, + ) + + @pytest.mark.parametrize( + "method_name", + ["NPUInt8LinearMethod", "NPUInt8OnlineLinearMethod"], + ) + def test_over_wide_layer_falls_back_to_unquantized(self, quant_config, method_name): + import vllm_omni.quantization.int8_config as int8_config + + method = getattr(int8_config, method_name)(quant_config) + layer = Module() + layer.quant_method = method + + self.create_weights(method, layer, [int8_config.NPU_QUANT_MATMUL_MAX_OUT_FEATURES + 1]) + + assert isinstance(layer.quant_method, UnquantizedLinearMethod) + assert layer.weight.dtype == torch.bfloat16 + + def test_layer_within_the_limit_keeps_int8(self, quant_config): + from vllm_omni.quantization.int8_config import ( + NPU_QUANT_MATMUL_MAX_OUT_FEATURES, + NPUInt8OnlineLinearMethod, + ) + + method = NPUInt8OnlineLinearMethod(quant_config) + layer = Module() + layer.quant_method = method + + self.create_weights(method, layer, [NPU_QUANT_MATMUL_MAX_OUT_FEATURES]) + + assert layer.quant_method is method + + def test_tensor_parallel_sharding_brings_a_wide_layer_back(self, quant_config): + from vllm_omni.quantization.int8_config import ( + NPU_QUANT_MATMUL_MAX_OUT_FEATURES, + NPUInt8OnlineLinearMethod, + ) + + global_out_features = 4 * NPU_QUANT_MATMUL_MAX_OUT_FEATURES + method = NPUInt8OnlineLinearMethod(quant_config) + layer = Module() + layer.quant_method = method + + self.create_weights(method, layer, [global_out_features // 4]) + + assert layer.quant_method is method + + +class TestOffloadAfterQuant: + def test_only_methods_advertising_the_capability_are_asked(self): + from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader + from vllm_omni.quantization.int8_config import DiffusionInt8Config, NPUInt8OnlineLinearMethod + + class MetaDeviceOnlyMethod: + """Stands in for upstream online FP8: lazy, but no offload hook.""" + + uses_meta_device = True + + model = torch.nn.Sequential(Module(), Module()) + lazy_int8 = NPUInt8OnlineLinearMethod( + DiffusionInt8Config(is_checkpoint_int8_serialized=False, activation_scheme="dynamic") + ) + model[0].quant_method = lazy_int8 + model[1].quant_method = MetaDeviceOnlyMethod() + + marked = DiffusersPipelineLoader._request_offload_after_quant(model) + + assert marked == 1 + assert lazy_int8._offload_after_quant + + class TestInt8LinearMethod: @pytest.fixture def mock_quant_config(self, mocker): diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index fd5a8e01e..68e152063 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -20,6 +20,21 @@ class AttentionBackend(ABC): # avoid a slower masked-attention plan when tail padding is not semantic. supports_prefix_kv_slicing: bool = False + # ``OmniPlatformEnum`` values this backend runs on; None means unrestricted. + # Platform resolution rejects an explicit selection outside this set, so a + # hardware-specific backend fails before the model is built. + supported_platforms: tuple[str, ...] | None = None + + @classmethod + def validate_available(cls) -> None: + """Raise if this backend's optional dependencies are missing. + + Called during platform resolution, i.e. before model construction, so a + backend that probes its kernel package lazily still reports the problem + while the user can still act on it. + """ + return None + @classmethod def supports_attention_mask(cls) -> bool: return False @@ -63,6 +78,23 @@ class QueryRange: global_start: int +@dataclass(frozen=True, slots=True) +class VideoTokenLayout: + """Where the video segment sits inside a packed multimodal sequence. + + A model that packs its sequence as ``[prefix | t*h*w video rows | padding]`` + publishes this so backends can recover spatiotemporal locality; the prefix + holds everything that is not video (text, visual conditions, audio). + Publishing it also asserts that any ``attn_mask`` masks only the trailing + padding, so ``prefix_len + t*h*w`` is the used length of the sequence. + + Plain ints, so reading it never forces a device-to-host sync. + """ + + prefix_len: int + latent_grid: tuple[int, int, int] + + @dataclass class AttentionMetadata: attn_mask: torch.Tensor | None = None @@ -95,6 +127,10 @@ class AttentionMetadata: full_attn_spans: list[list[tuple[int, int]]] | None = None query_ranges: tuple[QueryRange, ...] | None = None + # Geometry of the video segment for backends that exploit spatiotemporal + # locality (block-sparse selection, tiled masks). Dense backends ignore it. + video_layout: VideoTokenLayout | None = None + T = TypeVar("T", bound=AttentionMetadata) diff --git a/vllm_omni/diffusion/attention/backends/rainfusion_attn.py b/vllm_omni/diffusion/attention/backends/rainfusion_attn.py new file mode 100644 index 000000000..10ab55991 --- /dev/null +++ b/vllm_omni/diffusion/attention/backends/rainfusion_attn.py @@ -0,0 +1,366 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +import torch +from vllm.logger import init_logger +from vllm.model_executor.models.utils import extract_layer_index + +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, +) +from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionBackend +from vllm_omni.diffusion.config import get_current_diffusion_config_or_none +from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available + +logger = init_logger(__name__) + +# The rf_v2 kernel only implements a 128-token block. +_BLOCK_SIZE = 128 + +# Below this many video blocks, the pooling and gather that block selection adds +# cost more than the QK work it removes, so stay dense. +_MIN_VIDEO_BLOCKS = 32 + +# rf_v2's ``input_layout`` describes the caller's tensors, and everything below +# slices the sequence on dim 1. vLLM-Omni diffusion attention hands the impl +# [B, S, N, D], so that is the only layout this backend accepts. +_INPUT_LAYOUT = "BSND" + +# rf_v2 precision mode: 0 = high precision, 1 = high performance. Kept at the +# precise setting because the sparsity knob is the intended perf lever, and on +# A5 devices the kernel is routed to rf_v3, which overrides this anyway. +_INNER_PRECISE = 0 + +_WRONG_PLATFORM = ( + "RAINFUSION_ATTN runs the MindIE-SD rf_v2 kernel and is available on Ascend NPU only. " + "Select FLASH_ATTN or TORCH_SDPA on this platform." +) + +_MISSING_MINDIESD = ( + "RAINFUSION_ATTN requires MindIE-SD. Please install MindIE-SD to enable RainFusion sparse " + "attention on Ascend NPU. For installation details, see https://gitcode.com/Ascend/MindIE-SD " + "Otherwise, use FlashAttention by setting DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN" +) + + +def _try_extract_layer_index(prefix: str) -> int | None: + if not prefix: + return None + try: + return extract_layer_index(prefix) + except (AssertionError, ValueError): + return None + + +@dataclass(frozen=True) +class RainFusionConfig: + """Resolved RainFusion controls for one attention layer. + + ``sparsity`` is the nominal fraction of key blocks dropped per query block. + The realized sparsity is lower because rf_v2 always keeps the prefix rows and + the first-frame blocks. ``start_step`` and ``skip_layers`` are the accuracy + knobs: early denoise steps and specific DiT blocks stay dense. + """ + + sparsity: float = 0.0 + start_step: int = 0 + skip_layers: frozenset[int] = frozenset() + + @classmethod + def from_backend_kwargs(cls, backend_kwargs: dict | None) -> RainFusionConfig: + bk = backend_kwargs or {} + return cls( + sparsity=float(bk.get("sparsity", 0.0)), + start_step=int(bk.get("start_step", 0)), + skip_layers=frozenset(bk.get("skip_layers") or ()), + ) + + @property + def enabled(self) -> bool: + return self.sparsity > 0.0 + + +@dataclass(frozen=True) +class RainFusionPlan: + """Per-forward geometry handed to the rf_v2 kernel.""" + + prefix_len: int + used_len: int + latent_shape: list[int] + + +class RainFusionAttentionBackend(AttentionBackend): + accept_output_buffer: bool = True + supported_platforms: tuple[str, ...] = ("npu",) + + @classmethod + def validate_available(cls) -> None: + from importlib.util import find_spec + + if find_spec("mindiesd") is None: + raise ValueError(_MISSING_MINDIESD) + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [64, 96, 128, 192, 256] + + @staticmethod + def get_name() -> str: + return "RAINFUSION_ATTN" + + @staticmethod + def get_impl_cls() -> type[RainFusionAttentionImpl]: + return RainFusionAttentionImpl + + +class RainFusionAttentionImpl(AttentionImpl): + """Block-sparse video attention via MindIE-SD RainFusion (rf_v2) on Ascend NPU. + + Sparsity applies only to the video segment of a packed multimodal sequence, + whose extent the model publishes as ``AttentionMetadata.video_layout``. Every + other case — warmup denoise steps, exempt layers, a layer that does not declare + ``qkv_layout="BSND"``, sequences without a published video segment, video + segments too short to pay for block selection or not aligned to the kernel's + block size — delegates to FlashAttention, so a model can select this backend + unconditionally. + """ + + def __init__( + self, + num_heads: int, + head_size: int, + softmax_scale: float, + causal: bool = False, + num_kv_heads: int | None = None, + prefix: str = "", + qkv_layout: str | None = None, + backend_kwargs: dict[str, Any] | None = None, + **extra_impl_args, + ) -> None: + self.num_heads = num_heads + self.causal = causal + self.softmax_scale = softmax_scale + self.qkv_layout = qkv_layout + + self.rainfusion = RainFusionConfig.from_backend_kwargs(backend_kwargs) + self.layer_idx = _try_extract_layer_index(prefix) + + if self.rainfusion.enabled: + self._validate_parallel_config() + if causal: + raise ValueError( + "RAINFUSION_ATTN does not support causal attention: rf_v2 selects key " + "blocks by pooled relevance and cannot express a causal mask. Select " + "FLASH_ATTN for causal roles." + ) + if qkv_layout is not None and qkv_layout.upper() != _INPUT_LAYOUT: + raise ValueError( + f"RAINFUSION_ATTN needs {_INPUT_LAYOUT} tensors to locate the video segment along " + f"the sequence axis, but this layer declares qkv_layout={qkv_layout!r}. Select " + "FLASH_ATTN for this role." + ) + + self.dense_fallback = FlashAttentionBackend.get_impl_cls()( + num_heads=num_heads, + head_size=head_size, + softmax_scale=softmax_scale, + causal=causal, + num_kv_heads=num_kv_heads, + prefix=prefix, + qkv_layout=qkv_layout, + ) + + def _validate_parallel_config(self) -> None: + config = get_current_diffusion_config_or_none() + parallel_config = getattr(config, "parallel_config", None) + ring_degree = getattr(parallel_config, "ring_degree", 1) + if ring_degree > 1: + # Ring gives each rank a slice of the sequence, so block selection + # would score only local keys and the layer bypasses the backend + # entirely (see Attention._run_ring_attention). + raise ValueError( + "RAINFUSION_ATTN is not compatible with ring sequence parallelism " + f"(ring_degree={ring_degree}): rf_v2 needs the whole key sequence to rank " + "blocks. Use Ulysses SP (ring_degree=1) instead." + ) + + def forward_cuda( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + # ROCm and MUSA route through forward_cuda by default, so this covers them too. + raise NotImplementedError(_WRONG_PLATFORM) + + def forward_xpu( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + raise NotImplementedError(_WRONG_PLATFORM) + + def forward_npu( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + plan = self._resolve_plan(attn_metadata) + if plan is None: + return self.dense_fallback.forward_npu(query, key, value, attn_metadata) + return self._forward_sparse_npu(query, key, value, plan) + + def _resolve_plan(self, attn_metadata: AttentionMetadata | None) -> RainFusionPlan | None: + """Return the rf_v2 geometry, or None when this forward must stay dense.""" + rf = self.rainfusion + if not rf.enabled: + return None + if self.layer_idx is not None and self.layer_idx in rf.skip_layers: + return None + if is_forward_context_available(): + step_idx = get_forward_context().denoise_step_idx + if step_idx is not None and step_idx < rf.start_step: + return None + if self.qkv_layout is None: + # The sparse path reads the sequence off dim 1, which the tensors alone + # do not establish, and the dense fallback resolves an absent layout its + # own way. Sparsifying on an assumption would put the two paths on + # different axes, so an undeclared layout stays dense. + logger.warning_once( + "RAINFUSION_ATTN staying dense: this layer does not declare qkv_layout, and rf_v2 " + "needs %s to locate the video segment along the sequence axis. Set qkv_layout=%r on " + "the Attention layer to enable sparsity.", + _INPUT_LAYOUT, + _INPUT_LAYOUT, + ) + return None + + if attn_metadata is None: + return None + + layout = attn_metadata.video_layout + if layout is None: + logger.warning_once( + "RAINFUSION_ATTN staying dense: this attention role carries no video segment. The " + "model must publish AttentionMetadata.video_layout for the sequence to be sparsified." + ) + return None + max_seqlen_q = attn_metadata.extra.get("max_seqlen_q") + if max_seqlen_q is None: + logger.warning_once( + "RAINFUSION_ATTN staying dense: attention metadata is missing max_seqlen_q, so the " + "video segment cannot be confirmed to be the tail of packed document 0." + ) + return None + + prefix_len = int(layout.prefix_len) + latent_shape = [int(dim) for dim in layout.latent_grid] + # rf_v2 splits the sequence as [prefix | t*h*w video rows]. Document 0 of + # the packed sequence holds those rows; anything past it is alignment + # padding that rf_v2 must not see. + video_len = math.prod(latent_shape) + used_len = prefix_len + video_len + + if used_len != int(max_seqlen_q): + logger.warning_once( + "RAINFUSION_ATTN staying dense: prefix (%d) plus latent grid %s does not fill " + "packed document 0 (%d rows). rf_v2 requires the video segment to be its tail.", + prefix_len, + tuple(latent_shape), + int(max_seqlen_q), + ) + return None + if video_len < _MIN_VIDEO_BLOCKS * _BLOCK_SIZE: + logger.warning_once( + "RAINFUSION_ATTN staying dense: %d video rows is under the %d-row " + "(%d block) threshold where sparse selection pays off.", + video_len, + _MIN_VIDEO_BLOCKS * _BLOCK_SIZE, + _MIN_VIDEO_BLOCKS, + ) + return None + if video_len % _BLOCK_SIZE != 0: + # rf_v2 tiles the sequence into ceil((video + prefix) / 128) blocks but + # builds its block mask from the two segments pooled apart. The counts + # only agree when the video segment ends on a block boundary; otherwise + # a prefix block is forced on out of range, and the block straddling the + # seam holds real prefix rows that selection may drop. Padding the + # sequence to realign it is not an option because rf_v2 ignores + # attn_mask, so pad keys would take a share of every softmax denominator. + logger.warning_once( + "RAINFUSION_ATTN staying dense: latent grid %s gives %d video rows, which is not a " + "multiple of the %d-row kernel block. Choose a geometry whose latent_t * h * w is a " + "multiple of %d to enable sparsity.", + tuple(latent_shape), + video_len, + _BLOCK_SIZE, + _BLOCK_SIZE, + ) + return None + + logger.info_once( + "RAINFUSION_ATTN active: sparsity=%.2f, start_step=%d, exempt_layers=%d, " + "latent_grid=%s, prefix_rows=%d, video_rows=%d. Realized sparsity is lower than nominal " + "because prefix and first-frame blocks are always kept.", + rf.sparsity, + rf.start_step, + len(rf.skip_layers), + tuple(latent_shape), + prefix_len, + video_len, + ) + return RainFusionPlan( + prefix_len=prefix_len, + used_len=used_len, + latent_shape=latent_shape, + ) + + def _forward_sparse_npu( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + plan: RainFusionPlan, + ) -> torch.Tensor: + try: + from mindiesd import sparse_attention + except ImportError: + raise ImportError(_MISSING_MINDIESD) + + used = plan.used_len + q, k, v = (tensor[:, :used] for tensor in (query, key, value)) + # Ulysses has already gathered the full sequence onto this rank and split + # the heads, so read the head count off the tensor rather than num_heads. + out = sparse_attention( + q, + k, + v, + scale=self.softmax_scale, + head_num=query.shape[-2], + input_layout=_INPUT_LAYOUT, + inner_precise=_INNER_PRECISE, + sparse_type="rf_v2", + txt_len=plan.prefix_len, + block_size=_BLOCK_SIZE, + latent_shape_q=plan.latent_shape, + latent_shape_k=plan.latent_shape, + sparsity=self.rainfusion.sparsity, + ) + if used == query.shape[1]: + return out + padded = torch.zeros_like(query) + padded[:, :used] = out + return padded diff --git a/vllm_omni/diffusion/attention/backends/registry.py b/vllm_omni/diffusion/attention/backends/registry.py index 92d3b6926..fdfaf5f9c 100644 --- a/vllm_omni/diffusion/attention/backends/registry.py +++ b/vllm_omni/diffusion/attention/backends/registry.py @@ -65,6 +65,7 @@ class DiffusionAttentionBackendEnum(Enum, metaclass=_DiffusionBackendEnumMeta): FLASH_ATTN_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttentionHubBackend" FLASH_ATTN_3_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttention3HubBackend" TRTLLM_ATTN = "vllm_omni.diffusion.attention.backends.trtllm_attn.TrtllmAttentionBackend" + RAINFUSION_ATTN = "vllm_omni.diffusion.attention.backends.rainfusion_attn.RainFusionAttentionBackend" def get_path(self, include_classname: bool = True) -> str: """Get the class path for this backend (respects overrides). diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index e47b642ff..af205713e 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -1461,6 +1461,33 @@ class AttnQuantSpec: return self.dtype_qk is not None or self.dtype_vo is not None +# Backends that select key blocks instead of attending densely, and so accept +# a ``block_sparse`` spec. Each maps the same knobs onto its own kernel. +BLOCK_SPARSE_BACKENDS = frozenset({"RAINFUSION_ATTN"}) + + +@dataclass +class BlockSparseSpec: + """User-facing controls shared by block-sparse attention backends. + + ``sparsity`` is the nominal fraction of key blocks dropped per query block; + ``start_step`` keeps the first N denoise steps dense and ``skip_layers`` (an + index selector such as "0-3,38") exempts individual DiT blocks. Those two are + the accuracy knobs to trade back quality at a fixed ``sparsity``. + """ + + sparsity: float = 0.8 + start_step: int = 0 + skip_layers: str | list[int] | None = None + skip_layer_indices: set[int] | None = field(default=None, repr=False) + + def __post_init__(self) -> None: + self.sparsity = _in_range(self.sparsity, "block_sparse.sparsity", 0.0, 1.0) or 0.0 + if self.start_step < 0: + raise ValueError(f"block_sparse.start_step must be >= 0; got {self.start_step!r}.") + self.skip_layer_indices = parse_kv_cache_skip_selector(self.skip_layers) + + @dataclass class AttentionSpec: """Specifies a backend and its typed backend-specific config for one attention role.""" @@ -1468,6 +1495,7 @@ class AttentionSpec: backend: str skip_softmax: SkipSoftmaxSpec | None = None quant: AttnQuantSpec | None = None + block_sparse: BlockSparseSpec | None = None skip_calibration: dict | None = field(default=None, repr=False) def __post_init__(self) -> None: @@ -1475,6 +1503,7 @@ class AttentionSpec: raise TypeError(f"Expected str for AttentionSpec.backend, got {type(self.backend)!r}") self.skip_softmax = self._coerce(self.skip_softmax, SkipSoftmaxSpec, "skip_softmax") self.quant = self._coerce(self.quant, AttnQuantSpec, "quant") + self.block_sparse = self._coerce(self.block_sparse, BlockSparseSpec, "block_sparse") if self.skip_softmax is not None and self.backend.upper() != "TRTLLM_ATTN": raise ValueError( f"skip_softmax is only supported by the TRTLLM_ATTN backend, but backend={self.backend!r}. " @@ -1485,6 +1514,15 @@ class AttentionSpec: f"quant is only supported by the TRTLLM_ATTN and FLASHINFER_ATTN backends, but " f"backend={self.backend!r}. Remove quant or set a supported backend." ) + if self.backend.upper() in BLOCK_SPARSE_BACKENDS: + # Selecting the backend is the opt-in; without an explicit block the + # defaults apply rather than silently running dense. + self.block_sparse = self.block_sparse or BlockSparseSpec() + elif self.block_sparse is not None: + raise ValueError( + f"block_sparse is only supported by the {sorted(BLOCK_SPARSE_BACKENDS)} backends, but " + f"backend={self.backend!r}. Remove block_sparse or set a supported backend." + ) @staticmethod def _coerce(value: Any, cls: type, field_name: str) -> Any: @@ -1517,6 +1555,12 @@ class AttentionSpec: if q.flashinfer_backend is not None: quant_kw["flashinfer_backend"] = q.flashinfer_backend kw["quant"] = quant_kw + if self.block_sparse is not None: + bs = self.block_sparse + kw["sparsity"] = bs.sparsity + kw["start_step"] = bs.start_step + if bs.skip_layer_indices: + kw["skip_layers"] = sorted(bs.skip_layer_indices) return kw or None @@ -1586,7 +1630,7 @@ class AttentionConfig: normalized[role] = node return - spec_keys = {"backend", "skip_softmax", "quant"} + spec_keys = {"backend", "skip_softmax", "quant", "block_sparse"} node_dict = dict(node) node_keys = set(node_dict) if node_keys & spec_keys: diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index df7aa2523..d1a1ed15e 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -393,6 +393,13 @@ class DiffusersPipelineLoader: "DLO+AllGather active but model does not support mmap: loading weights via regular loader." ) logger.debug("Loading weights on %s ...", load_device) + if offload_after_quant: + marked = self._request_offload_after_quant(model) + if marked: + logger.info( + "Online quantization will return each of %d layers to CPU as it is quantized", + marked, + ) if load_format == "diffusers": cast(DiffusersAdapterPipeline, model).load_weights() else: @@ -406,6 +413,29 @@ class DiffusersPipelineLoader: self._apply_skip_softmax_calibration(model) return model.eval() + @staticmethod + def _request_offload_after_quant(model: nn.Module) -> int: + """Ask online-quant layers to return to host memory once quantized. + + The weights only visit the accelerator so the quant kernels can run on + them; ``load_model`` sends the model back to the host afterwards either + way. Without this the whole transformer accumulates on device until that + final move, which for MiniMax H3 is a ~43 GiB peak that no longer fits + beside a resident TP-sharded text encoder — even though layer-wise + offload means none of it is supposed to be resident at inference time. + + Only quant methods that advertise ``supports_offload_after_quant`` are + asked, since the implementation has to know when a layer is finished. + Deferring materialization to the ``meta`` device does not imply that. + """ + marked = 0 + for module in model.modules(): + quant_method = getattr(module, "quant_method", None) + if getattr(quant_method, "supports_offload_after_quant", False): + quant_method.enable_offload_after_quant() + marked += 1 + return marked + @staticmethod def _has_online_quant(model: nn.Module) -> bool: """Whether any layer uses an online-quant method that defers weight @@ -460,6 +490,13 @@ class DiffusersPipelineLoader: if quant_method is None or not isinstance(quant_method, QuantizeMethodBase): continue + # Layers finished during loading would only be staged onto the target + # device for a process call that immediately returns. That round trip + # is wasted work in general, and undoes the point of the offload for + # layers that already went back to the host. + if getattr(module, "_already_called_process_weights_after_loading", False): + continue + if has_online_quant: # Online quant may leave straggler params on the ``meta`` device. # Move only real (non-meta) params onto the target device for diff --git a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py index b051951a7..9b4402f30 100644 --- a/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py +++ b/vllm_omni/diffusion/models/minimax_h3/denoise_loop.py @@ -14,6 +14,9 @@ from typing import Any import torch +from vllm_omni.diffusion.attention.backends.abstract import VideoTokenLayout +from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx + from .scheduling_minimax_h3_euler_ancestral import ( minimax_h3_euler_eta0_step, minimax_h3_rf_v_to_x0, @@ -87,6 +90,13 @@ class MiniMaxH3DenoiseBranch: "max_seqlen_q": text_len, }, } + # Where the video segment sits in the packed sequence. Resolved to plain + # ints here so the attention layers never sync on it per step. + grid = packed["latent_grid"].tolist() + self.static_kwargs["video_token_layout"] = VideoTokenLayout( + prefix_len=int(packed["video_row_start"]), + latent_grid=(int(grid[0]), int(grid[1]), int(grid[2])), + ) def forward_kwargs( self, @@ -191,6 +201,9 @@ def minimax_h3_denoise_loop( for step in range(num_steps): step_cm = step_profiler(step) if step_profiler is not None else nullcontext() with step_cm: + # Publish the step index so step-gated attention features (e.g. the + # dense warmup of RAINFUSION_ATTN) can see where we are. + set_forward_context_denoise_step_idx(step) s_v, s_v_next = sigmas_video[step], sigmas_video[step + 1] s_a, s_a_next = sigmas_audio[step], sigmas_audio[step + 1] t_v, t_a = 1.0 - s_v, 1.0 - s_a @@ -236,6 +249,7 @@ def minimax_h3_denoise_loop( if on_step is not None: on_step(step, video_rows, audio_rows) + set_forward_context_denoise_step_idx(None) return video_rows, audio_rows 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 c48bbea48..d98c07545 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py @@ -26,7 +26,7 @@ from vllm.model_executor.layers.linear import ( ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig from vllm_omni.diffusion.distributed.sp_plan import ( @@ -134,6 +134,7 @@ _FORWARD_SUPPORTED_KWARGS = frozenset( "img_pos_for_infer_output_info", "packed_seq_params", "refiner_packed_seq_params", + "video_token_layout", } ) @@ -348,7 +349,11 @@ class MiniMaxH3Attention(nn.Module): head_size=self.head_dim, softmax_scale=self.softmax_scale, causal=False, + # Packed rows reach the impl as [B, S, N, D]. + qkv_layout="BSND", skip_sequence_parallel=skip_sequence_parallel, + role="self", + prefix=prefix, ) def _apply_rope(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: @@ -374,6 +379,7 @@ class MiniMaxH3Attention(nn.Module): cu_seqlens: torch.Tensor, max_seqlen: int, packed_total: int, + video_layout: VideoTokenLayout | None = None, ) -> torch.Tensor: """Run packed attention as a small eager island. @@ -407,6 +413,7 @@ class MiniMaxH3Attention(nn.Module): "max_seqlen_k": max_seqlen, "valid_kv_length": used, }, + video_layout=video_layout, ) return self.attention( q.unsqueeze(0), @@ -424,6 +431,7 @@ class MiniMaxH3Attention(nn.Module): max_seqlen: int, packed_total: int | None = None, sp_seq_lens: list[int] | None = None, + video_layout: VideoTokenLayout | None = None, ) -> torch.Tensor: """x: [T, hidden] packed thd rows -> [T, hidden]. @@ -463,6 +471,7 @@ class MiniMaxH3Attention(nn.Module): # backend receives the global sequence after all-to-all, so carry # its Python length explicitly instead of inferring it from q. packed_total=packed_total if packed_total is not None else q.shape[0], + video_layout=video_layout, ) out = out.reshape(total, self.num_heads * self.head_dim) out, _ = self.out_proj(out) @@ -642,6 +651,8 @@ class MiniMaxH3DiTBlock(nn.Module): super().__init__() self.norm1 = _norm(arch.hidden_size, eps=arch.norm_eps) self.norm2 = _norm(arch.hidden_size, eps=arch.norm_eps) + # The prefix also carries the block index that block-sparse attention + # backends match against their skip_layers selector. self.attn = MiniMaxH3Attention( arch, quant_config, @@ -672,6 +683,7 @@ class MiniMaxH3DiTBlock(nn.Module): max_seqlen: int, packed_total: int, sp_seq_lens: list[int] | None = None, + video_layout: VideoTokenLayout | None = None, ) -> torch.Tensor: """x: [T, H]; t_emb: [M, t_dim]; combined_indices: [T] (= inverse_indices * modality_num + token_tags.clamp(min=0)). @@ -699,6 +711,7 @@ class MiniMaxH3DiTBlock(nn.Module): max_seqlen=max_seqlen, packed_total=packed_total, sp_seq_lens=sp_seq_lens, + video_layout=video_layout, ) x = _modulate_gate(residual, gate_msa, h, combined_indices, dtype=_BF16_DTYPE) @@ -831,6 +844,10 @@ class MiniMaxH3DiTModel(nn.Module): }, "sp_gather": SequenceParallelOutput(gather_dim=0, expected_dims=2), } + # The checkpoint already stores qkv and the MLP gate/up as single tensors + # (see the reordering in load_weights), so there are no unfused names for + # quantization or LoRA to map onto. Address the fused layers directly, e.g. + # ignored_layers=["blocks.0.attn.qkv_proj"]. packed_modules_mapping = {} def _validate_tp_config(self, *, arch: MiniMaxH3DiTArchConfig, tp_size: int) -> None: @@ -1097,6 +1114,7 @@ class MiniMaxH3DiTModel(nn.Module): refiner_psp = _required_kwarg(kwargs, "refiner_packed_seq_params") refiner_cu = self._psp_field(refiner_psp, "refiner_packed_seq_params", "cu_seqlens_q").to(torch.int32) refiner_max = int(self._psp_field(refiner_psp, "refiner_packed_seq_params", "max_seqlen_q")) + video_layout = kwargs.get("video_token_layout") if x.dim() != 3 or x.shape[0] != 1: raise ValueError(f"x must be [1, S, C], got {list(x.shape)}") @@ -1145,6 +1163,7 @@ class MiniMaxH3DiTModel(nn.Module): cu_seqlens=cu_seqlens, max_seqlen=max_seqlen, packed_total=seq_len, + video_layout=video_layout, ) hidden = self.sp_gather(hidden) diff --git a/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py b/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py index 37353b459..a57eae700 100644 --- a/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py +++ b/vllm_omni/diffusion/models/minimax_h3/packed_sequence.py @@ -236,6 +236,9 @@ def minimax_h3_packed_sequence( "token_tags": token_tags, "cu_seqlens": cu, "document_id": doc, + # Extent of the video segment within the packed sequence. + "latent_grid": torch.tensor([latent_t, ph, pw], dtype=torch.int64), + "video_row_start": torch.tensor(video_sl.start, dtype=torch.int64), } @@ -554,6 +557,9 @@ def minimax_h3_packed_sequence_ref2va_blocks( "token_tags": token_tags, "cu_seqlens": cu, "document_id": doc, + # Extent of the video segment within the packed sequence. + "latent_grid": torch.tensor([latent_t, ph, pw], dtype=torch.int64), + "video_row_start": torch.tensor(video_sl.start, dtype=torch.int64), } diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index 48cc5b8ca..e28450c90 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -51,7 +51,14 @@ def _load_remote_component( class_reference, component_path, ) - return component_cls.from_pretrained(component_path) + # Build on the host regardless of the ambient default device. Online + # quantization wraps pipeline construction in a `with torch.device()` + # block for the DiT's quantized linears, and the checkpoint's own VAE code + # builds constants with ops that have no accelerator kernel (BigVGAN's + # anti-aliasing filters call torch.kaiser_window). Callers place the module + # explicitly right after this returns, so nothing depends on the context. + with torch.device("cpu"): + return component_cls.from_pretrained(component_path) class _AudioVAEDeterminismContext(AbstractContextManager): diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index 642d8ec02..4551126a7 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/vllm_omni/platforms/cuda/platform.py @@ -134,6 +134,7 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase): if selected_backend is not None: backend_upper = selected_backend.upper() + cls.validate_diffusion_attn_backend(backend_upper) if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): try: importlib.import_module("kernels") diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 99dafa2f2..fdd31ad21 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -127,6 +127,29 @@ class OmniPlatform(Platform): """ raise NotImplementedError + @classmethod + def validate_diffusion_attn_backend(cls, selected_backend: str) -> None: + """Reject an explicitly selected backend this platform cannot run. + + Platforms call this from ``get_diffusion_attn_backend_cls`` so that a + backend restricted to other hardware, or one whose kernel package is + missing, fails during resolution rather than at the first forward. + """ + from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum + + backend_upper = selected_backend.upper() + backend_cls = DiffusionAttentionBackendEnum[backend_upper].get_class() + + supported = backend_cls.supported_platforms + platform = cls._omni_enum.value + if supported is not None and platform not in supported: + raise ValueError( + f"The {backend_upper} diffusion attention backend runs on {', '.join(supported)} only, " + f"but the current platform is {platform}. Select a backend supported here, " + "such as FLASH_ATTN or TORCH_SDPA." + ) + backend_cls.validate_available() + @classmethod def supports_torch_inductor(cls) -> bool: """Check if the platform supports torch.compile with inductor backend.""" diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index 0774bca3a..02ce155ee 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/vllm_omni/platforms/musa/platform.py @@ -80,6 +80,7 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): if selected_backend is not None: backend_upper = selected_backend.upper() + cls.validate_diffusion_attn_backend(backend_upper) if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): logger.warning( "HuggingFace kernels-backed FlashAttention is " diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index bb0e1b436..0421814a0 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -130,6 +130,7 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform): if selected_backend is not None: backend_upper = selected_backend.upper() + cls.validate_diffusion_attn_backend(backend_upper) if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): logger.warning( "HuggingFace kernels-backed FlashAttention is " diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index 6ae65ed2b..630777058 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/vllm_omni/platforms/rocm/platform.py @@ -102,6 +102,7 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform): if selected_backend is not None: backend_upper = selected_backend.upper() + cls.validate_diffusion_attn_backend(backend_upper) if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): logger.warning( "HuggingFace kernels-backed FlashAttention is " diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index aa26cd82c..22cebeead 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/vllm_omni/platforms/xpu/platform.py @@ -53,6 +53,7 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform): if selected_backend is not None: backend_upper = selected_backend.upper() + cls.validate_diffusion_attn_backend(backend_upper) if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): logger.warning( "HuggingFace kernels-backed FlashAttention is " diff --git a/vllm_omni/quantization/int8_config.py b/vllm_omni/quantization/int8_config.py index 37d430047..ea096f9d5 100644 --- a/vllm_omni/quantization/int8_config.py +++ b/vllm_omni/quantization/int8_config.py @@ -50,9 +50,55 @@ if TYPE_CHECKING: # Dynamic quantization is supported first. ACTIVATION_SCHEMES = ["dynamic"] +# Ascend's npu_quant_matmul (QuantBatchMatmulV3) refuses a weight whose last +# dimension exceeds this, and it fails at the first forward rather than at load +# time. Layers still wider than the limit after TP sharding stay unquantized. +NPU_QUANT_MATMUL_MAX_OUT_FEATURES = 65535 + logger = init_logger(__name__) +def _fell_back_to_unquantized_npu( + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, +) -> bool: + """Swap a layer to unquantized weights when npu_quant_matmul cannot run its shape. + + The check uses the per-partition output size because that is what the kernel + actually sees; a layer over the limit on one rank can be within it at a + higher TP degree. Returns True when the layer was swapped and its weights + created, in which case the caller must not create its own. + """ + output_size_per_partition = sum(output_partition_sizes) + if output_size_per_partition <= NPU_QUANT_MATMUL_MAX_OUT_FEATURES: + return False + + logger.warning_once( + "Keeping a %d-wide linear unquantized: npu_quant_matmul rejects an output dimension past " + "%d. Tensor parallelism shrinks this per-rank dimension, so a higher TP degree brings such " + "layers back into range.", + output_size_per_partition, + NPU_QUANT_MATMUL_MAX_OUT_FEATURES, + ) + fallback = UnquantizedLinearMethod() + layer.quant_method = fallback + fallback.create_weights( + layer, + input_size_per_partition, + output_partition_sizes, + input_size, + output_size, + params_dtype, + **extra_weight_attrs, + ) + return True + + def create_weight_parameter( output_size_per_partition: int, input_size_per_partition: int, @@ -231,6 +277,21 @@ class LazyWeightMixin: uses_meta_device: bool = True + # This mixin knows when a layer's weight is final, so it can hand the layer + # back to the host right there. Loaders that intend to offload the whole + # model after loading opt in per layer via ``enable_offload_after_quant``. + supports_offload_after_quant: bool = True + _offload_after_quant: bool = False + + def enable_offload_after_quant(self) -> None: + """Return each layer to host memory as soon as it has been quantized. + + Caps the load-time device footprint at one layer instead of the whole + model. A quant method instance belongs to a single layer, so this is not + a global switch. + """ + self._offload_after_quant = True + def create_weights( self, layer: torch.nn.Module, @@ -287,6 +348,11 @@ class LazyWeightMixin: # anything layer._already_called_process_weights_after_loading = True + # This layer's weight is final, so nothing needs it on the + # accelerator until inference. + if self._offload_after_quant: + layer.to("cpu") + # Note that we keep `layer._loaded_numel` around just in case # there is logic added to vllm in the future which calls a # weight loader twice - we do not want to re-initialize in @@ -354,6 +420,11 @@ class NPUInt8LinearMethod(BaseInt8LinearMethod): def __init__(self, quant_config: DiffusionInt8Config): super().__init__(quant_config) + def create_weights(self, layer: torch.nn.Module, *args, **kwargs) -> None: + if _fell_back_to_unquantized_npu(layer, *args, **kwargs): + return + super().create_weights(layer, *args, **kwargs) + def process_weights_after_loading(self, layer: Module) -> None: layer.weight.data = layer.weight.data.t().contiguous() layer.weight_scale.data = layer.weight_scale.data.squeeze() @@ -424,6 +495,13 @@ class NPUInt8OnlineLinearMethod(LazyWeightMixin, NPUInt8LinearMethod): and quantized the weights during loading. """ + def create_weights(self, layer: torch.nn.Module, *args, **kwargs) -> None: + # NPUInt8LinearMethod's override is unreachable from here: LazyWeightMixin + # comes first in the MRO and does not call super(). + if _fell_back_to_unquantized_npu(layer, *args, **kwargs): + return + super().create_weights(layer, *args, **kwargs) + def process_weights_after_loading(self, layer: Module) -> None: if getattr(layer, "_already_called_process_weights_after_loading", False): return