343 lines
13 KiB
Python
Executable File
343 lines
13 KiB
Python
Executable File
"""CUTLASS based Fused MoE kernels."""
|
|
|
|
from typing import Optional, Tuple
|
|
|
|
import torch
|
|
|
|
from sglang.srt.utils import is_cuda, is_sm90_supported, is_sm100_supported
|
|
|
|
_is_cuda = is_cuda()
|
|
if _is_cuda:
|
|
from sgl_kernel import (
|
|
apply_shuffle_mul_sum,
|
|
es_fp8_blockwise_scaled_grouped_mm,
|
|
es_sm100_mxfp8_blockscaled_grouped_mm,
|
|
es_sm100_mxfp8_blockscaled_grouped_quant,
|
|
fp8_blockwise_scaled_grouped_mm,
|
|
prepare_moe_input,
|
|
shuffle_rows,
|
|
)
|
|
|
|
from sglang.kernels.ops.activation.activation import silu_and_mul
|
|
from sglang.kernels.ops.moe.shuffle_rows_with_scales import (
|
|
shuffle_rows_with_scales,
|
|
)
|
|
|
|
|
|
def cutlass_fused_experts_fp8(
|
|
a: torch.Tensor,
|
|
w1_q: torch.Tensor,
|
|
w2_q: torch.Tensor,
|
|
w1_scale: torch.Tensor,
|
|
w2_scale: torch.Tensor,
|
|
topk_weights: torch.Tensor,
|
|
topk_ids: torch.Tensor,
|
|
a1_strides: torch.Tensor,
|
|
c1_strides: torch.Tensor,
|
|
a2_strides: torch.Tensor,
|
|
c2_strides: torch.Tensor,
|
|
workspace: torch.Tensor,
|
|
a_ptrs: torch.Tensor,
|
|
b_ptrs: torch.Tensor,
|
|
out_ptrs: torch.Tensor,
|
|
a_scales_ptrs: torch.Tensor,
|
|
b_scales_ptrs: torch.Tensor,
|
|
expert_offsets: torch.Tensor,
|
|
problem_sizes1: torch.Tensor,
|
|
problem_sizes2: torch.Tensor,
|
|
use_fp8_blockscale: bool = True,
|
|
use_mxfp8: bool = False,
|
|
output: Optional[torch.Tensor] = None,
|
|
enable_es: Tuple[bool, bool] = (False, False),
|
|
) -> torch.Tensor:
|
|
"""Performs Fused MoE computation using CUTLASS-like kernels with FP8 weights and activations.
|
|
|
|
This function implements a Mixture of Experts (MoE) layer with a SwiGLU/SiLU
|
|
activation, leveraging custom kernels likely derived from CUTLASS principles
|
|
for grouped matrix multiplication (`fp8_blockwise_scaled_grouped_mm`) and
|
|
data preparation (`prepare_moe_input`, `silu_and_mul`).
|
|
|
|
It handles per-token routing, quantizes input activations to FP8 with
|
|
per-token scales, performs the expert computations using FP8 GEMMs with
|
|
pre-quantized FP8 weights (per-block scales), applies the SiLU activation,
|
|
and combines the results weighted by the router scores.
|
|
|
|
Args:
|
|
a (torch.Tensor): Input activations. Shape: `(m, k)`, where `m` is the total
|
|
number of tokens and `k` is the hidden size. Expected dtype: `torch.half`
|
|
or `torch.bfloat16`.
|
|
w1_q (torch.Tensor): Pre-quantized FP8 weight tensor for the first GEMM
|
|
(up-projection part of SwiGLU). Expected shape: `(E, k, n*2)`, where
|
|
`E` is the number of experts, `k` is the hidden size, and `n*2` is the
|
|
intermediate size (`I`). Expected dtype: `torch.float8_e4m3fn`.
|
|
Note: This shape implies weights are stored as (num_experts, hidden_size, intermediate_size).
|
|
w2_q (torch.Tensor): Pre-quantized FP8 weight tensor for the second GEMM
|
|
(down-projection). Expected shape: `(E, n, k)`, where `n` is half the
|
|
intermediate size (`I // 2`). Expected dtype: `torch.float8_e4m3fn`.
|
|
Note: This shape implies weights are stored as (num_experts, intermediate_size // 2, hidden_size).
|
|
w1_scale (torch.Tensor): Scales corresponding to `w1_q` (per-block scales).
|
|
Shape: `(E, num_blocks_n, num_blocks_k)`. Dtype: `torch.float32`.
|
|
w2_scale (torch.Tensor): Scales corresponding to `w2_q` (per-block scales).
|
|
Shape: `(E, num_blocks_k, num_blocks_n)`. Dtype: `torch.float32`.
|
|
topk_weights (torch.Tensor): Router weights for the selected top-k experts
|
|
for each token. Shape: `(m, topk)`. Dtype should ideally match `a`.
|
|
topk_ids (torch.Tensor): Indices of the selected top-k experts for each token.
|
|
Shape: `(m, topk)`. Dtype: `torch.int32`.
|
|
a1_strides (torch.Tensor): Stride information for the first GEMM's 'a' input.
|
|
Passed directly to the underlying kernel. Expected shape `(E,)`, dtype `torch.int64`.
|
|
Note: Its exact usage within `fp8_blockwise_scaled_grouped_mm` needs clarification
|
|
as it's passed as both a_stride and b_stride in the first call.
|
|
c1_strides (torch.Tensor): Stride information for the first GEMM's 'c' output.
|
|
Passed directly to the underlying kernel. Expected shape `(E,)`, dtype `torch.int64`.
|
|
a2_strides (torch.Tensor): Stride information for the second GEMM's 'a' input.
|
|
Passed directly to the underlying kernel. Expected shape `(E,)`, dtype `torch.int64`.
|
|
Note: Its exact usage within `fp8_blockwise_scaled_grouped_mm` needs clarification
|
|
as it's passed as both a_stride and b_stride in the second call.
|
|
c2_strides (torch.Tensor): Stride information for the second GEMM's 'c' output.
|
|
Passed directly to the underlying kernel. Expected shape `(E,)`, dtype `torch.int64`.
|
|
workspace (torch.Tensor): Reusable workspace for the underlying kernel.
|
|
a_ptrs (torch.Tensor): Pointers container for calculating offsets of the input activations for each expert.
|
|
b_ptrs (torch.Tensor): Pointers container for calculating offsets of the input weights for each expert.
|
|
out_ptrs (torch.Tensor): Pointers container for calculating offsets of the output activations for each expert.
|
|
a_scales_ptrs (torch.Tensor): Pointers container for calculating offsets of the input scales for each expert.
|
|
b_scales_ptrs (torch.Tensor): Pointers container for calculating offsets of the input scales for each expert.
|
|
use_fp8_blockscale (bool, optional): Flag indicating usage of FP8 with
|
|
block scaling. Currently, only `True` is supported. Defaults to `True`.
|
|
use_mxfp8 (bool, optional): Flag indicating usage of MXFP8 (UE8M0 scales)
|
|
with SM100 expert-specialization kernels. Defaults to `False`.
|
|
output (torch.Tensor, optional): Output tensor. If not provided, a new tensor will be created.
|
|
enable_es (tuple(bool, bool)): Flag indicating usage of expert specialization kernel for (up-projection, down-projection)
|
|
Returns:
|
|
torch.Tensor: The computed MoE layer output. Shape: `(m, k)`, dtype matches `a`.
|
|
|
|
Raises:
|
|
AssertionError: If input shapes, dtypes, or flags are inconsistent or unsupported.
|
|
NotImplementedError: If CUDA is not available or `sgl_kernel` is not properly installed.
|
|
"""
|
|
assert use_fp8_blockscale, "Only support fp8 blockscale for now"
|
|
assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
|
|
assert w1_q.dtype == torch.float8_e4m3fn
|
|
assert w2_q.dtype == torch.float8_e4m3fn
|
|
assert a.shape[1] == w1_q.shape[1], "Hidden size mismatch w1"
|
|
assert w1_q.shape[2] == w2_q.shape[1] * 2, "Hidden size mismatch w2"
|
|
assert w1_q.shape[0] == w2_q.shape[0], "Expert number mismatch"
|
|
assert w1_q.shape[0] == w2_q.shape[0], "Weights expert number mismatch"
|
|
assert w1_q.shape[0] == w1_scale.shape[0], "w1 scales expert number mismatch"
|
|
assert w1_q.shape[0] == w2_scale.shape[0], "w2 scales expert number mismatch"
|
|
assert a.dtype in [torch.half, torch.bfloat16], "Invalid output dtype"
|
|
|
|
if is_cuda:
|
|
from sglang.kernels.ops.quantization.fp8_kernel import (
|
|
sglang_per_token_group_quant_fp8,
|
|
)
|
|
es_up, es_down = enable_es
|
|
out_dtype = a.dtype
|
|
num_experts = w1_q.size(0)
|
|
m = a.size(0)
|
|
k = w1_q.size(1)
|
|
n = w2_q.size(1)
|
|
|
|
topk = topk_ids.size(1)
|
|
device = a.device
|
|
|
|
a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
|
|
c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
|
|
|
|
if use_mxfp8:
|
|
assert es_up and es_down, "MXFP8 requires expert-specialization for both GEMMs"
|
|
assert is_sm100_supported(), "MXFP8 requires SM100"
|
|
assert k % 32 == 0, "MXFP8 requires hidden size to be divisible by 32"
|
|
assert n % 32 == 0, "MXFP8 requires intermediate size to be divisible by 32"
|
|
assert w1_scale.dtype == torch.uint8, "MXFP8 w1_scale must be uint8"
|
|
assert w2_scale.dtype == torch.uint8, "MXFP8 w2_scale must be uint8"
|
|
expected_w1_scale_shape = (
|
|
num_experts,
|
|
w1_q.shape[1] // 32,
|
|
w1_q.shape[2],
|
|
)
|
|
expected_w2_scale_shape = (
|
|
num_experts,
|
|
w2_q.shape[1] // 32,
|
|
w2_q.shape[2],
|
|
)
|
|
assert (
|
|
w1_scale.shape == expected_w1_scale_shape
|
|
), f"MXFP8 w1_scale must be {expected_w1_scale_shape}, got {w1_scale.shape}"
|
|
assert (
|
|
w2_scale.shape == expected_w2_scale_shape
|
|
), f"MXFP8 w2_scale must be {expected_w2_scale_shape}, got {w2_scale.shape}"
|
|
|
|
mxfp8_blockscale_align = 128
|
|
total_tokens = m * topk
|
|
nonzero_experts = min(num_experts, total_tokens)
|
|
max_total = total_tokens + (mxfp8_blockscale_align - 1) * nonzero_experts
|
|
max_blockscale = (
|
|
(max_total + mxfp8_blockscale_align - 1) // mxfp8_blockscale_align
|
|
) * mxfp8_blockscale_align
|
|
|
|
blockscale_offsets = None
|
|
if use_mxfp8 and (es_up or es_down):
|
|
blockscale_offsets = torch.empty(
|
|
(num_experts + 1,), dtype=torch.int32, device=device
|
|
)
|
|
|
|
prepare_moe_input(
|
|
topk_ids,
|
|
expert_offsets,
|
|
problem_sizes1,
|
|
problem_sizes2,
|
|
a_map,
|
|
c_map,
|
|
num_experts,
|
|
n,
|
|
k,
|
|
blockscale_offsets,
|
|
)
|
|
|
|
if use_mxfp8 and es_up:
|
|
rep_a = shuffle_rows(a, a_map, (m * topk, k))
|
|
rep_a_q = torch.empty_like(rep_a, dtype=torch.float8_e4m3fn)
|
|
rep_a1_scales = torch.empty(
|
|
(max_blockscale, k // 32), dtype=torch.uint8, device=device
|
|
)
|
|
es_sm100_mxfp8_blockscaled_grouped_quant(
|
|
rep_a,
|
|
problem_sizes1,
|
|
expert_offsets[:-1],
|
|
blockscale_offsets[:-1],
|
|
rep_a_q,
|
|
rep_a1_scales,
|
|
)
|
|
else:
|
|
a_q, a1_scale = sglang_per_token_group_quant_fp8(a, 128)
|
|
# One gather for both: the scale rows are 1/32 of the value rows, so
|
|
# walking the map a second time for them was almost pure launch latency.
|
|
rep_a_q, rep_a1_scales = shuffle_rows_with_scales(
|
|
a_q, a1_scale, a_map, m * topk
|
|
)
|
|
|
|
c1 = torch.empty((m * topk, n * 2), device=device, dtype=out_dtype)
|
|
c2 = torch.empty((m * topk, k), device=device, dtype=out_dtype)
|
|
|
|
a_sf_layout = torch.empty((num_experts, 5), device=device, dtype=torch.int)
|
|
w_sf_layout = torch.empty((num_experts, 5), device=device, dtype=torch.int)
|
|
|
|
if is_sm90_supported() and es_up:
|
|
es_fp8_blockwise_scaled_grouped_mm(
|
|
c1,
|
|
rep_a_q,
|
|
w1_q,
|
|
rep_a1_scales,
|
|
w1_scale,
|
|
a1_strides,
|
|
a1_strides,
|
|
c1_strides,
|
|
problem_sizes1,
|
|
expert_offsets[:-1],
|
|
workspace,
|
|
)
|
|
elif use_mxfp8 and es_up:
|
|
es_sm100_mxfp8_blockscaled_grouped_mm(
|
|
c1,
|
|
rep_a_q,
|
|
w1_q,
|
|
rep_a1_scales,
|
|
w1_scale,
|
|
problem_sizes1,
|
|
expert_offsets[:-1],
|
|
blockscale_offsets[:-1],
|
|
)
|
|
else:
|
|
fp8_blockwise_scaled_grouped_mm(
|
|
c1,
|
|
a_ptrs,
|
|
b_ptrs,
|
|
out_ptrs,
|
|
a_scales_ptrs,
|
|
b_scales_ptrs,
|
|
rep_a_q,
|
|
w1_q,
|
|
rep_a1_scales,
|
|
w1_scale,
|
|
a1_strides,
|
|
a1_strides,
|
|
c1_strides,
|
|
a_sf_layout,
|
|
w_sf_layout,
|
|
problem_sizes1,
|
|
expert_offsets[:-1],
|
|
workspace,
|
|
)
|
|
|
|
intermediate = torch.empty((m * topk, n), device=device, dtype=out_dtype)
|
|
silu_and_mul(c1, intermediate)
|
|
|
|
if use_mxfp8 and es_down:
|
|
intemediate_q = torch.empty_like(intermediate, dtype=torch.float8_e4m3fn)
|
|
a2_scale = torch.empty(
|
|
(max_blockscale, n // 32), dtype=torch.uint8, device=device
|
|
)
|
|
es_sm100_mxfp8_blockscaled_grouped_quant(
|
|
intermediate,
|
|
problem_sizes2,
|
|
expert_offsets[:-1],
|
|
blockscale_offsets[:-1],
|
|
intemediate_q,
|
|
a2_scale,
|
|
)
|
|
else:
|
|
intemediate_q, a2_scale = sglang_per_token_group_quant_fp8(intermediate, 128)
|
|
|
|
if is_sm90_supported() and es_down:
|
|
es_fp8_blockwise_scaled_grouped_mm(
|
|
c2,
|
|
intemediate_q,
|
|
w2_q,
|
|
a2_scale,
|
|
w2_scale,
|
|
a2_strides,
|
|
a2_strides,
|
|
c2_strides,
|
|
problem_sizes2,
|
|
expert_offsets[:-1],
|
|
workspace,
|
|
)
|
|
elif use_mxfp8 and es_down:
|
|
es_sm100_mxfp8_blockscaled_grouped_mm(
|
|
c2,
|
|
intemediate_q,
|
|
w2_q,
|
|
a2_scale,
|
|
w2_scale,
|
|
problem_sizes2,
|
|
expert_offsets[:-1],
|
|
blockscale_offsets[:-1],
|
|
)
|
|
else:
|
|
fp8_blockwise_scaled_grouped_mm(
|
|
c2,
|
|
a_ptrs,
|
|
b_ptrs,
|
|
out_ptrs,
|
|
a_scales_ptrs,
|
|
b_scales_ptrs,
|
|
intemediate_q,
|
|
w2_q,
|
|
a2_scale,
|
|
w2_scale,
|
|
a2_strides,
|
|
a2_strides,
|
|
c2_strides,
|
|
a_sf_layout,
|
|
w_sf_layout,
|
|
problem_sizes2,
|
|
expert_offsets[:-1],
|
|
workspace,
|
|
)
|
|
|
|
if output is None:
|
|
output = torch.empty((m, k), device=device, dtype=out_dtype)
|
|
|
|
apply_shuffle_mul_sum(c2, output, c_map, topk_weights.to(out_dtype))
|
|
return output
|