[Attention] Refine TRTLLM attention support for MiniMax H3 (#5779)

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
Signed-off-by: lishunyang12 <lishunyang03@gmail.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: lishunyang12 <lishunyang03@gmail.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>

tests/diffusion/attention/test_trtllm_attn.py

@@ -8,6 +8,7 @@ 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.trtllm_attn import TrtllmAttentionImpl
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cuda]
@@ -142,6 +143,9 @@ def test_denoise_progress_mixin(monkeypatch):
     _Pipe().record_denoise_step(4)
     assert ctx.denoise_step_idx == 4 and ctx.denoise_timestep == pytest.approx(0.25)
 
+    _Pipe().record_denoise_step(5, normalized_timestep=0.7)
+    assert ctx.denoise_step_idx == 5 and ctx.denoise_timestep == pytest.approx(0.7)
+
 
 @pytest.mark.parametrize(
     "bk, field",
@@ -187,6 +191,7 @@ def test_mask_free_allowlist_gates_trtllm_default():
 
     assert meta("WanPipeline").attention_mask_free is True
     assert meta("WanVACEPipeline").attention_mask_free is True
+    assert meta("MiniMaxH3Pipeline").attention_mask_free is True
     assert meta("FluxPipeline").attention_mask_free is False
     assert meta(None).attention_mask_free is False
 
@@ -204,6 +209,91 @@ 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)
+    metadata = AttentionMetadata(
+        attn_mask=mask,
+        extra={
+            "cu_seqlens_q": cu_seqlens,
+            "cu_seqlens_k": cu_seqlens,
+            "max_seqlen_q": 6,
+            "max_seqlen_k": 6,
+        },
+    )
+    captured = {}
+
+    def fake_attention(**kwargs):
+        captured.update(kwargs)
+        return torch.zeros_like(kwargs["query"])
+
+    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)
+
+    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["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)
+
+
+def test_packed_metadata_must_be_complete():
+    b, s, h, d = 1, 8, 8, 128
+    q, k, v = (torch.zeros(b, s, h, d) for _ in range(3))
+    metadata = AttentionMetadata(extra={"cu_seqlens_q": torch.tensor([0, s], dtype=torch.int32)})
+
+    with pytest.raises(ValueError, match="Incomplete packed TRTLLM attention metadata"):
+        _impl().forward_cuda(q, k, v, metadata)
+
+
+def test_packed_metadata_must_cover_inputs(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], dtype=torch.int32)
+    metadata = AttentionMetadata(
+        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="must cover all Q/K/V tokens"):
+        _impl().forward_cuda(q, k, v, metadata)
+
+
 def test_propagate_calibration_045_schema(monkeypatch):
     import types
 
@@ -295,6 +385,30 @@ def test_bf16_dense_matches_sdpa():
     assert rel < 0.01, f"BF16 dense rel err {rel:.4f} too high"
 
 
+@requires_trtllm_attn
+def test_bf16_packed_padding_matches_sdpa():
+    torch.manual_seed(0)
+    b, s, used, h, d = 1, 256, 192, 8, 128
+    scale = 1.0 / math.sqrt(d)
+    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,
+        extra={
+            "cu_seqlens_q": cu_seqlens,
+            "cu_seqlens_k": cu_seqlens,
+            "max_seqlen_q": used,
+            "max_seqlen_k": used,
+        },
+    )
+
+    out = _impl().forward_cuda(q, k, v, metadata).float()
+    ref_used = _sdpa_ref(q[:, :used], k[:, :used], v[:, :used], scale)
+    used_rel = (out[:, :used] - ref_used).abs().mean() / ref_used.abs().mean()
+    assert used_rel < 0.01, f"BF16 packed valid-token rel err {used_rel:.4f} too high"
+    assert torch.count_nonzero(out[:, used:]) == 0
+
+
 requires_sage = pytest.mark.skipif(
     not (_has_trtllm_attn() and tg._sage_quantize_fn() is not None),
     reason="requires Blackwell SM100+ GPU with flashinfer >= 0.6.16rc1 (trtllm_sage_attention_quantize)",
@@ -321,6 +435,47 @@ def test_sage_quant_matches_sdpa(dtype_qk):
     assert rel < 0.1, f"SAGE({dtype_qk}) rel err {rel:.4f} too high"
 
 
+@requires_sage
+def test_sage_packed_non_aligned_length_matches_sdpa():
+    torch.manual_seed(0)
+    b, s, used, h, d = 1, 208, 194, 8, 128
+    scale = 1.0 / math.sqrt(d)
+    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,
+        extra={
+            "cu_seqlens_q": cu_seqlens,
+            "cu_seqlens_k": cu_seqlens,
+            "max_seqlen_q": used,
+            "max_seqlen_k": used,
+        },
+    )
+
+    out = _impl(quant={"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}).forward_cuda(q, k, v, metadata)
+    assert torch.isfinite(out).all()
+    ref = _sdpa_ref(q[:, :used], k[:, :used], v[:, :used], scale)
+    rel = (out[:, :used].float() - ref).abs().mean() / ref.abs().mean()
+    assert rel < 0.2, f"SAGE packed valid-token rel err {rel:.4f} too high"
+    assert torch.count_nonzero(out[:, used:]) == 0
+
+
+@requires_sage
+def test_sage_short_sequence_uses_dense_kernel():
+    torch.manual_seed(0)
+    b, s, h, d = 1, 14, 8, 128
+    scale = 1.0 / math.sqrt(d)
+    q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3))
+
+    impl = _impl(quant={"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16})
+    impl._sage_quantize_fn = lambda *args, **kwargs: pytest.fail("short sequence must not use SAGE")
+    out = impl.forward_cuda(q, k, v)
+    ref = _sdpa_ref(q, k, v, scale)
+    rel = (out.float() - ref).abs().mean() / ref.abs().mean()
+    assert torch.isfinite(out).all()
+    assert rel < 0.01, f"SAGE short-sequence dense fallback rel err {rel:.4f} too high"
+
+
 @requires_trtllm_attn
 def test_skip_tiny_threshold_matches_dense():
     torch.manual_seed(0)