| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import inspect |
| 4 | from collections.abc import Callable |
| 5 | from functools import wraps |
| 6 | |
| 7 | from vllm.distributed.kv_transfer import ( |
| 8 | get_kv_transfer_group, |
| 9 | has_kv_transfer_group, |
| 10 | is_v1_kv_transfer_group, |
| 11 | ) |
| 12 | from vllm.utils.torch_utils import _resolve_layer_name |
| 13 | |
| 14 | |
| 15 | def maybe_transfer_kv_layer(func: Callable) -> Callable: |
| 16 | """Decorator that handles KV layer transfer prior and after execution of |
| 17 | an attention layer, if enabled. Otherwise, the wrapper is a no-op. |
| 18 | |
| 19 | On entry: waits for the KV layer from the connector. |
| 20 | On exit: saves the KV layer to the connector. |
| 21 | """ |
| 22 | # Import at runtime to avoid circular dependency |
| 23 | from vllm.model_executor.layers.attention.attention import get_attention_context |
| 24 | |
| 25 | # Inspect the signature ONCE when the decorator is applied. |
| 26 | sig = inspect.signature(func) |
| 27 | param_names = list(sig.parameters.keys()) |
| 28 | |
| 29 | # Find the index of 'layer_name' parameter. |
| 30 | try: |
| 31 | layer_name_index = param_names.index("layer_name") |
| 32 | except ValueError as e: |
| 33 | raise TypeError( |
| 34 | f"Function {func.__name__} must have a 'layer_name' parameter" |
| 35 | ) from e |
| 36 | |
| 37 | @wraps(func) |
| 38 | def wrapper(*args, **kwargs): |
| 39 | if not has_kv_transfer_group() or not is_v1_kv_transfer_group(): |
| 40 | return func(*args, **kwargs) |
| 41 | |
| 42 | layer_name = _resolve_layer_name(args[layer_name_index]) |
| 43 | |
| 44 | # Extract attention context (metadata, layer, kv_cache, layer_slot_mapping) |
| 45 | attn_metadata, _, kv_cache, _ = get_attention_context(layer_name) |
| 46 | connector = get_kv_transfer_group() |
| 47 | if attn_metadata is None or not connector.has_connector_metadata(): |
| 48 | return func(*args, **kwargs) |
| 49 | |
| 50 | # Wait for KV layer on entry |
| 51 | connector.wait_for_layer_load(layer_name) |
| 52 | |
| 53 | # Execute the function |
| 54 | result = func(*args, **kwargs) |
| 55 | |
| 56 | # Save KV cache layer on exit |
| 57 | connector.save_kv_layer(layer_name, kv_cache, attn_metadata) |
| 58 | |
| 59 | return result |
| 60 | |
| 61 | return wrapper |