[Bugfix] Fix Fish Speech kvcache attention KV cache unpacking for vLLM 0.23.0 (#4428)

Signed-off-by: nagelanping <xiao.yi.ming@foxmail.com>

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,

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])