[Bugfix][RAINFUSION_ATTN] Fix padding-mask rejection on packed [real, pad] layouts and reuse query tail as output padding (#7235)

Signed-off-by: brandneway <gyuan4892@gmail.com>

tests/diffusion/attention/test_mindie_runtime.py

@@ -193,7 +193,11 @@ def test_sparse_quant_uses_public_runtime(runtime, monkeypatch, method):
     monkeypatch.setattr(rainfusion_attn, "_mindiesd_supports_precision", lambda: True)
     q = torch.randn(1, 4224, 2, 64, dtype=torch.bfloat16)
     out = sparse().forward_npu(q, q, q, video_metadata(kv_cache_dtype=method))
-    assert out.shape == q.shape and not out[:, 4096:].any()
+    # The sparse path reuses the query tail as the output padding instead of
+    # zero-filling a fresh buffer; rows past the 4096-row valid prefix are
+    # never read back downstream (the model unpads by cu_seqlens), so pin the
+    # passthrough rather than requiring zeros.
+    assert out.shape == q.shape and torch.equal(out[:, 4096:], q[:, 4096:])
     kwargs = runtime.sparse_attention.call_args.kwargs
     assert (kwargs["precision"], kwargs["sparse_type"], kwargs["inner_precise"]) == (method, "rf_v3", 4)
     runtime.quant_attention.assert_not_called()

tests/diffusion/attention/test_rainfusion_plan.py

@@ -159,6 +159,18 @@ def test_validate_available_accepts_new_mindiesd(monkeypatch):
     RainFusionAttentionBackend.validate_available()
 
 
+def test_prefix_kv_slicing_contract_is_pinned():
+    # Regression pin for the #5543 packed-padding crash: reverting the flag
+    # would pass CI while reintroducing the ValueError for non-64-aligned
+    # requests. Both halves of the contract must hold together — the impl
+    # trims [real, pad] packed tensors to the valid prefix itself (so it may
+    # advertise prefix slicing) and it never reads attn_mask (so it must keep
+    # refusing one), which is what lets the model skip materializing the
+    # padding mask that _assert_metadata_compatible would reject.
+    assert RainFusionAttentionBackend.supports_prefix_kv_slicing is True
+    assert RainFusionAttentionBackend.supports_attention_mask() is False
+
+
 @pytest.mark.parametrize("grid", [(4, 24, 40), (1, 24, 40)])
 def test_short_video_stays_dense(grid):
     assert make_impl()._resolve_plan(make_metadata(grid)) is None

tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py

@@ -1042,6 +1042,56 @@ def test_packed_attention_rejects_backends_without_multi_doc_capability(backend_
         )
 
 
+def test_rainfusion_packed_padding_stays_mask_free_on_unaligned_lengths():
+    """Regression for the #5543 follow-up crash fixed in #7235.
+
+    MiniMax-H3 pads every packed request to a multiple of 64 rows, so any
+    non-64-aligned length reaches this path with used < packed_total. Before
+    RAINFUSION_ATTN advertised supports_prefix_kv_slicing, the model built a
+    padding mask here that _assert_metadata_compatible then rejected at the
+    first sparse layer. This drives the real backend class (not a fake flag)
+    through the packed path; both the sparse dispatch and the dense fallback
+    share this metadata construction and neither ever reads attn_mask.
+    """
+    from vllm_omni.diffusion.attention.backends.rainfusion_attn import (
+        RainFusionAttentionBackend,
+    )
+    from vllm_omni.diffusion.models.minimax_h3.minimax_h3_transformer import (
+        MiniMaxH3Attention,
+    )
+
+    class FakeAttention(torch.nn.Module):
+        attn_backend = RainFusionAttentionBackend
+
+        def __init__(self):
+            super().__init__()
+            self.metadata = None
+
+        def forward(self, query, key, value, metadata):
+            self.metadata = metadata
+            return query
+
+    attention = object.__new__(MiniMaxH3Attention)
+    torch.nn.Module.__init__(attention)
+    attention.attention = FakeAttention()
+    q = torch.randn(8, 2, 4)
+
+    attention._run_packed_attention(
+        q,
+        q,
+        q,
+        # One request: 5 valid rows padded to the 8-row alignment boundary.
+        cu_seqlens=torch.tensor([0, 5, 8], dtype=torch.int32),
+        max_seqlen=5,
+        packed_total=8,
+    )
+
+    metadata = attention.attention.metadata
+    assert metadata is not None
+    assert metadata.attn_mask is None
+    assert metadata.extra["valid_kv_length"] == 5
+
+
 def test_reference_image_resize_contract():
     from PIL import Image
 

vllm_omni/diffusion/attention/backends/rainfusion_attn.py

@@ -135,6 +135,11 @@ class RainFusionPlan:
 class RainFusionAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
     supported_platforms: tuple[str, ...] = ("npu",)
+    # The impl trims [real, pad] packed tensors to the valid prefix with
+    # [:, :used] itself, so models may run the padded layout without
+    # materializing a padding mask (a forward that carries one anyway falls
+    # back dense in _resolve_plan).
+    supports_prefix_kv_slicing: bool = True
 
     @classmethod
     def validate_available(cls) -> None:
@@ -509,7 +514,12 @@ class RainFusionAttentionImpl(AttentionImpl):
 
         logger.info_once("RainFusion uses MindIE-SD sparse attention, precision=%s.", precision)
         used = plan.used_len
-        q, k, v = (tensor[:, :used] for tensor in (query, key, value))
+        # Split query once: the valid prefix drives the kernel, the padding
+        # tail is reused verbatim as the output padding below. All views,
+        # zero-copy. split_with_sizes always yields both parts, including an
+        # empty tail when the sequence needs no padding.
+        q, q_pad = torch.split(query, [used, query.shape[1] - used], dim=1)
+        k, v = key[:, :used], value[:, :used]
         # Ulysses has already gathered the full sequence onto this rank and split
         # the heads, so read the head count off the tensor rather than num_heads.
         common_kwargs: dict[str, object] = {
@@ -543,6 +553,9 @@ class RainFusionAttentionImpl(AttentionImpl):
             out = sparse_attention(q, k, v, **common_kwargs)
         if used == query.shape[1]:
             return out
-        padded = torch.zeros_like(query)
-        padded[:, :used] = out
-        return padded
+        # Only the first ``used`` rows are read back downstream, so the padding
+        # payload is free. Reuse the query tail split off above instead of
+        # zeroing a fresh buffer: a single cat write, no extra allocation, and
+        # always finite values (uninitialized memory could hold NaN bit
+        # patterns).
+        return torch.cat([out, q_pad], dim=1)