Support DP MLA (#1970)

2024-11-16 09:01:43 +00:00
parent 2f2e07439c
commit 976bc302e5
12 changed files with 395 additions and 63 deletions
@@ -28,9 +28,13 @@ class TritonAttnBackend(AttentionBackend):

        self.decode_attention_fwd = decode_attention_fwd
        self.extend_attention_fwd = extend_attention_fwd
-        self.num_head = (
-            model_runner.model_config.num_attention_heads // model_runner.tp_size
-        )
+
+        if model_runner.server_args.enable_dp_attention:
+            self.num_head = model_runner.model_config.num_attention_heads
+        else:
+            self.num_head = (
+                model_runner.model_config.num_attention_heads // model_runner.tp_size
+            )

        if global_server_args_dict.get("triton_attention_reduce_in_fp32", False):
            self.reduce_dtype = torch.float32