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

Fix NSA FP8 KV cache path for both-trtllm MHA one-shot (#18931)

Browse Source 
Co-authored-by: rainj-me <96632942+rainj-me@users.noreply.github.com>
This commit is contained in:
Mohammad Miadh Angkad
2026-02-20 22:00:09 +08:00
committed by GitHub
parent 8d789b5c3d
commit f23a23cc05
1 changed files with 8 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py
View File

@@ -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: