Fuse quantize and rope in trtllm_mla MTP (#10779)

2025-10-02 17:59:37 +08:00
parent d61615fe93
commit f35def8652
2 changed files with 37 additions and 5 deletions
--- a/python/sglang/srt/models/deepseek_v2.py
+++ b/python/sglang/srt/models/deepseek_v2.py
@@ -1399,7 +1399,10 @@ class DeepseekV2AttentionMLA(nn.Module):
        """
        return (
            self.current_attention_backend == "trtllm_mla"
-            and forward_batch.forward_mode.is_decode_or_idle()
+            and (
+                forward_batch.forward_mode.is_decode_or_idle()
+                or forward_batch.forward_mode.is_target_verify()
+            )
            and forward_batch.attn_backend.data_type == torch.float8_e4m3fn
        )