553 lines
23 KiB
Python
553 lines
23 KiB
Python
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
|
|
# Copyright 2023 The vLLM team.
|
|
#
|
|
# 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.
|
|
|
|
|
|
import torch
|
|
import torch_npu
|
|
from torch.nn.functional import pad
|
|
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
|
from vllm.triton_utils import HAS_TRITON
|
|
|
|
from vllm_ascend.ascend_forward_context import _EXTRA_CTX, MoECommType
|
|
from vllm_ascend.device.device_op import DeviceOperator
|
|
from vllm_ascend.device.mxfp_compat import (
|
|
ensure_mxfp8_moe_available,
|
|
)
|
|
from vllm_ascend.ops.activation import AscendSwigluOAIAndMul, AscendSwigluStepAndMul
|
|
from vllm_ascend.ops.fused_moe.moe_runtime_args import MoEMlpComputeInput
|
|
from vllm_ascend.quantization.quant_type import QuantType
|
|
from vllm_ascend.utils import (
|
|
dispose_tensor,
|
|
enable_custom_op,
|
|
get_ascend_device_type,
|
|
get_weight_prefetch_method,
|
|
)
|
|
|
|
ASCEND_DEVICE_TYPE = get_ascend_device_type()
|
|
|
|
|
|
def _custom_gmm_swiglu_enabled(fusion, dynamic_eplb):
|
|
return fusion and dynamic_eplb and enable_custom_op()
|
|
|
|
|
|
def cumsum_group_list(
|
|
group_list: torch.Tensor, src_list_type: int, dst_list_type: int, active_num: int = 0, expert_num: int = 0
|
|
) -> torch.Tensor:
|
|
if src_list_type not in [0, 1, 2]:
|
|
raise ValueError(f"group_list_type should be in [0, 1, 2], but received {src_list_type}")
|
|
|
|
if src_list_type == dst_list_type:
|
|
return group_list
|
|
if src_list_type == 1 and dst_list_type == 0:
|
|
return group_list.cumsum(dim=0)
|
|
if src_list_type == 0 and dst_list_type == 1:
|
|
group_diff = torch.diff(group_list)
|
|
new_group = torch.cat([group_list[0].unsqueeze(0), group_diff], dim=0)
|
|
return new_group
|
|
if src_list_type == 2 and dst_list_type == 0:
|
|
experts = pad(group_list[:, 0], (1, 0))
|
|
tokens = pad(group_list[:, 1].cumsum(dim=0), (1, 0))
|
|
cumsum_group_list = torch.full(
|
|
size=(expert_num,), fill_value=active_num, dtype=group_list.dtype, device=group_list.device
|
|
)
|
|
|
|
for i, (start, end) in enumerate(zip(experts[:-1], experts[1:])):
|
|
if end > start:
|
|
cumsum_group_list[start:end] = tokens[i]
|
|
|
|
return cumsum_group_list
|
|
raise NotImplementedError(
|
|
f"Conversion from src_list_type={src_list_type} to dst_list_type={dst_list_type} is not implemented yet. "
|
|
"This feature is under development."
|
|
)
|
|
|
|
|
|
def _require_single_tensor_for_swiglu_quant(
|
|
tensor_or_list: list[torch.Tensor] | torch.Tensor, *, name: str
|
|
) -> torch.Tensor:
|
|
if isinstance(tensor_or_list, list):
|
|
if len(tensor_or_list) != 1:
|
|
raise ValueError(f"{name} must be a tensor or a single-element list, but got {len(tensor_or_list)}.")
|
|
return tensor_or_list[0]
|
|
return tensor_or_list
|
|
|
|
|
|
def quant_apply_mlp(
|
|
hidden_states: torch.Tensor,
|
|
w1: list[torch.Tensor] | torch.Tensor,
|
|
w1_scale: list[torch.Tensor] | torch.Tensor,
|
|
w2: list[torch.Tensor] | torch.Tensor,
|
|
w2_scale: list[torch.Tensor] | torch.Tensor,
|
|
group_list: torch.Tensor,
|
|
group_list_type: int = 1,
|
|
dynamic_scale: torch.Tensor = None,
|
|
w1_scale_bias: torch.Tensor = None,
|
|
w2_scale_bias: torch.Tensor = None,
|
|
w1_offset: torch.Tensor | None = None,
|
|
w2_offset: torch.Tensor | None = None,
|
|
fusion: bool = False,
|
|
dynamic_eplb: bool = False,
|
|
use_mxfp_quant: bool = False,
|
|
mxfp_quant_dtype: QuantType | None = None,
|
|
act_quant_type: torch.dtype = torch.float8_e4m3fn,
|
|
weight_quant_type: torch.dtype | None = None,
|
|
scale_type: torch.dtype | None = None,
|
|
per_token_scale_type: torch.dtype | None = None,
|
|
use_bf16: bool = True,
|
|
activation: str | None = None,
|
|
swiglu_limit: float = 0.0,
|
|
use_w4a8_per_channel_gmm_swiglu: bool = False,
|
|
) -> torch.Tensor:
|
|
input_hidden_dtype = hidden_states.dtype
|
|
use_gmm_swiglu_quant_fusion = use_mxfp_quant or (fusion and not dynamic_eplb)
|
|
|
|
if use_mxfp_quant:
|
|
ensure_mxfp8_moe_available("MXFP MoE MLP path")
|
|
|
|
if w1_scale_bias is not None or w2_scale_bias is not None:
|
|
raise NotImplementedError("MXFP path does not support scale_bias yet.")
|
|
if w1_offset is not None or w2_offset is not None:
|
|
raise NotImplementedError("MXFP path does not support antiquant offset yet.")
|
|
|
|
if w1_offset is not None:
|
|
unquantized_hidden_states = hidden_states
|
|
quantized_hidden_states = None
|
|
elif mxfp_quant_dtype == QuantType.W4A16MXFP4:
|
|
quantized_hidden_states = None
|
|
pertoken_scale = None
|
|
elif dynamic_scale is None:
|
|
unquantized_hidden_states = hidden_states
|
|
hidden_states, pertoken_scale = DeviceOperator.npu_dynamic_quant(
|
|
hidden_states=hidden_states,
|
|
dynamic_scale=None,
|
|
act_quant_type=act_quant_type,
|
|
use_mxfp_quant=use_mxfp_quant,
|
|
)
|
|
dispose_tensor(unquantized_hidden_states)
|
|
quantized_hidden_states = None
|
|
else:
|
|
unquantized_hidden_states = None
|
|
pertoken_scale = (
|
|
DeviceOperator.maybe_normalize_mxfp_scale_layout(dynamic_scale) if use_mxfp_quant else dynamic_scale
|
|
)
|
|
quantized_hidden_states = hidden_states
|
|
|
|
bias1, bias2 = None, None
|
|
_output_dtype = w2_scale[0].dtype if isinstance(w2_scale, list) else w2_scale.dtype
|
|
|
|
weight_prefetch_method = get_weight_prefetch_method()
|
|
if weight_prefetch_method:
|
|
weight_prefetch_method.maybe_prefetch_moe_weight_postprocess(hidden_states)
|
|
is_mc2 = _EXTRA_CTX.moe_comm_type == MoECommType.MC2
|
|
if w1_scale_bias is None and w1_offset is None and is_mc2:
|
|
if _custom_gmm_swiglu_enabled(fusion, dynamic_eplb) and not use_mxfp_quant:
|
|
# gmm1: gate_up_proj & act_fn: swiglu
|
|
hidden_states, swiglu_out_scale, _ = torch.ops._C_ascend.grouped_matmul_swiglu_quant_weight_nz_tensor_list(
|
|
x=hidden_states,
|
|
weight=w1,
|
|
weight_scale=w1_scale,
|
|
x_scale=pertoken_scale,
|
|
group_list=cumsum_group_list(group_list, group_list_type, 0),
|
|
swiglu_limit=swiglu_limit,
|
|
)
|
|
elif use_gmm_swiglu_quant_fusion:
|
|
# gmm1: gate_up_proj & act_fn: swiglu
|
|
hidden_states, swiglu_out_scale, _ = DeviceOperator.npu_grouped_matmul_swiglu_quant(
|
|
x=hidden_states,
|
|
weight=_require_single_tensor_for_swiglu_quant(w1, name="w1"),
|
|
group_list=cumsum_group_list(group_list, group_list_type, 0),
|
|
weight_scale=_require_single_tensor_for_swiglu_quant(w1_scale, name="w1_scale"),
|
|
x_scale=pertoken_scale,
|
|
bias=None,
|
|
use_mxfp_quant=use_mxfp_quant,
|
|
act_quant_type=act_quant_type,
|
|
weight_quant_type=weight_quant_type,
|
|
swiglu_limit=swiglu_limit,
|
|
mxfp_quant_dtype=mxfp_quant_dtype,
|
|
)
|
|
if quantized_hidden_states is not None:
|
|
dispose_tensor(quantized_hidden_states)
|
|
else:
|
|
if w1_scale[0].dtype != torch.float32:
|
|
w1_scale[0] = w1_scale[0].to(torch.float32)
|
|
# gmm1: gate_up_proj
|
|
hidden_states = torch_npu.npu_grouped_matmul(
|
|
x=[hidden_states],
|
|
weight=w1,
|
|
split_item=3,
|
|
group_list_type=group_list_type,
|
|
group_type=0,
|
|
group_list=group_list,
|
|
output_dtype=torch.int32,
|
|
)[0]
|
|
if quantized_hidden_states is not None:
|
|
dispose_tensor(quantized_hidden_states)
|
|
# act_fn: swiglu
|
|
hidden_states, swiglu_out_scale = torch.ops._C_ascend.npu_dequant_swiglu_quant(
|
|
x=hidden_states,
|
|
weight_scale=w1_scale[0],
|
|
activation_scale=pertoken_scale,
|
|
bias=None,
|
|
quant_scale=None,
|
|
quant_offset=None,
|
|
group_index=cumsum_group_list(group_list, group_list_type, 1),
|
|
activate_left=True,
|
|
quant_mode=1,
|
|
)
|
|
before_gmm2_evt = torch.npu.current_stream().record_event()
|
|
# gmm2: down_proj
|
|
hidden_states = DeviceOperator.npu_grouped_matmul_gmm2(
|
|
hidden_states=hidden_states,
|
|
weight=w2,
|
|
weight_scale=w2_scale,
|
|
per_token_scale=swiglu_out_scale,
|
|
group_list=group_list,
|
|
group_list_type=group_list_type,
|
|
input_dtype=input_hidden_dtype,
|
|
act_quant_type=act_quant_type,
|
|
weight_quant_type=weight_quant_type,
|
|
scale_type=scale_type,
|
|
per_token_scale_type=per_token_scale_type,
|
|
use_bf16=use_bf16,
|
|
use_mxfp_quant=use_mxfp_quant,
|
|
bias=None,
|
|
fallback_output_dtype=w2_scale[0].dtype if isinstance(w2_scale, list) else w2_scale.dtype,
|
|
mxfp_quant_dtype=mxfp_quant_dtype,
|
|
)
|
|
elif w1_offset is not None:
|
|
# gmm1: gate_up_proj
|
|
hidden_states = torch_npu.npu_grouped_matmul(
|
|
x=[unquantized_hidden_states],
|
|
weight=[w1],
|
|
antiquant_scale=[w1_scale],
|
|
antiquant_offset=[w1_offset],
|
|
split_item=2,
|
|
group_list_type=group_list_type,
|
|
group_type=0,
|
|
group_list=group_list,
|
|
output_dtype=_output_dtype,
|
|
)[0]
|
|
dispose_tensor(unquantized_hidden_states)
|
|
# act_fn: swiglu
|
|
if activation == MoEActivation.SWIGLUSTEP:
|
|
hidden_states = AscendSwigluStepAndMul.swiglustep_forward(hidden_states, limit=swiglu_limit or 7.0)
|
|
else:
|
|
hidden_states = torch_npu.npu_swiglu(hidden_states)
|
|
before_gmm2_evt = torch.npu.current_stream().record_event()
|
|
# gmm2: down_proj
|
|
hidden_states = torch_npu.npu_grouped_matmul(
|
|
x=[hidden_states],
|
|
weight=[w2],
|
|
antiquant_scale=[w2_scale],
|
|
antiquant_offset=[w2_offset],
|
|
split_item=2,
|
|
group_list_type=group_list_type,
|
|
group_type=0,
|
|
group_list=group_list,
|
|
output_dtype=_output_dtype,
|
|
)[0]
|
|
else:
|
|
if w1_scale_bias is not None:
|
|
if group_list_type == 0:
|
|
group_list = torch.cat([group_list[:1], torch.diff(group_list, dim=0)])
|
|
group_list_type = 1
|
|
bias1 = w1_scale_bias
|
|
bias2 = w2_scale_bias
|
|
# TODO w4a8 scene: dynamic acquisition of dtype in the future
|
|
_output_dtype = torch.bfloat16
|
|
|
|
if use_w4a8_per_channel_gmm_swiglu and enable_custom_op() and activation != MoEActivation.SWIGLUSTEP:
|
|
hidden_states, swiglu_out_scale = torch.ops._C_ascend.grouped_matmul_swiglu_quant_v2(
|
|
x=hidden_states,
|
|
weight=w1,
|
|
weight_scale=w1_scale if isinstance(w1_scale, list) else [w1_scale],
|
|
x_scale=pertoken_scale,
|
|
group_list=group_list,
|
|
weight_assist_matrix=bias1,
|
|
dequant_mode=0,
|
|
group_list_type=group_list_type,
|
|
swiglu_limit=swiglu_limit,
|
|
)
|
|
elif _custom_gmm_swiglu_enabled(fusion, dynamic_eplb) and not use_mxfp_quant:
|
|
# gmm1: gate_up_proj & act_fn: swiglu
|
|
hidden_states, swiglu_out_scale, _ = torch.ops._C_ascend.grouped_matmul_swiglu_quant_weight_nz_tensor_list(
|
|
x=hidden_states,
|
|
weight=w1,
|
|
weight_scale=w1_scale,
|
|
x_scale=pertoken_scale,
|
|
group_list=cumsum_group_list(group_list, group_list_type, 0),
|
|
bias=bias1,
|
|
swiglu_limit=swiglu_limit,
|
|
)
|
|
elif use_gmm_swiglu_quant_fusion and activation != MoEActivation.SWIGLUSTEP:
|
|
hidden_states, swiglu_out_scale, _ = DeviceOperator.npu_grouped_matmul_swiglu_quant(
|
|
x=hidden_states,
|
|
weight=_require_single_tensor_for_swiglu_quant(w1, name="w1"),
|
|
group_list=cumsum_group_list(group_list, group_list_type, 0),
|
|
weight_scale=_require_single_tensor_for_swiglu_quant(w1_scale, name="w1_scale"),
|
|
x_scale=pertoken_scale,
|
|
bias=bias1,
|
|
use_mxfp_quant=use_mxfp_quant,
|
|
act_quant_type=act_quant_type,
|
|
weight_quant_type=weight_quant_type,
|
|
swiglu_limit=swiglu_limit,
|
|
mxfp_quant_dtype=mxfp_quant_dtype,
|
|
)
|
|
if quantized_hidden_states is not None:
|
|
dispose_tensor(quantized_hidden_states)
|
|
else:
|
|
w1_scale[0] = w1_scale[0].to(w2_scale[0].dtype)
|
|
# gmm1: gate_up_proj
|
|
hidden_states = torch_npu.npu_grouped_matmul(
|
|
x=[hidden_states],
|
|
weight=w1,
|
|
scale=w1_scale,
|
|
bias=bias1,
|
|
per_token_scale=[pertoken_scale],
|
|
split_item=2,
|
|
group_list_type=group_list_type,
|
|
group_type=0,
|
|
group_list=group_list,
|
|
output_dtype=_output_dtype,
|
|
)[0]
|
|
if quantized_hidden_states is not None:
|
|
dispose_tensor(quantized_hidden_states)
|
|
# act_fn: swiglu
|
|
if activation == MoEActivation.SWIGLUSTEP:
|
|
hidden_states = AscendSwigluStepAndMul.swiglustep_forward(hidden_states, limit=swiglu_limit or 7.0)
|
|
hidden_states, swiglu_out_scale = torch_npu.npu_dynamic_quant(hidden_states)
|
|
elif HAS_TRITON:
|
|
from vllm_ascend.ops.triton.activation.swiglu_quant import swiglu_quant
|
|
|
|
hidden_states, swiglu_out_scale = swiglu_quant(
|
|
hidden_states, group_list=group_list, group_list_type=group_list_type
|
|
)
|
|
else:
|
|
hidden_states = torch_npu.npu_swiglu(hidden_states)
|
|
hidden_states, swiglu_out_scale = torch_npu.npu_dynamic_quant(hidden_states)
|
|
before_gmm2_evt = torch.npu.current_stream().record_event()
|
|
# gmm2: down_proj
|
|
hidden_states = DeviceOperator.npu_grouped_matmul_gmm2(
|
|
hidden_states=hidden_states,
|
|
weight=w2,
|
|
weight_scale=w2_scale,
|
|
per_token_scale=swiglu_out_scale,
|
|
group_list=group_list,
|
|
group_list_type=group_list_type,
|
|
input_dtype=input_hidden_dtype,
|
|
act_quant_type=act_quant_type,
|
|
weight_quant_type=weight_quant_type,
|
|
scale_type=scale_type,
|
|
per_token_scale_type=per_token_scale_type,
|
|
use_bf16=use_bf16,
|
|
use_mxfp_quant=use_mxfp_quant,
|
|
bias=bias2,
|
|
fallback_output_dtype=_output_dtype,
|
|
mxfp_quant_dtype=mxfp_quant_dtype,
|
|
)
|
|
return hidden_states, before_gmm2_evt
|
|
|
|
|
|
def unquant_apply_mlp(
|
|
hidden_states: torch.Tensor,
|
|
w1: torch.Tensor,
|
|
w2: torch.Tensor,
|
|
group_list: torch.Tensor,
|
|
w1_bias: torch.Tensor = None,
|
|
w2_bias: torch.Tensor = None,
|
|
activation: str | None = None,
|
|
group_list_type: int = 1,
|
|
topk_scales: torch.Tensor | None = None,
|
|
need_trans: bool = True,
|
|
swiglu_limit: float = 0.0,
|
|
lora_context=None,
|
|
expanded_row_idx: torch.Tensor | None = None,
|
|
topk_ids: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
if need_trans:
|
|
w1 = w1.transpose(1, 2)
|
|
w2 = w2.transpose(1, 2)
|
|
|
|
gate_up_out = torch_npu.npu_grouped_matmul(
|
|
x=[hidden_states],
|
|
weight=[w1],
|
|
bias=[w1_bias.to(dtype=torch.float32)] if w1_bias is not None else None,
|
|
split_item=2,
|
|
group_list_type=group_list_type,
|
|
group_type=0,
|
|
group_list=group_list,
|
|
)[0]
|
|
|
|
# MoE LoRA: only attempt injection when an adapter wraps this layer and the
|
|
# comm method provided AllGather routing metadata (expanded_row_idx). Lazy
|
|
# import keeps the core MLP free of any LoRA dependency on the common path.
|
|
lora_routing = None
|
|
if lora_context is not None: # LoRA applied
|
|
if expanded_row_idx is None or topk_ids is None:
|
|
raise AssertionError(
|
|
"MoE LoRA requires expanded_row_idx and topk_ids metadata, "
|
|
"which are only available in AllGather communication mode. "
|
|
"Please ensure you are running in a supported configuration."
|
|
)
|
|
|
|
from vllm_ascend.lora.fused_moe import moe_lora_apply_w2, moe_lora_apply_w13
|
|
|
|
# LoRA w13 delta: applied to gate_up_out before activation, with the MLP
|
|
# input as the lora_a input (mirrors the base gate_up GMM above).
|
|
lora_routing = moe_lora_apply_w13(
|
|
lora_context,
|
|
gate_up_out=gate_up_out,
|
|
hidden_states=hidden_states,
|
|
expanded_row_idx=expanded_row_idx,
|
|
topk_ids=topk_ids,
|
|
)
|
|
|
|
if activation == MoEActivation.SWIGLUOAI:
|
|
num_experts, _, hidden_size = w1.shape
|
|
gate_up_out = AscendSwigluOAIAndMul.swiglu_oai_forward(gate_up_out.view(-1, hidden_size))
|
|
elif activation == MoEActivation.SWIGLUSTEP:
|
|
gate_up_out = AscendSwigluStepAndMul.swiglustep_forward(gate_up_out, limit=swiglu_limit or 7.0)
|
|
elif activation == MoEActivation.GELU:
|
|
gate, up = gate_up_out.chunk(2, dim=-1)
|
|
gate_up_out = torch.nn.functional.gelu(gate) * up
|
|
elif activation == MoEActivation.GELU_TANH:
|
|
gate, up = gate_up_out.chunk(2, dim=-1)
|
|
gate_up_out = torch.nn.functional.gelu(gate, approximate="tanh") * up
|
|
else:
|
|
if swiglu_limit > 0:
|
|
gate, up = gate_up_out.chunk(2, dim=-1)
|
|
gate.clamp_(max=swiglu_limit)
|
|
up.clamp_(min=-swiglu_limit, max=swiglu_limit)
|
|
gate_up_out = torch_npu.npu_swiglu(gate_up_out)
|
|
|
|
if topk_scales is not None:
|
|
gate_up_out *= topk_scales
|
|
|
|
hidden_states = torch_npu.npu_grouped_matmul(
|
|
x=[gate_up_out],
|
|
weight=[w2],
|
|
bias=[w2_bias.to(dtype=torch.float32)] if w2_bias is not None else None,
|
|
split_item=2,
|
|
group_list_type=group_list_type,
|
|
group_type=0,
|
|
group_list=group_list,
|
|
)[0]
|
|
|
|
# LoRA w2 delta: applied to the down-proj output, with the activation output
|
|
# as the lora_a input. Reuses the per-row routing computed for w13.
|
|
if lora_routing is not None:
|
|
moe_lora_apply_w2(
|
|
lora_context,
|
|
down_out=hidden_states,
|
|
silu_out=gate_up_out,
|
|
lora_routing=lora_routing,
|
|
)
|
|
return hidden_states, None
|
|
|
|
|
|
def unified_apply_mlp(*, mlp_compute_input: MoEMlpComputeInput) -> torch.Tensor:
|
|
"""
|
|
Unified MoE MLP entry.
|
|
Quant path is dispatched by DeviceOperator with explicit typed kernel flags.
|
|
"""
|
|
hidden_states = mlp_compute_input.hidden_states
|
|
group_list = mlp_compute_input.group_list
|
|
group_list_type = mlp_compute_input.group_list_type
|
|
dynamic_scale = mlp_compute_input.dynamic_scale
|
|
topk_scales = mlp_compute_input.topk_scales
|
|
w1 = mlp_compute_input.weights.w1
|
|
w2 = mlp_compute_input.weights.w2
|
|
w1_bias = mlp_compute_input.weights.w1_bias
|
|
w2_bias = mlp_compute_input.weights.w2_bias
|
|
w1_scale = mlp_compute_input.weights.w1_scale
|
|
w2_scale = mlp_compute_input.weights.w2_scale
|
|
w1_scale_bias = mlp_compute_input.weights.w1_scale_bias
|
|
w2_scale_bias = mlp_compute_input.weights.w2_scale_bias
|
|
w1_offset = mlp_compute_input.weights.w1_offset
|
|
w2_offset = mlp_compute_input.weights.w2_offset
|
|
activation = mlp_compute_input.activation
|
|
need_trans = mlp_compute_input.need_trans
|
|
dynamic_eplb = mlp_compute_input.dynamic_eplb
|
|
fusion = mlp_compute_input.fusion
|
|
swiglu_limit = mlp_compute_input.swiglu_limit
|
|
|
|
if not mlp_compute_input.quant.is_quant:
|
|
return unquant_apply_mlp(
|
|
hidden_states=hidden_states,
|
|
w1=w1,
|
|
w2=w2,
|
|
w1_bias=w1_bias,
|
|
w2_bias=w2_bias,
|
|
activation=activation,
|
|
group_list=group_list,
|
|
group_list_type=group_list_type,
|
|
topk_scales=topk_scales,
|
|
need_trans=need_trans,
|
|
swiglu_limit=swiglu_limit,
|
|
lora_context=mlp_compute_input.lora_context,
|
|
expanded_row_idx=mlp_compute_input.expanded_row_idx,
|
|
topk_ids=mlp_compute_input.topk_ids,
|
|
)
|
|
|
|
assert w1_scale is not None and w2_scale is not None
|
|
act_quant_type = torch.int8 if mlp_compute_input.quant.is_int_quant else torch.float8_e4m3fn
|
|
weight_quant_type = torch.float8_e4m3fn
|
|
scale_type = None
|
|
per_token_scale_type = None
|
|
use_bf16 = hidden_states.dtype == torch.bfloat16
|
|
use_mxfp_quant = mlp_compute_input.quant.is_mxfp
|
|
mxfp_quant_dtype = mlp_compute_input.quant.quant_type
|
|
|
|
if use_mxfp_quant:
|
|
mxfp = mlp_compute_input.quant.mxfp
|
|
assert mxfp is not None, "mlp_compute_input.quant.mxfp is required when quant_type is MXFP8."
|
|
act_quant_type = mxfp.act_quant_type or act_quant_type
|
|
if mxfp_quant_dtype == QuantType.W4A16MXFP4:
|
|
act_quant_type = mxfp.act_quant_type
|
|
weight_quant_type = mxfp.weight_quant_type or weight_quant_type
|
|
if mxfp_quant_dtype in [QuantType.W4A8MXFP, QuantType.W4A16MXFP4]:
|
|
weight_quant_type = mxfp.weight_quant_type
|
|
scale_type = mxfp.scale_dtype
|
|
per_token_scale_type = mxfp.per_token_scale_dtype
|
|
use_bf16 = mxfp.use_bf16
|
|
|
|
return quant_apply_mlp(
|
|
hidden_states=hidden_states,
|
|
w1=w1,
|
|
w1_scale=w1_scale,
|
|
w2=w2,
|
|
w2_scale=w2_scale,
|
|
group_list=group_list,
|
|
dynamic_scale=dynamic_scale,
|
|
group_list_type=group_list_type,
|
|
w1_scale_bias=w1_scale_bias,
|
|
w2_scale_bias=w2_scale_bias,
|
|
w1_offset=w1_offset,
|
|
w2_offset=w2_offset,
|
|
fusion=fusion,
|
|
dynamic_eplb=dynamic_eplb,
|
|
use_mxfp_quant=use_mxfp_quant,
|
|
mxfp_quant_dtype=mxfp_quant_dtype,
|
|
act_quant_type=act_quant_type,
|
|
weight_quant_type=weight_quant_type,
|
|
scale_type=scale_type,
|
|
per_token_scale_type=per_token_scale_type,
|
|
use_bf16=use_bf16,
|
|
activation=activation,
|
|
swiglu_limit=swiglu_limit,
|
|
use_w4a8_per_channel_gmm_swiglu=mlp_compute_input.quant.use_w4a8_per_channel_gmm_swiglu,
|
|
)
|