431 lines
16 KiB
Python
431 lines
16 KiB
Python
import math
|
|
from contextlib import contextmanager
|
|
from contextvars import ContextVar
|
|
from enum import Enum
|
|
from typing import Any
|
|
|
|
import torch
|
|
import vllm.envs as envs_vllm
|
|
from vllm.config import CUDAGraphMode, VllmConfig
|
|
from vllm.distributed import get_dp_group, get_ep_group, get_tensor_model_parallel_world_size
|
|
from vllm.forward_context import BatchDescriptor, get_forward_context, set_forward_context
|
|
from vllm.logger import logger
|
|
|
|
from vllm_ascend.ascend_config import get_ascend_config
|
|
from vllm_ascend.utils import (
|
|
AscendDeviceType,
|
|
enable_sp,
|
|
flashcomm2_enable,
|
|
get_ascend_device_type,
|
|
has_layer_idx,
|
|
is_drafter_moe_model,
|
|
is_moe_model,
|
|
speculative_enable_dispatch_gmm_combine_decode,
|
|
)
|
|
|
|
|
|
class MoECommType(Enum):
|
|
ALLGATHER = 0
|
|
MC2 = 1
|
|
ALLTOALL = 2
|
|
FUSED_MC2 = 3
|
|
|
|
|
|
_MRV2_IN_PROFILE_RUN: ContextVar[bool] = ContextVar("_MRV2_IN_PROFILE_RUN", default=False)
|
|
|
|
|
|
@contextmanager
|
|
def override_mrv2_in_profile_run(enabled: bool):
|
|
"""Override MRv2's extra profile-run marker for one forward path.
|
|
|
|
MRv2 builds the base forward context inside upstream vLLM, so Ascend's
|
|
platform hook cannot tell whether the current forward is the extra MC2
|
|
profile dummy run. A ContextVar keeps this MRv2-only state scoped to the
|
|
current forward path without adding default fallback behavior.
|
|
"""
|
|
token = _MRV2_IN_PROFILE_RUN.set(enabled)
|
|
try:
|
|
yield
|
|
finally:
|
|
_MRV2_IN_PROFILE_RUN.reset(token)
|
|
|
|
|
|
def get_mrv2_in_profile_run() -> bool:
|
|
return _MRV2_IN_PROFILE_RUN.get()
|
|
|
|
|
|
@contextmanager
|
|
def set_ascend_forward_context(
|
|
attn_metadata: Any,
|
|
vllm_config: VllmConfig,
|
|
num_tokens: int = 0,
|
|
num_tokens_across_dp: torch.Tensor | None = None,
|
|
in_profile_run: bool = False,
|
|
num_actual_tokens: int | None = None,
|
|
aclgraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
|
|
batch_descriptor: BatchDescriptor | None = None,
|
|
model_instance: torch.nn.Module = None,
|
|
is_draft_model=False,
|
|
skip_compiled: bool = False,
|
|
max_tokens_across_pcp: int = 0,
|
|
draft_attn_metadatas=None,
|
|
has_sinks=False,
|
|
input_ids=None,
|
|
eplb_heat_collection_status: bool = False,
|
|
):
|
|
"""A context manager that stores the current forward context,
|
|
can be attention metadata, etc.
|
|
We add some additional param into forward_context.
|
|
"""
|
|
forward_context_kwargs = {
|
|
"attn_metadata": attn_metadata,
|
|
"vllm_config": vllm_config,
|
|
"num_tokens": num_tokens,
|
|
"num_tokens_across_dp": num_tokens_across_dp,
|
|
"cudagraph_runtime_mode": aclgraph_runtime_mode,
|
|
"batch_descriptor": batch_descriptor,
|
|
"skip_compiled": skip_compiled,
|
|
}
|
|
with set_forward_context(**forward_context_kwargs):
|
|
forward_context = get_forward_context()
|
|
forward_context.draft_attn_metadatas = draft_attn_metadatas
|
|
|
|
forward_context.input_ids = input_ids
|
|
|
|
from vllm_ascend.ops.fused_moe.moe_comm_method import get_moe_comm_method
|
|
|
|
max_num_tokens = int(num_tokens_across_dp.max().item()) if num_tokens_across_dp is not None else num_tokens
|
|
moe_comm_type = select_moe_comm_method(max_num_tokens, vllm_config, is_draft_model)
|
|
|
|
forward_context.moe_comm_type = moe_comm_type
|
|
forward_context.moe_comm_method = get_moe_comm_method(moe_comm_type)
|
|
|
|
tp_world_size = get_tensor_model_parallel_world_size()
|
|
|
|
forward_context.in_profile_run = in_profile_run
|
|
|
|
# NOTE: This cannot be set using set_forward_context
|
|
# due to multiple warmups before actual capturing
|
|
forward_context.capturing = False
|
|
|
|
# TODO: remove it when fia merge in fiav2
|
|
forward_context.sinks = has_sinks
|
|
|
|
# TODO: remove it when torch_npu.npu_mm_reduce_scatter_base supports tp_size >= 16.
|
|
mmrs_fusion = tp_world_size <= 8
|
|
|
|
# set for sequence parallelism, 1000 is the batch size concurrency threshold
|
|
# for enabling the flashcomm_v1 or sequence_parallelism feature.
|
|
# Currently, it is an empirical value. In normal scenarios, if the concurrency
|
|
# exceeds this threshold, the performance benefits can be maximized.
|
|
# Conversely, if the concurrency is below the threshold,
|
|
# the performance may degrade due to the switching of communication methods.
|
|
|
|
# main model and drafter model may have different architecture
|
|
is_context_moe_model = is_drafter_moe_model(vllm_config) if is_draft_model else is_moe_model(vllm_config)
|
|
if is_context_moe_model:
|
|
flash_comm_v1_enabled = enable_sp(vllm_config) and num_tokens is not None
|
|
mmrs_fusion = False
|
|
elif is_draft_model:
|
|
# TODO: for dense drafter, `sp` is redundant and is not compatible with `dp` and `graph`.
|
|
# Disable it to avoid more problems.
|
|
flash_comm_v1_enabled = False
|
|
else:
|
|
flash_comm_v1_enabled = enable_sp(vllm_config) and num_tokens is not None and num_tokens > 1000
|
|
forward_context.mmrs_fusion = mmrs_fusion
|
|
forward_context.num_tokens = num_tokens
|
|
forward_context.flash_comm_v1_enabled = flash_comm_v1_enabled
|
|
# TODO(Levi-JQ): another PR to normalize the enabling logic for sp/fc2
|
|
forward_context.flashcomm_v2_enabled = flashcomm2_enable() and tp_world_size > 1 and num_tokens is not None
|
|
|
|
forward_context.pad_size = 0
|
|
if forward_context.flash_comm_v1_enabled or forward_context.flashcomm_v2_enabled:
|
|
pad_size = (tp_world_size - (num_tokens % tp_world_size)) % tp_world_size
|
|
forward_context.pad_size = pad_size
|
|
|
|
# set this for rope forward_oot using
|
|
forward_context.is_first_layer = True
|
|
|
|
# set layer_idx to enable optimization features that depend on this information.
|
|
# This is only applicable to models that contain these necessary attributes.
|
|
forward_context.layer_idx = None
|
|
if has_layer_idx(model_instance):
|
|
forward_context.layer_idx = model_instance.model.start_layer
|
|
|
|
forward_context.prefetch_mlp_gate_up_proj = False
|
|
forward_context.prefetch_mlp_down_proj = False
|
|
forward_context.model_instance = model_instance
|
|
forward_context.is_draft_model = is_draft_model
|
|
forward_context.is_draft_model_prefill = False
|
|
|
|
if num_tokens is None and attn_metadata is not None:
|
|
num_tokens = attn_metadata.num_actual_tokens
|
|
|
|
dp_world_size = get_dp_group().world_size
|
|
if dp_world_size > 1 and forward_context.dp_metadata is not None:
|
|
dp_meta = forward_context.dp_metadata
|
|
max_tokens_across_dp = dp_meta.num_tokens_across_dp_cpu.max().item()
|
|
if forward_context.flash_comm_v1_enabled or forward_context.flashcomm_v2_enabled:
|
|
padded_length = (max_tokens_across_dp + tp_world_size - 1) // tp_world_size * tp_world_size
|
|
pad_size = padded_length - num_tokens
|
|
forward_context.padded_length = padded_length
|
|
forward_context.pad_size = pad_size
|
|
else:
|
|
max_tokens_across_dp = num_tokens
|
|
|
|
forward_context.max_tokens_across_dp = max_tokens_across_dp
|
|
forward_context.max_tokens_across_pcp = max_tokens_across_pcp
|
|
|
|
forward_context.eplb_heat_collection_status = eplb_heat_collection_status
|
|
|
|
if num_tokens is not None:
|
|
if num_actual_tokens is None:
|
|
num_actual_tokens = num_tokens
|
|
# NOTE: token num which need to pad to when mc2
|
|
forward_context.padded_num_tokens = math.ceil(max_tokens_across_dp / tp_world_size) * tp_world_size
|
|
reserved_mc2_mask = get_mc2_mask()
|
|
if reserved_mc2_mask is not None:
|
|
mc2_mask = reserved_mc2_mask[: forward_context.padded_num_tokens]
|
|
mc2_mask[:num_actual_tokens] = True
|
|
mc2_mask[num_actual_tokens:] = False
|
|
forward_context.mc2_mask = mc2_mask
|
|
try:
|
|
yield
|
|
finally:
|
|
pass
|
|
|
|
|
|
_mc2_tokens_capacity: int | None = None
|
|
_reserved_mc2_mask: torch.Tensor | None = None
|
|
|
|
|
|
def set_mc2_tokens_capacity(vllm_config, max_num_reqs, uniform_decode_query_len):
|
|
global _mc2_tokens_capacity
|
|
if _mc2_tokens_capacity is not None:
|
|
return
|
|
if get_ascend_config().enable_prefill_mc2:
|
|
max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
|
|
elif vllm_config.compilation_config.cudagraph_capture_sizes:
|
|
max_num_tokens = vllm_config.compilation_config.max_cudagraph_capture_size
|
|
else:
|
|
max_num_tokens = max_num_reqs * uniform_decode_query_len
|
|
tp_size = vllm_config.parallel_config.tensor_parallel_size
|
|
# Use integer arithmetic for ceiling division.
|
|
num_tokens_per_tp_rank = (max_num_tokens + tp_size - 1) // tp_size
|
|
# NOTE: To save memory, we cap the max number of tokens to 512.
|
|
num_tokens_per_tp_rank = min(num_tokens_per_tp_rank, 512)
|
|
_mc2_tokens_capacity = num_tokens_per_tp_rank * tp_size
|
|
|
|
|
|
def get_mc2_tokens_capacity():
|
|
return _mc2_tokens_capacity
|
|
|
|
|
|
def set_mc2_mask(vllm_config, device):
|
|
global _reserved_mc2_mask
|
|
if _reserved_mc2_mask is not None:
|
|
return
|
|
if is_moe_model(vllm_config):
|
|
_reserved_mc2_mask = torch.zeros(
|
|
vllm_config.scheduler_config.max_num_batched_tokens, dtype=torch.bool, device=device
|
|
)
|
|
else:
|
|
_reserved_mc2_mask = None
|
|
|
|
|
|
def get_mc2_mask():
|
|
return _reserved_mc2_mask
|
|
|
|
|
|
def _select_a2_moe_comm_method(
|
|
num_tokens: int,
|
|
vllm_config: VllmConfig,
|
|
mc2_tokens_capacity: int,
|
|
) -> MoECommType:
|
|
num_experts = vllm_config.model_config.get_num_experts()
|
|
ep_world_size = (
|
|
vllm_config.parallel_config.world_size_across_dp // vllm_config.parallel_config.pipeline_parallel_size
|
|
)
|
|
num_experts_per_device = num_experts // ep_world_size
|
|
if num_experts_per_device <= 24 and ep_world_size >= 16 and num_tokens <= mc2_tokens_capacity:
|
|
return MoECommType.MC2
|
|
return MoECommType.ALLGATHER
|
|
|
|
|
|
def _select_a3_moe_comm_method(
|
|
num_tokens: int,
|
|
vllm_config: VllmConfig,
|
|
quant_type: str | None,
|
|
mc2_tokens_capacity: int,
|
|
enable_fused_mc2: int,
|
|
) -> MoECommType:
|
|
# TODO: drop the EP-size guard when dispatch_ffn_combine supports larger EP sizes
|
|
# TODO: drop speculative method guard when dispatch_gmm_combine_decode supports w16a16
|
|
dispatch_ffn_combine_enable = get_ep_group().world_size <= 32
|
|
if num_tokens <= mc2_tokens_capacity:
|
|
fused_decode_enable = enable_fused_mc2
|
|
if enable_fused_mc2 == 1:
|
|
fused_decode_enable = enable_fused_mc2 and dispatch_ffn_combine_enable
|
|
elif enable_fused_mc2 == 2:
|
|
fused_decode_enable = (
|
|
enable_fused_mc2
|
|
and speculative_enable_dispatch_gmm_combine_decode(vllm_config)
|
|
and quant_type == "w8a8_dynamic"
|
|
)
|
|
return MoECommType.FUSED_MC2 if fused_decode_enable else MoECommType.MC2
|
|
|
|
fused_prefill_enable = enable_fused_mc2
|
|
if enable_fused_mc2 == 1:
|
|
fused_prefill_enable = enable_fused_mc2 and dispatch_ffn_combine_enable
|
|
elif enable_fused_mc2 == 2:
|
|
fused_prefill_enable = False
|
|
return MoECommType.FUSED_MC2 if fused_prefill_enable else MoECommType.ALLTOALL
|
|
|
|
|
|
def _select_a5_moe_comm_method(
|
|
num_tokens: int,
|
|
vllm_config: VllmConfig,
|
|
mc2_tokens_capacity: int,
|
|
) -> MoECommType:
|
|
num_experts_per_tok = getattr(
|
|
vllm_config.model_config.hf_text_config,
|
|
"num_experts_per_tok",
|
|
getattr(vllm_config.model_config.hf_text_config, "top_k_experts", 1),
|
|
)
|
|
world_size = vllm_config.parallel_config.world_size_across_dp
|
|
if num_tokens <= mc2_tokens_capacity and world_size > 1:
|
|
return MoECommType.MC2
|
|
if world_size <= num_experts_per_tok:
|
|
return MoECommType.ALLGATHER
|
|
return MoECommType.ALLTOALL
|
|
|
|
|
|
def select_moe_comm_method(num_tokens: int, vllm_config: VllmConfig, is_draft_model=False) -> MoECommType | None:
|
|
"""Select the MoE communication method according to parallel settings,
|
|
device generation, token count, and quantization.
|
|
|
|
1. Non-MoE models return `None`.
|
|
2. Without expert parallel, fall back to all-gather.
|
|
3. On A2 with expert parallel, pick MC2 when tokens fit the MC2 capacity
|
|
and the DP size is large enough; otherwise use all-gather.
|
|
4. On A3 with expert parallel, prefer fused MC2 when using w8a8_dynamic
|
|
quantization with small EP size, no dynamic_eplb, and not in MTP
|
|
mode; otherwise use MC2 within capacity or all-to-all.
|
|
5. On 310P, always use all-gather.
|
|
6. On A5 with expert parallel, use MC2 when tokens fit the MC2 capacity
|
|
and the EP size is large enough; otherwise use all-gather when
|
|
EP size is smaller than num of topK experts or all-to-all.
|
|
|
|
Args:
|
|
num_tokens (int): The number of tokens in the current batch.
|
|
vllm_config (VllmConfig): Runtime configuration for the model.
|
|
is_draft_model (bool): Whether the model runs in MTP mode.
|
|
|
|
Raises:
|
|
ValueError: If the soc version is unsupported.
|
|
|
|
Returns:
|
|
MoECommType | None: The selected MoE communication method.
|
|
"""
|
|
if not is_moe_model(vllm_config):
|
|
return None
|
|
|
|
mc2_tokens_capacity = get_mc2_tokens_capacity()
|
|
soc_version = get_ascend_device_type()
|
|
quant_type = getattr(
|
|
vllm_config.model_config.hf_text_config,
|
|
"moe_quantize",
|
|
getattr(vllm_config.model_config.hf_text_config, "quantize", None),
|
|
)
|
|
|
|
if not vllm_config.parallel_config.enable_expert_parallel or get_ep_group().world_size == 1:
|
|
moe_comm_type = MoECommType.ALLGATHER
|
|
elif soc_version == AscendDeviceType.A2:
|
|
moe_comm_type = _select_a2_moe_comm_method(num_tokens, vllm_config, mc2_tokens_capacity)
|
|
elif soc_version == AscendDeviceType.A3:
|
|
moe_comm_type = _select_a3_moe_comm_method(
|
|
num_tokens,
|
|
vllm_config,
|
|
quant_type,
|
|
mc2_tokens_capacity,
|
|
get_ascend_config().enable_fused_mc2,
|
|
)
|
|
elif soc_version == AscendDeviceType.A5:
|
|
moe_comm_type = _select_a5_moe_comm_method(num_tokens, vllm_config, mc2_tokens_capacity)
|
|
elif soc_version == AscendDeviceType._310P:
|
|
moe_comm_type = MoECommType.ALLGATHER
|
|
|
|
else:
|
|
raise ValueError(f"Unsupported soc_version: {soc_version}")
|
|
logger.debug(
|
|
"MoE comm method selected: soc=%s, method=%s, num_tokens=%d, mc2_capacity=%s",
|
|
soc_version,
|
|
moe_comm_type,
|
|
num_tokens,
|
|
mc2_tokens_capacity,
|
|
)
|
|
return moe_comm_type
|
|
|
|
|
|
class _ExtraForwardContextProxy:
|
|
"""Unified forward-context access for v1/v2 model runners."""
|
|
|
|
extra_attrs = (
|
|
"capturing",
|
|
"moe_comm_type",
|
|
"moe_comm_method",
|
|
"mmrs_fusion",
|
|
"num_tokens",
|
|
"flash_comm_v1_enabled",
|
|
"flashcomm_v2_enabled",
|
|
"pad_size",
|
|
"padded_length",
|
|
"num_tokens_across_dp",
|
|
"mc2_mask",
|
|
"is_draft_model",
|
|
"is_draft_model_prefill",
|
|
"prefetch_mlp_gate_up_proj",
|
|
"prefetch_mlp_down_proj",
|
|
"model_instance",
|
|
"layer_idx",
|
|
"max_tokens_across_dp",
|
|
"max_tokens_across_pcp",
|
|
"num_accept_tokens",
|
|
"in_profile_run",
|
|
"padded_num_tokens",
|
|
"sinks",
|
|
"eplb_heat_collection_status",
|
|
)
|
|
|
|
def check_extra_attr(self, name: str):
|
|
if name not in self.extra_attrs:
|
|
raise AttributeError(
|
|
f"{name} is not extra forward context attribute, "
|
|
"please get/set it from vllm's _forward_context directly."
|
|
)
|
|
|
|
@staticmethod
|
|
def _ctx():
|
|
return get_forward_context()
|
|
|
|
def __getattr__(self, name: str) -> Any:
|
|
self.check_extra_attr(name)
|
|
ctx = self._ctx()
|
|
if envs_vllm.VLLM_USE_V2_MODEL_RUNNER:
|
|
# Unset known extras default to None so optional flags (e.g. `sinks`)
|
|
# can be read with truthiness checks before the V2 path populates them.
|
|
return ctx.additional_kwargs.get(name)
|
|
return getattr(ctx, name, None)
|
|
|
|
def __setattr__(self, name: str, value: Any) -> None:
|
|
self.check_extra_attr(name)
|
|
ctx = self._ctx()
|
|
if envs_vllm.VLLM_USE_V2_MODEL_RUNNER:
|
|
ctx.additional_kwargs[name] = value
|
|
else:
|
|
setattr(ctx, name, value)
|
|
|
|
|
|
# usage: from vllm_ascend.ascend_forward_context import _EXTRA_CTX
|
|
_EXTRA_CTX = _ExtraForwardContextProxy()
|