biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Performance] Split FlashAttn attention and cache update (#25954)

Browse Source 
Signed-off-by: ElizaWszola <ewszola@redhat.com>
Signed-off-by: mgoin <mgoin64@gmail.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Signed-off-by: Luka Govedič <luka.govedic@gmail.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
Co-authored-by: mgoin <mgoin64@gmail.com>
Co-authored-by: Varun Sundar Rabindranath <varunsundar08@gmail.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Co-authored-by: Luka Govedič <luka.govedic@gmail.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Luka Govedič <lgovedic@redhat.com>
This commit is contained in:
ElizaWszola
2026-01-24 02:28:06 +01:00
committed by GitHub
parent 0118cdcc02
commit a28b94e6ef
21 changed files with 458 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
tests/v1/attention/test_attention_backends.py
View File

@@ -13,6 +13,7 @@ from tests.v1.attention.utils import (
create_common_attn_metadata,
create_standard_kv_cache_spec,
create_vllm_config,
try_backend_includes_kv_cache_update,
try_get_attention_backend,
)
from vllm.config import ModelConfig
@@ -295,6 +296,10 @@ def run_attention_backend(
# Run forward pass
# NOTE: The query, key, and value are already shaped correctly
# in the calling test function.
if not try_backend_includes_kv_cache_update(actual_backend):
impl.do_kv_cache_update(
mock_layer, key, value, kv_cache, attn_metadata.slot_mapping
)
output = impl.forward(
mock_layer, query, key, value, kv_cache, attn_metadata, output=output
)

12
tests/v1/attention/utils.py
View File

@@ -130,6 +130,18 @@ def try_get_attention_backend(
raise AssertionError("unreachable") from None
def try_backend_includes_kv_cache_update(
backend: AttentionBackendEnum,
) -> bool:
"""Try to get the attention backend class, skipping test if not found."""
try:
backend_class = backend.get_class()
return backend_class.forward_includes_kv_cache_update
except ImportError as e:
pytest.skip(f"{backend.name} not available: {e}")
raise AssertionError("unreachable") from None
def create_standard_kv_cache_spec(vllm_config: VllmConfig) -> FullAttentionSpec:
"""Create a FullAttentionSpec from ModelParams only."""
return FullAttentionSpec(

2
tests/v1/kv_connector/unit/test_decode_bench_connector.py
View File

@@ -86,7 +86,7 @@ class DecodeBenchTestRunner:
self._block_hasher = get_request_block_hasher(block_size, sha256)
self._dummy_ctx: ForwardContext = ForwardContext(
no_compile_layers={}, attn_metadata={}, virtual_engine=0
no_compile_layers={}, attn_metadata={}, virtual_engine=0, slot_mapping={}
)
def new_request(self, token_ids: list[int]) -> Request:

8
tests/v1/kv_connector/unit/test_nixl_connector.py
View File

@@ -548,6 +548,7 @@ class TestNixlHandshake:
no_compile_layers={},
attn_metadata={},
virtual_engine=0,
slot_mapping={},
)
_before_load = time.perf_counter()
connector.start_load_kv(dummy_ctx)
@@ -618,6 +619,7 @@ class TestNixlHandshake:
no_compile_layers={},
attn_metadata={},
virtual_engine=0,
slot_mapping={},
)
_before_load = time.perf_counter()
connector.start_load_kv(dummy_ctx)
@@ -844,6 +846,7 @@ class TestNixlHandshake:
no_compile_layers={},
attn_metadata={},
virtual_engine=0,
slot_mapping={},
)
_before_load = time.perf_counter()
connector.start_load_kv(dummy_ctx)
@@ -1006,6 +1009,7 @@ def test_kv_connector_stats(default_vllm_config, dist_init):
no_compile_layers={},
attn_metadata={},
virtual_engine=0,
slot_mapping={},
)
connector.start_load_kv(dummy_ctx)
@@ -1767,6 +1771,7 @@ def test_aborted_request_removed_from_worker_in_batch(default_vllm_config, dist_
no_compile_layers={},
attn_metadata={},
virtual_engine=0,
slot_mapping={},
)
connector.start_load_kv(dummy_ctx)
@@ -1917,6 +1922,7 @@ def test_transfer_failure_logging(
no_compile_layers={},
attn_metadata={},
virtual_engine=0,
slot_mapping={},
)
# Capture logs from the nixl_connector logger specifically
@@ -2017,6 +2023,7 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init):
no_compile_layers={},
attn_metadata={},
virtual_engine=0,
slot_mapping={},
)
connector.start_load_kv(dummy_ctx)
@@ -2067,6 +2074,7 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)
no_compile_layers={},
attn_metadata={},
virtual_engine=0,
slot_mapping={},
)
connector.start_load_kv(dummy_ctx)

5
tests/v1/kv_connector/unit/test_offloading_connector.py
View File

@@ -209,7 +209,10 @@ class RequestRunner:
self._block_hasher = get_request_block_hasher(gpu_block_size, sha256)
self._dummy_ctx: ForwardContext = ForwardContext(
no_compile_layers={}, attn_metadata={}, virtual_engine=0
no_compile_layers={},
attn_metadata={},
virtual_engine=0,
slot_mapping={},
)
def new_request(self, token_ids: list[int]):

9
tests/v1/spec_decode/test_tree_attention.py
View File

@@ -9,6 +9,7 @@ import torch
from tests.v1.attention.utils import (
create_standard_kv_cache_spec,
create_vllm_config,
try_backend_includes_kv_cache_update,
try_get_attention_backend,
)
from vllm.config import ParallelConfig, SpeculativeConfig
@@ -120,6 +121,14 @@ def forward_attention(
key = k.view(-1, num_kv_heads, dim_per_head)
value = v.view(-1, num_kv_heads, dim_per_head)
output = torch.empty_like(query)
if not try_backend_includes_kv_cache_update(backend):
instance.do_kv_cache_update(
layer=layer,
key=key,
value=value,
kv_cache=kv_cache,
slot_mapping=attn_metadata.slot_mapping,
)
return instance.forward(
layer=layer,
query=query,