diff --git a/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py b/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py index 7050897c0..c27f8c323 100644 --- a/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py +++ b/benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py @@ -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 diff --git a/benchmark/kernels/fused_moe_triton/common_utils.py b/benchmark/kernels/fused_moe_triton/common_utils.py index ea669a1fa..865135008 100644 --- a/benchmark/kernels/fused_moe_triton/common_utils.py +++ b/benchmark/kernels/fused_moe_triton/common_utils.py @@ -95,6 +95,7 @@ def get_model_config( "DeepseekV4ForCausalLM", "Glm4MoeForCausalLM", "GlmMoeDsaForCausalLM", + "KimiVLForConditionalGeneration", "MistralLarge3ForCausalLM", ]: E = (config.n_routed_experts // ep_size) + ( diff --git a/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py b/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py index 536cc6913..8cfc19dbd 100644 --- a/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py +++ b/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py @@ -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() diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=128,N=768,device_name=NVIDIA_H200.json b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=128,N=768,device_name=NVIDIA_H200.json new file mode 100644 index 000000000..beaac7f64 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=128,N=768,device_name=NVIDIA_H200.json @@ -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 + } +} diff --git a/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 b/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 new file mode 100644 index 000000000..338011989 --- /dev/null +++ b/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 @@ -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 + } +} diff --git a/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 b/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 new file mode 100644 index 000000000..02e4a5457 --- /dev/null +++ b/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 @@ -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 + } +} diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_config.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_config.py index 2f1735cf4..210247f86 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_config.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe_triton_config.py @@ -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, diff --git a/test/registered/unit/layers/moe/test_fused_moe_common_utils.py b/test/registered/unit/layers/moe/test_fused_moe_common_utils.py new file mode 100644 index 000000000..bf28872d4 --- /dev/null +++ b/test/registered/unit/layers/moe/test_fused_moe_common_utils.py @@ -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"])) diff --git a/test/registered/unit/layers/moe/test_fused_moe_triton_config.py b/test/registered/unit/layers/moe/test_fused_moe_triton_config.py new file mode 100644 index 000000000..1b744b6ff --- /dev/null +++ b/test/registered/unit/layers/moe/test_fused_moe_triton_config.py @@ -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"]))