From 84bbd57e4d2f1210527587b8ba53e168902506cb Mon Sep 17 00:00:00 2001 From: Honghan Zhu Date: Fri, 24 Jul 2026 22:45:26 +0800 Subject: [PATCH] =?UTF-8?q?[Feature]=20Add=20AllGather-KV=20sequence-paral?= =?UTF-8?q?lel=20attention=20str=E2=80=A6=20(#4968)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: zhuhh97 Signed-off-by: dph97 Signed-off-by: zhuhh97 Co-authored-by: Claude Opus 4.7 Co-authored-by: dengyunyang <584797741@qq.com> --- tests/config/test_omni_config.py | 30 ++ .../diffusion/attention/test_attention_sp.py | 288 +++++++++++++++++- tests/diffusion/attention/test_flash_attn.py | 30 ++ .../attention/test_piecewise_attn.py | 87 ++++++ tests/diffusion/attention/test_sdpa.py | 57 ++++ .../test_parallel_state_sp_groups.py | 39 ++- .../test_image_kv_cache_manager.py | 2 + .../test_hunyuan_image3_pixel_accuracy.py | 48 ++- tests/test_config_factory.py | 25 ++ vllm_omni/config/omni_config.py | 8 +- vllm_omni/config/stage_config.py | 8 +- vllm_omni/deploy/hunyuan_image3_dit.yaml | 1 + .../diffusion/attention/backends/abstract.py | 9 + .../attention/backends/cudnn_attn.py | 6 +- .../attention/backends/flash_attn.py | 41 ++- .../attention/backends/flash_attn_hub.py | 8 +- .../diffusion/attention/backends/sdpa.py | 21 +- .../backends/utils/attn_runtime_selector.py | 40 +++ .../backends/utils/piecewise_attn.py | 91 ++++-- vllm_omni/diffusion/attention/layer.py | 17 ++ .../attention/parallel/allgather_kv.py | 182 +++++++++++ .../diffusion/attention/parallel/factory.py | 32 +- vllm_omni/diffusion/data.py | 29 +- .../distributed/group_coordinator.py | 20 +- .../diffusion/distributed/parallel_state.py | 90 ++++-- vllm_omni/diffusion/distributed/sp_plan.py | 20 +- .../hunyuan_image3_transformer.py | 31 +- vllm_omni/diffusion/registry.py | 29 +- vllm_omni/diffusion/utils/kv_utils.py | 13 + .../diffusion/worker/diffusion_worker.py | 1 + vllm_omni/engine/arg_utils.py | 1 + vllm_omni/engine/async_omni_engine.py | 4 +- vllm_omni/entrypoints/cli/serve.py | 8 + 33 files changed, 1181 insertions(+), 135 deletions(-) create mode 100644 tests/diffusion/attention/test_sdpa.py create mode 100644 vllm_omni/diffusion/attention/backends/utils/attn_runtime_selector.py create mode 100644 vllm_omni/diffusion/attention/parallel/allgather_kv.py diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 1da71b21d..523542232 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -431,6 +431,7 @@ def test_sub_config_fields_match_rfc_scopes(): "sequence_parallel_size", "ulysses_degree", "ring_degree", + "allgather_degree", "ulysses_mode", "cfg_parallel_size", "vae_patch_parallel_size", @@ -547,11 +548,40 @@ def test_from_pipeline_config_derives_sequence_parallel_size_from_degrees(tmp_pa assert stage.parallel_config.world_size == 6 +def test_from_pipeline_config_derives_sequence_parallel_size_from_allgather_degree(tmp_path): + deploy_path = tmp_path / "dreamzero_allgather_parallel.yaml" + deploy_path.write_text( + "\n".join( + [ + "pipeline: dreamzero", + "async_chunk: false", + "stages:", + " - stage_id: 0", + " parallel_config:", + " sequence_parallel_size: 99", + " allgather_degree: 2", + ] + ) + ) + + stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0) + + assert isinstance(stage, VllmOmniDiffusionStageConfig) + assert stage.parallel_config.allgather_degree == 2 + assert stage.parallel_config.sequence_parallel_size == 2 + assert stage.parallel_config.world_size == 2 + + def test_diffusion_parallel_config_rejects_cfg_parallel_size_outside_current_bound(): with pytest.raises(ValidationError): OmniStageDiffusionParallelConfig(cfg_parallel_size=4) +def test_diffusion_parallel_config_rejects_allgather_with_ulysses_or_ring(): + with pytest.raises(ValidationError): + OmniStageDiffusionParallelConfig(allgather_degree=2, ulysses_degree=2) + + def test_stage_realizations_use_stage_specific_parallel_config_types(): qwen_config = _from_pipeline_key("qwen3_tts") ar_stage = qwen_config.stage_by_id(0) diff --git a/tests/diffusion/attention/test_attention_sp.py b/tests/diffusion/attention/test_attention_sp.py index e5fb93d71..e32c4c5e4 100644 --- a/tests/diffusion/attention/test_attention_sp.py +++ b/tests/diffusion/attention/test_attention_sp.py @@ -31,7 +31,11 @@ import tempfile import pytest import torch +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata, QueryRange from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.attention.parallel.allgather_kv import ( + AllGatherKVParallelAttention, +) from vllm_omni.diffusion.config import set_current_diffusion_config from vllm_omni.diffusion.data import ( DiffusionParallelConfig, @@ -155,14 +159,267 @@ class TestMultiLayerAttentionModel(torch.nn.Module): return hidden_states +class _MockAllGatherSPGroup: + def __init__(self, *, rank: int, gather_chunks: list[list[torch.Tensor]]) -> None: + self.allgather_group = object() + self.allgather_world_size = len(gather_chunks[0]) + self.allgather_rank = rank + self._gather_chunks = list(gather_chunks) + self.gathered_input_shapes: list[tuple[int, ...]] = [] + + def all_gather(self, input_: torch.Tensor, dim: int = 0, separate_tensors: bool = False, group=None): + assert not separate_tensors + assert group is self.allgather_group + self.gathered_input_shapes.append(tuple(input_.shape)) + chunks = self._gather_chunks.pop(0) + return torch.cat(chunks, dim=dim) + + +def test_allgather_kv_slices_full_dense_mask_to_local_query_rows(): + rank = 1 + joint_len = 1 + img_seq_local = 2 + img_seq_full = 4 + key_chunks = [ + torch.full((1, img_seq_local, 1, 1), 10.0), + torch.full((1, img_seq_local, 1, 1), 20.0), + ] + value_chunks = [ + torch.full((1, img_seq_local, 1, 1), 30.0), + torch.full((1, img_seq_local, 1, 1), 40.0), + ] + strategy = AllGatherKVParallelAttention( + _MockAllGatherSPGroup(rank=rank, gather_chunks=[key_chunks, value_chunks]), + ) + + query = torch.zeros((1, img_seq_local, 1, 1)) + key = key_chunks[rank] + value = value_chunks[rank] + joint = torch.ones((1, joint_len, 1, 1)) + mask = torch.arange((joint_len + img_seq_full) * (joint_len + img_seq_full)) + mask = (mask % 2 == 0).view(1, 1, joint_len + img_seq_full, joint_len + img_seq_full) + metadata = AttentionMetadata( + attn_mask=mask, + joint_query=joint, + joint_key=joint, + joint_value=joint, + joint_strategy="front", + ) + + q_local, k_full, v_full, metadata_local, _ = strategy.pre_attention(query, key, value, metadata) + + expected_rows = torch.cat( + [ + mask[..., :joint_len, :], + mask[ + ..., + joint_len + rank * img_seq_local : joint_len + (rank + 1) * img_seq_local, + :, + ], + ], + dim=-2, + ) + assert q_local.shape[1] == joint_len + img_seq_local + assert k_full.shape[1] == joint_len + img_seq_full + assert v_full.shape[1] == joint_len + img_seq_full + torch.testing.assert_close(k_full[:, joint_len:], torch.cat(key_chunks, dim=1)) + torch.testing.assert_close(v_full[:, joint_len:], torch.cat(value_chunks, dim=1)) + assert metadata_local is not None + assert torch.equal(metadata_local.attn_mask, expected_rows) + + +def test_allgather_kv_slices_rear_joint_dense_mask(): + rank = 1 + joint_len = 1 + img_seq_local = 2 + img_seq_full = 4 + key_chunks = [torch.zeros((1, img_seq_local, 1, 1)) for _ in range(2)] + value_chunks = [torch.zeros_like(chunk) for chunk in key_chunks] + strategy = AllGatherKVParallelAttention( + _MockAllGatherSPGroup(rank=rank, gather_chunks=[key_chunks, value_chunks]), + ) + query = torch.zeros((1, img_seq_local, 1, 1)) + joint = torch.ones((1, joint_len, 1, 1)) + mask = torch.arange((img_seq_full + joint_len) ** 2) + mask = (mask % 2 == 0).view(1, 1, img_seq_full + joint_len, img_seq_full + joint_len) + metadata = AttentionMetadata( + attn_mask=mask, + joint_query=joint, + joint_key=joint, + joint_value=joint, + joint_strategy="rear", + ) + + q_local, _, _, metadata_local, _ = strategy.pre_attention(query, key_chunks[rank], value_chunks[rank], metadata) + + img_start = rank * img_seq_local + expected_rows = torch.cat( + [ + mask[..., img_start : img_start + img_seq_local, :], + mask[..., img_seq_full : img_seq_full + joint_len, :], + ], + dim=-2, + ) + assert q_local.shape[1] == img_seq_local + joint_len + assert metadata_local is not None + assert torch.equal(metadata_local.attn_mask, expected_rows) + + +def test_allgather_kv_rejects_invalid_joint_strategy(): + chunks = [[torch.zeros((1, 2, 1, 1)) for _ in range(2)] for _ in range(2)] + strategy = AllGatherKVParallelAttention(_MockAllGatherSPGroup(rank=0, gather_chunks=chunks)) + joint = torch.ones((1, 1, 1, 1)) + metadata = AttentionMetadata( + joint_query=joint, + joint_key=joint, + joint_value=joint, + joint_strategy="back", + ) + + with pytest.raises(ValueError, match="Unsupported joint_strategy"): + strategy.pre_attention(torch.zeros((1, 2, 1, 1)), chunks[0][0], chunks[1][0], metadata) + + +def test_allgather_kv_preserves_global_spans_and_sets_query_ranges(): + strategy = AllGatherKVParallelAttention( + _MockAllGatherSPGroup( + rank=1, + gather_chunks=[ + [torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))], + [torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))], + ], + ), + ) + query = torch.zeros((1, 2, 1, 1)) + key = torch.zeros((1, 2, 1, 1)) + value = torch.zeros((1, 2, 1, 1)) + joint = torch.ones((1, 1, 1, 1)) + metadata = AttentionMetadata( + joint_query=joint, + joint_key=joint, + joint_value=joint, + full_attn_spans=[ + [ + (0, 1), # joint span: kept on every rank. + (2, 5), # image span crossing rank 0 and rank 1 image shards. + ] + ], + ) + + _, _, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata) + + assert metadata_out is not None + assert metadata_out.full_attn_spans == [[(0, 1), (2, 5)]] + assert metadata_out.query_ranges == ( + QueryRange(0, 1, 0), + QueryRange(1, 3, 3), + ) + + +def test_allgather_kv_query_ranges_include_reused_prefix_offset(): + strategy = AllGatherKVParallelAttention( + _MockAllGatherSPGroup( + rank=1, + gather_chunks=[ + [torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))], + [torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))], + ], + ), + ) + query = torch.zeros((1, 2, 1, 1)) + key = torch.zeros((1, 2, 1, 1)) + value = torch.zeros((1, 2, 1, 1)) + joint_query = torch.ones((1, 1, 1, 1)) + # K/V retain two reused AR-prefix tokens that have already been removed + # from Q and from the attention mask's query rows. + joint_key = torch.ones((1, 3, 1, 1)) + joint_value = torch.ones((1, 3, 1, 1)) + mask = torch.arange(5 * 7).view(1, 1, 5, 7) + metadata = AttentionMetadata( + attn_mask=mask, + joint_query=joint_query, + joint_key=joint_key, + joint_value=joint_value, + full_attn_spans=[[(0, 7)]], + ) + + _, _, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata) + + assert metadata_out is not None + assert metadata_out.query_ranges == ( + QueryRange(0, 1, 2), + QueryRange(1, 3, 5), + ) + expected_mask = torch.cat([mask[..., :1, :], mask[..., 3:5, :]], dim=-2) + assert torch.equal(metadata_out.attn_mask, expected_mask) + + +def test_allgather_kv_allows_empty_full_attn_spans(): + strategy = AllGatherKVParallelAttention( + _MockAllGatherSPGroup( + rank=0, + gather_chunks=[ + [torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))], + [torch.zeros((1, 2, 1, 1)), torch.zeros((1, 2, 1, 1))], + ], + ), + ) + query = torch.zeros((1, 2, 1, 1)) + key = torch.zeros((1, 2, 1, 1)) + value = torch.zeros((1, 2, 1, 1)) + metadata = AttentionMetadata(full_attn_spans=[[]]) + + _, _, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata) + + assert metadata_out is not None + assert metadata_out.full_attn_spans == [[]] + assert metadata_out.query_ranges == (QueryRange(0, 2, 0),) + + +def test_allgather_kv_keeps_gathered_kv_compressed_for_gqa(): + rank = 0 + img_seq_local = 2 + kv_heads = 2 + repeat_num = 2 + q_heads = kv_heads * repeat_num + key_chunks = [ + torch.arange(1 * img_seq_local * kv_heads * 1, dtype=torch.float32).view(1, img_seq_local, kv_heads, 1), + torch.arange(100, 100 + 1 * img_seq_local * kv_heads * 1, dtype=torch.float32).view( + 1, img_seq_local, kv_heads, 1 + ), + ] + value_chunks = [chunk + 1000 for chunk in key_chunks] + sp_group = _MockAllGatherSPGroup(rank=rank, gather_chunks=[key_chunks, value_chunks]) + strategy = AllGatherKVParallelAttention(sp_group) + + query = torch.zeros((1, img_seq_local, q_heads, 1)) + _, k_full, v_full, _, _ = strategy.pre_attention(query, key_chunks[rank], value_chunks[rank], AttentionMetadata()) + + assert sp_group.gathered_input_shapes == [ + (1, img_seq_local, kv_heads, 1), + (1, img_seq_local, kv_heads, 1), + ] + assert k_full.shape == (1, img_seq_local * 2, kv_heads, 1) + assert v_full.shape == (1, img_seq_local * 2, kv_heads, 1) + expected_key = torch.cat(key_chunks, dim=1) + expected_value = torch.cat(value_chunks, dim=1) + torch.testing.assert_close(k_full, expected_key) + torch.testing.assert_close(v_full, expected_value) + + @pytest.mark.parametrize( "test_model_cls", [ TestMultiLayerAttentionModel, ], ) -@pytest.mark.parametrize("ulysses_degree", [2]) -@pytest.mark.parametrize("ring_degree", [2]) +@pytest.mark.parametrize( + ("ulysses_degree", "ring_degree", "allgather_degree"), + [ + pytest.param(2, 2, 1, id="ulysses-ring"), + pytest.param(1, 1, 2, id="allgather-kv"), + ], +) @pytest.mark.parametrize("batch_size", [2]) @pytest.mark.parametrize("seq_len", [16]) @pytest.mark.parametrize("num_heads", [8]) @@ -175,6 +432,7 @@ class TestMultiLayerAttentionModel(torch.nn.Module): def test_sequence_parallel( ulysses_degree: int, ring_degree: int, + allgather_degree: int, test_model_cls: type[torch.nn.Module], dtype: torch.dtype, causal: bool, @@ -186,8 +444,8 @@ def test_sequence_parallel( num_heads: int, head_size: int, ): - """Test Ulysses attention by comparing with and without SP enabled.""" - sequence_parallel_size = ulysses_degree * ring_degree + """Compare Ulysses/Ring and AllGather-KV SP against a single-rank run.""" + sequence_parallel_size = allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree # Skip if not enough GPUs available available_gpus = current_omni_platform.get_device_count() @@ -223,6 +481,7 @@ def test_sequence_parallel( use_compile, 1, # ulysses_degree = 1 1, # ring_degree = 1 + 1, # allgather_degree = 1 (baseline: no AllGather-KV SP) 1, # sequence_parallel_size = 1 baseline_output_file, model_state_file, @@ -233,7 +492,10 @@ def test_sequence_parallel( ) # Step 2: Run with SP enabled - print(f"\n[SP Test] Running with SP (ulysses_degree={ulysses_degree}, ring_degree={ring_degree})...") + print( + "\n[SP Test] Running with SP " + f"(ulysses={ulysses_degree}, ring={ring_degree}, allgather={allgather_degree})..." + ) torch.multiprocessing.spawn( ulysses_attention_on_test_model, args=( @@ -250,6 +512,7 @@ def test_sequence_parallel( use_compile, ulysses_degree, ring_degree, + allgather_degree, sequence_parallel_size, sp_output_file, model_state_file, @@ -345,6 +608,7 @@ def ulysses_attention_on_test_model( use_compile: bool, ulysses_degree: int, ring_degree: int, + allgather_degree: int, sequence_parallel_size: int, output_file: str, model_state_file: str, @@ -356,7 +620,11 @@ def ulysses_attention_on_test_model( RANDOM_SEED = 42 seed_everything(RANDOM_SEED) - mode_str = "Baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})" + if allgather_degree > 1: + sp_kind = f"allgather={allgather_degree}" + else: + sp_kind = f"ulysses={ulysses_degree}, ring={ring_degree}" + mode_str = "Baseline (no SP)" if is_baseline else f"SP ({sp_kind})" print(f"\n[{mode_str}] Rank {local_rank}/{world_size} - Random seed set to {RANDOM_SEED}") device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") @@ -384,6 +652,7 @@ def ulysses_attention_on_test_model( sequence_parallel_size=sequence_parallel_size, ulysses_degree=ulysses_degree, ring_degree=ring_degree, + allgather_degree=allgather_degree, cfg_parallel_size=1, ) @@ -400,6 +669,7 @@ def ulysses_attention_on_test_model( sequence_parallel_size=sequence_parallel_size, ulysses_degree=ulysses_degree, ring_degree=ring_degree, + allgather_degree=allgather_degree, tensor_parallel_size=1, pipeline_parallel_size=1, ) @@ -519,7 +789,11 @@ def ulysses_attention_on_test_model( with open(output_file, "wb") as f: pickle.dump(output_np, f) - mode_str = "baseline (no SP)" if is_baseline else f"SP (ulysses={ulysses_degree}, ring={ring_degree})" + if allgather_degree > 1: + sp_kind = f"allgather={allgather_degree}" + else: + sp_kind = f"ulysses={ulysses_degree}, ring={ring_degree}" + mode_str = "baseline (no SP)" if is_baseline else f"SP ({sp_kind})" print( f"\n[{mode_str}] ✓ Saved output with shape {full_output.shape}:\n" f" - batch_size={batch_size}, seq_len={seq_len}\n" diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py index a6fcf8fcb..5ffd6684a 100644 --- a/tests/diffusion/attention/test_flash_attn.py +++ b/tests/diffusion/attention/test_flash_attn.py @@ -297,6 +297,36 @@ def test_flash_attn_func_preferred_over_varlen(): print("✓ flash_attn_func forward works correctly!") +def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch): + calls = [] + + def fake_varlen_func(q, k, v, **kwargs): + assert q.shape[1] == 4 + assert k.shape[1] == 2 + assert v.shape[1] == 2 + calls.append(kwargs) + return q + + monkeypatch.setattr(fa, "flash_attn_func", None) + monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func) + monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True) + + impl = FlashAttentionImpl(num_heads=4, num_kv_heads=2, head_size=4, softmax_scale=0.5, causal=False) + query = torch.randn(1, 3, 4, 4) + key = torch.randn(1, 3, 2, 4) + value = torch.randn(1, 3, 2, 4) + metadata = AttentionMetadata(full_attn_spans=[[(0, 3)]]) + + output = impl.forward_cuda(query, key, value, metadata) + + assert output.shape == query.shape + assert len(calls) == 1 + assert calls[0]["max_seqlen_q"] == 3 + assert calls[0]["max_seqlen_k"] == 3 + assert calls[0]["causal"] is False + assert calls[0]["softmax_scale"] == 0.5 + + if __name__ == "__main__": print("Running FlashAttention Padding Tests...") print("=" * 60) diff --git a/tests/diffusion/attention/test_piecewise_attn.py b/tests/diffusion/attention/test_piecewise_attn.py index 656087623..00b98ac64 100644 --- a/tests/diffusion/attention/test_piecewise_attn.py +++ b/tests/diffusion/attention/test_piecewise_attn.py @@ -19,7 +19,11 @@ import pytest import torch import torch.nn.functional as F +import vllm_omni.diffusion.attention.backends.utils.piecewise_attn as piecewise_module +from vllm_omni.diffusion.attention.backends.abstract import QueryRange from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import ( + Segment, + build_segments, piecewise_attn, ) @@ -71,6 +75,12 @@ BATCH_CASES = [ ] +def test_build_segments_full_span_retains_global_kv_end(): + segments = build_segments([(2, 7)], query_offset=4, query_len=2) + + assert segments == [Segment(q_start=4, q_end=6, kv_end=7, mode="full")] + + @pytest.mark.parametrize("global_spans", SPAN_CASES) @pytest.mark.parametrize("q_range", Q_RANGE_CASES) @pytest.mark.parametrize("batch_size", BATCH_CASES) @@ -124,3 +134,80 @@ def test_piecewise_span_fully_before_qstart(): ) expected = _full_reference(query, key, value, global_spans, q_start, q_end, softmax_scale) torch.testing.assert_close(got, expected, atol=1e-5, rtol=1e-5) + + +def test_piecewise_noncontiguous_query_ranges_match_full_mask(): + torch.manual_seed(0) + B, H, D, Sk = 1, 2, 8, 8 + key = torch.randn(B, Sk, H, D, device=DEVICE) + value = torch.randn(B, Sk, H, D, device=DEVICE) + query = torch.randn(B, 4, H, D, device=DEVICE) + spans = [(2, 7)] + query_ranges = ( + QueryRange(0, 1, 0), + QueryRange(1, 4, 4), + ) + softmax_scale = D**-0.5 + + got = piecewise_attn( + query, + key, + value, + full_attn_spans=[spans], + softmax_scale=softmax_scale, + attn_func=_sdpa_attn_func, + query_ranges=query_ranges, + ) + expected = torch.cat( + [ + _full_reference(query[:, :1], key, value, spans, 0, 1, softmax_scale), + _full_reference(query[:, 1:], key, value, spans, 4, 7, softmax_scale), + ], + dim=1, + ) + torch.testing.assert_close(got, expected, atol=1e-5, rtol=1e-5) + + +def test_piecewise_rejects_incomplete_query_ranges(): + query = torch.zeros(1, 3, 1, 2) + key = torch.zeros(1, 3, 1, 2) + value = torch.zeros_like(key) + + with pytest.raises(ValueError, match="cover local query contiguously"): + piecewise_attn( + query, + key, + value, + full_attn_spans=[[]], + softmax_scale=1.0, + attn_func=_sdpa_attn_func, + query_ranges=(QueryRange(1, 3, 1),), + ) + + +def test_piecewise_combines_noncontiguous_query_ranges_with_one_cat(monkeypatch): + query = torch.randn(1, 4, 2, 8) + key = torch.randn(1, 8, 2, 8) + value = torch.randn_like(key) + cat_calls = 0 + original_cat = torch.cat + + def counted_cat(tensors, dim=0): + nonlocal cat_calls + cat_calls += 1 + return original_cat(tensors, dim=dim) + + monkeypatch.setattr(piecewise_module.torch, "cat", counted_cat) + + output = piecewise_attn( + query, + key, + value, + full_attn_spans=[[(0, 8)]], + softmax_scale=1.0, + attn_func=lambda q, k, v, **kwargs: q, + query_ranges=(QueryRange(0, 1, 0), QueryRange(1, 4, 4)), + ) + + assert cat_calls == 1 + torch.testing.assert_close(output, query) diff --git a/tests/diffusion/attention/test_sdpa.py b/tests/diffusion/attention/test_sdpa.py new file mode 100644 index 000000000..d02642517 --- /dev/null +++ b/tests/diffusion/attention/test_sdpa.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +import vllm_omni.diffusion.attention.backends.sdpa as sdpa_backend +from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_sdpa_expands_kv_when_native_gqa_kernel_is_unavailable(monkeypatch): + calls = [] + + def fake_sdpa(query, key, value, **kwargs): + calls.append((query.shape, key.shape, value.shape, kwargs)) + return query + + monkeypatch.setattr(sdpa_backend, "can_sdpa_use_fused_gqa", lambda *args: False) + monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa) + + impl = SDPAImpl(num_heads=4, num_kv_heads=2, head_size=8, softmax_scale=0.5) + output = impl.forward_cuda( + torch.randn(1, 3, 4, 8), + torch.randn(1, 3, 2, 8), + torch.randn(1, 3, 2, 8), + ) + + query_shape, key_shape, value_shape, kwargs = calls[0] + assert query_shape == key_shape == value_shape == (1, 4, 3, 8) + assert kwargs["enable_gqa"] is False + assert output.shape == (1, 3, 4, 8) + + +def test_sdpa_keeps_compressed_kv_when_native_gqa_kernel_is_available(monkeypatch): + calls = [] + + def fake_sdpa(query, key, value, **kwargs): + calls.append((query.shape, key.shape, value.shape, kwargs)) + return query + + monkeypatch.setattr(sdpa_backend, "can_sdpa_use_fused_gqa", lambda *args: True) + monkeypatch.setattr(torch.nn.functional, "scaled_dot_product_attention", fake_sdpa) + + impl = SDPAImpl(num_heads=4, num_kv_heads=2, head_size=8, softmax_scale=0.5) + output = impl.forward_cuda( + torch.randn(1, 3, 4, 8), + torch.randn(1, 3, 2, 8), + torch.randn(1, 3, 2, 8), + ) + + query_shape, key_shape, value_shape, kwargs = calls[0] + assert query_shape == (1, 4, 3, 8) + assert key_shape == value_shape == (1, 2, 3, 8) + assert kwargs["enable_gqa"] is True + assert output.shape == (1, 3, 4, 8) diff --git a/tests/diffusion/distributed/test_parallel_state_sp_groups.py b/tests/diffusion/distributed/test_parallel_state_sp_groups.py index 9dda5610b..c12882a22 100644 --- a/tests/diffusion/distributed/test_parallel_state_sp_groups.py +++ b/tests/diffusion/distributed/test_parallel_state_sp_groups.py @@ -40,7 +40,7 @@ def test_set_seq_parallel_pg_uses_explicit_sp_groups(rank, expected_ulysses, exp # tp=2, sp=2 -> SP groups are non-contiguous: [0,2] and [1,3] sp_group_ranks = RankGenerator(2, 2, 1, 1, 1, "tp-sp-pp-cfg-dp").get_ranks("sp") - ulysses_pg, ring_pg = set_seq_parallel_pg( + ulysses_pg, ring_pg, allgather_pg = set_seq_parallel_pg( sp_ulysses_degree=2, sp_ring_degree=1, rank=rank, @@ -50,6 +50,43 @@ def test_set_seq_parallel_pg_uses_explicit_sp_groups(rank, expected_ulysses, exp assert ulysses_pg.ranks == expected_ulysses assert ring_pg.ranks == expected_ring + assert allgather_pg.ranks == [rank] + + +@pytest.mark.cpu +@pytest.mark.parametrize("rank", [0, 1, 2, 3]) +def test_set_seq_parallel_pg_builds_independent_allgather_group(rank, monkeypatch): + created_groups: list[SimpleNamespace] = [] + monkeypatch.setattr(torch.distributed, "new_group", _fake_new_group_factory(created_groups)) + sp_group_ranks = [[0, 2], [1, 3]] + + ulysses_pg, ring_pg, allgather_pg = set_seq_parallel_pg( + sp_ulysses_degree=1, + sp_ring_degree=1, + sp_allgather_degree=2, + rank=rank, + world_size=4, + sp_group_ranks=sp_group_ranks, + ) + + assert ulysses_pg.ranks == [rank] + assert ring_pg.ranks == [rank] + assert allgather_pg.ranks == ([0, 2] if rank % 2 == 0 else [1, 3]) + + +@pytest.mark.cpu +def test_set_seq_parallel_pg_rejects_allgather_with_ulysses_or_ring(monkeypatch): + created_groups: list[SimpleNamespace] = [] + monkeypatch.setattr(torch.distributed, "new_group", _fake_new_group_factory(created_groups)) + + with pytest.raises(ValueError, match="mutually exclusive"): + set_seq_parallel_pg( + sp_ulysses_degree=2, + sp_ring_degree=1, + sp_allgather_degree=2, + rank=0, + world_size=4, + ) @pytest.mark.cpu diff --git a/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py b/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py index 8a35d8aaa..543b93a25 100644 --- a/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py +++ b/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py @@ -45,6 +45,8 @@ def patched_mgr_env(sp_size=1): target = _TRANSFORMER_MODULE patches = [ patch(f"{target}.get_sequence_parallel_world_size", return_value=sp_size), + patch(f"{target}.get_allgather_parallel_world_size", return_value=sp_size), + patch(f"{target}.get_ulysses_parallel_world_size", return_value=1, create=True), patch(f"{target}.get_sequence_parallel_rank", return_value=0), patch(f"{target}.Attention", MockAttention), ] diff --git a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py b/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py index 3f799c038..74fb14086 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py +++ b/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py @@ -3,6 +3,7 @@ from __future__ import annotations import base64 import copy import io +import logging import os import tempfile from pathlib import Path @@ -18,6 +19,8 @@ from tests.helpers.runtime import OmniServer pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] +logger = logging.getLogger(__name__) + MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct" SEED = 42 NUM_INFERENCE_STEPS = 50 @@ -53,11 +56,12 @@ _DEPLOY_CONFIG = { "parallel_config": { "pipeline_parallel_size": 1, "data_parallel_size": 1, - "tensor_parallel_size": 4, + "tensor_parallel_size": 1, "enable_expert_parallel": True, "sequence_parallel_size": 1, "ulysses_degree": 1, "ring_degree": 1, + "allgather_degree": 4, "cfg_parallel_size": 1, "vae_patch_parallel_size": 1, "use_hsdp": False, @@ -83,9 +87,47 @@ def _devices() -> str: def _write_deploy_config(path: Path) -> None: config = copy.deepcopy(_DEPLOY_CONFIG) devices = _devices() + num_devices = len(devices.split(",")) config["stages"][0]["devices"] = devices - config["stages"][0]["parallel_config"]["tensor_parallel_size"] = len(devices.split(",")) - path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False)) + parallel_config = config["stages"][0]["parallel_config"] + if parallel_config.get("allgather_degree", 1) > 1: + parallel_config["allgather_degree"] = num_devices + elif parallel_config.get("tensor_parallel_size", 1) > 1: + parallel_config["tensor_parallel_size"] = num_devices + + sequence_parallel_size = ( + parallel_config.get("allgather_degree", 1) + if parallel_config.get("allgather_degree", 1) > 1 + else parallel_config.get("ulysses_degree", 1) * parallel_config.get("ring_degree", 1) + ) + configured_world_size = ( + parallel_config.get("pipeline_parallel_size", 1) + * parallel_config.get("data_parallel_size", 1) + * parallel_config.get("tensor_parallel_size", 1) + * sequence_parallel_size + * parallel_config.get("cfg_parallel_size", 1) + ) + assert configured_world_size == num_devices, ( + f"Configured parallel world size ({configured_world_size}) does not match " + f"the number of devices ({num_devices}): {parallel_config}" + ) + + config_text = yaml.dump(config, default_flow_style=False, sort_keys=False) + path.write_text(config_text) + logger.info( + "HunyuanImage3 pixel accuracy deploy config:\n" + " path=%s\n" + " CUDA_VISIBLE_DEVICES=%s\n" + " devices=%s\n" + " tensor_parallel_size=%s\n" + " allgather_degree=%s\n%s", + path, + os.environ.get("CUDA_VISIBLE_DEVICES", ""), + devices, + parallel_config.get("tensor_parallel_size", 1), + parallel_config.get("allgather_degree", 1), + config_text, + ) def _run_vllm_omni_hunyuan_image3_online(*, model: str, deploy_config: str, output_path: Path) -> Image.Image: diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index 10f423c84..aa3ab604e 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -336,6 +336,31 @@ class TestStageConfig: assert "ring_degree" not in omega_config.engine_args assert "sequence_parallel_size" not in omega_config.engine_args + def test_to_omegaconf_diffusion_parallel_allgather_override_recomputes_sequence_parallel_size(self): + config = StageConfig( + stage_id=1, + model_stage="diffusion", + stage_type=StageType.DIFFUSION, + yaml_engine_args={ + "parallel_config": { + "sequence_parallel_size": 1, + "ulysses_degree": 1, + "ring_degree": 1, + "allgather_degree": 1, + } + }, + runtime_overrides={ + "allgather_degree": 2, + }, + ) + + omega_config = config.to_omegaconf() + + assert omega_config.engine_args.parallel_config.allgather_degree == 2 + assert omega_config.engine_args.parallel_config.sequence_parallel_size == 2 + assert "allgather_degree" not in omega_config.engine_args + assert "sequence_parallel_size" not in omega_config.engine_args + def test_to_omegaconf_diffusion_parallel_explicit_sequence_parallel_size_is_preserved(self): config = StageConfig( stage_id=1, diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index be89e1cae..cdcbafa79 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -109,6 +109,7 @@ class _ParallelConfigEngineOverrides(TypedDict, total=False): sequence_parallel_size: int ulysses_degree: int ring_degree: int + allgather_degree: int ulysses_mode: str cfg_parallel_size: int vae_patch_parallel_size: int @@ -360,6 +361,7 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): sequence_parallel_size: int = Field(default=1, ge=1, init=False) ulysses_degree: int = Field(default=1, ge=1) ring_degree: int = Field(default=1, ge=1) + allgather_degree: int = Field(default=1, ge=1) ulysses_mode: str = "strict" cfg_parallel_size: int = Field(default=1, ge=1, le=3) vae_patch_parallel_size: int = Field(default=1, ge=1) @@ -370,7 +372,11 @@ class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): hsdp_replicate_size: int = Field(default=1, ge=1) def __post_init__(self) -> None: - self.sequence_parallel_size = self.ulysses_degree * self.ring_degree + self.sequence_parallel_size = ( + self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree + ) + if self.allgather_degree > 1 and (self.ulysses_degree > 1 or self.ring_degree > 1): + raise ValueError("allgather_degree > 1 is mutually exclusive with ulysses_degree/ring_degree > 1") if self.ulysses_mode not in {"strict", "advanced_uaa"}: raise ValueError("ulysses_mode must be 'strict' or 'advanced_uaa'") if self.vae_parallel_mode not in {"tile", "spatial_shard_height", "spatial_shard_width"}: diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 2a8efe824..2f9749420 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -149,14 +149,17 @@ def _apply_diffusion_parallel_runtime_overrides( continue if parallel_config_dict is None: parallel_config_dict = {} - if key in ("ulysses_degree", "ring_degree"): + if key in ("ulysses_degree", "ring_degree", "allgather_degree"): degree_overridden = True parallel_config_dict[key] = runtime_overrides.pop(key) if parallel_config_dict is not None and degree_overridden and not sequence_parallel_explicit: ulysses_degree = parallel_config_dict.get("ulysses_degree") or 1 ring_degree = parallel_config_dict.get("ring_degree") or 1 - parallel_config_dict["sequence_parallel_size"] = ulysses_degree * ring_degree + allgather_degree = parallel_config_dict.get("allgather_degree") or 1 + parallel_config_dict["sequence_parallel_size"] = ( + allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree + ) if parallel_config_dict is not None: engine_args["parallel_config"] = parallel_config_dict @@ -366,6 +369,7 @@ class StageDeployConfig: ulysses_degree: int | None = None ulysses_mode: str | None = None ring_degree: int | None = None + allgather_degree: int | None = None sequence_parallel_size: int | None = None cfg_parallel_size: int | None = None vae_patch_parallel_size: int | None = None diff --git a/vllm_omni/deploy/hunyuan_image3_dit.yaml b/vllm_omni/deploy/hunyuan_image3_dit.yaml index 3c0ba1901..537afadc0 100644 --- a/vllm_omni/deploy/hunyuan_image3_dit.yaml +++ b/vllm_omni/deploy/hunyuan_image3_dit.yaml @@ -24,6 +24,7 @@ stages: sequence_parallel_size: 1 ulysses_degree: 1 ring_degree: 1 + allgather_degree: 1 cfg_parallel_size: 1 vae_patch_parallel_size: 1 use_hsdp: false diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index cd0408f3c..43c6f525c 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -14,6 +14,7 @@ class AttentionBackend(ABC): """Abstract class for diffusion attention backends.""" accept_output_buffer: bool = False + supports_piecewise_spans: bool = False @classmethod def supports_attention_mask(cls) -> bool: @@ -51,6 +52,13 @@ class AttentionBackend(ABC): return (not supported_head_sizes) or head_size in supported_head_sizes +@dataclass(frozen=True, slots=True) +class QueryRange: + local_start: int + local_end: int + global_start: int + + @dataclass class AttentionMetadata: attn_mask: torch.Tensor | None = None @@ -75,6 +83,7 @@ class AttentionMetadata: # Piecewise attention metadata (mixed causal/full masks). # full_attn_spans: per-sample [start, end) spans in global coordinates using full attention. full_attn_spans: list[list[tuple[int, int]]] | None = None + query_ranges: tuple[QueryRange, ...] | None = None T = TypeVar("T", bound=AttentionMetadata) diff --git a/vllm_omni/diffusion/attention/backends/cudnn_attn.py b/vllm_omni/diffusion/attention/backends/cudnn_attn.py index 357a2c0fd..78e428796 100644 --- a/vllm_omni/diffusion/attention/backends/cudnn_attn.py +++ b/vllm_omni/diffusion/attention/backends/cudnn_attn.py @@ -51,7 +51,6 @@ class CuDNNAttentionImpl(AttentionImpl): ) -> None: self.causal = causal self.softmax_scale = softmax_scale - self.requires_gqa = num_heads != num_kv_heads def forward_cuda( self, @@ -64,6 +63,7 @@ class CuDNNAttentionImpl(AttentionImpl): if attn_metadata: attention_mask = _maybe_reshape_attn_mask(query, key, attn_metadata.attn_mask, mask_mode="broadcast_k") + enable_gqa = query.shape[2] != key.shape[2] query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) # Pin cuDNN exclusively. A priority list like [CUDNN, FLASH, MATH] hits a # PyTorch SDPA dispatch quirk: when FLASH rejects a non-None attn_mask, @@ -85,7 +85,7 @@ class CuDNNAttentionImpl(AttentionImpl): dropout_p=0.0, is_causal=self.causal, scale=self.softmax_scale, - enable_gqa=self.requires_gqa, + enable_gqa=enable_gqa, ) except RuntimeError as e: if "No available kernel" not in str(e): @@ -102,6 +102,6 @@ class CuDNNAttentionImpl(AttentionImpl): dropout_p=0.0, is_causal=self.causal, scale=self.softmax_scale, - enable_gqa=self.requires_gqa, + enable_gqa=enable_gqa, ) return output.permute(0, 2, 1, 3) diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index 19ec49417..82d61c25d 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -21,6 +21,7 @@ logger = init_logger(__name__) class FlashAttentionBackend(AttentionBackend): accept_output_buffer: bool = True + supports_piecewise_spans: bool = True @classmethod def supports_attention_mask(cls) -> bool: @@ -82,6 +83,28 @@ class FlashAttentionImpl(AttentionImpl): def _flash_wrapper(q, k, v, *, attn_func, **kwargs): return FlashAttentionImpl._unwrap_flash_output(attn_func(q, k, v, **kwargs)) + @staticmethod + def _flash_varlen_wrapper(q, k, v, *, attn_func, causal, softmax_scale, **kwargs): + """Call a varlen-only FlashAttention backend for a dense segment.""" + del kwargs + batch_size, q_len = q.shape[:2] + k_len = k.shape[1] + cu_seqlens_q = torch.arange(0, (batch_size + 1) * q_len, q_len, dtype=torch.int32, device=q.device) + cu_seqlens_k = torch.arange(0, (batch_size + 1) * k_len, k_len, dtype=torch.int32, device=q.device) + out = attn_func( + q=q.flatten(0, 1), + k=k.flatten(0, 1), + v=v.flatten(0, 1), + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=q_len, + max_seqlen_k=k_len, + causal=causal, + softmax_scale=softmax_scale, + ) + out = FlashAttentionImpl._unwrap_flash_output(out) + return out.reshape(batch_size, q_len, *out.shape[1:]) + def _forward_varlen_masked( self, query: torch.Tensor, @@ -171,6 +194,7 @@ class FlashAttentionImpl(AttentionImpl): from vllm_omni.diffusion.attention.backends.utils.fa import ( HAS_FLASH_ATTN, flash_attn_func, + flash_attn_varlen_func, ) if not HAS_FLASH_ATTN: @@ -186,10 +210,18 @@ class FlashAttentionImpl(AttentionImpl): # Try piecewise attention if full_attn_spans is not None: logger.debug("Using piecewise Flash Attention for mixed causal/full mask") - attn_func = partial( - FlashAttentionImpl._flash_wrapper, - attn_func=flash_attn_func, - ) + if flash_attn_func is not None: + attn_func = partial( + FlashAttentionImpl._flash_wrapper, + attn_func=flash_attn_func, + ) + elif flash_attn_varlen_func is not None: + attn_func = partial( + FlashAttentionImpl._flash_varlen_wrapper, + attn_func=flash_attn_varlen_func, + ) + else: + raise ImportError("Piecewise FlashAttention requires a dense or varlen FlashAttention function") return piecewise_attn( query, @@ -198,6 +230,7 @@ class FlashAttentionImpl(AttentionImpl): full_attn_spans, self.softmax_scale, attn_func, + query_ranges=attn_metadata.query_ranges, ) if attention_mask is not None and torch.any(~attention_mask): diff --git a/vllm_omni/diffusion/attention/backends/flash_attn_hub.py b/vllm_omni/diffusion/attention/backends/flash_attn_hub.py index bce607888..73faa57a9 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn_hub.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn_hub.py @@ -12,9 +12,7 @@ from vllm_omni.diffusion.attention.backends.abstract import ( AttentionImpl, AttentionMetadata, ) -from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import ( - piecewise_attn, -) +from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import piecewise_attn logger = init_logger(__name__) @@ -78,6 +76,7 @@ def _run_varlen_dense( class FlashAttentionHubBackend(AttentionBackend): accept_output_buffer: bool = True + supports_piecewise_spans: bool = True @classmethod def supports_attention_mask(cls) -> bool: @@ -214,6 +213,7 @@ class FlashAttentionHubImpl(AttentionImpl): full_attn_spans, self.softmax_scale, attn_func, + query_ranges=attn_metadata.query_ranges, ) if attention_mask is not None and torch.any(~attention_mask): @@ -243,6 +243,7 @@ class FlashAttentionHubImpl(AttentionImpl): class FlashAttention3HubBackend(AttentionBackend): accept_output_buffer: bool = True + supports_piecewise_spans: bool = True @classmethod def supports_attention_mask(cls) -> bool: @@ -379,6 +380,7 @@ class FlashAttention3HubImpl(AttentionImpl): full_attn_spans, self.softmax_scale, attn_func, + query_ranges=attn_metadata.query_ranges, ) if attention_mask is not None and torch.any(~attention_mask): diff --git a/vllm_omni/diffusion/attention/backends/sdpa.py b/vllm_omni/diffusion/attention/backends/sdpa.py index f8d607e48..aec8eb50b 100644 --- a/vllm_omni/diffusion/attention/backends/sdpa.py +++ b/vllm_omni/diffusion/attention/backends/sdpa.py @@ -11,6 +11,7 @@ from vllm_omni.diffusion.attention.backends.abstract import ( AttentionImpl, AttentionMetadata, ) +from vllm_omni.diffusion.attention.backends.utils.attn_runtime_selector import can_sdpa_use_fused_gqa logger = init_logger(__name__) @@ -85,7 +86,6 @@ class SDPAImpl(AttentionImpl): ) -> None: self.causal = causal self.softmax_scale = softmax_scale - self.requires_gqa = num_heads != num_kv_heads if backend_kwargs: logger.warning("SDPAImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys())) @@ -103,7 +103,24 @@ class SDPAImpl(AttentionImpl): if attn_metadata: attention_mask = _maybe_reshape_attn_mask(query, key, attn_metadata.attn_mask, mask_mode=mask_mode) + enable_gqa = query.shape[2] != key.shape[2] query, key, value = (x.permute(0, 2, 1, 3) for x in (query, key, value)) + # Only the PyTorch SDPA backend needs this dispatch check. If SDPA + # cannot select a fused GQA kernel for the runtime shape/mask, expand + # K/V locally so it can use the better-supported equal-head path. + if enable_gqa and not can_sdpa_use_fused_gqa(query, key, value, attention_mask, self.causal): + if query.shape[1] % key.shape[1] != 0: + raise ValueError( + "GQA requires query heads to be a multiple of KV heads, " + f"got q_heads={query.shape[1]} and kv_heads={key.shape[1]}." + ) + repeat_num = query.shape[1] // key.shape[1] + key = key.repeat_interleave(repeat_num, dim=1) + value = value.repeat_interleave(repeat_num, dim=1) + enable_gqa = False + logger.debug( + "CUDA SDPA cannot use a fused native-GQA kernel for this shape; expanding K/V heads before SDPA." + ) output = torch.nn.functional.scaled_dot_product_attention( query, key, @@ -112,7 +129,7 @@ class SDPAImpl(AttentionImpl): dropout_p=0.0, is_causal=self.causal, scale=self.softmax_scale, - enable_gqa=self.requires_gqa, + enable_gqa=enable_gqa, ) out = output.permute(0, 2, 1, 3) return out diff --git a/vllm_omni/diffusion/attention/backends/utils/attn_runtime_selector.py b/vllm_omni/diffusion/attention/backends/utils/attn_runtime_selector.py new file mode 100644 index 000000000..84170dcd7 --- /dev/null +++ b/vllm_omni/diffusion/attention/backends/utils/attn_runtime_selector.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Runtime kernel-selection helpers for attention backends.""" + +import torch + + +def can_sdpa_use_fused_gqa( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor | None, + causal: bool, +) -> bool: + """Return whether PyTorch SDPA can use a fused kernel for this GQA call. + + This check is specific to the SDPA backend. Native GQA backends such as + FlashAttention should keep compressed K/V heads and do not use this helper. + """ + if not query.is_cuda: + return False + + try: + params = torch.backends.cuda.SDPAParams( + query, + key, + value, + attention_mask, + 0.0, + causal, + True, + ) + can_use_flash = torch.backends.cuda.can_use_flash_attention(params) + can_use_cudnn = getattr(torch.backends.cuda, "can_use_cudnn_attention", None) + return can_use_flash or (can_use_cudnn is not None and can_use_cudnn(params)) + except (AttributeError, RuntimeError, TypeError): + # Older PyTorch versions may not expose GQA-aware SDPA capability checks. + # Returning False lets the SDPA backend expand K/V and use the mature + # equal-head fused-kernel path instead. + return False diff --git a/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py b/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py index fbbd3c8e0..5d83d8dbf 100644 --- a/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py +++ b/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py @@ -9,17 +9,24 @@ follows FlashAttention's bottom-right convention (``K[:e]`` is attended by Per segment: - causal segment ``[s, e)``: ``attn(Q[:, s:e], K[:, :e], V[:, :e], causal=True)`` - - full-attn span ``[a, e)``: ``attn(Q[:, a:e], K[:, :e], V[:, :e], causal=False)`` + - full-attn span ``[a, b)`` intersecting the query range at ``[s, e)``: + ``attn(Q[:, s:e], K[:, :b], V[:, :b], causal=False)`` """ from __future__ import annotations -from typing import Literal, NamedTuple +from typing import TYPE_CHECKING, Literal, NamedTuple + +import torch + +if TYPE_CHECKING: + from vllm_omni.diffusion.attention.backends.abstract import QueryRange class Segment(NamedTuple): - start: int - end: int + q_start: int + q_end: int + kv_end: int mode: Literal["causal", "full"] @@ -30,30 +37,33 @@ def build_segments(full_attn_spans, query_offset, query_len): query_len: length of the query return: - List[Segment] in global coordinates, clipped to [query_offset, query_offset + query_len) + List[Segment] in global coordinates, clipped to + [query_offset, query_offset + query_len). Full-attention segments retain + the original span end as kv_end so a local query shard can attend past + its own boundary. """ q_start = query_offset q_end = query_offset + query_len - segs: list[Segment] = [] + segments: list[Segment] = [] cur = q_start - for a, e in full_attn_spans: + for span_start, span_end in full_attn_spans: # clip span to query range - a_clipped = max(a, q_start) - e_clipped = min(e, q_end) - if a_clipped >= e_clipped: + overlap_start = max(span_start, q_start) + overlap_end = min(span_end, q_end) + if overlap_start >= overlap_end: continue - if cur < a_clipped: - segs.append(Segment(cur, a_clipped, "causal")) - segs.append(Segment(a_clipped, e_clipped, "full")) - cur = e_clipped + if cur < overlap_start: + segments.append(Segment(cur, overlap_start, overlap_start, "causal")) + segments.append(Segment(overlap_start, overlap_end, span_end, "full")) + cur = overlap_end if cur < q_end: - segs.append(Segment(cur, q_end, "causal")) + segments.append(Segment(cur, q_end, q_end, "causal")) - return segs + return segments def _check_homogeneous( @@ -76,23 +86,40 @@ def piecewise_attn( full_attn_spans: list[list[tuple[int, int]]], softmax_scale: float, attn_func, + query_ranges: tuple[QueryRange, ...] | None = None, ): - B, Sq, H, D = query.shape _check_homogeneous(full_attn_spans) - - query_offset = key.shape[1] - Sq spans = full_attn_spans[0] - out = query.new_zeros(B, Sq, H, D) + if query_ranges is None: + query_len = query.shape[1] + ranges = ((0, query_len, key.shape[1] - query_len),) + else: + ranges = tuple((r.local_start, r.local_end, r.global_start) for r in query_ranges) - for s, e, mode in build_segments(spans, query_offset, Sq): - q_s = s - query_offset - q_e = e - query_offset - out_seg = attn_func( - query[:, q_s:q_e], - key[:, :e], - value[:, :e], - causal=(mode == "causal"), - softmax_scale=softmax_scale, - ) - out[:, q_s:q_e] = out_seg - return out + outputs = [] + covered = 0 + for local_start, local_end, global_start in ranges: + query_len = local_end - local_start + if local_start != covered or query_len < 0: + raise ValueError("query_ranges must cover local query contiguously") + for segment in build_segments(spans, global_start, query_len): + q_start = local_start + segment.q_start - global_start + q_end = local_start + segment.q_end - global_start + outputs.append( + attn_func( + query[:, q_start:q_end], + key[:, : segment.kv_end], + value[:, : segment.kv_end], + causal=(segment.mode == "causal"), + softmax_scale=softmax_scale, + ) + ) + covered = local_end + + if covered != query.shape[1]: + raise ValueError("query_ranges must cover the full local query") + if not outputs: + return torch.empty_like(query) + if len(outputs) == 1: + return outputs[0] + return torch.cat(outputs, dim=1) diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 0ea122273..a11aecbb8 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -140,6 +140,7 @@ class Attention(nn.Module): scatter_idx=scatter_idx, gather_idx=gather_idx, use_sync=use_sync, + causal=causal, ) # Fallback strategy when SP is not active (outside sharded regions) self._no_parallel_strategy = NoParallelAttention() @@ -282,6 +283,8 @@ class Attention(nn.Module): return out def _run_local_attention(self, query, key, value, attn_metadata): + self._assert_piecewise_compatible(attn_metadata) + if query.dtype == torch.float32: logger.warning_once( f"Only SDPA supports float32. Overriding user config {type(self.attention)} " @@ -292,6 +295,20 @@ class Attention(nn.Module): # Fallback to standard attention return self.attention.forward(query, key, value, attn_metadata) + def _assert_piecewise_compatible(self, attn_metadata: AttentionMetadata | None) -> None: + if attn_metadata is None or attn_metadata.full_attn_spans is None: + return + if attn_metadata.attn_mask is not None and attn_metadata.attn_mask.ndim == 4: + return + backend_name = self.attn_backend.get_name() + if not self.attn_backend.supports_piecewise_spans: + raise ValueError( + f"Attention backend '{backend_name}' does not support " + f"piecewise attention (full_attn_spans without a 4D attn_mask). " + f"Use a Flash backend (FLASH_ATTN / FLASH_ATTN_HUB / FLASH_ATTN_3_HUB), " + f"or provide a 4D attn_mask that encodes the mixed causal/full pattern." + ) + def _run_ring_attention(self, query, key, value, attn_metadata): # Delegate to RingParallelAttention strategy if available if self.ring_runner is not None: diff --git a/vllm_omni/diffusion/attention/parallel/allgather_kv.py b/vllm_omni/diffusion/attention/parallel/allgather_kv.py new file mode 100644 index 000000000..51a463bb1 --- /dev/null +++ b/vllm_omni/diffusion/attention/parallel/allgather_kv.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING + +import torch + +from vllm_omni.diffusion.attention.backends.abstract import QueryRange +from vllm_omni.diffusion.attention.parallel.base import ParallelAttentionContext +from vllm_omni.diffusion.distributed.group_coordinator import ( + SequenceParallelGroupCoordinator, +) + +if TYPE_CHECKING: + from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata + + +@dataclass(frozen=True, slots=True) +class _AllGatherKVCtx(ParallelAttentionContext): + pass + + +class AllGatherKVParallelAttention: + """Compute local Q against AllGathered K/V.""" + + def __init__( + self, + sp_group: SequenceParallelGroupCoordinator, + ) -> None: + self._sp_group = sp_group + self._allgather_group = sp_group.allgather_group + self._sp_size = sp_group.allgather_world_size + self._sp_rank = sp_group.allgather_rank + + @property + def enabled(self) -> bool: + return True + + @property + def name(self) -> str: + return "allgather_kv" + + def pre_attention( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None, + ): + joint_q = joint_k = joint_v = None + joint_strategy = "front" + if attn_metadata is not None: + joint_q = attn_metadata.joint_query + joint_k = attn_metadata.joint_key + joint_v = attn_metadata.joint_value + joint_strategy = attn_metadata.joint_strategy or "front" + if joint_strategy not in {"front", "rear"}: + raise ValueError(f"Unsupported joint_strategy: {joint_strategy!r}") + + k_img_full = self._sp_group.all_gather(key, dim=1, group=self._allgather_group) + v_img_full = self._sp_group.all_gather(value, dim=1, group=self._allgather_group) + + if joint_k is not None: + if joint_k.shape[2] != key.shape[2]: + raise ValueError( + "AllGather-KV SP expects joint_k to be GQA-compressed with the " + f"same num_kv_heads as the image shard, got joint_k.heads=" + f"{joint_k.shape[2]} vs key.heads={key.shape[2]}." + ) + if joint_strategy == "front": + k_full = torch.cat([joint_k, k_img_full], dim=1) + v_full = torch.cat([joint_v, v_img_full], dim=1) + else: + k_full = torch.cat([k_img_full, joint_k], dim=1) + v_full = torch.cat([v_img_full, joint_v], dim=1) + else: + k_full, v_full = k_img_full, v_img_full + + joint_len = joint_q.shape[1] if joint_q is not None else 0 + logical_q_full_len = joint_len + k_img_full.shape[1] + query_global_offset = k_full.shape[1] - logical_q_full_len + if query_global_offset < 0: + raise ValueError( + "AllGather-KV SP global K length is shorter than the logical Q length: " + f"k_len={k_full.shape[1]}, logical_q_len={logical_q_full_len}." + ) + + if joint_q is not None: + if joint_strategy == "front": + q_local = torch.cat([joint_q, query], dim=1) + else: + q_local = torch.cat([query, joint_q], dim=1) + else: + q_local = query + + attn_metadata = self._slice_attn_metadata_for_local_query( + attn_metadata, + q_local_len=q_local.shape[1], + img_seq_local=query.shape[1], + img_seq_full=k_img_full.shape[1], + joint_len=joint_len, + joint_strategy=joint_strategy, + query_global_offset=query_global_offset, + ) + + return q_local, k_full, v_full, attn_metadata, _AllGatherKVCtx(name=self.name) + + def _slice_attn_metadata_for_local_query( + self, + attn_metadata: AttentionMetadata | None, + *, + q_local_len: int, + img_seq_local: int, + img_seq_full: int, + joint_len: int, + joint_strategy: str, + query_global_offset: int, + ) -> AttentionMetadata | None: + """Slice global query metadata for the local query shard.""" + if attn_metadata is None: + return None + + img_start = self._sp_rank * img_seq_local + img_end = img_start + img_seq_local + if img_end > img_seq_full: + raise ValueError( + "AllGather-KV SP local image query range exceeds gathered image length: " + f"rank={self._sp_rank}, img_start={img_start}, img_end={img_end}, img_seq_full={img_seq_full}." + ) + + if joint_strategy == "front": + ranges = ( + *((QueryRange(0, joint_len, query_global_offset),) if joint_len else ()), + QueryRange(joint_len, q_local_len, query_global_offset + joint_len + img_start), + ) + else: + ranges = ( + QueryRange(0, img_seq_local, query_global_offset + img_start), + *((QueryRange(img_seq_local, q_local_len, query_global_offset + img_seq_full),) if joint_len else ()), + ) + + if attn_metadata.attn_mask is None: + return replace(attn_metadata, query_ranges=ranges) + + mask = attn_metadata.attn_mask + if mask.ndim != 4: + return replace(attn_metadata, query_ranges=ranges) + + if mask.shape[-2] == q_local_len: + local_mask = mask + else: + q_full_len = joint_len + img_seq_full + if mask.shape[-2] != q_full_len: + raise ValueError( + "AllGather-KV SP received an attention mask with incompatible Q length: " + f"mask_q={mask.shape[-2]}, expected local_q={q_local_len} or full_q={q_full_len} " + f"(joint_len={joint_len}, img_seq_local={img_seq_local}, img_seq_full={img_seq_full})." + ) + + parts = [ + mask[..., local_start : local_start + (r.local_end - r.local_start), :] + for r in ranges + for local_start in (r.global_start - query_global_offset,) + ] + local_mask = torch.cat(parts, dim=-2) + + if local_mask.shape[-2] != q_local_len: + raise ValueError( + "AllGather-KV SP produced an attention mask with incompatible local Q length: " + f"mask_q={local_mask.shape[-2]}, q_local={q_local_len}." + ) + return replace(attn_metadata, attn_mask=local_mask.contiguous(), query_ranges=ranges) + + def post_attention( + self, + attn_output: torch.Tensor, + ctx: ParallelAttentionContext | None, + ) -> torch.Tensor: + return attn_output diff --git a/vllm_omni/diffusion/attention/parallel/factory.py b/vllm_omni/diffusion/attention/parallel/factory.py index ce95d409a..76f7a78da 100644 --- a/vllm_omni/diffusion/attention/parallel/factory.py +++ b/vllm_omni/diffusion/attention/parallel/factory.py @@ -5,10 +5,19 @@ from __future__ import annotations from vllm.logger import init_logger -from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention, ParallelAttentionStrategy +from vllm_omni.diffusion.attention.parallel.allgather_kv import ( + AllGatherKVParallelAttention, +) +from vllm_omni.diffusion.attention.parallel.base import ( + NoParallelAttention, + ParallelAttentionStrategy, +) from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention from vllm_omni.diffusion.attention.parallel.ulysses import UlyssesParallelAttention -from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size, get_sp_group +from vllm_omni.diffusion.distributed.parallel_state import ( + get_sequence_parallel_world_size, + get_sp_group, +) from vllm_omni.diffusion.forward_context import get_forward_context logger = init_logger(__name__) @@ -19,6 +28,7 @@ def build_parallel_attention_strategy( scatter_idx: int, gather_idx: int, use_sync: bool, + causal: bool = False, ) -> ParallelAttentionStrategy: """Select a parallel attention strategy based on current diffusion config. @@ -36,6 +46,7 @@ def build_parallel_attention_strategy( ulysses_degree = getattr(p, "ulysses_degree", 1) ring_degree = getattr(p, "ring_degree", 1) + allgather_degree = getattr(p, "allgather_degree", 1) try: sp_group = get_sp_group() @@ -44,13 +55,26 @@ def build_parallel_attention_strategy( return NoParallelAttention() except Exception as e: # Log warning if SP is configured but group is not available - if ulysses_degree > 1 or ring_degree > 1: + if ulysses_degree > 1 or ring_degree > 1 or allgather_degree > 1: logger.warning( - f"SP configured (ulysses={ulysses_degree}, ring={ring_degree}) but SP group not available: {e}. " + f"SP configured (ulysses={ulysses_degree}, ring={ring_degree}, " + f"allgather={allgather_degree}) but SP group not available: {e}. " f"Falling back to NoParallelAttention. This may cause incorrect results." ) return NoParallelAttention() + if allgather_degree > 1: + if causal: + raise ValueError("AllGather-KV SP only supports non-causal attention.") + if ulysses_degree > 1 or ring_degree > 1: + raise ValueError( + f"AllGather-KV SP is mutually exclusive with Ulysses/Ring in v1 " + f"(got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, " + f"allgather_degree={allgather_degree})." + ) + logger.debug(f"Using AllGatherKVParallelAttention (allgather_degree={allgather_degree})") + return AllGatherKVParallelAttention(sp_group=sp_group) + # Ulysses (or Hybrid Ulysses+Ring) if ulysses_degree > 1: logger.debug(f"Using UlyssesParallelAttention (ulysses_degree={ulysses_degree})") diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index e131b833b..2d8ed2410 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -156,7 +156,9 @@ class DiffusionParallelConfig: """Enable expert parallelism for MoE layers (TP is still used for non-MoE layers).""" sequence_parallel_size: int | None = None - """Number of sequence parallel groups. sequence_parallel_size = ring_degree * ulysses_degree""" + """Number of sequence parallel groups. + sequence_parallel_size = ulysses_degree * ring_degree, or allgather_degree + when AllGather-KV is enabled.""" ulysses_degree: int = 1 """Number of GPUs used for ulysses sequence parallelism.""" @@ -164,6 +166,9 @@ class DiffusionParallelConfig: ring_degree: int = 1 """Number of GPUs used for ring sequence parallelism.""" + allgather_degree: int = 1 + """Number of GPUs used for AllGather-KV sequence parallelism (causal=False only).""" + ulysses_mode: str = "strict" """Ulysses sequence-parallel mode. @@ -225,6 +230,7 @@ class DiffusionParallelConfig: assert self.sequence_parallel_size > 0, "Sequence parallel size must be > 0" assert self.ulysses_degree > 0, "Ulysses degree must be > 0" assert self.ring_degree > 0, "Ring degree must be > 0" + assert self.allgather_degree > 0, "AllGather degree must be > 0" assert self.cfg_parallel_size > 0, "CFG parallel size must be > 0" assert self.cfg_parallel_size in [1, 2, 3], ( f"CFG parallel size must be 1, 2, or 3, but got {self.cfg_parallel_size}" @@ -234,9 +240,17 @@ class DiffusionParallelConfig: "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, " f"but got {self.vae_parallel_mode!r}." ) - assert self.sequence_parallel_size == self.ulysses_degree * self.ring_degree, ( - "Sequence parallel size must be equal to the product of ulysses degree and ring degree," - f" but got {self.sequence_parallel_size} != {self.ulysses_degree} * {self.ring_degree}" + if self.allgather_degree > 1: + assert self.ulysses_degree == 1 and self.ring_degree == 1, ( + "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. " + f"Got ulysses_degree={self.ulysses_degree}, ring_degree={self.ring_degree}, " + f"allgather_degree={self.allgather_degree}." + ) + expected_sp_size = ( + self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree + ) + assert self.sequence_parallel_size == expected_sp_size, ( + f"Sequence parallel size must be {expected_sp_size}, but got {self.sequence_parallel_size}" ) assert self.ulysses_mode in {"strict", "advanced_uaa"}, ( f"ulysses_mode must be one of {{'strict','advanced_uaa'}}, but got {self.ulysses_mode!r}." @@ -250,15 +264,16 @@ class DiffusionParallelConfig: def __post_init__(self) -> None: if self.sequence_parallel_size is None: - self.sequence_parallel_size = self.ulysses_degree * self.ring_degree + self.sequence_parallel_size = ( + self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree + ) # Calculate world_size from other parallelism dimensions other_parallel_world_size = ( self.pipeline_parallel_size * self.data_parallel_size * self.tensor_parallel_size - * self.ulysses_degree - * self.ring_degree + * self.sequence_parallel_size * self.cfg_parallel_size ) diff --git a/vllm_omni/diffusion/distributed/group_coordinator.py b/vllm_omni/diffusion/distributed/group_coordinator.py index 0a67fd941..14847007f 100644 --- a/vllm_omni/diffusion/distributed/group_coordinator.py +++ b/vllm_omni/diffusion/distributed/group_coordinator.py @@ -199,9 +199,15 @@ class GroupCoordinator: return input_ def all_gather( - self, input_: torch.Tensor, dim: int = 0, separate_tensors: bool = False + self, + input_: torch.Tensor, + dim: int = 0, + separate_tensors: bool = False, + group: ProcessGroup | None = None, ) -> torch.Tensor | list[torch.Tensor]: - world_size = self.world_size + if group is None: + group = self.device_group + world_size = torch.distributed.get_world_size(group) # Bypass the function if we are using only 1 GPU. if world_size == 1: return input_ @@ -214,7 +220,7 @@ class GroupCoordinator: input_size[0] *= world_size output_tensor = torch.empty(input_size, dtype=input_.dtype, device=input_.device) # All-gather. - torch.distributed.all_gather_into_tensor(output_tensor, input_.contiguous(), group=self.device_group) + torch.distributed.all_gather_into_tensor(output_tensor, input_.contiguous(), group=group) if dim != 0: input_size[0] //= world_size output_tensor = output_tensor.reshape( @@ -1002,6 +1008,7 @@ class SequenceParallelGroupCoordinator(GroupCoordinator): ulysses_group = kwargs.get("ulysses_group", None) ring_group = kwargs.get("ring_group", None) + allgather_group = kwargs.get("allgather_group", None) if ulysses_group is None: raise RuntimeError( "Please pass argument 'ulysses_group' when calling init func of SequenceParallelGroupCoordinator" @@ -1010,10 +1017,17 @@ class SequenceParallelGroupCoordinator(GroupCoordinator): raise RuntimeError( "Please pass argument 'ring_group' when calling init func of SequenceParallelGroupCoordinator" ) + if allgather_group is None: + raise RuntimeError( + "Please pass argument 'allgather_group' when calling init func of SequenceParallelGroupCoordinator" + ) self.ulysses_group = ulysses_group self.ring_group = ring_group + self.allgather_group = allgather_group self.ulysses_world_size = torch.distributed.get_world_size(self.ulysses_group) self.ulysses_rank = torch.distributed.get_rank(self.ulysses_group) self.ring_world_size = torch.distributed.get_world_size(self.ring_group) self.ring_rank = torch.distributed.get_rank(self.ring_group) + self.allgather_world_size = torch.distributed.get_world_size(self.allgather_group) + self.allgather_rank = torch.distributed.get_rank(self.allgather_group) diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 7b38d3775..4b531fec9 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/vllm_omni/diffusion/distributed/parallel_state.py @@ -303,6 +303,14 @@ def get_ring_parallel_rank(): return get_sp_group().ring_rank +def get_allgather_parallel_world_size(): + return get_sp_group().allgather_world_size + + +def get_allgather_parallel_rank(): + return get_sp_group().allgather_rank + + def get_expert_parallel_group_ranks() -> list[list[int]]: assert vllm_parallel_state._EP is not None, "expert parallel group is not initialized" assert _EXPERT_PARALLEL_GROUP_RANKS is not None, "expert parallel group ranks are not initialized" @@ -548,13 +556,12 @@ def set_seq_parallel_pg( world_size: int, use_ulysses_low: bool = True, sp_group_ranks: list[list[int]] | None = None, -) -> tuple[torch.distributed.ProcessGroup, torch.distributed.ProcessGroup]: + sp_allgather_degree: int = 1, +) -> tuple[torch.distributed.ProcessGroup, torch.distributed.ProcessGroup, torch.distributed.ProcessGroup]: """ - Initialize sequence-parallel Ulysses and Ring process groups. + Initialize Ulysses, Ring, and AllGather-KV process groups. - This builds sequence-parallel (SP) subgroups inside each data-parallel (DP) - slice. The SP group size is sp_ulysses_degree * sp_ring_degree, and - world_size must be divisible by that size. + AllGather-KV is mutually exclusive with Ulysses and Ring. Args: sp_ulysses_degree: Size of each Ulysses subgroup. @@ -569,7 +576,7 @@ def set_seq_parallel_pg( ranges. Returns: - ulyssess_pg (torch.distributed.ProcessGroup): The Ulysses process group + ulysses_pg (torch.distributed.ProcessGroup): The Ulysses process group for this rank. ring_pg (torch.distributed.ProcessGroup): The Ring process group for this rank. @@ -586,6 +593,31 @@ def set_seq_parallel_pg( - If sp_group_ranks is None, groups are auto-generated within each DP slice using offsets of size sp_size. """ + if sp_allgather_degree > 1: + if sp_ulysses_degree > 1 or sp_ring_degree > 1: + raise ValueError("AllGather-KV is mutually exclusive with Ulysses and Ring") + sp_size = sp_allgather_degree + if sp_group_ranks is None: + sp_group_ranks = [list(range(offset, offset + sp_size)) for offset in range(0, world_size, sp_size)] + if len(sp_group_ranks) * sp_size != world_size or any(len(ranks) != sp_size for ranks in sp_group_ranks): + raise ValueError(f"Invalid sp_group_ranks: expected {world_size // sp_size} groups of size {sp_size}.") + + ulysses_pg = ring_pg = allgather_pg = None + for group_ranks in sp_group_ranks: + group = torch.distributed.new_group(group_ranks) + if rank in group_ranks: + allgather_pg = group + for singleton_rank in range(world_size): + group = torch.distributed.new_group([singleton_rank]) + if rank == singleton_rank: + ulysses_pg = group + for singleton_rank in range(world_size): + group = torch.distributed.new_group([singleton_rank]) + if rank == singleton_rank: + ring_pg = group + assert ulysses_pg is not None and ring_pg is not None and allgather_pg is not None + return ulysses_pg, ring_pg, allgather_pg + sp_size = sp_ring_degree * sp_ulysses_degree dp_size = world_size // sp_size @@ -619,7 +651,7 @@ def set_seq_parallel_pg( ulysses_ranks = group_ranks[i * sp_ulysses_degree : (i + 1) * sp_ulysses_degree] group = torch.distributed.new_group(ulysses_ranks) if rank in ulysses_ranks: - ulyssess_pg = group + ulysses_pg = group local_ulysses = list(ulysses_ranks) for i in range(num_ring_pgs): ring_ranks = group_ranks[i::num_ring_pgs] @@ -639,7 +671,7 @@ def set_seq_parallel_pg( ulysses_ranks = group_ranks[i::num_ulysses_pgs] group = torch.distributed.new_group(ulysses_ranks) if rank in ulysses_ranks: - ulyssess_pg = group + ulysses_pg = group local_ulysses = list(ulysses_ranks) if local_sp_group is not None: logger.info( @@ -662,7 +694,7 @@ def set_seq_parallel_pg( ) group = torch.distributed.new_group(ulysses_ranks) if rank in ulysses_ranks: - ulyssess_pg = group + ulysses_pg = group for i in range(num_ring_pgs): ring_ranks = list(range(i + offset, sp_size + offset, num_ring_pgs)) @@ -683,9 +715,15 @@ def set_seq_parallel_pg( ulysses_ranks = list(range(i + offset, sp_size + offset, num_ulysses_pgs)) group = torch.distributed.new_group(ulysses_ranks) if rank in ulysses_ranks: - ulyssess_pg = group + ulysses_pg = group - return ulyssess_pg, ring_pg + allgather_pg = None + for singleton_rank in range(world_size): + group = torch.distributed.new_group([singleton_rank]) + if rank == singleton_rank: + allgather_pg = group + assert allgather_pg is not None + return ulysses_pg, ring_pg, allgather_pg def initialize_model_parallel( @@ -694,6 +732,7 @@ def initialize_model_parallel( sequence_parallel_size: int | None = None, ulysses_degree: int = 1, ring_degree: int = 1, + allgather_degree: int = 1, tensor_parallel_size: int = 1, pipeline_parallel_size: int = 1, fully_shard_degree: int = 1, @@ -710,9 +749,12 @@ def initialize_model_parallel( data_parallel_size: number of data parallelism groups. cfg_parallel_size: number of GPUs used for Classifier Free Guidance (CFG) parallelism. sequence_parallel_size: number of GPUs used for sequence parallelism. - sequence_parallel_size = ulysses_degree * ring_degree + Uses allgather_degree when AllGather-KV is enabled, otherwise + ulysses_degree * ring_degree. ulysses_degree: number of GPUs used for ulysses sequence parallelism. ring_degree: number of GPUs used for ring sequence parallelism. + allgather_degree: number of GPUs used for AllGather-KV sequence parallelism + (causal=False only). Mutually exclusive with ulysses/ring in v1. tensor_parallel_size: number of GPUs used for tensor parallelism. pipeline_parallel_size: number of GPUs used for pipeline parallelism. fully_shard_degree: number of GPUs used for fully sharded data parallelism (HSDP shard dimension). @@ -749,16 +791,22 @@ def initialize_model_parallel( world_size: int = torch.distributed.get_world_size() backend = backend or torch.distributed.get_backend(get_world_group().device_group) - if sequence_parallel_size is None: - sequence_parallel_size = ring_degree * ulysses_degree - logger.info( - f"sequence_parallel_size is not provided, using ring_degree * ulysses_degree = {sequence_parallel_size}" + if allgather_degree > 1: + assert ulysses_degree == 1 and ring_degree == 1, ( + "AllGather-KV (allgather_degree>1) is mutually exclusive with Ulysses/Ring in v1. " + f"Got ulysses_degree={ulysses_degree}, ring_degree={ring_degree}, " + f"allgather_degree={allgather_degree}." ) - if sequence_parallel_size != ring_degree * ulysses_degree: + expected_sequence_parallel_size = allgather_degree if allgather_degree > 1 else ring_degree * ulysses_degree + if sequence_parallel_size is None: + sequence_parallel_size = expected_sequence_parallel_size + logger.info("sequence_parallel_size is not provided, using %d", sequence_parallel_size) + + if sequence_parallel_size != expected_sequence_parallel_size: raise ValueError( - "sequence_parallel_size is not equal to ring_degree * ulysses_degree," - f" but got {sequence_parallel_size} != {ring_degree} * {ulysses_degree}" + f"sequence_parallel_size must be {expected_sequence_parallel_size} for the configured SP mode, " + f"but got {sequence_parallel_size}" ) dit_parallel_size = ( @@ -835,9 +883,10 @@ def initialize_model_parallel( global _SP assert _SP is None, "sequence parallel group is already initialized" - ulysses_pg, ring_pg = set_seq_parallel_pg( + ulysses_pg, ring_pg, allgather_pg = set_seq_parallel_pg( sp_ulysses_degree=ulysses_degree, sp_ring_degree=ring_degree, + sp_allgather_degree=allgather_degree, rank=get_world_group().rank_in_group, world_size=dit_parallel_size, sp_group_ranks=sp_group_ranks, @@ -849,6 +898,7 @@ def initialize_model_parallel( parallel_mode="sequence", ulysses_group=ulysses_pg, ring_group=ring_pg, + allgather_group=allgather_pg, ) if use_moe_parallel_mapping: # Diffusion normally uses its own SP group. Map it to vLLM PCP only for diff --git a/vllm_omni/diffusion/distributed/sp_plan.py b/vllm_omni/diffusion/distributed/sp_plan.py index ee7b1d66b..b8d248865 100644 --- a/vllm_omni/diffusion/distributed/sp_plan.py +++ b/vllm_omni/diffusion/distributed/sp_plan.py @@ -77,6 +77,7 @@ class SequenceParallelConfig: ulysses_degree: int = 1 ring_degree: int = 1 + allgather_degree: int = 1 convert_to_fp32: bool = True # Internal state - populated by setup() @@ -85,18 +86,18 @@ class SequenceParallelConfig: _device: torch.device | None = None def __post_init__(self) -> None: - if self.ulysses_degree < 1 or self.ring_degree < 1: - raise ValueError("`ulysses_degree` and `ring_degree` must be >= 1.") + if self.ulysses_degree < 1 or self.ring_degree < 1 or self.allgather_degree < 1: + raise ValueError("SP degrees must be >= 1.") - if self.ulysses_degree == 1 and self.ring_degree == 1: - raise ValueError( - "At least one of `ulysses_degree` or `ring_degree` must be > 1 to use sequence parallelism." - ) + if self.allgather_degree > 1 and (self.ulysses_degree > 1 or self.ring_degree > 1): + raise ValueError("AllGather-KV is mutually exclusive with Ulysses and Ring.") + if self.ulysses_degree == self.ring_degree == self.allgather_degree == 1: + raise ValueError("At least one SP degree must be > 1.") @property def sequence_parallel_size(self) -> int: """Total sequence parallel world size.""" - return self.ulysses_degree * self.ring_degree + return self.allgather_degree if self.allgather_degree > 1 else self.ulysses_degree * self.ring_degree def get_world_size(self) -> int: """Get the sequence parallel world size from parallel state. @@ -178,13 +179,12 @@ class SequenceParallelConfig: self._world_size = world_size self._device = device - expected_sp_size = self.ulysses_degree * self.ring_degree + expected_sp_size = self.sequence_parallel_size actual_sp_size = self.get_world_size() if expected_sp_size != actual_sp_size: raise ValueError( - f"Configuration mismatch: ulysses_degree ({self.ulysses_degree}) * " - f"ring_degree ({self.ring_degree}) = {expected_sp_size}, but " + f"Configuration mismatch: expected SP size {expected_sp_size}, but " f"actual sequence parallel world size is {actual_sp_size}." ) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index cae7913f2..5f72a607d 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -62,6 +62,7 @@ from vllm_omni.diffusion.attention.backends.abstract import ( from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig from vllm_omni.diffusion.distributed.parallel_state import ( + get_allgather_parallel_world_size, get_cfg_group, get_classifier_free_guidance_rank, get_classifier_free_guidance_world_size, @@ -78,6 +79,7 @@ from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step from vllm_omni.diffusion.layers.norm import RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import HunyuanFusedMoE +from vllm_omni.diffusion.utils.kv_utils import repeat_kv from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding logger = logging.getLogger(__name__) @@ -426,22 +428,6 @@ def apply_rotary_pos_emb( return q_embed, k_embed -def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: - """ - Equivalent to torch.repeat_interleave(x, dim=2, repeats=n_rep). - Input: (batch, seqlen, num_key_value_heads, head_dim) - Output: (batch, seqlen, num_attention_heads, head_dim) - """ - batch, slen, num_key_value_heads, head_dim = hidden_states.shape - - if n_rep == 1: - return hidden_states - - hidden_states = hidden_states[:, :, :, None, :].expand(batch, slen, num_key_value_heads, n_rep, head_dim) - - return hidden_states.reshape(batch, slen, num_key_value_heads * n_rep, head_dim) - - def default(value, default_value): return value if value is not None else default_value @@ -957,6 +943,7 @@ class ImageKVCacheManager: self._injected_ar_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None self.sp_size = get_sequence_parallel_world_size() + self.allgather_size = get_allgather_parallel_world_size() self.sp_rank = get_sequence_parallel_rank() self.attn = Attention( num_heads=self.num_heads, @@ -1130,6 +1117,7 @@ class ImageKVCacheManager: head_num_per_rank = query.shape[1] kv_head_num_per_rank = key.shape[1] repeat_num = head_num_per_rank // kv_head_num_per_rank + keep_kv_compressed = self.allgather_size > 1 head_dim = query.shape[2] query = query.reshape(bs, q_len, head_num_per_rank, head_dim) @@ -1171,11 +1159,12 @@ class ImageKVCacheManager: joint_text_query = query[:, :0, :, :] joint_text_key, joint_text_value = self._reuse_prompt_kv(key, value, seq_len, bs, shard_image_size) - key = repeat_kv(key, repeat_num) - value = repeat_kv(value, repeat_num) - if self.sp_size > 1: - joint_text_key = repeat_kv(joint_text_key, repeat_num) - joint_text_value = repeat_kv(joint_text_value, repeat_num) + if not keep_kv_compressed: + key = repeat_kv(key, repeat_num) + value = repeat_kv(value, repeat_num) + if self.sp_size > 1: + joint_text_key = repeat_kv(joint_text_key, repeat_num) + joint_text_value = repeat_kv(joint_text_value, repeat_num) attention_mask = attention_mask.contiguous() diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 41412424d..a0b1e9a7b 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -464,18 +464,25 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) - if plan is None: continue - # Create SP config - sp_config = SequenceParallelConfig( - ulysses_degree=od_config.parallel_config.ulysses_degree, - ring_degree=od_config.parallel_config.ring_degree, - ) + # AllGather-KV reuses the Ulysses sequence-sharding hooks. + allgather_degree = getattr(od_config.parallel_config, "allgather_degree", 1) + if allgather_degree > 1: + sp_config = SequenceParallelConfig( + allgather_degree=allgather_degree, + ) + mode = "allgather_kv" + else: + sp_config = SequenceParallelConfig( + ulysses_degree=od_config.parallel_config.ulysses_degree, + ring_degree=od_config.parallel_config.ring_degree, + ) + # Apply hooks according to the plan + mode = ( + "hybrid" + if sp_config.ulysses_degree > 1 and sp_config.ring_degree > 1 + else ("ulysses" if sp_config.ulysses_degree > 1 else "ring") + ) - # Apply hooks according to the plan - mode = ( - "hybrid" - if sp_config.ulysses_degree > 1 and sp_config.ring_degree > 1 - else ("ulysses" if sp_config.ulysses_degree > 1 else "ring") - ) logger.info( f"Applying sequence parallelism to {transformer.__class__.__name__} ({attr}) " f"(sp_size={sp_size}, mode={mode}, ulysses={sp_config.ulysses_degree}, ring={sp_config.ring_degree})" diff --git a/vllm_omni/diffusion/utils/kv_utils.py b/vllm_omni/diffusion/utils/kv_utils.py index 9ba231386..fc6ec6cc1 100644 --- a/vllm_omni/diffusion/utils/kv_utils.py +++ b/vllm_omni/diffusion/utils/kv_utils.py @@ -3,6 +3,19 @@ import torch +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """Expand GQA KV from num_kv_heads to num_heads. + + Input: (batch, seqlen, num_key_value_heads, head_dim) + Output: (batch, seqlen, num_key_value_heads * n_rep, head_dim) + """ + if n_rep == 1: + return hidden_states + batch, slen, num_key_value_heads, head_dim = hidden_states.shape + hidden_states = hidden_states[:, :, :, None, :].expand(batch, slen, num_key_value_heads, n_rep, head_dim) + return hidden_states.reshape(batch, slen, num_key_value_heads * n_rep, head_dim) + + def left_pad_stack( tensors: list[torch.Tensor], ) -> tuple[torch.Tensor, torch.Tensor | None]: diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index df3ee5910..948406c73 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -315,6 +315,7 @@ class DiffusionWorker: sequence_parallel_size=parallel_config.sequence_parallel_size, ulysses_degree=parallel_config.ulysses_degree, ring_degree=parallel_config.ring_degree, + allgather_degree=parallel_config.allgather_degree, tensor_parallel_size=parallel_config.tensor_parallel_size, pipeline_parallel_size=parallel_config.pipeline_parallel_size, fully_shard_degree=parallel_config.hsdp_shard_size if parallel_config.use_hsdp else 1, diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 19db55c63..2d049d73f 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -459,6 +459,7 @@ class OrchestratorArgs: ulysses_degree: int | None = None ulysses_mode: str = "strict" ring_degree: int | None = None + allgather_degree: int | None = None diffusion_quantization_config: str | None = None use_hsdp: bool = False hsdp_shard_size: int = -1 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index dbebd7e6d..7740f11e7 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -933,6 +933,7 @@ class AsyncOmniEngine: if parallel_config is None: ulysses_degree = normalized_kwargs.get("ulysses_degree") or 1 ring_degree = normalized_kwargs.get("ring_degree") or 1 + allgather_degree = normalized_kwargs.get("allgather_degree") or 1 ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict" sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size") pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 @@ -947,7 +948,7 @@ class AsyncOmniEngine: hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size", -1) hsdp_replicate_size = normalized_kwargs.get("hsdp_replicate_size", 1) if sequence_parallel_size is None: - sequence_parallel_size = ulysses_degree * ring_degree + sequence_parallel_size = allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree parallel_config = DiffusionParallelConfig( pipeline_parallel_size=pipeline_parallel_size, @@ -957,6 +958,7 @@ class AsyncOmniEngine: sequence_parallel_size=sequence_parallel_size, ulysses_degree=ulysses_degree, ring_degree=ring_degree, + allgather_degree=allgather_degree, ulysses_mode=ulysses_mode, cfg_parallel_size=cfg_parallel_size, vae_patch_parallel_size=vae_patch_parallel_size, diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 71483445a..d807c22ed 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -473,6 +473,14 @@ class OmniServeCommand(CLISubcommand): help="Ring Sequence Parallelism degree for diffusion models. " "Equivalent to setting DiffusionParallelConfig.ring_degree.", ) + omni_config_group.add_argument( + "--allgather-degree", + dest="allgather_degree", + type=int, + default=None, + help="AllGather-KV Sequence Parallelism degree for non-causal diffusion attention. " + "Equivalent to setting DiffusionParallelConfig.allgather_degree.", + ) omni_config_group.add_argument( "--diffusion-quantization-config", type=json.loads,