Add FP8 KV cache support for Qwen-Image-2.1 DiT (#7797)

recipes/Qwen/Qwen-Image-2.1.md

@@ -259,6 +259,58 @@ BF16 vision tower, so it is unaffected by text-encoder FP8. See
 [`docs/user_guide/quantization/fp8.md`](../../docs/user_guide/quantization/fp8.md)
 for the scope rules.
 
+#### FP8 prefix KV storage
+
+The cached prefix K/V can be stored in FP8 E4M3 (with per-token-per-head fp32
+scales), halving the prefix-cache memory — relevant for long prompts (up to
+8192 text tokens) and up to 4 condition images (~4096 latent tokens each).
+Prefill quantizes once; decode dequantizes per step, so attention compute is
+unchanged. Enable it through the model-specific `extras` config (e.g. in the
+deploy YAML's stage `extras` or as an `Omni(...)` engine kwarg):
+
+```yaml
+extras:
+  prefix_kv_cache_dtype: "fp8"
+```
+
+Accepted values: `"fp8"` / `"fp8_e4m3"` (K and V quantized), `"fp8_v"` (V
+only — K stays in the native dtype), and `None` / `"auto"` (default, native
+dtype — behavior unchanged). This is independent of
+`diffusion_kv_cache_dtype`, which quantizes attention Q/K/V *compute* per
+forward pass on supported backends.
+
+Measured on T2I 1024x1024 (seed 42, 50 steps, true CFG 4.0, PSNR vs. the
+bf16 baseline): **`"fp8"` 34.9 dB**, **`"fp8_v"` 40.9 dB** — same composition
+and semantics, with texture-level drift in fine detail for `"fp8"`, and
+visually indistinguishable output for `"fp8_v"`. The error is dominated by
+K quantization: post-RoPE keys are the precision-sensitive half of the
+cache, so `"fp8_v"` buys back ~6 dB at 75% (instead of 50%) of the original
+cache size. The residual error is inherent e4m3 precision accumulated
+coherently over the denoising trajectory; finer scale granularity
+(per-tensor/per-head/per-token were compared) or Hadamard-rotated V did not
+improve it. The saving scales with prefix length — at the limit
+(8192 text tokens + 4 condition images, ~24.6k prefix tokens) the prefix
+cache is ~12.9 GB per CFG branch in bf16, ~6.6 GB in `"fp8"` and ~9.7 GB in
+`"fp8_v"`, which is where this option matters. Treat it as an opt-in for
+memory-bound long-prompt / multi-image workloads, not a free lunch;
+`"fp8_v"` is the better default trade-off when quality matters.
+
+Note: a quantized prefix cache is not CUDA-graph capturable — requests using
+this option stay on the eager decode path (a warning is logged once).
+
+#### Why not the scheduler-managed paged KV (`DiffusionKVCacheMode.PAGED_SCHEDULER`)?
+
+Evaluated and deliberately not adopted (see the `qwen21-p0-kv` commit message
+for the full analysis): the paged scheduler targets *cross-request* KV
+management — a statically sized page pool, block sharing/dedup, and eviction —
+whereas this prefix cache lives exactly as long as one generation and is
+sized once per request. Migrating would require routing decode attention
+through the worker paged adapter (writing fresh target K/V into pages every
+step), per-request page tables, and re-validating the block-causal
+piecewise-span masking and Ulysses/CFG constraints, for no memory benefit over
+the per-request dense tensors. The FP8 storage above captures the actual
+memory win at ~1% of the integration cost.
+
 ## Known Limitations
 
 - Cache acceleration backends (`cache_dit`, `tea_cache`) are not supported;

vllm_omni/diffusion/models/qwen_image_21/decode_graph.py

@@ -244,6 +244,16 @@ class QwenImage21DecodeGraphManager:
         first = kv_cache[0].get(cache_branch)
         if first is None or "key" not in first:
             return
+        if any(name.endswith("_scale") for name in first):
+            # Quantized (FP8) prefix storage: the static K/V buffers assume a single
+            # dtype and the captured body would not see the dequant scales (a bf16 K +
+            # fp8 V mix would even copy_ fp8 into bf16 buffers, silently dropping the
+            # scales). Keep quantized-cache requests on the eager decode path.
+            logger.warning_once(
+                "Qwen-Image-2.1 CUDA graph decode: quantized prefix KV cache "
+                "(prefix_kv_cache_dtype) is not graph-capturable; falling back to eager decode."
+            )
+            return
         batch_size = first["key"].shape[0]
         key = (
             cache_branch,

vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py

@@ -1225,16 +1225,18 @@ class QwenImage21Pipeline(
             return [{} for _ in range(num_blocks)], True
         if phases == {"decode"}:
             # Decode only reads the cache, so a batched view is enough and the
-            # per-request caches stay untouched.
+            # per-request caches stay untouched. Merge every stored part — FP8
+            # cache entries carry per-request scale tensors alongside "key"/"value".
             num_blocks = len(caches[0])
             merged = []
             for block_idx in range(num_blocks):
                 branches: dict[str, dict[str, torch.Tensor]] = {}
                 branch_names = {branch for cache in caches for branch in cache[block_idx]}
                 for branch in branch_names:
+                    part_names = {part for cache in caches for part in cache[block_idx][branch]}
                     branches[branch] = {
                         part: torch.cat([cache[block_idx][branch][part] for cache in caches], dim=0)
-                        for part in ("key", "value")
+                        for part in sorted(part_names)
                     }
                 merged.append(branches)
             return merged, False

vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py

@@ -40,6 +40,46 @@ if TYPE_CHECKING:
 
 logger = init_logger(__name__)
 
+_FP8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
+
+# Accepted values for the `prefix_kv_cache_dtype` switch (see QwenImage21Transformer2DModel).
+_PREFIX_KV_FP8_ALIASES = {"fp8", "fp8_e4m3", "fp8_e4m3fn"}
+_PREFIX_KV_FP8_V_ALIASES = {"fp8_v", "fp8_e4m3_v"}
+_PREFIX_KV_FP8_V = "fp8_e4m3_v"
+
+
+def _quantize_prefix_kv_fp8(t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per-token-per-head symmetric FP8 E4M3 quantization of a (B, S, H, D) K/V tensor.
+
+    One fp32 scale per (batch, token, head): it tracks the heavy-tailed per-token magnitude of
+    post-RoPE keys and outlier-channel values far better than a per-tensor scale, and the scale
+    tensor stays mergeable across the batch dim for batched decode (`_assemble_kv_cache`).
+    """
+    scale = t.abs().amax(dim=-1, keepdim=True).float().clamp_min(1e-12) / _FP8_E4M3_MAX
+    quantized = (t / scale).clamp(-_FP8_E4M3_MAX, _FP8_E4M3_MAX).to(torch.float8_e4m3fn)
+    return quantized, scale
+
+
+def _dequantize_prefix_kv_fp8(quantized: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    return (quantized.float() * scale).to(dtype)
+
+
+def _normalize_prefix_kv_cache_dtype(value: Any) -> str | None:
+    """Normalize the `prefix_kv_cache_dtype` switch: None/"auto" = off, "fp8*" = FP8 E4M3 storage."""
+    if value is None or value == "auto":
+        return None
+    if value in _PREFIX_KV_FP8_ALIASES:
+        return "fp8_e4m3"
+    if value in _PREFIX_KV_FP8_V_ALIASES:
+        # V-only FP8 storage: K stays in the native dtype. Measured nearly lossless on this
+        # model (PSNR 40.9 dB vs bf16 vs 34.9 dB for K+V fp8) — post-RoPE K is the
+        # precision-sensitive half of the cache.
+        return _PREFIX_KV_FP8_V
+    raise ValueError(
+        f"Unknown prefix_kv_cache_dtype {value!r}; expected None, 'auto' or one of "
+        f"{sorted(_PREFIX_KV_FP8_ALIASES | _PREFIX_KV_FP8_V_ALIASES)}."
+    )
+
 
 def _apply_qwen_image21_rotary_emb(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
     """Rotate interleaved pairs with complex FP32 multiplication.
@@ -390,9 +430,11 @@ class QwenImage21Attention(nn.Module):
         eps: float = 1e-6,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
+        prefix_kv_cache_dtype: str | None = None,
     ):
         super().__init__()
         self.head_dim = dim_head
+        self.prefix_kv_cache_dtype = prefix_kv_cache_dtype
 
         self.to_qkv = QKVParallelLinear(
             hidden_size=dim,
@@ -460,7 +502,8 @@ class QwenImage21Attention(nn.Module):
             sp_decode (`bool`): KV-cache decode under sequence parallelism: the cached prefix K/V are replicated
                 and join through the Ulysses joint mechanism with an empty joint query.
             kv_cache (dict, *optional*): Per-block cache keyed by CFG branch (`"cond"` / `"uncond"`); each branch
-                stores `{"key", "value"}` of the timestep-independent prefix (text + condition images).
+                stores `{"key", "value"}` of the timestep-independent prefix (text + condition images). With
+                FP8 cache storage enabled, the branch additionally carries `{"key_scale", "value_scale"}`.
             cache_branch (`str`): CFG branch whose cache entry is read/written this forward.
             cache_write_len (`int`, *optional*): Prefill mode — cache the first `cache_write_len` K/V positions.
                 `None` with a populated branch means decode — cached prefix K/V are prepended instead.
@@ -487,12 +530,32 @@ class QwenImage21Attention(nn.Module):
             branch_cache = kv_cache.setdefault(cache_branch, {})
             if cache_write_len is not None:
                 # Prefill: cache the timestep-independent prefix K/V for later denoising steps.
-                # Own the prefix storage instead of retaining the full prefill tensors.
-                branch_cache["key"] = key[:, :cache_write_len].clone()
-                branch_cache["value"] = value[:, :cache_write_len].clone()
+                if self.prefix_kv_cache_dtype is not None:
+                    # FP8 storage: quantize once at prefill; decode dequantizes per step.
+                    # The fp32 scales ride along as extra branch entries so the pipeline's
+                    # batched decode merge can concatenate them like the K/V tensors. The
+                    # quantized tensors own their storage, so no clone is needed here.
+                    if self.prefix_kv_cache_dtype == "fp8_e4m3":
+                        prefix_key, branch_cache["key_scale"] = _quantize_prefix_kv_fp8(key[:, :cache_write_len])
+                        branch_cache["key"] = prefix_key
+                    else:
+                        # V-only FP8 ("fp8_e4m3_v"): K stays in the native dtype.
+                        branch_cache["key"] = key[:, :cache_write_len].clone()
+                    prefix_value, branch_cache["value_scale"] = _quantize_prefix_kv_fp8(value[:, :cache_write_len])
+                    branch_cache["value"] = prefix_value
+                else:
+                    # Own the prefix storage instead of retaining the full prefill tensors.
+                    branch_cache["key"] = key[:, :cache_write_len].clone()
+                    branch_cache["value"] = value[:, :cache_write_len].clone()
             else:
                 cached_key = branch_cache["key"]
                 cached_value = branch_cache["value"]
+                if self.prefix_kv_cache_dtype is not None:
+                    # Dequantize only entries that were actually quantized (scale present).
+                    if "key_scale" in branch_cache:
+                        cached_key = _dequantize_prefix_kv_fp8(cached_key, branch_cache["key_scale"], key.dtype)
+                    if "value_scale" in branch_cache:
+                        cached_value = _dequantize_prefix_kv_fp8(cached_value, branch_cache["value_scale"], value.dtype)
 
         metadata = copy.copy(attn_metadata) if attn_metadata is not None else None
         if sp_prefix_len > 0:
@@ -545,6 +608,7 @@ class QwenImage21TransformerBlock(nn.Module):
         eps: float = 1e-6,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
+        prefix_kv_cache_dtype: str | None = None,
     ):
         super().__init__()
         self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
@@ -555,6 +619,7 @@ class QwenImage21TransformerBlock(nn.Module):
             eps=eps,
             quant_config=quant_config,
             prefix=f"{prefix}.attn",
+            prefix_kv_cache_dtype=prefix_kv_cache_dtype,
         )
         self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
         self.img_mlp = QwenImage21SwiGLUFeedForward(
@@ -696,7 +761,9 @@ class QwenImage21Transformer2DModel(CachedTransformer):
       batch is homogeneous and unpadded, or a dense 4D boolean `attn_mask` otherwise.
     - `causal_condition` — text and condition-image tokens are modulated from `t = 0` instead of the sampled
       timestep, which also makes their activations timestep-independent and so cacheable across denoising steps
-      (`kv_cache`).
+      (`kv_cache`). The cached prefix K/V are stored in the running dtype by default, or in FP8 E4M3 (with
+      per-token-per-head fp32 scales) when `od_config.extras["prefix_kv_cache_dtype"]` is `"fp8"` (K and V)
+      or `"fp8_v"` (V only — K stays in the native dtype).
     """
 
     # the small and frequently-repeated block(s) of a model
@@ -763,6 +830,14 @@ class QwenImage21Transformer2DModel(CachedTransformer):
         self.quant_config = _enable_pattern_ignored_layers(quant_config)
         quant_config = self.quant_config
 
+        # Storage dtype for the cross-step prefix KV cache, from the supplementary
+        # model-specific config entry `extras["prefix_kv_cache_dtype"]` (None = bf16/fp32
+        # native storage, unchanged behavior; "fp8" = FP8 E4M3 quantized storage, halving
+        # prefix K/V memory). This is independent of `diffusion_kv_cache_dtype`, which
+        # quantizes attention Q/K/V *compute* per forward on supported backends.
+        extras = getattr(od_config, "extras", None) or {}
+        self.prefix_kv_cache_dtype = _normalize_prefix_kv_cache_dtype(extras.get("prefix_kv_cache_dtype"))
+
         self.pos_embed = QwenImage21Rope(theta=10000, axes_dim=list(axes_dims_rope))
         self.time_text_embed = QwenImage21TimestepProjEmbeddings(
             embedding_dim=self.inner_dim, prefix=f"{prefix}.time_text_embed"
@@ -802,6 +877,7 @@ class QwenImage21Transformer2DModel(CachedTransformer):
                     eps=eps,
                     quant_config=quant_config,
                     prefix=f"{prefix}.transformer_blocks.{i}",
+                    prefix_kv_cache_dtype=self.prefix_kv_cache_dtype,
                 )
                 for i in range(num_layers)
             ]