[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,
         },
     )