From f23a23cc05fca67d296426e4acb9ce69524f26e4 Mon Sep 17 00:00:00 2001
From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
Date: Fri, 20 Feb 2026 22:00:09 +0800
Subject: [PATCH] Fix NSA FP8 KV cache path for both-trtllm MHA one-shot
 (#18931)

Co-authored-by: rainj-me <96632942+rainj-me@users.noreply.github.com>
---
 .../attention_forward_methods/forward_mha.py             | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py
index 6eff360c4..ceb18afc7 100644
--- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py
+++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py
@@ -215,7 +215,14 @@ class DeepseekMHAForwardMixin:
             forward_batch.mha_one_shot
             and sum(forward_batch.extend_prefix_lens_cpu) != 0
         ):
-            if self.use_nsa and self.kv_cache_dtype == "fp8_e4m3":
+            if (
+                self.use_nsa
+                and self.kv_cache_dtype == "fp8_e4m3"
+                and (
+                    not get_global_server_args().nsa_decode_backend == "trtllm"
+                    or not get_global_server_args().nsa_prefill_backend == "trtllm"
+                )
+            ):
                 # FP8 path: dequantize NSA-specific FP8 format to BF16
                 kv_a, k_pe = self._get_mla_kv_buffer_from_fp8_for_nsa(forward_batch)
             else: