diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 2885a7562..a0ce0a435 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1279,7 +1279,9 @@ class ServerArgs: ) if self.kv_cache_dtype == "auto": - self.kv_cache_dtype = "fp8_e4m3" if major >= 10 else "bfloat16" + self.kv_cache_dtype = ( + "fp8_e4m3" if (major >= 10 and self.dp_size > 1) else "bfloat16" + ) logger.warning( f"Setting KV cache dtype to {self.kv_cache_dtype} for DeepSeek DSA on SM{major} device." ) @@ -1298,11 +1300,18 @@ class ServerArgs: self.nsa_prefill_backend = "tilelang" self.nsa_decode_backend = "tilelang" elif kv_cache_dtype == "fp8_e4m3": - # flashmla_auto dispatches to flashmla_sparse/flashmla_kv based on hardware and heuristics - if not user_set_prefill: - self.nsa_prefill_backend = "flashmla_auto" - if not user_set_decode: - self.nsa_decode_backend = "flashmla_kv" + if self.dp_size == 1 and major >= 10: + self.nsa_prefill_backend = "trtllm" + self.nsa_decode_backend = "trtllm" + logger.warning( + "Flashmla is not supported on Blackwell device without DP attention. Set NSA prefill/decode backends to trtllm, which runs fast but loses a little accuracy." + ) + else: + # flashmla_auto dispatches to flashmla_sparse/flashmla_kv based on hardware and heuristics + if not user_set_prefill: + self.nsa_prefill_backend = "flashmla_auto" + if not user_set_decode: + self.nsa_decode_backend = "flashmla_kv" else: # set prefill/decode backends based on hardware architecture. if major >= 10: