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

fix(server_args): keep NSA prefill/decode backends in the same family for fp8

Browse Source 
calculate_mla_kv_cache_dim only packs fp8 (override dim -> nsa_kv_cache_store_fp8=True, required by the flashmla_sparse dequant path) when NEITHER prefill nor decode backend is trtllm. The Blackwell-no-DP guard defaulted the unset side to trtllm, producing a mixed flashmla/trtllm pair that left fp8 KV unpacked and crashed flashmla_sparse ('kv must have dtype bf16'). Now an explicit flashmla choice on either side pulls the unset side to flashmla too.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit ed30312c3d599595e40ae99bc2832159b634cca6)
This commit is contained in:
孙远航
2026-06-08 13:22:34 +00:00
committed by laoyao0822
parent 98fdae3375
commit 134b88b9b4
1 changed files with 30 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
python/sglang/srt/server_args.py
View File

@@ -1498,12 +1498,36 @@ class ServerArgs:
self.nsa_prefill_backend = "tilelang"
self.nsa_decode_backend = "tilelang"
elif kv_cache_dtype == "fp8_e4m3":
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."
)
# trtllm and flashmla use INCOMPATIBLE MLA fp8 KV cache layouts:
# calculate_mla_kv_cache_dim() only packs fp8 (override dim ->
# nsa_kv_cache_store_fp8=True, which the flashmla_sparse dequant path
# requires) when NEITHER prefill nor decode backend is trtllm. So the
# prefill/decode pair must stay in the same family -- a mixed
# flashmla/trtllm pair leaves fp8 KV unpacked and crashes flashmla_sparse
# ("kv must be bf16"). If the user explicitly picks flashmla for either
# side (e.g. --nsa-prefill-backend flashmla_auto under CP), pull the
# unset side to flashmla too instead of defaulting it to trtllm.
user_wants_flashmla = (
(user_set_prefill and self.nsa_prefill_backend != "trtllm")
or (user_set_decode and self.nsa_decode_backend != "trtllm")
)
if self.dp_size == 1 and major >= 10 and not user_wants_flashmla:
# Default to trtllm on Blackwell without DP attention. The "no DP
# attention" limitation does NOT apply to context-parallel prefill
# (flashmla_sparse pads heads 64->128 and works under CP), which is
# why an explicit flashmla choice above takes this out of the trtllm
# branch entirely.
if not user_set_prefill:
self.nsa_prefill_backend = "trtllm"
if not user_set_decode:
self.nsa_decode_backend = "trtllm"
if not user_set_prefill or not user_set_decode:
logger.warning(
"Flashmla is not supported on Blackwell device without DP "
"attention. Defaulting unset NSA prefill/decode backends to "
"trtllm (fast, slightly less accurate); pass "
"--nsa-prefill-backend/--nsa-decode-backend to override."
)
else:
# flashmla_auto dispatches to flashmla_sparse/flashmla_kv based on hardware and heuristics
if not user_set_prefill: