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 5, 18:30 UTC
feat(attention): extract KV-cache update from FlashAttentionDiffKV ba… (#36466)
Signed-off-by: Prathmesh Bhatt <71340361+Prathmesh234@users.noreply.github.com>
Prathmesh Bhatt committed
Mar 30, 2026 at 16:16 UTC
93b3ec15859af6610cad00410f09f024bc73dc62
1 file changed
+34
-27
vllm/v1/attention/backends/flash_attn_diffkv.py
+34
-27
@@ -85,6 +85,40 @@ class FlashAttentionDiffKVBackend(FlashAttentionBackend):
85
86
87
class FlashAttentionDiffKVImpl(FlashAttentionImpl):
88
+ def do_kv_cache_update(
89
+ self,
90
+ layer: torch.nn.Module,
91
+ key: torch.Tensor,
92
+ value: torch.Tensor,
93
+ kv_cache: torch.Tensor,
94
+ slot_mapping: torch.Tensor,
95
+ ) -> None:
96
+ if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
97
+ # For encoder attention,
98
+ # we use direct Q, K, V tensors without caching
99
+ return
100
+
101
+ # Unlike standard FlashAttn which splits kv_cache via unbind(0),
102
+ # DiffKV packs K and V into a single tensor along the last dim:
103
+ # kv_cache shape: [num_blocks, block_size, num_kv_heads,
104
+ # head_size_k + head_size_v]
105
+ # The triton kernel handles this combined layout directly.
106
+ #
107
+ # NOTE(woosuk): key and value are padded while slot_mapping is
108
+ # not padded. However, we don't need to do key[:num_actual_tokens]
109
+ # and value[:num_actual_tokens] because the reshape_and_cache_flash
110
+ # op uses the slot_mapping's shape to determine the number of
111
+ # actual tokens.
112
+ triton_reshape_and_cache_flash_diffkv(
113
+ key,
114
+ value,
115
+ kv_cache,
116
+ slot_mapping,
117
+ self.kv_cache_dtype,
118
+ layer._k_scale,
119
+ layer._v_scale,
120
+ )
121
+
122
def forward(
123
self,
124
layer: torch.nn.Module,
@@ -157,33 +191,6 @@ class FlashAttentionDiffKVImpl(FlashAttentionImpl):
191
key_cache = kv_cache[..., : self.head_size]
192
value_cache = kv_cache[..., self.head_size :]
193
160
- # key and value may be None in the case of cross attention. They are
161
- # calculated once based on the output from the encoder and then cached
162
- # in KV cache.
163
- if (
164
- self.kv_sharing_target_layer_name is None
165
- and key is not None
166
- and value is not None
167
- ):
168
- # Reshape the input keys and values and store them in the cache.
169
- # Skip this if sharing KV cache with an earlier attention layer.
170
- # NOTE(woosuk): Here, key and value are padded while slot_mapping is
171
- # not padded. However, we don't need to do key[:num_actual_tokens]
172
- # and value[:num_actual_tokens] because the reshape_and_cache_flash
173
- # op uses the slot_mapping's shape to determine the number of
174
- # actual tokens.
175
-
176
- # kv_cache update for different head_size K and V
177
- triton_reshape_and_cache_flash_diffkv(
178
- key,
179
- value,
180
- kv_cache,
181
- attn_metadata.slot_mapping,
182
- self.kv_cache_dtype,
183
- layer._k_scale,
184
- layer._v_scale,
185
- )
186
-
194
if self.kv_cache_dtype.startswith("fp8"):
195
# queries are quantized in the attention layer
196
dtype = FlashAttentionBackend.get_fp8_dtype_for_flashattn(