[Bugfix][DSv4] Keep indexer scoring in breakable graphs (#52492)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
Lucas Wilkinson committed
Aug 16, 2026 at 23:01 UTC
292187dd8ca1b1bfa195f25b2886262527269999
1 file changed
+4
-1
vllm/models/deepseek_v4/attention.py
+4
-1
@@ -906,7 +906,10 @@ class DeepseekV4Indexer(nn.Module):
906
attn_metadata = get_forward_context().attn_metadata
907
if isinstance(attn_metadata, dict):
908
indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix])
909
- if indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens:
909
+ if (
910
+ indexer_metadata.max_seq_len // self.compress_ratio <= self.topk_tokens
911
+ and not torch.cuda.is_current_stream_capturing()
912
+ ):
913
# candidates num smaller than topk, every candidate is selected
914
# but we still need to build k cache
915
compressor(compressed_kv_score, positions, rotary_emb)