From e51fe6ec1b9a9a0e14bb1fdb296d61b6593b93c6 Mon Sep 17 00:00:00 2001 From: Xenoryn Date: Wed, 2 Sep 2026 08:24:51 +0800 Subject: [PATCH] [Bugfix][Diffusion] Keep the GQA ratio when padding Ulysses heads (#5716) Signed-off-by: Xenoryn Co-authored-by: Claude --- tests/diffusion/attention/test_ulysses_uaa.py | 444 ++++++++++++++++-- .../diffusion/attention/parallel/ulysses.py | 201 ++++++-- 2 files changed, 566 insertions(+), 79 deletions(-) diff --git a/tests/diffusion/attention/test_ulysses_uaa.py b/tests/diffusion/attention/test_ulysses_uaa.py index 77e00edda..7568a1f37 100644 --- a/tests/diffusion/attention/test_ulysses_uaa.py +++ b/tests/diffusion/attention/test_ulysses_uaa.py @@ -29,28 +29,296 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @pytest.mark.core_model @pytest.mark.cpu -def test_advanced_uaa_hybrid_rejects_padded_mqa_heads(monkeypatch) -> None: - monkeypatch.setattr( - "vllm_omni.diffusion.attention.parallel.ulysses.get_ulysses_mode", - lambda default: "advanced_uaa", +def test_uaa_gqa_head_padding_preserves_the_query_to_kv_ratio( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Q must be padded to a multiple of the *KV* padding, not of world_size. + + Rounding each up independently is the obvious implementation and is wrong: + 28Q/7KV at SP2 would become 28/8, i.e. 14/4 per rank, whose 3.5 ratio is not + a valid GQA shape. Deriving Q from the padded KV count gives 32/8 -> 16/4. + """ + from vllm_omni.diffusion.attention.parallel import ulysses + + class FakeGroup: + ulysses_group = object() + ulysses_world_size = 2 + ulysses_rank = 0 + ring_world_size = 1 + + observed = [] + + def fake_all_to_all(pg, tensor, **kwargs): + observed.append(kwargs["padded_head_cnt"]) + return tensor, tensor.shape[2] + + monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all) + monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa") + monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3]) + strategy = ulysses.UlyssesParallelAttention( + FakeGroup(), + scatter_idx=2, + gather_idx=1, + use_sync=False, ) + + with set_forward_context(): + strategy.pre_attention( + torch.zeros(1, 3, 28, 4), + torch.zeros(1, 3, 7, 4), + torch.zeros(1, 3, 7, 4), + None, + ) + + q_padded, k_padded, v_padded = observed + assert (q_padded, k_padded, v_padded) == (32, 8, 8) + # The per-rank shape must still be a valid GQA shape. + assert (q_padded // 2) % (k_padded // 2) == 0 + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_uaa_rejects_head_counts_that_are_not_a_gqa_shape( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm_omni.diffusion.attention.parallel import ulysses + + class FakeGroup: + ulysses_group = object() + ulysses_world_size = 2 + ulysses_rank = 0 + ring_world_size = 1 + + monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa") + monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3]) + strategy = ulysses.UlyssesParallelAttention( + FakeGroup(), + scatter_idx=2, + gather_idx=1, + use_sync=False, + ) + + with set_forward_context(), pytest.raises(ValueError, match="multiple of KV heads"): + strategy.pre_attention( + torch.zeros(1, 3, 10, 4), + torch.zeros(1, 3, 4, 4), + torch.zeros(1, 3, 4, 4), + None, + ) + + +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize( + "q_heads, kv_heads, world_size, expected_warn", + [ + # 28Q/7KV @ U=2 -> pad to 32Q, 32/28 = 1.14x, below threshold, no warn. + (28, 7, 2, False), + # 32Q/1KV @ U=8 -> pad KV to 8 -> pad Q to 256, 256/32 = 8x, warn. + (32, 1, 8, True), + # 16Q/2KV @ U=4 -> pad KV to 4 -> pad Q to 32, 32/16 = 2x, warn. + (16, 2, 4, True), + ], +) +def test_uaa_padding_ratio_warning_fires_when_blowup_is_large( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + q_heads: int, + kv_heads: int, + world_size: int, + expected_warn: bool, +) -> None: + """advanced_uaa must warn once when GQA padding materially inflates Q heads. + + MQA-style shapes at high ulysses_degree pad Q by U x, which grows attention + FLOPs and temporary VRAM by the same factor. We do not fall back (there is + no cheap alternative -- Ulysses needs KV divisible by U for the head-dim + split), but we surface a one-shot warning so the user can pick a + ulysses_degree that divides KV_heads if the overhead is unacceptable. + """ + from vllm_omni.diffusion.attention.parallel import ulysses + + # Reset the warn-once dedup set so parametrized cases do not shadow each other. + ulysses._uaa_pad_ratio_warned.clear() + + class FakeGroup: + ulysses_group = object() + ulysses_world_size = world_size + ulysses_rank = 0 + ring_world_size = 1 + + def fake_all_to_all(pg, tensor, **kwargs): + return tensor, tensor.shape[2] + + monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all) + monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa") + monkeypatch.setattr(ulysses, "_all_gather_int", lambda *a, **kw: [3] * world_size) + strategy = ulysses.UlyssesParallelAttention( + FakeGroup(), + scatter_idx=2, + gather_idx=1, + use_sync=False, + ) + + caplog.set_level("WARNING", logger="vllm_omni.diffusion.attention.parallel.ulysses") + with set_forward_context(): + strategy.pre_attention( + torch.zeros(1, 3, q_heads, 4), + torch.zeros(1, 3, kv_heads, 4), + torch.zeros(1, 3, kv_heads, 4), + None, + ) + + warn_msgs = [r.message for r in caplog.records if "GQA padding inflates" in r.message] + if expected_warn: + assert len(warn_msgs) == 1, warn_msgs + # Second call with the same shape must not re-warn (warn-once dedup). + with set_forward_context(): + strategy.pre_attention( + torch.zeros(1, 3, q_heads, 4), + torch.zeros(1, 3, kv_heads, 4), + torch.zeros(1, 3, kv_heads, 4), + None, + ) + warn_msgs = [r.message for r in caplog.records if "GQA padding inflates" in r.message] + assert len(warn_msgs) == 1, "warn-once dedup should suppress the second call" + else: + assert warn_msgs == [], f"unexpected warning for {q_heads}Q/{kv_heads}KV @ U={world_size}: {warn_msgs}" + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_uaa_joint_gqa_head_padding_preserves_the_query_to_kv_ratio( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Joint Q/K/V must also preserve the GQA ratio through Ulysses padding. + + MMDiT-style models (SD3, FLUX, Qwen-Image, Hunyuan-Video, ...) pass a + separate text stream via ``AttentionMetadata.joint_*``. Before the fix the + joint block padded K/V by Q's amount and integer-divided K/V by world_size, + which for 28Q/7KV at SP=2 produces mismatched joint vs main head counts and + the downstream ring-attention concat blows up. Post-fix the joint block + mirrors the main-path rule: pad KV to a world_size multiple, derive Q as + ``padded_kv * (Q // KV)``. + """ + from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata + from vllm_omni.diffusion.attention.parallel import ulysses + + class FakeGroup: + ulysses_group = object() + ulysses_world_size = 2 + ulysses_rank = 0 + ring_world_size = 1 + + def fake_all_to_all(pg, tensor, **kwargs): + # Simulate the head-dim sharding so the post-a2a joint concat has + # matching per-rank head counts on Q and joint_Q. + heads_per_rank = kwargs["padded_head_cnt"] // 2 + return tensor[:, :, :heads_per_rank, :].contiguous(), tensor.shape[2] + + monkeypatch.setattr(ulysses, "_ulysses_all_to_all_any_qkv", fake_all_to_all) + monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa") + monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3]) + strategy = ulysses.UlyssesParallelAttention( + FakeGroup(), + scatter_idx=2, + gather_idx=1, + use_sync=False, + ) + + attn_metadata = AttentionMetadata( + joint_query=torch.zeros(1, 5, 28, 4), + joint_key=torch.zeros(1, 5, 7, 4), + joint_value=torch.zeros(1, 5, 7, 4), + ) + + with set_forward_context(): + strategy.pre_attention( + torch.zeros(1, 3, 28, 4), + torch.zeros(1, 3, 7, 4), + torch.zeros(1, 3, 7, 4), + attn_metadata, + ) + + # 28Q/7KV at SP=2 -> pad to 32/8 -> per-rank slice is 16Q/4KV. + assert attn_metadata.joint_key.shape[-2] == 4 + assert attn_metadata.joint_value.shape[-2] == 4 + # Per-rank joint shape must remain a valid GQA shape (Q multiple of KV). + per_rank_joint_q = 32 // 2 + per_rank_joint_kv = attn_metadata.joint_key.shape[-2] + assert per_rank_joint_q % per_rank_joint_kv == 0 + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_uaa_rejects_joint_head_counts_that_are_not_a_gqa_shape( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata + from vllm_omni.diffusion.attention.parallel import ulysses + + class FakeGroup: + ulysses_group = object() + ulysses_world_size = 2 + ulysses_rank = 0 + ring_world_size = 1 + + monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa") + monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3]) + strategy = ulysses.UlyssesParallelAttention( + FakeGroup(), + scatter_idx=2, + gather_idx=1, + use_sync=False, + ) + + attn_metadata = AttentionMetadata( + joint_query=torch.zeros(1, 5, 10, 4), + joint_key=torch.zeros(1, 5, 4, 4), + joint_value=torch.zeros(1, 5, 4, 4), + ) + + with set_forward_context(), pytest.raises(ValueError, match="multiple of joint KV heads"): + strategy.pre_attention( + torch.zeros(1, 3, 8, 4), + torch.zeros(1, 3, 4, 4), + torch.zeros(1, 3, 4, 4), + attn_metadata, + ) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_advanced_uaa_hybrid_rejects_non_gqa_shapes(monkeypatch: pytest.MonkeyPatch) -> None: + """Hybrid Ulysses+Ring still rejects shapes that are not a GQA shape. + + Valid GQA shapes are safe under advanced_uaa head padding (the derived-Q + rule keeps real Q heads paired with real K/V heads after the split, see + pre_attention), but a query count that is not a multiple of the K/V count + cannot be padded while preserving the ratio, so it must be rejected. + """ + from vllm_omni.diffusion.attention.parallel import ulysses + sp_group = SimpleNamespace( ulysses_group=None, ulysses_world_size=2, ulysses_rank=0, ring_world_size=2, + ring_group=object(), ) + monkeypatch.setattr(ulysses, "get_ulysses_mode", lambda **kwargs: "advanced_uaa") + monkeypatch.setattr(ulysses, "_all_gather_int", lambda *args, **kwargs: [3, 3]) strategy = UlyssesParallelAttention( sp_group=sp_group, scatter_idx=2, gather_idx=1, use_sync=False, ) - query = torch.zeros(1, 2, 8, 4) - key = torch.zeros(1, 2, 1, 4) + query = torch.zeros(1, 2, 10, 4) + key = torch.zeros(1, 2, 4, 4) value = torch.zeros_like(key) - with pytest.raises(ValueError, match="does not support K/V head padding"): + with pytest.raises(ValueError, match="multiple of KV heads"): strategy.pre_attention(query, key, value, None) @@ -67,6 +335,9 @@ def _run_attention_case( ring_degree: int = 1, split_sizes: list[int] | None = None, sdp_kernel_mode: str = "math", + num_kv_heads: int | None = None, + has_joint: bool = False, + joint_strategy: str = "front", ) -> None: device = torch.device(f"{current_omni_platform.device_type}:{local_rank}") current_omni_platform.set_device(device) @@ -102,15 +373,33 @@ def _run_attention_case( head_size=head_size, causal=False, softmax_scale=1.0 / (head_size**0.5), + num_kv_heads=num_kv_heads, ).to(device=device, dtype=torch.float32) with np.load(input_file, allow_pickle=False) as payload: q_full = torch.from_numpy(payload["q"]).to(device=device) k_full = torch.from_numpy(payload["k"]).to(device=device) v_full = torch.from_numpy(payload["v"]).to(device=device) + if has_joint: + joint_q = torch.from_numpy(payload["joint_q"]).to(device=device) + joint_k = torch.from_numpy(payload["joint_k"]).to(device=device) + joint_v = torch.from_numpy(payload["joint_v"]).to(device=device) + attn_metadata = None if world_size == 1: - q, k, v = q_full, k_full, v_full + if has_joint: + # Baseline: SP is inactive so pre_attention won't concat joint for us. + # Build the full sequence directly and skip the joint metadata. + if joint_strategy == "front": + q = torch.cat([joint_q, q_full], dim=1) + k = torch.cat([joint_k, k_full], dim=1) + v = torch.cat([joint_v, v_full], dim=1) + else: + q = torch.cat([q_full, joint_q], dim=1) + k = torch.cat([k_full, joint_k], dim=1) + v = torch.cat([v_full, joint_v], dim=1) + else: + q, k, v = q_full, k_full, v_full else: if split_sizes is None: # NOTE: torch.chunk may return fewer than `world_size` chunks for some @@ -135,6 +424,16 @@ def _run_attention_case( # shard manually for testing. get_forward_context()._sp_shard_depth = 1 + if has_joint: + from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata + + attn_metadata = AttentionMetadata( + joint_query=joint_q, + joint_key=joint_k, + joint_value=joint_v, + joint_strategy=joint_strategy, + ) + if sdp_kernel_mode == "math": sdp_ctx = torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True) elif sdp_kernel_mode == "mem_efficient": @@ -143,25 +442,47 @@ def _run_attention_case( raise ValueError(f"Invalid sdp_kernel_mode: {sdp_kernel_mode!r}") with sdp_ctx: - out_local = attn(q, k, v).contiguous() + out_local = attn(q, k, v, attn_metadata=attn_metadata).contiguous() if world_size == 1: out_full = out_local else: - local_len = torch.tensor([out_local.shape[1]], device=device, dtype=torch.int64) + if has_joint: + joint_len = int(joint_q.shape[1]) + if joint_strategy == "front": + joint_out_local = out_local[:, :joint_len] + img_out_local = out_local[:, joint_len:] + else: + img_out_local = out_local[:, : out_local.shape[1] - joint_len] + joint_out_local = out_local[:, -joint_len:] + else: + joint_out_local = None + img_out_local = out_local + + local_len = torch.tensor([img_out_local.shape[1]], device=device, dtype=torch.int64) gathered_lens = [torch.empty_like(local_len) for _ in range(world_size)] torch.distributed.all_gather(gathered_lens, local_len) lens = [int(t.item()) for t in gathered_lens] max_len = max(lens) - if out_local.shape[1] < max_len: - pad = max_len - out_local.shape[1] - out_local = torch.nn.functional.pad(out_local, (0, 0, 0, 0, 0, pad)).contiguous() + if img_out_local.shape[1] < max_len: + pad = max_len - img_out_local.shape[1] + img_out_local = torch.nn.functional.pad(img_out_local, (0, 0, 0, 0, 0, pad)).contiguous() + else: + img_out_local = img_out_local.contiguous() - gathered = [torch.empty_like(out_local) for _ in range(world_size)] - torch.distributed.all_gather(gathered, out_local) + gathered = [torch.empty_like(img_out_local) for _ in range(world_size)] + torch.distributed.all_gather(gathered, img_out_local) if local_rank == 0: - out_full = torch.cat([t[:, : lens[i]].contiguous() for i, t in enumerate(gathered)], dim=1) + img_full = torch.cat([t[:, : lens[i]].contiguous() for i, t in enumerate(gathered)], dim=1) + if has_joint: + out_full = ( + torch.cat([joint_out_local, img_full], dim=1) + if joint_strategy == "front" + else torch.cat([img_full, joint_out_local], dim=1) + ) + else: + out_full = img_full else: out_full = None @@ -172,20 +493,43 @@ def _run_attention_case( @pytest.mark.parametrize( - "sp_world_size,seq_len,num_heads", + "sp_world_size,seq_len,joint_len,num_heads,num_kv_heads", [ - (2, 6, 3), # head_cnt not divisible by P=2 - (2, 5, 4), # seq_len not divisible by P=2 - (4, 9, 30), # Z-Image-like: head_cnt not divisible by P=4 - (4, 10, 8), # seq_len not divisible by P=4 + # MHA (num_kv_heads=None, joint_len=None) + (2, 6, None, 3, None), # head_cnt not divisible by P=2 + (2, 5, None, 4, None), # seq_len not divisible by P=2 + (4, 9, None, 30, None), # Z-Image-like: head_cnt not divisible by P=4 + (4, 10, None, 8, None), # seq_len not divisible by P=4 + # GQA (joint_len=None) + (2, 6, None, 28, 7), # BOOGU-like GQA: neither Q nor KV divisible by P=2 + (2, 5, None, 6, 3), # KV divisible by P, Q is not + (4, 8, None, 12, 3), # KV not divisible by P=4 + # Joint GQA (MMDiT-style text stream via AttentionMetadata.joint_*) + (2, 6, 5, 28, 7), # BOOGU-like GQA joint: KV=7 not divisible by P=2 + (2, 5, 4, 6, 3), # KV divisible by P, Q is not ], ) -def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_heads: int) -> None: +def test_ulysses_uaa_matches_baseline( + sp_world_size: int, + seq_len: int, + joint_len: int | None, + num_heads: int, + num_kv_heads: int | None, +) -> None: + """End-to-end SP-vs-single-GPU parity under UAA. + + Covers MHA, GQA and joint (MMDiT-style) shapes. The joint variant catches + the pre-fix bug where the joint block padded K/V by Q's head count and + integer-divided K/V by world_size, giving per-rank shapes that mismatched + the main-tensor slice (e.g. 14Q/3KV vs. 16Q/4KV for 28Q/7KV at SP=2). + """ if current_omni_platform.get_device_count() < sp_world_size: pytest.skip(f"Test requires {sp_world_size} GPUs") batch_size = 2 head_size = 8 + kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + has_joint = joint_len is not None base_init_method = get_distributed_init_method() sp_init_method = get_distributed_init_method() @@ -200,17 +544,23 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head try: torch.manual_seed(0) q = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32) - k = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32) - v = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32) - np.savez(input_file, q=q.numpy(), k=k.numpy(), v=v.numpy()) + k = torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32) + v = torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32) + payload = {"q": q.numpy(), "k": k.numpy(), "v": v.numpy()} + if has_joint: + payload["joint_q"] = torch.randn(batch_size, joint_len, num_heads, head_size, dtype=torch.float32).numpy() + payload["joint_k"] = torch.randn(batch_size, joint_len, kv_heads, head_size, dtype=torch.float32).numpy() + payload["joint_v"] = torch.randn(batch_size, joint_len, kv_heads, head_size, dtype=torch.float32).numpy() + np.savez(input_file, **payload) + + common_tail = (1, None, "math", num_kv_heads, has_joint) # Baseline (no SP) torch.multiprocessing.spawn( _run_attention_case, - args=(1, base_init_method, input_file, baseline_file, num_heads, head_size, 1, "strict"), + args=(1, base_init_method, input_file, baseline_file, num_heads, head_size, 1, "strict", *common_tail), nprocs=1, ) - # SP (Ulysses-P with UAA) torch.multiprocessing.spawn( _run_attention_case, @@ -223,15 +573,13 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head head_size, sp_world_size, "advanced_uaa", + *common_tail, ), nprocs=sp_world_size, ) - baseline = np.load(baseline_file, allow_pickle=False) - sp = np.load(sp_file, allow_pickle=False) - - baseline_t = torch.from_numpy(baseline) - sp_t = torch.from_numpy(sp) + baseline_t = torch.from_numpy(np.load(baseline_file, allow_pickle=False)) + sp_t = torch.from_numpy(np.load(sp_file, allow_pickle=False)) assert baseline_t.shape == sp_t.shape torch.testing.assert_close(sp_t, baseline_t, atol=1e-5, rtol=1e-5) finally: @@ -242,7 +590,17 @@ def test_ulysses_uaa_matches_baseline(sp_world_size: int, seq_len: int, num_head pass -def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None: +@pytest.mark.parametrize( + "num_heads,num_kv_heads,joint_len", + [ + (3, None, None), # MHA baseline (head_cnt not divisible by ulysses_degree=2) + (28, 7, None), # GQA: exercises Ring pytorch_attn_forward KV expansion (Q!=K heads) + (28, 7, 5), # GQA + MMDiT joint stream: joint K/V ride the padded Ulysses layout and Ring SDPA + ], +) +def test_ulysses_uaa_hybrid_ring_matches_baseline( + num_heads: int, num_kv_heads: int | None, joint_len: int | None +) -> None: sp_world_size = 4 ulysses_degree = 2 ring_degree = 2 @@ -253,9 +611,8 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None: batch_size = 2 head_size = 8 seq_len = 10 - # This is the positive Hybrid case: K/V heads must not require padding, - # which is unsupported by advanced_uaa when Ring is also enabled. - num_heads = 4 + kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + has_joint = joint_len is not None # Ensure ring ranks see equal post-Ulysses seq_len: # rank0/1 -> 3+2=5, rank2/3 -> 3+2=5 @@ -274,9 +631,14 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None: try: torch.manual_seed(0) q = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32) - k = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32) - v = torch.randn(batch_size, seq_len, num_heads, head_size, dtype=torch.float32) - np.savez(input_file, q=q.numpy(), k=k.numpy(), v=v.numpy()) + k = torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32) + v = torch.randn(batch_size, seq_len, kv_heads, head_size, dtype=torch.float32) + payload = {"q": q.numpy(), "k": k.numpy(), "v": v.numpy()} + if has_joint: + payload["joint_q"] = torch.randn(batch_size, joint_len, num_heads, head_size, dtype=torch.float32).numpy() + payload["joint_k"] = torch.randn(batch_size, joint_len, kv_heads, head_size, dtype=torch.float32).numpy() + payload["joint_v"] = torch.randn(batch_size, joint_len, kv_heads, head_size, dtype=torch.float32).numpy() + np.savez(input_file, **payload) # Baseline (no SP) torch.multiprocessing.spawn( @@ -293,6 +655,8 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None: 1, None, "mem_efficient", + num_kv_heads, + has_joint, ), nprocs=1, ) @@ -312,6 +676,8 @@ def test_ulysses_uaa_hybrid_ring_matches_baseline() -> None: ring_degree, split_sizes, "mem_efficient", + num_kv_heads, + has_joint, ), nprocs=sp_world_size, ) diff --git a/vllm_omni/diffusion/attention/parallel/ulysses.py b/vllm_omni/diffusion/attention/parallel/ulysses.py index 69f6a810e..4319db108 100644 --- a/vllm_omni/diffusion/attention/parallel/ulysses.py +++ b/vllm_omni/diffusion/attention/parallel/ulysses.py @@ -8,6 +8,7 @@ from dataclasses import dataclass import torch import torch.distributed as dist import torch.nn.functional as F +from vllm.logger import init_logger from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.parallel.base import ParallelAttentionContext @@ -15,6 +16,55 @@ from vllm_omni.diffusion.distributed.comm import SeqAllToAll4D from vllm_omni.diffusion.distributed.group_coordinator import SequenceParallelGroupCoordinator from vllm_omni.diffusion.forward_context import get_ulysses_mode +logger = init_logger(__name__) + +# When advanced_uaa pads Q by the GQA ratio, MQA/very-uneven-GQA shapes can +# inflate the query-head count substantially (worst case: MQA @ U=N pads Q from +# H to N*H, an N x blow-up in Q attention work and temporary VRAM). Warn once +# per (Q, KV, world_size) tuple when the ratio crosses this threshold so users +# can pick a friendlier ulysses_degree if the overhead is unacceptable. +_UAA_PAD_RATIO_WARN_THRESHOLD = 1.5 +_uaa_pad_ratio_warned: set[tuple[int, int, int, str]] = set() + + +def _maybe_warn_uaa_pad_ratio( + query_head_cnt: int, + kv_head_cnt: int, + padded_query_head_cnt: int, + ulysses_world_size: int, + tensor_label: str, +) -> None: + """Emit a one-shot warning when advanced_uaa padding materially inflates Q. + + The padding itself is mandatory for correctness (Ulysses splits along the + head dim, so KV must be a multiple of ulysses_world_size and Q must stay a + multiple of KV to remain a valid GQA shape). There is no cheap fallback, + but the caller may want to know when the inflation is large enough to + justify picking a different ulysses_degree. + """ + if query_head_cnt <= 0 or padded_query_head_cnt <= query_head_cnt: + return + ratio = padded_query_head_cnt / query_head_cnt + if ratio < _UAA_PAD_RATIO_WARN_THRESHOLD: + return + key = (query_head_cnt, kv_head_cnt, ulysses_world_size, tensor_label) + if key in _uaa_pad_ratio_warned: + return + _uaa_pad_ratio_warned.add(key) + logger.warning( + "Ulysses advanced_uaa GQA padding inflates %s heads %.2fx " + "(Q=%d, KV=%d, ulysses_degree=%d -> padded Q=%d). " + "Attention FLOPs and temporary VRAM for this tensor grow by the same " + "factor. If this is unacceptable, choose a ulysses_degree that divides " + "KV_heads (or the GCD of Q/KV) to reduce or eliminate padding.", + tensor_label, + ratio, + query_head_cnt, + kv_head_cnt, + ulysses_world_size, + padded_query_head_cnt, + ) + def _a2a_permute_enabled(enabled: bool, scatter_idx: int, gather_idx: int, world_size: int) -> bool: """Whether the fused permute-free all-to-all applies (strict layout only).""" @@ -60,6 +110,7 @@ def _ulysses_all_to_all_any_qkv( *, seq_lens: list[int], use_sync: bool, + padded_head_cnt: int | None = None, ) -> tuple[torch.Tensor, int]: """UAA forward all-to-all: (B, S_local, H, D) -> (B, S_global, H_local, D). @@ -72,7 +123,12 @@ def _ulysses_all_to_all_any_qkv( bsz, s_local, head_cnt, head_dim = x.shape orig_head_cnt = int(head_cnt) - padded_head_cnt = _ceil_div(orig_head_cnt, world_size) * world_size + if padded_head_cnt is None: + padded_head_cnt = _ceil_div(orig_head_cnt, world_size) * world_size + if padded_head_cnt < orig_head_cnt or padded_head_cnt % world_size != 0: + raise ValueError( + f"Invalid padded head count {padded_head_cnt} for original heads={orig_head_cnt}, world_size={world_size}." + ) head_pad = padded_head_cnt - orig_head_cnt if head_pad: x = F.pad(x, (0, 0, 0, head_pad)) @@ -223,21 +279,9 @@ class UlyssesParallelAttention: ulysses_world_size = self._sp_group.ulysses_world_size # advanced_uaa pads non-divisible head counts before the Ulysses - # all-to-all. Padding K/V is not valid in hybrid Ulysses+Ring: the Ring - # GQA/MQA path would repeat the padded zero heads as real K/V heads. - # Reject this layout until K/V replication is performed before Ulysses. - if mode == "advanced_uaa" and self._sp_group.ring_world_size > 1: - for name, tensor in (("key", key), ("value", value)): - kv_head_cnt = int(tensor.shape[2]) - if kv_head_cnt % ulysses_world_size != 0: - raise ValueError( - "ulysses_mode='advanced_uaa' with hybrid Ulysses+Ring " - "does not support K/V head padding. " - f"{name}_head_cnt={kv_head_cnt}, " - f"ulysses_degree={ulysses_world_size}. " - "Use ring_degree=1, choose a compatible ulysses_degree, " - "or replicate K/V heads before Ulysses." - ) + # all-to-all. This also holds in hybrid Ulysses+Ring: Q is derived + # from the padded K/V count by the GQA ratio (see below), so padded + # heads pair with padded heads; non-GQA layouts are rejected below. joint_tensor_query = joint_tensor_key = joint_tensor_value = None joint_strategy = "front" @@ -262,32 +306,61 @@ class UlyssesParallelAttention: # Slice joint_query for this Ulysses rank # joint_query is (B, S, H, D). We split H (dim 2). ulysses_rank = self._sp_group.ulysses_rank - joint_head_cnt = int(joint_tensor_query.shape[-2]) - joint_orig_head_cnt = joint_head_cnt + joint_q_head_cnt = int(joint_tensor_query.shape[-2]) + joint_k_head_cnt = int(joint_tensor_key.shape[-2]) + joint_v_head_cnt = int(joint_tensor_value.shape[-2]) + joint_orig_head_cnt = joint_q_head_cnt + + if joint_k_head_cnt != joint_v_head_cnt or joint_q_head_cnt % joint_k_head_cnt != 0: + raise ValueError( + "Ulysses joint attention requires equal K/V head counts and joint query " + "heads to be a multiple of joint KV heads, got " + f"joint_Q={joint_q_head_cnt}, joint_K={joint_k_head_cnt}, joint_V={joint_v_head_cnt}." + ) if mode == "advanced_uaa": - padded_joint_head_cnt = _ceil_div(joint_head_cnt, ulysses_world_size) * ulysses_world_size - joint_head_pad = padded_joint_head_cnt - joint_head_cnt - if joint_head_pad: - joint_tensor_query = F.pad(joint_tensor_query, (0, 0, 0, joint_head_pad)) - joint_tensor_key = F.pad(joint_tensor_key, (0, 0, 0, joint_head_pad)) - joint_tensor_value = F.pad(joint_tensor_value, (0, 0, 0, joint_head_pad)) - joint_head_cnt = padded_joint_head_cnt + # Pad joint KV to a world_size multiple, then derive joint Q by the + # GQA ratio so the ratio survives the head-dim split (mirrors the + # main-tensor path below; e.g. 28Q/7KV at SP2 becomes 32/8). + padded_joint_kv_head_cnt = _ceil_div(joint_k_head_cnt, ulysses_world_size) * ulysses_world_size + gqa_ratio = joint_q_head_cnt // joint_k_head_cnt + padded_joint_q_head_cnt = padded_joint_kv_head_cnt * gqa_ratio + _maybe_warn_uaa_pad_ratio( + joint_q_head_cnt, + joint_k_head_cnt, + padded_joint_q_head_cnt, + ulysses_world_size, + "joint_query", + ) + + joint_q_pad = padded_joint_q_head_cnt - joint_q_head_cnt + joint_kv_pad = padded_joint_kv_head_cnt - joint_k_head_cnt + if joint_q_pad: + joint_tensor_query = F.pad(joint_tensor_query, (0, 0, 0, joint_q_pad)) + if joint_kv_pad: + joint_tensor_key = F.pad(joint_tensor_key, (0, 0, 0, joint_kv_pad)) + joint_tensor_value = F.pad(joint_tensor_value, (0, 0, 0, joint_kv_pad)) + joint_q_head_cnt = padded_joint_q_head_cnt + joint_k_head_cnt = padded_joint_kv_head_cnt else: - if joint_head_cnt % ulysses_world_size != 0: - supported = _positive_divisors(joint_head_cnt) - raise ValueError( - "Ulysses-SP strict mode requires joint head_cnt divisible by ulysses_degree. " - f"joint_head_cnt={joint_head_cnt}, ulysses_degree={ulysses_world_size}. " - f"Try ulysses_degree in {supported}, or set ulysses_mode='advanced_uaa'." - ) + for name, cnt in ( + ("joint_query", joint_q_head_cnt), + ("joint_key", joint_k_head_cnt), + ("joint_value", joint_v_head_cnt), + ): + if cnt % ulysses_world_size != 0: + supported = _positive_divisors(cnt) + raise ValueError( + "Ulysses-SP strict mode requires joint head_cnt divisible by ulysses_degree. " + f"{name}_head_cnt={cnt}, ulysses_degree={ulysses_world_size}. " + f"Try ulysses_degree in {supported}, or set ulysses_mode='advanced_uaa'." + ) - attn_heads_per_ulysses_rank = joint_head_cnt // ulysses_world_size + attn_heads_per_ulysses_rank_q = joint_q_head_cnt // ulysses_world_size - # Note: We use the same heads for Q/K/V joint_tensor_query = joint_tensor_query[ ..., - attn_heads_per_ulysses_rank * ulysses_rank : attn_heads_per_ulysses_rank * (ulysses_rank + 1), + attn_heads_per_ulysses_rank_q * ulysses_rank : attn_heads_per_ulysses_rank_q * (ulysses_rank + 1), :, ] @@ -301,8 +374,9 @@ class UlyssesParallelAttention: if is_joint: # Slice joint key/value heads for this ulysses rank. - # Using same slicing logic as query - attn_heads_per_ulysses_rank_kv = joint_tensor_key.shape[-2] // ulysses_world_size + # Using the KV head count (post-pad in advanced_uaa) so the per-rank + # KV shape matches the main K/V slice and preserves the GQA ratio. + attn_heads_per_ulysses_rank_kv = joint_k_head_cnt // ulysses_world_size joint_tensor_key = joint_tensor_key[ ..., @@ -345,11 +419,58 @@ class UlyssesParallelAttention: "This typically means the input sequence was not evenly shardable across the ring. " "Try setting ring_degree=1, or choose a sequence length divisible by ring_degree." ) + + # Pad KV to a world_size multiple, then scale Q by the GQA ratio so + # the ratio survives the head-dim split (e.g. 28Q/7KV at SP2 becomes + # 32/8, i.e. 16/4 per rank -- padding each independently would give + # 14/4, which is not a valid GQA shape). + # + # Overhead: padded_Q = ceil(KV/U)*U * (Q/KV). This is exact for + # well-aligned shapes (0% overhead when U divides KV) but can be + # substantial for uneven GQA/MQA: 28Q/7KV @ U=2 pads to 32Q (+14%), + # 32Q/1KV @ U=8 pads to 256Q (8x). _maybe_warn_uaa_pad_ratio() + # surfaces a one-shot warning once the blow-up crosses + # _UAA_PAD_RATIO_WARN_THRESHOLD so callers can adjust + # ulysses_degree if the extra work / VRAM is unacceptable. + query_head_cnt = int(query.shape[2]) + kv_head_cnt = int(key.shape[2]) + if key.shape[2] != value.shape[2] or query_head_cnt % kv_head_cnt != 0: + raise ValueError( + "Ulysses GQA requires equal K/V head counts and query heads " + f"to be a multiple of KV heads, got Q={query_head_cnt}, " + f"K={kv_head_cnt}, V={int(value.shape[2])}." + ) + padded_kv_head_cnt = _ceil_div(kv_head_cnt, ulysses_world_size) * ulysses_world_size + padded_query_head_cnt = padded_kv_head_cnt * (query_head_cnt // kv_head_cnt) + _maybe_warn_uaa_pad_ratio( + query_head_cnt, + kv_head_cnt, + padded_query_head_cnt, + ulysses_world_size, + "query", + ) + query, orig_head_cnt = _ulysses_all_to_all_any_qkv( - self._ulysses_pg, query, seq_lens=seq_lens, use_sync=self._use_sync + self._ulysses_pg, + query, + seq_lens=seq_lens, + use_sync=self._use_sync, + padded_head_cnt=padded_query_head_cnt, + ) + key, _ = _ulysses_all_to_all_any_qkv( + self._ulysses_pg, + key, + seq_lens=seq_lens, + use_sync=self._use_sync, + padded_head_cnt=padded_kv_head_cnt, + ) + value, _ = _ulysses_all_to_all_any_qkv( + self._ulysses_pg, + value, + seq_lens=seq_lens, + use_sync=self._use_sync, + padded_head_cnt=padded_kv_head_cnt, ) - key, _ = _ulysses_all_to_all_any_qkv(self._ulysses_pg, key, seq_lens=seq_lens, use_sync=self._use_sync) - value, _ = _ulysses_all_to_all_any_qkv(self._ulysses_pg, value, seq_lens=seq_lens, use_sync=self._use_sync) else: # Strict mode: fail fast with actionable errors for head divisibility. for name, t in (("query", query), ("key", key), ("value", value)):