[Bugfix] Fix FlashInfer KV sharing with omitted K/V (#55864)
Signed-off-by: zhicheng Liu <lzc021026@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Zhicheng Liu committed
Sep 11, 2026 at 14:07 UTC
fbf51c70269bb0d8333e14c5788e628e33c72f5d
2 files changed
+14
-18
vllm/model_executor/models/gemma4_mtp.py
+2
-10
@@ -246,16 +246,8 @@ class Gemma4MTPAttention(nn.Module):
246
247
q, _ = self.rotary_emb(positions, q, None)
248
249
- # Attention reads K/V from the target's cache via KV sharing;
250
- # these dummy tensors are never consumed but required by the API.
251
- num_tokens = q.shape[0]
252
- kv_dummy = torch.empty(
253
- num_tokens,
254
- self.num_kv_heads * self.head_dim,
255
- dtype=q.dtype,
256
- device=q.device,
257
- )
258
- attn_output = self.attn(q, kv_dummy, kv_dummy)
249
+ # Attention reads K/V from the target's cache via KV sharing.
250
+ attn_output = self.attn(q, None, None)
251
output, _ = self.o_proj(attn_output)
252
return output
253
vllm/v1/attention/backends/flashinfer.py
+12
-8
@@ -1962,8 +1962,8 @@ class FlashInferImpl(AttentionImpl):
1962
self,
1963
layer: torch.nn.Module,
1964
query: torch.Tensor,
1965
- key: torch.Tensor,
1966
- value: torch.Tensor,
1965
+ key: torch.Tensor | None,
1966
+ value: torch.Tensor | None,
1967
kv_cache: torch.Tensor,
1968
attn_metadata: FlashInferMetadata,
1969
output: torch.Tensor,
@@ -1974,8 +1974,10 @@ class FlashInferImpl(AttentionImpl):
1974
1975
Args:
1976
query: shape = [num_tokens, num_heads, head_size]
1977
- key: shape = [num_tokens, num_kv_heads, head_size]
1978
- value: shape = [num_tokens, num_kv_heads, head_size]
1977
+ key: shape = [num_tokens, num_kv_heads, head_size], or None for a
1978
+ KV-sharing decoder layer.
1979
+ value: shape = [num_tokens, num_kv_heads, head_size], or None for a
1980
+ KV-sharing decoder layer.
1981
kv_cache: [num_blocks, num_kv_heads, block_size, 2*head_size]
1982
attn_metadata: Metadata for attention.
1983
Returns:
@@ -2066,8 +2068,6 @@ class FlashInferImpl(AttentionImpl):
2068
2069
# Inputs and outputs may be padded for CUDA graphs
2070
query = query[:num_actual_tokens]
2069
- key = key[:num_actual_tokens]
2070
- value = value[:num_actual_tokens]
2071
output_padded = output
2072
output = output[:num_actual_tokens]
2073
@@ -2158,6 +2158,10 @@ class FlashInferImpl(AttentionImpl):
2158
prefill_wrapper = attn_metadata.prefill.wrapper
2159
assert prefill_wrapper is not None
2160
if use_dcp:
2161
+ if key is None or value is None:
2162
+ raise NotImplementedError(
2163
+ "FlashInfer DCP prefill does not support KV-sharing layers"
2164
+ )
2165
assert isinstance(prefill_wrapper, BatchDCPPrefillWrapper)
2166
assert prefill_wrapper._context._window_left == self.window_left
2167
assert prefill_wrapper._context._logits_soft_cap == (
@@ -2176,8 +2180,8 @@ class FlashInferImpl(AttentionImpl):
2180
layer,
2181
prefill_query,
2182
kv_cache_tuple,
2179
- key[num_decode_tokens:],
2180
- value[num_decode_tokens:],
2183
+ key[num_decode_tokens:num_actual_tokens],
2184
+ value[num_decode_tokens:num_actual_tokens],
2185
out=output[num_decode_tokens:],
2186
)
2187
else: