fix(mla): keep DCP outputs head-major for safe BMM

tests/distributed/test_dcp_a2a.py

@@ -548,7 +548,9 @@ def test_b12x_lse_reduce_honors_token_cap(monkeypatch: pytest.MonkeyPatch):
     sentinel = torch.zeros(1)
 
     class _FakePool:
-        def lse_reduce_scatter(self, out, lse, *, is_lse_base_on_e):
+        def lse_reduce_scatter(
+            self, partial, lse, out=None, *, is_lse_base_on_e
+        ):
             return sentinel
 
     def fake_get_pool(
@@ -633,17 +635,19 @@ def fake_get_pool(
 
 
 @pytest.mark.skipif(torch.accelerator.device_count() < 1, reason="CUDA is required.")
-def test_b12x_lse_reduce_makes_views_contiguous(monkeypatch: pytest.MonkeyPatch):
-    """Head-sliced attention views must reach the PCIe pool contiguous."""
+def test_b12x_lse_reduce_preserves_supported_layouts(monkeypatch: pytest.MonkeyPatch):
+    """Preserve head-major input while materializing legacy head slices."""
     from vllm.v1.attention.ops import dcp_alltoall
 
     monkeypatch.setenv("VLLM_USE_B12X_DCP_A2A", "1")
     received: dict[str, Any] = {}
     sentinel = torch.zeros(1)
 
     class _FakePool:
-        def lse_reduce_scatter(self, out, lse, *, is_lse_base_on_e):
-            received.update(out=out, lse=lse)
+        def lse_reduce_scatter(
+            self, partial, lse, out=None, *, is_lse_base_on_e
+        ):
+            received.update(partial=partial, lse=lse, out=out)
             return sentinel
 
     monkeypatch.setattr(
@@ -671,8 +675,26 @@ def lse_reduce_scatter(self, out, lse, *, is_lse_base_on_e):
         query_head_dim=64,
     )
     assert result is sentinel
-    assert received["out"].is_contiguous()
+    assert received["partial"].is_contiguous()
     assert received["lse"].is_contiguous()
+    assert received["out"].movedim(0, 1).is_contiguous()
+
+    head_major_storage = torch.zeros(
+        16, 8, 64, dtype=torch.bfloat16, device="cuda"
+    )
+    head_major = head_major_storage.transpose(0, 1)[:4]
+    result = dcp_alltoall._try_b12x_dcp_lse_reduce(
+        head_major,
+        torch.zeros(4, 16, dtype=torch.float32, device="cuda"),
+        group,  # type: ignore[arg-type]
+        return_lse=False,
+        is_lse_base_on_e=True,
+        max_batch_size=8192,
+        query_head_dim=64,
+    )
+    assert result is sentinel
+    assert received["partial"] is head_major
+    assert received["out"].stride() == (64, 4 * 64, 1)
 
 
 def test_b12x_query_gather_requires_env(monkeypatch: pytest.MonkeyPatch):
@@ -778,7 +800,44 @@ def test_pack_unpack_combine_matches_reference(
             _assert_packed_a2a_close(actual_out, expected_out, dtype)
             torch.testing.assert_close(actual_lse, expected_lse, rtol=1e-4, atol=1e-4)
         else:
+            actual_out = actual
             _assert_packed_a2a_close(actual, expected_out, dtype)
+        assert actual_out.movedim(0, 1).is_contiguous()
+        assert not actual_out.is_contiguous()
+
+
+def test_cuda_reduce_scatter_can_preserve_head_major_output(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    from vllm.distributed.device_communicators import cuda_communicator
+
+    monkeypatch.setattr(
+        cuda_communicator,
+        "should_nccl_symm_mem_ag_rs",
+        lambda: False,
+    )
+
+    class FakePyNccl:
+        disabled = False
+
+        def reduce_scatter(self, output, input_):
+            output.copy_(input_[: output.shape[0]])
+
+    class FakeCommunicator:
+        world_size = 2
+        pynccl_comm = FakePyNccl()
+
+    input_storage = torch.arange(8 * 3 * 16, dtype=torch.bfloat16).view(8, 3, 16)
+    input_ = input_storage.movedim(0, 1)
+    actual = cuda_communicator.CudaCommunicator.reduce_scatter_head_major(
+        FakeCommunicator(), input_, dim=1
+    )
+
+    expected = input_[:, :4]
+    torch.testing.assert_close(actual, expected)
+    assert actual.shape == (3, 4, 16)
+    assert actual.stride() == (16, 3 * 16, 1)
+    assert actual.movedim(0, 1).is_contiguous()
 
 
 def _distributed_packed_a2a_worker(env: dict[str, str]) -> None:

tests/v1/attention/test_b12x_mla_dcp_workspace.py

@@ -102,6 +102,42 @@ def reduce_scatter_into(self, input_, output, dim):
     assert torch.equal(lse, corrected_lse[:, rank : rank + 1])
 
 
+def test_cp_lse_ag_out_rs_requests_head_major_output(monkeypatch):
+    corrected_storage = torch.arange(8 * 3 * 16, dtype=torch.bfloat16).view(
+        8, 3, 16
+    )
+    corrected = corrected_storage.movedim(0, 1)
+    corrected_lse = torch.zeros(3, 8, dtype=torch.float32)
+
+    monkeypatch.setattr(
+        common,
+        "_cp_lse_common",
+        lambda *args, **kwargs: (corrected, corrected_lse),
+    )
+
+    class FakeGroup:
+        rank_in_group = 0
+        world_size = 2
+
+        def reduce_scatter_head_major(self, input_, dim):
+            assert input_ is corrected
+            assert dim == 1
+            storage = torch.empty(4, 3, 16, dtype=input_.dtype)
+            output = storage.movedim(0, 1)
+            output.copy_(input_[:, :4])
+            return output
+
+    output = common.cp_lse_ag_out_rs(
+        corrected,
+        corrected_lse,
+        FakeGroup(),
+        head_major_output=True,
+    )
+
+    assert output.stride() == (16, 3 * 16, 1)
+    torch.testing.assert_close(output, corrected[:, :4])
+
+
 @pytest.mark.parametrize(
     ("tp_size", "dcp_size", "local_heads", "input_heads", "kernel_heads"),
     [

vllm/distributed/device_communicators/cuda_communicator.py

@@ -467,6 +467,46 @@ def reduce_scatter_into(
         )
         return output
 
+    def reduce_scatter_head_major(
+        self,
+        input_: torch.Tensor,
+        dim: int = -1,
+    ) -> torch.Tensor:
+        """Reduce-scatter heads without the final token-major copy."""
+        if self.world_size <= 1:
+            raise RuntimeError("head-major reduce-scatter requires world size > 1")
+        if dim < 0:
+            dim += input_.dim()
+        if dim != 1 or input_.ndim != 3:
+            raise ValueError(
+                "head-major reduce-scatter requires a rank-3 tensor on dim 1"
+            )
+
+        input_head_major = input_.movedim(0, dim).contiguous()
+        if input_head_major.shape[0] % self.world_size != 0:
+            raise ValueError("head count is not divisible by world size")
+        output_shape = (
+            input_head_major.shape[0] // self.world_size,
+            *input_head_major.shape[1:],
+        )
+
+        if should_nccl_symm_mem_ag_rs():
+            output_head_major = self._reduce_scatter_symm_mem(input_head_major)
+        else:
+            output_head_major = torch.empty(
+                output_shape,
+                dtype=input_head_major.dtype,
+                device=input_head_major.device,
+            )
+            pynccl_comm = self.pynccl_comm
+            if pynccl_comm is None or pynccl_comm.disabled:
+                raise RuntimeError(
+                    "head-major reduce-scatter requires an active PyNccl communicator"
+                )
+            pynccl_comm.reduce_scatter(output_head_major, input_head_major)
+
+        return output_head_major.movedim(0, dim)
+
     def reduce_scatterv(
         self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None
     ):

vllm/distributed/parallel_state.py

@@ -751,6 +751,30 @@ def reduce_scatter_into(
             raise RuntimeError("reduce_scatter_into did not preserve output identity")
         return output
 
+    def reduce_scatter_head_major(
+        self,
+        input_: torch.Tensor,
+        dim: int = -1,
+    ) -> torch.Tensor:
+        """Reduce-scatter and preserve a physically head-major output view."""
+        if self.world_size <= 1 or dim != 1:
+            raise RuntimeError(
+                "reduce_scatter_head_major requires DCP heads on dim 1"
+            )
+        if self.device_communicator is None:
+            raise RuntimeError(
+                "reduce_scatter_head_major requires a device communicator"
+            )
+        reduce_scatter_head_major = getattr(
+            self.device_communicator, "reduce_scatter_head_major", None
+        )
+        if not callable(reduce_scatter_head_major):
+            raise RuntimeError(
+                f"{type(self.device_communicator).__name__} does not support "
+                "head-major reduce-scatter"
+            )
+        return reduce_scatter_head_major(input_, dim)
+
     def reduce_scatterv(
         self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None
     ) -> torch.Tensor:

vllm/model_executor/layers/attention/mla_attention.py

@@ -1206,6 +1206,7 @@ def forward_impl(
                             lse,
                             get_dcp_group(),
                             is_lse_base_on_e=self.impl.lse_base_on_e,
+                            head_major_output=True,
                         )
 
             if project_before_merge:

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

@@ -976,10 +976,6 @@ def __init__(
         self.v_head_dim: int = mla_args.get("v_head_dim", 512)
         # GLM_NSA contract: q_head_dim = kv_lora_rank (512) + qk_rope (64) = 576.
         self.q_head_dim = self.kv_lora_rank + self.qk_rope_head_dim
-        self.force_contiguous_mla_bmm_input = True
-        self.force_contiguous_mla_bmm_weight = True
-        self.force_contiguous_mla_bmm_output = True
-
         # The indexer carries the shared buffer for normal layers and tests;
         # the explicitly-passed buffer covers backbone skip layers, whose
         # indexer is not constructed (see deepseek_v2.py).
@@ -998,6 +994,7 @@ def __init__(
         parallel_config = vllm_config.parallel_config
         self.dcp_workspace_non_dbo = not bool(parallel_config.enable_dbo)
         self.dcp_world_size = parallel_config.decode_context_parallel_size
+        self._head_major_mla_output = True
         self.tp_world_size = int(parallel_config.tensor_parallel_size)
         self.dcp_rank = 0
         if self.dcp_world_size > 1:
@@ -1138,6 +1135,7 @@ def _make_plan(
                     max_batch=int(max_batch),
                     max_chunks_per_row=self._num_splits_cap,
                     page_size=self.block_size,
+                    head_major_output=self._head_major_mla_output,
                 )
             )
 
@@ -1305,7 +1303,17 @@ def do_kv_cache_update(
         )
 
     def _borrow_workspaces(self) -> list[torch.Tensor]:
-        return current_workspace_manager().get_simultaneous(*self._workspace_specs)
+        workspaces = current_workspace_manager().get_simultaneous(
+            *self._workspace_specs
+        )
+        if self._pad_heads:
+            dense_storage = workspaces[1]
+            workspaces[1] = dense_storage.view(
+                self._input_num_heads,
+                self._max_batched,
+                self.kv_lora_rank,
+            ).transpose(0, 1)
+        return workspaces
 
     def _borrow_workspace_parts(
         self,
@@ -1341,7 +1349,7 @@ def _borrow_workspace_parts(
             != (self._max_batched, self._input_num_heads, self.kv_lora_rank)
             or dense_out_workspace.dtype != torch.bfloat16
             or dense_out_workspace.device != self.device
-            or not dense_out_workspace.is_contiguous()
+            or not dense_out_workspace.movedim(0, 1).is_contiguous()
         ):
             raise RuntimeError("B12X DCP prefill borrowed an invalid dense output")
         if (
@@ -1502,7 +1510,7 @@ def dcp_project_before_merge_in_workspace(
         if (
             tuple(attn_out.shape)
             != (num_tokens, self._input_num_heads, self.kv_lora_rank)
-            or not attn_out.is_contiguous()
+            or not attn_out.movedim(0, 1).is_contiguous()
             or attn_out.dtype != torch.bfloat16
             or tuple(w_uv.shape)
             != (self._input_num_heads, self.kv_lora_rank, self.v_head_dim)

vllm/v1/attention/ops/common.py

@@ -231,6 +231,7 @@ def cp_lse_ag_out_rs(
     ctx: CPTritonContext | None = None,
     return_lse: bool = False,
     is_lse_base_on_e=True,
+    head_major_output: bool = False,
 ):
     """
     cp_attn_out: [ B, H, D ]
@@ -239,7 +240,10 @@ def cp_lse_ag_out_rs(
     out, lse = _cp_lse_common(
         cp_attn_out, cp_attn_lse, cp_group, ctx=ctx, is_lse_base_on_e=is_lse_base_on_e
     )
-    out = cp_group.reduce_scatter(out, dim=1)
+    if head_major_output:
+        out = cp_group.reduce_scatter_head_major(out, dim=1)
+    else:
+        out = cp_group.reduce_scatter(out, dim=1)
 
     if return_lse:
         cp_num_heads = lse.shape[1] // cp_group.world_size

vllm/v1/attention/ops/dcp_alltoall.py

@@ -46,6 +46,21 @@
 ] = {}
 
 
+def _is_supported_bhd_layout(tensor: torch.Tensor) -> bool:
+    """Accept packed token-major or capacity-strided head-major BHD views."""
+    if tensor.ndim != 3 or int(tensor.stride(2)) != 1:
+        return False
+    batch, heads, head_dim = (int(value) for value in tensor.shape)
+    stride_batch, stride_head, _ = (int(value) for value in tensor.stride())
+    packed_token_major = (
+        stride_batch == heads * head_dim and stride_head == head_dim
+    )
+    capacity_strided_head_major = (
+        stride_batch == head_dim and stride_head >= batch * head_dim
+    )
+    return packed_token_major or capacity_strided_head_major
+
+
 @lru_cache(maxsize=1)
 def _load_b12x_dcp_a2a_pool() -> Any | None:
     try:
@@ -216,19 +231,24 @@ def _try_b12x_dcp_lse_reduce(
         )
         return None
 
-    # Sparse MLA backends can return head-sliced views (e.g. GLM TP6 pads
-    # 64 -> 66 heads and slices the kernel output back), and the PCIe pool
-    # requires contiguous operands. The NCCL packers take explicit strides,
-    # so only this fast path needs the copy; LSE is tiny and the output is
-    # already contiguous on unpadded head counts.
-    if not cp_attn_out.is_contiguous():
+    # The channel accepts packed token-major input and the capacity-strided
+    # head-major layout produced by B12X sparse MLA. Preserve either layout;
+    # only legacy padded-head slices need materialization.
+    if not _is_supported_bhd_layout(cp_attn_out):
         cp_attn_out = cp_attn_out.contiguous()
     if not cp_attn_lse.is_contiguous():
         cp_attn_lse = cp_attn_lse.contiguous()
 
+    reduced_storage = torch.empty(
+        (total_heads // world_size, batch, head_dim),
+        device=cp_attn_out.device,
+        dtype=cp_attn_out.dtype,
+    )
+    reduced = reduced_storage.transpose(0, 1)
     return pool.lse_reduce_scatter(
         cp_attn_out,
         cp_attn_lse,
+        out=reduced,
         is_lse_base_on_e=is_lse_base_on_e,
     )
 
@@ -789,11 +809,12 @@ def _dcp_a2a_unpack_combine(
     is_lse_base_on_e: bool,
 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
     world_size, num_tokens, h_per_rank, _ = recv_buffer.shape
-    out = torch.empty(
-        (num_tokens, h_per_rank, head_dim),
+    out_storage = torch.empty(
+        (h_per_rank, num_tokens, head_dim),
         device=recv_buffer.device,
         dtype=recv_buffer.dtype,
     )
+    out = out_storage.transpose(0, 1)
     out_lse = torch.empty(
         (num_tokens, h_per_rank) if return_lse else (1, 1),
         device=recv_buffer.device,