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

chore: update vlm moe config and tune scripts (#30866)

Browse Source
This commit is contained in:
Mick
2026-07-12 08:35:59 +08:00
committed by GitHub
parent 14bef7cd11
commit a358abd651
9 changed files with 389 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py
+9 -1
View File
@@ -140,13 +140,14 @@ class FlashInferFusedAllReduceParams:
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):
@@ -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"
)
@@ -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
benchmark/kernels/fused_moe_triton/common_utils.py
+1
View File
@@ -95,6 +95,7 @@ def get_model_config(
"DeepseekV4ForCausalLM",
"Glm4MoeForCausalLM",
"GlmMoeDsaForCausalLM",
"KimiVLForConditionalGeneration",
"MistralLarge3ForCausalLM",
]:
E = (config.n_routed_experts // ep_size) + (
benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py
+29 -6
View File
@@ -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
@@ -242,7 +243,6 @@ def benchmark_config(
@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)
@@ -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"])
@@ -419,9 +421,19 @@ def main(args: argparse.Namespace):
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"
)
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
@@ -522,7 +534,18 @@ if __name__ == "__main__":
)
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()
python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=128,N=768,device_name=NVIDIA_H200.json
+146
View File
@@ -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
}
}
python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=16,N=1408,device_name=NVIDIA_H100_80GB_HBM3.json
+34
View File
@@ -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
}
}
python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=32,N=768,device_name=NVIDIA_H100_80GB_HBM3.json
+30
View File
@@ -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
}
}
python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_config.py
+20 -3
View File
@@ -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,
test/registered/unit/layers/moe/test_fused_moe_common_utils.py
+61
View File
@@ -0,0 +1,61 @@
import importlib.util
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
import torch
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=3, suite="base-a-test-cpu")
def _load_common_utils():
source = (
Path(__file__).resolve().parents[5]
/ "benchmark/kernels/fused_moe_triton/common_utils.py"
)
spec = importlib.util.spec_from_file_location("fused_moe_common_utils", source)
module = importlib.util.module_from_spec(spec)
assert spec.loader is not None
spec.loader.exec_module(module)
return module
def test_get_model_config_supports_kimi_vl():
common_utils = _load_common_utils()
text_config = SimpleNamespace(
hidden_size=2048,
n_routed_experts=64,
num_experts_per_tok=6,
moe_intermediate_size=1408,
torch_dtype=torch.bfloat16,
)
model_config = SimpleNamespace(
architectures=["KimiVLForConditionalGeneration"],
text_config=text_config,
get_text_config=lambda: text_config,
)
with patch.object(common_utils, "get_config", return_value=model_config):
tuned_config = common_utils.get_model_config(
"moonshotai/Kimi-VL-A3B-Instruct", tp_size=4, ep_size=4
)
assert tuned_config == {
"num_experts": 16,
"topk": 6,
"hidden_size": 2048,
"shard_intermediate_size": 2816,
"dtype": torch.bfloat16,
"block_shape": None,
"architecture": "KimiVLForConditionalGeneration",
}
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main([__file__, "-v"]))
test/registered/unit/layers/moe/test_fused_moe_triton_config.py
+59
View File
@@ -0,0 +1,59 @@
import json
from pathlib import Path
from types import SimpleNamespace
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=3, suite="base-a-test-cpu")
from sglang.srt.layers.moe.moe_runner.triton_utils import fused_moe_triton_config
def test_h200_bf16_config_is_available_for_current_triton_runtime():
config_path = (
Path(fused_moe_triton_config.__file__).parent
/ "configs"
/ "triton_3_6_0"
/ "E=128,N=768,device_name=NVIDIA_H200.json"
)
assert config_path.is_file()
assert json.loads(config_path.read_text())["128"]["BLOCK_SIZE_M"] > 0
def test_down_moe_reuses_tuned_up_config_when_separate_config_is_absent(
monkeypatch, tmp_path
):
config_root = tmp_path / "configs" / "triton_3_6_0"
config_root.mkdir(parents=True)
tuned_config = {"128": {"BLOCK_SIZE_M": 64}}
(config_root / "up.json").write_text(json.dumps(tuned_config))
monkeypatch.setenv("SGLANG_MOE_CONFIG_DIR", str(tmp_path))
monkeypatch.setattr(fused_moe_triton_config.triton, "__version__", "3.6.0")
monkeypatch.setattr(
fused_moe_triton_config,
"get_server_args",
lambda: SimpleNamespace(enable_deterministic_inference=False),
)
monkeypatch.setattr(
fused_moe_triton_config,
"get_config_file_name",
lambda *args, down_moe=False, **kwargs: "down.json" if down_moe else "up.json",
)
fused_moe_triton_config.get_moe_configs.cache_clear()
try:
assert fused_moe_triton_config.get_moe_configs(
32, 768, None, down_moe=True
) == {128: {"BLOCK_SIZE_M": 64}}
finally:
fused_moe_triton_config.get_moe_configs.cache_clear()
if __name__ == "__main__":
import sys
import pytest
sys.exit(pytest.main([__file__, "-v"]))