# Copyright (c) ModelScope Contributors. All rights reserved.
"""vLLM-Ascend MoE patches used by SWIFT NPU rollout.

There are two independent responsibilities in this file:

* runtime routing: avoid the unstable custom non-quantized MoE routing op on
  stacks where vLLM-Ascend still dispatches that branch to
  ``aclnnMoeInitRoutingCustom``;
* weight sync: adapt 2D HF/Megatron MoE expert weights to the already-processed
  3D vLLM-Ascend expert parameter layout during GRPO colocate updates.

Both patches are guarded by vLLM-Ascend implementation checks and only touch the
specific MoE paths they need.
"""
from __future__ import annotations

import inspect
import torch

from swift.utils.logger import get_logger

logger = get_logger()

_VLLM_ASCEND_MOE_SYNC_LAYOUT_ATTR = '_swift_vllm_ascend_moe_weight_sync_layout'
_VLLM_ASCEND_MOE_SKIP_POST_LOAD_ATTR = '_swift_vllm_ascend_moe_skip_post_load'
_VLLM_ASCEND_MOE_PROCESSED_LAYOUT = 'megatron_processed'
_VLLM_ASCEND_MOE_PREPROCESSED_LAYOUT = 'fsdp2_preprocessed'
_QWEN_MOE_MODEL_TYPES = {'qwen3_moe', 'qwen3_5_moe'}


def _patch_vllm_ascend_device_op_nonquant_routing() -> None:
    """Use the stable torch-npu routing op for non-quantized MoE when needed.

    Some released vLLM-Ascend versions route the non-quantized MoE case
    (``scale is None`` and ``quant_mode == -1``) through
    ``npu_moe_init_routing_custom`` / ``aclnnMoeInitRoutingCustom``, which is
    not stable for the parameter combination used by Qwen-style MoE rollout.

    This is intentionally gated by implementation detection instead of a fixed
    version threshold: source builds or future/backported versions may already
    dispatch the non-quantized path to ``torch_npu.npu_moe_init_routing_v2``.
    When that fixed branch is present, skip patching and keep the upstream
    implementation intact.

    Do not probe the custom op by calling it first.  On Ascend, a missing custom
    binary can be reported asynchronously: even if Python catches the immediate
    RuntimeError and falls back, the failed launch can poison the stream and hang
    later at an unrelated event synchronization.  Therefore, when source
    inspection shows that the non-quantized branch still routes to the custom op,
    dispatch that branch directly to ``torch_npu.npu_moe_init_routing_v2``.
    """
    try:
        import torch_npu
        from vllm_ascend.device import device_op
    except (ImportError, AttributeError):
        return

    adaptor_cls = getattr(device_op, 'BaseDeviceAdaptor', None)
    if adaptor_cls is None:
        return
    origin_routing = getattr(adaptor_cls, 'npu_moe_init_routing', None)
    if origin_routing is None or getattr(origin_routing, '_swift_nonquant_routing_patched', False):
        return
    try:
        origin_source = inspect.getsource(origin_routing)
    except (OSError, TypeError):
        origin_source = ''
    if 'npu_moe_init_routing_v2' in origin_source and 'quant_mode == -1' in origin_source:
        return
    origin_signature = inspect.signature(origin_routing)
    routing_defaults = {
        'scale': None,
        'active_num': None,
        'expert_num': None,
        'expert_tokens_num_type': 1,
        'expert_tokens_num_flag': True,
        'active_expert_range': None,
        'quant_mode': -1,
    }
    missing_params = set(routing_defaults).difference(origin_signature.parameters)
    if missing_params:
        raise RuntimeError('Unsupported vLLM-Ascend npu_moe_init_routing signature: '
                           f'signature={origin_signature}, missing={sorted(missing_params)}.')

    def is_nonquant_routing(routing_kwargs) -> bool:
        return routing_kwargs['scale'] is None and routing_kwargs['quant_mode'] == -1

    def npu_moe_init_routing_v2(hidden_states, topk_ids, routing_kwargs):
        active_num = routing_kwargs['active_num']
        expert_num = routing_kwargs['expert_num']
        active_expert_range = routing_kwargs['active_expert_range']
        return torch_npu.npu_moe_init_routing_v2(
            hidden_states,
            topk_ids,
            scale=None,
            offset=None,
            active_num=0 if active_num is None else active_num,
            expert_capacity=-1,
            expert_num=expert_num,
            drop_pad_mode=0,
            expert_tokens_num_type=routing_kwargs['expert_tokens_num_type'],
            expert_tokens_num_flag=routing_kwargs['expert_tokens_num_flag'],
            active_expert_range=[0, expert_num] if active_expert_range is None else active_expert_range,
            quant_mode=routing_kwargs['quant_mode'],
            row_idx_type=0,
        )

    def patched_npu_moe_init_routing(hidden_states, topk_ids, *args, **kwargs):
        try:
            bound = origin_signature.bind(hidden_states, topk_ids, *args, **kwargs)
        except TypeError as e:
            raise RuntimeError('Failed to bind vLLM-Ascend npu_moe_init_routing arguments: '
                               f'signature={origin_signature}, args={args}, kwargs={kwargs}.') from e
        bound.apply_defaults()
        routing_kwargs = {key: bound.arguments.get(key, default) for key, default in routing_defaults.items()}

        if not is_nonquant_routing(routing_kwargs):
            return origin_routing(hidden_states, topk_ids, *args, **kwargs)
        logger.warning_once(
            'Using torch_npu.npu_moe_init_routing_v2 for vLLM-Ascend non-quantized MoE routing. '
            'The installed vLLM-Ascend implementation still dispatches this branch to '
            'npu_moe_init_routing_custom, whose missing custom-op binary fails asynchronously on this stack.')
        return npu_moe_init_routing_v2(hidden_states, topk_ids, routing_kwargs)

    patched_npu_moe_init_routing._swift_nonquant_routing_patched = True
    patched_npu_moe_init_routing._swift_origin = origin_routing
    adaptor_cls.npu_moe_init_routing = staticmethod(patched_npu_moe_init_routing)


def patch_vllm_ascend_moe_runtime() -> None:
    """Apply MoE runtime patches that are independent of GRPO weight sync."""
    _patch_vllm_ascend_device_op_nonquant_routing()


def _is_qwen_moe_model(model) -> bool:
    return getattr(getattr(model, 'config', None), 'model_type', None) in _QWEN_MOE_MODEL_TYPES


def configure_vllm_ascend_moe_weight_sync(vllm_model, train_model, *, is_fsdp2: bool) -> None:
    """Record the vLLM-Ascend MoE sync layout required by this training backend."""
    fsdp2_qwen_moe = is_fsdp2 and _is_qwen_moe_model(train_model)
    layout = _VLLM_ASCEND_MOE_PROCESSED_LAYOUT
    # Current vLLM-Ascend 0.18 non-quantized Qwen MoE forward keeps
    # ``need_trans=False`` and feeds ``w13_weight`` directly to
    # ``npu_grouped_matmul``.  After FSDP2 runtime sync, write Qwen MoE weights
    # directly into the runtime [hidden, I_tp] direction and skip checkpoint
    # post-load processing; otherwise post-load transposes them back to
    # [I_tp, hidden] and the first rollout fails with a hidden-size mismatch
    # such as 2048 vs 192/384.
    setattr(vllm_model, _VLLM_ASCEND_MOE_SYNC_LAYOUT_ATTR, layout)
    setattr(vllm_model, _VLLM_ASCEND_MOE_SKIP_POST_LOAD_ATTR, fsdp2_qwen_moe)


def configure_vllm_ascend_moe_preprocessed_weight_sync(vllm_model) -> None:
    """Record that reload writes the layout expected before vLLM-Ascend post-processing."""
    setattr(vllm_model, _VLLM_ASCEND_MOE_SYNC_LAYOUT_ATTR, _VLLM_ASCEND_MOE_PREPROCESSED_LAYOUT)
    setattr(vllm_model, _VLLM_ASCEND_MOE_SKIP_POST_LOAD_ATTR, False)


def use_vllm_ascend_moe_preprocessed_weight(vllm_model) -> bool:
    """Return whether runtime sync should write the pre-process MoE layout."""
    return getattr(vllm_model, _VLLM_ASCEND_MOE_SYNC_LAYOUT_ATTR,
                   _VLLM_ASCEND_MOE_PROCESSED_LAYOUT) == _VLLM_ASCEND_MOE_PREPROCESSED_LAYOUT


def should_skip_vllm_ascend_moe_post_load(vllm_model) -> bool:
    """Return whether vLLM post-load processing should be skipped after sync."""
    return bool(getattr(vllm_model, _VLLM_ASCEND_MOE_SKIP_POST_LOAD_ATTR, False))


def expand_fused_moe_expert_names_for_vllm_ascend(name: str):
    """Map Transformers fused Qwen MoE expert names to vLLM checkpoint names.

    FSDP2 can expose Qwen-style MoE expert weights as fused tensors:

        mlp.experts.gate_up_proj: [experts, 2 * intermediate, hidden]
        mlp.experts.down_proj   : [experts, hidden, intermediate]

    vLLM's Qwen MoE ``load_weights`` path expects checkpoint-style names such as
    ``mlp.experts.0.gate_proj.weight`` / ``up_proj`` / ``down_proj`` and maps
    those names onto its internal ``w13_weight`` / ``w2_weight`` parameters.
    Use expert 0 only as a name anchor; the paired vLLM-Ascend weight-loader
    patch below copies all local experts from the full 3D tensor.
    """
    gate_up_suffix = '.mlp.experts.gate_up_proj'
    down_suffix = '.mlp.experts.down_proj'
    if name.endswith(gate_up_suffix):
        prefix = name[:-len('gate_up_proj')]
        return [
            f'{prefix}0.gate_proj.weight',
            f'{prefix}0.up_proj.weight',
        ]
    if name.endswith(down_suffix):
        prefix = name[:-len('down_proj')]
        return [f'{prefix}0.down_proj.weight']
    return None


def expand_fused_moe_expert_weight_for_vllm_ascend(name: str, param):
    """Expand one FSDP2 fused Qwen MoE expert tensor for vLLM-Ascend weight sync."""
    if not isinstance(param, torch.Tensor) or param.dim() != 3:
        return None
    expanded_names = expand_fused_moe_expert_names_for_vllm_ascend(name)
    if expanded_names is None:
        return None
    if name.endswith('.mlp.experts.gate_up_proj'):
        gate_proj, up_proj = param.chunk(2, dim=1)
        return [
            (expanded_names[0], gate_proj.contiguous()),
            (expanded_names[1], up_proj.contiguous()),
        ]
    if name.endswith('.mlp.experts.down_proj'):
        return [(expanded_names[0], param)]
    return None


def patch_vllm_ascend_moe_expert_weight_loader(experts,
                                               name: str,
                                               param,
                                               *,
                                               load_preprocessed_weight: bool = False) -> None:
    """Patch one processed vLLM-Ascend MoE expert parameter loader.

    vLLM-Ascend transposes unquantized MoE weights after each model load
    so grouped matmul can consume them efficiently.  During GRPO weight sync,
    however, SWIFT can send regular HF/Megatron expert weights, for example:

        gate_proj/up_proj: [intermediate, hidden] -> w13_weight
        down_proj       : [hidden, intermediate] -> w2_weight

    FSDP2 Qwen MoE may expose the same weights as fused 3D tensors.  SWIFT
    expands those tensors to checkpoint-style gate/up/down names before calling
    vLLM ``load_weights``:

        gate_proj/up_proj: [experts, intermediate, hidden]
        down_proj       : [experts, hidden, intermediate]

    Full-weight server reload still writes the pre-processed layout and then
    calls ``process_weights_after_loading`` once, letting vLLM-Ascend transpose
    complete weights afterwards:

        w13_weight before process: [local_experts, 2 * intermediate_per_tp, hidden]
        w2_weight before process : [local_experts, hidden, intermediate_per_tp]

    Megatron colocate runtime sync loads into the already-processed layout used
    by the existing Megatron rollout path:

        w13_weight after process: [local_experts, hidden, 2 * intermediate_per_tp]
        w2_weight after process : [local_experts, intermediate_per_tp, hidden]

    ``load_preprocessed_weight`` selects the server full-reload target.  FSDP2
    Qwen MoE colocate runtime sync keeps the processed target and deliberately
    skips the post-load transpose because current vLLM-Ascend non-quantized
    grouped matmul consumes the [hidden, I_tp] direction in this path.

    This wrapper keeps the normal vLLM loader for initial checkpoint load,
    quantized experts, and non-Ascend backends.  It only handles the 3D
    vLLM-Ascend expert tensors when a 2D or fused 3D runtime-sync tensor is
    loaded into ``w13_weight`` or ``w2_weight``.
    """
    if 'w13_weight' not in name and 'w2_weight' not in name:
        return
    quant_method = getattr(experts, 'quant_method', None)
    quant_method_module = type(quant_method).__module__ if quant_method is not None else ''
    if not quant_method_module.startswith('vllm_ascend'):
        return

    def make_ascend_moe_weight_loader(experts, origin_weight_loader):

        def load_processed_ascend_weight(param, loaded_weight, weight_name, shard_id, expert_id, return_success=False):
            quant_method = getattr(experts, 'quant_method', None)
            quant_method_module = type(quant_method).__module__ if quant_method is not None else ''
            # Only the GRPO runtime-sync path needs special handling here.
            # SWIFT provides HF/Megatron tensors, while vLLM-Ascend stores MoE
            # experts as 3D per-local-expert tensors.  Initial checkpoint load
            # and other layouts continue to use the original vLLM loader.
            is_runtime_sync_into_processed_param = (
                param.data.dim() == 3 and loaded_weight.dim() in {2, 3}
                and quant_method_module.startswith('vllm_ascend'))
            if not is_runtime_sync_into_processed_param:
                return origin_weight_loader(param, loaded_weight, weight_name, shard_id, expert_id, return_success)

            is_w13_shard = shard_id in {'w1', 'w3'} and 'w13_weight' in weight_name
            is_w2_shard = shard_id == 'w2' and 'w2_weight' in weight_name

            loaded_expert_sample = loaded_weight[0] if loaded_weight.dim() == 3 else loaded_weight

            def prepare_fsdp2_preprocessed_target_layout():
                """FSDP2 path: write weights before vLLM-Ascend post-load processing."""
                if is_w13_shard and param.data.shape[1] == loaded_expert_sample.shape[-1]:
                    param.data = param.data.transpose(1, 2).contiguous()
                elif is_w2_shard and param.data.shape[2] == loaded_expert_sample.shape[0]:
                    param.data = param.data.transpose(1, 2).contiguous()

            def prepare_megatron_processed_target_layout():
                """Megatron path: write weights into vLLM-Ascend runtime layout."""
                if (is_w13_shard and param.data.shape[-1] == loaded_expert_sample.shape[-1]
                        and param.data.shape[-2] != loaded_expert_sample.shape[-1]):
                    param.data = param.data.transpose(1, 2).contiguous()
                elif (is_w2_shard and param.data.shape[-2] == loaded_expert_sample.shape[0]
                      and param.data.shape[-1] != loaded_expert_sample.shape[0]):
                    param.data = param.data.transpose(1, 2).contiguous()

            tp_rank = experts.tp_rank

            def copy_fsdp2_preprocessed_expert(local_expert_id: int, loaded_expert_weight) -> bool:
                """Copy FSDP2 fused expert weights into pre-process vLLM-Ascend layout."""
                param_data = param.data[local_expert_id]
                if is_w13_shard:
                    # Target: [2 * intermediate_per_tp, hidden].
                    shard_size = param_data.shape[0] // 2
                    loaded_expert_weight = loaded_expert_weight.narrow(0, shard_size * tp_rank, shard_size)
                    offset = 0 if shard_id == 'w1' else shard_size
                    param_data[offset:offset + shard_size].copy_(loaded_expert_weight.contiguous())
                    return True

                if is_w2_shard:
                    # Target: [hidden, intermediate_per_tp].
                    shard_size = param_data.shape[1]
                    loaded_expert_weight = loaded_expert_weight.narrow(1, shard_size * tp_rank, shard_size)
                    param_data.copy_(loaded_expert_weight.contiguous())
                    return True

                return False

            def copy_megatron_processed_expert(local_expert_id: int, loaded_expert_weight) -> bool:
                """Copy Megatron/HF expert shards into processed vLLM-Ascend layout."""
                param_data = param.data[local_expert_id]
                if is_w13_shard:
                    # Target: [hidden, 2 * intermediate_per_tp].
                    shard_size = param_data.shape[1] // 2
                    loaded_expert_weight = loaded_expert_weight.narrow(0, shard_size * tp_rank, shard_size)
                    offset = 0 if shard_id == 'w1' else shard_size
                    param_data[:, offset:offset + shard_size].copy_(loaded_expert_weight.transpose(0, 1).contiguous())
                    return True

                if is_w2_shard:
                    # Target: [intermediate_per_tp, hidden].
                    shard_size = param_data.shape[0]
                    loaded_expert_weight = loaded_expert_weight.narrow(1, shard_size * tp_rank, shard_size)
                    param_data.copy_(loaded_expert_weight.transpose(0, 1).contiguous())
                    return True

                return False

            if load_preprocessed_weight:
                prepare_fsdp2_preprocessed_target_layout()
                copy_one_expert = copy_fsdp2_preprocessed_expert
            else:
                prepare_megatron_processed_target_layout()
                copy_one_expert = copy_megatron_processed_expert

            if loaded_weight.dim() == 3:
                copied = False
                for global_expert_id, loaded_expert_weight in enumerate(loaded_weight):
                    local_expert_id = experts._map_global_expert_id_to_local_expert_id(global_expert_id)
                    if local_expert_id == -1:
                        continue
                    copied = copy_one_expert(local_expert_id, loaded_expert_weight) or copied
                return copied if return_success else None

            local_expert_id = experts._map_global_expert_id_to_local_expert_id(expert_id)
            if local_expert_id == -1:
                return False if return_success else None

            if copy_one_expert(local_expert_id, loaded_weight):
                return True if return_success else None

            return origin_weight_loader(param, loaded_weight, weight_name, shard_id, expert_id, return_success)

        load_processed_ascend_weight._swift_ascend_moe_weight_loader = True
        load_processed_ascend_weight._swift_origin_weight_loader = origin_weight_loader
        load_processed_ascend_weight._swift_load_preprocessed_weight = load_preprocessed_weight
        return load_processed_ascend_weight

    if not hasattr(experts, 'weight_loader'):
        return
    weight_loader = getattr(param, 'weight_loader', experts.weight_loader)
    origin_weight_loader = getattr(weight_loader, '_swift_origin_weight_loader', weight_loader)
    if (not getattr(weight_loader, '_swift_ascend_moe_weight_loader', False)
            or getattr(weight_loader, '_swift_load_preprocessed_weight', None) != load_preprocessed_weight):
        param.weight_loader = make_ascend_moe_weight_loader(experts, origin_weight_loader)


__all__ = [
    'configure_vllm_ascend_moe_preprocessed_weight_sync',
    'configure_vllm_ascend_moe_weight_sync',
    'expand_fused_moe_expert_names_for_vllm_ascend',
    'expand_fused_moe_expert_weight_for_vllm_ascend',
    'patch_vllm_ascend_moe_expert_weight_loader',
    'patch_vllm_ascend_moe_runtime',
    'should_skip_vllm_ascend_moe_post_load',
    'use_vllm_ascend_moe_preprocessed_weight',
]
