Nsa trtllm mla sparse fp8 support with Deepseek v3.2 NVFP4 (#18389)

2026-02-16 09:29:54 +08:00
parent 8290171f52
commit 0ffd0a3995
10 changed files with 352 additions and 183 deletions
@@ -1493,6 +1493,12 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin):
        """
        Check if we should skip rope and do fused rope+quantize for TRTLLM MLA decode in fp8_e4m3 path.
        """
+        if self.current_attention_backend == "nsa":
+            return (
+                get_global_server_args().nsa_decode_backend == "trtllm"
+                or get_global_server_args().nsa_prefill_backend == "trtllm"
+            ) and forward_batch.attn_backend.kv_cache_dtype == torch.float8_e4m3fn
+
        return (
            self.current_attention_backend == "trtllm_mla"
            and (