diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md index 094a0f3a9..b6fd13b12 100644 --- a/docs/user_guide/diffusion/attention_backends.md +++ b/docs/user_guide/diffusion/attention_backends.md @@ -283,23 +283,20 @@ blocks — use Ulysses SP (`ring_degree=1`). ### Which geometries run sparse -Sparsity is only applied when the video segment is a **multiple of 128 rows**, where the row count -is `latent_t × (height / 32) × (width / 32)`. Otherwise `rf_v2`'s block mask and the kernel's own -tiling disagree on the block count, and the block straddling the seam mixes video and prefix rows -that selection may then drop. Any resolution still runs — an unaligned geometry falls back to dense -attention and logs `RAINFUSION_ATTN staying dense` with the row count it computed — but it gets no -speedup, so pick an aligned geometry when you want one. +MindIE-SD rf_v2 handles arbitrary video grids by spatially rearranging video first and, when +necessary, promoting a real-video suffix to its always-kept prefix segment. It never pads: rf_v2 +does not consume an attention mask for padded keys. The remaining video rows stay sparse, while the +promoted rows remain visible to every query and key. -Alignment is necessary, not sufficient, for good quality. `rf_v2` groups video positions into 8x8 -spatial tiles (two tiles fill one 128-row block), which is what makes a selected block a compact -patch of the frame. When the latent `h` or `w` is not a multiple of 8, the leftover rows or columns -are peeled off and appended as a flat run instead of being tiled, so they get pooled with spatially -distant positions and selection can no longer rank them meaningfully. This is silent, and invisible -to a `sparsity=0` check, because a fully populated mask does not care how blocks are grouped. +For example, MiniMax-H3 at 1344x768 has grid `(62, 24, 42)`: its 62496 video rows leave a 32-row +128-block residue, and its two non-8-wide latent columns form a 2928-row rearranged suffix. +MindIE-SD promotes 2976 rows, leaving 59520 sparse rows (`465 × 128`). This keeps the block mask +and kernel tiling aligned, at the cost of retaining more blocks. -"Multiple of 8" on the latent grid means **width and height that are multiples of 256**. Aligned -resolutions off that grid still run sparse; they just lose more fidelity at the same `sparsity`, -which you can buy back with `start_step`. +8x8-aligned spatial grids remain preferable for performance. A latent `h` or `w` not divisible by 8 +creates an always-kept residual suffix and lowers realized sparsity. "Multiple of 8" on the latent +grid means **width and height that are multiples of 256**. This path requires a MindIE-SD release +that implements protected video tails for rf_v2. ## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell) diff --git a/tests/diffusion/attention/test_rainfusion_plan.py b/tests/diffusion/attention/test_rainfusion_plan.py index e88a973a9..f079a1758 100644 --- a/tests/diffusion/attention/test_rainfusion_plan.py +++ b/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" diff --git a/vllm_omni/diffusion/attention/backends/rainfusion_attn.py b/vllm_omni/diffusion/attention/backends/rainfusion_attn.py index 10ab55991..28b032712 100644 --- a/vllm_omni/diffusion/attention/backends/rainfusion_attn.py +++ b/vllm_omni/diffusion/attention/backends/rainfusion_attn.py @@ -128,9 +128,9 @@ class RainFusionAttentionImpl(AttentionImpl): whose extent the model publishes as ``AttentionMetadata.video_layout``. Every other case — warmup denoise steps, exempt layers, a layer that does not declare ``qkv_layout="BSND"``, sequences without a published video segment, video - segments too short to pay for block selection or not aligned to the kernel's - block size — delegates to FlashAttention, so a model can select this backend - unconditionally. + segments too short to pay for block selection — delegates to FlashAttention, + so a model can select this backend unconditionally. MindIE-SD handles an + irregular video tail internally, retaining it outside the sparse blocks. """ def __init__( @@ -292,25 +292,6 @@ class RainFusionAttentionImpl(AttentionImpl): _MIN_VIDEO_BLOCKS, ) return None - if video_len % _BLOCK_SIZE != 0: - # rf_v2 tiles the sequence into ceil((video + prefix) / 128) blocks but - # builds its block mask from the two segments pooled apart. The counts - # only agree when the video segment ends on a block boundary; otherwise - # a prefix block is forced on out of range, and the block straddling the - # seam holds real prefix rows that selection may drop. Padding the - # sequence to realign it is not an option because rf_v2 ignores - # attn_mask, so pad keys would take a share of every softmax denominator. - logger.warning_once( - "RAINFUSION_ATTN staying dense: latent grid %s gives %d video rows, which is not a " - "multiple of the %d-row kernel block. Choose a geometry whose latent_t * h * w is a " - "multiple of %d to enable sparsity.", - tuple(latent_shape), - video_len, - _BLOCK_SIZE, - _BLOCK_SIZE, - ) - return None - logger.info_once( "RAINFUSION_ATTN active: sparsity=%.2f, start_step=%d, exempt_layers=%d, " "latent_grid=%s, prefix_rows=%d, video_rows=%d. Realized sparsity is lower than nominal "