223 lines
10 KiB
Python
223 lines
10 KiB
Python
#
|
||
# Copyright (c) 2026 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.
|
||
#
|
||
"""Ascend MoE-LoRA wrapper (v1).
|
||
|
||
Design (see plan in conversation history):
|
||
|
||
- Inherits weight allocation / set_lora / slice helpers from upstream
|
||
FusedMoEWithLoRA. Only the injection mechanism differs: upstream wraps
|
||
Triton modular kernel internals (`TritonExperts.activation` / `moe_sum`),
|
||
which do not exist on Ascend. We instead wrap the per-layer
|
||
`quant_method.apply` and, inside it, temporarily swap the active
|
||
`MoECommMethod._apply_mlp` so the LoRA delta is added on permuted
|
||
activations between the grouped GMMs.
|
||
|
||
- Per-layer ownership is critical: `_MoECommMethods` is a module-level
|
||
singleton shared by all 48 MoE layers. If we wrapped `_apply_mlp` at
|
||
init time, layer N+1 would compose on top of layer N's wrapper and
|
||
every forward would stack all layers' LoRA deltas. We bracket the swap
|
||
inside `apply_wrapper` so only the active layer is in effect.
|
||
|
||
- v1 deliberately limits scope to: unquant + AllGather + TP-only +
|
||
no shared experts + no FusedMC2 + no dynamic EPLB. These are the exact
|
||
conditions under which `Qwen3-30B-A3B-Thinking-2507` runs cleanly with
|
||
TP=4 EP=1 on 4×64GB. Other paths assert early so users get a clear
|
||
error rather than silently wrong outputs.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import torch
|
||
from torch import nn
|
||
from vllm import envs
|
||
from vllm.distributed.parallel_state import (
|
||
get_tensor_model_parallel_rank,
|
||
get_tensor_model_parallel_world_size,
|
||
)
|
||
from vllm.lora.layers.base import BaseLayerWithLoRA
|
||
from vllm.lora.layers.fused_moe import FusedMoE3DWithLoRA, FusedMoEWithLoRA
|
||
from vllm.lora.layers.utils import _get_lora_device
|
||
|
||
import vllm_ascend.envs as envs_ascend
|
||
|
||
|
||
def _assert_ascend_moe_lora_supported(base_layer: nn.Module) -> None:
|
||
if getattr(base_layer, "use_ep", False):
|
||
raise AssertionError(
|
||
"Ascend MoE LoRA v1 does not support expert parallelism. "
|
||
"Launch with `--enable-expert-parallel=false` and use TP only "
|
||
"(e.g. TP=4 for Qwen3-30B-A3B on 4x64GB)."
|
||
)
|
||
if getattr(base_layer, "dynamic_eplb", False):
|
||
raise AssertionError(
|
||
"Ascend MoE LoRA v1 is incompatible with dynamic EPLB "
|
||
"(expert migration would break the per-expert LoRA layout)."
|
||
)
|
||
if int(envs_ascend.VLLM_ASCEND_ENABLE_FUSED_MC2) != 0:
|
||
raise AssertionError(
|
||
"Ascend MoE LoRA v1 cannot patch FusedMC2 path "
|
||
"(dispatch_ffn_combine is a single fused C++ op). "
|
||
"Set VLLM_ASCEND_ENABLE_FUSED_MC2=0."
|
||
)
|
||
if getattr(base_layer, "_shared_experts", None) is not None:
|
||
raise AssertionError(
|
||
"Ascend MoE LoRA v1 does not wrap the shared_experts path "
|
||
"(it runs outside quant_method.apply). The target model "
|
||
"Qwen3-30B-A3B-Thinking-2507 has no shared experts; models "
|
||
"like DeepSeek-V3 are not yet supported."
|
||
)
|
||
if getattr(base_layer, "multistream_overlap_gate", False):
|
||
raise AssertionError(
|
||
"multistream_overlap_gate=True interleaves quant_method.apply "
|
||
"calls on multiple streams; the MoE LoRA path has not been "
|
||
"validated under this overlap. Disable it for MoE LoRA."
|
||
)
|
||
|
||
|
||
def _recover_moe_lora_routing(lora_context, expanded_row_idx, topk_ids):
|
||
"""Recover per-permuted-row (expert_id, lora_slot) for the dispatched rows.
|
||
|
||
npu_moe_init_routing semantics (verified empirically): ``expanded_row_idx``
|
||
is indexed by the ORIGINAL flat (token, k) position and gives where that
|
||
pair landed in the expert-sorted array -- not the reverse. So recovering
|
||
"which (token, k) pair does sorted row i hold" needs the inverse permutation
|
||
of ``expanded``, not a direct gather by it. ``argsort`` output shape ==
|
||
input shape (value-independent), so this stays graph-capturable -- no
|
||
``.item()``/data-dependent host sync.
|
||
"""
|
||
top_k = lora_context.top_k
|
||
expanded = torch.abs(expanded_row_idx)
|
||
inv_perm = torch.argsort(expanded)
|
||
expert_per_row = topk_ids.reshape(-1)[inv_perm].to(torch.long)
|
||
|
||
# token_lora_indices is a 1D LongTensor sized to max_num_batched_tokens
|
||
# (host-known constant). Clamping defensively to the last index is a no-op
|
||
# in normal operation but keeps the gather graph-safe.
|
||
orig_token = inv_perm // top_k
|
||
token_lora_indices = lora_context.punica_wrapper.token_lora_indices
|
||
orig_token = orig_token.clamp_(max=token_lora_indices.numel() - 1)
|
||
lora_per_row = token_lora_indices[orig_token]
|
||
return expert_per_row, lora_per_row
|
||
|
||
|
||
def moe_lora_apply_w13(lora_context, *, gate_up_out, hidden_states, expanded_row_idx, topk_ids):
|
||
"""Add the w13 LoRA delta into ``gate_up_out`` (in place), before activation.
|
||
|
||
Called from ``unquant_apply_mlp`` right after the base gate_up GMM. Returns
|
||
the recovered per-row routing so the w2 delta can reuse it.
|
||
"""
|
||
routing = _recover_moe_lora_routing(lora_context, expanded_row_idx, topk_ids)
|
||
expert_per_row, lora_per_row = routing
|
||
lora_context.punica_wrapper.add_lora_fused_moe(
|
||
y=gate_up_out,
|
||
x=hidden_states,
|
||
lora_a_stacked=lora_context.w13_lora_a_stacked,
|
||
lora_b_stacked=lora_context.w13_lora_b_stacked,
|
||
expert_ids=expert_per_row,
|
||
adapter_enabled=lora_context.adapter_enabled,
|
||
token_lora_mapping=lora_per_row,
|
||
)
|
||
return routing
|
||
|
||
|
||
def moe_lora_apply_w2(lora_context, *, down_out, silu_out, lora_routing):
|
||
"""Add the w2 LoRA delta into ``down_out`` (in place), after the down GMM.
|
||
|
||
Reuses the per-row routing computed by ``moe_lora_apply_w13``; ``silu_out``
|
||
is the activation output that fed the base down GMM.
|
||
"""
|
||
expert_per_row, lora_per_row = lora_routing
|
||
lora_context.punica_wrapper.add_lora_fused_moe(
|
||
y=down_out,
|
||
x=silu_out,
|
||
lora_a_stacked=lora_context.w2_lora_a_stacked,
|
||
lora_b_stacked=lora_context.w2_lora_b_stacked,
|
||
expert_ids=expert_per_row,
|
||
adapter_enabled=lora_context.adapter_enabled,
|
||
token_lora_mapping=lora_per_row,
|
||
)
|
||
|
||
|
||
class AscendFusedMoEWithLoRA(FusedMoEWithLoRA):
|
||
"""Ascend-native MoE-LoRA wrapper.
|
||
|
||
Reuses upstream weight allocation, set_lora, reset_lora, and slicing.
|
||
Instead of the GPU modular-kernel injection, it publishes a per-layer
|
||
``MoELoRAContext`` onto the base layer (``_ascend_moe_lora_context``).
|
||
The Ascend unquant MoE path threads that context through
|
||
``MoEFusedExpertsInput`` -> ``MoEMlpComputeInput`` and applies the LoRA
|
||
delta natively inside ``unquant_apply_mlp`` (see
|
||
``moe_lora_apply_w13`` / ``moe_lora_apply_w2`` below) -- no runtime
|
||
monkey-patch of ``comm._apply_mlp``.
|
||
"""
|
||
|
||
def __init__(self, base_layer: nn.Module) -> None:
|
||
# Skip FusedMoEWithLoRA.__init__: it immediately asserts Triton
|
||
# internals and calls _inject_lora_into_fused_moe which is GPU-only.
|
||
BaseLayerWithLoRA.__init__(self)
|
||
self.base_layer = base_layer
|
||
_assert_ascend_moe_lora_supported(base_layer)
|
||
self.tp_size = get_tensor_model_parallel_world_size()
|
||
self.tp_rank = get_tensor_model_parallel_rank()
|
||
self.device = _get_lora_device(base_layer)
|
||
self._enable_aux_cuda_stream = envs.VLLM_LORA_ENABLE_DUAL_STREAM
|
||
self.moe_config = base_layer.moe_config
|
||
self._w13_slices = 2 if base_layer.moe_config.is_act_and_mul else 1
|
||
|
||
# ------------------------------------------------------------------
|
||
# Mapping
|
||
# ------------------------------------------------------------------
|
||
def set_mapping(self, punica_wrapper):
|
||
# Upstream FusedMoEWithLoRA.set_mapping (vllm v0.22.0+) chains into
|
||
# ``self._moe_kernel.fused_experts.set_lora_context(...)``, but
|
||
# ``_moe_kernel`` is only set by the GPU modular-kernel path that we
|
||
# deliberately skip in __init__. We instead build the per-layer
|
||
# MoELoRAContext (now that punica_wrapper is available) and publish it
|
||
# on the module that ``AscendUnquantizedFusedMoEMethod.apply`` reads via
|
||
# ``getattr(layer, "_ascend_moe_lora_context", None)`` -- the base layer
|
||
# itself on 0.23.0, but ``base_layer.routed_experts`` on main (there the
|
||
# runner *is* the layer and it calls apply with ``layer=routed_experts``).
|
||
# The context holds stable references (the in-place-updated LoRA stacks,
|
||
# adapter_enabled and the punica wrapper), so building it once here is
|
||
# sufficient.
|
||
BaseLayerWithLoRA.set_mapping(self, punica_wrapper)
|
||
self.base_layer.set_lora_context(self._build_lora_context())
|
||
|
||
|
||
class AscendFusedMoE3DWithLoRA(AscendFusedMoEWithLoRA, FusedMoE3DWithLoRA):
|
||
"""For checkpoints that already fuse w1+w3 into a 3D weight (single slice)."""
|
||
|
||
def __init__(self, base_layer: nn.Module) -> None:
|
||
AscendFusedMoEWithLoRA.__init__(self, base_layer)
|
||
# Override: 3D MoE LoRA uses a single w13 slice.
|
||
self._w13_slices = 1
|
||
|
||
|
||
# ----------------------------------------------------------------------
|
||
# Upstream compatibility shim: vllm/lora/model_manager.py:create_dummy_lora
|
||
# branches on `module.__class__.__name__ == "FusedMoEWithLoRA"` (and the
|
||
# 3D variant). Without this override, our subclasses would skip the
|
||
# pack_moe path and hit the generic pack() fallback, which produces a
|
||
# flat list of N_experts * 3 sub-LoRAs -- `set_lora` then fails with
|
||
# "too many values to unpack (expected 3)".
|
||
#
|
||
# Overriding only __name__ keeps the actual class object distinct (so
|
||
# isinstance / type identity / debugging are unaffected) but lets the
|
||
# upstream string compare hit our objects.
|
||
# ----------------------------------------------------------------------
|
||
AscendFusedMoEWithLoRA.__name__ = "FusedMoEWithLoRA"
|
||
AscendFusedMoE3DWithLoRA.__name__ = "FusedMoE3DWithLoRA"
|