Skip to content

Tune VLM MoE paths #30866

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mickqian merged 3 commits into from
Jul 12, 2026
Merged

Tune VLM MoE paths #30866

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
13b6b27
tune VLM MoE paths
mickqian Jul 11, 2026
88248d5
fix standalone MoE config test entrypoint
mickqian Jul 11, 2026
239c35c
fix standalone MoE common-utils test entrypoint
mickqian Jul 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -140,13 +140,14 @@ def __init__(
world_size: int,
use_fp32_lamport: bool = False,
max_token_num: int = 1024,
fp32_acc: bool = True,
):
self.rank = rank
self.world_size = world_size
self.use_fp32_lamport = use_fp32_lamport
self.trigger_completion_at_end = True
self.launch_with_pdl = True
self.fp32_acc = True
self.fp32_acc = fp32_acc
self.max_token_num = max_token_num

def get_trtllm_fused_allreduce_kwargs(self):
Expand Down Expand Up @@ -1134,6 +1135,12 @@ def main():
action="store_true",
help="Disable oneshot mode for FlashInfer operations",
)
parser.add_argument(
"--fp32-acc",
action=argparse.BooleanOptionalAction,
default=True,
help="Use FP32 accumulation in FlashInfer fused all-reduce (default: enabled).",
)
parser.add_argument(
"--warmup", type=int, default=5, help="Number of warmup iterations"
)
Expand Down Expand Up @@ -1238,6 +1245,7 @@ def main():
rank=rank,
world_size=world_size,
max_token_num=max_num_token,
fp32_acc=args.fp32_acc,
)

# Collect all results for markdown export
Expand Down
1 change: 1 addition & 0 deletions benchmark/kernels/fused_moe_triton/common_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,7 @@ def get_model_config(
"DeepseekV4ForCausalLM",
"Glm4MoeForCausalLM",
"GlmMoeDsaForCausalLM",
"KimiVLForConditionalGeneration",
"MistralLarge3ForCausalLM",
]:
E = (config.n_routed_experts // ep_size) + (
Expand Down
35 changes: 29 additions & 6 deletions benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Adapted from https://github.com/vllm-project/vllm/blob/main/benchmarks/kernels/benchmark_moe.py
import argparse
import json
import time
from contextlib import nullcontext
from datetime import datetime
Expand Down Expand Up @@ -242,7 +243,6 @@ def run():

@ray.remote(num_gpus=1)
class BenchmarkWorker:

def __init__(self, seed: int, server_args: ServerArgs) -> None:
torch.set_default_device(get_device())
torch.get_device_module().manual_seed_all(0)
Expand Down Expand Up @@ -398,10 +398,12 @@ def main(args: argparse.Namespace):
use_int4_w4a16 = args.dtype == "int4_w4a16"
per_channel_quant = args.per_channel_quant

if args.batch_size is None:
batch_sizes = get_default_batch_sizes()
else:
if args.batch_sizes is not None:
batch_sizes = args.batch_sizes
elif args.batch_size is not None:
batch_sizes = [args.batch_size]
else:
batch_sizes = get_default_batch_sizes()

ray.init()
num_gpus = int(ray.available_resources()["GPU"])
Expand All @@ -419,9 +421,19 @@ def _distribute(method: str, inputs: List[Any]) -> List[Any]:
return ray.get(outputs)

if args.tune:
search_space = get_configs_compute_bound()
if args.search_space_file:
with open(args.search_space_file) as f:
search_space = json.load(f)
if not isinstance(search_space, list) or not all(
isinstance(config, dict) for config in search_space
):
raise ValueError(
"--search-space-file must contain a JSON list of configs"
)
Comment on lines +424 to +432

@gemini-code-assist gemini-code-assist Bot Jul 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

When loading a custom search space file via --search-space-file, there are two potential issues:

  1. Opening the file without an explicit encoding can lead to UnicodeDecodeError on platforms where the default encoding is not UTF-8 (e.g., Windows).
  2. If the JSON file contains configurations that are missing required Triton keys (such as BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, etc.), it will cause cryptic KeyErrors later during filtering or sorting.

We should open the file with encoding="utf-8" and validate that each configuration dictionary contains all the required keys.

        if args.search_space_file:
            with open(args.search_space_file, encoding="utf-8") as f:
                search_space = json.load(f)
            if not isinstance(search_space, list) or not all(
                isinstance(config, dict) for config in search_space
            ):
                raise ValueError(
                    "--search-space-file must contain a JSON list of configs"
                )
            required_keys = {"BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K", "GROUP_SIZE_M", "num_warps", "num_stages"}
            for i, config in enumerate(search_space):
                if not required_keys.issubset(config.keys()):
                    raise ValueError(
                        f"Config at index {i} in --search-space-file is missing required keys: {required_keys - config.keys()}"
                    )

else:
search_space = get_configs_compute_bound()
if block_shape is not None:
block_n, block_k = block_shape[0], block_shape[1]
block_k = block_shape[1]
search_space = [
config
for config in search_space
Expand Down Expand Up @@ -522,7 +534,18 @@ def _distribute(method: str, inputs: List[Any]) -> List[Any]:
)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--batch-size", type=int, required=False)
parser.add_argument(
"--batch-sizes",
type=int,
nargs="+",
help="Tune or benchmark an explicit set of token counts in parallel.",
)
parser.add_argument("--tune", action="store_true")
parser.add_argument(
"--search-space-file",
type=str,
help="JSON file containing an explicit list of Triton configs to evaluate with --tune.",
)
parser.add_argument("--disable-shared-experts-fusion", action="store_true")
args = parser.parse_args()

Expand Down
146 changes: 146 additions & 0 deletions ...moe/moe_runner/triton_utils/configs/triton_3_6_0/E=128,N=768,device_name=NVIDIA_H200.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,146 @@
{
"1": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 5
},
"2": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"4": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2
},
"8": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"16": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 256,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 2
},
"24": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 4
},
"32": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2
},
"48": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 3
},
"64": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 2
},
"96": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 2
},
"128": {
"BLOCK_SIZE_M": 16,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 2
},
"256": {
"BLOCK_SIZE_M": 32,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 128,
"GROUP_SIZE_M": 16,
"num_warps": 4,
"num_stages": 2
},
"512": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 4,
"num_stages": 3
},
"1024": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 16,
"num_warps": 8,
"num_stages": 4
},
"1536": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 4
},
"2048": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 4
},
"3072": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 4
},
"4096": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 16,
"num_warps": 8,
"num_stages": 4
}
}
34 changes: 34 additions & 0 deletions ...nner/triton_utils/configs/triton_3_6_0/E=16,N=1408,device_name=NVIDIA_H100_80GB_HBM3.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
{
"128": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 3
},
"256": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 128,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 3
},
"512": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 4
},
"1024": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 4
}
}
30 changes: 30 additions & 0 deletions ...unner/triton_utils/configs/triton_3_6_0/E=32,N=768,device_name=NVIDIA_H100_80GB_HBM3.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
{
"128": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 32,
"GROUP_SIZE_M": 8
},
"256": {
"BLOCK_SIZE_M": 64,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 32,
"GROUP_SIZE_M": 8
},
"512": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 64,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 3
},
"1024": {
"BLOCK_SIZE_M": 128,
"BLOCK_SIZE_N": 256,
"BLOCK_SIZE_K": 64,
"GROUP_SIZE_M": 1,
"num_warps": 8,
"num_stages": 4
}
}
23 changes: 20 additions & 3 deletions python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,12 +139,29 @@ def get_moe_configs(
# 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 when down_moe is False
# When down_moe is True, we will try to use the config for down_moe=False
if down_moe:
# A separate down-projection config enables the TMA path, but it is
# optional. Reuse a tuned up-projection config when it is absent so
# the second GEMM does not silently fall back to the heuristic.
up_configs = get_moe_configs(
E,
N,
dtype,
block_n,
block_k,
per_channel_quant=per_channel_quant,
down_moe=False,
)
if up_configs is not None:
logger.warning(
"Down MoE config file not found at %s; reusing the tuned "
"up-projection config without TMA. Performance might be sub-optimal.",
config_file_path,
)
return up_configs
logger.warning(
(
"Using MoE kernel config with down_moe=False. Performance might be sub-optimal! "
"Using default MoE kernel config. Performance might be sub-optimal! "
"Config file not found at %s, you can create them with https://github.com/sgl-project/sglang/tree/main/benchmark/kernels/fused_moe_triton"
),
config_file_path,
Expand Down
Loading
Loading