diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 678b66eab8..ceb52e5d32 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -578,6 +578,7 @@ def sparse_attn_indexer( and num_rows <= 32 and logits.stride(0) % 4 == 0 # TMA 16-byte alignment and current_platform.has_device_capability(90) + and not current_platform.is_device_capability_family(120) ) use_persistent_topk = current_platform.is_cuda() and topk_tokens in ( 512,