sparse_attn_indexer.patch 667 B

1234567891011121314151617
  1. --- a/vllm/sparse_attn_indexer.py
  2. +++ b/vllm/sparse_attn_indexer.py
  3. @@ -622,10 +622,10 @@
  4. and current_platform.has_device_capability(90)
  5. and not current_platform.is_device_capability_family(120)
  6. )
  7. - use_persistent_topk = current_platform.is_cuda() and topk_tokens in (
  8. - 512,
  9. - 1024,
  10. - 2048,
  11. + use_persistent_topk = (
  12. + current_platform.is_cuda()
  13. + and topk_tokens in (512, 1024, 2048)
  14. + and not current_platform.is_device_capability_family(120)
  15. )
  16. if use_cooperative_topk:
  17. workspace_manager = current_workspace_manager()