[Bugfix][MiniMax-H3] Fix text-encoder TP group tiling and NPU packed varlen under AllGather-KV SP

tests/diffusion/attention/test_flash_attn.py

@@ -17,7 +17,7 @@
 import torch
 
 from tests.helpers.mark import hardware_test
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, QueryRange
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
 from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
 from vllm_omni.diffusion.attention.backends.utils import fa  # noqa: E402
@@ -628,6 +628,75 @@ def test_resolve_packed_seq_rejects_malformed_metadata(case):
     assert impl._resolve_packed_seq_npu(q, q, extra) is None
 
 
+# --- Test group A2: AllGather-KV (CP) local-Q/global-KV generalization -------
+#
+# With AllGather-KV sequence parallelism the backend sees a local query shard
+# against the global K/V. The strategy records the shard's position as a
+# single QueryRange; the resolver must then derive the shard's [real, pad]
+# boundary pair from the global real length instead of rejecting the contract
+# (used_q is global and exceeds the local shard length by construction).
+
+
+@pytest.mark.cpu
+def test_resolve_packed_seq_cp_local_query_shards():
+    """Each AllGather-KV rank resolves its own [real, pad] Q boundaries."""
+    impl = _npu_impl()
+    # Global packed layout: 5 real + 3 pad = 8 tokens; CP x2 splits Q into two
+    # 4-row shards. max_seqlen_* describe the GLOBAL real document length.
+    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
+    extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 5, "max_seqlen_k": 5}
+    k_global = torch.randn(1, 8, 2, 4)
+    q_shard = torch.randn(1, 4, 2, 4)
+
+    # Rank 0 owns global rows [0, 4): all real, so the local pad document is
+    # empty (zero-length tail).
+    rank0_ranges = (QueryRange(local_start=0, local_end=4, global_start=0),)
+    assert impl._resolve_packed_seq_npu(q_shard, k_global, extra, rank0_ranges) == ([4, 4], [5, 8])
+
+    # Rank 1 owns global rows [4, 8): 1 real + 3 pad rows.
+    rank1_ranges = (QueryRange(local_start=0, local_end=4, global_start=4),)
+    assert impl._resolve_packed_seq_npu(q_shard, k_global, extra, rank1_ranges) == ([1, 4], [5, 8])
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "ranges",
+    [
+        # Joint layouts produce multiple ranges; the CP generalization is only
+        # defined for a single packed query range.
+        (
+            QueryRange(local_start=0, local_end=2, global_start=0),
+            QueryRange(local_start=2, local_end=4, global_start=4),
+        ),
+        # Shard starts beyond the global real document: nothing real locally.
+        (QueryRange(local_start=0, local_end=4, global_start=6),),
+        # Shard starts exactly at the real boundary: zero real rows locally.
+        (QueryRange(local_start=0, local_end=4, global_start=5),),
+    ],
+)
+def test_resolve_packed_seq_cp_rejects_unsupported_ranges(ranges):
+    impl = _npu_impl()
+    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
+    extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 5, "max_seqlen_k": 5}
+    q_shard = torch.randn(1, 4, 2, 4)
+    k_global = torch.randn(1, 8, 2, 4)
+
+    assert impl._resolve_packed_seq_npu(q_shard, k_global, extra, ranges) is None
+
+
+@pytest.mark.cpu
+def test_resolve_packed_seq_without_ranges_keeps_square_contract():
+    """No query_ranges means no SP slicing: the square contract still applies
+    and a global-length used_q against a local-length query is rejected."""
+    impl = _npu_impl()
+    cu = torch.tensor([0, 5, 8], dtype=torch.int32)
+    extra = {"cu_seqlens_q": cu, "cu_seqlens_k": cu, "max_seqlen_q": 5, "max_seqlen_k": 5}
+    q_shard = torch.randn(1, 4, 2, 4)
+    k_global = torch.randn(1, 8, 2, 4)
+
+    assert impl._resolve_packed_seq_npu(q_shard, k_global, extra) is None
+
+
 # --- Test group B: env dispatch in forward_fa_npu (current behavior) --------
 
 

tests/diffusion/models/minimax_h3/test_minimax_h3_encoder_group.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""CPU tests for the text-encoder TP group layout.
+
+``_text_encoder_group_ranks`` decides the process groups behind
+``--text-encoder-tp-size``.  It must tile the DiT world into contiguous
+groups so that every rank holds group membership (``GroupCoordinator``
+asserts on ranks outside all groups), while keeping the first group exactly
+``[0, tp_size)`` so the "first tp_size ranks encode" contract is unchanged.
+"""
+
+import pytest
+
+from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
+    _text_encoder_group_ranks,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.mark.parametrize(
+    ("dit_world", "tp_size", "expected"),
+    [
+        # tp_size == world: a single group covering every rank, identical to
+        # the previous init_world_group([0..tp)) behavior.
+        (4, 4, [[0, 1, 2, 3]]),
+        (8, 8, [[0, 1, 2, 3, 4, 5, 6, 7]]),
+        # tp_size < world: tiled groups, every rank is a member of exactly one.
+        (16, 8, [[0, 1, 2, 3, 4, 5, 6, 7], [8, 9, 10, 11, 12, 13, 14, 15]]),
+        (16, 4, [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]]),
+        (16, 2, [[2 * i, 2 * i + 1] for i in range(8)]),
+        (4, 2, [[0, 1], [2, 3]]),
+        (8, 4, [[0, 1, 2, 3], [4, 5, 6, 7]]),
+    ],
+)
+def test_text_encoder_group_ranks_tile_dit_world(dit_world, tp_size, expected):
+    group_ranks = _text_encoder_group_ranks(dit_world, tp_size)
+
+    assert group_ranks == expected
+    # Every DiT rank is a member of exactly one group (GroupCoordinator assert).
+    assert sorted(rank for group in group_ranks for rank in group) == list(range(dit_world))
+    # The encoding group contract is unchanged: the first tp_size ranks encode.
+    assert group_ranks[0] == list(range(tp_size))
+
+
+@pytest.mark.parametrize(("dit_world", "tp_size"), [(16, 3), (16, 6), (8, 3), (4, 3), (2, 3)])
+def test_text_encoder_group_ranks_rejects_non_divisible(dit_world, tp_size):
+    with pytest.raises(ValueError, match="must divide the DiT group size"):
+        _text_encoder_group_ranks(dit_world, tp_size)
+
+
+def test_text_encoder_group_ranks_rejects_zero_tp():
+    with pytest.raises(ValueError, match=">= 1"):
+        _text_encoder_group_ranks(16, 0)

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -8,7 +8,12 @@
 import torch
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend, AttentionImpl, AttentionMetadata
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionBackend,
+    AttentionImpl,
+    AttentionMetadata,
+    QueryRange,
+)
 from vllm_omni.diffusion.attention.backends.sdpa import _maybe_reshape_attn_mask
 from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
     piecewise_attn,
@@ -20,6 +25,18 @@
 logger = init_logger(__name__)
 
 
+def _packed_query_ranges(attn_metadata: AttentionMetadata | None) -> tuple[QueryRange, ...] | None:
+    """Query ranges for the packed resolvers, or None outside SP slicing.
+
+    The AllGather-KV strategies record where the local query shard sits in the
+    global sequence via ``query_ranges``; the packed paths use that to
+    generalize the square [real, pad] contract to local-Q/global-KV.
+    """
+    if attn_metadata is None or attn_metadata.query_ranges is None:
+        return None
+    return tuple(attn_metadata.query_ranges)
+
+
 @cache
 def _get_npu_compressed_causal_mask(device: torch.device) -> torch.Tensor:
     """Return the shared block mask used by NPU right-down causal attention."""
@@ -666,9 +683,21 @@ def forward_fa_npu(
         extra = attn_metadata.extra if attn_metadata else {}
         if extra.get("npu_attn_varlen", False):
             if os.environ.get("MINDIE_SD_FA_TYPE") == "ascend_laser_attention":
-                out = self._forward_prefix_kv_slice_npu(query, key, value, extra)
+                out = self._forward_prefix_kv_slice_npu(
+                    query,
+                    key,
+                    value,
+                    extra,
+                    query_ranges=_packed_query_ranges(attn_metadata),
+                )
             else:
-                out = self._forward_varlen_packed_npu(query, key, value, extra)
+                out = self._forward_varlen_packed_npu(
+                    query,
+                    key,
+                    value,
+                    extra,
+                    query_ranges=_packed_query_ranges(attn_metadata),
+                )
             if out is not None:
                 return out
         if attention_mask is None and extra.get("npu_attn_varlen", False):
@@ -677,13 +706,19 @@ def forward_fa_npu(
             # paths above declined (contract mismatch), so rebuild the padding
             # mask here to keep the masked fallback correct.
             used = extra.get("valid_kv_length")
-            if not isinstance(used, int) or not 0 < used <= query.shape[1]:
+            # The mask marks the padding suffix of the K/V domain. With
+            # AllGather-KV SP the K/V domain is the global sequence while the
+            # query is a local shard, so the domain is key.shape[1] (equal to
+            # query.shape[1] in the square non-SP case, where behavior is
+            # unchanged).
+            kv_len = key.shape[1]
+            if not isinstance(used, int) or not 0 < used <= kv_len:
                 raise ValueError(
                     "npu_attn_varlen packed metadata is unusable and no attn_mask "
-                    f"was constructed (valid_kv_length={used!r}, seq_len={query.shape[1]}); "
+                    f"was constructed (valid_kv_length={used!r}, kv_len={kv_len}); "
                     "refusing to run unmasked attention over padding rows."
                 )
-            attention_mask = torch.arange(query.shape[1], device=query.device)[None] < used
+            attention_mask = torch.arange(kv_len, device=query.device)[None] < used
 
         # NPU aclnnFlashAttentionScore requires mask shape to be one of:
         # [B, N, Sq, Skv], [B, 1, Sq, Skv], [1, 1, Sq, Skv], or [Sq, Skv]
@@ -707,6 +742,7 @@ def _resolve_packed_seq_npu(
         query: torch.Tensor,
         key: torch.Tensor,
         extra: dict,
+        query_ranges: tuple[QueryRange, ...] | None = None,
     ) -> tuple[list[int], list[int]] | None:
         """Resolve packed document boundaries (cumulative end offsets per doc).
 
@@ -718,6 +754,13 @@ def _resolve_packed_seq_npu(
             padding document as a strict suffix;
           - max_seqlen_q/k are Python ints equal to the real document length
             (so boundaries are derived without any device sync).
+
+        With AllGather-KV sequence parallelism, the kernel sees local Q
+        against global K/V. When ``query_ranges`` carries the single packed
+        range produced by the AllGather-KV strategy, the local Q rows are a
+        contiguous slice of the same global real document, so the [real, pad]
+        contract generalizes: seq_q becomes that slice's [local_real,
+        local_total] pair while seq_k keeps the global [real, pad] pair.
         """
         if self.causal:
             return None
@@ -734,15 +777,22 @@ def _resolve_packed_seq_npu(
         # .shape is host-side metadata: counting documents never syncs.
         if cu_q.shape[0] != cu_k.shape[0] or cu_q.shape[0] > 3:
             return None
-        if not (0 < used_q <= total_q) or not (0 < used_k <= total_k):
+        if not (0 < used_q) or not (0 < used_k <= total_k):
+            return None
+        if query_ranges is not None and len(query_ranges) == 1:
+            q_start = query_ranges[0].global_start
+            local_real_q = min(used_q - q_start, total_q)
+            if 0 <= q_start < used_q and local_real_q > 0:
+                return [local_real_q, total_q], [used_k, total_k]
             return None
-        if used_q == total_q and used_k == total_k:
-            # No padding document: one full-length document.
-            return [total_q], [total_k]
-        if cu_q.shape[0] == 3 and used_q >= total_q - used_q and used_k >= total_k - used_k:
-            # [real, pad] packing: the real document must be the longer one
-            # (consistent with the max_seqlen naming).
-            return [used_q, total_q], [used_k, total_k]
+        if used_q <= total_q:
+            if used_q == total_q and used_k == total_k:
+                # No padding document: one full-length document.
+                return [total_q], [total_k]
+            if cu_q.shape[0] == 3 and used_q >= total_q - used_q and used_k >= total_k - used_k:
+                # [real, pad] packing: the real document must be the longer one
+                # (consistent with the max_seqlen naming).
+                return [used_q, total_q], [used_k, total_k]
         return None
 
     def _forward_varlen_packed_npu(
@@ -751,13 +801,14 @@ def _forward_varlen_packed_npu(
         key: torch.Tensor,
         value: torch.Tensor,
         extra: dict,
+        query_ranges: tuple[QueryRange, ...] | None = None,
     ) -> torch.Tensor | None:
         """Packed varlen attention on NPU via mindiesd attention_forward_varlen.
 
         Returns None (caller falls back to the mask path) when the packed
         contract does not hold; see _resolve_packed_seq_npu.
         """
-        resolved = self._resolve_packed_seq_npu(query, key, extra)
+        resolved = self._resolve_packed_seq_npu(query, key, extra, query_ranges)
         if resolved is None:
             return None
         seq_q, seq_k = resolved
@@ -785,6 +836,7 @@ def _forward_prefix_kv_slice_npu(
         key: torch.Tensor,
         value: torch.Tensor,
         extra: dict,
+        query_ranges: tuple[QueryRange, ...] | None = None,
     ) -> torch.Tensor | None:
         """Mask-free attention by slicing K/V to the valid prefix (zero-copy
         views), then mindiesd attention_forward without a mask.
@@ -805,7 +857,7 @@ def _forward_prefix_kv_slice_npu(
         square, and the output is scaled back. Absent or invalid factor means
         no pre-scaling.
         """
-        resolved = self._resolve_packed_seq_npu(query, key, extra)
+        resolved = self._resolve_packed_seq_npu(query, key, extra, query_ranges)
         if resolved is None:
             return None
         _, seq_k = resolved

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -28,7 +28,8 @@
     RequestScopedCacheDiTRuntime,
 )
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.parallel_state import get_world_group, init_world_group
+from vllm_omni.diffusion.distributed.group_coordinator import GroupCoordinator
+from vllm_omni.diffusion.distributed.parallel_state import get_world_group
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import DenoiseProgressMixin
 from vllm_omni.diffusion.model_loader.diffusers_loader import (
@@ -587,6 +588,29 @@ def __init__(self, rank: int) -> None:
         self.device_group = None
 
 
+def _text_encoder_group_ranks(dit_world_size: int, tp_size: int) -> list[list[int]]:
+    """Tiled encoder TP groups so every DiT rank joins exactly one group.
+
+    ``GroupCoordinator`` asserts that every constructing rank is a member of
+    one of the groups, so with ``tp_size < dit_world_size`` a single
+    ``[0, tp_size)`` group crashes the ranks outside it.  Tiling the DiT world
+    into ``dit_world_size // tp_size`` contiguous groups keeps every rank a
+    member; only the first group (ranks < tp_size) actually encodes, the rest
+    just hold a valid group handle and receive the world-group broadcast of
+    the encoder output.
+    """
+    if tp_size < 1:
+        raise ValueError(f"text_encoder_tp_size must be >= 1, got {tp_size}")
+    if dit_world_size % tp_size != 0:
+        raise ValueError(
+            f"text_encoder_tp_size ({tp_size}) must divide the DiT group size "
+            f"({dit_world_size}): the encoder TP group tiles the DiT world into "
+            "equal contiguous groups, and ranks outside the first group still "
+            "need group membership."
+        )
+    return [list(range(offset, offset + tp_size)) for offset in range(0, dit_world_size, tp_size)]
+
+
 class MiniMaxH3Pipeline(
     nn.Module,
     DenoiseProgressMixin,
@@ -1211,21 +1235,21 @@ def _resolve_task(
     def _build_text_encoder_group(self, text_encoder_tp_size: int) -> Any:
         """Create the encoder tensor-parallel process group.
 
-        The encoder group covers the first ``text_encoder_tp_size`` DiT ranks
-        (the DiT group is always global ranks ``[0, dit_world)``).  Every rank
-        participates in ``new_group`` so the collective completes; ranks
-        outside the group never run encoder collectives.  For a single-rank
-        encoder we return a lightweight placeholder so non-encoder ranks do
-        not need to join a ``GroupCoordinator`` that would assert on ranks
-        outside the group.
+        The encoder TP group tiles the DiT world into contiguous groups of
+        ``text_encoder_tp_size`` ranks (see ``_text_encoder_group_ranks``);
+        only the first group (ranks < tp_size) runs the encoder, and its
+        output is broadcast to every DiT rank over the world group.  Every
+        rank must hold group membership because ``GroupCoordinator`` asserts
+        on ranks outside all groups.  For a single-rank encoder we return a
+        lightweight placeholder instead.
         """
         if text_encoder_tp_size == 1:
             return _SingleRankEncoderGroup(rank=self._dit_rank)
-        ranks = list(range(text_encoder_tp_size))
-        return init_world_group(
-            ranks=ranks,
+        _, _, dit_world = _dit_rank_world()
+        return GroupCoordinator(
+            group_ranks=_text_encoder_group_ranks(dit_world, text_encoder_tp_size),
             local_rank=envs.LOCAL_RANK,
-            backend=current_omni_platform.dist_backend,
+            torch_distributed_backend=current_omni_platform.dist_backend,
         )
 
     def _encoder_group_broadcast_tensor(