fa MLA cp support

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Author: LucasWilkinson

vllm/v1/attention/backends/mla/flashattn_mla.py

@@ -172,6 +172,7 @@ def _build_decode(self, block_table_tensor: torch.Tensor,
 
 
 class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):
+    can_return_lse_for_decode: bool = True
 
     def __init__(
             self,
@@ -239,7 +240,7 @@ def _forward_decode(
         # to prevent invalid grid configuration during graph capture.
         max_seqlen_q = max(attn_metadata.decode.max_query_len, 1)
 
-        o = flash_attn_varlen_func(
+        attn_out = flash_attn_varlen_func(
             q=q_pe,
             k=k_pe_cache.unsqueeze(-2),  # Add head dim of 1
             v=kv_c_cache.unsqueeze(-2),  # Add head dim of 1
@@ -251,9 +252,15 @@ def _forward_decode(
             block_table=attn_metadata.decode.block_table,
             softmax_scale=self.scale,
             causal=True,
+            return_softmax_lse=self.need_to_return_lse_for_decode,
             fa_version=3,  # only version 3 is supported
             scheduler_metadata=attn_metadata.decode.scheduler_metadata,
             num_splits=attn_metadata.decode.max_num_splits,
         )
-
-        return self._v_up_proj(o)
+        
+        if self.need_to_return_lse_for_decode:
+            o, lse = attn_out
+            return o, lse
+        else:
+            o = attn_out
+            return o, None

vllm/v1/worker/gpu_model_runner.py

@@ -440,9 +440,6 @@ def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None:
             return
 
         if self.reorder_batch_threshold is not None:
-            if self.dcp_world_size > 1:
-                assert self.reorder_batch_threshold == 1, \
-                    "DCP not support reorder_batch_threshold > 1 now."
             reorder_batch_to_split_decodes_and_prefills(
                 self.input_batch,
                 scheduler_output,