| 123456789101112131415 |
- --- 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),
|
