Skip to content

Commit d02f3ba

Browse files
xuanyuanminzhengxuanyuanminzheng
authored
[Feature] Add TritonMoEMethod for BF16 MoE inference (#7815)
1 parent 478c9fa commit d02f3ba

5 files changed

Lines changed: 1151 additions & 55 deletions

File tree

‎fastdeploy/model_executor/layers/moe/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
CutlassW4AFP8MoEMethod,
1818
CutlassWeightOnlyMoEMethod,
1919
)
20-
from .fused_moe_triton_backend import TritonWeightOnlyMoEMethod
20+
from .fused_moe_triton_backend import TritonMoEMethod, TritonWeightOnlyMoEMethod
2121
from .moe import FusedMoE
2222

2323
__all__ = [
@@ -26,4 +26,5 @@
2626
CutlassW4AFP8MoEMethod,
2727
FusedMoE,
2828
TritonWeightOnlyMoEMethod,
29+
TritonMoEMethod,
2930
]

‎fastdeploy/model_executor/layers/moe/fused_moe_triton_backend.py‎

Lines changed: 255 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,17 @@
2020
from paddle import nn
2121

2222
import fastdeploy
23+
from fastdeploy.model_executor.layers.moe.moe import get_moe_scores
24+
from fastdeploy.model_executor.layers.moe.triton_moe_kernels import (
25+
fused_moe_kernel_bf16,
26+
fused_moe_kernel_paddle,
27+
)
28+
from fastdeploy.model_executor.layers.quantization.fp8_utils import (
29+
fused_stack_transpose_quant,
30+
quant_weight_ue8m0,
31+
transform_scale_ue8m0,
32+
)
33+
from fastdeploy.model_executor.layers.quantization.ops import scaled_fp8_quant
2334
from fastdeploy.model_executor.layers.utils import get_tensor
2435
from fastdeploy.model_executor.utils import (
2536
TensorTracker,
@@ -31,21 +42,15 @@
3142
from fastdeploy.platforms import current_platform
3243
from fastdeploy.utils import ceil_div, register_custom_python_op
3344

34-
from ..quantization.quant_base import QuantMethodBase
35-
3645
try:
37-
from fastdeploy.model_executor.ops.gpu import tritonmoe_preprocess_func
46+
import triton.language as tl
3847

39-
from .triton_moe_kernels import fused_moe_kernel_paddle
48+
from fastdeploy.model_executor.ops.gpu import tritonmoe_preprocess_func
4049
except ImportError:
4150
pass
42-
from fastdeploy.model_executor.layers.moe.moe import get_moe_scores
43-
from fastdeploy.model_executor.layers.quantization.fp8_utils import (
44-
fused_stack_transpose_quant,
45-
quant_weight_ue8m0,
46-
transform_scale_ue8m0,
47-
)
48-
from fastdeploy.model_executor.layers.quantization.ops import scaled_fp8_quant
51+
52+
from ..quantization.quant_base import QuantMethodBase
53+
from .fused_moe_backend_base import UnquantizedFusedMoEMethod
4954

5055

5156
class TritonWeightOnlyMoEMethod(QuantMethodBase):
@@ -778,8 +783,8 @@ def apply(
778783
stride_am=x_q.strides[0],
779784
stride_ak=x_q.strides[1],
780785
stride_be=layer.up_gate_proj_weight.strides[0],
781-
stride_bk=layer.up_gate_proj_weight.strides[2],
782-
stride_bn=layer.up_gate_proj_weight.strides[1],
786+
stride_bk=layer.up_gate_proj_weight.strides[1],
787+
stride_bn=layer.up_gate_proj_weight.strides[2],
783788
stride_cm=up_gate_proj_out.strides[0],
784789
stride_cn=up_gate_proj_out.strides[1],
785790
#
@@ -1846,3 +1851,240 @@ def apply(
18461851
self.quant_config,
18471852
topk_ids_hookfunc,
18481853
)
1854+
1855+
1856+
class TritonMoEMethod(UnquantizedFusedMoEMethod):
1857+
"""
1858+
Use Triton Group Gemm (BF16 unquantized) to compute Fused MoE.
1859+
1860+
Activated via: export FD_MOE_BACKEND=triton
1861+
Weight layout (CUDA path): [E, K, 2N] for up_gate_proj, [E, N, K] for down_proj.
1862+
This matches UnquantizedFusedMoEMethod.create_weights layout on CUDA.
1863+
"""
1864+
1865+
def __init__(self, quant_config=None):
1866+
super().__init__(quant_config)
1867+
1868+
def process_loaded_weights(self, layer: nn.Layer, state_dict):
1869+
"""Stack individual expert weights into the stacked parameter."""
1870+
up_gate_proj_weights, down_proj_weights, _, _ = layer.extract_moe_ffn_weights(state_dict)
1871+
layer.up_gate_proj_weight.set_value(paddle.stack(up_gate_proj_weights, axis=0))
1872+
layer.down_proj_weight.set_value(paddle.stack(down_proj_weights, axis=0))
1873+
1874+
def _get_default_config(self, M: int, E: int) -> dict:
1875+
"""
1876+
Heuristic tile config for BF16 MoE, ported verbatim from vLLM's
1877+
`get_default_config` (bf16/fp16 non-block_shape branch).
1878+
See vllm/model_executor/layers/fused_moe/fused_moe.py:1273-1319.
1879+
1880+
M: number of tokens (A.size(0) in vLLM), i.e. pre-expansion token count.
1881+
E: number of (local) experts.
1882+
"""
1883+
1884+
# Tile sizes scale with batch: small batches are memory-bound
1885+
# (favor tall-K tiles), large batches are compute-bound (favor
1886+
# large M/N tiles with more warps).
1887+
if M <= 32:
1888+
block_m = 16
1889+
elif M <= 96:
1890+
block_m = 32
1891+
elif M <= 512:
1892+
block_m = 64
1893+
else:
1894+
block_m = 128
1895+
1896+
block_n = 64 if M <= 64 else 128
1897+
1898+
block_k = 64
1899+
1900+
# Grouping adjacent M-blocks lets them share weight tiles in L2.
1901+
# Only helps when there are enough M-blocks per expert to group;
1902+
# with many experts each one sees few tokens so grouping is useless.
1903+
tokens_per_expert = M // max(E, 1)
1904+
group_m = 16 if tokens_per_expert > 128 else 1
1905+
1906+
# Large batches have enough blocks to saturate the GPU, so we
1907+
# use more warps per block to increase arithmetic intensity.
1908+
num_warps = 4 if M <= 128 else 8
1909+
1910+
num_stages = 4 if M <= 32 else 3
1911+
1912+
return {
1913+
"BLOCK_SIZE_M": block_m,
1914+
"BLOCK_SIZE_N": block_n,
1915+
"BLOCK_SIZE_K": block_k,
1916+
"GROUP_SIZE_M": group_m,
1917+
"num_warps": num_warps,
1918+
"num_stages": num_stages,
1919+
}
1920+
1921+
def apply_tp(
1922+
self,
1923+
layer: nn.Layer,
1924+
x: paddle.Tensor,
1925+
gate: nn.Layer,
1926+
topk_ids_hookfunc: Callable = None,
1927+
fc1_latent_proj: nn.Layer = None,
1928+
fc2_latent_proj: nn.Layer = None,
1929+
) -> paddle.Tensor:
1930+
"""
1931+
BF16 Triton Fused MoE forward.
1932+
1933+
Pipeline:
1934+
1. Gate + topk routing
1935+
2. tritonmoe_preprocess -> sorted_token_ids, expert_ids, num_tokens_post_padded
1936+
3. fused_moe_kernel_bf16 GEMM1: [tokens*topk, K] x [E, K, 2N] -> [tokens*topk, 2N]
1937+
4. SwiGLU activation
1938+
5. fused_moe_kernel_bf16 GEMM2: [tokens*topk, N] x [E, N, K] -> [tokens*topk, K]
1939+
(with MUL_ROUTED_WEIGHT=True to fuse router weight multiplication)
1940+
6. Reshape + sum over topk dim
1941+
"""
1942+
token_num = x.shape[0]
1943+
if token_num == 0:
1944+
return paddle.zeros([token_num, layer.hidden_size], dtype=x.dtype)
1945+
1946+
top_k = layer.top_k
1947+
num_local_experts = layer.num_local_experts
1948+
moe_intermediate_size = layer.moe_intermediate_size
1949+
hidden_size = layer.hidden_size
1950+
1951+
# --- 1. Routing ---
1952+
gate_out = gate(x)
1953+
1954+
if layer.topk_method == "noaux_tc":
1955+
use_fused = not fastdeploy.envs.FD_ENABLE_RL and current_platform.is_cuda()
1956+
if not use_fused:
1957+
gate_out = gate_out.cast("float32")
1958+
1959+
_, topk_weights, topk_ids = get_moe_scores(
1960+
gate_out,
1961+
layer.n_group,
1962+
layer.topk_group,
1963+
top_k,
1964+
layer.routed_scaling_factor,
1965+
layer.gate_correction_bias,
1966+
getattr(layer, "renormalize", True),
1967+
use_fused_cast=use_fused,
1968+
topk_reduce_func=getattr(layer, "topk_reduce_func", None),
1969+
)
1970+
else:
1971+
gate_out = gate_out.cast("float32")
1972+
topk_ids, topk_weights = fastdeploy.model_executor.ops.gpu.moe_topk_select(
1973+
gate_out,
1974+
layer.gate_correction_bias,
1975+
top_k,
1976+
True, # apply_norm_weight
1977+
False,
1978+
)
1979+
1980+
if topk_ids_hookfunc is not None:
1981+
topk_ids_hookfunc(topk_ids=topk_ids)
1982+
1983+
# --- 2. Preprocess: sort tokens by expert assignment ---
1984+
num_token_expert_pairs = token_num * top_k
1985+
# vLLM convention: pass num_tokens (pre-expansion), NOT tokens*top_k.
1986+
cfg = self._get_default_config(token_num, num_local_experts)
1987+
1988+
sorted_token_ids, expert_ids, num_tokens_post_padded = tritonmoe_preprocess_func(
1989+
topk_ids, num_local_experts, cfg["BLOCK_SIZE_M"]
1990+
)
1991+
max_possible_num_post_padded = sorted_token_ids.shape[0]
1992+
1993+
# --- 3. GEMM1: hidden -> up_gate (BF16 x BF16 -> BF16) ---
1994+
# up_gate_proj_weight layout: [E, hidden_size, inter*2] => stride_be, stride_bk, stride_bn
1995+
up_gate_proj_out = paddle.empty(
1996+
[num_token_expert_pairs, moe_intermediate_size * 2],
1997+
dtype=x.dtype,
1998+
)
1999+
grid1 = (
2000+
ceil_div(max_possible_num_post_padded, cfg["BLOCK_SIZE_M"])
2001+
* ceil_div(moe_intermediate_size * 2, cfg["BLOCK_SIZE_N"]),
2002+
)
2003+
fused_moe_kernel_bf16[grid1](
2004+
x,
2005+
layer.up_gate_proj_weight,
2006+
up_gate_proj_out,
2007+
None, # topk_weights_ptr (no weight mul on GEMM1)
2008+
sorted_token_ids,
2009+
expert_ids,
2010+
num_tokens_post_padded,
2011+
N=moe_intermediate_size * 2,
2012+
K=hidden_size,
2013+
EM=max_possible_num_post_padded,
2014+
num_valid_tokens=num_token_expert_pairs,
2015+
stride_am=x.strides[0],
2016+
stride_ak=x.strides[1],
2017+
stride_be=layer.up_gate_proj_weight.strides[0],
2018+
stride_bk=layer.up_gate_proj_weight.strides[1],
2019+
stride_bn=layer.up_gate_proj_weight.strides[2],
2020+
stride_cm=up_gate_proj_out.strides[0],
2021+
stride_cn=up_gate_proj_out.strides[1],
2022+
BLOCK_SIZE_M=cfg["BLOCK_SIZE_M"],
2023+
BLOCK_SIZE_N=cfg["BLOCK_SIZE_N"],
2024+
BLOCK_SIZE_K=cfg["BLOCK_SIZE_K"],
2025+
GROUP_SIZE_M=cfg["GROUP_SIZE_M"],
2026+
MUL_ROUTED_WEIGHT=False,
2027+
top_k=top_k,
2028+
compute_type=tl.bfloat16,
2029+
even_Ks=(hidden_size % cfg["BLOCK_SIZE_K"] == 0),
2030+
num_warps=cfg["num_warps"],
2031+
num_stages=cfg["num_stages"],
2032+
)
2033+
2034+
# --- 4. SwiGLU activation ---
2035+
down_proj_input = paddle.incubate.nn.functional.swiglu(up_gate_proj_out)
2036+
2037+
# --- 5. GEMM2: inter -> hidden, fuse router weight multiplication ---
2038+
# down_proj_weight layout: [E, moe_intermediate_size, hidden_size] => stride_be, stride_bk, stride_bn
2039+
down_proj_out = paddle.empty(
2040+
(num_token_expert_pairs, hidden_size),
2041+
dtype=x.dtype,
2042+
)
2043+
grid2 = (
2044+
ceil_div(max_possible_num_post_padded, cfg["BLOCK_SIZE_M"]) * ceil_div(hidden_size, cfg["BLOCK_SIZE_N"]),
2045+
)
2046+
fused_moe_kernel_bf16[grid2](
2047+
down_proj_input,
2048+
layer.down_proj_weight,
2049+
down_proj_out,
2050+
topk_weights,
2051+
sorted_token_ids,
2052+
expert_ids,
2053+
num_tokens_post_padded,
2054+
N=hidden_size,
2055+
K=moe_intermediate_size,
2056+
EM=max_possible_num_post_padded,
2057+
num_valid_tokens=num_token_expert_pairs,
2058+
stride_am=down_proj_input.strides[0],
2059+
stride_ak=down_proj_input.strides[1],
2060+
stride_be=layer.down_proj_weight.strides[0],
2061+
stride_bk=layer.down_proj_weight.strides[1],
2062+
stride_bn=layer.down_proj_weight.strides[2],
2063+
stride_cm=down_proj_out.strides[0],
2064+
stride_cn=down_proj_out.strides[1],
2065+
BLOCK_SIZE_M=cfg["BLOCK_SIZE_M"],
2066+
BLOCK_SIZE_N=cfg["BLOCK_SIZE_N"],
2067+
BLOCK_SIZE_K=cfg["BLOCK_SIZE_K"],
2068+
GROUP_SIZE_M=cfg["GROUP_SIZE_M"],
2069+
MUL_ROUTED_WEIGHT=True,
2070+
top_k=1,
2071+
compute_type=tl.bfloat16,
2072+
even_Ks=(moe_intermediate_size % cfg["BLOCK_SIZE_K"] == 0),
2073+
num_warps=cfg["num_warps"],
2074+
num_stages=cfg["num_stages"],
2075+
)
2076+
2077+
# --- 6. Reduce over topk ---
2078+
down_proj_out.reshape_([token_num, top_k, hidden_size])
2079+
out = down_proj_out.sum(axis=1)
2080+
return out
2081+
2082+
def apply_ep_prefill(
2083+
self, layer, x, gate, topk_ids_hookfunc=None, shared_experts=None, fc1_latent_proj=None, fc2_latent_proj=None
2084+
):
2085+
raise NotImplementedError("TritonMoEMethod does not support EP prefill yet.")
2086+
2087+
def apply_ep_decode(
2088+
self, layer, x, gate, topk_ids_hookfunc=None, shared_experts=None, fc1_latent_proj=None, fc2_latent_proj=None
2089+
):
2090+
raise NotImplementedError("TritonMoEMethod does not support EP decode yet.")

‎fastdeploy/model_executor/layers/moe/moe.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,11 @@ def get_moe_method(layer=None):
5353
"""
5454

5555
if current_platform.is_cuda():
56+
moe_backend = envs.FD_MOE_BACKEND.lower()
57+
if moe_backend == "triton":
58+
from .fused_moe_triton_backend import TritonMoEMethod
59+
60+
return TritonMoEMethod(None)
5661
from .fused_moe_cutlass_backend import CutlassMoEMethod
5762

5863
return CutlassMoEMethod(None)

0 commit comments

Comments
 (0)