[DSA] Set trtllm kernels as nsa default for Blackwell (#21914)

This commit is contained in:
Baizhou Zhang
2026-04-02 00:22:27 -07:00
committed by GitHub
parent f30df723bf
commit fbc1f92453
+2 -7
View File
@@ -1424,15 +1424,10 @@ class ServerArgs:
) )
if self.kv_cache_dtype == "auto": if self.kv_cache_dtype == "auto":
# TODO: Temporarily set default dtype on B200 as bfloat16 to avoid performance regression. if major >= 10:
# TODO: Remove this after the performance regression is fixed. (Ref: https://github.com/sgl-project/sglang/issues/21291)
if quantization == "modelopt_fp4" and major >= 10 and self.dp_size > 1:
self.kv_cache_dtype = "fp8_e4m3" self.kv_cache_dtype = "fp8_e4m3"
else: else:
self.kv_cache_dtype = "bfloat16" self.kv_cache_dtype = "bfloat16"
# self.kv_cache_dtype = (
# "fp8_e4m3" if (major >= 10 and self.dp_size > 1) else "bfloat16"
# )
logger.warning( logger.warning(
f"Setting KV cache dtype to {self.kv_cache_dtype} for DeepSeek DSA on SM{major} device." f"Setting KV cache dtype to {self.kv_cache_dtype} for DeepSeek DSA on SM{major} device."
) )
@@ -1463,7 +1458,7 @@ class ServerArgs:
self.nsa_prefill_backend = "tilelang" self.nsa_prefill_backend = "tilelang"
self.nsa_decode_backend = "tilelang" self.nsa_decode_backend = "tilelang"
elif kv_cache_dtype == "fp8_e4m3": elif kv_cache_dtype == "fp8_e4m3":
if self.dp_size == 1 and major >= 10: if major >= 10:
self.nsa_prefill_backend = "trtllm" self.nsa_prefill_backend = "trtllm"
self.nsa_decode_backend = "trtllm" self.nsa_decode_backend = "trtllm"
else: else: