189 lines
8.8 KiB
Python
189 lines
8.8 KiB
Python
# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# This file is a part of the vllm-ascend project.
|
|
#
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
from vllm.distributed.parallel_state import get_tp_group
|
|
from vllm.forward_context import get_forward_context
|
|
from vllm.model_executor.layers.fused_moe.routed_experts_capturer import RoutedExpertsCapturer
|
|
|
|
from vllm_ascend.ascend_forward_context import _EXTRA_CTX, MoECommType
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:
|
|
"""Capture expert routing decisions for a specific layer.
|
|
|
|
Under data parallelism, ``topk_ids`` may have four different batch
|
|
layouts depending on where the DP combine happens and whether
|
|
Sequence Parallelism (SP) is active for the MoE layer:
|
|
- ``n == total`` (naive dispatch): all DP ranks' tokens are
|
|
concatenated before routing; we slice out this rank's span
|
|
using the cumulative per-rank counts.
|
|
- ``n == token_num_per_dp`` (modular-kernel path): DP combine
|
|
happens inside ``quant_method.apply``; ``select_experts`` only
|
|
ever sees this rank's tokens, so we take the whole tensor.
|
|
- ``n == total_with_padding`` (padded all-gather path): tokens are
|
|
padded to max_tokens before all-gather across DP group; each
|
|
DP rank occupies a contiguous block of size max_tokens, and we
|
|
extract only the actual tokens for this rank (skip padding).
|
|
When all DP ranks have equal token counts, ``total == total_with_padding``,
|
|
so the naive dispatch branch fires instead (equivalent result).
|
|
- ``n == ceil(token_num_per_dp / tp_size)`` (SP + modular-kernel
|
|
path): tokens were split along dim=0 across the TP group by
|
|
``_sequence_parallel_context``
|
|
(``moe_runner_base.py:_sequence_parallel_context``), so each
|
|
TP rank only sees its shard. We all-gather along dim=0 to
|
|
reconstruct this DP rank's full routing tensor. SP pads with
|
|
ceil-div (see ``_compute_sp_num_tokens`` in
|
|
``forward_context.py``), so the gathered tensor may contain a
|
|
few trailing padding rows which are trimmed by the downstream
|
|
``[:token_num_per_dp]`` slice.
|
|
|
|
Args:
|
|
layer_id: The layer index.
|
|
topk_ids: Tensor of shape (batch_size, num_routed_experts).
|
|
"""
|
|
|
|
ctx = get_forward_context()
|
|
if ctx.dp_metadata is None: # single dp
|
|
start_loc = 0
|
|
end_loc = topk_ids.shape[0]
|
|
token_num_per_dp = topk_ids.shape[0]
|
|
else: # multi dp
|
|
num_tokens_dp = ctx.dp_metadata.num_tokens_across_dp_cpu
|
|
token_num_per_dp = int(num_tokens_dp[self.dp_rank].item())
|
|
total = int(num_tokens_dp.sum().item())
|
|
n = topk_ids.shape[0]
|
|
|
|
# Calculate total with padding for all-gather scenario.
|
|
# When tokens are padded to max_tokens before all-gather across DP group,
|
|
# the total size becomes max_tokens * dp_size.
|
|
# Example: DP0 has 5 tokens, DP1 has 7 tokens, max_tokens=7.
|
|
# After padding: DP0 has 7 tokens, DP1 has 7 tokens.
|
|
# After all-gather: total_with_padding = 7 * 2 = 14.
|
|
max_tokens = int(num_tokens_dp.max().item())
|
|
total_with_padding = max_tokens * len(num_tokens_dp)
|
|
|
|
if n == total:
|
|
# Naive dispatch: all DP ranks' tokens concatenated
|
|
# before routing. This rank owns tokens
|
|
# [end_loc - token_num_per_dp, end_loc).
|
|
cumsum = torch.cumsum(num_tokens_dp, dim=0)
|
|
end_loc = int(cumsum[self.dp_rank].item())
|
|
start_loc = end_loc - token_num_per_dp
|
|
elif n == token_num_per_dp:
|
|
# Modular-kernel path: DP combine happens inside
|
|
# quant_method.apply; select_experts only sees this
|
|
# rank's tokens, take the whole tensor.
|
|
start_loc = 0
|
|
end_loc = token_num_per_dp
|
|
elif n == total_with_padding:
|
|
# NOTE(Ronald1995): When all DP ranks have equal token counts,
|
|
# total == total_with_padding, so the first branch (n == total)
|
|
# fires instead. This overlap is intentional since both branches
|
|
# produce equivalent results in that case.
|
|
|
|
# Padded all-gather path: tokens are padded to max_tokens before
|
|
# all-gather across DP group. Each DP rank occupies a contiguous
|
|
# block of size max_tokens. Extract only the actual tokens for
|
|
# this rank (skip padding).
|
|
# Example: dp_rank=0, max_tokens=7, token_num_per_dp=5.
|
|
# start_loc = 0 * 7 = 0
|
|
# end_loc = 0 + 5 = 5 (only first 5 tokens are valid)
|
|
|
|
start_loc = self.dp_rank * max_tokens
|
|
end_loc = start_loc + token_num_per_dp
|
|
elif (
|
|
self.tp_size > 1
|
|
and n != token_num_per_dp
|
|
and (
|
|
# all2all scenario use tensor split, different tp rank have different
|
|
# size of tokens.
|
|
n == (token_num_per_dp + self.tp_size - 1) // self.tp_size
|
|
or n == token_num_per_dp // self.tp_size
|
|
# mc2 scenario will pad dp tokens to max_tokens and then ceil-div.
|
|
or n == (max_tokens + self.tp_size - 1) // self.tp_size
|
|
)
|
|
):
|
|
# SP + modular-kernel path. All-gather across the TP
|
|
# group along dim=0 to reconstruct the full per-DP-rank
|
|
# tensor; keep only the first ``token_num_per_dp`` rows
|
|
# (trailing rows are SP ceil-div padding). The TP group
|
|
# is always initialized on real rollout workers, and
|
|
# every rank in the group reaches this branch in
|
|
# lockstep (bind is per-FusedMoE layer, SP is a global
|
|
# condition), so a bare all_gather here will not
|
|
# deadlock -- let it raise if the precondition is
|
|
# violated rather than skip silently.
|
|
#
|
|
# ``topk_ids`` is already whatever the router produced
|
|
# (typically int32/int64, both supported by NCCL); the
|
|
# downstream ``device_buffer[...] = topk_ids[...]``
|
|
# setitem narrows into int32 automatically.
|
|
|
|
# NOTE(Ronald1995): if total_num_per_dp == max_tokens,
|
|
# it will be both all2all and mc2 scenario.
|
|
# but we fires all2all scenario first.
|
|
# the result will be the same.
|
|
# all2all scenario in vllm-ascend.
|
|
if _EXTRA_CTX.moe_comm_type == MoECommType.ALLTOALL:
|
|
gather_topk_ids_shape = (
|
|
(token_num_per_dp, topk_ids.shape[1])
|
|
if token_num_per_dp >= self.tp_size
|
|
else (self.tp_size, topk_ids.shape[1])
|
|
)
|
|
# mc2 scenario in vllm-ascend
|
|
else:
|
|
gather_topk_ids_shape = (n * self.tp_size, topk_ids.shape[1])
|
|
|
|
gather_topk_ids = torch.empty(
|
|
gather_topk_ids_shape,
|
|
dtype=topk_ids.dtype,
|
|
device=topk_ids.device,
|
|
)
|
|
split_topk_ids = torch.tensor_split(gather_topk_ids, self.tp_size, dim=0)
|
|
dist.all_gather(list(split_topk_ids), topk_ids, get_tp_group().device_group)
|
|
topk_ids = gather_topk_ids
|
|
start_loc = 0
|
|
end_loc = token_num_per_dp
|
|
else:
|
|
sp_expected = (token_num_per_dp + self.tp_size - 1) // self.tp_size if self.tp_size > 0 else -1
|
|
raise AssertionError(
|
|
"RoutedExpertsCapturer: unexpected topk_ids batch "
|
|
f"dim {n} (expected {total}, {token_num_per_dp}, "
|
|
f"{total_with_padding}, or {sp_expected} for "
|
|
f"dp_rank={self.dp_rank}, tp_size={self.tp_size})"
|
|
)
|
|
|
|
# Defensive: model may expose more layers than the capture buffer
|
|
# was sized for (unusual, but guards against miss-config).
|
|
if layer_id >= self.device_buffer.shape[1]:
|
|
return
|
|
|
|
self.device_buffer[:token_num_per_dp, layer_id, :] = topk_ids[start_loc:end_loc, :]
|
|
|
|
|
|
RoutedExpertsCapturer.capture = capture
|