[Bugfix][DCP] Handle sparse MLA metadata after DCP Manager refactor (#52377)
Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com>
cjackal committed
Aug 25, 2026 at 00:02 UTC
22099afc6423efa2c6dac58ecaf0c5cbe652ece2
2 files changed
+62
-11
vllm/model_executor/layers/attention/mla_attention.py
+15
-9
@@ -824,6 +824,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
824
)
825
num_mqa_tokens = attn_metadata.num_decode_tokens
826
num_mha_tokens = q.size(0) - num_mqa_tokens
827
+ use_mha = True
828
829
if self.impl.is_sparse and num_mha_tokens > 0:
830
prefill = getattr(attn_metadata, "prefill", None)
@@ -992,16 +993,21 @@ class MLAAttention(nn.Module, AttentionLayerBase):
993
if self.impl.dcp_world_size > 1:
994
assert lse is not None
995
assert self.dcp_manager is not None
995
- seq_lens = (
996
- attn_metadata.decode.seq_lens
997
- if attn_metadata.decode is not None
998
- else cast(torch.Tensor, attn_metadata.seq_lens)[ # type: ignore[attr-defined]
999
- : attn_metadata.num_decodes
996
+ decode_metadata = getattr(attn_metadata, "decode", None)
997
+ if not use_mha:
998
+ seq_lens = cast(torch.Tensor, attn_metadata.seq_lens) # type: ignore[attr-defined]
999
+ query_start_loc = attn_metadata.query_start_loc
1000
+ else:
1001
+ seq_lens = (
1002
+ decode_metadata.seq_lens
1003
+ if decode_metadata is not None
1004
+ else cast(torch.Tensor, attn_metadata.seq_lens)[ # type: ignore[attr-defined]
1005
+ : attn_metadata.num_decodes
1006
+ ]
1007
+ )
1008
+ query_start_loc = attn_metadata.query_start_loc[
1009
+ : attn_metadata.num_decodes + 1
1010
]
1001
- )
1002
- query_start_loc = attn_metadata.query_start_loc[
1003
- : attn_metadata.num_decodes + 1
1004
- ]
1011
attn_out = self.dcp_manager.combine(
1012
attn_out,
1013
lse,
vllm/v1/attention/ops/dcp.py
+47
-2
@@ -1282,7 +1282,8 @@ class MLADCPManager:
1282
if direct_workspace is not None:
1283
logger.info_once("Using direct symmetric-memory DCP A2A for MLA.")
1284
return functools.partial(
1285
- direct_workspace.lse_reduce,
1285
+ self._direct_workspace_combine,
1286
+ direct_workspace,
1287
is_lse_base_on_e=is_lse_base_on_e,
1288
)
1289
@@ -1299,6 +1300,35 @@ class MLADCPManager:
1300
is_lse_base_on_e=is_lse_base_on_e,
1301
)
1302
1303
+ def _direct_workspace_combine(
1304
+ self,
1305
+ direct_workspace: DirectDCPA2AWorkspace,
1306
+ partial_output: torch.Tensor,
1307
+ partial_lse: torch.Tensor,
1308
+ is_lse_base_on_e: bool,
1309
+ seq_lens: torch.Tensor | None = None,
1310
+ query_start_loc: torch.Tensor | None = None,
1311
+ ) -> torch.Tensor:
1312
+ # Forced MQA path pass all batch tokens (including prefill) into combine,
1313
+ # which may exceed the direct symmetric-memory workspace. Fall back to
1314
+ # the nccl a2a combine for those cases.
1315
+ if partial_output.shape[0] <= direct_workspace.max_num_tokens:
1316
+ return direct_workspace.lse_reduce(
1317
+ partial_output,
1318
+ partial_lse,
1319
+ is_lse_base_on_e,
1320
+ seq_lens,
1321
+ query_start_loc,
1322
+ )
1323
+ return dcp_a2a_lse_reduce(
1324
+ partial_output,
1325
+ partial_lse,
1326
+ cp_group=self.group,
1327
+ is_lse_base_on_e=is_lse_base_on_e,
1328
+ seq_lens=seq_lens,
1329
+ query_start_loc=query_start_loc,
1330
+ )
1331
+
1332
def _init_query_gather(
1333
self,
1334
num_heads: int,
@@ -1317,7 +1347,10 @@ class MLADCPManager:
1347
)
1348
if direct_workspace is not None:
1349
logger.info_once("Using direct symmetric-memory DCP query gather for MLA.")
1320
- return direct_workspace.gather
1350
+ return functools.partial(
1351
+ self._direct_workspace_query_gather,
1352
+ direct_workspace,
1353
+ )
1354
return self._gather_query
1355
1356
def _gather_query(self, query: torch.Tensor) -> torch.Tensor:
@@ -1326,6 +1359,18 @@ class MLADCPManager:
1359
query = reserve_query_head_storage(query, self.padded_num_heads)
1360
return query
1361
1362
+ def _direct_workspace_query_gather(
1363
+ self,
1364
+ direct_workspace: DirectDCPQGatherWorkspace,
1365
+ query: torch.Tensor,
1366
+ ) -> torch.Tensor:
1367
+ # Forced MQA path can be taken for long sparse MLA prefills, whose batch size
1368
+ # may exceed the direct symmetric-memory workspace. Fall back to allgather
1369
+ # for those cases.
1370
+ if query.shape[0] <= direct_workspace.max_num_tokens:
1371
+ return direct_workspace.gather(query)
1372
+ return self._gather_query(query)
1373
+
1374
def init_kv_gather(
1375
self,
1376
workspace: torch.Tensor,