Fix NSA Bug in Centralize NSA Dispatch Logic (#14245)

This commit is contained in:
YAMY
2025-12-01 13:18:18 -08:00
committed by GitHub
parent e7d6027e4a
commit 1f2b84d28d

View File

@@ -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]