[Bugfix] Treat kv_cache_dtype=auto as unset for ring attention (#3622)

Signed-off-by: Lancer <maruixiang6688@gmail.com>

tests/diffusion/attention/test_attention_config.py

@@ -420,3 +420,102 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
         assert attn.use_ring is True
         assert attn.ring_runner is not None
         assert attn.ring_runner.attn_backend_pref == "TORCH_SDPA"
+
+
+class TestDiffusionKvCacheQuantization:
+    @staticmethod
+    def _install_attention_init_stubs(monkeypatch):
+        class _FakeAttentionImpl:
+            def __init__(self, **kwargs):
+                self.kwargs = kwargs
+
+            def forward(self, query, key, value, attn_metadata=None):
+                return query
+
+        class _FakeBackend:
+            @staticmethod
+            def get_name() -> str:
+                return "FAKE_BACKEND"
+
+            @staticmethod
+            def get_impl_cls():
+                return _FakeAttentionImpl
+
+            @staticmethod
+            def supports_kv_cache_dtype(kv_cache_dtype, platform_key) -> bool:
+                return True
+
+        class _FakeRingParallelAttention:
+            def __init__(self, sp_group, attn_backend_pref=None):
+                self.sp_group = sp_group
+                self.attn_backend_pref = attn_backend_pref
+
+        monkeypatch.setattr(
+            layer_mod,
+            "get_attn_backend_for_role",
+            lambda role, head_size, attention_config=None, role_category=None: (_FakeBackend, None),
+        )
+        monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
+        monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())
+        monkeypatch.setattr(layer_mod, "get_sp_group", lambda: SimpleNamespace(ring_group="ring-group"))
+        monkeypatch.setattr(layer_mod, "RingParallelAttention", _FakeRingParallelAttention)
+        monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
+
+    def test_diffusion_kv_cache_dtype_none_does_not_trigger_ring_quantization_error(self, monkeypatch):
+        self._install_attention_init_stubs(monkeypatch)
+        od_config = SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(),
+            parallel_config=SimpleNamespace(ring_degree=2),
+            diffusion_kv_cache_dtype=None,
+            diffusion_kv_cache_skip_step_indices=None,
+            diffusion_kv_cache_skip_layer_indices=None,
+        )
+
+        with set_current_diffusion_config(od_config):
+            attn = Attention(
+                num_heads=4,
+                head_size=64,
+                causal=False,
+                softmax_scale=1.0,
+            )
+
+        assert attn._kv_cache_dtype is None
+
+    def test_diffusion_kv_cache_dtype_auto_does_not_trigger_ring_quantization_error(self, monkeypatch):
+        self._install_attention_init_stubs(monkeypatch)
+        od_config = SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(),
+            parallel_config=SimpleNamespace(ring_degree=2),
+            diffusion_kv_cache_dtype="auto",
+            diffusion_kv_cache_skip_step_indices=None,
+            diffusion_kv_cache_skip_layer_indices=None,
+        )
+
+        with set_current_diffusion_config(od_config):
+            attn = Attention(
+                num_heads=4,
+                head_size=64,
+                causal=False,
+                softmax_scale=1.0,
+            )
+
+        assert attn._kv_cache_dtype is None
+
+    def test_diffusion_kv_cache_dtype_fp8_raises_with_ring_attention(self, monkeypatch):
+        self._install_attention_init_stubs(monkeypatch)
+        od_config = SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(),
+            parallel_config=SimpleNamespace(ring_degree=2),
+            diffusion_kv_cache_dtype="fp8",
+            diffusion_kv_cache_skip_step_indices=None,
+            diffusion_kv_cache_skip_layer_indices=None,
+        )
+
+        with set_current_diffusion_config(od_config):
+            with pytest.raises(ValueError, match="KV quantization is not compatible with ring attention"):
+                Attention(
+                    num_heads=4,
+                    head_size=64,
+                    causal=False,
+                    softmax_scale=1.0,
+                )

vllm_omni/diffusion/attention/layer.py

@@ -171,9 +171,13 @@ class Attention(nn.Module):
     def _init_kv_cache_quantization(self, config) -> None:
         if config is None:
             return
-        dtype = config.diffusion_kv_cache_dtype
+        dtype = getattr(config, "diffusion_kv_cache_dtype", None)
+        if dtype == "auto":
+            dtype = None
+        parallel_config = getattr(config, "parallel_config", None)
+        ring_degree = getattr(parallel_config, "ring_degree", 1)
         if dtype:
-            if config.parallel_config.ring_degree > 1:
+            if ring_degree > 1:
                 raise ValueError(
                     "KV quantization is not compatible with ring attention "
                     "(ring_degree > 1). Ring kernels do not propagate quantization descale "
@@ -190,8 +194,8 @@ class Attention(nn.Module):
                 )
                 dtype = None
         self._kv_cache_dtype = dtype
-        self._kv_cache_skip_steps = config.diffusion_kv_cache_skip_step_indices
-        self._kv_cache_skip_layers = config.diffusion_kv_cache_skip_layer_indices
+        self._kv_cache_skip_steps = getattr(config, "diffusion_kv_cache_skip_step_indices", None)
+        self._kv_cache_skip_layers = getattr(config, "diffusion_kv_cache_skip_layer_indices", None)
 
     def _should_apply_kv_cache_quant(self) -> bool:
         skip_steps = self._kv_cache_skip_steps