430 lines
18 KiB
Python
430 lines
18 KiB
Python
|
import functools
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
from typing import Any, Dict, Optional, Tuple
|
||
|
|
from loguru import logger
|
||
|
|
import torch
|
||
|
|
import ixformer.inference.functions as ops
|
||
|
|
|
||
|
|
CHUNK_SIZE = int(os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", "65536"))
|
||
|
|
|
||
|
|
def fused_topk(
|
||
|
|
hidden_states: torch.Tensor,
|
||
|
|
gating_output: torch.Tensor,
|
||
|
|
topk: int,
|
||
|
|
renormalize: bool,
|
||
|
|
):
|
||
|
|
assert hidden_states.shape[0] == gating_output.shape[0], (
|
||
|
|
"Number of tokens mismatch")
|
||
|
|
|
||
|
|
M, _ = hidden_states.shape
|
||
|
|
|
||
|
|
topk_weights = torch.empty(M,
|
||
|
|
topk,
|
||
|
|
dtype=torch.float32,
|
||
|
|
device=hidden_states.device)
|
||
|
|
topk_ids = torch.empty(M,
|
||
|
|
topk,
|
||
|
|
dtype=torch.int32,
|
||
|
|
device=hidden_states.device)
|
||
|
|
token_expert_indicies = torch.empty(M,
|
||
|
|
topk,
|
||
|
|
dtype=torch.int32,
|
||
|
|
device=hidden_states.device)
|
||
|
|
ops.vllm_moe_topk_softmax(
|
||
|
|
topk_weights,
|
||
|
|
topk_ids,
|
||
|
|
token_expert_indicies,
|
||
|
|
gating_output.float(), # TODO(woosuk): Optimize this.
|
||
|
|
)
|
||
|
|
del token_expert_indicies # Not used. Will be used in the future.
|
||
|
|
|
||
|
|
if renormalize:
|
||
|
|
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
||
|
|
return topk_weights, topk_ids
|
||
|
|
|
||
|
|
# This is used by the Deepseek-V2 model
|
||
|
|
def grouped_topk(hidden_states: torch.Tensor,
|
||
|
|
gating_output: torch.Tensor,
|
||
|
|
topk: int,
|
||
|
|
renormalize: bool,
|
||
|
|
num_expert_group: int = 0,
|
||
|
|
topk_group: int = 0):
|
||
|
|
|
||
|
|
assert hidden_states.shape[0] == gating_output.shape[0], (
|
||
|
|
"Number of tokens mismatch")
|
||
|
|
|
||
|
|
scores = torch.softmax(gating_output, dim=-1)
|
||
|
|
num_token = scores.shape[0]
|
||
|
|
group_scores = scores.view(num_token, num_expert_group,
|
||
|
|
-1).max(dim=-1).values # [n, n_group]
|
||
|
|
group_idx = torch.topk(group_scores, k=topk_group, dim=-1,
|
||
|
|
sorted=False)[1] # [n, top_k_group]
|
||
|
|
group_mask = torch.zeros_like(group_scores) # [n, n_group]
|
||
|
|
group_mask.scatter_(1, group_idx, 1) # [n, n_group]
|
||
|
|
score_mask = group_mask.unsqueeze(-1).expand(
|
||
|
|
num_token, num_expert_group,
|
||
|
|
scores.shape[-1] // num_expert_group).reshape(num_token, -1) # [n, e]
|
||
|
|
tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) # [n, e]
|
||
|
|
topk_weights, topk_ids = torch.topk(tmp_scores,
|
||
|
|
k=topk,
|
||
|
|
dim=-1,
|
||
|
|
sorted=False)
|
||
|
|
|
||
|
|
if renormalize:
|
||
|
|
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
||
|
|
return topk_weights, topk_ids
|
||
|
|
|
||
|
|
def get_config_file_name(E: int, N: int, dtype: Optional[str]) -> str:
|
||
|
|
device_name = torch.cuda.get_device_name().replace(" ", "_")
|
||
|
|
dtype_selector = "" if not dtype else f",dtype={dtype}"
|
||
|
|
return f"E={E},N={N},device_name={device_name}{dtype_selector}.json"
|
||
|
|
|
||
|
|
@functools.lru_cache
|
||
|
|
def get_moe_configs(E: int, N: int,
|
||
|
|
dtype: Optional[str]) -> Optional[Dict[int, Any]]:
|
||
|
|
"""
|
||
|
|
Return optimized configurations for the fused MoE kernel.
|
||
|
|
|
||
|
|
The return value will be a dictionary that maps an irregular grid of
|
||
|
|
batch sizes to configurations of the fused_moe kernel. To evaluate the
|
||
|
|
kernel on a given batch size bs, the closest batch size in the grid should
|
||
|
|
be picked and the associated configuration chosen to invoke the kernel.
|
||
|
|
"""
|
||
|
|
|
||
|
|
# First look up if an optimized configuration is available in the configs
|
||
|
|
# directory
|
||
|
|
json_file_name = get_config_file_name(E, N, dtype)
|
||
|
|
|
||
|
|
config_file_path = os.path.join(
|
||
|
|
os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name)
|
||
|
|
if os.path.exists(config_file_path):
|
||
|
|
with open(config_file_path) as f:
|
||
|
|
logger.info("Using configuration from %s for MoE layer.",
|
||
|
|
config_file_path)
|
||
|
|
# If a configuration has been found, return it
|
||
|
|
return {int(key): val for key, val in json.load(f).items()}
|
||
|
|
|
||
|
|
# If no optimized configuration is available, we will use the default
|
||
|
|
# configuration
|
||
|
|
return None
|
||
|
|
|
||
|
|
def get_default_config(
|
||
|
|
M: int,
|
||
|
|
E: int,
|
||
|
|
N: int,
|
||
|
|
K: int,
|
||
|
|
topk: int,
|
||
|
|
dtype: Optional[str],
|
||
|
|
) -> Dict[str, int]:
|
||
|
|
config = {
|
||
|
|
'BLOCK_SIZE_M': 64,
|
||
|
|
'BLOCK_SIZE_N': 64,
|
||
|
|
'BLOCK_SIZE_K': 32,
|
||
|
|
'GROUP_SIZE_M': 8
|
||
|
|
}
|
||
|
|
if M <= E:
|
||
|
|
config = {
|
||
|
|
'BLOCK_SIZE_M': 16,
|
||
|
|
'BLOCK_SIZE_N': 32,
|
||
|
|
'BLOCK_SIZE_K': 64,
|
||
|
|
'GROUP_SIZE_M': 1
|
||
|
|
}
|
||
|
|
numel = M * topk
|
||
|
|
if numel <= 64:
|
||
|
|
config['BLOCK_SIZE_M'] = 32
|
||
|
|
elif numel <= 1024:
|
||
|
|
config['BLOCK_SIZE_M'] = 64
|
||
|
|
else:
|
||
|
|
config['BLOCK_SIZE_M'] = 256
|
||
|
|
return config
|
||
|
|
|
||
|
|
def try_get_optimal_moe_config(
|
||
|
|
w1_shape: Tuple[int, ...],
|
||
|
|
w2_shape: Tuple[int, ...],
|
||
|
|
top_k: int,
|
||
|
|
dtype: Optional[str],
|
||
|
|
M: int,
|
||
|
|
override_config: Optional[Dict[str, Any]] = None,
|
||
|
|
):
|
||
|
|
if override_config:
|
||
|
|
config = override_config
|
||
|
|
else:
|
||
|
|
# First try to load optimal config from the file
|
||
|
|
E, _, N = w2_shape
|
||
|
|
configs = get_moe_configs(E, N, dtype)
|
||
|
|
|
||
|
|
if configs:
|
||
|
|
# If an optimal configuration map has been found, look up the
|
||
|
|
# optimal config
|
||
|
|
config = configs[min(configs.keys(), key=lambda x: abs(x - M))]
|
||
|
|
else:
|
||
|
|
# Else use the default config
|
||
|
|
config = get_default_config(M, E, N, w1_shape[2], top_k, dtype)
|
||
|
|
return config
|
||
|
|
|
||
|
|
def moe_align_block_size(
|
||
|
|
topk_ids: torch.Tensor, block_size: int,
|
||
|
|
num_experts: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||
|
|
"""
|
||
|
|
Aligns the token distribution across experts to be compatible with block
|
||
|
|
size for matrix multiplication.
|
||
|
|
|
||
|
|
Parameters:
|
||
|
|
- topk_ids: A tensor of shape [total_tokens, top_k] representing the
|
||
|
|
top-k expert indices for each token.
|
||
|
|
- block_size: The block size used in block matrix multiplication.
|
||
|
|
- num_experts: The total number of experts.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
- sorted_token_ids: A tensor containing the sorted token indices according
|
||
|
|
to their allocated expert.
|
||
|
|
- expert_ids: A tensor indicating the assigned expert index for each block.
|
||
|
|
- num_tokens_post_padded: The total number of tokens after padding,
|
||
|
|
ensuring divisibility by block_size.
|
||
|
|
|
||
|
|
This function pads the number of tokens that each expert needs to process
|
||
|
|
so that it is divisible by block_size.
|
||
|
|
Padding ensures that during block matrix multiplication, the dimensions
|
||
|
|
align correctly.
|
||
|
|
|
||
|
|
Example:
|
||
|
|
Given topk_ids = [[2, 3, 4], [1, 2, 4], [1, 3, 4], [1, 2, 3]],
|
||
|
|
block_size = 4, and num_experts = 4:
|
||
|
|
- We initially have 12 tokens (after repeating 'top_k' times) and 4 experts,
|
||
|
|
with each expert needing to process 3 tokens.
|
||
|
|
- As block_size is 4, we pad 1 token for each expert.
|
||
|
|
- First, flatten topk_ids to [2, 3, 4, 1, 2, 4, 1, 3, 4, 1, 2, 3].
|
||
|
|
- Then append padding tokens [12, 12, 12, 12] for each block.
|
||
|
|
- After sorting by expert index, we obtain token_ids
|
||
|
|
[3, 6, 9, 12, 0, 4, 10, 12, 1, 7, 11, 12, 2, 5, 8, 12].
|
||
|
|
Tokens 12 are non-existent (padding) and are ignored in
|
||
|
|
the subsequent matrix multiplication.
|
||
|
|
- The padding ensures that the total number of tokens is now divisible
|
||
|
|
by block_size for proper block matrix operations.
|
||
|
|
"""
|
||
|
|
max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
|
||
|
|
sorted_ids = torch.empty((max_num_tokens_padded, ),
|
||
|
|
dtype=torch.int32,
|
||
|
|
device=topk_ids.device)
|
||
|
|
sorted_ids.fill_(topk_ids.numel())
|
||
|
|
# max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
|
||
|
|
max_num_m_blocks = topk_ids.numel() + num_experts
|
||
|
|
expert_ids = torch.empty((max_num_m_blocks, ),
|
||
|
|
dtype=torch.int32,
|
||
|
|
device=topk_ids.device)
|
||
|
|
num_tokens_post_pad = torch.empty((1),
|
||
|
|
dtype=torch.int32,
|
||
|
|
device=topk_ids.device)
|
||
|
|
ops.vllm_moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids,
|
||
|
|
expert_ids, num_tokens_post_pad)
|
||
|
|
return sorted_ids, expert_ids, num_tokens_post_pad
|
||
|
|
|
||
|
|
def invoke_fused_moe_kernel(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,
|
||
|
|
A_scale: Optional[torch.Tensor],
|
||
|
|
B_scale: Optional[torch.Tensor],
|
||
|
|
topk_weights: torch.Tensor, topk_ids: torch.Tensor,
|
||
|
|
sorted_token_ids: torch.Tensor,
|
||
|
|
expert_ids: torch.Tensor,
|
||
|
|
num_tokens_post_padded: torch.Tensor,
|
||
|
|
mul_routed_weight: bool, top_k: int,
|
||
|
|
config: Dict[str, Any], compute_type: torch.dtype,
|
||
|
|
use_fp8: bool) -> None:
|
||
|
|
ops.vllm_invoke_fused_moe_kernel(A, B, C, topk_weights, topk_ids,
|
||
|
|
sorted_token_ids,expert_ids, num_tokens_post_padded,
|
||
|
|
mul_routed_weight, top_k, config['BLOCK_SIZE_M'])
|
||
|
|
|
||
|
|
def fused_experts(hidden_states: torch.Tensor,
|
||
|
|
w1: torch.Tensor,
|
||
|
|
w2: torch.Tensor,
|
||
|
|
topk_weights: torch.Tensor,
|
||
|
|
topk_ids: torch.Tensor,
|
||
|
|
inplace: bool = False,
|
||
|
|
override_config: Optional[Dict[str, Any]] = None,
|
||
|
|
use_fp8: bool = False,
|
||
|
|
w1_scale: Optional[torch.Tensor] = None,
|
||
|
|
w2_scale: Optional[torch.Tensor] = None,
|
||
|
|
a1_scale: Optional[torch.Tensor] = None,
|
||
|
|
a2_scale: Optional[torch.Tensor] = None):
|
||
|
|
# Check constraints.
|
||
|
|
assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
|
||
|
|
assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
|
||
|
|
assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
|
||
|
|
assert w1.is_contiguous(), "Expert weights1 must be contiguous"
|
||
|
|
assert w2.is_contiguous(), "Expert weights2 must be contiguous"
|
||
|
|
assert hidden_states.dtype in [
|
||
|
|
torch.float32, torch.float16, torch.bfloat16
|
||
|
|
]
|
||
|
|
|
||
|
|
num_tokens, _ = hidden_states.shape
|
||
|
|
E, N, _ = w1.shape
|
||
|
|
# We execute the fused_moe kernel in chunks to circumvent this issue:
|
||
|
|
# https://github.com/vllm-project/vllm/issues/5938
|
||
|
|
M = min(num_tokens, CHUNK_SIZE)
|
||
|
|
|
||
|
|
get_config_func = functools.partial(
|
||
|
|
try_get_optimal_moe_config,
|
||
|
|
w1.shape,
|
||
|
|
w2.shape,
|
||
|
|
topk_ids.shape[1],
|
||
|
|
"float8" if use_fp8 else None,
|
||
|
|
override_config=override_config,
|
||
|
|
)
|
||
|
|
|
||
|
|
config = get_config_func(M)
|
||
|
|
|
||
|
|
intermediate_cache1 = torch.empty((M, topk_ids.shape[1], N),
|
||
|
|
device=hidden_states.device,
|
||
|
|
dtype=hidden_states.dtype)
|
||
|
|
intermediate_cache2 = torch.empty((M * topk_ids.shape[1], N // 2),
|
||
|
|
device=hidden_states.device,
|
||
|
|
dtype=hidden_states.dtype)
|
||
|
|
intermediate_cache3 = torch.empty((M, topk_ids.shape[1], w2.shape[1]),
|
||
|
|
device=hidden_states.device,
|
||
|
|
dtype=hidden_states.dtype)
|
||
|
|
|
||
|
|
compute_type = (torch.bfloat16
|
||
|
|
if hidden_states.dtype == torch.bfloat16 else torch.float16)
|
||
|
|
|
||
|
|
if inplace:
|
||
|
|
out_hidden_states = hidden_states
|
||
|
|
else:
|
||
|
|
out_hidden_states = torch.empty_like(hidden_states)
|
||
|
|
|
||
|
|
for chunk in range((num_tokens // CHUNK_SIZE) + 1):
|
||
|
|
begin_chunk_idx, end_chunk_idx = (chunk * CHUNK_SIZE,
|
||
|
|
min((chunk + 1) * CHUNK_SIZE,
|
||
|
|
num_tokens))
|
||
|
|
curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
|
||
|
|
tokens_in_chunk, _ = curr_hidden_states.shape
|
||
|
|
|
||
|
|
if tokens_in_chunk == 0:
|
||
|
|
break
|
||
|
|
|
||
|
|
if tokens_in_chunk < CHUNK_SIZE and chunk > 0:
|
||
|
|
# Adjust the intermediate cache size and config for the last
|
||
|
|
# chunk. Note that in most cases we only have one chunk
|
||
|
|
# so the cache size and config are already set correctly and
|
||
|
|
# do not need to be adjusted.
|
||
|
|
intermediate_cache1 = intermediate_cache1[:tokens_in_chunk]
|
||
|
|
intermediate_cache2 = intermediate_cache2[:tokens_in_chunk]
|
||
|
|
intermediate_cache3 = intermediate_cache3[:tokens_in_chunk]
|
||
|
|
config = get_config_func(tokens_in_chunk)
|
||
|
|
|
||
|
|
curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
|
||
|
|
curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
|
||
|
|
|
||
|
|
sorted_token_ids, expert_ids, num_tokens_post_padded = (
|
||
|
|
moe_align_block_size(curr_topk_ids, config['BLOCK_SIZE_M'], E))
|
||
|
|
|
||
|
|
invoke_fused_moe_kernel(curr_hidden_states,
|
||
|
|
w1,
|
||
|
|
intermediate_cache1,
|
||
|
|
a1_scale,
|
||
|
|
w1_scale,
|
||
|
|
curr_topk_weights,
|
||
|
|
curr_topk_ids,
|
||
|
|
sorted_token_ids,
|
||
|
|
expert_ids,
|
||
|
|
num_tokens_post_padded,
|
||
|
|
False,
|
||
|
|
topk_ids.shape[1],
|
||
|
|
config,
|
||
|
|
compute_type=compute_type,
|
||
|
|
use_fp8=use_fp8)
|
||
|
|
|
||
|
|
ops.silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
|
||
|
|
|
||
|
|
invoke_fused_moe_kernel(intermediate_cache2,
|
||
|
|
w2,
|
||
|
|
intermediate_cache3,
|
||
|
|
a2_scale,
|
||
|
|
w2_scale,
|
||
|
|
curr_topk_weights,
|
||
|
|
curr_topk_ids,
|
||
|
|
sorted_token_ids,
|
||
|
|
expert_ids,
|
||
|
|
num_tokens_post_padded,
|
||
|
|
True,
|
||
|
|
1,
|
||
|
|
config,
|
||
|
|
compute_type=compute_type,
|
||
|
|
use_fp8=use_fp8)
|
||
|
|
|
||
|
|
torch.sum(intermediate_cache3.view(*intermediate_cache3.shape),
|
||
|
|
dim=1,
|
||
|
|
out=out_hidden_states[begin_chunk_idx:end_chunk_idx])
|
||
|
|
return out_hidden_states
|
||
|
|
|
||
|
|
def fused_moe(
|
||
|
|
hidden_states: torch.Tensor,
|
||
|
|
w1: torch.Tensor,
|
||
|
|
w2: torch.Tensor,
|
||
|
|
gating_output: torch.Tensor,
|
||
|
|
topk: int,
|
||
|
|
renormalize: bool,
|
||
|
|
inplace: bool = False,
|
||
|
|
override_config: Optional[Dict[str, Any]] = None,
|
||
|
|
use_grouped_topk: bool = False,
|
||
|
|
num_expert_group: Optional[int] = None,
|
||
|
|
topk_group: Optional[int] = None,
|
||
|
|
use_fp8: bool = False,
|
||
|
|
w1_scale: Optional[torch.Tensor] = None,
|
||
|
|
w2_scale: Optional[torch.Tensor] = None,
|
||
|
|
a1_scale: Optional[torch.Tensor] = None,
|
||
|
|
a2_scale: Optional[torch.Tensor] = None,
|
||
|
|
) -> torch.Tensor:
|
||
|
|
"""
|
||
|
|
This function computes a Mixture of Experts (MoE) layer using two sets of
|
||
|
|
weights, w1 and w2, and top-k gating mechanism.
|
||
|
|
|
||
|
|
Parameters:
|
||
|
|
- hidden_states (torch.Tensor): The input tensor to the MoE layer.
|
||
|
|
- w1 (torch.Tensor): The first set of expert weights.
|
||
|
|
- w2 (torch.Tensor): The second set of expert weights.
|
||
|
|
- gating_output (torch.Tensor): The output of the gating operation
|
||
|
|
(before softmax).
|
||
|
|
- topk (int): The number of top-k experts to select.
|
||
|
|
- renormalize (bool): If True, renormalize the top-k weights to sum to 1.
|
||
|
|
- inplace (bool): If True, perform the operation in-place.
|
||
|
|
Defaults to False.
|
||
|
|
- override_config (Optional[Dict[str, Any]]): Optional override
|
||
|
|
for the kernel configuration.
|
||
|
|
- num_expert_group: Optional[int]: additional parameter for grouped_topk
|
||
|
|
- topk_group: Optional[int]: additional parameter for grouped_topk
|
||
|
|
- use_grouped_topk: If True, use grouped_topk instead of fused_topk
|
||
|
|
note: Deepseekv2 model uses grouped_topk
|
||
|
|
- use_fp8 (bool): If True, use fp8 arithmetic to compute the inner
|
||
|
|
products for w1 and w2. Defaults to False.
|
||
|
|
- w1_scale (Optional[torch.Tensor]): Optional scale to be used for
|
||
|
|
w1.
|
||
|
|
- w2_scale (Optional[torch.Tensor]): Optional scale to be used for
|
||
|
|
w2.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
- torch.Tensor: The output tensor after applying the MoE layer.
|
||
|
|
"""
|
||
|
|
# Check constraints.
|
||
|
|
assert gating_output.shape[1] == w1.shape[0], "Number of experts mismatch"
|
||
|
|
|
||
|
|
if use_grouped_topk:
|
||
|
|
assert num_expert_group is not None and topk_group is not None
|
||
|
|
topk_weights, topk_ids = grouped_topk(hidden_states, gating_output,
|
||
|
|
topk, renormalize,
|
||
|
|
num_expert_group, topk_group)
|
||
|
|
else:
|
||
|
|
topk_weights, topk_ids = fused_topk(hidden_states, gating_output, topk,
|
||
|
|
renormalize)
|
||
|
|
|
||
|
|
return fused_experts(hidden_states,
|
||
|
|
w1,
|
||
|
|
w2,
|
||
|
|
topk_weights,
|
||
|
|
topk_ids,
|
||
|
|
inplace=inplace,
|
||
|
|
override_config=override_config,
|
||
|
|
use_fp8=use_fp8,
|
||
|
|
w1_scale=w1_scale,
|
||
|
|
w2_scale=w2_scale,
|
||
|
|
a1_scale=a1_scale,
|
||
|
|
a2_scale=a2_scale)
|