anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Bugfix][Diffusion] Keep the GQA ratio when padding Ulysses heads (#5716)

Browse Source 
Signed-off-by: Xenoryn <timberh2000@outlook.com>
Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
XenorynandClaude authored and GitHub committed 2026-09-02 08:24:51 +08:00
1 parent 5f49848536
commit e51fe6ec1b
2 files changed
+566 -79
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/attention/test_ulysses_uaa.py
+405 -39
View File
@@ -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,
)
vllm_omni/diffusion/attention/parallel/ulysses.py
+161 -40
View File
@@ -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)):