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 8, 16:30 UTC
move checks out of `unified_kv_cache_update` custom op (#33943)
Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>
Rohan Potdar committed
Feb 7, 2026 at 07:30 UTC
de3869bb4db76658bc4775ba4ba5fa6ef546237a
7 files changed
+80
-101
vllm/model_executor/layers/attention/attention.py
+14
-6
@@ -422,9 +422,15 @@ class Attention(nn.Module, AttentionLayerBase):
422
key = key.view(-1, self.num_kv_heads, self.head_size)
423
if value is not None:
424
value = value.view(-1, self.num_kv_heads, self.head_size_v)
425
+ kv_cache_dummy_dep = None
426
if self.use_direct_call:
426
- kv_cache_dummy_dep = None
427
- if not self.attn_backend.forward_includes_kv_cache_update:
427
+ # Skip this if sharing KV cache with an earlier attention layer.
428
+ if (
429
+ not self.attn_backend.forward_includes_kv_cache_update
430
+ and self.kv_sharing_target_layer_name is None
431
+ and key is not None
432
+ and value is not None
433
+ ):
434
kv_cache_dummy_dep = unified_kv_cache_update(
435
key, value, self.layer_name
436
)
@@ -437,10 +443,12 @@ class Attention(nn.Module, AttentionLayerBase):
443
kv_cache_dummy_dep=kv_cache_dummy_dep,
444
)
445
else:
440
- kv_cache_dummy_dep = None
441
- if not self.attn_backend.forward_includes_kv_cache_update and (
442
- # torch can only dispatch custom op if a tensor is passed
443
- key is not None or value is not None
446
+ # Skip this if sharing KV cache with an earlier attention layer.
447
+ if (
448
+ not self.attn_backend.forward_includes_kv_cache_update
449
+ and self.kv_sharing_target_layer_name is None
450
+ and key is not None
451
+ and value is not None
452
):
453
kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
454
key, value, self.layer_name
vllm/model_executor/layers/attention/cross_attention.py
+3
@@ -136,6 +136,9 @@ def create_cross_attention_backend(
136
if (
137
not underlying_attn_backend.forward_includes_kv_cache_update
138
and attn_metadata is not None
139
+ and layer.kv_sharing_target_layer_name is None
140
+ and key is not None
141
+ and value is not None
142
):
143
self.do_kv_cache_update(
144
layer, key, value, kv_cache, attn_metadata.slot_mapping
vllm/model_executor/models/whisper_causal.py
+3
@@ -172,6 +172,9 @@ def create_whisper_attention_backend_with_block_pooling(
172
if (
173
not underlying_attn_backend.forward_includes_kv_cache_update
174
and attn_metadata is not None
175
+ and layer.kv_sharing_target_layer_name is None
176
+ and key is not None
177
+ and value is not None
178
):
179
self.do_kv_cache_update(
180
layer, key, value, kv_cache, attn_metadata.slot_mapping
vllm/v1/attention/backends/flash_attn.py
-10
@@ -771,16 +771,6 @@ class FlashAttentionImpl(AttentionImpl):
771
# we use direct Q, K, V tensors without caching
772
return
773
774
- # key and value may be None in the case of cross attention. They are
775
- # calculated once based on the output from the encoder and then cached
776
- # in KV cache.
777
- if (
778
- self.kv_sharing_target_layer_name is not None
779
- or key is None
780
- or value is None
781
- ):
782
- return
783
-
774
key_cache, value_cache = kv_cache.unbind(0)
775
776
# Reshape the input keys and values and store them in the cache.
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+11
-20
@@ -196,23 +196,14 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
196
):
197
key_cache, value_cache = kv_cache.unbind(0)
198
199
- # key and value may be None in the case of cross attention. They are
200
- # calculated once based on the output from the encoder and then cached
201
- # in KV cache.
202
- if (
203
- self.kv_sharing_target_layer_name is None
204
- and key is not None
205
- and value is not None
206
- ):
207
- # Reshape the input keys and values and store them in the cache.
208
- # Skip this if sharing KV cache with an earlier attention layer.
209
- ops.reshape_and_cache_flash(
210
- key,
211
- value,
212
- key_cache,
213
- value_cache,
214
- slot_mapping,
215
- self.kv_cache_dtype,
216
- layer._k_scale,
217
- layer._v_scale,
218
- )
199
+ # Reshape the input keys and values and store them in the cache.
200
+ ops.reshape_and_cache_flash(
201
+ key,
202
+ value,
203
+ key_cache,
204
+ value_cache,
205
+ slot_mapping,
206
+ self.kv_cache_dtype,
207
+ layer._k_scale,
208
+ layer._v_scale,
209
+ )
vllm/v1/attention/backends/rocm_attn.py
+32
-42
@@ -383,45 +383,35 @@ class RocmAttentionImpl(AttentionImpl):
383
kv_cache, self.num_kv_heads, self.head_size
384
)
385
386
- # key and value may be None in the case of cross attention. They are
387
- # calculated once based on the output from the encoder and then cached
388
- # in KV cache.
389
- if (
390
- self.kv_sharing_target_layer_name is None
391
- and key is not None
392
- and value is not None
393
- ):
394
- # Reshape the input keys and values and store them in the cache.
395
- # Skip this if sharing KV cache with an earlier attention layer.
396
-
397
- # Get the actual block_size from value_cache
398
- # value_cache shape: [num_blocks, num_heads, head_size, block_size]
399
- block_size = value_cache.shape[3]
400
- # Determine if it is a power of 2
401
- is_pow2 = block_size > 0 and (block_size & (block_size - 1) == 0)
402
-
403
- if is_pow2:
404
- # Normal 16, 32, 64, etc., use vLLM native HIP C++ logic
405
- PagedAttention.write_to_paged_cache(
406
- key,
407
- value,
408
- key_cache,
409
- value_cache,
410
- slot_mapping,
411
- self.kv_cache_dtype,
412
- layer._k_scale,
413
- layer._v_scale,
414
- )
415
- else:
416
- # Case B: Non-standard blocks (e.g., 544 in Qwen3),
417
- # force using our modified Triton logic
418
- triton_reshape_and_cache_flash(
419
- key,
420
- value,
421
- key_cache,
422
- value_cache,
423
- slot_mapping,
424
- self.kv_cache_dtype,
425
- layer._k_scale,
426
- layer._v_scale,
427
- )
386
+ # Reshape the input keys and values and store them in the cache.
387
+ # Get the actual block_size from value_cache
388
+ # value_cache shape: [num_blocks, num_heads, head_size, block_size]
389
+ block_size = value_cache.shape[3]
390
+ # Determine if it is a power of 2
391
+ is_pow2 = block_size > 0 and (block_size & (block_size - 1) == 0)
392
+
393
+ if is_pow2:
394
+ # Normal 16, 32, 64, etc., use vLLM native HIP C++ logic
395
+ PagedAttention.write_to_paged_cache(
396
+ key,
397
+ value,
398
+ key_cache,
399
+ value_cache,
400
+ slot_mapping,
401
+ self.kv_cache_dtype,
402
+ layer._k_scale,
403
+ layer._v_scale,
404
+ )
405
+ else:
406
+ # Case B: Non-standard blocks (e.g., 544 in Qwen3),
407
+ # force using our modified Triton logic
408
+ triton_reshape_and_cache_flash(
409
+ key,
410
+ value,
411
+ key_cache,
412
+ value_cache,
413
+ slot_mapping,
414
+ self.kv_cache_dtype,
415
+ layer._k_scale,
416
+ layer._v_scale,
417
+ )
vllm/v1/attention/backends/triton_attn.py
+17
-23
@@ -579,26 +579,20 @@ class TritonAttentionImpl(AttentionImpl):
579
# For decoder and cross-attention, use KV cache as before
580
key_cache, value_cache = kv_cache.unbind(1)
581
582
- if (
583
- self.kv_sharing_target_layer_name is None
584
- and key is not None
585
- and value is not None
586
- ):
587
- # Reshape the input keys and values and store them in the cache.
588
- # Skip this if sharing KV cache with an earlier attention layer.
589
- if self.kv_cache_dtype.startswith("fp8"):
590
- key_cache = key_cache.view(self.fp8_dtype)
591
- value_cache = value_cache.view(self.fp8_dtype)
592
- # triton kernel does not support uint8 kv_cache
593
- # (because some explicit casts (e.g. float8_e4m3fnuz)
594
- # are not supported)
595
- triton_reshape_and_cache_flash(
596
- key,
597
- value,
598
- key_cache,
599
- value_cache,
600
- slot_mapping,
601
- self.kv_cache_dtype,
602
- layer._k_scale,
603
- layer._v_scale,
604
- )
582
+ # Reshape the input keys and values and store them in the cache.
583
+ if self.kv_cache_dtype.startswith("fp8"):
584
+ key_cache = key_cache.view(self.fp8_dtype)
585
+ value_cache = value_cache.view(self.fp8_dtype)
586
+ # triton kernel does not support uint8 kv_cache
587
+ # (because some explicit casts (e.g. float8_e4m3fnuz)
588
+ # are not supported)
589
+ triton_reshape_and_cache_flash(
590
+ key,
591
+ value,
592
+ key_cache,
593
+ value_cache,
594
+ slot_mapping,
595
+ self.kv_cache_dtype,
596
+ layer._k_scale,
597
+ layer._v_scale,
598
+ )