Fix dp run error with fp8-kv enable in high concurrency test (#15241)

Co-authored-by: wunhuang <wunhuang@amd.com>
2025-12-17 11:49:18 -08:00
parent 8451e22758
commit 5d087891c9
1 changed files with 1 additions and 7 deletions
@@ -411,13 +411,7 @@ def handle_attention_trtllm_mla(attn, forward_batch):

 def handle_attention_aiter(attn, forward_batch):
    if forward_batch.forward_mode.is_extend_without_speculative():
-        if is_dp_attention_enabled():
-            if sum(forward_batch.extend_prefix_lens_cpu) == 0:
-                return AttnForwardMethod.MHA
-            else:
-                return AttnForwardMethod.MLA
-        else:
-            return AttnForwardMethod.MHA
+        return AttnForwardMethod.MHA
    else:
        return AttnForwardMethod.MLA