222
vllm_ascend/lora/fused_moe.py
Normal file
222
vllm_ascend/lora/fused_moe.py
Normal file
@@ -0,0 +1,222 @@
|
||||
#
|
||||
# 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"
|
||||
Reference in New Issue
Block a user