# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Unit tests for Helios per-component sequence splitting. These tests verify the manual sharding logic that replaces the broken rope/blocks.0 hooks in the Helios ``_sp_plan``. They exercise the *real* ``HeliosTransformer3DModel._sp_split_seq`` (the parallel-state getters it imports lazily inside the method body are monkeypatched), so the model-side helper cannot drift silently from what is tested. For the multi-rank mechanism around the helper (manual ``_sp_shard_depth`` bump, attn1/attn2 parallel strategy, ws=2 vs ws=1 output equivalence) see ``test_helios_usp_2rank.py``. Coverage: 1. USP disabled (ws == 1) returns the tensor unchanged. 2. Divisible split gives each rank the expected slice. 3. Non-divisible split pads with replicated tokens (no crash). 4. Per-component split: history + current tokens are distributed so that *both* ranks receive some current tokens (no mosaic regression). 5. Edge cases: 1-D tensor, empty seq, 4-D tensor. """ from collections.abc import Iterator from contextlib import contextmanager import pytest import torch import vllm_omni.diffusion.distributed.parallel_state as sp_parallel_state from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @contextmanager def _fake_sp_state(ws: int, rank: int) -> Iterator[None]: """Patch the parallel-state getters used by the real helper. ``_sp_split_seq`` imports these names from the module *inside* its function body, so replacing the module attributes takes effect. """ prev_ws = sp_parallel_state.get_sequence_parallel_world_size prev_rank = sp_parallel_state.get_sequence_parallel_rank sp_parallel_state.get_sequence_parallel_world_size = lambda: ws sp_parallel_state.get_sequence_parallel_rank = lambda: rank try: yield finally: sp_parallel_state.get_sequence_parallel_world_size = prev_ws sp_parallel_state.get_sequence_parallel_rank = prev_rank def _sp_split_seq(x: torch.Tensor, ws: int, rank: int) -> torch.Tensor: """Call the real model-side helper without constructing a model.""" with _fake_sp_state(ws, rank): return HeliosTransformer3DModel._sp_split_seq(None, x) # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- class TestSpSplitSeqDisabled: """When USP is off (ws == 1) the helper must be a no-op.""" def test_ws1_returns_unchanged(self): x = torch.randn(2, 100, 64) out = _sp_split_seq(x, ws=1, rank=0) assert torch.equal(out, x) assert out is x class TestSpSplitSeqDivisible: """Divisible sequences: each rank gets exactly seq_len / ws tokens.""" def test_each_rank_gets_correct_slice(self): for ws in [2, 4, 8]: seq_len = ws * 10 x = torch.arange(seq_len * 4).float().reshape(1, seq_len, 4) for rank in range(ws): out = _sp_split_seq(x, ws=ws, rank=rank) n = seq_len // ws expected = x[:, rank * n : (rank + 1) * n, :] assert out.shape == (1, n, 4) assert torch.equal(out, expected) def test_output_is_contiguous(self): x = torch.randn(1, 20, 8) out = _sp_split_seq(x, ws=2, rank=0) assert out.is_contiguous() class TestSpSplitSeqNonDivisible: """Non-divisible sequences are padded with replicated tokens.""" def test_non_divisible_pads_instead_of_crash(self): """400x272 → 9*25*17 = 3825 tokens, USP=2 → not divisible. Should pad to 3826 and split, not crash.""" x = torch.randn(1, 3825, 8) for rank in range(2): out = _sp_split_seq(x, ws=2, rank=rank) # 3825 → pad 1 → 3826 → /2 = 1913 assert out.shape[1] == 1913 def test_non_divisible_pads_with_replicated_last_token(self): """Verify the padding is the replicated last token, not zeros.""" x = torch.randn(1, 7, 4) # 7 % 2 = 1, pad 1 last_token = x[:, -1:, :].clone() out = _sp_split_seq(x, ws=2, rank=0) # rank 0 gets tokens [0:4] (first 4 of the padded 8) assert out.shape == (1, 4, 4) # rank 1 gets tokens [4:8] — token 7 (the padded one) = last token out1 = _sp_split_seq(x, ws=2, rank=1) assert out1.shape == (1, 4, 4) assert torch.equal(out1[:, -1:, :], last_token) def test_non_divisible_7_4(self): """7 tokens, ws=4 → pad 1 → 8 → /4 = 2 per rank.""" x = torch.randn(1, 7, 8) for rank in range(4): out = _sp_split_seq(x, ws=4, rank=rank) assert out.shape[1] == 2 def test_non_divisible_15_4(self): """15 tokens, ws=4 → pad 1 → 16 → /4 = 4 per rank.""" x = torch.randn(1, 15, 8) for rank in range(4): out = _sp_split_seq(x, ws=4, rank=rank) assert out.shape[1] == 4 class TestSpSplitSeqPerComponent: """Simulate Helios forward: split current + history components so that *both* ranks receive some current tokens (no mosaic regression).""" def test_both_ranks_get_current_tokens(self): """history=2400, current=540, ws=2. After per-component split, each rank must have > 0 current tokens.""" current = torch.randn(1, 540, 64) history = torch.randn(1, 2400, 64) for rank in range(2): cur_local = _sp_split_seq(current, ws=2, rank=rank) hist_local = _sp_split_seq(history, ws=2, rank=rank) assert cur_local.shape[1] > 0, f"Rank {rank} received 0 current tokens" assert cur_local.shape[1] == 270 # 540 / 2 assert hist_local.shape[1] == 1200 # 2400 / 2 original_context_length = cur_local.shape[1] assert original_context_length > 0 # no mosaic regression def test_gather_restores_full_current(self): """After splitting current into 2 ranks, simulating proj_out gather must restore the full current sequence length so that the original unpatchify reshape succeeds.""" current = torch.randn(1, 540, 64) local_chunks = [_sp_split_seq(current, ws=2, rank=rank) for rank in range(2)] gathered = torch.cat(local_chunks, dim=1) assert gathered.shape[1] == 540 # full current restored def test_gather_with_padding_truncates_correctly(self): """Non-divisible current (3825) → pad to 3826 → split → gather → 3826. Slice to 3825 for unpatchify.""" current = torch.randn(1, 3825, 64) local_chunks = [_sp_split_seq(current, ws=2, rank=rank) for rank in range(2)] gathered = torch.cat(local_chunks, dim=1) assert gathered.shape[1] == 3826 # padded size # Simulate the slice before unpatchify expected_seq = 3825 if gathered.shape[1] > expected_seq: gathered = gathered[:, :expected_seq, :] assert gathered.shape[1] == 3825 # correct after slice def test_rotary_emb_aligned_with_hidden_states(self): """rotary_emb and hidden_states must be split at the same positions so that tokens align inside the attention layer.""" hidden = torch.randn(1, 540, 5120) rotary = torch.randn(1, 540, 256) for rank in range(2): h_local = _sp_split_seq(hidden, ws=2, rank=rank) r_local = _sp_split_seq(rotary, ws=2, rank=rank) assert h_local.shape[1] == r_local.shape[1] def test_whole_split_causes_mosaic_but_per_component_does_not(self): """Demonstrates why per-component split is needed: with whole split (simulated), rank 0 gets 0 current tokens; with per-component split, both ranks get current tokens.""" history = torch.randn(1, 2400, 64) current = torch.randn(1, 540, 64) full = torch.cat([history, current], dim=1) # 2940 # Whole split (the OLD broken behavior) whole_n = full.shape[1] // 2 # 1470 rank0_whole = full[:, :whole_n, :] rank0_current = rank0_whole.shape[1] - 2400 # negative → 0 current assert rank0_current <= 0, "Whole split should give rank 0 no current" # Per-component split (the NEW correct behavior) for rank in range(2): cur_local = _sp_split_seq(current, ws=2, rank=rank) assert cur_local.shape[1] == 270 # both ranks have current class TestSpSplitSeqEdgeCases: """Edge cases that should not crash or silently misbehave.""" def test_1d_tensor_skipped(self): """1-D tensors must be returned unchanged (dim < 2 guard).""" x = torch.randn(100) out = _sp_split_seq(x, ws=2, rank=0) assert torch.equal(out, x) def test_empty_seq_skipped(self): """Empty sequence dimension must be returned unchanged.""" x = torch.randn(1, 0, 64) out = _sp_split_seq(x, ws=2, rank=0) assert out.shape == x.shape def test_4d_tensor_split_on_dim1(self): """4-D tensor [B, S, H, D] should split along dim=1.""" x = torch.randn(1, 20, 8, 32) out = _sp_split_seq(x, ws=2, rank=1) assert out.shape == (1, 10, 8, 32)