--- a/vllm/sparse_attn_indexer_kpool.py +++ b/vllm/sparse_attn_indexer_kpool.py @@ -807,7 +807,11 @@ else: topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens] - if current_platform.is_cuda() and select_k in (512, 1024, 2048): + if ( + current_platform.is_cuda() + and select_k in (512, 1024, 2048) + and not current_platform.is_device_capability_family(120) + ): workspace_manager = current_workspace_manager() (topk_workspace,) = workspace_manager.get_simultaneous( ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),