diff --git a/tests/diffusion/attention/test_piecewise_attn.py b/tests/diffusion/attention/test_piecewise_attn.py new file mode 100644 index 000000000..656087623 --- /dev/null +++ b/tests/diffusion/attention/test_piecewise_attn.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end test for ``piecewise_attn`` (CPU). + +Verify that running attention in segments (causal outside full-attn spans, +bidirectional inside full-attn spans) matches running a single full SDPA call +with the equivalent 2D attention mask. + +Covers: + * batch size = 1 and batch size > 1 (homogeneous CFG-like batch) + * query length == key length (full prefill) + * query length < key length (decode-like tail slice) + * various full-attn-span layouts (none / start / middle / end / multi) +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import ( + piecewise_attn, +) + +DEVICE = torch.device("cpu") + + +def _sdpa_attn_func(q, k, v, causal, softmax_scale): + q_ = q.transpose(1, 2) + k_ = k.transpose(1, 2) + v_ = v.transpose(1, 2) + attn_mask = None + if causal: + Sq, Sk = q_.shape[-2], k_.shape[-2] + i = torch.arange(Sq, device=q.device).unsqueeze(1) + j = torch.arange(Sk, device=q.device).unsqueeze(0) + attn_mask = j <= (i + (Sk - Sq)) + out = F.scaled_dot_product_attention(q_, k_, v_, attn_mask=attn_mask, scale=softmax_scale) + return out.transpose(1, 2).contiguous() + + +def _full_reference(query, key, value, global_spans, q_start, q_end, softmax_scale): + """Build a full 2D mask with global spans and compute reference output.""" + Sk = key.shape[1] + mask = torch.tril(torch.ones(Sk, Sk, dtype=torch.bool, device=key.device)) + for a, e in global_spans: + mask[a:e, :e] = True + mask_q = mask[q_start:q_end, :] + q_ = query.transpose(1, 2) + k_ = key.transpose(1, 2) + v_ = value.transpose(1, 2) + out = F.scaled_dot_product_attention(q_, k_, v_, attn_mask=mask_q, scale=softmax_scale) + return out.transpose(1, 2).contiguous() + + +SPAN_CASES = [ + pytest.param([], id="no-spans"), + pytest.param([(0, 10)], id="span-at-start"), + pytest.param([(10, 30), (54, 64)], id="multi-spans"), +] + +Q_RANGE_CASES = [ + pytest.param((0, 64), id="q_eq_k"), # Sq == Sk (prefill) + pytest.param((53, 64), id="q_lt_k"), # Sq < Sk (decode-like) +] + +BATCH_CASES = [ + pytest.param(1, id="B1"), + pytest.param(2, id="B2"), +] + + +@pytest.mark.parametrize("global_spans", SPAN_CASES) +@pytest.mark.parametrize("q_range", Q_RANGE_CASES) +@pytest.mark.parametrize("batch_size", BATCH_CASES) +def test_piecewise_matches_full(global_spans, q_range, batch_size): + torch.manual_seed(0) + H, D, Sk = 2, 16, 64 + q_start, q_end = q_range + Sq = q_end - q_start + + key = torch.randn(batch_size, Sk, H, D, device=DEVICE) + value = torch.randn(batch_size, Sk, H, D, device=DEVICE) + query = torch.randn(batch_size, Sq, H, D, device=DEVICE) + + full_attn_spans = [list(global_spans) for _ in range(batch_size)] + softmax_scale = 1.0 / (D**0.5) + + got = piecewise_attn( + query, + key, + value, + full_attn_spans=full_attn_spans, + softmax_scale=softmax_scale, + attn_func=_sdpa_attn_func, + ) + 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_span_fully_before_qstart(): + """Spans entirely before query region produce pure causal attention.""" + torch.manual_seed(0) + B, H, D, Sk = 1, 2, 16, 30 + q_start, q_end = 15, 30 + Sq = q_end - q_start + + key = torch.randn(B, Sk, H, D, device=DEVICE) + value = torch.randn(B, Sk, H, D, device=DEVICE) + query = torch.randn(B, Sq, H, D, device=DEVICE) + + global_spans = [(5, 10)] + full_attn_spans = [list(global_spans) for _ in range(B)] + softmax_scale = 1.0 / (D**0.5) + + got = piecewise_attn( + query, + key, + value, + full_attn_spans=full_attn_spans, + softmax_scale=softmax_scale, + attn_func=_sdpa_attn_func, + ) + 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) diff --git a/vllm_omni/diffusion/attention/backends/abstract.py b/vllm_omni/diffusion/attention/backends/abstract.py index 6dd785e66..f702dc650 100644 --- a/vllm_omni/diffusion/attention/backends/abstract.py +++ b/vllm_omni/diffusion/attention/backends/abstract.py @@ -68,6 +68,10 @@ class AttentionMetadata: # Opaque backend-specific per-forward parameters (e.g. block masks, KV indices). # Backends MUST silently ignore unknown keys. + # 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 + T = TypeVar("T", bound=AttentionMetadata) diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index 341398446..a61254694 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from functools import partial import torch from vllm.logger import init_logger @@ -10,6 +11,9 @@ from vllm_omni.diffusion.attention.backends.abstract import ( AttentionImpl, AttentionMetadata, ) +from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import ( + piecewise_attn, +) logger = init_logger(__name__) @@ -59,6 +63,10 @@ class FlashAttentionImpl(AttentionImpl): # FA3 may return (out, lse), FA2 returns out return out[0] if isinstance(out, tuple) else out + @staticmethod + def _flash_wrapper(q, k, v, *, attn_func, **kwargs): + return FlashAttentionImpl._unwrap_flash_output(attn_func(q, k, v, **kwargs)) + def _forward_varlen_masked( self, query: torch.Tensor, @@ -156,6 +164,24 @@ class FlashAttentionImpl(AttentionImpl): ) attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None + full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None + + # 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, + ) + + return piecewise_attn( + query, + key, + value, + full_attn_spans, + self.softmax_scale, + attn_func, + ) if attention_mask is not None and torch.any(~attention_mask): return self._forward_varlen_masked( diff --git a/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py b/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py new file mode 100644 index 000000000..fbbd3c8e0 --- /dev/null +++ b/vllm_omni/diffusion/attention/backends/utils/piecewise_attn.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Piecewise attention for mixed causal / full (bidirectional) masks. + +Dispatches each segment as a separate attention call whose causal flag +follows FlashAttention's bottom-right convention (``K[:e]`` is attended by +``Q[s:e]``, with causal alignment anchored at the bottom-right corner). + +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)`` +""" + +from __future__ import annotations + +from typing import Literal, NamedTuple + + +class Segment(NamedTuple): + start: int + end: int + mode: Literal["causal", "full"] + + +def build_segments(full_attn_spans, query_offset, query_len): + """ + full_attn_spans: list of (start, end) half-open spans in global coordinates + query_offset: starting position of query in the global sequence + query_len: length of the query + + return: + List[Segment] in global coordinates, clipped to [query_offset, query_offset + query_len) + """ + q_start = query_offset + q_end = query_offset + query_len + + segs: list[Segment] = [] + cur = q_start + + for a, e 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: + 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 < q_end: + segs.append(Segment(cur, q_end, "causal")) + + return segs + + +def _check_homogeneous( + full_attn_spans: list[list[tuple[int, int]]], +) -> None: + """Assert all samples share identical spans.""" + if len(full_attn_spans) > 1: + ref = full_attn_spans[0] + for i, s in enumerate(full_attn_spans[1:], 1): + if s != ref: + raise ValueError( + f"piecewise_attn requires homogeneous batch: sample 0 spans {ref} != sample {i} spans {s}" + ) + + +def piecewise_attn( + query, # (B, Sq, H, D) + key, + value, + full_attn_spans: list[list[tuple[int, int]]], + softmax_scale: float, + attn_func, +): + 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) + + 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 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 06cc6cfde..0125b319a 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -1054,9 +1054,12 @@ class ImageKVCacheManager: attention_mask = attention_mask.contiguous() + full_attn_spans = kwargs.get("full_attn_spans", None) + if self.sp_size <= 1: attn_metadata = AttentionMetadata( attn_mask=attention_mask, + full_attn_spans=full_attn_spans, ) else: attn_metadata = AttentionMetadata( @@ -1065,6 +1068,7 @@ class ImageKVCacheManager: joint_value=joint_text_value, joint_strategy="front", attn_mask=attention_mask, + full_attn_spans=full_attn_spans, ) attn_output = self.attn(query, key, value, attn_metadata) attn_output = attn_output.reshape(bs * q_len, head_num_per_rank, head_dim) @@ -2246,6 +2250,7 @@ class HunyuanImage3Model(nn.Module): num_image_tokens: int | None = None, gen_timestep_scatter_index: torch.Tensor | None = None, uncond_cfg_prefill: bool = False, + full_attn_spans: list[list[tuple[int, int]]] | None = None, ) -> tuple | BaseModelOutputWithPast: output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( @@ -2354,6 +2359,7 @@ class HunyuanImage3Model(nn.Module): shard_image_size=shard_image_size, shard_padding_size=shard_padding_size, uncond_cfg_prefill=uncond_cfg_prefill, + full_attn_spans=full_attn_spans, ) hidden_states = layer_outputs[0] @@ -2594,6 +2600,10 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): if key in model_kwargs and model_kwargs[key] is not None: model_kwargs[key] = model_kwargs[key][s] + # List[List[...]] per-sample metadata indexed along the CFG batch dim + if isinstance(model_kwargs.get("full_attn_spans"), list): + model_kwargs["full_attn_spans"] = model_kwargs["full_attn_spans"][s.start : s.stop] + # custom_pos_emb: tuple of (cos, sin) if "custom_pos_emb" in model_kwargs and model_kwargs["custom_pos_emb"] is not None: cos, sin = model_kwargs["custom_pos_emb"] @@ -2724,6 +2734,9 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): query_lens=[prefill_query_len], seq_lens=[prefill_seq_len], num_image_tokens=0, + full_attn_spans=model_kwargs["full_attn_spans"][batch_slice] + if model_kwargs.get("full_attn_spans") + else None, ) # ========================================================== diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 3a4d0d2c6..7a8be0745 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import logging -import os from collections.abc import Iterable from typing import Any @@ -340,11 +339,6 @@ class HunyuanImage3Pipeline( ) ] quant_config = od_config.quantization_config - os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA" - logger.info( - "Setting attention backend to TORCH_SDPA. " - "HunyuanImage3Pipeline only supports TORCH_SDPA to handle mixed causal and full attention." - ) self.model = HunyuanImage3Model(self.hf_config, quant_config=quant_config) self.transformer = self.model self.vae = AutoencoderKLConv3D.from_config(self.hf_config.vae) @@ -962,10 +956,18 @@ class HunyuanImage3Pipeline( tokenizer_output.joint_image_slices[i] + tokenizer_output.gen_image_slices[i] for i in range(bsz) ] attention_mask = torch.ones(seq_len, seq_len, dtype=torch.bool).tril(diagonal=0).repeat(bsz, 1, 1) + full_attn_spans: list[list[tuple[int, int]]] = [[] for _ in range(bsz)] for i in range(bsz): for j, image_slice in enumerate(batch_image_slices[i]): attention_mask[i, image_slice, image_slice] = True + start = image_slice.start if image_slice.start is not None else 0 + stop = image_slice.stop if image_slice.stop is not None else seq_len + assert start < stop, f"Invalid image slice: {image_slice}" + full_attn_spans[i].append((int(start), int(stop))) + if full_attn_spans[i]: + full_attn_spans[i].sort(key=lambda x: x[0]) attention_mask = attention_mask.unsqueeze(1) + model_kwargs["full_attn_spans"] = full_attn_spans return attention_mask def prepare_inputs_for_generation( @@ -1009,6 +1011,7 @@ class HunyuanImage3Pipeline( "query_lens": kwargs.get("query_lens"), "seq_lens": kwargs.get("seq_lens"), "num_image_tokens": kwargs.get("num_image_tokens"), + "full_attn_spans": kwargs.get("full_attn_spans"), } ) return model_inputs @@ -1027,6 +1030,8 @@ class HunyuanImage3Pipeline( "custom_pos_emb": model_kwargs["custom_pos_emb"], "num_image_tokens": model_kwargs["num_image_tokens"], } + if "full_attn_spans" in model_kwargs: + updated_model_kwargs["full_attn_spans"] = model_kwargs["full_attn_spans"] # update past_key_values keeping its naming used in model code for possible_cache_name in ALL_CACHE_NAMES: @@ -1057,6 +1062,7 @@ class HunyuanImage3Pipeline( torch.arange(bsz), model_kwargs["gen_timestep_scatter_index"][:, -1] ].unsqueeze(-1) updated_model_kwargs["position_ids"] = torch.cat([timestep_position_ids, position_ids], dim=1) + # attention mask mask_list = [] for attention_mask_i, position_ids_i in zip( @@ -1156,6 +1162,7 @@ class HunyuanImage3Pipeline( seq_lens: list[int] | None = None, num_image_tokens: int | None = None, uncond_cfg_prefill: bool = False, + full_attn_spans: list[list[tuple[int, int]]] | None = None, ) -> tuple | CausalMMOutputWithPast: return_dict = return_dict if return_dict is not None else self.config.use_return_dict # Sanity Check of Inputs @@ -1251,6 +1258,7 @@ class HunyuanImage3Pipeline( num_image_tokens=num_image_tokens, gen_timestep_scatter_index=gen_timestep_scatter_index, uncond_cfg_prefill=uncond_cfg_prefill, + full_attn_spans=full_attn_spans, ) hidden_states = outputs[0] diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index ad008c397..d673ad166 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -995,6 +995,11 @@ class OmniKVTransferManager: logger.info(f"Skip receiving KV cache for {request_id} (need_recv_cache=False)") return None, 0 + # Skip during warmup dummy run — no sender is available. + if request_id == "dummy_req_id": + logger.info("Skip receiving KV cache for dummy warmup request") + return None, 0 + timeout = self.config.recv_timeout start_time = time.time() poll_interval = 0.01