[Bugfix] Fix MiniMax-H3 packed padding with Ring attention

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

recipes/MiniMaxAI/MiniMax-H3.md

@@ -1157,8 +1157,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

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -905,6 +905,64 @@ def forward(self, query, key, value, metadata):
     )
 
 
+def test_ring_packed_attention_publishes_prefix_length_without_global_mask():
+    attention = _fake_packed_attention("CUDNN_ATTN")
+    attention.attention.use_ring = True
+    q = torch.randn(4, 2, 4)
+
+    attention._run_packed_attention(
+        q,
+        q,
+        q,
+        cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
+        max_seqlen=5,
+        packed_total=8,
+    )
+
+    metadata = attention.attention.metadata
+    assert metadata.attn_mask is None
+    assert metadata.extra["valid_kv_length"] == 5
+
+
+def test_ring_packed_attention_keeps_mask_when_runtime_sp_is_inactive(monkeypatch):
+    from vllm_omni.diffusion.models.minimax_h3 import minimax_h3_transformer
+
+    class InactiveContext:
+        sp_active = False
+
+    attention = _fake_packed_attention("CUDNN_ATTN")
+    attention.attention.use_ring = True
+    monkeypatch.setattr(
+        minimax_h3_transformer,
+        "is_forward_context_available",
+        lambda: True,
+        raising=False,
+    )
+    monkeypatch.setattr(
+        minimax_h3_transformer,
+        "get_forward_context",
+        InactiveContext,
+        raising=False,
+    )
+    q = torch.randn(8, 2, 4)
+
+    attention._run_packed_attention(
+        q,
+        q,
+        q,
+        cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
+        max_seqlen=5,
+        packed_total=8,
+    )
+
+    metadata = attention.attention.metadata
+    assert torch.equal(
+        metadata.attn_mask,
+        torch.tensor([[True, True, True, True, True, False, False, False]]),
+    )
+    assert metadata.extra["npu_attn_varlen"] is True
+
+
 def _fake_packed_attention(
     backend_name: str,
     *,
@@ -926,6 +984,10 @@ def get_name() -> str:
         def supports_multi_doc_packed_varlen(cls) -> bool:
             return supports_multi_doc
 
+        @classmethod
+        def supports_packed_mask_free(cls) -> bool:
+            return False
+
     class FakeAttention(torch.nn.Module):
         attn_backend = FakeBackend
 

vllm_omni/diffusion/attention/backends/abstract.py

@@ -208,8 +208,9 @@ class AttentionMetadata:
     #     variable-length query/key sequences for FlashAttention.
     #   "max_seqlen_q" / "max_seqlen_k": maximum sequence lengths paired with
     #     the packed cu_seqlens tensors.
-    #   "valid_kv_length": int — contiguous valid K/V prefix length for a
-    #     backend that advertises supports_prefix_kv_slicing.
+    #   "valid_kv_length": int — contiguous valid K/V prefix length. Local
+    #     backends may slice it directly; ring backends trim each circulated
+    #     K/V block from the same global prefix.
     #   "npu_attn_varlen": bool — model opt-in for the NPU packed varlen path
     #     (TND npu_fusion_attention driven by cu_seqlens, mask never read).
     #     Requires the [real, pad] two-document packing contract; see

vllm_omni/diffusion/attention/backends/ring/ring_utils.py

@@ -1,4 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 # Copyright (c) 2024, Jiarui Fang.
 # Adapted from https://github.com/feifeibear/long-context-attention
 
@@ -9,7 +10,33 @@
 
 LseLayout = Literal["bhs", "bsh"]
 
-__all__ = ["update_out_and_lse", "flatten_varlen_lse", "unflatten_varlen_lse"]
+__all__ = [
+    "ring_kv_block_valid_length",
+    "update_out_and_lse",
+    "flatten_varlen_lse",
+    "unflatten_varlen_lse",
+]
+
+
+def ring_kv_block_valid_length(
+    valid_kv_length: int | None,
+    block_size: int,
+    block_rank: int,
+    world_size: int,
+) -> int:
+    """Return the valid prefix length in one globally ordered ring K/V block."""
+    if valid_kv_length is None:
+        return block_size
+    if isinstance(valid_kv_length, bool) or not isinstance(valid_kv_length, int):
+        raise ValueError("valid_kv_length must be a Python integer")
+    total_length = block_size * world_size
+    if not 0 < valid_kv_length <= total_length:
+        raise ValueError(
+            "valid_kv_length must be within the global ring K/V sequence, "
+            f"got {valid_kv_length} for length {total_length}"
+        )
+    block_start = block_rank * block_size
+    return max(0, min(block_size, valid_kv_length - block_start))
 
 
 def _normalize_lse(block_lse: torch.Tensor, block_out: torch.Tensor, lse_layout: LseLayout) -> torch.Tensor: