diff --git a/tests/model_executor/models/qwen2_5_omni/test_token2wav_block_attention.py b/tests/model_executor/models/qwen2_5_omni/test_token2wav_block_attention.py new file mode 100644 index 00000000000..665282ffcd4 --- /dev/null +++ b/tests/model_executor/models/qwen2_5_omni/test_token2wav_block_attention.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch +import torch.nn.functional as F + +from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav import block_sparse_attention + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +BLOCK_SIZE = 24 + + +def _dense_masked_reference(query, key, value, look_backward_block, look_ahead_block): + """The original Token2Wav DiT attention: dense SDPA with a block-window mask.""" + seq_len = query.shape[2] + block_indices = torch.arange(seq_len) // BLOCK_SIZE + block_diff = block_indices[None, :] - block_indices[:, None] + mask = (block_diff >= -look_backward_block) & (block_diff <= look_ahead_block) + return F.scaled_dot_product_attention(query, key, value, attn_mask=mask) + + +# seq_len covers: shorter than one block, a multiple of the block size, and a +# partial last block. +@pytest.mark.parametrize("seq_len", [10, 96, 101]) +@pytest.mark.parametrize(("look_backward_block", "look_ahead_block"), [(0, 0), (1, 0), (0, 1)]) +def test_block_sparse_attention_matches_dense_masked_sdpa(seq_len, look_backward_block, look_ahead_block): + torch.manual_seed(0) + query, key, value = (torch.randn(2, 4, seq_len, 16) for _ in range(3)) + + out = block_sparse_attention( + query, + key, + value, + block_size=BLOCK_SIZE, + look_backward_block=look_backward_block, + look_ahead_block=look_ahead_block, + ) + ref = _dense_masked_reference(query, key, value, look_backward_block, look_ahead_block) + + assert out.shape == ref.shape + torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5) diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py index dacb47b88f5..9af7ea34845 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_token2wav.py @@ -11,7 +11,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import ( Qwen2_5OmniBigVGANConfig, Qwen2_5OmniDiTConfig, @@ -533,6 +532,81 @@ def rotate_half_codec(x): return q_embed, k_embed +def block_sparse_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + block_size: int, + look_backward_block: int = 0, + look_ahead_block: int = 0, +) -> torch.Tensor: + """Attention where each block of ``block_size`` queries attends to its own + block plus ``look_backward_block`` blocks before and ``look_ahead_block`` + blocks after it. + + This is equivalent to dense SDPA with the mask + ``(block_j - block_i >= -look_backward_block) & (block_j - block_i <= look_ahead_block)``, + but only computes the kept key blocks instead of materializing an (n, n) + mask for every layer. q/k/v and the result are ``(batch, heads, seq, dim)``. + """ + batch_size, num_heads, seq_len, head_dim = query.shape + num_blocks = -(-seq_len // block_size) + pad = num_blocks * block_size - seq_len + if pad: + query, key, value = (F.pad(t, (0, 0, 0, pad)) for t in (query, key, value)) + blocked_shape = (batch_size, num_heads, num_blocks, block_size, head_dim) + query = query.view(blocked_shape) + key = key.view(blocked_shape) + value = value.view(blocked_shape) + + # Some key is masked only when the last block is padded or a neighbour + # block is shifted past either end, both known from these Python ints, so + # the mask decision needs no device sync. + has_neighbours = look_backward_block > 0 or look_ahead_block > 0 + key_valid = None + if pad or has_neighbours: + # key_valid[b, j] is False for padded positions (only in the last block). + key_valid = torch.ones(num_blocks, block_size, dtype=torch.bool, device=query.device) + if pad: + key_valid[-1, block_size - pad :] = False + + if not has_neighbours: + # Own block only (most layers): attend to the blocked keys directly. + keys, values, valid = key, value, key_valid + else: + keys, values, valids = [], [], [] + for offset in range(-look_backward_block, look_ahead_block + 1): + if offset == 0: + keys.append(key) + values.append(value) + valids.append(key_valid) + continue + # Key block i + offset for each query block i; blocks past either + # end are zero-filled and masked out. + shifted_key = torch.zeros_like(key) + shifted_value = torch.zeros_like(value) + valid = torch.zeros_like(key_valid) + if offset > 0: + shifted_key[:, :, : num_blocks - offset] = key[:, :, offset:] + shifted_value[:, :, : num_blocks - offset] = value[:, :, offset:] + valid[: num_blocks - offset] = key_valid[offset:] + else: + shifted_key[:, :, -offset:] = key[:, :, : num_blocks + offset] + shifted_value[:, :, -offset:] = value[:, :, : num_blocks + offset] + valid[-offset:] = key_valid[: num_blocks + offset] + keys.append(shifted_key) + values.append(shifted_value) + valids.append(valid) + keys = torch.cat(keys, dim=3) + values = torch.cat(values, dim=3) + valid = torch.cat(valids, dim=1) + # Every query block keeps at least its own block, so no row is fully masked. + mask = None if valid is None else valid.view(1, 1, num_blocks, 1, -1) + + out = F.scaled_dot_product_attention(query, keys, values, attn_mask=mask) + return out.reshape(batch_size, num_heads, num_blocks * block_size, head_dim)[:, :, :seq_len] + + class DiTAttention(nn.Module): def __init__(self, config: Qwen2_5OmniDiTConfig, prefix: str = ""): super().__init__() @@ -559,7 +633,9 @@ def forward( self, hidden_states, # noised input x position_embeddings=None, # rotary position embedding for x - attention_mask=None, + block_size: int | None = None, + look_backward_block: int = 0, + look_ahead_block: int = 0, ) -> torch.Tensor: batch_size = hidden_states.shape[0] @@ -579,19 +655,17 @@ def forward( cos, sin = position_embeddings query[:, :1], key[:, :1] = apply_rotary_pos_emb(query[:, :1], key[:, :1], cos, sin) - attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] - attention_weights, _ = attention_interface( - self, + # The DiT only attends within a window of mel blocks, so run block + # sparse attention instead of a dense (n, n) masked SDPA. + attention_weights = block_sparse_attention( query, key, value, - attention_mask=attention_mask, - is_causal=False, + block_size=block_size, + look_backward_block=look_backward_block, + look_ahead_block=look_ahead_block, ) - - # mask. e.g. inference got a batch with different target durations, - # mask out the padding - attention_weights = attention_weights.reshape(batch_size, -1, self.heads * head_dim) + attention_weights = attention_weights.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim) attention_weights = attention_weights.to(query.dtype) # linear proj @@ -613,7 +687,7 @@ def __init__(self, config: Qwen2_5OmniDiTConfig, look_ahead_block=0, look_backwa self.ff = DiTMLP(dim=config.hidden_size, mult=config.ff_mult, dropout=config.dropout) def forward( - self, hidden_states, timestep, position_embeddings=None, block_diff=None + self, hidden_states, timestep, position_embeddings=None, block_size=None ): # x: noised input, t: time embedding # pre-norm & modulation for attention input norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(hidden_states, emb=timestep) @@ -622,8 +696,9 @@ def forward( attn_output = self.attn( hidden_states=norm, position_embeddings=position_embeddings, - attention_mask=(block_diff >= -float(self.look_backward_block)) - & (block_diff <= float(self.look_ahead_block)), + block_size=block_size, + look_backward_block=self.look_backward_block, + look_ahead_block=self.look_ahead_block, ) # process attention output for input x @@ -948,16 +1023,6 @@ def __init__(self, config: Qwen2_5OmniDiTConfig): self.norm_out = Qwen2_5_OmniAdaLayerNormZero_Final(config.hidden_size) # final modulation self.proj_out = nn.Linear(config.hidden_size, config.mel_dim) - def _create_block_diff(self, hidden_states): - batch, seq_len = hidden_states.shape[0], hidden_states.shape[1] - block_indices = torch.arange(seq_len, device=hidden_states.device) // self.block_size # [seq_length] - - block_i = block_indices.unsqueeze(1) # [seq_length, 1] - block_j = block_indices.unsqueeze(0) # [1, seq_length] - block_diff = block_j - block_i # (n, n) - - return block_diff.expand(batch, self.num_attention_heads, seq_len, seq_len) - def forward( self, hidden_states, @@ -990,7 +1055,6 @@ def forward( # Compute positional encodings position_embeddings = self.rotary_embed(hidden_states) - blockwise_difference = self._create_block_diff(hidden_states) # Transformer blocks for transformer_block in self.transformer_blocks: @@ -998,7 +1062,7 @@ def forward( hidden_states, time_embedding, position_embeddings=position_embeddings, - block_diff=blockwise_difference, + block_size=self.block_size, ) hidden_states = self.norm_out(hidden_states, time_embedding)