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

[Bugfix][Cosmos3] Auto-pad non-divisible Ulysses GEN sequences (#6918)

Browse Source 
Signed-off-by: Rahul Steiger <rsteiger@nvl72d188-T17.cm.cluster>
Co-authored-by: Rahul Steiger <rsteiger@nvl72d188-T17.cm.cluster>
This commit is contained in:
Rahul SteigerandRahul Steiger authored and GitHub committed 2026-09-07 08:58:28 +00:00
1 parent 85b0bad2d3
commit f080f41e03
3 files changed
+155 -73
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/attention/parallel/ulysses.py
+8 -7
View File
@@ -608,15 +608,15 @@ class UlyssesParallelAttention:
if attn_metadata.attn_mask is None:
assert attn_metadata.joint_attn_mask is not None
attn_metadata.attn_mask = torch.ones(
[query.shape[0], query.shape[1] - attn_metadata.joint_attn_mask.shape[1]],
[key.shape[0], key.shape[1] - attn_metadata.joint_attn_mask.shape[1]],
dtype=torch.bool,
device=query.device,
device=key.device,
)
elif attn_metadata.joint_attn_mask is None:
attn_metadata.joint_attn_mask = torch.ones(
[query.shape[0], query.shape[1] - attn_metadata.attn_mask.shape[1]],
[key.shape[0], key.shape[1] - attn_metadata.attn_mask.shape[1]],
dtype=torch.bool,
device=query.device,
device=key.device,
)
attn_metadata.attn_mask = (
torch.cat([attn_metadata.joint_attn_mask, attn_metadata.attn_mask], dim=1)
@@ -625,9 +625,10 @@ class UlyssesParallelAttention:
)
if attn_metadata.attn_mask is not None:
# the final attn_mask is ready, the length should be aligedn with query length
assert attn_metadata.attn_mask.shape[1] == query.shape[1], (
f"attn_mask length: {attn_metadata.attn_mask.shape[1]} != query length: {query.shape[1]}"
# A 2D mask describes keys. Joint attention may provide K/V-only
# context, so key and query lengths are not necessarily equal.
assert attn_metadata.attn_mask.shape[1] == key.shape[1], (
f"attn_mask length: {attn_metadata.attn_mask.shape[1]} != key length: {key.shape[1]}"
)
attn_metadata.attn_mask = attn_metadata.attn_mask.bool().contiguous()
return query, key, value, attn_metadata, ctx