[Kernel] Refresh FlashInfer attention; Add quantized attention support (Blackwell QK16/V8) (#5344)

Signed-off-by: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com>

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)

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)