From a4debbd5bd0608999314ad68e6463fb1d76f5d4e Mon Sep 17 00:00:00 2001 From: khluu Date: Sat, 2 May 2026 00:52:06 -0700 Subject: [PATCH] Revert "Temporary disable persistent topk (#41442)" This reverts commit 05ebca525009a4fe3dc89ff53de6469cb2ac0800. --- vllm/model_executor/layers/sparse_attn_indexer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 6bca26c4b9f..ca82f2feb7e 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -320,7 +320,7 @@ def sparse_attn_indexer( num_rows = logits.shape[0] topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens] - if current_platform.is_cuda() and topk_tokens in (512, 2048): + if current_platform.is_cuda() and topk_tokens in (512, 1024, 2048): workspace_manager = current_workspace_manager() (topk_workspace,) = workspace_manager.get_simultaneous( ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),