chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
YAMY
2025-12-01 13:18:18 -08:00
committed by GitHub
parent e7d6027e4a
commit 1f2b84d28d
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/nsa_backend.py
+3 -2
View File
@@ -1442,8 +1442,9 @@ class NativeSparseAttnBackend(AttentionBackend):
# Requirements: H200/B200, short sequences, supported dtype, fits in chunk # Requirements: H200/B200, short sequences, supported dtype, fits in chunk
self.use_mha = ( 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 max_kv_len <= self.nsa_index_topk # Short enough for MHA
and forward_batch.token_to_kv_pool.dtype and forward_batch.token_to_kv_pool.dtype
in [torch.bfloat16, torch.float8_e4m3fn] in [torch.bfloat16, torch.float8_e4m3fn]