diff --git a/tests/diffusion/attention/test_mindie_runtime.py b/tests/diffusion/attention/test_mindie_runtime.py index bd9234a5f..40e51b21d 100644 --- a/tests/diffusion/attention/test_mindie_runtime.py +++ b/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() diff --git a/tests/diffusion/attention/test_rainfusion_plan.py b/tests/diffusion/attention/test_rainfusion_plan.py index 6b364f1f1..f78866c31 100644 --- a/tests/diffusion/attention/test_rainfusion_plan.py +++ b/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 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 4e261b3e2..49197d315 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_contract.py +++ b/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 diff --git a/vllm_omni/diffusion/attention/backends/rainfusion_attn.py b/vllm_omni/diffusion/attention/backends/rainfusion_attn.py index 24a32bd31..d4219fa05 100644 --- a/vllm_omni/diffusion/attention/backends/rainfusion_attn.py +++ b/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)