diff --git a/tests/diffusion/attention/test_trtllm_attn.py b/tests/diffusion/attention/test_trtllm_attn.py index 76daf9d3b..336f7d037 100644 --- a/tests/diffusion/attention/test_trtllm_attn.py +++ b/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, }, ) diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py index 32e1518dc..000e02f55 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/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(): diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py index 33aec8c0f..63c3a42f1 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_step_execution.py +++ b/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)) diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index 63b3b9cdb..9493e305a 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/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) diff --git a/vllm_omni/diffusion/attention/backends/trtllm_attn.py b/vllm_omni/diffusion/attention/backends/trtllm_attn.py index c63280fc8..967e81c6e 100644 --- a/vllm_omni/diffusion/attention/backends/trtllm_attn.py +++ b/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 " diff --git a/vllm_omni/diffusion/models/minimax_h3/batched_packing.py b/vllm_omni/diffusion/models/minimax_h3/batched_packing.py index e32b61667..4d41bde61 100644 --- a/vllm_omni/diffusion/models/minimax_h3/batched_packing.py +++ b/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 diff --git a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py b/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py index 509190740..6bbc2a11e 100644 --- a/vllm_omni/diffusion/models/minimax_h3/minimax_h3_transformer.py +++ b/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,