[Model][MiniMax-H3] Multi-GPU serving fixes and layer-wise offload component selection

docs/design/feature/sequence_parallel.md

@@ -77,7 +77,9 @@ Use it when plain Ulysses-SP would otherwise fail because:
 
 - `auto_pad=True` pads sequence tokens in `_sp_plan` and requires attention backends that support `attention_mask`.
 - `advanced_uaa` does not depend on mask-based token padding inside Ulysses attention. It is therefore a better fit for non-divisible head counts and uneven Ulysses shard sizes.
-- `auto_pad=True` remains incompatible with Ring attention because the ring backend does not consume `attention_mask`.
+- `auto_pad=True` remains incompatible with Ring attention because the ring backend does not consume arbitrary
+  `attention_mask` values. A model with contiguous suffix padding may instead publish `valid_kv_length`; Ring trims
+  each circulated K/V block against that global valid prefix.
 - `advanced_uaa` is still experimental and hybrid mode remains limited by Ring's equal-shape requirement.
 
 ---
@@ -416,7 +418,7 @@ _sp_plan = {
 | Constraint | Description |
 |------------|-------------|
 | **Attention Backend Compatibility** | The attention backends must support `attention_mask`. Currently only `TORCH_SDPA` and `FLASH_ATTN` (default for diffusion models) are supported. |
-| **Ring Attention Limitation** | Ring attention does not support `attention_mask`. Therefore, when using `auto_pad=True`, the combination of Ulysses + Ring attention is not feasible. |
+| **Ring Attention Limitation** | Ring attention does not support arbitrary `attention_mask` values, so `auto_pad=True` remains incompatible. Contiguous suffix padding is supported when the model publishes the global `valid_kv_length` prefix. |
 
 1. Enable `auto_pad=True` for all sequence-dimension inputs in `_sp_plan`:
 ```python

docs/user_guide/diffusion/regional_compilation.md

@@ -34,5 +34,34 @@ that provide their own `setup_compile()` implementation manage their compilation
 policy independently. Compilation is lazy, so backend or graph errors can first
 surface on the initial request.
 
+## Pinning one packed shape (MiniMax-H3)
+
+Compiled regions are keyed by input shape. MiniMax-H3 packs a request into a
+row count that follows the request -- a longer prompt, more frames, a larger
+frame -- and rounds it up to the next 64-row boundary, so two requests can land
+on two different packed lengths.
+
+Under the default `diffusion_compile_dynamic: true` a new packed length is
+absorbed by the dynamic shape, so it is not by itself a recompilation. It still
+hands the backend a shape it has not autotuned, which is the risk this knob
+exists to remove: a fixed-shape deployment (`--no-diffusion-compile-dynamic`)
+recompiles the transformer blocks per shape, and a first-seen large shape can
+push autotuning into its worst-case memory use.
+
+A request pins the length with `extra_args["pad_seq_len"]`:
+
+```python
+sampling_params = SamplingParams(extra_args={"pad_seq_len": 54080})
+```
+
+The value must be a positive multiple of 64 and must cover the rows the request
+actually uses; the packed sequence is then padded to it instead of to the next
+64-row boundary. Pick a bucket that covers every request dimension the
+deployment accepts -- prompt length, frame count, frame size and reference
+blocks all change the used row count. The server logs the effective length as
+`MiniMax H3 packed sequence: ... pad_seq_len=... used=... seq_len=...` whenever
+a request pins it. The padding rows are masked out, so their cost is the
+attention and feed-forward work on those rows.
+
 Use `--enforce-eager` to disable the model runner's generic compile setup.
 Pipelines that compile internally define their own eager-mode behavior.

recipes/MiniMaxAI/MiniMax-H3.md

@@ -1240,8 +1240,9 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
 - `--cfg-parallel-size > 1` is rejected by design (CFG-distilled, no negative branch).
 - VAE patch parallelism requires size 1 or the full DiT group size and supports the
   H3 native `tile` mode only.
-- A U2 x Ring2 hybrid currently fails with an attention-mask length mismatch; use
-  pure Ulysses.
+- Ulysses x Ring hybrid attention supports H3's single-request contiguous suffix
+  padding. Arbitrary attention masks and multi-request packed batches remain
+  unsupported on the Ring path.
 - Online FP8 with DLO AllGather temporarily materializes the complete FP8 model
   in host memory on every rank during startup before retaining only each rank's
   shard. Size startup host memory for that transient peak.

tests/diffusion/attention/test_ring_padding.py

@@ -0,0 +1,221 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from dataclasses import dataclass
+from typing import Any
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.attention.backends import ring_flash_attn, ring_pytorch_attn
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType
+from vllm_omni.diffusion.attention.backends.ring.ring_utils import ring_kv_block_valid_length
+from vllm_omni.diffusion.attention.parallel import ring as ring_parallel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _SingleRankComm:
+    rank = 0
+    world_size = 1
+
+
+@dataclass(frozen=True)
+class _FakeSequenceParallelGroup:
+    ring_group: object
+
+
+def _fake_attention_result(q, _k, _v, **_kwargs):
+    return torch.zeros_like(q), torch.zeros(q.shape[0], q.shape[2], q.shape[1])
+
+
+@pytest.mark.parametrize(("block_rank", "expected"), [(0, 4), (1, 1), (2, 0)])
+def test_global_prefix_is_partitioned_across_ring_blocks(block_rank, expected):
+    assert ring_kv_block_valid_length(5, 4, block_rank, 3) == expected
+
+
+def test_absent_prefix_length_keeps_the_full_ring_block():
+    assert ring_kv_block_valid_length(None, 4, 2, 3) == 4
+
+
+def test_pytorch_ring_trims_padding_from_each_kv_block(monkeypatch):
+    captured: dict[str, Any] = {}
+
+    def fake_attention(q, k, v, **kwargs):
+        captured.update(k=k, v=v, **kwargs)
+        return _fake_attention_result(q, k, v, **kwargs)
+
+    monkeypatch.setattr(ring_pytorch_attn, "RingComm", lambda _group: _SingleRankComm())
+    monkeypatch.setattr(ring_pytorch_attn, "pytorch_attn_forward", fake_attention)
+    qkv = torch.randn(1, 4, 2, 8)
+
+    output = ring_pytorch_attn.ring_pytorch_attn_func(
+        qkv,
+        qkv,
+        qkv,
+        group=object(),
+        valid_kv_length=3,
+    )
+
+    assert output.shape == qkv.shape
+    assert captured["k"].shape[1] == 3
+    assert captured["v"].shape[1] == 3
+
+
+def test_pytorch_ring_trims_circulated_blocks_by_source_rank(monkeypatch):
+    observed_lengths = []
+    # Each ring step sends K and V separately.
+    blocks = [torch.randn(1, 4, 2, 8) for _ in range(4)]
+
+    class FakeRingComm:
+        rank = 2
+        world_size = 3
+
+        def __init__(self, _group):
+            self._next_block = 0
+
+        def send_recv(self, _tensor):
+            block = blocks[self._next_block]
+            self._next_block += 1
+            return block
+
+        def commit(self):
+            pass
+
+        def wait(self):
+            pass
+
+    def fake_attention(q, k, v, **kwargs):
+        observed_lengths.append(k.shape[1])
+        return _fake_attention_result(q, k, v, **kwargs)
+
+    monkeypatch.setattr(ring_pytorch_attn, "RingComm", FakeRingComm)
+    monkeypatch.setattr(ring_pytorch_attn, "pytorch_attn_forward", fake_attention)
+    qkv = torch.randn(1, 4, 2, 8)
+
+    ring_pytorch_attn.ring_pytorch_attn_func(
+        qkv,
+        qkv,
+        qkv,
+        group=object(),
+        valid_kv_length=5,
+    )
+
+    # Rank 2 starts with global block 2, then receives blocks 1 and 0.
+    # The first block is empty and therefore does not launch attention.
+    assert observed_lengths == [1, 4]
+
+
+def test_flash_ring_trims_padding_from_each_kv_block(monkeypatch):
+    captured: dict[str, Any] = {}
+
+    def fake_attention(q, k, v, **kwargs):
+        captured.update(k=k, v=v, **kwargs)
+        return _fake_attention_result(q, k, v, **kwargs)
+
+    monkeypatch.setattr(ring_flash_attn, "RingComm", lambda _group: _SingleRankComm())
+    monkeypatch.setattr(ring_flash_attn, "select_flash_attn_impl", lambda *_args, **_kwargs: fake_attention)
+    qkv = torch.randn(1, 4, 2, 8)
+
+    output, _ = ring_flash_attn.ring_flash_attn_forward(
+        object(),
+        qkv,
+        qkv,
+        qkv,
+        softmax_scale=0.25,
+        causal=False,
+        attn_type=AttnType.FA,
+        valid_kv_length=3,
+    )
+
+    assert output.shape == qkv.shape
+    assert captured["k"].shape[1] == 3
+    assert captured["v"].shape[1] == 3
+
+
+def test_ring_dispatch_forwards_global_valid_kv_length(monkeypatch):
+    captured: dict[str, Any] = {}
+
+    def fake_ring_attention(*_args, **kwargs):
+        captured.update(kwargs)
+        return "output"
+
+    monkeypatch.setattr(ring_pytorch_attn, "ring_pytorch_attn_func", fake_ring_attention)
+    strategy = ring_parallel.RingParallelAttention(
+        _FakeSequenceParallelGroup(ring_group=object()),
+        attn_backend_pref="sdpa",
+    )
+    metadata = AttentionMetadata(extra={"valid_kv_length": 3})
+    query = torch.randn(1, 4, 2, 8)
+
+    assert strategy.run_attention(query, query, query, metadata) == "output"
+    assert captured["valid_kv_length"] == 3
+
+
+def test_pytorch_ring_rejects_causal_with_a_valid_prefix(monkeypatch):
+    """Trimming a K/V block under causal=True would silently shift the diagonal.
+
+    The query block keeps the full padded length while the circulated K/V block
+    is trimmed, and a causal mask over a rectangular tile is bottom-right
+    aligned, so the combination must be refused rather than computed.
+    """
+    monkeypatch.setattr(ring_pytorch_attn, "RingComm", lambda _group: _SingleRankComm())
+    monkeypatch.setattr(ring_pytorch_attn, "pytorch_attn_forward", _fake_attention_result)
+    qkv = torch.randn(1, 4, 2, 8)
+
+    with pytest.raises(ValueError, match="valid_kv_length is not supported with causal=True"):
+        ring_pytorch_attn.ring_pytorch_attn_func(
+            qkv,
+            qkv,
+            qkv,
+            causal=True,
+            group=object(),
+            valid_kv_length=3,
+        )
+
+
+def test_flash_ring_rejects_causal_with_a_valid_prefix(monkeypatch):
+    """Same refusal on the Flash-Attention ring path."""
+    monkeypatch.setattr(ring_flash_attn, "RingComm", lambda _group: _SingleRankComm())
+    monkeypatch.setattr(
+        ring_flash_attn,
+        "select_flash_attn_impl",
+        lambda *_args, **_kwargs: _fake_attention_result,
+    )
+    qkv = torch.randn(1, 4, 2, 8)
+
+    with pytest.raises(ValueError, match="valid_kv_length is not supported with causal=True"):
+        ring_flash_attn.ring_flash_attn_forward(
+            object(),
+            qkv,
+            qkv,
+            qkv,
+            softmax_scale=0.25,
+            causal=True,
+            attn_type=AttnType.FA,
+            valid_kv_length=3,
+        )
+
+
+def test_ring_paths_still_accept_causal_without_a_valid_prefix(monkeypatch):
+    """The guard is scoped to the padded case: plain causal Ring is unchanged."""
+    monkeypatch.setattr(ring_flash_attn, "RingComm", lambda _group: _SingleRankComm())
+    monkeypatch.setattr(
+        ring_flash_attn,
+        "select_flash_attn_impl",
+        lambda *_args, **_kwargs: _fake_attention_result,
+    )
+    qkv = torch.randn(1, 4, 2, 8)
+
+    output, _ = ring_flash_attn.ring_flash_attn_forward(
+        object(),
+        qkv,
+        qkv,
+        qkv,
+        softmax_scale=0.25,
+        causal=True,
+        attn_type=AttnType.FA,
+    )
+
+    assert output.shape == qkv.shape