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, 16:30 UTC
[Performance] Extract KV cache update op from flashinfer forward (#35422)
Signed-off-by: ElizaWszola <ewszola@redhat.com>
ElizaWszola committed
Feb 26, 2026 at 22:29 UTC
98217b09f9ce22429ce35badfa1d50e1f4fe4137
1 file changed
+37
-25
vllm/v1/attention/backends/flashinfer.py
+37
-25
@@ -381,6 +381,8 @@ class FlashInferBackend(AttentionBackend):
381
return "HND"
382
return None
383
384
+ forward_includes_kv_cache_update: bool = False
385
+
386
387
@dataclass
388
class FIPrefill:
@@ -1330,32 +1332,15 @@ class FlashInferImpl(AttentionImpl):
1332
1333
num_actual_tokens = attn_metadata.num_actual_tokens
1334
1333
- if self.kv_sharing_target_layer_name is None:
1334
- # Reshape the input keys and values and store them in the cache.
1335
- # Skip this if sharing KV cache with an earlier attention layer.
1336
- # NOTE(woosuk): Here, key and value are padded while slot_mapping is
1337
- # not padded. However, we don't need to do key[:num_actual_tokens]
1338
- # and value[:num_actual_tokens] because the reshape_and_cache_flash
1339
- # op uses the slot_mapping's shape to determine the number of
1340
- # actual tokens.
1341
- torch.ops._C_cache_ops.reshape_and_cache_flash(
1342
- key,
1343
- value,
1344
- kv_cache[:, 0],
1345
- kv_cache[:, 1],
1346
- attn_metadata.slot_mapping,
1347
- self.kv_cache_dtype,
1348
- layer._k_scale,
1349
- layer._v_scale,
1335
+ # The FlashInfer api requires data to be in fp8_e4m3 or fp8_e5m2
1336
+ # to process the cache when the kv_cache_dtype is fp8
1337
+ if self.kv_sharing_target_layer_name is None and self.kv_cache_dtype.startswith(
1338
+ "fp8"
1339
+ ):
1340
+ torch_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
1341
+ self.kv_cache_dtype
1342
)
1351
-
1352
- # The FlashInfer api requires data to be in fp8_e4m3 or fp8_e5m2
1353
- # to process the cache when the kv_cache_dtype is fp8
1354
- if self.kv_cache_dtype.startswith("fp8"):
1355
- torch_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
1356
- self.kv_cache_dtype
1357
- )
1358
- kv_cache = kv_cache.view(torch_dtype)
1343
+ kv_cache = kv_cache.view(torch_dtype)
1344
1345
# Inputs and outputs may be padded for CUDA graphs
1346
query = query[:num_actual_tokens]
@@ -1599,6 +1584,33 @@ class FlashInferImpl(AttentionImpl):
1584
)
1585
return output_padded
1586
1587
+ def do_kv_cache_update(
1588
+ self,
1589
+ layer: torch.nn.Module,
1590
+ key: torch.Tensor,
1591
+ value: torch.Tensor,
1592
+ kv_cache: torch.Tensor,
1593
+ slot_mapping: torch.Tensor,
1594
+ ) -> None:
1595
+ if self.kv_sharing_target_layer_name is None:
1596
+ # Reshape the input keys and values and store them in the cache.
1597
+ # Skip this if sharing KV cache with an earlier attention layer.
1598
+ # NOTE(woosuk): Here, key and value are padded while slot_mapping is
1599
+ # not padded. However, we don't need to do key[:num_actual_tokens]
1600
+ # and value[:num_actual_tokens] because the reshape_and_cache_flash
1601
+ # op uses the slot_mapping's shape to determine the number of
1602
+ # actual tokens.
1603
+ torch.ops._C_cache_ops.reshape_and_cache_flash(
1604
+ key,
1605
+ value,
1606
+ kv_cache[:, 0],
1607
+ kv_cache[:, 1],
1608
+ slot_mapping,
1609
+ self.kv_cache_dtype,
1610
+ layer._k_scale,
1611
+ layer._v_scale,
1612
+ )
1613
+
1614
1615
def fast_plan_decode(
1616
self, # decode wrapper