data: port complete MoE + xllm layer call chains from upstream repos
MoE call chain from ds_vllm (vllm-project/vllm latest): ex_engine/moe/ — 20 files, 8736 lines - modular_kernel.py (1630 lines) — base classes for modular MoE - experts/fused_batched_moe.py (972 lines) — NaiveBatchedExperts - prepare_finalize/batched.py (171 lines) — token grouping by expert - topk_weight_and_reduce.py (176 lines) — scatter-add finalize - fused_moe.py (1740 lines) — main fused_moe dispatch - config.py (1407 lines) — FusedMoEQuantConfig - activation.py, utils.py, layer.py, etc. xllm layer code (jd-opensource/xllm): ex_engine/xllm_layers/ — 39 files, 5859 lines - ilu/fused_moe.cpp (797 lines) — production ixformer 7-step MoE pipeline - ilu/attention.cpp (189 lines) — paged_attention + flash_attn bridge - npu_torch/qwen3_gated_delta_net_base.cpp (576 lines) — GDN reference - common/rms_norm.cpp, rotary_embedding.cpp, activation.cpp, dense_mlp.cpp xllm ILU kernels — synced 10 files to upstream (diffs from prior edits) These are reference implementations, NOT hand-written. Source repos: vllm-project/vllm, jd-opensource/xllm
This commit is contained in:
Claude
0
ex_engine/moe/experts/__init__.py
Normal file
0
ex_engine/moe/experts/__init__.py
Normal file
170
ex_engine/moe/experts/fallback.py
Normal file
170
ex_engine/moe/experts/fallback.py
Normal file
@@ -0,0 +1,170 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import torch
|
||||
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
||||
from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
|
||||
|
||||
|
||||
class FallbackExperts(mk.FusedMoEExpertsModular, ABC):
|
||||
"""Base class for runtime dispatching of expert implementations."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
experts: mk.FusedMoEExpertsModular,
|
||||
fallback_experts: mk.FusedMoEExpertsModular,
|
||||
):
|
||||
super().__init__(
|
||||
moe_config=experts.moe_config, quant_config=experts.quant_config
|
||||
)
|
||||
self.fallback_experts = fallback_experts
|
||||
self.experts = experts
|
||||
|
||||
@staticmethod
|
||||
def get_clses() -> tuple[
|
||||
type[mk.FusedMoEExpertsModular],
|
||||
type[mk.FusedMoEExpertsModular],
|
||||
]:
|
||||
"""
|
||||
Get the cls for the experts and fallback experts.
|
||||
|
||||
Subclasses should implement this method, so that
|
||||
we have a consistent way to call the _supports_*
|
||||
class methods below.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Subclasses must return the cls for the experts and fallback experts."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def activation_format(
|
||||
cls: type["FallbackExperts"],
|
||||
) -> mk.FusedMoEActivationFormat:
|
||||
experts_cls, fallback_cls = cls.get_clses()
|
||||
assert experts_cls.activation_format() == fallback_cls.activation_format()
|
||||
return experts_cls.activation_format()
|
||||
|
||||
@classmethod
|
||||
def _supports_current_device(cls) -> bool:
|
||||
experts_cls, fallback_cls = cls.get_clses()
|
||||
return (
|
||||
experts_cls._supports_current_device()
|
||||
and fallback_cls._supports_current_device()
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _supports_no_act_and_mul(cls) -> bool:
|
||||
experts_cls, fallback_cls = cls.get_clses()
|
||||
return (
|
||||
experts_cls._supports_no_act_and_mul()
|
||||
and fallback_cls._supports_no_act_and_mul()
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _supports_quant_scheme(
|
||||
cls,
|
||||
weight_key: QuantKey | None,
|
||||
activation_key: QuantKey | None,
|
||||
) -> bool:
|
||||
experts_cls, fallback_cls = cls.get_clses()
|
||||
return experts_cls._supports_quant_scheme(
|
||||
weight_key, activation_key
|
||||
) and fallback_cls._supports_quant_scheme(weight_key, activation_key)
|
||||
|
||||
@classmethod
|
||||
def _supports_activation(cls, activation: MoEActivation) -> bool:
|
||||
experts_cls, fallback_cls = cls.get_clses()
|
||||
return experts_cls._supports_activation(
|
||||
activation
|
||||
) and fallback_cls._supports_activation(activation)
|
||||
|
||||
@classmethod
|
||||
def _supports_parallel_config(
|
||||
cls, moe_parallel_config: FusedMoEParallelConfig
|
||||
) -> bool:
|
||||
experts_cls, fallback_cls = cls.get_clses()
|
||||
return experts_cls._supports_parallel_config(
|
||||
moe_parallel_config
|
||||
) and fallback_cls._supports_parallel_config(moe_parallel_config)
|
||||
|
||||
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
|
||||
e_war = self.experts.finalize_weight_and_reduce_impl()
|
||||
fbe_war = self.fallback_experts.finalize_weight_and_reduce_impl()
|
||||
is_dge_war = e_war is not None
|
||||
is_fbe_war = fbe_war is not None
|
||||
|
||||
if is_dge_war and is_fbe_war:
|
||||
assert e_war == fbe_war, (
|
||||
"Both implementations should agree on WeightAndReduce impls. "
|
||||
f"Got e_war: {e_war}, and fbe_war: {fbe_war}"
|
||||
)
|
||||
|
||||
if e_war is not None:
|
||||
return e_war
|
||||
assert fbe_war is not None
|
||||
return fbe_war
|
||||
|
||||
@abstractmethod
|
||||
def workspace_shapes(
|
||||
self,
|
||||
M: int,
|
||||
N: int,
|
||||
K: int,
|
||||
topk: int,
|
||||
global_num_experts: int,
|
||||
local_num_experts: int,
|
||||
expert_tokens_meta: mk.ExpertTokensMetadata | None,
|
||||
activation: MoEActivation,
|
||||
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def _select_experts_impl(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
w1: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
) -> mk.FusedMoEExpertsModular:
|
||||
raise NotImplementedError
|
||||
|
||||
def apply(
|
||||
self,
|
||||
output: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
w1: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
activation: MoEActivation,
|
||||
global_num_experts: int,
|
||||
expert_map: torch.Tensor | None,
|
||||
a1q_scale: torch.Tensor | None,
|
||||
a2_scale: torch.Tensor | None,
|
||||
workspace13: torch.Tensor,
|
||||
workspace2: torch.Tensor,
|
||||
expert_tokens_meta: mk.ExpertTokensMetadata | None,
|
||||
apply_router_weight_on_input: bool,
|
||||
):
|
||||
experts = self._select_experts_impl(hidden_states, w1, w2)
|
||||
experts.apply(
|
||||
output,
|
||||
hidden_states,
|
||||
w1,
|
||||
w2,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
activation,
|
||||
global_num_experts,
|
||||
expert_map,
|
||||
a1q_scale,
|
||||
a2_scale,
|
||||
workspace13,
|
||||
workspace2,
|
||||
expert_tokens_meta,
|
||||
apply_router_weight_on_input,
|
||||
)
|
||||
972
ex_engine/moe/experts/fused_batched_moe.py
Normal file
972
ex_engine/moe/experts/fused_batched_moe.py
Normal file
@@ -0,0 +1,972 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Fused batched MoE kernel."""
|
||||
|
||||
import torch
|
||||
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEParallelConfig,
|
||||
FusedMoEQuantConfig,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import try_get_optimal_moe_config
|
||||
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
|
||||
TopKWeightAndReduceDelegate,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.utils import (
|
||||
_resize_cache,
|
||||
moe_kernel_quantize_input,
|
||||
normalize_batched_scales_shape,
|
||||
swiglu_limit_func,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
QuantKey,
|
||||
group_broadcast,
|
||||
kFp8Dynamic128Sym,
|
||||
kFp8DynamicTensorSym,
|
||||
kFp8DynamicTokenSym,
|
||||
kFp8Static128BlockSym,
|
||||
kFp8StaticChannelSym,
|
||||
kFp8StaticTensorSym,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import tl, triton
|
||||
|
||||
|
||||
@triton.jit
|
||||
def moe_mmk(
|
||||
a_ptrs,
|
||||
b_ptrs,
|
||||
K,
|
||||
expert_id,
|
||||
a_scale_ptr,
|
||||
b_scale_ptr,
|
||||
# The stride variables represent how much to increase the ptr by when
|
||||
# moving by 1 element in a particular dimension. E.g. `stride_am` is
|
||||
# how much to increase `a_ptr` by to get the element one row down
|
||||
# (A has M rows).
|
||||
stride_ak: tl.int64,
|
||||
stride_bk: tl.int64,
|
||||
stride_ase: tl.int64,
|
||||
stride_asm: tl.int64,
|
||||
stride_ask: tl.int64,
|
||||
stride_bse: tl.int64,
|
||||
stride_bsk: tl.int64,
|
||||
stride_bsn: tl.int64,
|
||||
# Offsets and masks
|
||||
offs_m,
|
||||
offs_n,
|
||||
offs_bn,
|
||||
mask_m,
|
||||
# Block size for block-wise quantization
|
||||
group_n: tl.constexpr,
|
||||
group_k: tl.constexpr,
|
||||
# Meta-parameters
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
compute_type: tl.constexpr,
|
||||
use_w8a8: tl.constexpr,
|
||||
use_w8a16: tl.constexpr,
|
||||
per_act_token_quant: tl.constexpr,
|
||||
):
|
||||
offs_k = tl.arange(0, BLOCK_K)
|
||||
|
||||
if use_w8a16:
|
||||
b_scale_ptrs = (
|
||||
b_scale_ptr + expert_id * stride_bse + offs_n[None, :] * stride_bsn
|
||||
)
|
||||
b_scale = tl.load(b_scale_ptrs)
|
||||
|
||||
if use_w8a8:
|
||||
# block-wise
|
||||
if group_k > 0 and group_n > 0:
|
||||
a_scale_ptrs = a_scale_ptr + offs_m * stride_asm
|
||||
offs_bsn = offs_bn // group_n
|
||||
b_scale_ptrs = b_scale_ptr + offs_bsn * stride_bsn
|
||||
|
||||
# per act token
|
||||
elif per_act_token_quant:
|
||||
# Load per-token scale for activations
|
||||
a_scale_ptrs = a_scale_ptr + offs_m * stride_asm
|
||||
a_scale = tl.load(a_scale_ptrs, mask=mask_m, other=0.0)[:, None]
|
||||
|
||||
b_scale_ptrs = b_scale_ptr + offs_bn[None, :] * stride_bsn
|
||||
b_scale = tl.load(b_scale_ptrs)
|
||||
|
||||
# tensor-wise
|
||||
else:
|
||||
a_scale = tl.load(a_scale_ptr)
|
||||
b_scale = tl.load(b_scale_ptr)
|
||||
|
||||
# -----------------------------------------------------------
|
||||
# Iterate to compute a block of the C matrix.
|
||||
# We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
|
||||
# of fp32 values for higher accuracy.
|
||||
# `accumulator` will be converted back to fp16 after the loop.
|
||||
accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
|
||||
for k in range(0, tl.cdiv(K, BLOCK_K)):
|
||||
# Load the next block of A and B, generate a mask by checking the
|
||||
# K dimension.
|
||||
a = tl.load(
|
||||
a_ptrs,
|
||||
mask=mask_m[:, None] & (offs_k[None, :] < K - k * BLOCK_K),
|
||||
other=0.0,
|
||||
)
|
||||
b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0)
|
||||
# We accumulate along the K dimension.
|
||||
if use_w8a16:
|
||||
accumulator = tl.dot(a, b.to(compute_type), acc=accumulator)
|
||||
elif use_w8a8:
|
||||
if group_k > 0 and group_n > 0:
|
||||
k_start = k * BLOCK_K
|
||||
offs_ks = k_start // group_k
|
||||
a_scale = tl.load(
|
||||
a_scale_ptrs + offs_ks * stride_ask, mask=mask_m, other=0.0
|
||||
)
|
||||
b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk)
|
||||
|
||||
accumulator += tl.dot(a, b) * a_scale[:, None] * b_scale[None, :]
|
||||
else:
|
||||
# acc used to enable fp8_fast_accum
|
||||
accumulator = tl.dot(a, b, acc=accumulator)
|
||||
else:
|
||||
accumulator += tl.dot(a, b)
|
||||
|
||||
# Advance the ptrs to the next K block.
|
||||
a_ptrs += BLOCK_K * stride_ak
|
||||
b_ptrs += BLOCK_K * stride_bk
|
||||
|
||||
if use_w8a16:
|
||||
accumulator = (accumulator * b_scale).to(compute_type)
|
||||
elif use_w8a8:
|
||||
if group_k > 0 and group_n > 0:
|
||||
accumulator = accumulator.to(compute_type)
|
||||
else:
|
||||
accumulator = (accumulator * a_scale * b_scale).to(compute_type)
|
||||
else:
|
||||
accumulator = accumulator.to(compute_type)
|
||||
|
||||
return accumulator
|
||||
|
||||
|
||||
@triton.jit
|
||||
def expert_triton_kernel(
|
||||
a_ptr, # [max_tokens, K]
|
||||
b_ptr, # [K, N]
|
||||
c_ptr, # [max_tokens, N]
|
||||
expert_id,
|
||||
compute_type: tl.constexpr,
|
||||
# Dimensions
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
# Quantization data
|
||||
a_scale_ptr,
|
||||
b_scale_ptr,
|
||||
b_zp_ptr,
|
||||
# strides
|
||||
stride_am: tl.int64,
|
||||
stride_ak: tl.int64,
|
||||
stride_bk: tl.int64,
|
||||
stride_bn: tl.int64,
|
||||
stride_cm: tl.int64,
|
||||
stride_cn: tl.int64,
|
||||
stride_ase: tl.int64,
|
||||
stride_asm: tl.int64,
|
||||
stride_ask: tl.int64,
|
||||
stride_bse: tl.int64,
|
||||
stride_bsk: tl.int64,
|
||||
stride_bsn: tl.int64,
|
||||
# offsets
|
||||
offs_bn,
|
||||
# Blockwise quantization data
|
||||
group_n,
|
||||
group_k,
|
||||
# Quantization schemes
|
||||
use_fp8_w8a8: tl.constexpr,
|
||||
use_int8_w8a16: tl.constexpr,
|
||||
per_act_token_quant: tl.constexpr,
|
||||
# Kernel config
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
):
|
||||
offs_m = tl.arange(0, BLOCK_M)
|
||||
offs_n = tl.arange(0, BLOCK_N) % N
|
||||
offs_k = tl.arange(0, BLOCK_K)
|
||||
mask_m = offs_m < M
|
||||
|
||||
# Make grids of a + b pointers
|
||||
a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
|
||||
b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn
|
||||
|
||||
accumulator = moe_mmk(
|
||||
a_ptrs,
|
||||
b_ptrs,
|
||||
K,
|
||||
expert_id,
|
||||
a_scale_ptr,
|
||||
b_scale_ptr,
|
||||
# The stride variables represent how much to increase the ptr by when
|
||||
# moving by 1 element in a particular dimension. E.g. `stride_am` is
|
||||
# how much to increase `a_ptr` by to get the element one row down
|
||||
# (A has M rows).
|
||||
stride_ak,
|
||||
stride_bk,
|
||||
stride_ase,
|
||||
stride_asm,
|
||||
stride_ask,
|
||||
stride_bse,
|
||||
stride_bsk,
|
||||
stride_bsn,
|
||||
# Offsets and masks
|
||||
offs_m,
|
||||
offs_n,
|
||||
offs_bn,
|
||||
mask_m,
|
||||
# Block size for block-wise quantization
|
||||
group_n,
|
||||
group_k,
|
||||
# Meta-parameters
|
||||
BLOCK_M,
|
||||
BLOCK_N,
|
||||
BLOCK_K,
|
||||
compute_type,
|
||||
use_fp8_w8a8,
|
||||
use_int8_w8a16,
|
||||
per_act_token_quant,
|
||||
)
|
||||
|
||||
# store in C
|
||||
offs_cn = tl.arange(0, BLOCK_N)
|
||||
c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_cn[None, :] * stride_cn
|
||||
c_mask = mask_m[:, None] & (offs_cn[None, :] < N)
|
||||
tl.store(c_ptrs, accumulator, mask=c_mask)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def batched_triton_kernel(
|
||||
a_ptr, # [E, max_num_tokens, K]
|
||||
b_ptr, # [E, K, N]
|
||||
c_ptr, # [E, max_num_tokens, N]
|
||||
expert_num_tokens, # [E]
|
||||
compute_type: tl.constexpr,
|
||||
# Dimensions
|
||||
max_num_tokens,
|
||||
K,
|
||||
N,
|
||||
# Quantization data
|
||||
a_scale_ptr,
|
||||
b_scale_ptr,
|
||||
b_zp_ptr,
|
||||
# The stride variables represent how much to increase the ptr by when
|
||||
# moving by 1 element in a particular dimension. E.g. `stride_am` is
|
||||
# how much to increase `a_ptr` by to get the element one row down
|
||||
# (A has M rows).
|
||||
stride_ae: tl.int64,
|
||||
stride_am: tl.int64,
|
||||
stride_ak: tl.int64,
|
||||
stride_be: tl.int64,
|
||||
stride_bk: tl.int64,
|
||||
stride_bn: tl.int64,
|
||||
stride_ce: tl.int64,
|
||||
stride_cm: tl.int64,
|
||||
stride_cn: tl.int64,
|
||||
stride_ase: tl.int64,
|
||||
stride_asm: tl.int64,
|
||||
stride_ask: tl.int64,
|
||||
stride_bse: tl.int64,
|
||||
stride_bsk: tl.int64,
|
||||
stride_bsn: tl.int64,
|
||||
# Blockwise quantization data
|
||||
group_n: tl.constexpr,
|
||||
group_k: tl.constexpr,
|
||||
# Quantization schemes
|
||||
use_fp8_w8a8: tl.constexpr,
|
||||
use_int8_w8a16: tl.constexpr,
|
||||
per_act_token_quant: tl.constexpr,
|
||||
# Kernel config
|
||||
BLOCK_M: tl.constexpr,
|
||||
BLOCK_N: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
):
|
||||
expert_id = tl.program_id(axis=0)
|
||||
e_num_tokens = tl.load(expert_num_tokens + expert_id)
|
||||
if e_num_tokens == 0:
|
||||
# Early exit
|
||||
return
|
||||
|
||||
# axis 1 is M_blocks * N_blocks
|
||||
pid_mn = tl.program_id(axis=1)
|
||||
# num_pid_m = tl.cdiv(max_num_tokens, BLOCK_M)
|
||||
num_pid_n = tl.cdiv(N, BLOCK_N)
|
||||
pid_m = pid_mn // num_pid_n
|
||||
pid_n = pid_mn % num_pid_n
|
||||
|
||||
cta_m_start = pid_m * BLOCK_M
|
||||
cta_n_start = pid_n * BLOCK_N
|
||||
if cta_m_start >= e_num_tokens:
|
||||
# Early exit
|
||||
return
|
||||
|
||||
cta_m_size = min(BLOCK_M, e_num_tokens - cta_m_start)
|
||||
cta_n_size = min(BLOCK_N, N - cta_n_start)
|
||||
|
||||
a_ptr = a_ptr + expert_id * stride_ae + cta_m_start * stride_am
|
||||
b_ptr = b_ptr + expert_id * stride_be + cta_n_start * stride_bn
|
||||
c_ptr = (
|
||||
c_ptr
|
||||
+ expert_id * stride_ce
|
||||
+ cta_m_start * stride_cm
|
||||
+ cta_n_start * stride_cn
|
||||
)
|
||||
|
||||
offs_bn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N).to(tl.int64)) % N
|
||||
|
||||
if use_fp8_w8a8:
|
||||
a_scale_ptr = a_scale_ptr + expert_id * stride_ase
|
||||
b_scale_ptr = b_scale_ptr + expert_id * stride_bse
|
||||
|
||||
# block-wise
|
||||
if group_k > 0 and group_n > 0 or per_act_token_quant:
|
||||
a_scale_ptr = a_scale_ptr + cta_m_start * stride_asm
|
||||
|
||||
expert_triton_kernel(
|
||||
a_ptr,
|
||||
b_ptr,
|
||||
c_ptr,
|
||||
expert_id,
|
||||
compute_type,
|
||||
cta_m_size, # M
|
||||
cta_n_size, # N
|
||||
K, # K
|
||||
a_scale_ptr,
|
||||
b_scale_ptr,
|
||||
b_zp_ptr,
|
||||
# Strides
|
||||
stride_am,
|
||||
stride_ak,
|
||||
stride_bk,
|
||||
stride_bn,
|
||||
stride_cm,
|
||||
stride_cn,
|
||||
stride_ase,
|
||||
stride_asm,
|
||||
stride_ask,
|
||||
stride_bse,
|
||||
stride_bsk,
|
||||
stride_bsn,
|
||||
# offsets
|
||||
offs_bn,
|
||||
# Blockwise quantization data
|
||||
group_n,
|
||||
group_k,
|
||||
# Quantization schemes
|
||||
use_fp8_w8a8,
|
||||
use_int8_w8a16,
|
||||
per_act_token_quant,
|
||||
# Kernel config
|
||||
BLOCK_M,
|
||||
BLOCK_N,
|
||||
BLOCK_K,
|
||||
)
|
||||
|
||||
|
||||
def invoke_moe_batched_triton_kernel(
|
||||
A: torch.Tensor, # [E, max_tokens, K]
|
||||
B: torch.Tensor, # [E, N, K]
|
||||
C: torch.Tensor, # [E, max_tokens, N]
|
||||
expert_num_tokens: torch.Tensor, # [E]
|
||||
compute_type: tl.dtype,
|
||||
# Quantization data
|
||||
A_scale: torch.Tensor | None,
|
||||
B_scale: torch.Tensor | None,
|
||||
B_zp: torch.Tensor,
|
||||
# Quantization schemes
|
||||
use_fp8_w8a8: bool,
|
||||
use_int8_w8a16: bool,
|
||||
use_int4_w4a16: bool,
|
||||
config: dict[str, int],
|
||||
per_act_token_quant: bool,
|
||||
block_shape: list[int] | None = None,
|
||||
):
|
||||
assert not use_int4_w4a16
|
||||
max_num_tokens = A.size(1)
|
||||
K = A.size(2)
|
||||
N = C.size(2)
|
||||
|
||||
BLOCK_M = config["BLOCK_SIZE_M"]
|
||||
BLOCK_N = config["BLOCK_SIZE_N"]
|
||||
BLOCK_K = config["BLOCK_SIZE_K"]
|
||||
|
||||
grid = (
|
||||
expert_num_tokens.size(0),
|
||||
triton.cdiv(max_num_tokens, BLOCK_M) * triton.cdiv(B.size(1), BLOCK_N),
|
||||
)
|
||||
|
||||
A_scale = normalize_batched_scales_shape(A_scale, expert_num_tokens.shape[0])
|
||||
|
||||
if B_scale is not None and B_scale.ndim == 1:
|
||||
assert B_scale.numel() == expert_num_tokens.shape[0]
|
||||
B_scale = B_scale.view(-1, 1, 1)
|
||||
|
||||
assert A_scale is None or A_scale.ndim == 3, (
|
||||
f"{0 if A_scale is None else A_scale.shape}"
|
||||
)
|
||||
assert B_scale is None or B_scale.ndim == 1 or B_scale.ndim == 3, (
|
||||
f"{0 if B_scale is None else B_scale.shape}"
|
||||
)
|
||||
|
||||
if B_scale is not None:
|
||||
if B_scale.ndim == 1:
|
||||
stride_bse = 1
|
||||
stride_bsk = 0
|
||||
stride_bsn = 0
|
||||
else:
|
||||
stride_bse = B_scale.stride(0)
|
||||
stride_bsk = B_scale.stride(2)
|
||||
stride_bsn = B_scale.stride(1)
|
||||
|
||||
else:
|
||||
stride_bse = 0
|
||||
stride_bsk = 0
|
||||
stride_bsn = 0
|
||||
|
||||
if A_scale is not None:
|
||||
stride_ase = A_scale.stride(0)
|
||||
stride_asm = A_scale.stride(1)
|
||||
stride_ask = A_scale.stride(2)
|
||||
else:
|
||||
stride_ase = 0
|
||||
stride_asm = 0
|
||||
stride_ask = 0
|
||||
|
||||
batched_triton_kernel[grid](
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
expert_num_tokens,
|
||||
compute_type,
|
||||
# Dimensions
|
||||
max_num_tokens,
|
||||
K,
|
||||
N,
|
||||
# Quantization data
|
||||
A_scale,
|
||||
B_scale,
|
||||
B_zp,
|
||||
# Strides
|
||||
A.stride(0),
|
||||
A.stride(1),
|
||||
A.stride(2),
|
||||
B.stride(0),
|
||||
B.stride(2),
|
||||
B.stride(1),
|
||||
C.stride(0),
|
||||
C.stride(1),
|
||||
C.stride(2),
|
||||
stride_ase,
|
||||
stride_asm,
|
||||
stride_ask,
|
||||
stride_bse,
|
||||
stride_bsk,
|
||||
stride_bsn,
|
||||
# Blockwise quantization data
|
||||
0 if block_shape is None else block_shape[0],
|
||||
0 if block_shape is None else block_shape[1],
|
||||
# Quantization schemes
|
||||
use_fp8_w8a8,
|
||||
use_int8_w8a16,
|
||||
per_act_token_quant,
|
||||
# Kernel config
|
||||
BLOCK_M=BLOCK_M,
|
||||
BLOCK_N=BLOCK_N,
|
||||
BLOCK_K=BLOCK_K,
|
||||
)
|
||||
|
||||
|
||||
class NaiveBatchedExperts(mk.FusedMoEExpertsModular):
|
||||
"""
|
||||
A reference MoE expert class that operates on expert batched format,
|
||||
i.e. E x max_num_tokens x K. This is the format that the batched
|
||||
dispatch/combine kernels use.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
moe_config: FusedMoEConfig,
|
||||
quant_config: FusedMoEQuantConfig,
|
||||
max_num_tokens: int,
|
||||
num_dispatchers: int,
|
||||
):
|
||||
super().__init__(
|
||||
moe_config=moe_config,
|
||||
quant_config=quant_config,
|
||||
max_num_tokens=max_num_tokens,
|
||||
num_dispatchers=num_dispatchers,
|
||||
)
|
||||
assert not self.quant_config.use_int8_w8a8, "NYI"
|
||||
assert not self.quant_config.use_int8_w8a16, "NYI"
|
||||
assert not self.quant_config.use_int4_w4a16, "NYI"
|
||||
assert self.quant_config.ocp_mx_scheme is None, "NYI"
|
||||
|
||||
@staticmethod
|
||||
def activation_format() -> mk.FusedMoEActivationFormat:
|
||||
return mk.FusedMoEActivationFormat.BatchedExperts
|
||||
|
||||
@staticmethod
|
||||
def _supports_current_device() -> bool:
|
||||
raise NotImplementedError(
|
||||
"NaiveBatchedExperts is not yet used by an Oracle. "
|
||||
"This method should not be called."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _supports_no_act_and_mul() -> bool:
|
||||
raise NotImplementedError(
|
||||
"NaiveBatchedExperts is not yet used by an Oracle. "
|
||||
"This method should not be called."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _supports_quant_scheme(
|
||||
weight_key: QuantKey | None,
|
||||
activation_key: QuantKey | None,
|
||||
) -> bool:
|
||||
raise NotImplementedError(
|
||||
"NaiveBatchedExperts is not yet used by an Oracle. "
|
||||
"This method should not be called."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _supports_activation(activation: MoEActivation) -> bool:
|
||||
raise NotImplementedError(
|
||||
"NaiveBatchedExperts is not yet used by an Oracle. "
|
||||
"This method should not be called."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
|
||||
raise NotImplementedError(
|
||||
"NaiveBatchedExperts is not yet used by an Oracle. "
|
||||
"This method should not be called."
|
||||
)
|
||||
|
||||
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
|
||||
# Let PrepareAndFinalize::finalize() decide the impl.
|
||||
return TopKWeightAndReduceDelegate()
|
||||
|
||||
def workspace_shapes(
|
||||
self,
|
||||
M: int,
|
||||
N: int,
|
||||
K: int,
|
||||
topk: int,
|
||||
global_num_experts: int,
|
||||
local_num_experts: int,
|
||||
expert_tokens_meta: mk.ExpertTokensMetadata | None,
|
||||
activation: MoEActivation,
|
||||
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
|
||||
assert self.num_dispatchers is not None
|
||||
assert self.max_num_tokens is not None
|
||||
num_dp = self.num_dispatchers
|
||||
num_experts = local_num_experts
|
||||
workspace13 = (num_experts, self.max_num_tokens * num_dp, K)
|
||||
workspace2 = (self.max_num_tokens * num_dp, N)
|
||||
output = workspace13
|
||||
return (workspace13, workspace2, output)
|
||||
|
||||
def dequant(self, t: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
|
||||
assert self.quant_config.is_quantized
|
||||
f32 = torch.float32
|
||||
if self.quant_config.is_per_act_token or self.quant_config.is_per_tensor:
|
||||
return t.to(f32) * scale
|
||||
else:
|
||||
return t.to(f32) * group_broadcast(scale, t.shape)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
output: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
w1: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
activation: MoEActivation,
|
||||
global_num_experts: int,
|
||||
expert_map: torch.Tensor | None,
|
||||
a1q_scale: torch.Tensor | None,
|
||||
a2_scale: torch.Tensor | None,
|
||||
workspace13: torch.Tensor,
|
||||
workspace2: torch.Tensor,
|
||||
expert_tokens_meta: mk.ExpertTokensMetadata | None,
|
||||
apply_router_weight_on_input: bool,
|
||||
):
|
||||
assert hidden_states.dim() == 3
|
||||
assert expert_tokens_meta is not None
|
||||
expert_num_tokens = expert_tokens_meta.expert_num_tokens
|
||||
|
||||
num_local_experts = w1.size(0)
|
||||
assert num_local_experts == w1.size(0), f"{num_local_experts} == {w1.size(0)}"
|
||||
|
||||
N = w1.size(1) // 2
|
||||
|
||||
for expert in range(num_local_experts):
|
||||
# Indexing expert_num_tokens doesn't work w/cudagraphs or inductor
|
||||
if (
|
||||
torch.compiler.is_compiling()
|
||||
or torch.cuda.is_current_stream_capturing()
|
||||
):
|
||||
num = hidden_states.shape[1]
|
||||
else:
|
||||
num = int(expert_num_tokens[expert].item())
|
||||
|
||||
if num == 0:
|
||||
continue
|
||||
|
||||
tmp = _resize_cache(workspace2, (num, N))
|
||||
|
||||
if self.quant_config.is_quantized:
|
||||
assert a1q_scale is not None and self.w1_scale is not None
|
||||
input = self.dequant(hidden_states[expert, :, :], a1q_scale[expert])
|
||||
w1_dq = self.dequant(w1[expert], self.w1_scale[expert])
|
||||
input = input[:num] @ w1_dq.transpose(0, 1)
|
||||
else:
|
||||
input = hidden_states[expert, :num, :] @ w1[expert].transpose(0, 1)
|
||||
|
||||
self.activation(activation, tmp, input.to(tmp.dtype))
|
||||
|
||||
if self.quant_config.is_quantized:
|
||||
assert self.w2_scale is not None
|
||||
w2_dq = self.dequant(w2[expert], self.w2_scale[expert])
|
||||
else:
|
||||
w2_dq = w2[expert]
|
||||
|
||||
output[expert, :num, :] = tmp @ w2_dq.transpose(0, 1).to(tmp.dtype)
|
||||
|
||||
|
||||
def batched_moe_kernel_quantize_input(
|
||||
A: torch.Tensor,
|
||||
A_scale: torch.Tensor | None,
|
||||
num_tokens: int,
|
||||
E: int,
|
||||
N: int,
|
||||
expert_num_tokens: torch.Tensor,
|
||||
qtype: torch.dtype | None,
|
||||
per_act_token_quant: bool,
|
||||
block_shape: list[int] | None = None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
if torch.compiler.is_compiling() or torch.cuda.is_current_stream_capturing():
|
||||
# Note: this does a bunch of extra work because expert_num_tokens is
|
||||
# ignored but it does support torch.compile + cudagraphs.
|
||||
hidden_dim = A.size(-1)
|
||||
assert A_scale is None or A_scale.ndim <= 2, (
|
||||
f"{A_scale.shape if A_scale is not None else None}"
|
||||
)
|
||||
A_q, A_q_scale = moe_kernel_quantize_input(
|
||||
A.view(-1, hidden_dim), A_scale, qtype, per_act_token_quant, block_shape
|
||||
)
|
||||
A_q = A_q.view(E, -1, hidden_dim)
|
||||
A_q_scale = normalize_batched_scales_shape(A_q_scale, E)
|
||||
|
||||
return A_q, A_q_scale
|
||||
elif qtype is None:
|
||||
return A, normalize_batched_scales_shape(A_scale, E)
|
||||
else:
|
||||
A_q = torch.empty_like(A, dtype=qtype)
|
||||
|
||||
if per_act_token_quant:
|
||||
assert block_shape is None
|
||||
scale_shape = (E, num_tokens, 1)
|
||||
elif block_shape is not None:
|
||||
_, block_k = block_shape
|
||||
k_tiles = (A.shape[-1] + block_k - 1) // block_k
|
||||
scale_shape = (E, num_tokens, k_tiles)
|
||||
else:
|
||||
scale_shape = (E, 1, 1)
|
||||
|
||||
A_q_scale = torch.zeros(scale_shape, dtype=torch.float32, device=A.device)
|
||||
|
||||
num_experts = expert_num_tokens.numel()
|
||||
|
||||
A_scale = normalize_batched_scales_shape(A_scale, num_experts)
|
||||
|
||||
for e in range(E):
|
||||
num_tokens = int(expert_num_tokens[e].item())
|
||||
if num_tokens > 0:
|
||||
if A_scale is not None:
|
||||
scales = A_scale[e, : min(num_tokens, A_scale.shape[1])]
|
||||
else:
|
||||
scales = None
|
||||
A_q[e, :num_tokens], tmp_scale = moe_kernel_quantize_input(
|
||||
A[e, :num_tokens],
|
||||
scales,
|
||||
qtype,
|
||||
per_act_token_quant,
|
||||
block_shape,
|
||||
)
|
||||
assert tmp_scale is not None
|
||||
A_q_scale[e, : tmp_scale.shape[0]] = tmp_scale
|
||||
|
||||
return A_q, A_q_scale
|
||||
|
||||
|
||||
class BatchedTritonExperts(mk.FusedMoEExpertsModular):
|
||||
"""
|
||||
A Triton based MoE expert class that operates on expert batched format,
|
||||
i.e. E x max_num_tokens x K. This is the format that the batched
|
||||
dispatch/combine kernels use.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
moe_config: FusedMoEConfig,
|
||||
quant_config: FusedMoEQuantConfig,
|
||||
max_num_tokens: int,
|
||||
num_dispatchers: int,
|
||||
):
|
||||
super().__init__(
|
||||
moe_config=moe_config,
|
||||
quant_config=quant_config,
|
||||
max_num_tokens=max_num_tokens,
|
||||
num_dispatchers=num_dispatchers,
|
||||
)
|
||||
assert not self.quant_config.use_int8_w8a8, "NYI"
|
||||
assert not self.quant_config.use_int8_w8a16, "NYI"
|
||||
assert not self.quant_config.use_int4_w4a16, "NYI"
|
||||
assert self.quant_config.ocp_mx_scheme is None, "NYI"
|
||||
|
||||
@staticmethod
|
||||
def activation_format() -> mk.FusedMoEActivationFormat:
|
||||
return mk.FusedMoEActivationFormat.BatchedExperts
|
||||
|
||||
@staticmethod
|
||||
def _supports_current_device() -> bool:
|
||||
return current_platform.is_cuda_alike()
|
||||
|
||||
@staticmethod
|
||||
def _supports_no_act_and_mul() -> bool:
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _supports_quant_scheme(
|
||||
weight_key: QuantKey | None,
|
||||
activation_key: QuantKey | None,
|
||||
) -> bool:
|
||||
p = current_platform
|
||||
if p.is_rocm():
|
||||
from vllm.platforms.rocm import on_gfx9
|
||||
|
||||
is_rocm_on_gfx9 = on_gfx9()
|
||||
else:
|
||||
is_rocm_on_gfx9 = False
|
||||
|
||||
device_supports_fp8 = is_rocm_on_gfx9 or (
|
||||
p.is_cuda() and p.has_device_capability((8, 9))
|
||||
)
|
||||
|
||||
supported: list[tuple[QuantKey | None, QuantKey | None]] = [(None, None)]
|
||||
if device_supports_fp8:
|
||||
supported += [
|
||||
(kFp8Static128BlockSym, kFp8Dynamic128Sym),
|
||||
(kFp8StaticChannelSym, kFp8DynamicTokenSym),
|
||||
(kFp8StaticTensorSym, kFp8DynamicTokenSym),
|
||||
(kFp8StaticTensorSym, kFp8StaticTensorSym),
|
||||
(kFp8StaticTensorSym, kFp8DynamicTensorSym),
|
||||
]
|
||||
return (weight_key, activation_key) in supported
|
||||
|
||||
@staticmethod
|
||||
def _supports_activation(activation: MoEActivation) -> bool:
|
||||
return activation in [
|
||||
MoEActivation.SILU,
|
||||
MoEActivation.GELU,
|
||||
MoEActivation.GELU_TANH,
|
||||
MoEActivation.SWIGLUOAI,
|
||||
MoEActivation.SILU_NO_MUL,
|
||||
MoEActivation.GELU_NO_MUL,
|
||||
MoEActivation.GELU_TANH_NO_MUL,
|
||||
MoEActivation.RELU2_NO_MUL,
|
||||
]
|
||||
|
||||
@staticmethod
|
||||
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
|
||||
return True
|
||||
|
||||
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
|
||||
# Let PrepareAndFinalize::finalize() decide the impl.
|
||||
return TopKWeightAndReduceDelegate()
|
||||
|
||||
def activation(
|
||||
self, activation: MoEActivation, output: torch.Tensor, input: torch.Tensor
|
||||
) -> None:
|
||||
gemm1_clamp_limit = self.quant_config.gemm1_clamp_limit
|
||||
if activation == MoEActivation.SILU and gemm1_clamp_limit is not None:
|
||||
swiglu_limit_func(output, input, float(gemm1_clamp_limit))
|
||||
return
|
||||
|
||||
super().activation(activation, output, input)
|
||||
|
||||
def workspace_shapes(
|
||||
self,
|
||||
M: int,
|
||||
N: int,
|
||||
K: int,
|
||||
topk: int,
|
||||
global_num_experts: int,
|
||||
local_num_experts: int,
|
||||
expert_tokens_meta: mk.ExpertTokensMetadata | None,
|
||||
activation: MoEActivation,
|
||||
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
|
||||
assert self.num_dispatchers is not None
|
||||
assert self.max_num_tokens is not None
|
||||
num_dp = self.num_dispatchers
|
||||
num_experts = local_num_experts
|
||||
max_num_tokens = self.max_num_tokens
|
||||
activation_out_dim = self.adjust_N_for_activation(N, activation)
|
||||
workspace13 = (num_experts, max_num_tokens * num_dp, max(K, N))
|
||||
workspace2 = (num_experts, max_num_tokens * num_dp, activation_out_dim)
|
||||
output = (num_experts, max_num_tokens * num_dp, K)
|
||||
return (workspace13, workspace2, output)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
output: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
w1: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
activation: MoEActivation,
|
||||
global_num_experts: int,
|
||||
expert_map: torch.Tensor | None,
|
||||
a1q_scale: torch.Tensor | None,
|
||||
a2_scale: torch.Tensor | None,
|
||||
workspace13: torch.Tensor,
|
||||
workspace2: torch.Tensor,
|
||||
expert_tokens_meta: mk.ExpertTokensMetadata | None,
|
||||
apply_router_weight_on_input: bool,
|
||||
):
|
||||
# Check constraints.
|
||||
if self.quant_config.use_int4_w4a16:
|
||||
assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch"
|
||||
else:
|
||||
assert hidden_states.size(-1) == w1.size(2), (
|
||||
f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}"
|
||||
)
|
||||
|
||||
assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
|
||||
assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
|
||||
assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
|
||||
assert hidden_states.dtype in [
|
||||
torch.float32,
|
||||
torch.float16,
|
||||
torch.bfloat16,
|
||||
torch.float8_e4m3fn,
|
||||
torch.float8_e4m3fnuz,
|
||||
]
|
||||
assert expert_tokens_meta is not None
|
||||
|
||||
expert_num_tokens = expert_tokens_meta.expert_num_tokens
|
||||
|
||||
E, max_num_tokens, N, K, top_k_num = self.moe_problem_size(
|
||||
hidden_states, w1, w2, topk_ids
|
||||
)
|
||||
|
||||
assert w1.size(0) == E
|
||||
assert w2.size(0) == E
|
||||
|
||||
config_dtype = self.quant_config.config_name(hidden_states.dtype)
|
||||
|
||||
config = try_get_optimal_moe_config(
|
||||
w1.size(),
|
||||
w2.size(),
|
||||
top_k_num,
|
||||
config_dtype,
|
||||
max_num_tokens,
|
||||
block_shape=self.block_shape,
|
||||
)
|
||||
|
||||
if hidden_states.dtype == torch.bfloat16:
|
||||
compute_type = tl.bfloat16
|
||||
elif hidden_states.dtype == torch.float16:
|
||||
compute_type = tl.float16
|
||||
elif hidden_states.dtype == torch.float32:
|
||||
compute_type = tl.float32
|
||||
elif hidden_states.dtype == current_platform.fp8_dtype():
|
||||
compute_type = tl.bfloat16
|
||||
else:
|
||||
raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")
|
||||
|
||||
# We can reuse the memory between these because by the time we need
|
||||
# cache3, we're done with cache1
|
||||
intermediate_cache1 = _resize_cache(workspace13, (E, max_num_tokens, N))
|
||||
activation_out_dim = self.adjust_N_for_activation(N, activation)
|
||||
intermediate_cache2 = _resize_cache(
|
||||
workspace2, (E, max_num_tokens, activation_out_dim)
|
||||
)
|
||||
|
||||
# TODO(bnell): should this be done for any quantized type?
|
||||
if self.quant_config.use_fp8_w8a8:
|
||||
intermediate_cache1.fill_(0)
|
||||
|
||||
a1q_scale = normalize_batched_scales_shape(a1q_scale, E)
|
||||
|
||||
# MM1
|
||||
invoke_moe_batched_triton_kernel(
|
||||
A=hidden_states,
|
||||
B=w1,
|
||||
C=intermediate_cache1,
|
||||
expert_num_tokens=expert_num_tokens,
|
||||
compute_type=compute_type,
|
||||
A_scale=a1q_scale,
|
||||
B_scale=self.w1_scale,
|
||||
B_zp=self.w1_zp,
|
||||
use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
|
||||
use_int8_w8a16=self.quant_config.use_int8_w8a16,
|
||||
use_int4_w4a16=self.quant_config.use_int4_w4a16,
|
||||
config=config,
|
||||
per_act_token_quant=self.per_act_token_quant,
|
||||
block_shape=self.block_shape,
|
||||
)
|
||||
|
||||
intermediate_cache2.fill_(0)
|
||||
|
||||
# TODO (bnell): use triton utility from batched deep gemm.
|
||||
self.activation(
|
||||
activation,
|
||||
intermediate_cache2.view(-1, activation_out_dim),
|
||||
intermediate_cache1.view(-1, N),
|
||||
)
|
||||
|
||||
qintermediate_cache2, a2q_scale = batched_moe_kernel_quantize_input(
|
||||
intermediate_cache2,
|
||||
a2_scale,
|
||||
max_num_tokens,
|
||||
E,
|
||||
N,
|
||||
expert_num_tokens,
|
||||
self.quant_dtype,
|
||||
self.per_act_token_quant,
|
||||
self.block_shape,
|
||||
)
|
||||
|
||||
invoke_moe_batched_triton_kernel(
|
||||
A=qintermediate_cache2,
|
||||
B=w2,
|
||||
C=output,
|
||||
expert_num_tokens=expert_num_tokens,
|
||||
compute_type=compute_type,
|
||||
A_scale=a2q_scale,
|
||||
B_scale=self.w2_scale,
|
||||
B_zp=self.w2_zp,
|
||||
use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
|
||||
use_int8_w8a16=self.quant_config.use_int8_w8a16,
|
||||
use_int4_w4a16=self.quant_config.use_int4_w4a16,
|
||||
config=config,
|
||||
per_act_token_quant=self.per_act_token_quant,
|
||||
block_shape=self.block_shape,
|
||||
)
|
||||
Reference in New Issue
Block a user