[Feature][MiniMax-H3] Complete VSA and Ulysses support (#6909)

Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>

docs/configuration/environment_variables.md

@@ -179,7 +179,7 @@ collection, examples, bug reports, or logs.
 
 ## Model-specific variables
 
-The audit found 56 variables read by a single model or pipeline family. They are
+The audit found 58 variables read by a single model or pipeline family. They are
 not listed as public usage options here because doing so would turn implementation
 escape hatches into an accidental compatibility contract.
 
@@ -188,7 +188,7 @@ Every audited model-specific name has a migration disposition in the
 
 | Disposition | Count | Required outcome |
 | --- | ---: | --- |
-| Promote | 33 | Move a stable setting into typed stage or model configuration. |
+| Promote | 35 | Move a stable setting into typed stage or model configuration. |
 | Request scope | 6 | Move request-varying behavior into a declared request-option schema. |
 | External | 0 | Retain only when a supported third-party library owns the contract. |
 | Internalize | 11 | Keep a debug or diagnostic switch out of public documentation and configuration. |

docs/user_guide/diffusion/attention_backends.md

@@ -20,7 +20,7 @@ For the internal selector, registry, and platform contract, see
 | Install and use SageAttention 2.2 or SageAttention3 | [SageAttention](attention_backends/sage.md) |
 | Match training or rollout kernels loaded from Hugging Face | [Hugging Face Hub Backends](attention_backends/huggingface_hub.md) |
 | Use block-sparse video attention on Ascend NPU | [RainFusion](attention_backends/rainfusion.md) |
-| Use FastVideo VSA with FastWan2.2-TI2V-5B on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
+| Use FastVideo VSA with FastWan2.2-TI2V-5B or FastH3 MiniMax-H3 on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.md) |
 
 ## Backend options
 
@@ -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 on CUDA | [FastVideo VSA](attention_backends/fastvideo_vsa.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) |
 
 ## Configuration
 

docs/user_guide/diffusion/attention_backends/fastvideo_vsa.md

@@ -7,11 +7,18 @@ attention only against the selected top-k blocks.
 
 The supported checkpoint provides both text-to-video and image-to-video modes through `Wan22Pipeline`; the separate Wan I2V-14B, S2V, and VACE pipelines are outside this backend's supported scope.
 
+MiniMax-H3 served with a FastH3 VSA adapter uses a second, H3-specific route
+through the same backend. See the
+[MiniMax-H3 recipe](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) for its geometry,
+supported topologies, and commands.
+
 VSA is a CUDA-only, explicitly selected backend. It requires the
 `fastvideo-kernel` package and currently supports non-causal self-attention
 with equal query and key/value sequence lengths. Unsupported shapes, masks,
 dtypes, sequence-parallel execution, or kernel failures fall back to
-`TORCH_SDPA` and emit a warning with the reason.
+`TORCH_SDPA` and emit a warning with the reason. On the Wan route, an active
+sequence-parallel context is one of those fallbacks; the H3 route supports pure
+Ulysses and rejects ring or all-gather sequence parallelism at startup.
 
 ## Enable the backend
 
@@ -119,7 +126,8 @@ the backend guarantees sparse execution:
 - `route=SDPA` or `FASTVIDEO_VSA falling back to SDPA: ...` means dense SDPA
   executed; the warning includes the reason.
 
-The current kernel requires CUDA tensors in FP16 or BF16, 256-token blocks,
+The Wan route requires CUDA tensors in FP16 or BF16, 256-token blocks,
 standard `head_size**-0.5` scaling, equal Q/K/V head counts, no attention mask,
-and no active sequence-parallel context. NPU and XPU paths do not execute the
-FastVideo VSA CUDA kernel.
+and no active sequence-parallel context. The MiniMax-H3 route uses 64-token
+`(4, 4, 4)` blocks and supports pure Ulysses. NPU and XPU paths do not execute
+the FastVideo VSA CUDA kernel.

recipes/MiniMaxAI/MiniMax-H3.md

@@ -962,13 +962,27 @@ mostly base H3 weights on a four-step schedule. Offload is refused for the same
 reason - `--enable-cpu-offload`, `--enable-layerwise-offload` and
 `--enable-distributed-layerwise-offload` all bypass the fusion, so they fail fast.
 
-> [!NOTE]
-> Only the **Dense / Data-Free** variant is supported today. The three VSA
-> variants carry `.set_weight` compression gates for modules that do not exist
-> in vLLM-Omni's dense H3 attention, and the release states dense attention is
-> not a supported substitute for them, so they are refused at startup rather
-> than silently run wrong. Video Sparse Attention for H3's packed
-> `[text | cond | audio | video]` sequence is follow-up work.
+The VSA variants are supported through FastVideo's external kernel. Install a
+`fastvideo-kernel` build that provides the `fastvideo_kernel` Python module,
+then add the following flags to the same command:
+
+```bash
+--diffusion-attention-backend FASTVIDEO_VSA \
+--fastvideo-vsa-topk 64
+```
+
+FastH3 VSA applies its learned `.set_weight` compression gates to the complete
+packed `[text | cond | audio | video]` document using the official H3 geometry:
+text/condition/audio prefix tiles never cross segment boundaries, and target
+video rows use `(4, 4, 4)` 3-D tiles (64 tokens). Prefix queries remain dense;
+video queries select all prefix tiles plus the configured top-k video tiles.
+Pure Ulysses sequence parallelism is supported: the learned gate follows the
+same sequence-to-head all-to-all as Q/K/V before VSA runs. Ring and all-gather
+SP remain unsupported because they do not present a complete packed sequence
+to each block-sparse kernel rank.
+
+The Dense / Data-Free variant does not require `fastvideo-kernel` and should be
+served with a dense attention backend.
 
 Measured on 8x NVIDIA B300 with USP8, VAE patch-parallel 8, `TRTLLM_ATTN`, at
 1344x768, 4.4 s, seed 1101, one warmup excluded and two runs recorded:

tests/config/test_environment_variables.py

@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PUBLIC_OMNI: 24,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
-        EnvironmentVariableCategory.MODEL_SPECIFIC: 57,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 58,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
         EnvironmentVariableCategory.INTERNAL: 2,
     }
@@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC
     )
     assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
-        ModelEnvironmentVariableDisposition.PROMOTE: 34,
+        ModelEnvironmentVariableDisposition.PROMOTE: 35,
         ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
         ModelEnvironmentVariableDisposition.INTERNALIZE: 11,

tests/diffusion/attention/test_fastvideo_vsa.py

@@ -1,6 +1,7 @@
 # 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
@@ -8,9 +9,17 @@ from typing import Any
 import pytest
 import torch
 
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+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,
@@ -23,6 +32,12 @@ def test_fastvideo_vsa_backend_is_registered():
     assert DiffusionAttentionBackendEnum.FASTVIDEO_VSA.get_path().endswith("fastvideo_vsa.FastVideoVSABackend")
 
 
+def test_fastvideo_vsa_reports_missing_optional_kernel(monkeypatch):
+    monkeypatch.setattr("importlib.util.find_spec", lambda name: None)
+    with pytest.raises(ImportError, match="fastvideo-kernel"):
+        FastVideoVSABackend.validate_available()
+
+
 def test_fastvideo_vsa_tiles_3d_sequence_and_untiles(monkeypatch):
     calls: dict[str, Any] = {}
     fake_module = types.ModuleType("fastvideo_kernel")
@@ -201,3 +216,190 @@ 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()

tests/diffusion/attention/test_ulysses_uaa.py

@@ -322,6 +322,50 @@ def test_advanced_uaa_hybrid_rejects_non_gqa_shapes(monkeypatch: pytest.MonkeyPa
         strategy.pre_attention(query, key, value, None)
 
 
+def test_strict_ulysses_reshards_vsa_gate_with_qkv(monkeypatch) -> None:
+    from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.attention.parallel.ulysses.get_ulysses_mode",
+        lambda default: "strict",
+    )
+    calls: list[torch.Tensor] = []
+
+    def fake_all_to_all(pg, tensor, scatter_idx, gather_idx, use_sync):
+        del pg, scatter_idx, gather_idx, use_sync
+        calls.append(tensor)
+        return tensor.reshape(tensor.shape[0], tensor.shape[1] * 2, tensor.shape[2] // 2, tensor.shape[3])
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.attention.parallel.ulysses.SeqAllToAll4D.apply",
+        fake_all_to_all,
+    )
+    strategy = UlyssesParallelAttention(
+        sp_group=SimpleNamespace(
+            ulysses_group=object(),
+            ulysses_world_size=2,
+            ulysses_rank=0,
+            ring_world_size=1,
+        ),
+        scatter_idx=2,
+        gather_idx=1,
+        use_sync=False,
+    )
+    query = torch.arange(24).reshape(1, 3, 4, 2)
+    key = query + 10
+    value = query + 20
+    gate = query + 100
+    metadata = AttentionMetadata(extra={"gate_compress": gate})
+
+    query_out, _, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata)
+
+    assert all(actual is expected for actual, expected in zip(calls, (query, key, value, gate), strict=True))
+    assert query_out.shape == (1, 6, 2, 2)
+    assert metadata_out is metadata
+    assert metadata.extra["gate_compress"].shape == query_out.shape
+    torch.testing.assert_close(metadata.extra["gate_compress"], gate.reshape_as(query_out))
+
+
 def _run_attention_case(
     local_rank: int,
     world_size: int,

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -795,6 +795,10 @@ def test_cudnn_packed_attention_uses_python_length_without_padding_mask():
 
 
 def test_packed_attention_skips_mask_for_packed_mask_free_backend():
+    from vllm_omni.diffusion.attention.backends.abstract import (
+        VideoTokenLayout,
+        VideoTokenSpan,
+    )
     from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
         MiniMaxH3Attention,
     )
@@ -821,6 +825,10 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
     torch.nn.Module.__init__(attention)
     attention.attention = FakeAttention()
     q = torch.randn(8, 2, 4)
+    video_layout = VideoTokenLayout(
+        used_len=5,
+        video_spans=(VideoTokenSpan(start=3, latent_grid=(1, 1, 2), role="target"),),
+    )
 
     attention._run_packed_attention(
         q,
@@ -829,6 +837,9 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
         cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
         max_seqlen=5,
         packed_total=8,
+        video_layout=video_layout,
+        vsa_prefix_segments=(1, 2),
+        gate_compress=torch.ones_like(q),
     )
 
     metadata = attention.attention.metadata
@@ -836,6 +847,11 @@ def test_packed_attention_skips_mask_for_packed_mask_free_backend():
     assert metadata.attn_mask is None
     assert metadata.extra["valid_kv_length"] == 5
     assert metadata.extra["npu_attn_varlen"] is True
+    assert metadata.video_layout is video_layout
+    assert metadata.extra["vsa_h3_prefix_segments"] == (1, 2)
+    assert "vsa_h3_video_shape" not in metadata.extra
+    assert "vsa_h3_target_start" not in metadata.extra
+    assert metadata.extra["gate_compress"].shape == (1, *q.shape)
     packed_padding = metadata.packed_padding
     assert packed_padding is not None
     assert packed_padding.q_length == 5

tests/diffusion/models/minimax_h3/test_minimax_h3_fasth3.py

@@ -437,6 +437,83 @@ def test_adopting_the_contract_refuses_a_vsa_variant_and_ref2va(tmp_path):
         _fusion(tmp_path).check_serving_contract(partition="ref2va", **contract)
 
 
+def test_a_vsa_variant_accepts_pure_ulysses_and_refuses_ring(tmp_path):
+    sparse = tmp_path / "vsa" / "adapter_model.safetensors"
+    _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
+    od_config = SimpleNamespace(
+        diffusion_attention_backend="FASTVIDEO_VSA",
+        parallel_config=SimpleNamespace(sequence_parallel_size=8, ulysses_degree=8, ring_degree=1),
+    )
+    fusion = _load(sparse.parent)
+    fusion.check_serving_contract(partition="fl2va", od_config=od_config, video_shift=12.0, audio_shift=3.0)
+
+    od_config.parallel_config.ring_degree = 2
+    with pytest.raises(ValueError, match="pure Ulysses"):
+        fusion.check_serving_contract(partition="fl2va", od_config=od_config, video_shift=12.0, audio_shift=3.0)
+
+
+def test_a_vsa_variant_accepts_normalized_attention_config(tmp_path):
+    sparse = tmp_path / "vsa" / "adapter_model.safetensors"
+    _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
+    od_config = SimpleNamespace(
+        diffusion_attention_config=SimpleNamespace(default=SimpleNamespace(backend="FASTVIDEO_VSA")),
+        parallel_config=SimpleNamespace(sequence_parallel_size=1),
+    )
+    _load(sparse.parent).check_serving_contract(
+        partition="fl2va", od_config=od_config, video_shift=12.0, audio_shift=3.0
+    )
+
+
+def test_a_vsa_variant_reads_the_backend_the_dit_will_actually_resolve(tmp_path):
+    # The 50 DiT blocks carry attention role "self", so a per_role entry - not
+    # the default - decides whether the compression gates ever reach a sparse
+    # kernel.
+    sparse = tmp_path / "vsa" / "adapter_model.safetensors"
+    _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
+    fusion = _load(sparse.parent)
+    contract = {"partition": "fl2va", "video_shift": 12.0, "audio_shift": 3.0}
+
+    dense_dit = SimpleNamespace(
+        diffusion_attention_config=SimpleNamespace(
+            default=SimpleNamespace(backend="FASTVIDEO_VSA"),
+            per_role={"self": SimpleNamespace(backend="TRTLLM_ATTN")},
+        ),
+        parallel_config=SimpleNamespace(sequence_parallel_size=1),
+    )
+    with pytest.raises(ValueError, match="TRTLLM_ATTN"):
+        fusion.check_serving_contract(od_config=dense_dit, **contract)
+
+    sparse_dit = SimpleNamespace(
+        diffusion_attention_config=SimpleNamespace(
+            default=None,
+            per_role={"self": SimpleNamespace(backend="FASTVIDEO_VSA")},
+        ),
+        parallel_config=SimpleNamespace(sequence_parallel_size=1),
+    )
+    fusion.check_serving_contract(od_config=sparse_dit, **contract)
+
+
+def test_a_gate_that_never_reached_the_model_is_refused(tmp_path):
+    # load_weights only logs a skip for a parameter the model does not have, so
+    # yielding the gate is not evidence it arrived.
+    sparse = tmp_path / "vsa" / "adapter_model.safetensors"
+    _write_adapter(sparse, tensors={"transformer_blocks.0.attn.to_gate_compress.set_weight": torch.ones((2, 2))})
+    fusion = _load(sparse.parent)
+    checkpoint = [
+        ("blocks.0.attn.qkv_proj.weight", torch.zeros((3 * _INNER, _HIDDEN))),
+        ("blocks.0.attn.out_proj.weight", torch.zeros((_HIDDEN, _INNER))),
+        ("blocks.0.mlp.fc1.weight", torch.zeros((2 * _FFN, _HIDDEN))),
+        ("blocks.0.mlp.fc2.weight", torch.zeros((_HIDDEN, _FFN))),
+    ]
+    gate = "blocks.0.attn.to_gate_compress.weight"
+    streamed = {name for name, _ in fusion.apply(checkpoint)}
+    assert gate in streamed
+
+    with pytest.raises(FastH3AdapterError, match="never reached the model"):
+        fusion.validate_fully_applied(streamed - {gate})
+    fusion.validate_fully_applied(streamed)
+
+
 def test_a_full_rank_delta_on_a_fused_parameter_is_refused(tmp_path):
     # Only low-rank factors are placed into H3's fused QKV and gate/up layouts,
     # so a .diff aimed at one would otherwise be added transposed.

tests/diffusion/models/minimax_h3/test_minimax_h3_parallel.py

@@ -67,6 +67,33 @@ def test_packed_attention_is_a_regional_compile_boundary():
     assert getattr(MiniMaxH3Attention._run_packed_attention, "_torchdynamo_disable", False)
 
 
+def test_denoise_branch_keeps_fast_h3_prefix_geometry_model_local():
+    from vllm_omni.diffusion.models.minimax_h3.denoise_loop import (
+        MiniMaxH3DenoiseBranch,
+    )
+
+    packed = {
+        "seq_len": torch.tensor(9),
+        "img_pos": torch.arange(4),
+        "audio_pos": torch.empty(0, dtype=torch.long),
+        "text_pos": torch.arange(2),
+        "update_mask": torch.ones(4, dtype=torch.bool),
+        "cu_seqlens": torch.tensor([0, 9, 9], dtype=torch.int32),
+        "img_position_ids": torch.zeros(9, 3, dtype=torch.long),
+        "video_spans": ({"start": 5, "latent_grid": (1, 2, 2), "role": "target"},),
+    }
+    branch = MiniMaxH3DenoiseBranch(
+        packed=packed,
+        text_embeddings=torch.zeros(2, 4),
+        token_tags=torch.tensor([0, 0, 1, 1, 1, 2, 2, 2, 2]),
+        device=torch.device("cpu"),
+    )
+
+    layout = branch.static_kwargs["video_token_layout"]
+    assert not hasattr(layout, "prefix_segments")
+    assert branch.static_kwargs["packed_seq_params"]["vsa_prefix_segments"] == (2, 3)
+
+
 def test_h3_fused_rope_matches_reference_and_preserves_unrotated_dims():
     from vllm_omni.diffusion.layers.rope import RotaryEmbedding
     from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (

vllm_omni/config/environment_variable_inventory.py

@@ -147,6 +147,7 @@ _MODEL_PROMOTE = (
     "DYNIN_REMOTE_CODE_LOCAL_FILES_ONLY",
     "DYNIN_REMOTE_CODE_REPO_ID",
     "DYNIN_REMOTE_CODE_REVISION",
+    "FASTVIDEO_VSA_SM100A",
     "HIGGS_AUDIO_TOKENIZER_PATH",
     "INTERNVLA_A1_COSMOS_DECODER_PATH",
     "INTERNVLA_A1_COSMOS_DIR",

vllm_omni/diffusion/attention/backends/fastvideo_vsa.py

@@ -4,7 +4,9 @@
 from __future__ import annotations
 
 import functools
+import importlib.util
 import math
+import os
 from collections.abc import Mapping
 from typing import Any
 
@@ -18,6 +20,7 @@ from vllm_omni.diffusion.attention.backends.abstract import (
 )
 from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
@@ -85,6 +88,77 @@ if not hasattr(torch.ops.vllm_omni, "fastvideo_vsa_bshd"):
 _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],
@@ -152,6 +226,81 @@ 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
@@ -179,6 +328,24 @@ def _preserve_vsa_all_blocks(attn_metadata: AttentionMetadata | None) -> bool:
 class FastVideoVSABackend(AttentionBackend):
     accept_output_buffer: bool = True
 
+    @classmethod
+    def supports_packed_mask_free(cls) -> bool:
+        # FastVideo accepts variable-sized edge blocks. This lets packed
+        # [real, pad] inputs run on their valid prefix without materializing an
+        # attention mask; the implementation restores the ignored pad rows.
+        # Only forward_cuda honours packed_padding: every other platform hands
+        # the tensors straight to SDPA, which reads attn_mask and nothing else,
+        # so the pad rows would be attended as real keys.
+        return current_omni_platform.is_cuda()
+
+    @classmethod
+    def validate_available(cls) -> None:
+        if importlib.util.find_spec("fastvideo_kernel") is None:
+            raise ImportError(
+                "FASTVIDEO_VSA requires the optional fastvideo-kernel package "
+                "(the installation must provide the 'fastvideo_kernel' Python module)."
+            )
+
     @staticmethod
     def get_supported_head_sizes() -> list[int]:
         # FastVideo VSA is intended for video DiT head sizes such as 64/128.
@@ -257,8 +424,31 @@ class FastVideoVSAImpl(AttentionImpl):
         attn_metadata: AttentionMetadata | None,
         reason: str,
     ) -> torch.Tensor:
+        """Run dense SDPA, honouring the mask-free packed contract this backend claims.
+
+        ``supports_packed_mask_free`` tells the model it may skip building the
+        padding mask, so on this path there is nothing to stop SDPA attending
+        the structural pad rows as real keys. Slice to the valid prefix instead
+        and leave the pad rows zeroed, exactly as the VSA path does.
+        """
         logger.warning_once("FASTVIDEO_VSA falling back to SDPA: %s", reason)
-        return self.sdpa_fallback.forward(query, key, value, attn_metadata)
+        packed = attn_metadata.packed_padding if attn_metadata is not None else None
+        if attn_metadata is None or packed is None or attn_metadata.attn_mask is not None:
+            return self.sdpa_fallback.forward(query, key, value, attn_metadata)
+        q_length = min(int(packed.q_length), query.shape[1])
+        kv_length = min(int(packed.kv_length), key.shape[1])
+        output = self.sdpa_fallback.forward(
+            query[:, :q_length], key[:, :kv_length], value[:, :kv_length], attn_metadata
+        )
+        if q_length == query.shape[1]:
+            return output
+        restored = torch.zeros(
+            (output.shape[0], query.shape[1], output.shape[2], output.shape[3]),
+            device=output.device,
+            dtype=output.dtype,
+        )
+        restored[:, :q_length] = output
+        return restored
 
     def _fallback_reason(
         self,
@@ -314,6 +504,93 @@ class FastVideoVSAImpl(AttentionImpl):
                 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,
@@ -321,9 +598,42 @@ class FastVideoVSAImpl(AttentionImpl):
         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]
+
+        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(query, key, value, attn_metadata, reason)
+            return self._fallback(original_query, original_key, original_value, attn_metadata, reason)
 
         seq_len = query.shape[1]
         dit_seq_shape = _get_vsa_dit_seq_shape(attn_metadata)
@@ -379,6 +689,8 @@ class FastVideoVSAImpl(AttentionImpl):
             gate_compress = _get_gate_compress(attn_metadata)
             if gate_compress is None:
                 gate_compress = torch.zeros_like(query)
+            elif valid_seq_len != original_seq_len:
+                gate_compress = gate_compress[:, :valid_seq_len]
             elif gate_compress.shape != query.shape:
                 raise ValueError(f"gate_compress shape {gate_compress.shape} must match query shape {query.shape}")
             gate_tiled = torch.zeros_like(query_tiled)
@@ -397,11 +709,22 @@ class FastVideoVSAImpl(AttentionImpl):
                 self.block_size[1],
                 self.block_size[2],
             )
-            return output[:, untile_combined_index].contiguous()
+            output = output[:, untile_combined_index].contiguous()
+            if valid_seq_len == original_seq_len:
+                return output
+            restored = torch.zeros(
+                (query.shape[0], original_seq_len, query.shape[2], query.shape[3]),
+                device=query.device,
+                dtype=query.dtype,
+            )
+            restored[:, :valid_seq_len] = output
+            return restored
         except Exception as exc:
             if not self.fallback_on_error:
                 raise
-            return self._fallback(query, key, value, attn_metadata, f"VSA kernel failed: {exc}")
+            return self._fallback(
+                original_query, original_key, original_value, attn_metadata, f"VSA kernel failed: {exc}"
+            )
 
     def forward_npu(
         self,

vllm_omni/diffusion/attention/parallel/ulysses.py

@@ -277,6 +277,15 @@ class UlyssesParallelAttention:
     ):
         mode = get_ulysses_mode(default="strict")
         ulysses_world_size = self._sp_group.ulysses_world_size
+        gate_compress = None
+        if attn_metadata is not None:
+            candidate = attn_metadata.extra.get("gate_compress")
+            if isinstance(candidate, torch.Tensor):
+                if candidate.shape != query.shape:
+                    raise ValueError(
+                        f"gate_compress must match pre-Ulysses query shape, got {candidate.shape} vs {query.shape}"
+                    )
+                gate_compress = candidate
 
         # advanced_uaa pads non-divisible head counts before the Ulysses
         # all-to-all. This also holds in hybrid Ulysses+Ring: Q is derived
@@ -471,6 +480,16 @@ class UlyssesParallelAttention:
                 use_sync=self._use_sync,
                 padded_head_cnt=padded_kv_head_cnt,
             )
+            if gate_compress is not None:
+                # gate_compress mirrors the query layout, so it must follow the
+                # same GQA-derived head padding as the query all-to-all.
+                gate_compress, _ = _ulysses_all_to_all_any_qkv(
+                    self._ulysses_pg,
+                    gate_compress,
+                    seq_lens=seq_lens,
+                    use_sync=self._use_sync,
+                    padded_head_cnt=padded_query_head_cnt,
+                )
         else:
             # Strict mode: fail fast with actionable errors for head divisibility.
             for name, t in (("query", query), ("key", key), ("value", value)):
@@ -496,6 +515,8 @@ class UlyssesParallelAttention:
                 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)
+                if gate_compress is not None:
+                    gate_compress = ulysses_qkv_fwd(gate_compress, group_name, ulysses_world_size)
             else:
                 query = SeqAllToAll4D.apply(
                     self._ulysses_pg, query, self._scatter_idx, self._gather_idx, self._use_sync
@@ -504,6 +525,10 @@ class UlyssesParallelAttention:
                 value = SeqAllToAll4D.apply(
                     self._ulysses_pg, value, self._scatter_idx, self._gather_idx, self._use_sync
                 )
+                if gate_compress is not None:
+                    gate_compress = SeqAllToAll4D.apply(
+                        self._ulysses_pg, gate_compress, self._scatter_idx, self._gather_idx, self._use_sync
+                    )
             seq_lens = []
             local_seq_len = 0
             orig_head_cnt = 0
@@ -531,6 +556,12 @@ class UlyssesParallelAttention:
                 key = torch.cat([key, joint_tensor_key], dim=1)
                 value = torch.cat([value, joint_tensor_value], dim=1)
 
+        if gate_compress is not None and attn_metadata is not None:
+            # The VSA compression gate follows the same S<->H all-to-all as
+            # Q/K/V so every rank sees the full packed sequence for its local
+            # head shard, matching FastVideo's DistributedAttention_VSA.
+            attn_metadata.extra["gate_compress"] = gate_compress
+
         ctx = _UlyssesCtx(
             name=self.name,
             ulysses_pg=self._ulysses_pg,
@@ -558,6 +589,7 @@ class UlyssesParallelAttention:
                     attn_metadata.attn_mask = None
                 else:
                     if attn_metadata.attn_mask is None:
+                        assert attn_metadata.joint_attn_mask is not None
                         attn_metadata.attn_mask = torch.ones(
                             [query.shape[0], query.shape[1] - attn_metadata.joint_attn_mask.shape[1]],
                             dtype=torch.bool,

vllm_omni/diffusion/models/minimax_h3/denoise_loop.py

@@ -136,7 +136,23 @@ class MiniMaxH3DenoiseBranch:
                 )
                 for span in raw_spans
             )
-            self.static_kwargs["video_token_layout"] = VideoTokenLayout(used_len=int(cu[1]), video_spans=spans)
+            target_start = next(span.start for span in reversed(spans) if span.role == "target")
+            prefix_tags = token_tags[:target_start]
+            prefix_segments: list[int] = []
+            if prefix_tags.numel():
+                # Run-lengths preserve modality/reference boundaries without
+                # carrying CUDA tensors into every attention layer.
+                boundaries = torch.nonzero(prefix_tags[1:] != prefix_tags[:-1]).flatten().tolist()
+                starts = [0, *(index + 1 for index in boundaries)]
+                stops = [*(index + 1 for index in boundaries), target_start]
+                prefix_segments = [stop - start for start, stop in zip(starts, stops, strict=True)]
+            self.static_kwargs["video_token_layout"] = VideoTokenLayout(
+                used_len=int(cu[1]),
+                video_spans=spans,
+            )
+            # FastH3-specific prefix geometry belongs to MiniMax-H3's packed
+            # attention contract, not the shared VideoTokenLayout interface.
+            self.static_kwargs["packed_seq_params"]["vsa_prefix_segments"] = tuple(prefix_segments)
         else:
             grid = packed["latent_grid"].tolist()
             self.static_kwargs["video_token_layout"] = VideoTokenLayout(

vllm_omni/diffusion/models/minimax_h3/fasth3.py

@@ -114,6 +114,7 @@ _BLOCK_TARGETS = {
     "attn.to_k": ("attn.qkv_proj", _QKV_K),
     "attn.to_v": ("attn.qkv_proj", _QKV_V),
     "attn.to_out.0": ("attn.out_proj", _PLAIN),
+    "attn.to_gate_compress": ("attn.to_gate_compress", _PLAIN),
     "ff.net.0.proj": ("mlp.fc1", _SWAP_HALVES),
     "ff.net.2": ("mlp.fc2", _PLAIN),
     "adaln_proj.linear": ("adaln_proj.linear", _PLAIN),
@@ -121,6 +122,11 @@ _BLOCK_TARGETS = {
     "norm2": ("norm2", _PLAIN),
 }
 
+# The attention role MiniMaxH3Attention gives its 50 DiT blocks. The
+# compression gates live on exactly these layers, so this is the role whose
+# resolved backend decides whether the artifact runs sparse.
+_H3_DIT_ATTENTION_ROLE = "self"
+
 # Adapter block prefix -> native block prefix.
 _BLOCK_PREFIXES = (
     ("token_refiner.refiner_blocks.", "token_refiner.blocks."),
@@ -128,6 +134,26 @@ _BLOCK_PREFIXES = (
 )
 
 
+def _resolve_dit_attention_backend(od_config: Any) -> str:
+    """The backend the 50-block H3 DiT will actually resolve to.
+
+    The DiT's attention layers carry role ``"self"``, so a ``per_role`` entry
+    overrides the default for exactly the layers the compression gates live on.
+    Reading only the default would accept a config that runs the sparse student
+    dense, and reject a per-role-only config that is correct.
+    """
+    attention_config = getattr(od_config, "diffusion_attention_config", None)
+    per_role = getattr(attention_config, "per_role", None) or {}
+    spec = per_role.get(_H3_DIT_ATTENTION_ROLE)
+    if spec is not None:
+        return str(getattr(spec, "backend", "") or "").upper()
+    backend = str(getattr(od_config, "diffusion_attention_backend", "") or "").upper()
+    if backend:
+        return backend
+    default_spec = getattr(attention_config, "default", None)
+    return str(getattr(default_spec, "backend", "") or "").upper()
+
+
 class FastH3AdapterError(ValueError):
     """The artifact is a FastH3 adapter, but it cannot be applied as one."""
 
@@ -305,11 +331,16 @@ class FastH3WeightFusion:
         patches: Mapping[str, _ParamPatch],
         head_dim: int,
         requires_vsa: bool,
+        injections: Mapping[str, torch.Tensor] | None = None,
     ) -> None:
         self._source = source
         self._patches = dict(patches)
         self._head_dim = head_dim
         self.requires_vsa = requires_vsa
+        # Parameters the base checkpoint does not carry, assigned into the
+        # stream instead of fused onto an existing weight.
+        self._injections = dict(injections or {})
+        self._injected: set[str] = set()
         self._applied: set[str] = set()
         self._device: torch.device | None = None
 
@@ -352,6 +383,7 @@ class FastH3WeightFusion:
 
         patches: dict[str, _ParamPatch] = {}
         gate_tensors: list[str] = []
+        injections: dict[str, torch.Tensor] = {}
         unmapped: list[str] = []
         blocks_seen: dict[str, set[int]] = {}
         counted = {"low_rank_tensors": 0, "diff_tensors": 0}
@@ -365,12 +397,20 @@ class FastH3WeightFusion:
                     unmapped.append(name)
                     continue
                 module, role = split
-                if role == "set_weight":
-                    # A VSA compression gate: a module the base transformer does
-                    # not have, so there is nothing to fuse it into.
-                    gate_tensors.append(name)
-                    continue
                 target = _resolve_native_target(module)
+                if role == "set_weight":
+                    # A VSA compression gate. The base transformer has no such
+                    # parameter, so this is assigned into the stream rather than
+                    # fused onto an existing weight.
+                    if target is None:
+                        unmapped.append(name)
+                        continue
+                    native_module, _, block = target
+                    if block is not None:
+                        blocks_seen.setdefault(block[0], set()).add(block[1])
+                    gate_tensors.append(name)
+                    injections[f"{native_module}.weight"] = checkpoint.get_tensor(name)
+                    continue
                 if target is None:
                     unmapped.append(name)
                     continue
@@ -422,6 +462,7 @@ class FastH3WeightFusion:
             patches=patches,
             head_dim=head_dim,
             requires_vsa=bool(gate_tensors),
+            injections=injections,
         )
         logger.info(
             "FastH3 adapter %s: rank=%s, parameters patched=%d, low-rank=%s, diff=%s, set_weight=%d",
@@ -510,24 +551,48 @@ class FastH3WeightFusion:
             # server would serve base H3 weights on the student's ladder.
             raise FastH3AdapterError(f"{self._source} has already been fused into this checkpoint")
         for name, weight in weights:
+            if name in self._injections:
+                raise FastH3AdapterError(
+                    f"the checkpoint already provides {name}, which this adapter assigns; "
+                    "assigning it would discard the checkpoint's own weight"
+                )
             yield name, self.fuse(name, weight)
+        # The VSA gates have no counterpart in the base checkpoint, so they join
+        # the stream after it rather than being folded into one of its tensors.
+        for name, weight in self._injections.items():
+            self._injected.add(name)
+            yield name, weight
 
-    def validate_fully_applied(self) -> None:
+    def validate_fully_applied(self, loaded: Iterable[str] | None = None) -> None:
         """Close the fusion: every edit must have met its parameter.
 
         A silently unapplied delta is the failure mode that matters here: the
         model would load and generate, just not as the distilled student. The
         weights are loaded once, so the mapped payloads are dropped afterwards
         rather than held for the life of the process.
+
+        ``loaded`` is the set of parameter names ``load_weights`` actually
+        consumed. A gate is assigned rather than fused, so it lands on a module
+        the base transformer does not have; if that module was never built,
+        ``load_weights`` only logs a skip and the server would serve a
+        zero-initialized gate. Yielding a tensor is not evidence it arrived, so
+        the injections are closed against that set when it is available.
         """
         missing = sorted(set(self._patches) - self._applied)
         if missing:
             raise FastH3AdapterError(
                 f"FastH3 adapter edits {len(missing)} parameters the checkpoint never provided: {missing[:5]}"
             )
+        arrived = self._injected if loaded is None else set(loaded)
+        uninjected = sorted(set(self._injections) - arrived)
+        if uninjected:
+            raise FastH3AdapterError(
+                f"FastH3 adapter assigns {len(uninjected)} parameters that never reached the model: {uninjected[:5]}"
+            )
         for patch in self._patches.values():
             patch.low_rank.clear()
             patch.diff = None
+        self._injections.clear()
 
     def check_serving_contract(
         self,
@@ -555,11 +620,22 @@ class FastH3WeightFusion:
                 f"{sorted(offloads)}. Serve it without offload."
             )
         if self.requires_vsa:
-            raise ValueError(
-                f"{self.source} is a Video Sparse Attention variant of FastH3, and its "
-                "compression gates have no counterpart in vLLM-Omni's dense H3 attention. Use the "
-                "release's Dense variant until VSA lands for H3's packed sequence."
-            )
+            backend = _resolve_dit_attention_backend(od_config)
+            if backend != "FASTVIDEO_VSA":
+                raise ValueError(
+                    f"{self.source} is a Video Sparse Attention variant of FastH3. Its compression "
+                    "gates only mean anything to the VSA kernel, and any other backend would run it "
+                    f"as dense attention on a student distilled for 90% sparsity (got {backend or 'default'}). "
+                    "Serve it with --diffusion-attention-backend FASTVIDEO_VSA."
+                )
+            parallel_config = getattr(od_config, "parallel_config", None)
+            ring_degree = int(getattr(parallel_config, "ring_degree", 1) or 1)
+            allgather_degree = int(getattr(parallel_config, "allgather_degree", 1) or 1)
+            if ring_degree != 1 or allgather_degree != 1:
+                raise ValueError(
+                    "FastH3 VSA supports local attention or pure Ulysses sequence parallelism; "
+                    "ring/all-gather SP does not give the block-sparse kernel the complete packed sequence."
+                )
         logger.info(
             "FastH3 adapter active: sigma points %s for %d transformer forwards, "
             "flow_shift=%g, audio_flow_shift=%g, tasks=%s",

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -422,6 +422,13 @@ class MiniMaxH3Attention(nn.Module):
             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.
+        self.to_gate_compress: ColumnParallelLinear | None = None
+        self._gate_hidden_size = arch.hidden_size
+        self._gate_quant_config = quant_config
+        self._gate_prefix = f"{prefix}.to_gate_compress"
         self.attention = Attention(
             num_heads=self.num_heads,
             num_kv_heads=self.num_kv_heads,
@@ -436,6 +443,26 @@ class MiniMaxH3Attention(nn.Module):
             prefix=prefix,
         )
 
+    def enable_vsa_gate(self) -> None:
+        """Build the VSA compression gate this attention would otherwise lack.
+
+        Called before ``load_weights`` so the artifact's ``.set_weight`` tensor
+        has a parameter to land on. Zero-initialized like the Wan VSA layers, so
+        a gate that never receives weights degrades to sparse-only selection
+        rather than to garbage.
+        """
+        if self.to_gate_compress is not None:
+            return
+        self.to_gate_compress = ColumnParallelLinear(
+            self._gate_hidden_size,
+            self.total_num_heads * self.head_dim,
+            bias=False,
+            params_dtype=_BF16_DTYPE,
+            quant_config=self._gate_quant_config,
+            prefix=self._gate_prefix,
+        )
+        nn.init.zeros_(self.to_gate_compress.weight)
+
     def _apply_rope(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
         """Rotate the first rot_dim head dims; pass the rest through.
 
@@ -461,6 +488,8 @@ class MiniMaxH3Attention(nn.Module):
         packed_total: int,
         num_requests: int = 1,
         video_layout: VideoTokenLayout | None = None,
+        vsa_prefix_segments: tuple[int, ...] = (),
+        gate_compress: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Run packed attention as a small eager island.
 
@@ -539,6 +568,16 @@ class MiniMaxH3Attention(nn.Module):
                 # (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.
+                **(
+                    {"vsa_h3_prefix_segments": vsa_prefix_segments}
+                    if gate_compress is not None and video_layout is not None and video_layout.video_spans
+                    else {}
+                ),
             },
             video_layout=video_layout,
         )
@@ -560,6 +599,7 @@ class MiniMaxH3Attention(nn.Module):
         num_requests: int = 1,
         sp_seq_lens: list[int] | None = None,
         video_layout: VideoTokenLayout | None = None,
+        vsa_prefix_segments: tuple[int, ...] = (),
     ) -> torch.Tensor:
         """x: [T, hidden] packed thd rows -> [T, hidden].
 
@@ -593,6 +633,15 @@ class MiniMaxH3Attention(nn.Module):
                 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)
+
         # Each request contributes a document for its rows plus one for any
         # nonempty alignment padding. Local/Ulysses backends unpad it, while
         # Ring keeps aligned rows for fixed-size P2P buffers.
@@ -608,6 +657,8 @@ class MiniMaxH3Attention(nn.Module):
             packed_total=packed_total if packed_total is not None else q.shape[0],
             num_requests=num_requests,
             video_layout=video_layout,
+            vsa_prefix_segments=vsa_prefix_segments,
+            gate_compress=gate_compress,
         )
         out = out.reshape(total, self.num_heads * self.head_dim)
         out, _ = self.out_proj(out)
@@ -826,6 +877,7 @@ class MiniMaxH3DiTBlock(nn.Module):
         num_requests: int = 1,
         sp_seq_lens: list[int] | None = None,
         video_layout: VideoTokenLayout | None = None,
+        vsa_prefix_segments: tuple[int, ...] = (),
     ) -> torch.Tensor:
         """x: [T, H]; t_emb: [M, t_dim]; combined_indices: [T]
         (= inverse_indices * modality_num + token_tags.clamp(min=0)).
@@ -861,6 +913,7 @@ class MiniMaxH3DiTBlock(nn.Module):
             num_requests=num_requests,
             sp_seq_lens=sp_seq_lens,
             video_layout=video_layout,
+            vsa_prefix_segments=vsa_prefix_segments,
         )
         x, h = indexed_gate_rms_norm_scale_shift(
             residual,
@@ -1130,6 +1183,7 @@ class MiniMaxH3DiTModel(nn.Module):
         self.sp_prepare = MiniMaxH3SPPrepare()
         self.local_sp_prepare = MiniMaxH3SPPrepare()
         self.sp_gather = MiniMaxH3SPGather()
+        self.vsa_gates_enabled = False
         self.final_layer = MiniMaxH3FinalLayer(
             arch,
             quant_config,
@@ -1137,6 +1191,20 @@ class MiniMaxH3DiTModel(nn.Module):
         )
         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.
+        """
+        if self.vsa_gates_enabled:
+            return
+        for block in self.blocks:
+            block.attn.enable_vsa_gate()
+        self.vsa_gates_enabled = True
+
     def _mark_missing_params_required(self) -> None:
         for _, param in self.named_parameters():
             param.missing_param_init = "error"
@@ -1399,6 +1467,7 @@ class MiniMaxH3DiTModel(nn.Module):
         # attention never reads cu_seqlens scalars off the device; a producer
         # that omits it is packing a single request.
         num_requests = int(self._psp_optional(psp, "num_requests", 1))
+        vsa_prefix_segments = tuple(int(length) for length in self._psp_optional(psp, "vsa_prefix_segments", ()))
         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"))
@@ -1480,6 +1549,7 @@ class MiniMaxH3DiTModel(nn.Module):
                 packed_total=seq_len,
                 num_requests=num_requests,
                 video_layout=video_layout,
+                vsa_prefix_segments=vsa_prefix_segments,
             )
         if local_len == seq_len:
             hidden = self.sp_gather(hidden)

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -959,6 +959,13 @@ class MiniMaxH3Pipeline(
             )
 
         self._fasth3 = resolve_fasth3_fusion(od_config, self.transformer)
+        if self._fasth3 is not None and self._fasth3.requires_vsa:
+            # The artifact assigns a compression gate per DiT block, so those
+            # modules have to exist before load_weights streams them in. Only
+            # the ``transformer.`` stream is fused, and ``check_task`` admits
+            # T2VA only, so the Ref2VA DiT would carry 50 gates that nothing
+            # ever fills or reads.
+            self.transformer.enable_vsa_gates()
         if self._fasth3 is not None:
             self._fasth3.check_serving_contract(
                 partition=self.partition,
@@ -1070,6 +1077,7 @@ class MiniMaxH3Pipeline(
 
         loaded_with_prefix: set[str] = set()
         loaded_prefixes: set[str] = set()
+        transformer_loaded: set[str] = set()
         for prefix, grouped_weights in groupby(weights, key=source_prefix):
             if prefix in loaded_prefixes:
                 raise ValueError(f"MiniMax-H3 weight source {prefix!r} is not contiguous")
@@ -1081,6 +1089,8 @@ class MiniMaxH3Pipeline(
                 # point where the checkpoint's fused QKV/MLP layouts are intact.
                 stream = self._fasth3.apply(stream)
             loaded = component.load_weights(stream)
+            if prefix == "transformer.":
+                transformer_loaded = set(loaded)
             if prefix != "text_encoder.":
                 component.post_load_weights()
             loaded_with_prefix.update(prefix + name for name in loaded)
@@ -1094,7 +1104,9 @@ class MiniMaxH3Pipeline(
                 continue
             loaded_with_prefix.update(f"{component_name}.{name}" for name, _ in component.named_parameters())
         if self._fasth3 is not None:
-            self._fasth3.validate_fully_applied()
+            # load_weights only warns on a parameter the model does not have, so
+            # close the adapter against what the DiT actually consumed.
+            self._fasth3.validate_fully_applied(transformer_loaded)
         return loaded_with_prefix
 
     @property