[DSA] Set trtllm kernels as nsa default for Blackwell (#21914)
This commit is contained in:
@@ -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:
|
||||||
|
|||||||
Reference in New Issue
Block a user