diff --git a/recipes/Qwen/Qwen-Image-2.1.md b/recipes/Qwen/Qwen-Image-2.1.md index f833ac58e..31c43002e 100644 --- a/recipes/Qwen/Qwen-Image-2.1.md +++ b/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; diff --git a/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py b/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py index 5c840044d..1f568780a 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/decode_graph.py +++ b/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, diff --git a/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py b/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py index ca5bc594f..aecee9368 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/pipeline_qwen_image_21.py +++ b/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 diff --git a/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py b/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py index 7200cf7c7..07cf31645 100644 --- a/vllm_omni/diffusion/models/qwen_image_21/qwen_image_21_transformer.py +++ b/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) ]