From 134b88b9b40321a580e34228665f0ef3c0ae63e1 Mon Sep 17 00:00:00 2001
From: leavelet <leavelet@163.com>
Date: Mon, 8 Jun 2026 13:22:34 +0000
Subject: [PATCH] fix(server_args): keep NSA prefill/decode backends in the
 same family for fp8

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)
---
 python/sglang/srt/server_args.py | 36 ++++++++++++++++++++++++++------
 1 file changed, 30 insertions(+), 6 deletions(-)

diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index ddfa97df6..aafc7d337 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -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: