From ab459b36c26d5c7903f623156d88ab2ae06f0a97 Mon Sep 17 00:00:00 2001 From: RuQing Xu <7891482+xrq-phys@users.noreply.github.com> Date: Fri, 7 Aug 2026 17:14:26 +0900 Subject: [PATCH] [Kernel] Refresh FlashInfer attention; Add quantized attention support (Blackwell QK16/V8) (#5344) Signed-off-by: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com> --- .../diffusion/attention_backends.md | 25 +- .../attention/backends/flashinfer_attn.py | 272 +++++++++++++++--- 2 files changed, 259 insertions(+), 38 deletions(-) diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md index 7dc35b367..094a0f3a9 100644 --- a/docs/user_guide/diffusion/attention_backends.md +++ b/docs/user_guide/diffusion/attention_backends.md @@ -17,7 +17,7 @@ The full set of backends and their platform defaults is in the **Backend Options | `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). Dense BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Packed paths can provide `cu_seqlens` directly. Supports optional **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. | | `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. | | `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). | -| `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. | +| `FLASHINFER_ATTN` | Uses FlashInfer's batch-prefill wrapper and supports mixed Q/K and V input dtypes through backend-specific configuration. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer` >= 0.6.16rc1 for mixed-dtype configurations. | | `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. | | `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. | | `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. | @@ -84,7 +84,7 @@ A backend that needs configuration exposes it as a typed field on the spec: When constructing `OmniDiffusionConfig` directly: ```python -from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec, OmniDiffusionConfig config = OmniDiffusionConfig( diffusion_attention_config=AttentionConfig( @@ -99,6 +99,27 @@ config = OmniDiffusionConfig( A plain dict is also accepted and normalized to `AttentionConfig`. +#### FlashInfer QK16/V8 quantized attention + +`FLASHINFER_ATTN` accepts quantization options through `AttentionSpec.quant`. For QK16/V8: + +```python +config = OmniDiffusionConfig( + diffusion_attention_config=AttentionConfig( + default=AttentionSpec( + backend="FLASHINFER_ATTN", + quant=AttnQuantSpec( + dtype_qk="bfloat16", + dtype_vo="fp8_e4m3", + ), + ), + ), + ..., +) +``` + +`dtype_qk` configures the Q and K input dtypes. `dtype_vo` configures the V input dtype. + ## Platform Defaults ### Blackwell (sm_100 / sm_103 / sm_120 / sm_121) diff --git a/vllm_omni/diffusion/attention/backends/flashinfer_attn.py b/vllm_omni/diffusion/attention/backends/flashinfer_attn.py index 1d2fa893b..188875818 100644 --- a/vllm_omni/diffusion/attention/backends/flashinfer_attn.py +++ b/vllm_omni/diffusion/attention/backends/flashinfer_attn.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass + import torch +from packaging.version import InvalidVersion, Version from vllm.logger import init_logger from vllm_omni.diffusion.attention.backends.abstract import ( @@ -13,7 +16,8 @@ from vllm_omni.diffusion.attention.backends.abstract import ( logger = init_logger(__name__) try: - from flashinfer.prefill import single_prefill_with_kv_cache + import flashinfer + from flashinfer.prefill import BatchPrefillWithRaggedKVCacheWrapper HAS_FLASHINFER = True except Exception as e: @@ -29,10 +33,8 @@ class FlashInferAttentionBackend(AttentionBackend): @classmethod def supports_attention_mask(cls) -> bool: - # FlashInfer single_prefill_with_kv_cache accepts a ``custom_mask`` kwarg - # (2D boolean, ``True`` = keep) for non-causal attention. See reviewer - # comment on #3079 and the flashinfer docs: - # https://docs.flashinfer.ai/generated/flashinfer.prefill.single_prefill_with_kv_cache.html + # Non-CuTe batch-prefill backends accept boolean custom masks. CuTe-DSL + # falls back to SDPA for nontrivial masks. return True @staticmethod @@ -52,6 +54,27 @@ class FlashInferAttentionBackend(AttentionBackend): class FlashInferAttentionImpl(AttentionImpl): + _QK_DTYPES = {torch.float16, torch.bfloat16} + _VO_DTYPES = {torch.float16, torch.bfloat16, torch.float8_e4m3fn} + + @dataclass(frozen=True) + class _WrapperPlanKey: + batch_size: int + qo_len: int + kv_len: int + num_q_heads: int + num_kv_heads: int + head_dim_qk: int + head_dim_k: int + head_dim_vo: int + q_dtype: torch.dtype + k_dtype: torch.dtype + v_dtype: torch.dtype + o_dtype: torch.dtype + causal: bool + softmax_scale: float + has_custom_mask: bool + def __init__( self, num_heads: int, @@ -60,10 +83,92 @@ class FlashInferAttentionImpl(AttentionImpl): causal: bool = False, num_kv_heads: int | None = None, prefix: str = "", + backend_kwargs: dict | None = None, **extra_impl_args, ) -> None: self.causal = causal self.softmax_scale = softmax_scale + self.device = torch.device("cuda", torch.accelerator.current_device_index()) + backend_kwargs = backend_kwargs or {} + quant = backend_kwargs.get("quant") or {} + self.dtype_qk = self._check_dtype(quant.get("dtype_qk"), "dtype_qk", self._QK_DTYPES) + self.dtype_vo = self._check_dtype(quant.get("dtype_vo"), "dtype_vo", self._VO_DTYPES) + requested_backend = quant.get("flashinfer_backend", "auto") + + if not HAS_FLASHINFER: + raise ImportError("FLASHINFER_ATTN backend requires flashinfer") + + self._check_flashinfer_version() + + self.flashinfer_backend = self._select_backend(requested_backend, self.device) + workspace_size = 0 if self.flashinfer_backend == "cute-dsl" else 128 * 1024 * 1024 + self._workspace = torch.empty( + workspace_size, + device=self.device, + dtype=torch.uint8, + ) + self._wrapper = BatchPrefillWithRaggedKVCacheWrapper( + self._workspace, + kv_layout="NHD", + backend=self.flashinfer_backend, + ) + self._qo_indptr: torch.Tensor | None = None + self._kv_indptr: torch.Tensor | None = None + self._plan_key: FlashInferAttentionImpl._WrapperPlanKey | None = None + + if self.dtype_qk is not None or self.dtype_vo is not None: + logger.info_once( + "FLASHINFER_ATTN dtype override: Q/K=%s, V=%s.", + self.dtype_qk, + self.dtype_vo, + ) + + logger.info_once( + "FLASHINFER_ATTN initialized backend=%s on %s.", + self.flashinfer_backend, + self.device, + ) + + def _check_flashinfer_version(self) -> None: + if self.dtype_qk == self.dtype_vo: + return + try: + flashinfer_version = Version(flashinfer.__version__) + except (AttributeError, InvalidVersion): + return + if flashinfer_version <= Version("0.6.15"): + raise RuntimeError( + f"FlashInfer {flashinfer_version} is too old for reliable mixed " + f"QK/V dtype attention (Q/K={self.dtype_qk}, V={self.dtype_vo}); " + "install flashinfer >= 0.6.16rc1." + ) + + @staticmethod + def _select_backend(requested_backend: str, device: torch.device) -> str: + if requested_backend != "auto": + return requested_backend + major, _minor = torch.cuda.get_device_capability(device) + if major >= 10: + return "cute-dsl" + if major >= 9: + return "fa3" + return "fa2" + + @classmethod + def _check_dtype( + cls, + dtype: torch.dtype | str | None, + option_name: str, + allowed: set[torch.dtype], + ) -> torch.dtype | None: + if dtype is None: + return None + if isinstance(dtype, str): + dtype = torch.float8_e4m3fn if dtype == "fp8_e4m3" else getattr(torch, dtype) + if dtype not in allowed: + choices = ", ".join(sorted(str(item) for item in allowed)) + raise ValueError(f"Unsupported {option_name}={dtype}; expected one of: {choices}") + return dtype @staticmethod def _pack_mask_for_flashinfer( @@ -128,6 +233,126 @@ class FlashInferAttentionImpl(AttentionImpl): ) return impl.forward_cuda(query, key, value, attn_metadata) + @torch.compiler.disable + def _plan_wrapper( + self, + key: _WrapperPlanKey, + flat_mask: torch.Tensor | None, + ) -> None: + self._wrapper.plan( + self._qo_indptr, + self._kv_indptr, + key.num_q_heads, + key.num_kv_heads, + key.head_dim_qk, + head_dim_vo=key.head_dim_vo, + custom_mask=flat_mask, + causal=key.causal, + sm_scale=key.softmax_scale, + q_data_type=key.q_dtype, + # FlashInfer versions newer than 0.6.15 can dispatch K and V with + # independent runtime dtypes even though plan() names this K/V. + kv_data_type=key.k_dtype, + o_data_type=key.o_dtype, + ) + + @torch.compiler.disable + def _make_indptr(self, batch_size: int, sequence_length: int) -> torch.Tensor: + return ( + torch.arange( + batch_size + 1, + device=self.device, + dtype=torch.int32, + ) + * sequence_length + ) + + def _ensure_plan( + self, + key: _WrapperPlanKey, + flat_mask: torch.Tensor | None, + ) -> None: + key_changed = key != self._plan_key + if not key_changed and flat_mask is None: + return + + if key_changed: + self._qo_indptr = self._make_indptr(key.batch_size, key.qo_len) + self._kv_indptr = self._make_indptr(key.batch_size, key.kv_len) + + # A custom mask's contents may change without its shape changing, so + # copy/re-plan it for every masked invocation. + self._plan_wrapper(key, flat_mask) + self._plan_key = key + + def _run_batch_prefill( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + custom_mask: torch.Tensor | None, + ) -> torch.Tensor: + if query.device != self.device or key.device != self.device or value.device != self.device: + raise ValueError( + "FLASHINFER_ATTN inputs must remain on the layer initialization " + f"device {self.device}; got Q={query.device}, K={key.device}, " + f"V={value.device}" + ) + + batch_size, qo_len, num_q_heads, head_dim_qk = query.shape + kv_len = key.shape[1] + num_kv_heads = key.shape[2] + head_dim_k = key.shape[3] + head_dim_vo = value.shape[3] + + q = query.reshape(batch_size * qo_len, num_q_heads, head_dim_qk) + k = key.reshape(batch_size * kv_len, num_kv_heads, head_dim_k) + v = value.reshape(batch_size * kv_len, num_kv_heads, head_dim_vo) + if self.dtype_qk is not None: + q = q.to(self.dtype_qk) + k = k.to(self.dtype_qk) + if self.dtype_vo is not None: + v = v.to(self.dtype_vo) + + flat_mask = None + if custom_mask is not None: + if custom_mask.dim() == 2: + custom_mask = custom_mask.unsqueeze(0).expand(batch_size, -1, -1) + flat_mask = custom_mask.contiguous().view(-1) + + self._ensure_plan( + FlashInferAttentionImpl._WrapperPlanKey( + batch_size=batch_size, + qo_len=qo_len, + kv_len=kv_len, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + head_dim_qk=head_dim_qk, + head_dim_k=head_dim_k, + head_dim_vo=head_dim_vo, + q_dtype=q.dtype, + k_dtype=k.dtype, + v_dtype=v.dtype, + o_dtype=query.dtype, + causal=self.causal, + softmax_scale=self.softmax_scale, + has_custom_mask=flat_mask is not None, + ), + flat_mask, + ) + out = self._run_wrapper(q, k, v) + out = out.reshape(batch_size, qo_len, num_q_heads, head_dim_vo) + return out.to(query.dtype) if out.dtype != query.dtype else out + + @torch.compiler.disable + def _run_wrapper( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + ) -> torch.Tensor: + return self._wrapper.run(query, key, value) + def forward_cuda( self, query: torch.Tensor, @@ -135,18 +360,7 @@ class FlashInferAttentionImpl(AttentionImpl): value: torch.Tensor, attn_metadata: AttentionMetadata | None = None, ) -> torch.Tensor: - if not HAS_FLASHINFER: - raise ImportError( - "FLASHINFER_ATTN backend requires flashinfer. " - "Install it or set DIFFUSION_ATTENTION_BACKEND to another backend." - ) - - # Try the custom_mask path; if the mask can't be packed into the - # (qo_len, kv_len) layout FlashInfer expects, fall back to SDPA - # instead of risking an illegal-memory-access crash in the kernel. - # Input layout is (B, S, H, D); FlashInfer dense prefill takes (S, H, D). batch_size = query.shape[0] - custom_mask = None if attn_metadata is not None and attn_metadata.attn_mask is not None: try: @@ -156,28 +370,14 @@ class FlashInferAttentionImpl(AttentionImpl): qo_len=query.shape[1], kv_len=key.shape[1], ) - except ValueError as e: - logger.debug("Falling back to SDPA for mask path: %s", e) + except ValueError as error: + logger.debug("Falling back to SDPA for mask path: %s", error) return self._sdpa_fallback(query, key, value, attn_metadata) - # FlashInfer cannot combine causal masking with a custom_mask; rather - # than silently dropping the explicit mask (diverging from SDPA), let - # SDPA handle the causal+mask case correctly. if custom_mask is not None and self.causal: - logger.debug("causal=True with explicit attn_mask; deferring to SDPA") return self._sdpa_fallback(query, key, value, attn_metadata) - outputs = [] - for b in range(batch_size): - kwargs: dict = { - "sm_scale": self.softmax_scale, - "causal": self.causal, - "return_lse": False, - } - if custom_mask is not None: - kwargs["custom_mask"] = custom_mask if custom_mask.dim() == 2 else custom_mask[b] - out = single_prefill_with_kv_cache(query[b], key[b], value[b], **kwargs) - outputs.append(out) + if custom_mask is not None and self.flashinfer_backend == "cute-dsl": + logger.debug("CuTe DSL does not support custom masks; deferring to SDPA") + return self._sdpa_fallback(query, key, value, attn_metadata) - if batch_size == 1: - return outputs[0].unsqueeze(0) - return torch.stack(outputs, dim=0) + return self._run_batch_prefill(query, key, value, custom_mask)