diff --git a/vllm/model_executor/layers/sparse_attn_indexer_kpool.py b/vllm/model_executor/layers/sparse_attn_indexer_kpool.py index 2a314042deef..74060f1f88a0 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer_kpool.py +++ b/vllm/model_executor/layers/sparse_attn_indexer_kpool.py @@ -792,6 +792,7 @@ def sparse_attn_indexer_kpool( decode_metadata.schedule_metadata, max_model_len=max_model_len, clean_logits=False, + indices=decode_metadata.indices, ) num_rows = logits.shape[0] # kpool: logits are pool-granular -> select topk_tokens//kpool pools,