minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[kernel] add triton moe TMA up support (#33559)

Browse Source 
Co-authored-by: undefined <zhouchen.arrebol@jd.com>
Co-authored-by: xq25478 <xq25478@qq.com>
Co-authored-by: xieminghe.simon <xieminghe.simon@jd.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
This commit is contained in:
xieminghe1
2026-08-13 13:29:32 +08:00
committed by GitHub
co-authored by undefined xq25478 xieminghe.simon Xiaoyu Zhang
parent 69bf601e3c
commit ef7208d41d
5 changed files with 669 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/moe_runner/triton.py
+3
View File
@@ -138,6 +138,7 @@ class TritonRunnerCore(MoeRunnerCore):
running_state["config"],
running_state.get("down_config"),
running_state.get("down_moe_use_tma", False),
running_state.get("up_moe_use_tma", False),
b1=quant_info.b13,
b2=quant_info.b2,
use_fp8_w8a8=quant_info.use_fp8_w8a8,
@@ -280,6 +281,7 @@ def pre_permute_standard_to_triton(
config,
down_config,
down_moe_use_tma,
up_moe_use_tma,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
@@ -299,6 +301,7 @@ def pre_permute_standard_to_triton(
running_state["config"] = config
running_state["down_config"] = down_config
running_state["down_moe_use_tma"] = down_moe_use_tma
running_state["up_moe_use_tma"] = up_moe_use_tma
return TritonRunnerInput(
hidden_states=hidden_states,
python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=512,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128].json
+164
View File
@@ -0,0 +1,164 @@
{
"1": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 5,
"USE_TMA": false
},
"2": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 5,
"USE_TMA": false
},
"4": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 5,
"USE_TMA": false
},
"8": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 5,
"USE_TMA": false
},
"16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 5,
"USE_TMA": false
},
"24": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 5,
"USE_TMA": true
},
"32": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 5,
"USE_TMA": false
},
"48": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 5,
"USE_TMA": true
},
"64": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 5,
"USE_TMA": true
},
"96": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"128": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"256": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 4,
"USE_TMA": true
},
"512": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"1024": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"1536": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"2048": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 4,
"USE_TMA": false
},
"3072": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": false
},
"4096": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 4,
"USE_TMA": false
}
}
python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=512,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128, 128]_down.json
+164
View File
@@ -0,0 +1,164 @@
{
"1": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"2": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"4": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"8": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"24": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"32": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"48": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"64": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"96": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"128": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"256": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 64,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"512": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3,
"USE_TMA": true
},
"1024": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"1536": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"2048": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"3072": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
},
"4096": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 32,
"num_warps": 4,
"num_stages": 2,
"USE_TMA": true
}
}
python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py
+29 -6
View File
@@ -7,6 +7,7 @@
from __future__ import annotations
import functools
import logging
from typing import TYPE_CHECKING, Any, Dict, List, Optional
import torch
@@ -93,6 +94,8 @@ if not _is_cuda and not _is_hip and not _is_xpu:
padding_size = get_moe_padding_size(_use_aiter)
logger = logging.getLogger(__name__)
def _use_moe_sum_reduce_torch_compile(num_tokens: int) -> bool:
return num_tokens <= 32 and not is_batch_invariant_mode_enabled()
@@ -363,7 +366,7 @@ def swiglu_no_interleaved_with_alpha_and_limit(x, gemm1_alpha, gemm1_limit):
@functools.lru_cache()
def _down_moe_use_tma():
def _moe_support_tma():
return support_tensor_descriptor()
@@ -409,11 +412,26 @@ def _prepare_fused_moe_run(
per_channel_quant=per_channel_quant,
return_down_config=True,
)
down_moe_use_tma = (
_down_moe_use_tma()
and down_config is not None
and down_config.pop("USE_TMA", False)
)
# Copy config to avoid mutating the lru_cached dict returned by
# get_moe_configs; we pop USE_TMA below.
config = dict(config)
# Up-projection TMA is opt-in: only enabled when the up config file
# explicitly carries "USE_TMA": true (produced by tuning). By default the
# existing up config files do not contain this key, so existing users are
# unaffected unless they re-tune with the updated script.
up_tma_requested = config.pop("USE_TMA", False)
up_moe_use_tma = _moe_support_tma() and up_tma_requested
if up_moe_use_tma:
logger.warning_once(
"Up MoE TMA is enabled (USE_TMA=true in the up-projection config). "
"This requires a config produced by the updated tuning script. "
)
down_tma_requested = down_config is not None and down_config.pop("USE_TMA", False)
down_moe_use_tma = _moe_support_tma() and down_tma_requested
if down_moe_use_tma:
logger.warning_once(
"Down MoE TMA is enabled (USE_TMA=true in the down-projection config)."
)
sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
topk_ids, config["BLOCK_SIZE_M"], E
@@ -423,6 +441,7 @@ def _prepare_fused_moe_run(
config,
down_config,
down_moe_use_tma,
up_moe_use_tma,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
@@ -441,6 +460,7 @@ def _fused_moe_kernel_sequence(
config: Dict[str, Any],
down_config: Optional[Dict[str, Any]],
down_moe_use_tma: bool,
up_moe_use_tma: bool,
*,
b1: Optional[torch.Tensor],
b2: Optional[torch.Tensor],
@@ -568,6 +588,7 @@ def _fused_moe_kernel_sequence(
per_channel_quant=per_channel_quant,
block_shape=block_shape,
c_sorted=down_moe_use_tma,
b_use_tma=up_moe_use_tma,
filter_expert=filter_expert,
)
@@ -927,6 +948,7 @@ def fused_experts_impl(
config,
down_config,
down_moe_use_tma,
up_moe_use_tma,
sorted_token_ids,
expert_ids,
num_tokens_post_padded,
@@ -955,6 +977,7 @@ def fused_experts_impl(
config,
down_config,
down_moe_use_tma,
up_moe_use_tma,
b1=b1,
b2=b2,
use_fp8_w8a8=use_fp8_w8a8,