[perf] Add mask-free TRTLLM packed-padding path (#6542)

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>

tests/diffusion/attention/test_trtllm_attn.py

@@ -1,14 +1,18 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import math
+from unittest.mock import Mock
 
 import pytest
 import torch
 import torch.nn.functional as F
 
 from vllm_omni.diffusion.attention.backends import trtllm_attn as tg
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionMetadata,
+    PackedPaddingMetadata,
+)
 from vllm_omni.diffusion.attention.backends.trtllm_attn import TrtllmAttentionImpl
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
@@ -35,6 +39,15 @@ def _impl(**bk):
     return TrtllmAttentionImpl(8, 128, 1.0 / math.sqrt(128), causal=False, num_kv_heads=8, backend_kwargs=bk)
 
 
+def _packed_padding(cu_seqlens_q, cu_seqlens_k, q_length, kv_length):
+    return PackedPaddingMetadata(
+        q_length=q_length,
+        kv_length=kv_length,
+        cu_seqlens_q=cu_seqlens_q[:2],
+        cu_seqlens_k=cu_seqlens_k[:2],
+    )
+
+
 def test_skip_config_pure_resolution():
     from vllm_omni.diffusion.attention.backends.trtllm_attn import SkipSoftmaxConfig
 
@@ -75,7 +88,7 @@ def test_quant_quantize_requires_flashinfer_routine(monkeypatch):
 def test_quant_quantize_calls_routine_and_shapes_sfs():
     from vllm_omni.diffusion.attention.backends.trtllm_attn import QuantConfig
 
-    captured = {}
+    captured: dict[str, object] = {}
 
     def fake_quantize(q, k, v, q_block_size, k_block_size, qk_quant_dtype):
         captured.update(q_block_size=q_block_size, k_block_size=k_block_size, qk_quant_dtype=qk_quant_dtype)
@@ -209,62 +222,80 @@ def test_masked_layer_raises():
         _impl().forward_cuda(q, k, v, _Meta())
 
 
-def test_packed_metadata_trims_padding_mask(monkeypatch):
-    b, s, h, d = 1, 8, 8, 128
-    q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
-    mask = torch.tensor([[True, True, True, True, True, True, False, False]])
-    cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32)
+def test_generic_packed_metadata_preserves_ragged_batch(monkeypatch):
+    b, h, d = 1, 8, 128
+    q = torch.zeros(b, 5, h, d)
+    k, v = (torch.zeros(b, 7, h, d) for _ in range(2))
     metadata = AttentionMetadata(
-        attn_mask=mask,
         extra={
-            "cu_seqlens_q": cu_seqlens,
-            "cu_seqlens_k": cu_seqlens,
-            "max_seqlen_q": 6,
-            "max_seqlen_k": 6,
+            "cu_seqlens_q": torch.tensor([0, 2, 5], dtype=torch.int32),
+            "cu_seqlens_k": torch.tensor([0, 3, 7], dtype=torch.int32),
+            "max_seqlen_q": 3,
+            "max_seqlen_k": 4,
+        }
+    )
+    fake_attention = Mock(return_value=torch.zeros_like(q.reshape(5, h, d)))
+
+    monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
+    monkeypatch.setattr(tg, "trtllm_ragged_attention_deepseek", fake_attention)
+    monkeypatch.setattr(TrtllmAttentionImpl, "_get_workspace", classmethod(lambda cls, device: torch.empty(0)))
+
+    _impl().forward_cuda(q, k, v, metadata)
+    captured = fake_attention.call_args.kwargs
+
+    assert captured["query"].shape[0] == 5
+    assert captured["key"].shape[0] == 7
+    assert captured["value"].shape[0] == 7
+    assert captured["batch_size"] == 2
+    assert captured["seq_lens"].tolist() == [3, 4]
+    assert tg.TrtllmAttentionBackend.supports_multi_doc_packed_varlen()
+
+
+def test_packed_padding_metadata_trims_without_mask(monkeypatch):
+    q_len, kv_len = 8, 9
+    valid_q_tokens, valid_kv_tokens = 5, 6
+    b, h, d = 1, 8, 128
+    q = torch.zeros(b, q_len, h, d)
+    k, v = (torch.zeros(b, kv_len, h, d) for _ in range(2))
+    cu_seqlens_q = torch.tensor([0, valid_q_tokens, q_len], dtype=torch.int32)
+    cu_seqlens_k = torch.tensor([0, valid_kv_tokens, kv_len], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        packed_padding=_packed_padding(
+            cu_seqlens_q,
+            cu_seqlens_k,
+            valid_q_tokens,
+            valid_kv_tokens,
+        ),
+        extra={
+            "cu_seqlens_q": cu_seqlens_q,
+            "cu_seqlens_k": cu_seqlens_k,
+            "max_seqlen_q": valid_q_tokens,
+            "max_seqlen_k": valid_kv_tokens,
+            "valid_kv_length": valid_kv_tokens,
         },
     )
-    captured = {}
-
-    def fake_attention(**kwargs):
-        captured.update(kwargs)
-        return torch.zeros_like(kwargs["query"])
+    fake_attention = Mock(return_value=torch.ones(valid_q_tokens, h, d))
 
     monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
     monkeypatch.setattr(tg, "trtllm_ragged_attention_deepseek", fake_attention)
     monkeypatch.setattr(TrtllmAttentionImpl, "_get_workspace", classmethod(lambda cls, device: torch.empty(0)))
 
     out = _impl().forward_cuda(q, k, v, metadata)
+    captured = fake_attention.call_args.kwargs
 
+    assert tg.TrtllmAttentionBackend.supports_packed_mask_free()
     assert out.shape == q.shape
-    assert captured["query"].shape[0] == 6
-    assert captured["key"].shape[0] == 6
-    assert captured["value"].shape[0] == 6
+    assert captured["query"].shape[0] == valid_q_tokens
+    assert captured["key"].shape[0] == valid_kv_tokens
+    assert captured["value"].shape[0] == valid_kv_tokens
     assert captured["batch_size"] == 1
-    assert captured["seq_lens"].tolist() == [6]
-    assert captured["cum_seq_lens_q"].tolist() == [0, 6]
-    assert captured["cum_seq_lens_kv"].tolist() == [0, 6]
-    assert captured["max_q_len"] == 6
-    assert captured["max_kv_len"] == 6
-    assert torch.count_nonzero(out[:, 6:]) == 0
-
-
-def test_packed_metadata_rejects_non_prefix_mask(monkeypatch):
-    b, s, h, d = 1, 8, 8, 128
-    q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
-    cu_seqlens = torch.tensor([0, 6, 8], dtype=torch.int32)
-    metadata = AttentionMetadata(
-        attn_mask=torch.tensor([[True, True, False, True, True, True, False, False]]),
-        extra={
-            "cu_seqlens_q": cu_seqlens,
-            "cu_seqlens_k": cu_seqlens,
-            "max_seqlen_q": 6,
-            "max_seqlen_k": 6,
-        },
-    )
-    monkeypatch.setattr(tg, "HAS_FLASHINFER", True)
-
-    with pytest.raises(ValueError, match="prefix-valid"):
-        _impl().forward_cuda(q, k, v, metadata)
+    assert captured["seq_lens"].tolist() == [valid_kv_tokens]
+    assert captured["cum_seq_lens_q"].tolist() == [0, valid_q_tokens]
+    assert captured["cum_seq_lens_kv"].tolist() == [0, valid_kv_tokens]
+    assert captured["max_q_len"] == valid_q_tokens
+    assert captured["max_kv_len"] == valid_kv_tokens
+    assert torch.count_nonzero(out[:, :valid_q_tokens] != 1) == 0
+    assert torch.count_nonzero(out[:, valid_q_tokens:]) == 0
 
 
 def test_packed_metadata_must_be_complete():
@@ -393,12 +424,13 @@ def test_bf16_packed_padding_matches_sdpa():
     q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3))
     cu_seqlens = torch.tensor([0, used, s], dtype=torch.int32, device="cuda")
     metadata = AttentionMetadata(
-        attn_mask=torch.arange(s, device="cuda")[None] < used,
+        packed_padding=_packed_padding(cu_seqlens, cu_seqlens, used, used),
         extra={
             "cu_seqlens_q": cu_seqlens,
             "cu_seqlens_k": cu_seqlens,
             "max_seqlen_q": used,
             "max_seqlen_k": used,
+            "valid_kv_length": used,
         },
     )
 
@@ -443,12 +475,13 @@ def test_sage_packed_non_aligned_length_matches_sdpa():
     q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3))
     cu_seqlens = torch.tensor([0, used, s], dtype=torch.int32, device="cuda")
     metadata = AttentionMetadata(
-        attn_mask=torch.arange(s, device="cuda")[None] < used,
+        packed_padding=_packed_padding(cu_seqlens, cu_seqlens, used, used),
         extra={
             "cu_seqlens_q": cu_seqlens,
             "cu_seqlens_k": cu_seqlens,
             "max_seqlen_q": used,
             "max_seqlen_k": used,
+            "valid_kv_length": used,
         },
     )
 

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -774,6 +774,12 @@ 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
+    packed_padding = metadata.packed_padding
+    assert packed_padding is not None
+    assert packed_padding.q_length == 5
+    assert packed_padding.kv_length == 5
+    assert packed_padding.cu_seqlens_q.tolist() == [0, 5]
+    assert packed_padding.cu_seqlens_k.tolist() == [0, 5]
 
 
 def test_packed_attention_keeps_padding_mask_for_other_backends():

tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """MiniMax H3 step-wise execution (continuous batching) contract tests.
 
 These run on CPU against a stand-in DiT, so they cover the packing, the
@@ -31,6 +31,8 @@ class _SegmentMeanModel:
         x = kwargs["x"][0]
         audio_x = kwargs["audio_x"][0]
         bounds = kwargs["packed_seq_params"]["cu_seqlens_q"].tolist()
+        if kwargs["packed_seq_params"].get("num_requests", 1) > 1:
+            assert all(start < stop for start, stop in zip(bounds[:-1], bounds[1:]))
         img_pos = kwargs["img_pos_info"]["position_ids"]
         audio_pos = kwargs["audio_pos_info"]["position_ids"]
 
@@ -158,9 +160,9 @@ def test_batched_step_execution_matches_independent_requests():
 
     model = _SegmentMeanModel()
     specs = [
-        dict(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=6),
         # Exactly 64 packed rows exercises the no-padding-tail boundary case.
         dict(text_len=46, latent_t=7, latent_h=2, latent_w=4, audio_t=2, seed=7),
+        dict(text_len=9, latent_t=2, latent_h=4, latent_w=6, audio_t=3, seed=6),
     ]
     # Different step counts, so the batch composition changes mid-flight.
     schedules = [(_sigmas(6, 12.0), _sigmas(6, 3.0)), (_sigmas(4, 12.0), _sigmas(4, 3.0))]
@@ -178,7 +180,7 @@ def test_batched_step_execution_matches_independent_requests():
     states = []
     for index, (spec, (sigmas_video, sigmas_audio)) in enumerate(zip(specs, schedules)):
         branch, video_rows, audio_rows = _make_branch(**spec)
-        if index == 1:
+        if index == 0:
             assert branch.used_len == branch.seq_len
         states.append(_make_state(f"req-{index}", model, branch, video_rows, audio_rows, sigmas_video, sigmas_audio))
 

vllm_omni/diffusion/attention/backends/abstract.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
@@ -27,12 +27,13 @@ class AttentionBackend(ABC):
 
     @classmethod
     def supports_packed_mask_free(cls) -> bool:
-        """Whether packed attention never reads attn_mask on this platform.
+        """Whether [real, pad] packed layouts can run without attn_mask.
 
         When True, models that pack a [real, pad] two-document layout and
-        carry cu_seqlens/max_seqlen in ``AttentionMetadata.extra`` may skip
-        constructing the padding mask entirely. Backends whose mask-free
-        behavior is platform-dependent must check current_omni_platform.
+        provide ``AttentionMetadata.packed_padding`` alongside the packed
+        cu_seqlens/max_seqlen metadata may skip constructing the padding mask
+        entirely. Backends whose mask-free behavior is platform-dependent must
+        check current_omni_platform.
         """
         return False
 
@@ -141,6 +142,21 @@ class VideoTokenLayout:
     latent_grid: tuple[int, int, int]
 
 
+@dataclass(frozen=True, slots=True)
+class PackedPaddingMetadata:
+    """Producer-validated mask-free view of padding in a [real, pad] packing.
+
+    The cumulative-length tensors are canonical two-element ``[0, length]``
+    views. Consumers may use them without reading device scalars because the
+    producer owns the packing.
+    """
+
+    q_length: int
+    kv_length: int
+    cu_seqlens_q: torch.Tensor
+    cu_seqlens_k: torch.Tensor
+
+
 @dataclass
 class AttentionMetadata:
     attn_mask: torch.Tensor | None = None
@@ -188,6 +204,10 @@ class AttentionMetadata:
     # locality (block-sparse selection, tiled masks). Dense backends ignore it.
     video_layout: VideoTokenLayout | None = None
 
+    # Canonical mask-free view of structural suffix padding. Backends that do
+    # not advertise supports_packed_mask_free ignore it.
+    packed_padding: PackedPaddingMetadata | None = None
+
 
 T = TypeVar("T", bound=AttentionMetadata)
 

vllm_omni/diffusion/attention/backends/trtllm_attn.py

@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import functools
 import inspect
 import math
 from dataclasses import dataclass, replace
+from typing import NamedTuple, cast
 
 import torch
 from vllm.logger import init_logger
@@ -13,6 +14,7 @@ from vllm_omni.diffusion.attention.backends.abstract import (
     AttentionBackend,
     AttentionImpl,
     AttentionMetadata,
+    PackedPaddingMetadata,
 )
 
 logger = init_logger(__name__)
@@ -42,8 +44,9 @@ class SkipSoftmaxConfig:
         return cls(
             threshold=_validate_control(bk.get("skip_softmax_threshold"), "skip_softmax_threshold", 0.0, None),
             target_sparsity=_validate_control(bk.get("target_sparsity"), "target_sparsity", 0.0, 1.0),
-            disabled_until_timestep=_validate_control(
-                bk.get("disabled_until_timestep", 0.0), "disabled_until_timestep", 0.0, 1.0
+            disabled_until_timestep=cast(
+                float,
+                _validate_control(bk.get("disabled_until_timestep", 0.0), "disabled_until_timestep", 0.0, 1.0),
             ),
         )
 
@@ -131,7 +134,7 @@ class QuantConfig:
         return self.dtype_qk is not None
 
     def quantize(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, quantize_fn):
-        qk_quant_dtype = _QK_QUANT_DTYPES[self.dtype_qk]
+        qk_quant_dtype = _QK_QUANT_DTYPES[cast(str, self.dtype_qk)]
         q_q, k_q, v_q, q_sfs, k_sfs, v_sfs = quantize_fn(
             q,
             k,
@@ -145,6 +148,18 @@ class QuantConfig:
         return q_q, k_q, v_q, sage_attn_sfs, num_elts_per_sage_attn_blk
 
 
+class _PackedLayout(NamedTuple):
+    """Active token ranges and kernel metadata for a flattened ragged batch."""
+
+    q_tokens: int
+    kv_tokens: int
+    cu_seqlens_q: torch.Tensor
+    cu_seqlens_kv: torch.Tensor
+    batch_size: int
+    seq_lens: torch.Tensor
+    known_min_kv_length: int | None = None
+
+
 def _workspace_bytes() -> int:
     import vllm.envs as envs
 
@@ -154,6 +169,14 @@ def _workspace_bytes() -> int:
 class TrtllmAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
 
+    @classmethod
+    def supports_packed_mask_free(cls) -> bool:
+        return True
+
+    @classmethod
+    def supports_multi_doc_packed_varlen(cls) -> bool:
+        return True
+
     @staticmethod
     def get_supported_head_sizes() -> list[int]:
         return [128]
@@ -250,6 +273,97 @@ class TrtllmAttentionImpl(AttentionImpl):
             cls._workspace = ws
         return ws
 
+    @staticmethod
+    def _prepare_packed_padding_layout(
+        physical_batch: int,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        packed_padding: PackedPaddingMetadata,
+        extra: dict,
+    ) -> _PackedLayout:
+        """Prepare a [real, pad] layout from an explicit host-side boundary."""
+        valid_q_tokens = packed_padding.q_length
+        valid_kv_tokens = packed_padding.kv_length
+        max_q_len = extra["max_seqlen_q"]
+        max_kv_len = extra["max_seqlen_k"]
+        lengths = (valid_q_tokens, valid_kv_tokens, max_q_len, max_kv_len)
+        if any(isinstance(value, bool) or not isinstance(value, int) for value in lengths):
+            raise ValueError("Mask-free packed TRTLLM attention lengths must be Python integers")
+        if physical_batch != 1:
+            raise ValueError("Mask-free packed TRTLLM attention requires a single packed batch")
+
+        total_q_tokens = q.shape[0]
+        total_kv_tokens = k.shape[0]
+        if not 0 < valid_q_tokens <= total_q_tokens:
+            raise ValueError(
+                "PackedPaddingMetadata.q_length must be within the packed Q sequence, "
+                f"got {valid_q_tokens} for length {total_q_tokens}"
+            )
+        if not 0 < valid_kv_tokens <= total_kv_tokens:
+            raise ValueError(
+                "PackedPaddingMetadata.kv_length must be within the packed K/V sequence, "
+                f"got {valid_kv_tokens} for length {total_kv_tokens}"
+            )
+        if max_q_len != valid_q_tokens or max_kv_len != valid_kv_tokens:
+            raise ValueError("Mask-free packed TRTLLM attention lengths must match max_seqlen_q/k")
+        published_kv_length = extra.get("valid_kv_length")
+        if published_kv_length is not None:
+            if isinstance(published_kv_length, bool) or not isinstance(published_kv_length, int):
+                raise ValueError("valid_kv_length must be a Python integer")
+            if published_kv_length != valid_kv_tokens:
+                raise ValueError("PackedPaddingMetadata.kv_length must match valid_kv_length")
+
+        cu_seq_lens_q = packed_padding.cu_seqlens_q
+        cu_seq_lens_kv = packed_padding.cu_seqlens_k
+        if cu_seq_lens_q.dtype != torch.int32 or cu_seq_lens_kv.dtype != torch.int32:
+            raise ValueError("Mask-free packed TRTLLM attention requires int32 cu_seqlens")
+        if cu_seq_lens_q.device != q.device or cu_seq_lens_kv.device != k.device:
+            raise ValueError("Mask-free packed TRTLLM attention metadata must be on the Q/K device")
+        if cu_seq_lens_q.shape != (2,) or cu_seq_lens_kv.shape != (2,):
+            raise ValueError("Mask-free packed TRTLLM attention requires canonical two-element cu_seqlens")
+        return _PackedLayout(
+            q_tokens=valid_q_tokens,
+            kv_tokens=valid_kv_tokens,
+            cu_seqlens_q=cu_seq_lens_q,
+            cu_seqlens_kv=cu_seq_lens_kv,
+            batch_size=1,
+            # Canonical packed-padding metadata is [0, valid_kv_tokens], so this slice
+            # is the one-element sequence-length view expected by the kernel.
+            seq_lens=cu_seq_lens_kv[1:],
+            known_min_kv_length=valid_kv_tokens,
+        )
+
+    @staticmethod
+    def _prepare_generic_packed_layout(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        cu_seq_lens_q: torch.Tensor,
+        cu_seq_lens_kv: torch.Tensor,
+    ) -> _PackedLayout:
+        """Validate and normalize a general ragged batch."""
+        if int(cu_seq_lens_q[-1].item()) != q.shape[0] or int(cu_seq_lens_kv[-1].item()) != k.shape[0]:
+            raise ValueError("Packed TRTLLM attention cu_seqlens must cover all Q/K/V tokens")
+
+        # Packed producers may retain a zero-length trailing segment when the
+        # input is already aligned. It is not a real sequence and is removed
+        # before deriving the batch size and per-sequence lengths.
+        while (
+            cu_seq_lens_q.numel() > 2
+            and int(cu_seq_lens_q[-1].item()) == int(cu_seq_lens_q[-2].item())
+            and int(cu_seq_lens_kv[-1].item()) == int(cu_seq_lens_kv[-2].item())
+        ):
+            cu_seq_lens_q = cu_seq_lens_q[:-1]
+            cu_seq_lens_kv = cu_seq_lens_kv[:-1]
+
+        return _PackedLayout(
+            q_tokens=q.shape[0],
+            kv_tokens=k.shape[0],
+            cu_seqlens_q=cu_seq_lens_q,
+            cu_seqlens_kv=cu_seq_lens_kv,
+            batch_size=cu_seq_lens_q.numel() - 1,
+            seq_lens=(cu_seq_lens_kv[1:] - cu_seq_lens_kv[:-1]).to(dtype=torch.int32).contiguous(),
+        )
+
     def forward_cuda(
         self,
         query: torch.Tensor,
@@ -266,12 +380,16 @@ class TrtllmAttentionImpl(AttentionImpl):
         has_packed_metadata = len(present_packed_keys) == len(packed_keys)
 
         attn_mask = getattr(attn_metadata, "attn_mask", None) if attn_metadata is not None else None
-        if attn_mask is not None and not has_packed_metadata:
+        packed_padding = getattr(attn_metadata, "packed_padding", None) if attn_metadata is not None else None
+        if packed_padding is not None and not isinstance(packed_padding, PackedPaddingMetadata):
+            raise ValueError("packed_padding must be PackedPaddingMetadata")
+        if packed_padding is not None and not has_packed_metadata:
+            raise ValueError("PackedPaddingMetadata requires complete packed TRTLLM attention metadata")
+        if attn_mask is not None:
             raise ValueError(
-                "TRTLLM_ATTN does not support attn_mask (Wan sets one only under SP with "
-                "mask_sp_padding=True and a non-divisible seq len). Either set "
-                "parallel_config.mask_sp_padding=False, or select a mask-capable backend "
-                "(e.g. CUDNN_ATTN / TORCH_SDPA)."
+                "TRTLLM_ATTN does not support attn_mask. Represent structural suffix padding "
+                "with packed-padding metadata, or select a mask-capable backend such as "
+                "CUDNN_ATTN or TORCH_SDPA."
             )
 
         if not HAS_FLASHINFER:
@@ -282,7 +400,6 @@ class TrtllmAttentionImpl(AttentionImpl):
 
         physical_batch, q_len, num_q_heads, head_dim = query.shape
         kv_len, num_kv_heads = key.shape[1], key.shape[2]
-
         device = query.device
 
         q = query.reshape(physical_batch * q_len, num_q_heads, head_dim).contiguous()
@@ -290,6 +407,7 @@ class TrtllmAttentionImpl(AttentionImpl):
         v = value.reshape(physical_batch * kv_len, num_kv_heads, head_dim).contiguous()
         output_tokens = q.shape[0]
 
+        known_min_kv_length: int | None = kv_len
         if has_packed_metadata:
             cu_seq_lens_q = extra["cu_seqlens_q"]
             cu_seq_lens_kv = extra["cu_seqlens_k"]
@@ -297,48 +415,36 @@ class TrtllmAttentionImpl(AttentionImpl):
                 raise ValueError("Packed TRTLLM attention cu_seqlens tensors must be one-dimensional")
             if cu_seq_lens_q.numel() != cu_seq_lens_kv.numel() or cu_seq_lens_q.numel() < 2:
                 raise ValueError("Packed TRTLLM attention requires matching non-empty Q and KV sequence batches")
-            if int(cu_seq_lens_q[-1].item()) != q.shape[0] or int(cu_seq_lens_kv[-1].item()) != k.shape[0]:
-                raise ValueError("Packed TRTLLM attention cu_seqlens must cover all Q/K/V tokens")
 
-            # Drop trailing empty sequences retained by packed producers when
-            # alignment adds no padding tokens.
-            while (
-                cu_seq_lens_q.numel() > 2
-                and int(cu_seq_lens_q[-1].item()) == int(cu_seq_lens_q[-2].item())
-                and int(cu_seq_lens_kv[-1].item()) == int(cu_seq_lens_kv[-2].item())
-            ):
-                cu_seq_lens_q = cu_seq_lens_q[:-1]
-                cu_seq_lens_kv = cu_seq_lens_kv[:-1]
+            # PackedPaddingMetadata is an optional producer-published shortcut
+            # for structural suffix padding. Inputs without it retain the
+            # general cu_seqlens path for arbitrary ragged Q/KV batches.
+            if packed_padding is not None:
+                packed_layout = self._prepare_packed_padding_layout(
+                    physical_batch,
+                    q,
+                    k,
+                    packed_padding,
+                    extra,
+                )
+            else:
+                # Generic metadata can describe any number of real sequences,
+                # not only a single [real, pad] pair.
+                packed_layout = self._prepare_generic_packed_layout(
+                    q,
+                    k,
+                    cu_seq_lens_q,
+                    cu_seq_lens_kv,
+                )
 
-            # A packed prefix mask denotes structural suffix padding, not another
-            # document. Exclude it before quantization and attention; in particular,
-            # SAGE block quantization must not span the valid/padding boundary.
-            if attn_mask is not None:
-                if q.shape[0] != k.shape[0] or attn_mask.numel() != q.shape[0]:
-                    raise ValueError("Packed TRTLLM prefix masks require equal Q/K lengths")
-                flat_mask = attn_mask.reshape(-1).to(dtype=torch.bool)
-                valid_tokens = int(flat_mask.sum().item())
-                expected_mask = torch.arange(flat_mask.numel(), device=device) < valid_tokens
-                if not torch.equal(flat_mask, expected_mask):
-                    raise ValueError("TRTLLM_ATTN only supports prefix-valid packed attention masks")
-
-                q_boundary = (cu_seq_lens_q == valid_tokens).nonzero(as_tuple=False)
-                kv_boundary = (cu_seq_lens_kv == valid_tokens).nonzero(as_tuple=False)
-                if q_boundary.numel() != 1 or kv_boundary.numel() != 1:
-                    raise ValueError("Packed TRTLLM prefix length must be a Q/K sequence boundary")
-                q_end = int(q_boundary.item()) + 1
-                kv_end = int(kv_boundary.item()) + 1
-                if q_end != kv_end:
-                    raise ValueError("Packed TRTLLM prefix masks require matching Q/K sequence batches")
-
-                q = q[:valid_tokens]
-                k = k[:valid_tokens]
-                v = v[:valid_tokens]
-                cu_seq_lens_q = cu_seq_lens_q[:q_end].contiguous()
-                cu_seq_lens_kv = cu_seq_lens_kv[:kv_end].contiguous()
-
-            batch = cu_seq_lens_q.numel() - 1
-            seq_lens = (cu_seq_lens_kv[1:] - cu_seq_lens_kv[:-1]).to(dtype=torch.int32).contiguous()
+            q = q[: packed_layout.q_tokens]
+            k = k[: packed_layout.kv_tokens]
+            v = v[: packed_layout.kv_tokens]
+            cu_seq_lens_q = packed_layout.cu_seqlens_q
+            cu_seq_lens_kv = packed_layout.cu_seqlens_kv
+            batch = packed_layout.batch_size
+            seq_lens = packed_layout.seq_lens
+            known_min_kv_length = packed_layout.known_min_kv_length
             max_q_len = int(extra["max_seqlen_q"])
             max_kv_len = int(extra["max_seqlen_k"])
         else:
@@ -361,7 +467,13 @@ class TrtllmAttentionImpl(AttentionImpl):
         sage_kwargs: dict = {}
         # The SAGE kernel requires every KV sequence to contain at least one full
         # quantization block. Small auxiliary attention sites use the dense kernel.
-        use_sage = self.quant.enabled and bool(torch.all(seq_lens >= self.quant.k_block_size).item())
+        use_sage = False
+        if self.quant.enabled:
+            if known_min_kv_length is not None:
+                sage_lengths_supported = known_min_kv_length >= self.quant.k_block_size
+            else:
+                sage_lengths_supported = bool(torch.all(seq_lens >= self.quant.k_block_size).item())
+            use_sage = sage_lengths_supported
         if self.quant.enabled and not use_sage:
             message = (
                 f"TRTLLM_ATTN SAGE quantization is configured for attention role {self.role!r}, but at least one "

vllm_omni/diffusion/models/minimax_h3/batched_packing.py

@@ -1,12 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Batched multi-request packing for MiniMax H3 step-wise execution.
 
 Request mode forwards exactly one packed sequence per denoise step. Step mode
 (continuous batching) may hold several requests at once, so this module
 concatenates their packed layouts into a single sequence and rebuilds
-``cu_seqlens`` with one document per request plus that request's
-alignment-padding tail. Attention therefore never crosses a request boundary,
-while every request keeps its own RoPE coordinates, token tags, and timesteps.
+``cu_seqlens`` with one document per request plus any nonempty alignment-padding
+tail. Attention therefore never crosses a request boundary, while every request
+keeps its own RoPE coordinates, token tags, and timesteps.
 A batch of one is not rebuilt at all -- it delegates to the request-mode layout.
 
 Row order is the request order, so the concatenated video velocity returned by
@@ -93,12 +94,13 @@ def minimax_h3_batched_forward_kwargs(
         timesteps[audio_pos[branch.audio_update_mask_dev]] = float(t_audio[index])
         timesteps[audio_pos[~branch.audio_update_mask_dev]] = float(audio_ref_cond_timesteps[index])
 
-        # One document for the request's real rows and one for its padding tail,
-        # matching the single-request layout. The tail bound is emitted even when
-        # the 64-row alignment left it empty, so that a request always
-        # contributes exactly two bounds regardless of its length.
-        cu_bounds.append(seq_offset + branch.used_len)
-        cu_bounds.append(seq_offset + branch.seq_len)
+        # Empty documents are omitted because not every varlen kernel accepts
+        # repeated interior boundaries.
+        real_end = seq_offset + branch.used_len
+        padded_end = seq_offset + branch.seq_len
+        cu_bounds.append(real_end)
+        if padded_end != real_end:
+            cu_bounds.append(padded_end)
         max_seqlen = max(max_seqlen, branch.used_len, branch.seq_len - branch.used_len)
 
         text_offset += branch.text_len

vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py

@@ -27,7 +27,11 @@ from vllm.model_executor.layers.linear import (
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, VideoTokenLayout
+from vllm_omni.diffusion.attention.backends.abstract import (
+    AttentionMetadata,
+    PackedPaddingMetadata,
+    VideoTokenLayout,
+)
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import (
     indexed_gate,
@@ -472,6 +476,7 @@ class MiniMaxH3Attention(nn.Module):
                 f"max_seqlen must be within the packed sequence, got {max_seqlen} for length {packed_total}"
             )
         attn_mask = None
+        mask_free_packed_padding = False
         if num_requests > 1:
             # A step-mode batch packs one document per request, so its valid
             # rows are block-diagonal rather than a prefix: neither a KV prefix
@@ -498,14 +503,25 @@ class MiniMaxH3Attention(nn.Module):
             # supports_packed_mask_free: backend consumes the packed metadata
             # without ever reading attn_mask (CUDA packed varlen, NPU
             # npu_attn_varlen opt-in with its own fallback rebuild).
-            no_mask = not getattr(self.attention, "use_ring", False) and (
-                self.attention.attn_backend.supports_prefix_kv_slicing
-                or self.attention.attn_backend.supports_packed_mask_free()
+            use_ring = getattr(self.attention, "use_ring", False)
+            mask_free_packed_padding = not use_ring and self.attention.attn_backend.supports_packed_mask_free()
+            no_mask = not use_ring and (
+                self.attention.attn_backend.supports_prefix_kv_slicing or mask_free_packed_padding
             )
             if used < packed_total and not no_mask:
                 attn_mask = torch.arange(packed_total, device=q.device)[None] < used
         metadata = AttentionMetadata(
             attn_mask=attn_mask,
+            packed_padding=(
+                PackedPaddingMetadata(
+                    q_length=used,
+                    kv_length=used,
+                    cu_seqlens_q=cu_seqlens[:2],
+                    cu_seqlens_k=cu_seqlens[:2],
+                )
+                if mask_free_packed_padding
+                else None
+            ),
             extra={
                 "cu_seqlens_q": cu_seqlens,
                 "cu_seqlens_k": cu_seqlens,
@@ -575,9 +591,9 @@ class MiniMaxH3Attention(nn.Module):
                 self.q_norm.variance_epsilon,
             )
 
-        # Each request contributes a document for its rows plus one for its
-        # alignment padding. Local/Ulysses backends unpad it, while Ring keeps
-        # aligned rows for fixed-size P2P buffers.
+        # 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.
         out = self._run_packed_attention(
             q,
             k,