A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 09:30 UTC
[Misc] Refactor Attention kv transfer methods into decorator (#27816)
Signed-off-by: NickLucche <nlucches@redhat.com> Signed-off-by: Mark McLoughlin <markmc@redhat.com> Co-authored-by: Mark McLoughlin <markmc@redhat.com>
Nicolò Lucchesi committed
Nov 12, 2025 at 17:05 UTC
728a9eb70ee30b1ab355a98f7e19fb81a0a7b873
2 files changed
+99
-76
vllm/attention/layer.py
+39
-76
@@ -15,14 +15,10 @@ from vllm.attention.backends.abstract import AttentionBackend, MLAAttentionImpl
15
from vllm.attention.backends.registry import AttentionBackendEnum
16
from vllm.attention.selector import get_attn_backend
17
from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target
18
+from vllm.attention.utils.kv_transfer_utils import maybe_transfer_kv_layer
19
from vllm.config import CacheConfig, get_current_vllm_config
20
from vllm.config.multimodal import MultiModalConfig
21
from vllm.config.vllm import VllmConfig
21
-from vllm.distributed.kv_transfer import (
22
- get_kv_transfer_group,
23
- has_kv_transfer_group,
24
- is_v1_kv_transfer_group,
25
-)
22
from vllm.forward_context import ForwardContext, get_forward_context
23
from vllm.logger import init_logger
24
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
@@ -842,41 +838,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):
838
)
839
840
845
-def wait_for_kv_layer_from_connector(layer_name: str):
846
- if not has_kv_transfer_group() or not is_v1_kv_transfer_group():
847
- return
848
-
849
- connector = get_kv_transfer_group()
850
- if not connector.has_connector_metadata():
851
- return
852
-
853
- forward_context: ForwardContext = get_forward_context()
854
- attn_metadata = forward_context.attn_metadata
855
- if attn_metadata is None:
856
- return
857
- assert isinstance(attn_metadata, dict)
858
- connector.wait_for_layer_load(layer_name)
859
-
860
-
861
-def maybe_save_kv_layer_to_connector(
862
- layer_name: str,
863
- kv_cache_layer: list[torch.Tensor],
864
-):
865
- if not has_kv_transfer_group() or not is_v1_kv_transfer_group():
866
- return
867
-
868
- connector = get_kv_transfer_group()
869
- if not connector.has_connector_metadata():
870
- return
871
-
872
- forward_context: ForwardContext = get_forward_context()
873
- attn_metadata = forward_context.attn_metadata
874
- if attn_metadata is None:
875
- return
876
- assert isinstance(attn_metadata, dict)
877
- connector.save_kv_layer(layer_name, kv_cache_layer, attn_metadata[layer_name])
878
-
879
-
841
def maybe_calc_kv_scales(
842
query: torch.Tensor,
843
key: torch.Tensor,
@@ -911,23 +872,46 @@ direct_register_custom_op(
872
)
873
874
914
-def unified_attention(
915
- query: torch.Tensor,
916
- key: torch.Tensor,
917
- value: torch.Tensor,
875
+def get_attention_context(
876
layer_name: str,
919
-) -> torch.Tensor:
920
- wait_for_kv_layer_from_connector(layer_name)
877
+) -> tuple[dict | object | None, Attention | MLAAttention, torch.Tensor]:
878
+ """Extract attention context for a given layer.
879
+
880
+ This helper function extracts the attention metadata, attention layer
881
+ instance, and KV cache tensor for a specific layer.
882
+
883
+ Args:
884
+ layer_name: The name/identifier of the attention layer.
885
886
+ Returns:
887
+ A tuple containing:
888
+ - attn_metadata: Attention metadata for this specific layer, or None if
889
+ no metadata available
890
+ - attn_layer: The attention layer instance (Attention or MLAAttention)
891
+ - kv_cache: The KV cache tensor for current virtual engine
892
+
893
+ Note: attn_metadata may be None, but attn_layer and kv_cache are always
894
+ extracted from the forward context.
895
+ """
896
forward_context: ForwardContext = get_forward_context()
897
attn_metadata = forward_context.attn_metadata
898
if isinstance(attn_metadata, dict):
899
attn_metadata = attn_metadata[layer_name]
926
- self = forward_context.no_compile_layers[layer_name]
927
- kv_cache = self.kv_cache[forward_context.virtual_engine]
900
+ attn_layer: Attention | MLAAttention = forward_context.no_compile_layers[layer_name]
901
+ kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
902
+ return attn_metadata, attn_layer, kv_cache
903
+
904
+
905
+@maybe_transfer_kv_layer
906
+def unified_attention(
907
+ query: torch.Tensor,
908
+ key: torch.Tensor,
909
+ value: torch.Tensor,
910
+ layer_name: str,
911
+) -> torch.Tensor:
912
+ attn_metadata, self, kv_cache = get_attention_context(layer_name)
913
output = self.impl.forward(self, query, key, value, kv_cache, attn_metadata)
914
930
- maybe_save_kv_layer_to_connector(layer_name, kv_cache)
915
return output
916
917
@@ -947,6 +931,7 @@ direct_register_custom_op(
931
)
932
933
934
+@maybe_transfer_kv_layer
935
def unified_attention_with_output(
936
query: torch.Tensor,
937
key: torch.Tensor,
@@ -956,13 +941,7 @@ def unified_attention_with_output(
941
output_scale: torch.Tensor | None = None,
942
output_block_scale: torch.Tensor | None = None,
943
) -> None:
959
- wait_for_kv_layer_from_connector(layer_name)
960
- forward_context: ForwardContext = get_forward_context()
961
- attn_metadata = forward_context.attn_metadata
962
- if isinstance(attn_metadata, dict):
963
- attn_metadata = attn_metadata[layer_name]
964
- self = forward_context.no_compile_layers[layer_name]
965
- kv_cache = self.kv_cache[forward_context.virtual_engine]
944
+ attn_metadata, self, kv_cache = get_attention_context(layer_name)
945
self.impl.forward(
946
self,
947
query,
@@ -975,8 +954,6 @@ def unified_attention_with_output(
954
output_block_scale=output_block_scale,
955
)
956
978
- maybe_save_kv_layer_to_connector(layer_name, kv_cache)
979
-
957
958
def unified_attention_with_output_fake(
959
query: torch.Tensor,
@@ -998,23 +975,16 @@ direct_register_custom_op(
975
)
976
977
978
+@maybe_transfer_kv_layer
979
def unified_mla_attention(
980
q: torch.Tensor,
981
kv_c_normed: torch.Tensor,
982
k_pe: torch.Tensor,
983
layer_name: str,
984
) -> torch.Tensor:
1007
- wait_for_kv_layer_from_connector(layer_name)
1008
-
1009
- forward_context: ForwardContext = get_forward_context()
1010
- attn_metadata = forward_context.attn_metadata
1011
- if isinstance(attn_metadata, dict):
1012
- attn_metadata = attn_metadata[layer_name]
1013
- self: MLAAttention = forward_context.no_compile_layers[layer_name]
1014
- kv_cache = self.kv_cache[forward_context.virtual_engine]
985
+ attn_metadata, self, kv_cache = get_attention_context(layer_name)
986
output = self.impl.forward(self, q, kv_c_normed, k_pe, kv_cache, attn_metadata)
987
1017
- maybe_save_kv_layer_to_connector(layer_name, kv_cache)
988
return output
989
990
@@ -1036,6 +1006,7 @@ direct_register_custom_op(
1006
)
1007
1008
1009
+@maybe_transfer_kv_layer
1010
def unified_mla_attention_with_output(
1011
q: torch.Tensor,
1012
kv_c_normed: torch.Tensor,
@@ -1045,13 +1016,7 @@ def unified_mla_attention_with_output(
1016
output_scale: torch.Tensor | None = None,
1017
output_block_scale: torch.Tensor | None = None,
1018
) -> None:
1048
- wait_for_kv_layer_from_connector(layer_name)
1049
- forward_context: ForwardContext = get_forward_context()
1050
- attn_metadata = forward_context.attn_metadata
1051
- if isinstance(attn_metadata, dict):
1052
- attn_metadata = attn_metadata[layer_name]
1053
- self: MLAAttention = forward_context.no_compile_layers[layer_name]
1054
- kv_cache = self.kv_cache[forward_context.virtual_engine]
1019
+ attn_metadata, self, kv_cache = get_attention_context(layer_name)
1020
self.impl.forward(
1021
self,
1022
q,
@@ -1064,8 +1029,6 @@ def unified_mla_attention_with_output(
1029
output_block_scale=output_block_scale,
1030
)
1031
1067
- maybe_save_kv_layer_to_connector(layer_name, kv_cache)
1068
-
1032
1033
def unified_mla_attention_with_output_fake(
1034
q: torch.Tensor,
vllm/attention/utils/kv_transfer_utils.py
new
+60
@@ -0,0 +1,60 @@
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
+
13
+
14
+def maybe_transfer_kv_layer(func: Callable) -> Callable:
15
+ """Decorator that handles KV layer transfer prior and after execution of
16
+ an attention layer, if enabled. Otherwise, the wrapper is a no-op.
17
+
18
+ On entry: waits for the KV layer from the connector.
19
+ On exit: saves the KV layer to the connector.
20
+ """
21
+ # Import at runtime to avoid circular dependency
22
+ from vllm.attention.layer import get_attention_context
23
+
24
+ # Inspect the signature ONCE when the decorator is applied.
25
+ sig = inspect.signature(func)
26
+ param_names = list(sig.parameters.keys())
27
+
28
+ # Find the index of 'layer_name' parameter.
29
+ try:
30
+ layer_name_index = param_names.index("layer_name")
31
+ except ValueError as e:
32
+ raise TypeError(
33
+ f"Function {func.__name__} must have a 'layer_name' parameter"
34
+ ) from e
35
+
36
+ @wraps(func)
37
+ def wrapper(*args, **kwargs):
38
+ if not has_kv_transfer_group() or not is_v1_kv_transfer_group():
39
+ return func(*args, **kwargs)
40
+
41
+ layer_name: str = args[layer_name_index]
42
+
43
+ # Extract attention context (layer-specific metadata, layer, and kv_cache)
44
+ attn_metadata, attn_layer, kv_cache = get_attention_context(layer_name)
45
+ connector = get_kv_transfer_group()
46
+ if attn_metadata is None or not connector.has_connector_metadata():
47
+ return func(*args, **kwargs)
48
+
49
+ # Wait for KV layer on entry
50
+ connector.wait_for_layer_load(layer_name)
51
+
52
+ # Execute the function
53
+ result = func(*args, **kwargs)
54
+
55
+ # Save KV cache layer on exit
56
+ connector.save_kv_layer(layer_name, kv_cache, attn_metadata)
57
+
58
+ return result
59
+
60
+ return wrapper