anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
0ae9bdd8c2eec90659b39bc73072f34b0d0e5fae
vllm-omni/tests/attention/test_fish_kvcache_attn.py
T
tzhouamandClaude Opus 5 c1e32fc01d rebase: close the remaining vLLM 0.29 drift found by CI 3009 
Five further removals/renames, each traced from a failing 3009 job:

- StageEngineCoreProcManager mirrors the parent __init__ rather than calling
  super(), and 0.29 added _request_shutdown_timeout there. The inherited
  shutdown() read it, raised AttributeError, and left the engine core
  subprocesses running, so the test process hung at interpreter exit until
  the 15 minute job timeout. That is the whole of the Engine&Model Executor
  exit-124: its pytest reported 567 passed, then sat idle for 663s. The
  failure surfaced only as a StagePool warning, never as a test failure.

- PromptUpdateDetails.select_text is gone; prompt updates are token oriented
  now. MiniCPM-o 4.5 kept three text-based image/video/audio replacements,
  which crashed stage-0 startup and account for the whole MiniCPM-o cluster.
  They encode the text and select the unk placeholder by token id instead.

- Qwen3-Omni: 0.29 marks audio_feature_lengths keep_on_cpu, while the audio
  tower derives device placement from feature_lens, so the mask comparison
  mixed cuda and cpu tensors. Moved explicitly, matching upstream
  Qwen3OmniMoeThinker._process_audio_input.

- AttentionBackend.get_kv_cache_shape is gone; the per-layer shape comes from
  the spec via compute_layer_kv_cache_shape_bytes, whose content dim is in
  bytes.

- WeightsMapper.get_unstacked_mapper is renamed to get_rename_mapper.

tests/engine/test_async_omni_engine_abort.py::test_abort fails here both with
and without these changes, so it is environmental rather than a regression.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0146Xu5Sgj4vXxLxvLjV5YPj
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
2026-09-10 00:41:53 +08:00

1268 lines
45 KiB
Python
Raw Blame History

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import pytest
import torch
from vllm_omni.attention import fish_kvcache_attn, fish_kvcache_backend
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class _FakeImpl:
dcp_world_size = 1
alibi_slopes = None
sliding_window = None
scale = 1.0
def __init__(self):
self.original_calls = 0
def forward(
self,
layer,
query,
key,
value,
kv_cache,
attn_metadata,
output,
output_scale=None,
output_block_scale=None,
):
del layer, query, key, value, kv_cache, attn_metadata, output_scale, output_block_scale
self.original_calls += 1
output.fill_(2)
return output
class _FakeAttentionLayer:
def __init__(self):
self.impl = _FakeImpl()
class _FakeSelfAttn:
def __init__(self):
self.attn = _FakeAttentionLayer()
class _FakeLayer:
def __init__(self):
self.self_attn = _FakeSelfAttn()
class _FakeModel:
def __init__(self):
self.layers = [_FakeLayer()]
_DEFAULT_UPPER_BOUND = object()
def _metadata(
*,
batch_size: int = 2,
num_actual_tokens: int | None = None,
max_query_len: int = 1,
max_seq_len: int = 128,
use_cascade: bool = False,
seq_lens_cpu_upper_bound: torch.Tensor | None | object = _DEFAULT_UPPER_BOUND,
seq_lens: torch.Tensor | None = None,
):
if num_actual_tokens is None:
num_actual_tokens = batch_size
if seq_lens is None:
seq_lens = torch.ones((batch_size,), dtype=torch.int32)
if seq_lens_cpu_upper_bound is _DEFAULT_UPPER_BOUND and max_query_len == 1:
seq_lens_cpu_upper_bound = seq_lens.detach().cpu()
elif seq_lens_cpu_upper_bound is _DEFAULT_UPPER_BOUND:
seq_lens_cpu_upper_bound = None
return type(
"Metadata",
(),
{
"use_cascade": use_cascade,
"num_actual_tokens": num_actual_tokens,
"block_table": torch.zeros((batch_size, 8), dtype=torch.int32),
"seq_lens": seq_lens,
"seq_lens_cpu_upper_bound": seq_lens_cpu_upper_bound,
"max_query_len": max_query_len,
"max_seq_len": max_seq_len,
},
)()
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: (num_blocks, num_kv_heads, block_size,
# 2*head_size) with K and V packed in the last dimension.
return torch.zeros((num_blocks, num_kv_heads, block_size, 2 * head_size), dtype=dtype)
def test_fish_kvcache_enabled_by_default(monkeypatch):
monkeypatch.delenv("VLLM_OMNI_FISH_KVCACHE_ATTN", raising=False)
assert fish_kvcache_attn.is_fish_kvcache_attn_enabled()
assert fish_kvcache_backend._fish_kvcache_enabled()
assert not fish_kvcache_attn.is_fish_kvcache_attn_required()
@pytest.mark.parametrize("value", ["", "1", "true", "yes", "on", "required"])
def test_fish_kvcache_enabled_values_are_consistent(monkeypatch, value):
monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", value)
assert fish_kvcache_attn.is_fish_kvcache_attn_enabled()
assert fish_kvcache_backend._fish_kvcache_enabled()
assert fish_kvcache_attn.is_fish_kvcache_attn_required() is (value == "required")
@pytest.mark.parametrize("value", ["0", "false", "no", "off", "disabled", "disable"])
def test_fish_kvcache_disabled_values_are_consistent(monkeypatch, value):
monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", value)
assert not fish_kvcache_attn.is_fish_kvcache_attn_enabled()
assert not fish_kvcache_backend._fish_kvcache_enabled()
assert not fish_kvcache_attn.is_fish_kvcache_attn_required()
@pytest.mark.parametrize("sliding_window", [None, (-1, -1)])
def test_fish_kvcache_attn_guard_accepts_decode_only_shape(monkeypatch, sliding_window):
monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
monkeypatch.setattr(fish_kvcache_attn, "is_available", lambda: True)
assert fish_kvcache_attn.can_use_fish_kvcache_attn(
query=torch.empty((4, 32, 128), dtype=torch.float16),
key_cache=torch.empty((8, 16, 8, 128), dtype=torch.float16),
value_cache=torch.empty((8, 16, 8, 128), dtype=torch.float16),
block_table=torch.zeros((4, 8), dtype=torch.int32),
seq_lens=torch.ones((4,), dtype=torch.int32),
max_query_len=1,
max_seq_len=128,
dcp_world_size=1,
use_cascade=False,
alibi_slopes=None,
sliding_window=sliding_window,
)
def test_fish_kvcache_attn_guard_rejects_unsupported_block_size(monkeypatch):
monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
monkeypatch.setattr(fish_kvcache_attn, "is_available", lambda: True)
assert not fish_kvcache_attn.can_use_fish_kvcache_attn(
query=torch.empty((4, 32, 128), dtype=torch.float16),
key_cache=torch.empty((8, 32, 8, 128), dtype=torch.float16),
value_cache=torch.empty((8, 32, 8, 128), dtype=torch.float16),
block_table=torch.zeros((4, 8), dtype=torch.int32),
seq_lens=torch.ones((4,), dtype=torch.int32),
max_query_len=1,
max_seq_len=128,
dcp_world_size=1,
use_cascade=False,
alibi_slopes=None,
sliding_window=None,
)
def test_fish_kvcache_backend_wraps_only_model_instance(monkeypatch):
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_backend, "can_use_fish_kvcache_attn", lambda **_: True)
hits = []
def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len):
del key_cache, value_cache, block_table, seq_lens, scale, max_seq_len
hits.append(tuple(query.shape))
out.fill_(1)
return out
monkeypatch.setattr(fish_kvcache_backend, "fish_decode_kvcache_attn", fake_decode)
model = _FakeModel()
other_impl = _FakeImpl()
assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
assert not hasattr(other_impl, "_fish_kvcache_attn_installed")
metadata = _metadata()
query = torch.zeros((2, 32, 128), dtype=torch.float16)
output = torch.zeros_like(query)
kv_cache = _decode_kv_cache()
result = model.layers[0].self_attn.attn.impl.forward(
None,
query,
None,
None,
kv_cache,
metadata,
output,
)
assert result is output
assert hits == [(2, 32, 128)]
assert model.layers[0].self_attn.attn.impl.original_calls == 0
assert torch.equal(output, torch.ones_like(output))
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, dtype=torch.float16):
# Source of truth for the paged KV cache layout the backend must unpack.
# vLLM 0.29 removed AttentionBackend.get_kv_cache_shape; the per-layer shape
# now comes from the spec, with the content dim expressed in bytes.
from vllm.v1.kv_cache_interface import FullAttentionSpec, compute_layer_kv_cache_shape_bytes
spec = FullAttentionSpec(
block_size=block_size,
num_kv_heads=num_kv_heads,
head_size=head_size,
dtype=dtype,
)
num_pages, num_heads, num_states, content_bytes = compute_layer_kv_cache_shape_bytes(spec, num_blocks)
return (num_pages, num_heads, num_states, content_bytes // dtype.itemsize)
def test_fish_kvcache_backend_unbinds_kv_on_vllm_cache_layout(monkeypatch):
# vLLM >=0.23.0 lays out the KV cache as (num_blocks, num_kv_heads,
# block_size, 2*head_size) with K and V packed in the last dimension.
# The backend must split on the content dim 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)
# Upstream now returns a 4-D packed shape: (B, H, N, 2*D).
assert cache_shape == (num_blocks, num_kv_heads, block_size, 2 * 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 transpose+split fix: unbind(1) on the old 5-D layout
# raised "too many values to unpack" on the vLLM 0.23.0 layout, and the
# upstream now packs K/V into the last dimension (4-D shape). The
# transpose(1, 2) + split correctly unpacks into separate key/value tensors
# 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):
# transpose(1, 2) on the interleaved (B, H, N, 2*D) 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: transpose(1,2) + split gives the right 4-D shape but
# a non-contiguous view.
key_cache, value_cache = kv_cache.transpose(1, 2).split(head_size, dim=-1)
assert tuple(key_cache.shape) == (num_blocks, block_size, num_kv_heads, head_size)
assert not key_cache.is_contiguous()
# transpose+split no longer crashes (unbind(1) 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)
monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", lambda **_: False)
model = _FakeModel()
assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
wrapped_forward = model.layers[0].self_attn.attn.impl.forward
assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 0
assert model.layers[0].self_attn.attn.impl.forward is wrapped_forward
query = torch.zeros((2, 32, 128), dtype=torch.float16)
output = torch.zeros_like(query)
kv_cache = _decode_kv_cache()
model.layers[0].self_attn.attn.impl.forward(
None,
query,
None,
None,
kv_cache,
_metadata(),
output,
)
assert model.layers[0].self_attn.attn.impl.original_calls == 1
def test_fish_kvcache_backend_falls_back_to_original_forward_on_guard_miss(monkeypatch):
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_backend, "can_use_fish_kvcache_attn", lambda **_: False)
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, 128), dtype=torch.float16)
output = torch.zeros_like(query)
kv_cache = _decode_kv_cache()
result = model.layers[0].self_attn.attn.impl.forward(
None,
query,
None,
None,
kv_cache,
_metadata(),
output,
)
assert result is output
assert model.layers[0].self_attn.attn.impl.original_calls == 1
assert decode_calls == []
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_uses_decode_seq_lens_upper_bound(monkeypatch):
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)
guard_max_seq_lens = []
def fake_can_use(**kwargs):
guard_max_seq_lens.append(kwargs["max_seq_len"])
return True
decode_max_seq_lens = []
def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len):
del query, key_cache, value_cache, block_table, seq_lens, scale
decode_max_seq_lens.append(max_seq_len)
out.fill_(1)
return out
monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", fake_can_use)
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, 128), dtype=torch.float16)
output = torch.zeros_like(query)
kv_cache = _decode_kv_cache()
model.layers[0].self_attn.attn.impl.forward(
None,
query,
None,
None,
kv_cache,
_metadata(
max_query_len=1,
max_seq_len=16384,
seq_lens_cpu_upper_bound=torch.tensor([512, 768], dtype=torch.int32),
),
output,
)
assert guard_max_seq_lens == [1024]
assert decode_max_seq_lens == [1024]
def test_fish_kvcache_backend_uses_cpu_upper_bound_above_metadata_max(monkeypatch):
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)
guard_max_seq_lens = []
def fake_can_use(**kwargs):
guard_max_seq_lens.append(kwargs["max_seq_len"])
return True
monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", fake_can_use)
monkeypatch.setattr(
fish_kvcache_backend,
"fish_decode_kvcache_attn",
lambda query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len: out.fill_(1),
)
model = _FakeModel()
assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
query = torch.zeros((2, 32, 128), dtype=torch.float16)
output = torch.zeros_like(query)
kv_cache = _decode_kv_cache()
model.layers[0].self_attn.attn.impl.forward(
None,
query,
None,
None,
kv_cache,
_metadata(
max_query_len=1,
max_seq_len=64,
seq_lens_cpu_upper_bound=torch.tensor([1024, 1024], dtype=torch.int32),
),
output,
)
assert guard_max_seq_lens == [1024]
def test_fish_kvcache_backend_slices_upper_bound_by_active_batch(monkeypatch):
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)
guard_max_seq_lens = []
def fake_can_use(**kwargs):
guard_max_seq_lens.append(kwargs["max_seq_len"])
return True
def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len):
del query, key_cache, value_cache, block_table, seq_lens, scale, max_seq_len
out.fill_(1)
return out
monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", fake_can_use)
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, 128), dtype=torch.float16)
output = torch.zeros_like(query)
kv_cache = _decode_kv_cache()
model.layers[0].self_attn.attn.impl.forward(
None,
query,
None,
None,
kv_cache,
_metadata(
batch_size=4,
num_actual_tokens=2,
max_query_len=1,
max_seq_len=16384,
seq_lens_cpu_upper_bound=torch.tensor([512, 768, 4096, 4096], dtype=torch.int32),
),
output,
)
assert guard_max_seq_lens == [1024]
def test_fish_kvcache_backend_falls_back_on_short_upper_bound_shape(monkeypatch):
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)
guard_max_seq_lens = []
def fake_can_use(**kwargs):
guard_max_seq_lens.append(kwargs["max_seq_len"])
return True
def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len):
del query, key_cache, value_cache, block_table, seq_lens, scale, max_seq_len
out.fill_(1)
return out
monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", fake_can_use)
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, 128), dtype=torch.float16)
output = torch.zeros_like(query)
kv_cache = _decode_kv_cache()
result = model.layers[0].self_attn.attn.impl.forward(
None,
query,
None,
None,
kv_cache,
_metadata(
batch_size=2,
max_query_len=1,
max_seq_len=16384,
seq_lens_cpu_upper_bound=torch.tensor([512], dtype=torch.int32),
),
output,
)
assert result is output
assert guard_max_seq_lens == []
assert model.layers[0].self_attn.attn.impl.original_calls == 1
assert torch.equal(output, torch.full_like(output, 2))
def test_fish_kvcache_backend_falls_back_on_non_cpu_upper_bound(monkeypatch):
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)
guard_max_seq_lens = []
def fake_can_use(**kwargs):
guard_max_seq_lens.append(kwargs["max_seq_len"])
return True
def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len):
del query, key_cache, value_cache, block_table, seq_lens, scale, max_seq_len
out.fill_(1)
return out
monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", fake_can_use)
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, 128), dtype=torch.float16)
output = torch.zeros_like(query)
kv_cache = _decode_kv_cache()
result = model.layers[0].self_attn.attn.impl.forward(
None,
query,
None,
None,
kv_cache,
_metadata(
batch_size=2,
max_query_len=1,
max_seq_len=16384,
seq_lens_cpu_upper_bound=torch.empty((2,), device="meta", dtype=torch.int32),
),
output,
)
assert result is output
assert guard_max_seq_lens == []
assert model.layers[0].self_attn.attn.impl.original_calls == 1
assert torch.equal(output, torch.full_like(output, 2))
def test_fish_kvcache_backend_required_rejects_non_cpu_upper_bound(monkeypatch):
monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "required")
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)
model = _FakeModel()
assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
query = torch.zeros((2, 32, 128), dtype=torch.float16)
output = torch.zeros_like(query)
kv_cache = _decode_kv_cache()
with pytest.raises(RuntimeError, match="must be a CPU tensor"):
model.layers[0].self_attn.attn.impl.forward(
None,
query,
None,
None,
kv_cache,
_metadata(
batch_size=2,
max_query_len=1,
max_seq_len=16384,
seq_lens_cpu_upper_bound=torch.empty((2,), device="meta", dtype=torch.int32),
),
output,
)
def test_fish_kvcache_backend_required_rejects_underestimated_upper_bound(monkeypatch):
monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "required")
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)
model = _FakeModel()
assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
query = torch.zeros((2, 32, 128), dtype=torch.float16)
output = torch.zeros_like(query)
kv_cache = _decode_kv_cache()
with pytest.raises(RuntimeError, match="underestimates"):
model.layers[0].self_attn.attn.impl.forward(
None,
query,
None,
None,
kv_cache,
_metadata(
max_query_len=1,
max_seq_len=16384,
seq_lens=torch.tensor([512, 2048], dtype=torch.int32),
seq_lens_cpu_upper_bound=torch.tensor([512, 768], dtype=torch.int32),
),
output,
)
def test_fish_kvcache_backend_required_rejects_missing_upper_bound(monkeypatch):
monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "required")
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)
model = _FakeModel()
assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
query = torch.zeros((2, 32, 128), dtype=torch.float16)
output = torch.zeros_like(query)
kv_cache = _decode_kv_cache()
with pytest.raises(RuntimeError, match="requires seq_lens_cpu_upper_bound"):
model.layers[0].self_attn.attn.impl.forward(
None,
query,
None,
None,
kv_cache,
_metadata(max_query_len=1, max_seq_len=16384, seq_lens_cpu_upper_bound=None),
output,
)
def test_fish_kvcache_backend_required_rejects_zero_installed_layers(monkeypatch):
monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "required")
monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
with pytest.raises(RuntimeError, match="installed 0 attention layers"):
fish_kvcache_backend.install_fish_kvcache_attn_backend(type("EmptyModel", (), {"layers": []})())
def test_fish_kvcache_backend_required_raises_on_guard_miss(monkeypatch):
monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "required")
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 **_: False)
model = _FakeModel()
assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
query = torch.zeros((2, 32, 128), dtype=torch.float16)
output = torch.zeros_like(query)
kv_cache = _decode_kv_cache()
with pytest.raises(RuntimeError, match="required"):
model.layers[0].self_attn.attn.impl.forward(
None,
query,
None,
None,
kv_cache,
_metadata(),
output,
)
def test_fish_kvcache_backend_required_allows_prefill_fallback(monkeypatch):
fish_kvcache_backend.reset_fish_kvcache_attn_stats()
monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "required")
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 **_: False)
model = _FakeModel()
assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1
query = torch.zeros((2, 32, 128), dtype=torch.float16)
output = torch.zeros_like(query)
kv_cache = _decode_kv_cache()
result = model.layers[0].self_attn.attn.impl.forward(
None,
query,
None,
None,
kv_cache,
_metadata(max_query_len=2, max_seq_len=2),
output,
)
assert result is output
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"] == {"non_decode": 1}
def test_fish_kvcache_backend_falls_back_without_cpu_upper_bound(monkeypatch):
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)
guard_max_seq_lens = []
def fake_can_use(**kwargs):
guard_max_seq_lens.append(kwargs["max_seq_len"])
return True
decode_max_seq_lens = []
def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len):
del query, key_cache, value_cache, block_table, seq_lens, scale
decode_max_seq_lens.append(max_seq_len)
out.fill_(1)
return out
monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", fake_can_use)
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, 128), dtype=torch.float16)
output = torch.zeros_like(query)
kv_cache = _decode_kv_cache()
result = model.layers[0].self_attn.attn.impl.forward(
None,
query,
None,
None,
kv_cache,
_metadata(max_query_len=1, max_seq_len=16384, seq_lens_cpu_upper_bound=None),
output,
)
assert result is output
assert guard_max_seq_lens == []
assert decode_max_seq_lens == []
assert model.layers[0].self_attn.attn.impl.original_calls == 1
assert torch.equal(output, torch.full_like(output, 2))
def test_fish_kvcache_attn_guard_accepts_decode_metadata_above_small_cap(monkeypatch):
monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
monkeypatch.setattr(fish_kvcache_attn, "is_available", lambda: True)
assert fish_kvcache_attn.can_use_fish_kvcache_attn(
query=torch.empty((4, 32, 128), dtype=torch.float16),
key_cache=torch.empty((1024, 16, 8, 128), dtype=torch.float16),
value_cache=torch.empty((1024, 16, 8, 128), dtype=torch.float16),
block_table=torch.zeros((4, 1024), dtype=torch.int32),
seq_lens=torch.ones((4,), dtype=torch.int32),
max_query_len=1,
max_seq_len=16384,
dcp_world_size=1,
use_cascade=False,
alibi_slopes=None,
sliding_window=None,
)
def test_fish_kvcache_attn_guard_rejects_seq_len_above_block_table_capacity(monkeypatch):
monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
monkeypatch.setattr(fish_kvcache_attn, "is_available", lambda: True)
assert not fish_kvcache_attn.can_use_fish_kvcache_attn(
query=torch.empty((4, 32, 128), dtype=torch.float16),
key_cache=torch.empty((8, 16, 8, 128), dtype=torch.float16),
value_cache=torch.empty((8, 16, 8, 128), dtype=torch.float16),
block_table=torch.zeros((4, 8), dtype=torch.int32),
seq_lens=torch.ones((4,), dtype=torch.int32),
max_query_len=1,
max_seq_len=129,
dcp_world_size=1,
use_cascade=False,
alibi_slopes=None,
sliding_window=None,
)
def test_fish_kvcache_decode_reuses_long_workspace(monkeypatch):
fish_kvcache_attn._WORKSPACE_CACHE.clear()
calls = []
def fake_decode(
query,
key_cache,
value_cache,
block_table,
seq_lens,
out,
scale,
max_seq_len,
small_path_max_seq_len,
long_split_tokens,
partial_m,
partial_l,
partial_acc,
):
del query, key_cache, value_cache, block_table, seq_lens, scale, max_seq_len
del small_path_max_seq_len, long_split_tokens
calls.append((partial_m.data_ptr(), partial_l.data_ptr(), partial_acc.data_ptr(), tuple(partial_acc.shape)))
return out
monkeypatch.setattr(
fish_kvcache_attn._triton_backend(),
"fish_decode_kvcache_attn_triton",
fake_decode,
)
q = torch.zeros((2, 4, 128), dtype=torch.float16)
k_cache = torch.zeros((128, 16, 2, 128), dtype=torch.float16)
v_cache = torch.zeros_like(k_cache)
block_table = torch.zeros((2, 128), dtype=torch.int32)
seq_lens = torch.full((2,), 2048, dtype=torch.int32)
out = torch.empty_like(q)
for _ in range(2):
fish_kvcache_attn.fish_decode_kvcache_attn(
q,
k_cache,
v_cache,
block_table,
seq_lens,
out,
scale=128**-0.5,
max_seq_len=2048,
)
assert len(calls) == 2
assert calls[0] == calls[1]
assert calls[0][3] == (2, 2, 4, 128)
def test_fish_kvcache_decode_workspace_miss_rejects_cuda_graph_capture(monkeypatch):
fish_kvcache_attn._WORKSPACE_CACHE.clear()
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
q = torch.zeros((2, 4, 128), dtype=torch.float16)
with pytest.raises(RuntimeError, match="workspace was not prewarmed"):
fish_kvcache_attn._get_decode_workspace(q, 2048)
def test_fish_kvcache_backend_prewarms_capture_workspaces(monkeypatch):
fish_kvcache_attn._WORKSPACE_CACHE.clear()
monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
text_config = type(
"TextConfig",
(),
{"num_attention_heads": 4, "head_dim": 128},
)()
hf_config = type("HFConfig", (), {"text_config": text_config})()
model_config = type(
"ModelConfig",
(),
{
"hf_config": hf_config,
"max_model_len": 2048,
"model_arch": "FishSpeechSlowARForConditionalGeneration",
},
)()
assert (
fish_kvcache_backend.prewarm_fish_kvcache_attn_capture_workspaces(
model_config=model_config,
device=torch.device("cpu"),
dtype=torch.float16,
capture_sizes=[1, 4],
)
== 3
)
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
q = torch.zeros((4, 4, 128), dtype=torch.float16)
partial_m, partial_l, partial_acc = fish_kvcache_attn._get_decode_workspace(q, 2048)
assert tuple(partial_m.shape) == (2, 4, 4)
assert tuple(partial_l.shape) == (2, 4, 4)
assert tuple(partial_acc.shape) == (2, 4, 4, 128)
def _reference_decode_attention(
q: torch.Tensor,
k_cache: torch.Tensor,
v_cache: torch.Tensor,
block_table: torch.Tensor,
seq_lens: torch.Tensor,
scale: float,
) -> torch.Tensor:
block_size = int(k_cache.shape[1])
num_q_heads = int(q.shape[1])
num_kv_heads = int(k_cache.shape[2])
outputs = []
for batch_idx in range(q.shape[0]):
seq_len = int(seq_lens[batch_idx].item())
if seq_len <= 0:
outputs.append(torch.zeros_like(q[batch_idx], dtype=torch.float32))
continue
k_rows = []
v_rows = []
for logical_block in range((seq_len + block_size - 1) // block_size):
physical_block = int(block_table[batch_idx, logical_block].item())
start = logical_block * block_size
take = min(block_size, seq_len - start)
k_rows.append(k_cache[physical_block, :take])
v_rows.append(v_cache[physical_block, :take])
k = torch.cat(k_rows, dim=0).to(torch.float32)
v = torch.cat(v_rows, dim=0).to(torch.float32)
per_head = []
for q_head in range(num_q_heads):
kv_head = q_head // (num_q_heads // num_kv_heads)
scores = torch.matmul(k[:, kv_head, :], q[batch_idx, q_head].to(torch.float32)) * scale
weights = torch.softmax(scores, dim=0)
per_head.append(torch.matmul(weights, v[:, kv_head, :]))
outputs.append(torch.stack(per_head, dim=0))
return torch.stack(outputs, dim=0).to(q.dtype)
@pytest.mark.parametrize("seq_len", [64, 128, 256, 512, 1024, 1025, 2048])
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@pytest.mark.parametrize("batch_size", [1, 3])
@pytest.mark.skipif(
not torch.cuda.is_available() or not fish_kvcache_attn.is_available(),
reason="Fish kvcache Triton attention is not available",
)
def test_fish_kvcache_native_op_matches_reference(seq_len, dtype, batch_size):
torch.manual_seed(0)
device = torch.device("cuda")
block_size = 16
num_q_heads = 4
num_kv_heads = 2
head_dim = 128
max_blocks = (seq_len + block_size - 1) // block_size
num_blocks = batch_size * max_blocks + 3
scale = head_dim**-0.5
q = torch.randn((batch_size, num_q_heads, head_dim), device=device, dtype=dtype)
k_cache = torch.randn((num_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
v_cache = torch.randn((num_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
block_table = torch.empty((batch_size, max_blocks), device=device, dtype=torch.int32)
for batch_idx in range(batch_size):
pages = torch.arange(batch_idx * max_blocks, (batch_idx + 1) * max_blocks, device=device, dtype=torch.int32)
block_table[batch_idx] = torch.flip(pages, dims=[0])
seq_lens = torch.full((batch_size,), seq_len, device=device, dtype=torch.int32)
out = torch.empty_like(q)
fish_kvcache_attn.fish_decode_kvcache_attn(
q,
k_cache,
v_cache,
block_table,
seq_lens,
out,
scale=scale,
max_seq_len=seq_len,
)
expected = _reference_decode_attention(q, k_cache, v_cache, block_table, seq_lens, scale)
torch.testing.assert_close(out, expected, atol=3e-2, rtol=3e-2)
@pytest.mark.skipif(
not torch.cuda.is_available() or not fish_kvcache_attn.is_available(),
reason="Fish kvcache Triton attention is not available",
)
def test_fish_kvcache_native_op_zero_seq_len_writes_zero():
device = torch.device("cuda")
dtype = torch.float16
block_size = 16
num_q_heads = 4
num_kv_heads = 2
head_dim = 128
max_seq_len = 512
max_blocks = max_seq_len // block_size
scale = head_dim**-0.5
q = torch.randn((1, num_q_heads, head_dim), device=device, dtype=dtype)
k_cache = torch.randn((max_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
v_cache = torch.randn_like(k_cache)
block_table = torch.arange(max_blocks, device=device, dtype=torch.int32).reshape(1, max_blocks)
seq_lens = torch.zeros((1,), device=device, dtype=torch.int32)
out = torch.full_like(q, float("nan"))
fish_kvcache_attn.fish_decode_kvcache_attn(
q,
k_cache,
v_cache,
block_table,
seq_lens,
out,
scale=scale,
max_seq_len=max_seq_len,
)
assert torch.isfinite(out).all()
assert torch.equal(out, torch.zeros_like(out))
@pytest.mark.skipif(
not torch.cuda.is_available() or not fish_kvcache_attn.is_available(),
reason="Fish kvcache Triton attention is not available",
)
def test_fish_kvcache_native_op_handles_mixed_short_and_long_seq_lens():
torch.manual_seed(0)
device = torch.device("cuda")
dtype = torch.float16
block_size = 16
num_q_heads = 4
num_kv_heads = 2
head_dim = 128
seq_lens = torch.tensor([512, 2048, 1025], device=device, dtype=torch.int32)
max_seq_len = int(seq_lens.max().item())
batch_size = int(seq_lens.numel())
max_blocks = (max_seq_len + block_size - 1) // block_size
num_blocks = batch_size * max_blocks
scale = head_dim**-0.5
q = torch.randn((batch_size, num_q_heads, head_dim), device=device, dtype=dtype)
k_cache = torch.randn((num_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
v_cache = torch.randn((num_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
block_table = torch.empty((batch_size, max_blocks), device=device, dtype=torch.int32)
for batch_idx in range(batch_size):
pages = torch.arange(batch_idx * max_blocks, (batch_idx + 1) * max_blocks, device=device, dtype=torch.int32)
block_table[batch_idx] = pages
out = torch.empty_like(q)
fish_kvcache_attn.fish_decode_kvcache_attn(
q,
k_cache,
v_cache,
block_table,
seq_lens,
out,
scale=scale,
max_seq_len=max_seq_len,
)
expected = _reference_decode_attention(q, k_cache, v_cache, block_table, seq_lens, scale)
torch.testing.assert_close(out, expected, atol=3e-2, rtol=3e-2)
@pytest.mark.skipif(
not torch.cuda.is_available() or not fish_kvcache_attn.is_available(),
reason="Fish kvcache Triton attention is not available",
)
def test_fish_kvcache_native_op_handles_zero_short_and_long_seq_lens():
torch.manual_seed(0)
device = torch.device("cuda")
dtype = torch.float16
block_size = 16
num_q_heads = 4
num_kv_heads = 2
head_dim = 128
seq_lens = torch.tensor([0, 512, 2048], device=device, dtype=torch.int32)
max_seq_len = int(seq_lens.max().item())
batch_size = int(seq_lens.numel())
max_blocks = (max_seq_len + block_size - 1) // block_size
num_blocks = batch_size * max_blocks
scale = head_dim**-0.5
q = torch.randn((batch_size, num_q_heads, head_dim), device=device, dtype=dtype)
k_cache = torch.randn((num_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
v_cache = torch.randn((num_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
block_table = torch.empty((batch_size, max_blocks), device=device, dtype=torch.int32)
for batch_idx in range(batch_size):
pages = torch.arange(batch_idx * max_blocks, (batch_idx + 1) * max_blocks, device=device, dtype=torch.int32)
block_table[batch_idx] = pages
out = torch.full_like(q, float("nan"))
fish_kvcache_attn.fish_decode_kvcache_attn(
q,
k_cache,
v_cache,
block_table,
seq_lens,
out,
scale=scale,
max_seq_len=max_seq_len,
)
expected = _reference_decode_attention(q, k_cache, v_cache, block_table, seq_lens, scale)
assert torch.isfinite(out).all()
torch.testing.assert_close(out, expected, atol=3e-2, rtol=3e-2)
@pytest.mark.skipif(
not torch.cuda.is_available() or not fish_kvcache_attn.is_available(),
reason="Fish kvcache Triton attention is not available",
)
def test_fish_kvcache_native_op_can_be_captured_by_cuda_graph():
device = torch.device("cuda")
dtype = torch.float16
q = torch.randn((2, 4, 128), device=device, dtype=dtype)
k_cache = torch.randn((8, 16, 2, 128), device=device, dtype=dtype)
v_cache = torch.randn((8, 16, 2, 128), device=device, dtype=dtype)
block_table = torch.arange(8, device=device, dtype=torch.int32).reshape(2, 4)
seq_lens = torch.full((2,), 64, device=device, dtype=torch.int32)
out = torch.empty_like(q)
expected = torch.empty_like(q)
fish_kvcache_attn.fish_decode_kvcache_attn(
q,
k_cache,
v_cache,
block_table,
seq_lens,
expected,
scale=128**-0.5,
max_seq_len=64,
)
torch.accelerator.synchronize()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
fish_kvcache_attn.fish_decode_kvcache_attn(
q,
k_cache,
v_cache,
block_table,
seq_lens,
out,
scale=128**-0.5,
max_seq_len=64,
)
graph.replay()
torch.accelerator.synchronize()
torch.testing.assert_close(out, expected, atol=3e-2, rtol=3e-2)
@pytest.mark.skipif(
not torch.cuda.is_available() or not fish_kvcache_attn.is_available(),
reason="Fish kvcache Triton attention is not available",
)
def test_fish_kvcache_native_long_op_can_be_captured_by_cuda_graph():
device = torch.device("cuda")
dtype = torch.float16
seq_len = 2048
block_size = 16
num_q_heads = 4
num_kv_heads = 2
head_dim = 128
max_blocks = (seq_len + block_size - 1) // block_size
q = torch.randn((2, num_q_heads, head_dim), device=device, dtype=dtype)
k_cache = torch.randn((2 * max_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
v_cache = torch.randn((2 * max_blocks, block_size, num_kv_heads, head_dim), device=device, dtype=dtype)
block_table = torch.arange(2 * max_blocks, device=device, dtype=torch.int32).reshape(2, max_blocks)
seq_lens = torch.full((2,), seq_len, device=device, dtype=torch.int32)
out = torch.empty_like(q)
expected = torch.empty_like(q)
fish_kvcache_attn.fish_decode_kvcache_attn(
q,
k_cache,
v_cache,
block_table,
seq_lens,
expected,
scale=head_dim**-0.5,
max_seq_len=seq_len,
)
torch.accelerator.synchronize()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
fish_kvcache_attn.fish_decode_kvcache_attn(
q,
k_cache,
v_cache,
block_table,
seq_lens,
out,
scale=head_dim**-0.5,
max_seq_len=seq_len,
)
graph.replay()
torch.accelerator.synchronize()
torch.testing.assert_close(out, expected, atol=3e-2, rtol=3e-2)
Reference in new issue View Git Blame Copy Permalink