[Bugfix] rainfusion support all shape video (#6000)

Signed-off-by: Fan <fan@FandeMacBook-Air.local>
Co-authored-by: Fan <fan@FandeMacBook-Air.local>

tests/diffusion/attention/test_rainfusion_plan.py

@@ -5,10 +5,9 @@
 
 The grids here are the ones MiniMax-H3 FL2VA actually produces for an 8.7s clip:
 1280x768 gives a 62x24x40 latent grid whose 59520 video rows land on the 128-row
-kernel block, and 1344x768 gives 62x24x42 whose 62496 rows straddle it. Only the
-first can run sparse: rf_v2 tiles the sequence as a whole but pools its block
-mask from the video and prefix segments separately, so the two only agree on the
-block count when the video segment ends on a boundary.
+kernel block, and 1344x768 gives 62x24x42 whose 62496 rows do not. Both are
+handed to MindIE-SD: its rf_v2 preprocessing promotes an irregular video tail
+to the always-kept prefix segment before generating the block mask.
 """
 
 import dataclasses
@@ -57,10 +56,13 @@ def test_block_aligned_video_segment_runs_sparse():
     assert plan.latent_shape == list(ALIGNED_GRID)
 
 
-def test_misaligned_video_segment_stays_dense():
-    # Realigning the sequence would mean padding it, and rf_v2 ignores attn_mask:
-    # pad keys would take a share of every real query's softmax denominator.
-    assert make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID)) is None
+def test_misaligned_video_segment_runs_sparse_with_updated_mindiesd():
+    plan = make_impl()._resolve_plan(make_metadata(MISALIGNED_GRID))
+
+    assert plan is not None
+    assert plan.prefix_len == PREFIX_ROWS
+    assert plan.used_len == PREFIX_ROWS + 62496
+    assert plan.latent_shape == list(MISALIGNED_GRID)
 
 
 @pytest.mark.parametrize("prefix_len", [1, 127, _BLOCK_SIZE, 710, 900])
@@ -155,3 +157,28 @@ def test_fully_populated_mask_reproduces_dense_attention():
     ).transpose(1, 2)
     error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
     assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"
+
+
+@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="rf_v2 runs on Ascend NPU only.")
+def test_irregular_video_tail_reproduces_dense_attention_with_updated_mindiesd():
+    """vLLM forwards the full irregular grid; MindIE-SD owns its tail handling."""
+    torch.manual_seed(1)
+    prefix_len, grid = 384, (4, 16, 10)
+    video_len = grid[0] * grid[1] * grid[2]
+    used = prefix_len + video_len
+    heads, head_dim = 4, 128
+    q, k, v = (torch.randn(1, used, heads, head_dim, dtype=torch.bfloat16, device="npu") for _ in range(3))
+
+    impl = make_impl()
+    impl.rainfusion = dataclasses.replace(impl.rainfusion, sparsity=0.0)
+    plan = RainFusionPlan(prefix_len=prefix_len, used_len=used, latent_shape=list(grid))
+    out = impl._forward_sparse_npu(q, k, v, plan)
+
+    reference = torch.nn.functional.scaled_dot_product_attention(
+        q.transpose(1, 2),
+        k.transpose(1, 2),
+        v.transpose(1, 2),
+        scale=impl.softmax_scale,
+    ).transpose(1, 2)
+    error = (out.float() - reference.float()).abs().mean() / reference.float().abs().mean()
+    assert error < 2e-3, f"mean relative error {error:.4%} against dense attention"