diff --git a/tests/attention/test_fish_kvcache_attn.py b/tests/attention/test_fish_kvcache_attn.py index 89b43e627..c82fd1858 100644 --- a/tests/attention/test_fish_kvcache_attn.py +++ b/tests/attention/test_fish_kvcache_attn.py @@ -92,6 +92,18 @@ def _metadata( )() +def _decode_kv_cache( + *, + num_blocks: int = 8, + block_size: int = 16, + num_kv_heads: int = 8, + head_size: int = 128, + dtype: torch.dtype = torch.float16, +) -> torch.Tensor: + # vLLM >=0.23.0 KV cache layout: key/value live on dim 1. + return torch.zeros((num_blocks, 2, block_size, num_kv_heads, head_size), dtype=dtype) + + def test_fish_kvcache_enabled_by_default(monkeypatch): monkeypatch.delenv("VLLM_OMNI_FISH_KVCACHE_ATTN", raising=False) @@ -182,7 +194,7 @@ def test_fish_kvcache_backend_wraps_only_model_instance(monkeypatch): metadata = _metadata() query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() result = model.layers[0].self_attn.attn.impl.forward( None, @@ -201,6 +213,110 @@ def test_fish_kvcache_backend_wraps_only_model_instance(monkeypatch): assert fish_kvcache_backend.get_fish_kvcache_attn_stats()["small_hit_count"] == 1 +def _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size): + # Source of truth for the paged KV cache layout the backend must unpack. + from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend + + return FlashAttentionBackend.get_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size) + + +def test_fish_kvcache_backend_unbinds_kv_on_vllm_cache_layout(monkeypatch): + # vLLM >=0.23.0 lays the KV cache out as (num_blocks, 2, block_size, + # num_kv_heads, head_size). The backend must unbind on dim 1 so each of + # key/value comes back as the 4-D (num_blocks, block_size, num_kv_heads, + # head_size) tensor the decode kernel expects. + monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1") + monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True) + monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None) + monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", lambda **_: True) + + num_blocks, block_size, num_kv_heads, head_size = 8, 16, 8, 128 + cache_shape = _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size) + assert cache_shape == (num_blocks, 2, block_size, num_kv_heads, head_size) + + captured = {} + + def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len): + del query, block_table, seq_lens, scale, max_seq_len + captured["key"] = tuple(key_cache.shape) + captured["value"] = tuple(value_cache.shape) + out.fill_(1) + return out + + monkeypatch.setattr(fish_kvcache_backend, "fish_decode_kvcache_attn", fake_decode) + + model = _FakeModel() + assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1 + + query = torch.zeros((2, 32, head_size), dtype=torch.float16) + output = torch.zeros_like(query) + kv_cache = torch.zeros(cache_shape, dtype=torch.float16) + + model.layers[0].self_attn.attn.impl.forward(None, query, None, None, kv_cache, _metadata(), output) + + expected = (num_blocks, block_size, num_kv_heads, head_size) + assert captured["key"] == expected + assert captured["value"] == expected + + +def test_fish_kvcache_vllm_cache_unbinds_then_falls_back_until_contiguous(monkeypatch): + # End-to-end with the real vLLM cache layout and the real (unmocked) guard. + # + # The point of the unbind(1) fix: unbind(0) raised "too many values to + # unpack" on the 5-D vLLM 0.23.0 layout, so the request crashed outright. + # unbind(1) makes the request *work* again -- it unpacks cleanly and the + # backend gracefully falls back to the original attention with correct + # output. This test guards that fix and must NOT depend on the fast path + # firing. + # + # It also documents a known limitation (tracked separately): unbind(1) on + # the interleaved (num_blocks, 2, ...) layout yields *non-contiguous* + # key/value views, which the guard rejects on the is_contiguous() check. + # So the dimensions are compatible (4-D, block_size=16, head_size=128) -- + # the only thing keeping the fast path from engaging is contiguity, not + # shape. Making the fast path actually fire would require a .contiguous() + # copy and is out of scope here. + fish_kvcache_backend.reset_fish_kvcache_attn_stats() + monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1") + monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True) + monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None) + monkeypatch.setattr(fish_kvcache_attn, "is_available", lambda: True) + + num_blocks, block_size, num_kv_heads, head_size = 8, 16, 8, 128 + cache_shape = _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size) + + decode_calls = [] + + def fake_decode(*args, **kwargs): + del args, kwargs + decode_calls.append(True) + + monkeypatch.setattr(fish_kvcache_backend, "fish_decode_kvcache_attn", fake_decode) + + model = _FakeModel() + assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1 + + query = torch.zeros((2, 32, head_size), dtype=torch.float16) + output = torch.zeros_like(query) + kv_cache = torch.zeros(cache_shape, dtype=torch.float16) + + # Root cause record: unbind(1) gives the right 4-D shape but a non-contiguous view. + key_cache, value_cache = kv_cache.unbind(1) + assert tuple(key_cache.shape) == (num_blocks, block_size, num_kv_heads, head_size) + assert not key_cache.is_contiguous() + + # unbind(1) no longer crashes (unbind(0) would have raised here) and the + # request completes via the correct fallback path. + result = model.layers[0].self_attn.attn.impl.forward(None, query, None, None, kv_cache, _metadata(), output) + + assert result is output + assert decode_calls == [] + assert model.layers[0].self_attn.attn.impl.original_calls == 1 + assert torch.equal(output, torch.full_like(output, 2)) + stats = fish_kvcache_backend.get_fish_kvcache_attn_stats() + assert stats["fallback_count_by_reason"] == {"guard_miss": 1} + + def test_fish_kvcache_backend_install_is_idempotent(monkeypatch): monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1") monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True) @@ -215,7 +331,7 @@ def test_fish_kvcache_backend_install_is_idempotent(monkeypatch): query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() model.layers[0].self_attn.attn.impl.forward( None, @@ -250,7 +366,7 @@ def test_fish_kvcache_backend_falls_back_to_original_forward_on_guard_miss(monke query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() result = model.layers[0].self_attn.attn.impl.forward( None, @@ -297,7 +413,7 @@ def test_fish_kvcache_backend_uses_decode_seq_lens_upper_bound(monkeypatch): query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() model.layers[0].self_attn.attn.impl.forward( None, @@ -340,7 +456,7 @@ def test_fish_kvcache_backend_uses_cpu_upper_bound_above_metadata_max(monkeypatc query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() model.layers[0].self_attn.attn.impl.forward( None, @@ -383,7 +499,7 @@ def test_fish_kvcache_backend_slices_upper_bound_by_active_batch(monkeypatch): query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() model.layers[0].self_attn.attn.impl.forward( None, @@ -428,7 +544,7 @@ def test_fish_kvcache_backend_falls_back_on_short_upper_bound_shape(monkeypatch) query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() result = model.layers[0].self_attn.attn.impl.forward( None, @@ -475,7 +591,7 @@ def test_fish_kvcache_backend_falls_back_on_non_cpu_upper_bound(monkeypatch): query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() result = model.layers[0].self_attn.attn.impl.forward( None, @@ -509,7 +625,7 @@ def test_fish_kvcache_backend_required_rejects_non_cpu_upper_bound(monkeypatch): query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() with pytest.raises(RuntimeError, match="must be a CPU tensor"): model.layers[0].self_attn.attn.impl.forward( @@ -539,7 +655,7 @@ def test_fish_kvcache_backend_required_rejects_underestimated_upper_bound(monkey query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() with pytest.raises(RuntimeError, match="underestimates"): model.layers[0].self_attn.attn.impl.forward( @@ -569,7 +685,7 @@ def test_fish_kvcache_backend_required_rejects_missing_upper_bound(monkeypatch): query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() with pytest.raises(RuntimeError, match="requires seq_lens_cpu_upper_bound"): model.layers[0].self_attn.attn.impl.forward( @@ -603,7 +719,7 @@ def test_fish_kvcache_backend_required_raises_on_guard_miss(monkeypatch): query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() with pytest.raises(RuntimeError, match="required"): model.layers[0].self_attn.attn.impl.forward( @@ -629,7 +745,7 @@ def test_fish_kvcache_backend_required_allows_prefill_fallback(monkeypatch): query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() result = model.layers[0].self_attn.attn.impl.forward( None, @@ -675,7 +791,7 @@ def test_fish_kvcache_backend_falls_back_without_cpu_upper_bound(monkeypatch): query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() result = model.layers[0].self_attn.attn.impl.forward( None, diff --git a/vllm_omni/attention/fish_kvcache_backend.py b/vllm_omni/attention/fish_kvcache_backend.py index 3f3e9bdc6..0178f722e 100644 --- a/vllm_omni/attention/fish_kvcache_backend.py +++ b/vllm_omni/attention/fish_kvcache_backend.py @@ -233,7 +233,12 @@ def _forward_with_fish_kvcache( if attn_metadata is not None and not attn_metadata.use_cascade: num_actual_tokens = attn_metadata.num_actual_tokens - key_cache, value_cache = kv_cache.unbind(0) + # vLLM >=0.23.0 lays out the KV cache as + # (num_blocks, 2, block_size, num_kv_heads, head_size), so key/value + # live on dim 1 (matches FlashAttentionImpl's kv_cache.unbind(1)). + # The pre-0.23 layout had them on dim 0; unbind(0) here would split + # along num_blocks and raise "too many values to unpack". + key_cache, value_cache = kv_cache.unbind(1) q = query[:num_actual_tokens] out = output[:num_actual_tokens] active_batch = int(q.shape[0])