[Feature] hunyuanimage support flash attn (#2981)

Signed-off-by: dengyunyang <584797741@qq.com>

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)

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)
 

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(

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

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,
         )
 
     # ==========================================================

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]
 

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