[Bugfix][DCP] Handle sparse MLA metadata after DCP Manager refactor (#52377)

Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com>

vllm/model_executor/layers/attention/mla_attention.py

@@ -824,6 +824,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
         )
         num_mqa_tokens = attn_metadata.num_decode_tokens
         num_mha_tokens = q.size(0) - num_mqa_tokens
+        use_mha = True
 
         if self.impl.is_sparse and num_mha_tokens > 0:
             prefill = getattr(attn_metadata, "prefill", None)
@@ -992,16 +993,21 @@ class MLAAttention(nn.Module, AttentionLayerBase):
             if self.impl.dcp_world_size > 1:
                 assert lse is not None
                 assert self.dcp_manager is not None
-                seq_lens = (
-                    attn_metadata.decode.seq_lens
-                    if attn_metadata.decode is not None
-                    else cast(torch.Tensor, attn_metadata.seq_lens)[  # type: ignore[attr-defined]
-                        : attn_metadata.num_decodes
+                decode_metadata = getattr(attn_metadata, "decode", None)
+                if not use_mha:
+                    seq_lens = cast(torch.Tensor, attn_metadata.seq_lens)  # type: ignore[attr-defined]
+                    query_start_loc = attn_metadata.query_start_loc
+                else:
+                    seq_lens = (
+                        decode_metadata.seq_lens
+                        if decode_metadata is not None
+                        else cast(torch.Tensor, attn_metadata.seq_lens)[  # type: ignore[attr-defined]
+                            : attn_metadata.num_decodes
+                        ]
+                    )
+                    query_start_loc = attn_metadata.query_start_loc[
+                        : attn_metadata.num_decodes + 1
                     ]
-                )
-                query_start_loc = attn_metadata.query_start_loc[
-                    : attn_metadata.num_decodes + 1
-                ]
                 attn_out = self.dcp_manager.combine(
                     attn_out,
                     lse,

vllm/v1/attention/ops/dcp.py

@@ -1282,7 +1282,8 @@ class MLADCPManager:
         if direct_workspace is not None:
             logger.info_once("Using direct symmetric-memory DCP A2A for MLA.")
             return functools.partial(
-                direct_workspace.lse_reduce,
+                self._direct_workspace_combine,
+                direct_workspace,
                 is_lse_base_on_e=is_lse_base_on_e,
             )
 
@@ -1299,6 +1300,35 @@ class MLADCPManager:
             is_lse_base_on_e=is_lse_base_on_e,
         )
 
+    def _direct_workspace_combine(
+        self,
+        direct_workspace: DirectDCPA2AWorkspace,
+        partial_output: torch.Tensor,
+        partial_lse: torch.Tensor,
+        is_lse_base_on_e: bool,
+        seq_lens: torch.Tensor | None = None,
+        query_start_loc: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # Forced MQA path pass all batch tokens (including prefill) into combine,
+        # which may exceed the direct symmetric-memory workspace. Fall back to
+        # the nccl a2a combine for those cases.
+        if partial_output.shape[0] <= direct_workspace.max_num_tokens:
+            return direct_workspace.lse_reduce(
+                partial_output,
+                partial_lse,
+                is_lse_base_on_e,
+                seq_lens,
+                query_start_loc,
+            )
+        return dcp_a2a_lse_reduce(
+            partial_output,
+            partial_lse,
+            cp_group=self.group,
+            is_lse_base_on_e=is_lse_base_on_e,
+            seq_lens=seq_lens,
+            query_start_loc=query_start_loc,
+        )
+
     def _init_query_gather(
         self,
         num_heads: int,
@@ -1317,7 +1347,10 @@ class MLADCPManager:
         )
         if direct_workspace is not None:
             logger.info_once("Using direct symmetric-memory DCP query gather for MLA.")
-            return direct_workspace.gather
+            return functools.partial(
+                self._direct_workspace_query_gather,
+                direct_workspace,
+            )
         return self._gather_query
 
     def _gather_query(self, query: torch.Tensor) -> torch.Tensor:
@@ -1326,6 +1359,18 @@ class MLADCPManager:
             query = reserve_query_head_storage(query, self.padded_num_heads)
         return query
 
+    def _direct_workspace_query_gather(
+        self,
+        direct_workspace: DirectDCPQGatherWorkspace,
+        query: torch.Tensor,
+    ) -> torch.Tensor:
+        # Forced MQA path can be taken for long sparse MLA prefills, whose batch size
+        # may exceed the direct symmetric-memory workspace. Fall back to allgather
+        # for those cases.
+        if query.shape[0] <= direct_workspace.max_num_tokens:
+            return direct_workspace.gather(query)
+        return self._gather_query(query)
+
     def init_kv_gather(
         self,
         workspace: torch.Tensor,