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 15, 22:15 UTC
Add comments on accessing `kv_cache` and `attn_metadata` (#13887)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Feb 26, 2025 at 10:41 UTC
0ecdd98031f9a0ea7941f9785e3ab2d928f4f2df
1 file changed
+13
vllm/attention/layer.py
+13
@@ -47,6 +47,10 @@ class Attention(nn.Module):
47
attn_type: str = AttentionType.DECODER,
48
**extra_impl_args,
49
) -> None:
50
+ """
51
+ The KV cache is stored inside this class and is accessed via
52
+ `self.kv_cache`.
53
+ """
54
super().__init__()
55
if per_layer_sliding_window is not None:
56
# per-layer sliding window
@@ -155,6 +159,15 @@ class Attention(nn.Module):
159
key: torch.Tensor,
160
value: torch.Tensor,
161
) -> torch.Tensor:
162
+ """
163
+ The KV cache is stored inside this class and is accessed via
164
+ `self.kv_cache`.
165
+
166
+ Attention metadata (`attn_metadata`) is set using a context manager in
167
+ the model runner's `execute_model` method. It is accessed via forward
168
+ context using
169
+ `vllm.forward_context.get_forward_context().attn_metadata`.
170
+ """
171
if self.calculate_kv_scales:
172
attn_metadata = get_forward_context().attn_metadata
173
if attn_metadata.enable_kv_scales_calculation: