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 7, 01:30 UTC
feat(attention): extract KV-cache update from FlexAttention backend (#36263)
Signed-off-by: cong-or <conchubhar.gannon@gmail.com>
cong-or committed
Mar 9, 2026 at 03:40 UTC
747431044df6b15c7b359b5720cc7368c662c232
1 file changed
+25
-11
vllm/v1/attention/backends/flex_attention.py
+25
-11
@@ -82,6 +82,8 @@ class FlexAttentionBackend(AttentionBackend):
82
]
83
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"]
84
85
+ forward_includes_kv_cache_update: bool = False
86
+
87
@staticmethod
88
def get_name() -> str:
89
return "FLEX_ATTENTION"
@@ -827,6 +829,29 @@ class FlexAttentionImpl(AttentionImpl):
829
assert tensor.ndim == 3
830
return tensor[None, :, :, :]
831
832
+ def do_kv_cache_update(
833
+ self,
834
+ layer: torch.nn.Module,
835
+ key: torch.Tensor,
836
+ value: torch.Tensor,
837
+ kv_cache: torch.Tensor,
838
+ slot_mapping: torch.Tensor,
839
+ ) -> None:
840
+ if self.attn_type == AttentionType.ENCODER_ONLY:
841
+ return
842
+
843
+ key_cache, value_cache = kv_cache.unbind(0)
844
+ torch.ops._C_cache_ops.reshape_and_cache_flash(
845
+ key,
846
+ value,
847
+ key_cache,
848
+ value_cache,
849
+ slot_mapping,
850
+ self.kv_cache_dtype,
851
+ layer._k_scale,
852
+ layer._v_scale,
853
+ )
854
+
855
def forward(
856
self,
857
layer: torch.nn.Module,
@@ -908,17 +933,6 @@ class FlexAttentionImpl(AttentionImpl):
933
assert self.attn_type == AttentionType.DECODER
934
key_cache, value_cache = kv_cache.unbind(0)
935
911
- torch.ops._C_cache_ops.reshape_and_cache_flash(
912
- key,
913
- value,
914
- key_cache,
915
- value_cache,
916
- attn_metadata.slot_mapping,
917
- self.kv_cache_dtype,
918
- layer._k_scale,
919
- layer._v_scale,
920
- )
921
-
936
# View out the block_size dim
937
key_cache = key_cache.view(-1, self.num_kv_heads, self.head_size)
938
value_cache = value_cache.view(-1, self.num_kv_heads, self.head_size)