[Core] Add explicit layer parallel plans

Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

tests/distributed/test_layer_parallel.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import FrozenInstanceError
+
+import pytest
+
+from vllm.distributed.layer_parallel import (
+    LayerParallelPlan,
+    LayerType,
+    ParallelAxis,
+    ParallelGroupType,
+    clear_layer_parallel_config,
+    get_layer_parallel_config,
+    init_layer_parallel_config,
+)
+
+
+@pytest.fixture(autouse=True)
+def _reset_layer_parallel_config():
+    clear_layer_parallel_config()
+    yield
+    clear_layer_parallel_config()
+
+
+def _axis(
+    world_size: int,
+    rank: int,
+    group: ParallelGroupType = ParallelGroupType.TENSOR,
+) -> ParallelAxis:
+    return ParallelAxis(world_size=world_size, rank=rank, group=group)
+
+
+def test_default_layer_uses_default_plan():
+    tensor_axis = _axis(8, 5)
+    default_plan = LayerParallelPlan(input=tensor_axis, output=tensor_axis)
+    init_layer_parallel_config(default_plan)
+
+    assert get_layer_parallel_config() == default_plan
+    assert get_layer_parallel_config(LayerType.ATTENTION) == default_plan
+
+
+def test_layer_override_can_reshard_between_input_and_output():
+    tensor_axis = _axis(8, 5)
+    attention_axis = _axis(2, 1, ParallelGroupType.ATTENTION_TENSOR)
+    default_plan = LayerParallelPlan(input=tensor_axis, output=tensor_axis)
+    attention_plan = LayerParallelPlan(
+        input=attention_axis,
+        output=tensor_axis,
+    )
+    init_layer_parallel_config(
+        default_plan,
+        {LayerType.ATTENTION: attention_plan},
+    )
+
+    assert get_layer_parallel_config() == default_plan
+    assert get_layer_parallel_config(LayerType.ATTENTION) == attention_plan
+    assert attention_plan.reshards_output
+    assert attention_plan.get_output_size(input_size=16) == 4
+
+
+def test_plan_rejects_non_divisible_output_size():
+    plan = LayerParallelPlan(
+        input=_axis(2, 0, ParallelGroupType.ATTENTION_TENSOR),
+        output=_axis(8, 0),
+    )
+
+    with pytest.raises(ValueError, match="Global input size"):
+        plan.get_output_size(input_size=3)
+
+
+@pytest.mark.parametrize(
+    "world_size, rank",
+    [
+        (0, 0),
+        (2, -1),
+        (2, 2),
+    ],
+)
+def test_axis_rejects_invalid_size_or_rank(world_size: int, rank: int):
+    with pytest.raises(ValueError):
+        _axis(world_size, rank)
+
+
+def test_plan_is_immutable():
+    axis = _axis(4, 2)
+    plan = LayerParallelPlan(input=axis, output=axis)
+
+    with pytest.raises(FrozenInstanceError):
+        plan.input = _axis(2, 0)

tests/kernels/core/test_minimax_reduce_rms.py

@@ -188,7 +188,7 @@ def test_minimax_qk_norm_triton_fallback(
     ``/ tp_world`` scaling without needing multiple ranks -- ``hidden_dims``
     are the per-rank q/k segment widths.
     """
-    monkeypatch.setattr(rms_norm_tp, "_all_reduce_variance", lambda v: v)
+    monkeypatch.setattr(rms_norm_tp, "_all_reduce_variance", lambda v, *_: v)
 
     q_size, kv_size = hidden_dims
     device = "cuda"

tests/test_config.py

@@ -362,6 +362,55 @@ def test_async_scheduling_with_pipeline_parallelism_is_allowed():
     assert cfg.scheduler_config.async_scheduling is True
 
 
+@pytest.mark.parametrize("attention_tp_size", [1, 2])
+def test_attention_tp_requires_matching_dcp(attention_tp_size: int):
+    tp_size = 4
+    expected_dcp_size = tp_size // attention_tp_size
+
+    config = ParallelConfig(
+        tensor_parallel_size=tp_size,
+        tensor_parallel_size_attention=attention_tp_size,
+        decode_context_parallel_size=expected_dcp_size,
+    )
+
+    assert config.attention_tp_size == attention_tp_size
+
+    with pytest.raises(ValidationError, match="decode_context_parallel_size"):
+        ParallelConfig(
+            tensor_parallel_size=tp_size,
+            tensor_parallel_size_attention=attention_tp_size,
+            decode_context_parallel_size=1,
+        )
+
+
+def test_attention_tp_equal_to_tp_is_noop():
+    config = ParallelConfig(
+        tensor_parallel_size=4,
+        tensor_parallel_size_attention=4,
+        decode_context_parallel_size=2,
+    )
+
+    assert config.attention_tp_size == config.tensor_parallel_size
+
+
+def test_attention_tp_rejects_model_without_layer_plan_support():
+    model_config = SimpleNamespace(
+        model_arch_config=SimpleNamespace(total_num_attention_heads=32),
+        architectures=["UnsupportedForCausalLM"],
+        registry=SimpleNamespace(
+            is_layer_parallel_supported_model=lambda *args: False,
+        ),
+    )
+    parallel_config = ParallelConfig(
+        tensor_parallel_size=4,
+        tensor_parallel_size_attention=2,
+        decode_context_parallel_size=2,
+    )
+
+    with pytest.raises(ValueError, match="ATTENTION layer parallel plan"):
+        ModelConfig.verify_with_parallel_config(model_config, parallel_config)
+
+
 @dataclass
 class _TestConfigFields:
     a: int

vllm/config/model.py

@@ -24,6 +24,7 @@ from vllm.config.pooler import PoolerConfig
 from vllm.config.quantization import QuantizationConfigArgs
 from vllm.config.scheduler import RunnerType
 from vllm.config.utils import config, getattr_iter
+from vllm.distributed.layer_parallel import LayerType
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.tasks import PoolingTask, ScoreType, SupportedTask
@@ -1216,6 +1217,17 @@ class ModelConfig:
                 f"({tensor_parallel_size})."
             )
 
+        attention_tp_size = parallel_config.attention_tp_size
+        if attention_tp_size != tensor_parallel_size and not (
+            self.registry.is_layer_parallel_supported_model(
+                self.architectures, self, LayerType.ATTENTION
+            )
+        ):
+            raise ValueError(
+                "Attention tensor parallelism requires a model that supports "
+                "the ATTENTION layer parallel plan."
+            )
+
         if parallel_config.enable_expert_parallel:
             self._verify_with_expert_parallelism()
 
@@ -1229,7 +1241,11 @@ class ModelConfig:
             )
 
         decode_context_parallel_size = parallel_config.decode_context_parallel_size
-        if decode_context_parallel_size > 1 and not self.use_mla:
+        if (
+            decode_context_parallel_size > 1
+            and not self.use_mla
+            and attention_tp_size == tensor_parallel_size
+        ):
             total_num_kv_heads = self.get_total_num_kv_heads()
             assert tensor_parallel_size > total_num_kv_heads, (
                 f"tensor parallel size {tensor_parallel_size} must be greater "
@@ -1314,15 +1330,16 @@ class ModelConfig:
             return 1
 
         total_num_kv_heads = self.get_total_num_kv_heads()
+        tensor_parallel_size = parallel_config.attention_tp_size
         # If tensor parallelism is used, we divide the number of KV heads by
         # the tensor parallel size. We will replicate the KV heads in the
         # case where the number of KV heads is smaller than the tensor
         # parallel size so each GPU has at least one KV head.
-        return max(1, total_num_kv_heads // parallel_config.tensor_parallel_size)
+        return max(1, total_num_kv_heads // tensor_parallel_size)
 
     def get_num_attention_heads(self, parallel_config: ParallelConfig) -> int:
         num_heads = self.model_arch_config.total_num_attention_heads
-        return num_heads // parallel_config.tensor_parallel_size
+        return num_heads // parallel_config.attention_tp_size
 
     def get_num_experts(self) -> int:
         return self.model_arch_config.num_experts

vllm/config/parallel.py

@@ -121,6 +121,8 @@ class ParallelConfig:
     """Number of pipeline parallel groups."""
     tensor_parallel_size: int = Field(default=1, ge=1)
     """Number of tensor parallel groups."""
+    tensor_parallel_size_attention: int | None = Field(default=None, gt=0)
+    """Input tensor-parallel size for explicitly migrated attention layers."""
     prefill_context_parallel_size: int = Field(default=1, ge=1)
     """Number of prefill context parallel groups."""
     data_parallel_size: int = Field(default=1, ge=1)
@@ -511,8 +513,37 @@ class ParallelConfig:
                 "dcp_comm_backend='a2a' requires decode_context_parallel_size > 1."
             )
 
+        if self.tensor_parallel_size_attention is not None:
+            attention_tp = self.tensor_parallel_size_attention
+            if attention_tp > self.tensor_parallel_size:
+                raise ValueError(
+                    "tensor_parallel_size_attention cannot exceed "
+                    f"tensor_parallel_size: {attention_tp} > "
+                    f"{self.tensor_parallel_size}."
+                )
+            if self.tensor_parallel_size % attention_tp != 0:
+                raise ValueError(
+                    "tensor_parallel_size must be divisible by "
+                    f"tensor_parallel_size_attention: {self.tensor_parallel_size} "
+                    f"% {attention_tp} != 0."
+                )
+            required_dcp = self.tensor_parallel_size // attention_tp
+            if (
+                attention_tp != self.tensor_parallel_size
+                and self.decode_context_parallel_size != required_dcp
+            ):
+                raise ValueError(
+                    "decode_context_parallel_size must equal "
+                    "tensor_parallel_size / tensor_parallel_size_attention "
+                    f"({required_dcp}), got {self.decode_context_parallel_size}."
+                )
+
         return self
 
+    @property
+    def attention_tp_size(self) -> int:
+        return self.tensor_parallel_size_attention or self.tensor_parallel_size
+
     @property
     def world_size_across_dp(self) -> int:
         """world_size_across_dp is TPxPPxDP, it is the size of the world

vllm/config/vllm.py

@@ -2066,6 +2066,7 @@ class VllmConfig:
             f"download_dir={self.load_config.download_dir!r}, "
             f"load_format={self.load_config.load_format}, "
             f"tensor_parallel_size={self.parallel_config.tensor_parallel_size}, "  # noqa
+            f"tensor_parallel_size_attention={self.parallel_config.tensor_parallel_size_attention}, "  # noqa
             f"pipeline_parallel_size={self.parallel_config.pipeline_parallel_size}, "  # noqa
             f"data_parallel_size={self.parallel_config.data_parallel_size}, "  # noqa
             f"decode_context_parallel_size={self.parallel_config.decode_context_parallel_size}, "  # noqa

vllm/distributed/layer_parallel.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Per-layer input and output parallel plans."""
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from enum import Enum, auto
+
+
+class LayerType(Enum):
+    """Semantic layer types understood by parallel policies."""
+
+    DEFAULT = auto()
+    ATTENTION = auto()
+
+
+class ParallelGroupType(Enum):
+    """Runtime process groups that can back a parallel axis."""
+
+    TENSOR = auto()
+    ATTENTION_TENSOR = auto()
+
+
+@dataclass(frozen=True)
+class ParallelAxis:
+    """A rank's position along one parallel axis."""
+
+    world_size: int
+    rank: int
+    group: ParallelGroupType
+
+    def __post_init__(self) -> None:
+        if self.world_size < 1:
+            raise ValueError(f"world_size must be positive, got {self.world_size}")
+        if not 0 <= self.rank < self.world_size:
+            raise ValueError(f"rank must be in [0, {self.world_size}), got {self.rank}")
+
+
+@dataclass(frozen=True)
+class LayerParallelPlan:
+    """Input and output layouts for a layer."""
+
+    input: ParallelAxis
+    output: ParallelAxis
+
+    @property
+    def reshards_output(self) -> bool:
+        return self.input != self.output
+
+    def get_output_size(self, input_size: int) -> int:
+        global_size = input_size * self.input.world_size
+        if global_size % self.output.world_size != 0:
+            raise ValueError(
+                "Global input size must be divisible by the output parallel size: "
+                f"{global_size=}, output_size={self.output.world_size}"
+            )
+        return global_size // self.output.world_size
+
+
+_SINGLE_RANK_AXIS = ParallelAxis(1, 0, ParallelGroupType.TENSOR)
+_SINGLE_RANK_PLAN = LayerParallelPlan(
+    input=_SINGLE_RANK_AXIS,
+    output=_SINGLE_RANK_AXIS,
+)
+_default_plan = _SINGLE_RANK_PLAN
+_layer_plans: dict[LayerType, LayerParallelPlan] = {}
+
+
+def init_layer_parallel_config(
+    default_plan: LayerParallelPlan,
+    overrides: Mapping[LayerType, LayerParallelPlan] | None = None,
+) -> None:
+    """Install the process-local parallel policy after group initialization."""
+    global _default_plan, _layer_plans
+    _default_plan = default_plan
+    _layer_plans = dict(overrides or {})
+
+
+def get_layer_parallel_config(
+    layer_type: LayerType = LayerType.DEFAULT,
+) -> LayerParallelPlan:
+    """Resolve ``layer_type`` to a fully specified parallel plan."""
+    return _layer_plans.get(layer_type, _default_plan)
+
+
+def clear_layer_parallel_config() -> None:
+    """Restore the single-rank default policy."""
+    init_layer_parallel_config(_SINGLE_RANK_PLAN)

vllm/distributed/parallel_state.py

@@ -46,6 +46,13 @@ import vllm.envs as envs
 from vllm.distributed.device_communicators.base_device_communicator import (
     DeviceCommunicatorBase,
 )
+from vllm.distributed.layer_parallel import (
+    LayerParallelPlan,
+    LayerType,
+    ParallelAxis,
+    ParallelGroupType,
+    init_layer_parallel_config,
+)
 from vllm.distributed.utils import (
     StatelessProcessGroup,
     get_cached_tcp_store_client,
@@ -1370,6 +1377,15 @@ def get_tp_group() -> GroupCoordinator:
     return _TP
 
 
+_ATTENTION_TP: GroupCoordinator | None = None
+
+
+def get_parallel_group(group_type: ParallelGroupType) -> GroupCoordinator:
+    if group_type is ParallelGroupType.ATTENTION_TENSOR:
+        return _ATTENTION_TP or get_tp_group()
+    return get_tp_group()
+
+
 _DCP: GroupCoordinator | None = None
 
 
@@ -1810,6 +1826,31 @@ def initialize_model_parallel(
         group_name="tp",
     )
 
+    attention_tp_size = parallel_config.attention_tp_size
+    global _ATTENTION_TP
+    assert _ATTENTION_TP is None, (
+        "attention tensor parallel group is already initialized"
+    )
+    if attention_tp_size != tensor_model_parallel_size:
+        dcp_size = tensor_model_parallel_size // attention_tp_size
+        tp_ranks = (
+            local_all_ranks.view(-1, tensor_model_parallel_size)
+            if enable_elastic_ep
+            else all_ranks.view(-1, tensor_model_parallel_size)
+        )
+        group_ranks = (
+            tp_ranks.reshape(-1, attention_tp_size, dcp_size)
+            .transpose(1, 2)
+            .reshape(-1, attention_tp_size)
+            .unbind(0)
+        )
+        _ATTENTION_TP = init_model_parallel_group(
+            [ranks.tolist() for ranks in group_ranks],
+            get_world_group().local_rank,
+            backend,
+            group_name="attention_tp",
+        )
+
     # Build the DCP model-parallel groups.
     global _DCP
     assert _DCP is None, "decode context model parallel group is already initialized"
@@ -1939,6 +1980,23 @@ def initialize_model_parallel(
     # If no EP group needed, _EP remains None
     # If no EPLB group needed, _EPLB remains None
 
+    tensor_axis = ParallelAxis(
+        _TP.world_size, _TP.rank_in_group, ParallelGroupType.TENSOR
+    )
+    overrides: dict[LayerType, LayerParallelPlan] = {}
+    if _ATTENTION_TP is not None:
+        overrides[LayerType.ATTENTION] = LayerParallelPlan(
+            input=ParallelAxis(
+                _ATTENTION_TP.world_size,
+                _ATTENTION_TP.rank_in_group,
+                ParallelGroupType.ATTENTION_TENSOR,
+            ),
+            output=tensor_axis,
+        )
+    init_layer_parallel_config(
+        LayerParallelPlan(input=tensor_axis, output=tensor_axis), overrides
+    )
+
     logger.info_once(
         "rank %s in world size %s is assigned as "
         "DP rank %s, PP rank %s, PCP rank %s, "
@@ -2008,6 +2066,8 @@ def prepare_communication_buffer_for_model(model: torch.nn.Module):
     """
     if _TP is not None:
         _TP.prepare_communication_buffer_for_model(model)
+    if _ATTENTION_TP is not None:
+        _ATTENTION_TP.prepare_communication_buffer_for_model(model)
     if _PCP is not None:
         _PCP.prepare_communication_buffer_for_model(model)
     if _PP is not None:
@@ -2046,6 +2106,11 @@ def get_node_count() -> int:
 
 def destroy_model_parallel():
     """Set the groups to none and destroy them."""
+    global _ATTENTION_TP
+    if _ATTENTION_TP:
+        _ATTENTION_TP.destroy()
+    _ATTENTION_TP = None
+
     global _TP
 
     if _TP:
@@ -2082,6 +2147,10 @@ def destroy_model_parallel():
         _EPLB.destroy()
     _EPLB = None
 
+    from vllm.distributed.layer_parallel import clear_layer_parallel_config
+
+    clear_layer_parallel_config()
+
 
 def destroy_distributed_environment():
     global _WORLD, _NODE_COUNT

vllm/engine/arg_utils.py

@@ -471,6 +471,9 @@ class EngineArgs:
     numa_bind_cpus: list[str] | None = ParallelConfig.numa_bind_cpus
     device_ids: list[int | str] | None = None
     tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
+    tensor_parallel_size_attention: int | None = (
+        ParallelConfig.tensor_parallel_size_attention
+    )
     prefill_context_parallel_size: int = ParallelConfig.prefill_context_parallel_size
     decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size
     dcp_comm_backend: DCPCommBackend = ParallelConfig.dcp_comm_backend
@@ -1011,6 +1014,11 @@ class EngineArgs:
         parallel_group.add_argument(
             "--tensor-parallel-size", "-tp", **parallel_kwargs["tensor_parallel_size"]
         )
+        parallel_group.add_argument(
+            "--tensor-parallel-size-attention",
+            "-tpa",
+            **parallel_kwargs["tensor_parallel_size_attention"],
+        )
         parallel_group.add_argument(
             "--decode-context-parallel-size",
             "-dcp",
@@ -2094,6 +2102,7 @@ class EngineArgs:
         parallel_config = ParallelConfig(
             pipeline_parallel_size=self.pipeline_parallel_size,
             tensor_parallel_size=self.tensor_parallel_size,
+            tensor_parallel_size_attention=self.tensor_parallel_size_attention,
             prefill_context_parallel_size=self.prefill_context_parallel_size,
             data_parallel_size=self.data_parallel_size,
             data_parallel_rank=self.data_parallel_rank or 0,

vllm/model_executor/layers/attention/__init__.py

@@ -1,7 +1,10 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-from vllm.model_executor.layers.attention.attention import Attention
+from vllm.model_executor.layers.attention.attention import (
+    Attention,
+    HeteroParallelAttention,
+)
 from vllm.model_executor.layers.attention.chunked_local_attention import (
     ChunkedLocalAttention,
 )
@@ -24,6 +27,7 @@ __all__ = [
     "ChunkedLocalAttention",
     "CrossAttention",
     "EncoderOnlyAttention",
+    "HeteroParallelAttention",
     "MLAAttention",
     "MMEncoderAttention",
     "PrefillPrefixLMAttention",

vllm/model_executor/layers/attention/attention.py

@@ -10,6 +10,11 @@ import vllm.envs as envs
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
 from vllm.config import CacheConfig, get_current_vllm_config
 from vllm.config.vllm import VllmConfig
+from vllm.distributed.layer_parallel import (
+    LayerParallelPlan,
+    LayerType,
+    get_layer_parallel_config,
+)
 from vllm.forward_context import ForwardContext, get_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention.kv_transfer_utils import (
@@ -332,6 +337,7 @@ class Attention(nn.Module, AttentionLayerBase):
         self.layer_name = prefix
 
         self.num_heads = num_heads
+        self.output_num_heads = num_heads
         self.head_size = head_size
         self.head_size_v = self.head_size if head_size_v is None else head_size_v
         self.num_kv_heads = num_kv_heads
@@ -524,14 +530,16 @@ class Attention(nn.Module, AttentionLayerBase):
             # Handle both 2D [num_tokens, hidden] and
             # 3D [num_tokens, heads, head_dim] query
             num_tokens = query.shape[0]
-            output_shape = torch.Size((num_tokens, self.num_heads * self.head_size_v))
+            output_shape = torch.Size(
+                (num_tokens, self.output_num_heads * self.head_size_v)
+            )
         output = torch.empty(output_shape, dtype=output_dtype, device=query.device)
         hidden_size = output_shape[-1]
         # Reshape the query, key, and value tensors.
         # NOTE(woosuk): We do this outside the custom op to minimize the
         # CPU overheads from the non-CUDA-graph regions.
         query = query.view(-1, self.num_heads, self.head_size)
-        output = output.view(-1, self.num_heads, self.head_size_v)
+        output = output.view(-1, self.output_num_heads, self.head_size_v)
         if key is not None:
             key = key.view(-1, self.num_kv_heads, self.head_size)
         if value is not None:
@@ -688,6 +696,28 @@ class Attention(nn.Module, AttentionLayerBase):
             )
 
 
+class HeteroParallelAttention(Attention):
+    """Attention with an explicit input-to-output parallel plan."""
+
+    def __init__(
+        self,
+        *args: Any,
+        parallel_plan: LayerParallelPlan | None = None,
+        **kwargs: Any,
+    ) -> None:
+        self.parallel_plan = parallel_plan or get_layer_parallel_config(
+            LayerType.ATTENTION
+        )
+        super().__init__(*args, **kwargs)
+
+        if not self.attn_backend.supports_layer_parallel_plan(self.parallel_plan):
+            raise ValueError(
+                f"{self.attn_backend.get_name()} does not support layer parallel "
+                f"plan {self.parallel_plan}."
+            )
+        self.output_num_heads = self.parallel_plan.get_output_size(self.num_heads)
+
+
 def maybe_calc_kv_scales(
     query: torch.Tensor,
     key: torch.Tensor,

vllm/model_executor/layers/linear.py

@@ -19,6 +19,10 @@ from vllm.distributed import (
     tensor_model_parallel_all_gather,
     tensor_model_parallel_all_reduce,
 )
+from vllm.distributed.layer_parallel import (
+    LayerParallelPlan,
+    get_layer_parallel_config,
+)
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import PluggableLayer
 from vllm.model_executor.layers.batch_invariant import (
@@ -982,6 +986,7 @@ class QKVParallelLinear(ColumnParallelLinear):
         return_bias: bool = True,
         disable_tp: bool = False,
         v_head_size: int | None = None,
+        parallel_plan: LayerParallelPlan | None = None,
     ):
         self.hidden_size = hidden_size
         self.head_size = head_size
@@ -990,15 +995,28 @@ class QKVParallelLinear(ColumnParallelLinear):
         if total_num_kv_heads is None:
             total_num_kv_heads = total_num_heads
         self.total_num_kv_heads = total_num_kv_heads
-        # Divide the weight matrix along the last dimension.
         tp_size = get_tensor_model_parallel_world_size() if not disable_tp else 1
-        self.num_heads = divide(self.total_num_heads, tp_size)
-        if tp_size >= self.total_num_kv_heads:
+        plan = parallel_plan or get_layer_parallel_config()
+        if not disable_tp and plan.output.world_size != tp_size:
+            raise ValueError(
+                "QKV output plan must match tensor parallelism: "
+                f"plan={plan.output.world_size}, tp={tp_size}"
+            )
+
+        input_tp_size = plan.input.world_size
+        input_tp_rank = plan.input.rank
+        if disable_tp:
+            input_tp_size = 1
+            input_tp_rank = 0
+
+        self.num_heads = divide(self.total_num_heads, input_tp_size)
+        if input_tp_size >= self.total_num_kv_heads:
             self.num_kv_heads = 1
-            self.num_kv_head_replicas = divide(tp_size, self.total_num_kv_heads)
+            self.num_kv_head_replicas = divide(input_tp_size, self.total_num_kv_heads)
         else:
-            self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)
+            self.num_kv_heads = divide(self.total_num_kv_heads, input_tp_size)
             self.num_kv_head_replicas = 1
+        self._qkv_tp_rank = input_tp_rank
         input_size = self.hidden_size
         self.output_sizes = [
             self.num_heads * self.head_size * tp_size,  # q_proj
@@ -1046,6 +1064,18 @@ class QKVParallelLinear(ColumnParallelLinear):
         }
         return shard_size_mapping.get(loaded_shard_id)
 
+    def split_qkv(
+        self, qkv: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return qkv.split(
+            [
+                self.num_heads * self.head_size,
+                self.num_kv_heads * self.head_size,
+                self.num_kv_heads * self.v_head_size,
+            ],
+            dim=-1,
+        )
+
     def _load_fused_module_from_checkpoint(
         self, param: BasevLLMParameter, loaded_weight: torch.Tensor
     ):
@@ -1111,11 +1141,16 @@ class QKVParallelLinear(ColumnParallelLinear):
                 # works correctly while preserving the parameter shape.
                 for idx in range(param.data.shape[0]):
                     param.load_qkv_weight(
-                        loaded_weight=loaded_weight, shard_id=idx, tp_rank=self.tp_rank
+                        loaded_weight=loaded_weight,
+                        shard_id=idx,
+                        tp_rank=self._qkv_tp_rank,
                     )
                 return
             elif type(param) in (RowvLLMParameter, BasevLLMParameter):
-                param.load_qkv_weight(loaded_weight=loaded_weight, tp_rank=self.tp_rank)
+                param.load_qkv_weight(
+                    loaded_weight=loaded_weight,
+                    tp_rank=self._qkv_tp_rank,
+                )
                 return
             # TODO: @dsikka - move to parameter.py
             self._load_fused_module_from_checkpoint(param, loaded_weight)
@@ -1139,7 +1174,7 @@ class QKVParallelLinear(ColumnParallelLinear):
             shard_id=loaded_shard_id,
             shard_offset=shard_offset,
             shard_size=shard_size,
-            tp_rank=self.tp_rank,
+            tp_rank=self._qkv_tp_rank,
         )
 
     def weight_loader(
@@ -1297,9 +1332,9 @@ class QKVParallelLinear(ColumnParallelLinear):
 
             param_data = param_data.narrow(output_dim, shard_offset, shard_size)
             if loaded_shard_id == "q":
-                shard_rank = self.tp_rank
+                shard_rank = self._qkv_tp_rank
             else:
-                shard_rank = self.tp_rank // self.num_kv_head_replicas
+                shard_rank = self._qkv_tp_rank // self.num_kv_head_replicas
             start_idx = shard_rank * shard_size
 
             if not is_sharded_weight:

vllm/model_executor/layers/minimax_rms_norm/rms_norm_tp.py

@@ -6,11 +6,15 @@ from functools import partial
 import torch
 from torch import nn
 
-from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
+from vllm.distributed.layer_parallel import (
+    LayerParallelPlan,
+    ParallelGroupType,
+    get_layer_parallel_config,
+)
 from vllm.distributed.parallel_state import (
+    get_parallel_group,
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
-    get_tp_group,
 )
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp
@@ -26,8 +30,11 @@ MINIMAX_QK_NORM_MAX_TOKEN_NUM = 2048
 _MINIMAX_FUSED_AR_RMS_QK = getattr(torch.ops._C, "minimax_allreduce_rms_qk", None)
 
 
-def _all_reduce_variance(var: torch.Tensor) -> torch.Tensor:
-    """All-reduce a per-token variance tensor across the TP group.
+def _all_reduce_variance(
+    var: torch.Tensor,
+    reduce_group_type: int,
+) -> torch.Tensor:
+    """All-reduce a per-token variance tensor across the configured group.
 
     Variance is accumulated in fp32 for numerical stability. The FlashInfer
     fused all-reduce caches a single global workspace keyed to the model's
@@ -37,7 +44,8 @@ def _all_reduce_variance(var: torch.Tensor) -> torch.Tensor:
     a flattened (1D) view keeps these fp32 reductions on custom all-reduce /
     pynccl, both of which handle fp32 correctly.
     """
-    return tensor_model_parallel_all_reduce(var.flatten()).view_as(var)
+    group = get_parallel_group(ParallelGroupType(reduce_group_type))
+    return group.all_reduce(var.flatten()).view_as(var)
 
 
 @triton.jit
@@ -132,6 +140,7 @@ def _minimax_qk_norm_tp_eager(
     kv_size: int,
     tp_world: int,
     eps: float,
+    reduce_group_type: int = ParallelGroupType.TENSOR.value,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Pure-torch reference path used when Triton is unavailable."""
     q, k, _ = qkv.split([q_size, kv_size, kv_size], dim=-1)
@@ -142,7 +151,7 @@ def _minimax_qk_norm_tp_eager(
     k_var = k.pow(2).mean(dim=-1, keepdim=True)
 
     qk_var = torch.cat([q_var, k_var], dim=-1)
-    qk_var = _all_reduce_variance(qk_var) / tp_world
+    qk_var = _all_reduce_variance(qk_var, reduce_group_type) / tp_world
     q_var, k_var = qk_var.chunk(2, dim=-1)
     q = q * torch.rsqrt(q_var + eps) * q_weight
     k = k * torch.rsqrt(k_var + eps) * k_weight
@@ -158,6 +167,7 @@ def _minimax_qk_norm_tp_fallback(
     tp_rank: int,
     tp_world: int,
     eps: float,
+    reduce_group_type: int = ParallelGroupType.TENSOR.value,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """All-reduce + QK RMSNorm without the Lamport fused kernel.
 
@@ -169,7 +179,14 @@ def _minimax_qk_norm_tp_fallback(
     """
     if not HAS_TRITON:
         return _minimax_qk_norm_tp_eager(
-            qkv, q_weight, k_weight, q_size, kv_size, tp_world, eps
+            qkv,
+            q_weight,
+            k_weight,
+            q_size,
+            kv_size,
+            tp_world,
+            eps,
+            reduce_group_type,
         )
 
     num_tokens = qkv.shape[0]
@@ -184,7 +201,7 @@ def _minimax_qk_norm_tp_fallback(
 
     # All-reduce sums the per-shard means; the /tp_world that turns this back
     # into the global mean is folded into the apply kernel's rsqrt below.
-    qk_var = _all_reduce_variance(qk_var)
+    qk_var = _all_reduce_variance(qk_var, reduce_group_type)
 
     q_out = torch.empty(num_tokens, q_size, dtype=qkv.dtype, device=qkv.device)
     k_out = torch.empty(num_tokens, kv_size, dtype=qkv.dtype, device=qkv.device)
@@ -215,6 +232,7 @@ def _minimax_qk_norm_fusion(
     tp_world: int,
     eps: float,
     workspace: torch.Tensor | None,
+    reduce_group_type: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     assert qkv.ndim == 2
     num_tokens = qkv.shape[0]
@@ -236,7 +254,15 @@ def _minimax_qk_norm_fusion(
             eps,
         )
     return _minimax_qk_norm_tp_fallback(
-        qkv, q_weight, k_weight, q_size, kv_size, tp_rank, tp_world, eps
+        qkv,
+        q_weight,
+        k_weight,
+        q_size,
+        kv_size,
+        tp_rank,
+        tp_world,
+        eps,
+        reduce_group_type,
     )
 
 
@@ -250,6 +276,7 @@ def _minimax_qk_norm_fusion_fake(
     tp_world: int,
     eps: float,
     workspace: torch.Tensor | None,
+    reduce_group_type: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     assert qkv.ndim == 2
     num_tokens = qkv.shape[0]
@@ -276,10 +303,14 @@ class MiniMaxText01RMSNormTP(CustomOp):
         *,
         weight_shard_world_size: int | None = None,
         weight_shard_rank: int | None = None,
+        parallel_plan: LayerParallelPlan | None = None,
     ) -> None:
         super().__init__()
-        self.tp_world = get_tensor_model_parallel_world_size()
-        self.tp_rank = get_tensor_model_parallel_rank()
+        plan = parallel_plan or get_layer_parallel_config()
+        self.reduce_group = get_parallel_group(plan.input.group)
+        self.reduce_group_type = plan.input.group.value
+        self.tp_world = plan.input.world_size
+        self.tp_rank = plan.input.rank
         self.weight_shard_world = weight_shard_world_size or self.tp_world
         self.weight_shard_rank = (
             self.tp_rank if weight_shard_rank is None else weight_shard_rank
@@ -309,7 +340,7 @@ class MiniMaxText01RMSNormTP(CustomOp):
                     rank=self.tp_rank,
                     world_size=self.tp_world,
                     max_tokens=MINIMAX_QK_NORM_MAX_TOKEN_NUM,
-                    process_group=get_tp_group().cpu_group,
+                    process_group=self.reduce_group.cpu_group,
                 )
             except Exception as e:
                 logger.warning_once(
@@ -345,7 +376,9 @@ class MiniMaxText01RMSNormTP(CustomOp):
         x = x.to(torch.float32)
         variance = x.pow(2).mean(dim=-1, keepdim=True, dtype=torch.float32)
         if self.tp_world > 1:
-            variance = _all_reduce_variance(variance) / self.tp_world
+            variance = (
+                _all_reduce_variance(variance, self.reduce_group_type) / self.tp_world
+            )
         x = x * torch.rsqrt(variance + self.variance_epsilon)
         x = (x * self.weight).to(orig_dtype)
         return x
@@ -389,7 +422,7 @@ class MiniMaxText01RMSNormTP(CustomOp):
 
         assert q_norm.variance_epsilon == k_norm.variance_epsilon
         # Case 0： tp_size=1
-        if get_tensor_model_parallel_world_size() == 1:
+        if q_norm.tp_world == 1:
             q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
             q, k = MiniMaxText01RMSNormTP.forward_qk(q_norm, k_norm, q, k)
             return q, k, v
@@ -404,6 +437,7 @@ class MiniMaxText01RMSNormTP(CustomOp):
             q_norm.tp_world,
             q_norm.variance_epsilon,
             q_norm.workspace,
+            q_norm.reduce_group_type,
         )
         _, _, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
         return q, k, v

vllm/model_executor/models/interfaces.py

@@ -29,6 +29,7 @@ from torch import Tensor
 from transformers.models.whisper.tokenization_whisper import LANGUAGES
 from typing_extensions import Self, TypeIs
 
+from vllm.distributed.layer_parallel import LayerType
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.utils.collection_utils import common_prefix
@@ -70,6 +71,19 @@ The output embeddings must be one of the following formats:
 """
 
 
+class SupportsHeteroLayerParallelism:
+    """Marker for models migrated to explicit layer parallel plans."""
+
+    supported_layer_types: ClassVar[frozenset[LayerType]] = frozenset(
+        {LayerType.ATTENTION}
+    )
+
+
+def get_supported_layer_types(model: type[object] | object) -> tuple[str, ...]:
+    layer_types = getattr(model, "supported_layer_types", frozenset())
+    return tuple(sorted(layer_type.name for layer_type in layer_types))
+
+
 class StreamingTranscriptionPostProcessor:
     """Stateful streaming post-processor for transcription deltas."""
 

vllm/model_executor/models/llama.py

@@ -33,11 +33,15 @@ from transformers import LlamaConfig
 
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, VllmConfig
-from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
+from vllm.distributed import get_pp_group
+from vllm.distributed.layer_parallel import (
+    LayerType,
+    get_layer_parallel_config,
+)
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.attention import (
-    Attention,
     EncoderOnlyAttention,
+    HeteroParallelAttention,
 )
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
@@ -61,6 +65,7 @@ from .interfaces import (
     LocalArgmaxMixin,
     SupportsEagle,
     SupportsEagle3,
+    SupportsHeteroLayerParallelism,
     SupportsLoRA,
     SupportsPP,
     SupportsQuant,
@@ -137,27 +142,17 @@ class LlamaAttention(nn.Module):
         super().__init__()
         layer_idx = extract_layer_index(prefix)
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
         self.total_num_heads = num_heads
-        assert self.total_num_heads % tp_size == 0
-        self.num_heads = self.total_num_heads // tp_size
-        self.total_num_kv_heads = num_kv_heads
-        if self.total_num_kv_heads >= tp_size:
-            # Number of KV heads is greater than TP size, so we partition
-            # the KV heads across multiple tensor parallel GPUs.
-            assert self.total_num_kv_heads % tp_size == 0
-        else:
-            # Number of KV heads is less than TP size, so we replicate
-            # the KV heads across multiple tensor parallel GPUs.
-            assert tp_size % self.total_num_kv_heads == 0
-        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+        self.total_num_kv_heads = num_kv_heads or num_heads
 
         head_dim = getattr(config, "head_dim", None)
         self.head_dim = head_dim or self.hidden_size // self.total_num_heads
-        self.q_size = self.num_heads * self.head_dim
-        self.kv_size = self.num_kv_heads * self.head_dim
         self.scaling = self.head_dim**-0.5
         self.max_position_embeddings = max_position_embeddings
+        is_encoder_only = attn_type == AttentionType.ENCODER_ONLY
+        parallel_plan = get_layer_parallel_config(
+            LayerType.DEFAULT if is_encoder_only else LayerType.ATTENTION
+        )
 
         self.qkv_proj = QKVParallelLinear(
             hidden_size=hidden_size,
@@ -167,7 +162,10 @@ class LlamaAttention(nn.Module):
             bias=bias,
             quant_config=quant_config,
             prefix=f"{prefix}.qkv_proj",
+            parallel_plan=parallel_plan,
         )
+        self.num_heads = self.qkv_proj.num_heads
+        self.num_kv_heads = self.qkv_proj.num_kv_heads
 
         self.o_proj = RowParallelLinear(
             input_size=self.total_num_heads * self.head_dim,
@@ -200,11 +198,7 @@ class LlamaAttention(nn.Module):
             if is_sliding:
                 sliding_window = config.sliding_window
 
-        attn_cls = (
-            EncoderOnlyAttention
-            if attn_type == AttentionType.ENCODER_ONLY
-            else Attention
-        )
+        attn_cls = EncoderOnlyAttention if is_encoder_only else HeteroParallelAttention
 
         self.attn = attn_cls(
             self.num_heads,
@@ -224,7 +218,7 @@ class LlamaAttention(nn.Module):
         hidden_states: torch.Tensor,
     ) -> torch.Tensor:
         qkv, _ = self.qkv_proj(hidden_states)
-        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q, k, v = self.qkv_proj.split_qkv(qkv)
         q, k = self.rotary_emb(positions, q, k)
         attn_output = self.attn(q, k, v)
         output, _ = self.o_proj(attn_output)
@@ -451,6 +445,7 @@ class LlamaForCausalLM(
     SupportsEagle,
     SupportsEagle3,
     SupportsQuant,
+    SupportsHeteroLayerParallelism,
 ):
     hf_to_vllm_mapper = LlamaModel.hf_to_vllm_mapper
     # LoRA specific attributes

vllm/model_executor/models/minimax_m2.py

@@ -35,10 +35,13 @@ from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, ModelConfig, VllmConfig
 from vllm.distributed import (
     get_pp_group,
-    get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
 )
-from vllm.model_executor.layers.attention import Attention
+from vllm.distributed.layer_parallel import (
+    LayerType,
+    get_layer_parallel_config,
+)
+from vllm.model_executor.layers.attention import HeteroParallelAttention
 from vllm.model_executor.layers.fused_moe import (
     FusedMoE,
 )
@@ -58,7 +61,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
 )
 from vllm.sequence import IntermediateTensors
 
-from .interfaces import EagleModelMixin, SupportsEagle3, SupportsLoRA, SupportsPP
+from .interfaces import (
+    EagleModelMixin,
+    SupportsEagle3,
+    SupportsHeteroLayerParallelism,
+    SupportsLoRA,
+    SupportsPP,
+)
 from .utils import (
     AutoWeightsLoader,
     PPMissingLayer,
@@ -155,23 +164,11 @@ class MiniMaxM2Attention(nn.Module):
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
-        tp_size = get_tensor_model_parallel_world_size()
+        parallel_plan = get_layer_parallel_config(LayerType.ATTENTION)
+        input_axis = parallel_plan.input
         self.total_num_heads = num_heads
-        assert self.total_num_heads % tp_size == 0
-        self.num_heads = self.total_num_heads // tp_size
         self.total_num_kv_heads = num_kv_heads
-        if self.total_num_kv_heads >= tp_size:
-            # Number of KV heads is greater than TP size, so we partition
-            # the KV heads across multiple tensor parallel GPUs.
-            assert self.total_num_kv_heads % tp_size == 0
-        else:
-            # Number of KV heads is less than TP size, so we replicate
-            # the KV heads across multiple tensor parallel GPUs.
-            assert tp_size % self.total_num_kv_heads == 0
-        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
         self.head_dim = head_dim or (hidden_size // self.total_num_heads)
-        self.q_size = self.num_heads * self.head_dim
-        self.kv_size = self.num_kv_heads * self.head_dim
         self.scaling = self.head_dim**-0.5
         self.max_position_embeddings = max_position_embeddings
 
@@ -183,7 +180,12 @@ class MiniMaxM2Attention(nn.Module):
             bias=qkv_bias,
             quant_config=quant_config,
             prefix=f"{prefix}.qkv_proj",
+            parallel_plan=parallel_plan,
         )
+        self.num_heads = self.qkv_proj.num_heads
+        self.num_kv_heads = self.qkv_proj.num_kv_heads
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
 
         self.o_proj = RowParallelLinear(
             self.total_num_heads * self.head_dim,
@@ -203,7 +205,7 @@ class MiniMaxM2Attention(nn.Module):
             max_position=max_position_embeddings,
             rope_parameters=rope_parameters,
         )
-        self.attn = Attention(
+        self.attn = HeteroParallelAttention(
             self.num_heads,
             self.head_dim,
             self.scaling,
@@ -215,22 +217,26 @@ class MiniMaxM2Attention(nn.Module):
         )
 
         self.q_norm = MiniMaxText01RMSNormTP(
-            self.head_dim * self.total_num_heads, eps=rms_norm_eps
+            self.head_dim * self.total_num_heads,
+            eps=rms_norm_eps,
+            parallel_plan=parallel_plan,
         )
-        if self.total_num_kv_heads >= tp_size:
+        if self.total_num_kv_heads >= input_axis.world_size:
             self.k_norm = MiniMaxText01RMSNormTP(
-                self.head_dim * self.total_num_kv_heads, eps=rms_norm_eps
+                self.head_dim * self.total_num_kv_heads,
+                eps=rms_norm_eps,
+                parallel_plan=parallel_plan,
             )
         else:
             # KV heads are replicated across TP ranks; shard k_norm weight by
             # total_num_kv_heads rather than tp_size to avoid incorrect sharding.
-            num_kv_head_replicas = tp_size // self.total_num_kv_heads
+            num_kv_head_replicas = input_axis.world_size // self.total_num_kv_heads
             self.k_norm = MiniMaxText01RMSNormTP(
                 self.head_dim * self.total_num_kv_heads,
                 eps=rms_norm_eps,
                 weight_shard_world_size=self.total_num_kv_heads,
-                weight_shard_rank=get_tensor_model_parallel_rank()
-                // num_kv_head_replicas,
+                weight_shard_rank=input_axis.rank // num_kv_head_replicas,
+                parallel_plan=parallel_plan,
             )
 
     def forward(
@@ -429,7 +435,13 @@ class MiniMaxM2Model(nn.Module, EagleModelMixin):
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
 
-class MiniMaxM2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):
+class MiniMaxM2ForCausalLM(
+    nn.Module,
+    SupportsLoRA,
+    SupportsPP,
+    SupportsEagle3,
+    SupportsHeteroLayerParallelism,
+):
     hf_to_vllm_mapper = MiniMaxM2Model.hf_to_vllm_mapper
     packed_modules_mapping = {
         "qkv_proj": [

vllm/model_executor/models/registry.py

@@ -29,6 +29,7 @@ from vllm.config import (
     iter_architecture_defaults,
     try_match_architecture_defaults,
 )
+from vllm.distributed.layer_parallel import LayerType
 from vllm.logger import init_logger
 from vllm.logging_utils import logtime
 from vllm.tasks import ScoreType
@@ -45,6 +46,7 @@ else:
 
 
 from .interfaces import (
+    get_supported_layer_types,
     has_inner_state,
     has_noops,
     is_attention_free,
@@ -766,6 +768,7 @@ class _ModelInfo:
     requires_raw_input_tokens: bool
     supports_multimodal_encoder_tp_data: bool
     supports_pp: bool
+    supported_layer_types: tuple[str, ...]
     has_inner_state: bool
     is_attention_free: bool
     is_hybrid: bool
@@ -793,6 +796,7 @@ class _ModelInfo:
                 model
             ),
             supports_pp=supports_pp(model),
+            supported_layer_types=get_supported_layer_types(model),
             has_inner_state=has_inner_state(model),
             is_attention_free=is_attention_free(model),
             is_hybrid=is_hybrid(model),
@@ -1344,6 +1348,15 @@ class _ModelRegistry:
         model_cls, _ = self.inspect_model_cls(architectures, model_config)
         return model_cls.supports_pp
 
+    def is_layer_parallel_supported_model(
+        self,
+        architectures: str | list[str],
+        model_config: ModelConfig,
+        layer_type: LayerType,
+    ) -> bool:
+        model_cls, _ = self.inspect_model_cls(architectures, model_config)
+        return layer_type.name in model_cls.supported_layer_types
+
     def model_has_inner_state(
         self,
         architectures: str | list[str],

vllm/model_executor/parameter.py

@@ -180,6 +180,7 @@ class _ColumnvLLMParameter(BasevLLMParameter):
         shard_size: int = kwargs["shard_size"]
         shard_id: str = kwargs["shard_id"]
         num_heads: int = kwargs["num_heads"]
+        tp_rank: int = kwargs.get("tp_rank", self.tp_rank)
 
         # TODO: move these to PackedColumnParameter and PackedvLLMParameter
         if (
@@ -191,7 +192,7 @@ class _ColumnvLLMParameter(BasevLLMParameter):
             )
 
         param_data = self.data
-        shard_id_int = self.tp_rank if shard_id == "q" else self.tp_rank // num_heads
+        shard_id_int = tp_rank if shard_id == "q" else tp_rank // num_heads
         param_data = param_data.narrow(self.output_dim, shard_offset, shard_size)
         loaded_weight = loaded_weight.narrow(
             self.output_dim, shard_id_int * shard_size, shard_size

vllm/v1/attention/backend.py

@@ -21,6 +21,7 @@ from vllm.utils.torch_utils import np_to_pinned_tensor
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
     from vllm.config.cache import CacheDType
+    from vllm.distributed.layer_parallel import LayerParallelPlan
     from vllm.model_executor.layers.linear import ColumnParallelLinear
     from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
     from vllm.platforms.interface import DeviceCapability
@@ -66,6 +67,10 @@ class AttentionBackend(ABC):
     # Does attention's forward() include kv cache update?
     forward_includes_kv_cache_update: bool = True
 
+    @classmethod
+    def supports_layer_parallel_plan(cls, plan: "LayerParallelPlan") -> bool:
+        return not plan.reshards_output
+
     @staticmethod
     def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
         return [MultipleOf(1)]

vllm/v1/attention/backends/flash_attn.py

@@ -47,6 +47,11 @@ from vllm.config import (
     get_layers_from_vllm_config,
 )
 from vllm.config.cache import CacheDType
+from vllm.distributed.layer_parallel import (
+    LayerParallelPlan,
+    LayerType,
+    get_layer_parallel_config,
+)
 from vllm.distributed.parallel_state import get_dcp_group
 from vllm.logger import init_logger
 from vllm.platforms.interface import DeviceCapability
@@ -82,6 +87,12 @@ class FlashAttentionBackend(AttentionBackend):
 
     forward_includes_kv_cache_update: bool = False
 
+    @classmethod
+    def supports_layer_parallel_plan(cls, plan: LayerParallelPlan) -> bool:
+        if plan.reshards_output and get_flash_attn_version() not in (3, 4):
+            return False
+        return plan.output.world_size % plan.input.world_size == 0
+
     @classmethod
     def get_preferred_block_size(cls, default_block_size: int) -> int:
         if current_platform.is_xpu():
@@ -362,7 +373,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
         self.cache_config = vllm_config.cache_config
         self.compilation_config = vllm_config.compilation_config
         self.attention_config = vllm_config.attention_config
-
         self.num_heads_q = self.model_config.get_num_attention_heads(
             self.parallel_config
         )
@@ -384,6 +394,13 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
             self.dcp_world_size = 1
             self.dcp_rank = 0
 
+        plan = get_layer_parallel_config(LayerType.ATTENTION)
+        self.dcp_context_num_heads = (
+            self.num_heads_q
+            if plan.reshards_output
+            else self.num_heads_q * self.dcp_world_size
+        )
+
         self.cp_kv_cache_interleave_size = (
             self.parallel_config.cp_kv_cache_interleave_size
         )
@@ -422,6 +439,13 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
                 dtype=torch.int32,
                 device=self.device,
             )
+            max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+            current_workspace_manager().get_simultaneous(
+                (
+                    (max_num_tokens, self.dcp_context_num_heads, self.headdim),
+                    self.model_config.dtype,
+                ),
+            )
 
         # Sliding window size to be used with the AOT scheduler will be
         # populated on first build() call.
@@ -510,7 +534,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
                     batch_size=batch_size,
                     max_seqlen_q=max_query_len,
                     max_seqlen_k=max_seq_len,
-                    num_heads_q=self.num_heads_q * self.dcp_world_size,
+                    num_heads_q=self.dcp_context_num_heads,
                     num_heads_kv=self.num_heads_kv,
                     headdim=self.headdim,
                     cache_seqlens=seqlens,
@@ -915,6 +939,7 @@ class FlashAttentionImpl(AttentionImpl):
 
             if self.dcp_world_size > 1:
                 self._forward_with_dcp(
+                    layer,
                     query[:num_actual_tokens],
                     key[:num_actual_tokens],
                     value[:num_actual_tokens],
@@ -1102,6 +1127,7 @@ class FlashAttentionImpl(AttentionImpl):
 
     def _forward_with_dcp(
         self,
+        layer: torch.nn.Module,
         query: torch.Tensor,
         key: torch.Tensor,
         value: torch.Tensor,
@@ -1120,14 +1146,29 @@ class FlashAttentionImpl(AttentionImpl):
         cu_seqlens_q = attn_metadata.query_start_loc
         max_seqlen_q = attn_metadata.max_query_len
         block_table = attn_metadata.block_table
+        parallel_plan = getattr(layer, "parallel_plan", None)
+        reshard_output = parallel_plan is not None and parallel_plan.reshards_output
+        if reshard_output:
+            output_heads = output.shape[1]
+            assert query.shape[1] == output_heads * self.dcp_world_size
+            head_start = get_dcp_group().rank_in_group * output_heads
+            output_head_slice = slice(head_start, head_start + output_heads)
 
         query = query.contiguous()
         if attn_metadata.max_dcp_context_kv_len == 0:
+            local_output = output
+            if reshard_output:
+                (local_output,) = current_workspace_manager().get_simultaneous(
+                    (
+                        (query.shape[0], query.shape[1], self.head_size),
+                        self._dcp_dtype,
+                    ),
+                )
             flash_attn_varlen_func(
                 q=query,
                 k=key,
                 v=value,
-                out=output,
+                out=local_output,
                 cu_seqlens_q=cu_seqlens_q,
                 max_seqlen_q=max_seqlen_q,
                 cu_seqlens_k=cu_seqlens_q,
@@ -1146,30 +1187,38 @@ class FlashAttentionImpl(AttentionImpl):
                 v_descale=v_descale,
                 num_splits=attn_metadata.max_num_splits,
             )
+            if reshard_output:
+                output.copy_(local_output[:, output_head_slice, :])
             return output
 
-        query_across_dcp = get_dcp_group().all_gather(query, dim=1)
+        context_query = (
+            query if reshard_output else get_dcp_group().all_gather(query, dim=1)
+        )
         sliding_window_size = (
             list(self.sliding_window) if self.sliding_window is not None else None
         )
-        n = query_across_dcp.shape[0]
+        n = context_query.shape[0]
         num_reqs = cu_seqlens_q.shape[0] - 1
         num_decodes = attn_metadata.num_decode_reqs
         num_context_prefills = attn_metadata.num_prefill_reqs
         num_decode_tokens = attn_metadata.num_decode_tokens
         num_context_prefill_tokens = attn_metadata.num_prefill_tokens
-        split_dcp_context = should_split_fa2_dcp_context_attention(
-            self.vllm_flash_attn_version,
-            max_seqlen_q,
-            num_reqs,
-            num_decodes,
-            num_context_prefills,
+        split_dcp_context = (
+            not reshard_output
+            and should_split_fa2_dcp_context_attention(
+                self.vllm_flash_attn_version,
+                max_seqlen_q,
+                num_reqs,
+                num_decodes,
+                num_context_prefills,
+            )
         )
         dcp_context_out_tokens = max(n, self._dcp_max_num_tokens)
         dcp_context_out_spec = (
             (
                 dcp_context_out_tokens,
-                self.num_heads * self.dcp_world_size,
+                self.num_heads,
+                self.dcp_world_size,
                 self.head_size,
             ),
             self._dcp_dtype,
@@ -1187,7 +1236,7 @@ class FlashAttentionImpl(AttentionImpl):
             assert self.vllm_flash_attn_version is not None
             context_attn_out, context_lse = run_split_fa2_dcp_context_attention(
                 flash_attn_varlen_func,
-                query_across_dcp,
+                context_query,
                 key_cache,
                 value_cache,
                 dcp_context_out,
@@ -1214,7 +1263,7 @@ class FlashAttentionImpl(AttentionImpl):
             )
         else:
             context_attn_out, context_lse = flash_attn_varlen_func(
-                q=query_across_dcp,
+                q=context_query,
                 k=key_cache,
                 v=value_cache,
                 out=dcp_context_out,
@@ -1236,6 +1285,9 @@ class FlashAttentionImpl(AttentionImpl):
                 v_descale=v_descale,
                 num_splits=attn_metadata.max_num_splits,
             )
+        context_num_heads = context_query.shape[1]
+        context_attn_out = context_attn_out[:, :context_num_heads]
+        context_lse = context_lse[:context_num_heads]
         # FA returns LSE in shape [ H, B ] but DCP combine wants [ B, H ]
         context_attn_out_cor, context_lse_cor = self.dcp_combine(
             context_attn_out,
@@ -1245,11 +1297,16 @@ class FlashAttentionImpl(AttentionImpl):
         )
         context_lse_cor = context_lse_cor.transpose(0, 1).contiguous()
 
+        query_output = output
+        if reshard_output:
+            (query_output,) = current_workspace_manager().get_simultaneous(
+                ((query.shape[0], query.shape[1], self.head_size), self._dcp_dtype),
+            )
         query_attn_out, query_lse = flash_attn_varlen_func(
             q=query,
             k=key,
             v=value,
-            out=output,
+            out=query_output,
             cu_seqlens_q=cu_seqlens_q,
             max_seqlen_q=max_seqlen_q,
             cu_seqlens_k=cu_seqlens_q,
@@ -1266,6 +1323,9 @@ class FlashAttentionImpl(AttentionImpl):
             v_descale=v_descale,
             num_splits=attn_metadata.max_num_splits,
         )
+        if reshard_output:
+            query_attn_out = query_attn_out[:, output_head_slice, :].contiguous()
+            query_lse = query_lse[output_head_slice, :]
         assert context_attn_out_cor.shape == query_attn_out.shape
         assert context_lse_cor.shape == query_lse.shape
         merge_attn_states(