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, 04:00 UTC
[Performance] Extract KV-cache update from TreeAttention backend (#35384)
Signed-off-by: dorhuri123 <dor.huri1@live.biu.ac.il>
Dor Huri committed
Mar 6, 2026 at 02:22 UTC
ebed80a7c8c652ff43b5bd910c8fe35d73bfa786
1 file changed
+28
-19
vllm/v1/attention/backends/tree_attn.py
+28
-19
@@ -31,6 +31,7 @@ logger = init_logger(__name__)
31
class TreeAttentionBackend(AttentionBackend):
32
accept_output_buffer: bool = True
33
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
34
+ forward_includes_kv_cache_update: bool = False
35
36
@staticmethod
37
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
@@ -326,6 +327,33 @@ class TreeAttentionImpl(AttentionImpl):
327
"TreeAttentionImpl."
328
)
329
330
+ def do_kv_cache_update(
331
+ self,
332
+ layer: torch.nn.Module,
333
+ key: torch.Tensor,
334
+ value: torch.Tensor,
335
+ kv_cache: torch.Tensor,
336
+ slot_mapping: torch.Tensor,
337
+ ) -> None:
338
+ key_cache, value_cache = kv_cache.unbind(0)
339
+
340
+ # Reshape the input keys and values and store them in the cache.
341
+ # NOTE(woosuk): Here, key and value are padded while slot_mapping is
342
+ # not padded. However, we don't need to do key[:num_actual_tokens]
343
+ # and value[:num_actual_tokens] because the reshape_and_cache_flash
344
+ # op uses the slot_mapping's shape to determine the number of
345
+ # actual tokens.
346
+ ops.reshape_and_cache_flash(
347
+ key,
348
+ value,
349
+ key_cache,
350
+ value_cache,
351
+ slot_mapping,
352
+ self.kv_cache_dtype,
353
+ layer._k_scale,
354
+ layer._v_scale,
355
+ )
356
+
357
def forward(
358
self,
359
layer: torch.nn.Module,
@@ -361,26 +389,7 @@ class TreeAttentionImpl(AttentionImpl):
389
# Profiling run.
390
return output.fill_(0)
391
364
- # Cache the input KVs.
392
key_cache, value_cache = kv_cache.unbind(0)
366
- if self.kv_sharing_target_layer_name is None:
367
- # Reshape the input keys and values and store them in the cache.
368
- # Skip this if sharing KV cache with an earlier attention layer.
369
- # NOTE(woosuk): Here, key and value are padded while slot_mapping is
370
- # not padded. However, we don't need to do key[:num_actual_tokens]
371
- # and value[:num_actual_tokens] because the reshape_and_cache_flash
372
- # op uses the slot_mapping's shape to determine the number of
373
- # actual tokens.
374
- ops.reshape_and_cache_flash(
375
- key,
376
- value,
377
- key_cache,
378
- value_cache,
379
- attn_metadata.slot_mapping,
380
- self.kv_cache_dtype,
381
- layer._k_scale,
382
- layer._v_scale,
383
- )
393
394
num_actual_tokens = attn_metadata.num_actual_tokens
395
num_decode_tokens = attn_metadata.num_decode_tokens