From bbc00f9f86e5bf54633737bedb7964ea4003e37d Mon Sep 17 00:00:00 2001 From: lyj-jjj Date: Fri, 15 May 2026 11:22:59 +0800 Subject: [PATCH] [BugFix] fix(omni): isolate diffusion KV-cache dtype from vLLM --kv-cache-dtype #3585 (#3596) Signed-off-by: lyj-jjj Co-authored-by: Cursor --- .../quantization/quantized_kvcache.md | 60 +++++++++++-------- .../image_to_video/image_to_video.py | 24 ++++---- vllm_omni/diffusion/attention/layer.py | 6 +- vllm_omni/diffusion/data.py | 30 +++++++--- vllm_omni/engine/async_omni_engine.py | 37 ++++++------ vllm_omni/entrypoints/cli/serve.py | 25 ++++---- 6 files changed, 104 insertions(+), 78 deletions(-) diff --git a/docs/user_guide/quantization/quantized_kvcache.md b/docs/user_guide/quantization/quantized_kvcache.md index 62815eb4b..e560032d3 100644 --- a/docs/user_guide/quantization/quantized_kvcache.md +++ b/docs/user_guide/quantization/quantized_kvcache.md @@ -8,15 +8,20 @@ vLLM-Omni supports online FP8 quantization for eligible diffusion Flash Attention (FA) to reduce FA latency while keeping model weights in their original dtype. -This feature is configured through `kv_cache_dtype`, matching the option name -used by vLLM's language-model KV-cache quantization. In vLLM-Omni diffusion -pipelines, however, it is a runtime FA path: Q/K/V tensors are dynamically -quantized before the attention operator. It does not quantize model weights and -is separate from [FP8 W8A8](fp8.md), [Int8 W8A8](int8.md), or pre-quantized -checkpoint formats. +This feature is configured through `diffusion_kv_cache_dtype` on +`OmniDiffusionConfig` (CLI: `--diffusion-kv-cache-dtype`). It is intentionally +**not** the same as vLLM's `--kv-cache-dtype`, which controls autoregressive +language-model KV cache storage and defaults to `"auto"`. Diffusion FA +quantization uses the dedicated diffusion flags so omni serve does not inherit +that default. -If `kv_cache_dtype` is not set, behavior is unchanged and attention runs in the -native dtype. +In vLLM-Omni diffusion pipelines, this is a runtime FA path: Q/K/V tensors are +dynamically quantized before the attention operator. It does not quantize model +weights and is separate from [FP8 W8A8](fp8.md), [Int8 W8A8](int8.md), or +pre-quantized checkpoint formats. + +If `diffusion_kv_cache_dtype` is not set, behavior is unchanged and attention +runs in the native dtype. ## Hardware Support @@ -30,8 +35,8 @@ native dtype. Legend: `✅` supported, `❌` unsupported. FP8 FA is currently implemented only for the NPU Flash Attention backend. Other -backends do not support `kv_cache_dtype="fp8"` for diffusion attention and fall -back to native dtype execution. +backends do not support `diffusion_kv_cache_dtype="fp8"` for diffusion attention +and fall back to native dtype execution. ## Model Type Support @@ -40,7 +45,7 @@ back to native dtype execution. | Model | Scope | Status | Notes | |-------|-------|--------|-------| | Wan2.2 | Eligible DiT full-attention FA on Ascend NPU | Tested | Compare quality and latency against a BF16 baseline before production use | -| Other diffusion models | Eligible DiT full-attention FA on Ascend NPU | Not tested | You can try `kv_cache_dtype="fp8"`; tune `kv_cache_skip_steps` and `kv_cache_skip_layers` when higher precision is needed | +| Other diffusion models | Eligible DiT full-attention FA on Ascend NPU | Not tested | You can try `diffusion_kv_cache_dtype="fp8"`; tune `diffusion_kv_cache_skip_steps` and `diffusion_kv_cache_skip_layers` when higher precision is needed | ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) @@ -50,8 +55,8 @@ guide documents support. ### Multi-Stage Diffusion Model (BAGEL, GLM-Image) Not tested. If the diffusion stage uses the same NPU Flash Attention backend, -`kv_cache_dtype` may apply in theory; validate quality and latency for each -stage and model. +`diffusion_kv_cache_dtype` may apply in theory; validate quality and latency for +each stage and model. ## Configuration @@ -67,15 +72,15 @@ python examples/offline_inference/image_to_video/image_to_video.py \ --num-inference-steps 4 \ --ulysses-degree 4 \ --vae-patch-parallel-size 4 \ - --kv-cache-dtype fp8 \ - --kv-cache-skip-steps "0,1" \ - --kv-cache-skip-layers "0-2" + --diffusion-kv-cache-dtype fp8 \ + --diffusion-kv-cache-skip-steps "0,1" \ + --diffusion-kv-cache-skip-layers "0-2" ``` Online serving: ```bash -vllm serve --omni --kv-cache-dtype fp8 +vllm serve --omni --diffusion-kv-cache-dtype fp8 ``` Stage config: @@ -86,18 +91,23 @@ stage_args: stage_type: diffusion engine_args: model_stage: dit - kv_cache_dtype: "fp8" - kv_cache_skip_steps: "0,1" - kv_cache_skip_layers: "0-2" + diffusion_kv_cache_dtype: "fp8" + diffusion_kv_cache_skip_steps: "0,1" + diffusion_kv_cache_skip_layers: "0-2" ``` +Legacy YAML keys `kv_cache_dtype`, `kv_cache_skip_steps`, and +`kv_cache_skip_layers` are still accepted when constructing +`OmniDiffusionConfig` (for example via `from_kwargs`); prefer the `diffusion_*` +names for new configs. + ## Parameters | Parameter | Type | Default | Description | |-----------|------|---------|-------------| -| `kv_cache_dtype` | str \| None | `None` | Set to `"fp8"` to enable dynamic FP8 FA on supported attention backends | -| `kv_cache_skip_steps` | str \| None | `None` | Denoising step selector to keep in native dtype, for example `"0,1,4-6"` | -| `kv_cache_skip_layers` | str \| None | `None` | Transformer layer selector to keep in native dtype, for example `"0-2,10"` | +| `diffusion_kv_cache_dtype` | str \| None | `None` | Set to `"fp8"` to enable dynamic FP8 FA on supported attention backends | +| `diffusion_kv_cache_skip_steps` | str \| None | `None` | Denoising step selector to keep in native dtype, for example `"0,1,4-6"` | +| `diffusion_kv_cache_skip_layers` | str \| None | `None` | Transformer layer selector to keep in native dtype, for example `"0-2,10"` | Selectors use comma-separated integers and inclusive ranges. Listed steps or layers skip FP8 FA; all other eligible full-attention forwards use the FP8 path. @@ -106,9 +116,9 @@ layers skip FP8 FA; all other eligible full-attention forwards use the FP8 path. 1. Compare generated images or videos against a BF16 baseline with the same seed, prompt, resolution, frame count, and denoising steps. -2. Use `kv_cache_skip_steps` for denoising steps where quality is more +2. Use `diffusion_kv_cache_skip_steps` for denoising steps where quality is more sensitive. -3. Use `kv_cache_skip_layers` for transformer layers that show visible quality +3. Use `diffusion_kv_cache_skip_layers` for transformer layers that show visible quality regressions. 4. Report both latency and quality results when enabling this option for a new model. For image or video models, include visual comparison and quantitative diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 84fbf2a94..4964f10c4 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -103,22 +103,22 @@ def parse_args() -> argparse.Namespace: help="Sampling solver for Wan2.2 pipelines. Use 'euler' for Lightning/Distill setups.", ) parser.add_argument( - "--kv-cache-dtype", + "--diffusion-kv-cache-dtype", type=str, default=None, - help="Config-level KV cache dtype (e.g. float8_e4m3fn).", + help="Diffusion attention KV cache dtype (e.g. float8_e4m3fn). Separate from vLLM --kv-cache-dtype.", ) parser.add_argument( - "--kv-cache-skip-steps", + "--diffusion-kv-cache-skip-steps", type=str, default=None, - help="Config-level KV-cache quantization skip-step selector, e.g. '0-9,20,25-30'.", + help="Diffusion KV-cache quantization skip-step selector, e.g. '0-9,20,25-30'.", ) parser.add_argument( - "--kv-cache-skip-layers", + "--diffusion-kv-cache-skip-layers", type=str, default=None, - help="Config-level KV-cache quantization skip-layer selector, e.g. '0,1,4-8'.", + help="Diffusion KV-cache quantization skip-layer selector, e.g. '0,1,4-8'.", ) parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).") parser.add_argument("--fps", type=int, default=None, help="Frames per second for the output video.") @@ -334,9 +334,9 @@ def main(): vae_use_tiling=args.vae_use_tiling, boundary_ratio=args.boundary_ratio, flow_shift=args.flow_shift, - kv_cache_dtype=args.kv_cache_dtype, - kv_cache_skip_steps=args.kv_cache_skip_steps, - kv_cache_skip_layers=args.kv_cache_skip_layers, + diffusion_kv_cache_dtype=args.diffusion_kv_cache_dtype, + diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps, + diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers, enable_cpu_offload=args.enable_cpu_offload, parallel_config=parallel_config, enforce_eager=args.enforce_eager, @@ -361,9 +361,9 @@ def main(): print(f" Inference steps: {args.num_inference_steps}") print(f" Frames: {args.num_frames}") print(f" Solver: {args.sample_solver}") - print(f" kv_cache_dtype(config): {args.kv_cache_dtype}") - print(f" kv_cache_skip_steps(config): {args.kv_cache_skip_steps}") - print(f" kv_cache_skip_layers(config): {args.kv_cache_skip_layers}") + print(f" diffusion_kv_cache_dtype(config): {args.diffusion_kv_cache_dtype}") + print(f" diffusion_kv_cache_skip_steps(config): {args.diffusion_kv_cache_skip_steps}") + print(f" diffusion_kv_cache_skip_layers(config): {args.diffusion_kv_cache_skip_layers}") print( f" Parallel configuration: cfg_parallel_size={args.cfg_parallel_size}," f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}" diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 4ab139d4b..5cbc39eb1 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -171,7 +171,7 @@ class Attention(nn.Module): def _init_kv_cache_quantization(self, config) -> None: if config is None: return - dtype = config.kv_cache_dtype + dtype = config.diffusion_kv_cache_dtype if dtype: if config.parallel_config.ring_degree > 1: raise ValueError( @@ -190,8 +190,8 @@ class Attention(nn.Module): ) dtype = None self._kv_cache_dtype = dtype - self._kv_cache_skip_steps = config.kv_cache_skip_step_indices - self._kv_cache_skip_layers = config.kv_cache_skip_layer_indices + self._kv_cache_skip_steps = config.diffusion_kv_cache_skip_step_indices + self._kv_cache_skip_layers = config.diffusion_kv_cache_skip_layer_indices def _should_apply_kv_cache_quant(self) -> bool: skip_steps = self._kv_cache_skip_steps diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 9cbc7b877..be95349b1 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -568,18 +568,18 @@ class OmniDiffusionConfig: # has already resolved to vLLM's ModelOpt FP8 linear method. force_cutlass_fp8: bool = False - # KV cache dtype for attention. Aligned with upstream vLLM's --kv-cache-dtype. + # Diffusion attention KV cache dtype (not vLLM's --kv-cache-dtype for AR models). # None = native dtype (no quantization). # "fp8" = dynamic FP8 (float8_e4m3fn) quantization per forward pass. # On Hopper+FA3: native FP8 attention (memory + compute savings). # On other backends: no benefit, backends skip quantization. - kv_cache_dtype: str | None = None + diffusion_kv_cache_dtype: str | None = None # Optional skip selectors for KV-cache quantization. Format: "0-9,20,25-30". # Listed steps/layers skip quantization; others keep quantized execution. - kv_cache_skip_steps: str | None = None - kv_cache_skip_layers: str | None = None - kv_cache_skip_step_indices: set[int] | None = None - kv_cache_skip_layer_indices: set[int] | None = None + diffusion_kv_cache_skip_steps: str | None = None + diffusion_kv_cache_skip_layers: str | None = None + diffusion_kv_cache_skip_step_indices: set[int] | None = None + diffusion_kv_cache_skip_layer_indices: set[int] | None = None # Diffusion pipeline Profiling config enable_diffusion_pipeline_profiler: bool = False @@ -726,8 +726,8 @@ class OmniDiffusionConfig: # Match vLLM's config flow: parse entrypoint shorthands before the # config object is built, and keep a single runtime truth source. self.diffusion_attention_config = build_attention_config(self.diffusion_attention_config) - self.kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.kv_cache_skip_steps) - self.kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.kv_cache_skip_layers) + self.diffusion_kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_steps) + self.diffusion_kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_layers) if self.max_cpu_loras is None: self.max_cpu_loras = 1 @@ -874,6 +874,20 @@ class OmniDiffusionConfig: else: kwargs.pop("quantization", None) + # Renamed from kv_cache_* to avoid clashing with vLLM's --kv-cache-dtype. + if kwargs.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in kwargs: + kwargs["diffusion_kv_cache_dtype"] = kwargs.pop("kv_cache_dtype") + else: + kwargs.pop("kv_cache_dtype", None) + if kwargs.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in kwargs: + kwargs["diffusion_kv_cache_skip_steps"] = kwargs.pop("kv_cache_skip_steps") + else: + kwargs.pop("kv_cache_skip_steps", None) + if kwargs.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in kwargs: + kwargs["diffusion_kv_cache_skip_layers"] = kwargs.pop("kv_cache_skip_layers") + else: + kwargs.pop("kv_cache_skip_layers", None) + # Handle "diffusion_attention_backend" shorthand: merge into # diffusion_attention_config before field filtering. diffusion_attn_backend = kwargs.pop("diffusion_attention_backend", None) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 35b323487..579ef0667 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1535,9 +1535,9 @@ class AsyncOmniEngine: "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4), "quantization": kwargs.get("quantization", None), - "kv_cache_dtype": kwargs.get("kv_cache_dtype", None), - "kv_cache_skip_steps": kwargs.get("kv_cache_skip_steps", None), - "kv_cache_skip_layers": kwargs.get("kv_cache_skip_layers", None), + "diffusion_kv_cache_dtype": kwargs.get("diffusion_kv_cache_dtype", None), + "diffusion_kv_cache_skip_steps": kwargs.get("diffusion_kv_cache_skip_steps", None), + "diffusion_kv_cache_skip_layers": kwargs.get("diffusion_kv_cache_skip_layers", None), **({"diffusion_attention_config": attention_config} if attention_config is not None else {}), "force_cutlass_fp8": bool(kwargs.get("force_cutlass_fp8", False)), "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False), @@ -1717,24 +1717,27 @@ class AsyncOmniEngine: if quantization is not None: if not hasattr(cfg.engine_args, "quantization") or cfg.engine_args.quantization is None: cfg.engine_args.quantization = quantization - kv_cache_dtype = kwargs.get("kv_cache_dtype") - if kv_cache_dtype is not None: - if not hasattr(cfg.engine_args, "kv_cache_dtype") or cfg.engine_args.kv_cache_dtype is None: - cfg.engine_args.kv_cache_dtype = kv_cache_dtype - kv_cache_skip_steps = kwargs.get("kv_cache_skip_steps") - if kv_cache_skip_steps is not None: + diffusion_kv_cache_dtype = kwargs.get("diffusion_kv_cache_dtype") + if diffusion_kv_cache_dtype is not None: if ( - not hasattr(cfg.engine_args, "kv_cache_skip_steps") - or cfg.engine_args.kv_cache_skip_steps is None + not hasattr(cfg.engine_args, "diffusion_kv_cache_dtype") + or cfg.engine_args.diffusion_kv_cache_dtype is None ): - cfg.engine_args.kv_cache_skip_steps = kv_cache_skip_steps - kv_cache_skip_layers = kwargs.get("kv_cache_skip_layers") - if kv_cache_skip_layers is not None: + cfg.engine_args.diffusion_kv_cache_dtype = diffusion_kv_cache_dtype + diffusion_kv_cache_skip_steps = kwargs.get("diffusion_kv_cache_skip_steps") + if diffusion_kv_cache_skip_steps is not None: if ( - not hasattr(cfg.engine_args, "kv_cache_skip_layers") - or cfg.engine_args.kv_cache_skip_layers is None + not hasattr(cfg.engine_args, "diffusion_kv_cache_skip_steps") + or cfg.engine_args.diffusion_kv_cache_skip_steps is None ): - cfg.engine_args.kv_cache_skip_layers = kv_cache_skip_layers + cfg.engine_args.diffusion_kv_cache_skip_steps = diffusion_kv_cache_skip_steps + diffusion_kv_cache_skip_layers = kwargs.get("diffusion_kv_cache_skip_layers") + if diffusion_kv_cache_skip_layers is not None: + if ( + not hasattr(cfg.engine_args, "diffusion_kv_cache_skip_layers") + or cfg.engine_args.diffusion_kv_cache_skip_layers is None + ): + cfg.engine_args.diffusion_kv_cache_skip_layers = diffusion_kv_cache_skip_layers except Exception as e: logger.warning("Failed to inject LoRA config for stage: %s", e) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index d714d1f53..f2689d4b5 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -472,26 +472,25 @@ class OmniServeCommand(CLISubcommand): default=None, help="Scheduler flow_shift for video models (e.g., 5.0 for 720p, 12.0 for 480p).", ) - # vLLM already registers --kv-cache-dtype for the serve parser. Keep - # this fallback only for older vLLM versions where the option is absent. - if "--kv-cache-dtype" not in serve_parser._option_string_actions: - omni_config_group.add_argument( - "--kv-cache-dtype", - type=str, - default=None, - help="Config-level KV cache dtype (e.g. fp8).", - ) + # Diffusion KV-cache quantization uses dedicated flags so we do not reuse + # vLLM's --kv-cache-dtype (AR cache dtype, default "auto"). omni_config_group.add_argument( - "--kv-cache-skip-steps", + "--diffusion-kv-cache-dtype", type=str, default=None, - help="Config-level KV-cache quantization skip-step selector, e.g. '0-9,20,25-30'.", + help="Diffusion attention KV cache dtype (e.g. fp8). Separate from vLLM --kv-cache-dtype.", ) omni_config_group.add_argument( - "--kv-cache-skip-layers", + "--diffusion-kv-cache-skip-steps", type=str, default=None, - help="Config-level KV-cache quantization skip-layer selector, e.g. '0,1,4-8'.", + help="Diffusion KV-cache quantization skip-step selector, e.g. '0-9,20,25-30'.", + ) + omni_config_group.add_argument( + "--diffusion-kv-cache-skip-layers", + type=str, + default=None, + help="Diffusion KV-cache quantization skip-layer selector, e.g. '0,1,4-8'.", ) omni_config_group.add_argument( "--cfg-parallel-size",