From 109bba8cdf322da5664a6b96bba8f74fb22f92b4 Mon Sep 17 00:00:00 2001 From: Lancer Date: Wed, 20 May 2026 13:18:57 +0800 Subject: [PATCH] [Bugfix] Treat kv_cache_dtype=auto as unset for ring attention (#3622) Signed-off-by: Lancer --- .../attention/test_attention_config.py | 99 +++++++++++++++++++ vllm_omni/diffusion/attention/layer.py | 12 ++- 2 files changed, 107 insertions(+), 4 deletions(-) diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index 5dea43b91..a4c67851a 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/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, + ) diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 5cbc39eb1..e0fb8d264 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/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