Fix NSA Bug in Centralize NSA Dispatch Logic (#14245)
This commit is contained in:
@@ -1442,8 +1442,9 @@ class NativeSparseAttnBackend(AttentionBackend):
|
||||
|
||||
# Requirements: H200/B200, short sequences, supported dtype, fits in chunk
|
||||
self.use_mha = (
|
||||
device_sm == 90
|
||||
or (device_sm >= 100 and device_sm < 110) # SM90/SM100f only
|
||||
(
|
||||
device_sm == 90 or (device_sm >= 100 and device_sm < 110)
|
||||
) # SM90/SM100 only
|
||||
and max_kv_len <= self.nsa_index_topk # Short enough for MHA
|
||||
and forward_batch.token_to_kv_pool.dtype
|
||||
in [torch.bfloat16, torch.float8_e4m3fn]
|
||||
|
||||
Reference in New Issue
Block a user