[BugFix] fix(omni): isolate diffusion KV-cache dtype from vLLM --kv-cache-dtype #3585 (#3596)

Signed-off-by: lyj-jjj <liuyingjun5@huawei.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

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 <your-model> --omni --kv-cache-dtype fp8
+vllm serve <your-model> --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

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}"

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

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)

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)
 

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",