minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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
co-authored by rainj-me
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
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py
+8 -1
View File
@@ -215,7 +215,14 @@ class DeepseekMHAForwardMixin:
forward_batch.mha_one_shot forward_batch.mha_one_shot
and sum(forward_batch.extend_prefix_lens_cpu) != 0 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 # FP8 path: dequantize NSA-specific FP8 format to BF16
kv_a, k_pe = self._get_mla_kv_buffer_from_fp8_for_nsa(forward_batch) kv_a, k_pe = self._get_mla_kv_buffer_from_fp8_for_nsa(forward_batch)
else: else: