add fbgemm moe grouped gemm kernel benchmark

[BBuf]

H100 Benchmark FBGEMM GroupedGEMM Results

When running benchmarks with triton==3.2.0, the following warning appears: we can't use warp-specialized features, but persistent kernels and TMA load/store remain available.

/home/ubuntu/bbuf/sglang/benchmark/kernels/fbgemm/fbgemm_grouped_gemm.py:1104: UserWarning: Warp specialization is disabled as the Triton build in current environment doesn't have such support. Please build from https://github.com/facebookexperimental/triton/tree/ws-3.2.x to enable it for best performance on Nvidia's SM90 GPUs.

Qwen2-57B-A14B-Instruct BF16 W8A8 TP4

python3 benchmark/kernels/fbgemm/benchmark_fbgemm_grouped_gemm.py --model Qwen/Qwen2-57B-A14B-Instruct --tp-size 4

grouped-gemm-performance:
    batch_size  FBGEMM Grouped GEMM BF16  SGLang Grouped GEMM BF16
0          1.0                  0.032352                  0.022272
1          2.0                  0.032096                  0.022080
2          4.0                  0.032640                  0.021984
3          8.0                  0.031840                  0.021472
4         16.0                  0.030832                  0.021536
5         32.0                  0.032192                  0.021632
6         64.0                  0.393504                  0.595008
7        128.0                  0.393872                  0.598048
8        256.0                  0.394848                  0.589760
9        512.0                  0.397488                  0.605888
10      1024.0                  0.401248                  0.581952
11      2048.0                  0.407232                  0.559232
12      4096.0                  0.416368                  0.717936

Qwen2-57B-A14B-Instruct FP8 W8A8 TP4

python3 benchmark/kernels/fbgemm/benchmark_fbgemm_grouped_gemm.py --model Qwen/Qwen2-57B-A14B-Instruct --tp-size 4 --use-fp8-w8a8 

    batch_size  FBGEMM Grouped GEMM FP8  SGLang Grouped GEMM FP8
0          1.0                 0.042560                 0.022336
1          2.0                 0.041312                 0.022128
2          4.0                 0.040384                 0.022240
3          8.0                 0.041184                 0.022016
4         16.0                 0.040128                 0.022816
5         32.0                 0.014272                 0.021440
6         64.0                 0.212832                 0.595040
7        128.0                 0.211328                 0.598688
8        256.0                 0.211776                 0.590992
9        512.0                 0.213504                 0.606304
10      1024.0                 0.216864                 0.582624
11      2048.0                 0.220512                 0.558128
12      4096.0                 0.227296                 0.718848

meta-llama/Llama-4-Scout-17B-16E-Instruct FP16 TP8

python3 benchmark/kernels/fbgemm/benchmark_fbgemm_grouped_gemm.py --model meta-llama/Llama-4-Scout-17B-16E-Instruct --tp-size 8 

grouped-gemm-performance:
    batch_size  FBGEMM Grouped GEMM BF16  SGLang Grouped GEMM BF16
0          1.0                  0.034592                  0.022816
1          2.0                  0.033440                  0.022016
2          4.0                  0.033984                  0.022400
3          8.0                  0.324592                  0.532960
4         16.0                  0.321024                  0.516960
5         32.0                  0.322736                  0.695840
6         64.0                  0.321184                  0.607008
7        128.0                  0.321264                  0.475136
8        256.0                  0.321984                  0.419232
9        512.0                  0.325728                  0.363392
10      1024.0                  0.339616                  0.693824
11      2048.0                  0.396928                  1.383792
12      4096.0                  0.732640                  2.761792

meta-llama/Llama-4-Scout-17B-16E-Instruct FP8 TP8

python3 benchmark/kernels/fbgemm/benchmark_fbgemm_grouped_gemm.py --model meta-llama/Llama-4-Scout-17B-16E-Instruct --tp-size 8 --use-fp8-w8a8

grouped-gemm-performance:
    batch_size  FBGEMM Grouped GEMM FP8  SGLang Grouped GEMM FP8
0          1.0                 0.042336                 0.020592
1          2.0                 0.006464                 0.013536
2          4.0                 0.006464                 0.014112
3          8.0                 0.171712                 0.531744
4         16.0                 0.170944                 0.518208
5         32.0                 0.170432                 0.693952
6         64.0                 0.172704                 0.608352
7        128.0                 0.173248                 0.475200
8        256.0                 0.175040                 0.420544
9        512.0                 0.178400                 0.367200
10      1024.0                 0.196736                 0.697968
11      2048.0                 0.230688                 1.385600
12      4096.0                 0.383872                 2.766432

The conclusion is that FBGEMM can achieve significant performance improvements over SGLang's grouped GEMM implementation for MoE models. This kernel can be directly applied to SGLang's EP-MoE grouped GEMM kernel to boost performance under fp16/bf16 and per-tensor quantized fp8 conditions.

Limitation

The current limitation is that warp-specialized kernels appear unavailable without compiling the Meta-specific Triton version. Additionally, this kernel currently only supports fp16/bf16 and per-tensor quantized fp8w8a8—further modifications would be needed for DeepSeek compatibility.

[gemini-code-assist]

Hello @BBuf, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

Summary of Changes

Hello! gemini-code-assist here, providing a summary of this pull request. This PR introduces a benchmark script and the necessary kernel code to compare the performance of FBGEMM's grouped GEMM implementation against the existing SGLang grouped GEMM kernel, specifically targeting Mixture-of-Experts (MoE) models. The goal is to evaluate if the FBGEMM kernel can offer performance improvements for MoE layers in SGLang, particularly under BF16 and FP8 W8A8 conditions. The PR includes the FBGEMM kernel code itself and a detailed benchmark script that tests various batch sizes and model configurations, presenting the results and highlighting the potential performance gains observed with FBGEMM.

Highlights

• New Benchmark Script: Adds a new Python script (benchmark/fbgemm/benchmark_fbgemm_grouped_gemm.py) to benchmark FBGEMM's grouped GEMM kernel against SGLang's existing implementation.
• FBGEMM Kernel Integration: Includes the FBGEMM grouped GEMM kernel code (benchmark/fbgemm/fbgemm_grouped_gemm.py), likely copied from the PyTorch FBGEMM project, supporting both standard (BF16/FP16) and FP8 row-wise quantized inputs.
• Performance Comparison: The benchmark results provided in the description show that the FBGEMM kernel can achieve significant performance improvements over the current SGLang grouped GEMM for various model sizes and data types, especially at larger batch sizes.
• Correctness Verification: Includes a test script (benchmark/fbgemm/test_grouped_gemm.py) to verify the numerical correctness of the integrated FBGEMM kernel against the SGLang implementation for different group configurations (uniform and non-uniform) and data types.
• Triton Warp Specialization: Notes a limitation where Triton's warp-specialized features might not be available without a specific Triton build, which could impact potential performance gains from the FBGEMM kernel's advanced features.

Changelog

Click here to see the changelog

• benchmark/fbgemm/benchmark_fbgemm_grouped_gemm.py
  • Added a new script for benchmarking FBGEMM grouped GEMM.
  • Includes logic to load model configurations from Hugging Face transformers.
  • Functions to create test data for BF16 and FP8 W8A8.
  • Uses Triton's perf_report for benchmarking across different batch sizes.
  • Compares FBGEMM and SGLang kernels.
  • Includes an optional correctness verification step.
• benchmark/fbgemm/fbgemm_grouped_gemm.py
  • Added the FBGEMM grouped GEMM kernel implementation.
  • Includes Triton JIT kernels for standard and FP8 row-wise GEMM.
  • Supports TMA load/store and warp specialization (conditional on Triton build).
  • Includes Python wrappers grouped_gemm and grouped_gemm_fp8_rowwise.
• benchmark/fbgemm/test_grouped_gemm.py
  • Added a new script for testing the correctness of the FBGEMM grouped GEMM kernel.
  • Includes functions to create uniform and non-uniform group configurations.
  • Tests BF16 correctness against SGLang's kernel.
  • Tests FP8 functionality (shape and dtype) for the FBGEMM kernel.
  • Uses pytest for test execution.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a benchmark for comparing FBGEMM's grouped GEMM kernel against SGLang's existing implementation, along with the FBGEMM kernel code itself (adapted from the FBGEMM repository) and unit tests. The benchmark results provided in the PR description are informative and demonstrate potential performance benefits with FBGEMM, especially for FP8 and larger batch sizes.

The code for the benchmark script, the FBGEMM kernels, and the tests are generally well-structured. The use of Triton for kernel implementation is appropriate for high-performance computing. The PR also clearly outlines current limitations, such as the dependency on a specific Triton build for warp-specialized features.

I've identified a few areas for improvement, primarily concerning the robustness of the benchmark data generation and long-term maintainability of the Triton kernel code. Addressing these would enhance the reliability and future upkeep of this valuable contribution.

Summary of Findings

• Benchmark Data Generation Correctness: The m_sizes calculation in create_test_data and create_fp8_test_data does not correctly handle cases where batch_size is not divisible by num_groups. This can lead to some input tokens not being processed, affecting benchmark accuracy. The seg_indptr calculation also needs to be updated accordingly.
• Model Configuration Robustness: The get_model_config function's fallback for unknown model architectures might use incorrect parameters if default attribute names don't apply, potentially leading to benchmarks with unintended model shapes.
• Triton Kernel Maintainability: The FBGEMM Triton kernel code (fbgemm_grouped_gemm.py) contains significant duplication across different kernel versions, which could pose long-term maintainability challenges. This is acknowledged by existing TODO comments.

Merge Readiness

This pull request adds valuable benchmarking capabilities and high-performance FBGEMM kernels. The benchmark results are promising. However, there are a few issues that should be addressed to ensure the benchmark's correctness and the code's long-term maintainability:

1. Critical: The generation of m_sizes and seg_indptr in the benchmark script needs to be robust to cases where batch_size is not divisible by num_groups to ensure all input data is processed.
2. Consideration: The model configuration fallback logic could be made more informative to users.
3. Consideration: The acknowledged code duplication in Triton kernels is a maintainability concern for the future.

I am unable to approve pull requests. Given the high-severity issue related to benchmark data generation, I recommend that these changes be made before merging to ensure the reliability of the benchmark. The medium-severity items should also be considered for improvement.

[gemini-code-assist]

The current calculation of m_sizes assumes batch_size is perfectly divisible by num_groups. If not, sum(m_sizes) (where m_sizes elements are batch_size // num_groups) would be less than batch_size. This means not all input tokens from x would be processed, potentially leading to incorrect or misleading benchmark results for certain configurations.

Could m_sizes be adjusted to ensure all batch_size tokens are distributed among the groups, for example, by distributing any remainder tokens? This would make the benchmark more robust for arbitrary batch_size and num_groups combinations.

    base_tokens_per_group = batch_size // num_groups
    remainder_tokens = batch_size % num_groups
    m_sizes_list = [base_tokens_per_group] * num_groups
    for i in range(remainder_tokens):
        m_sizes_list[i] += 1
    m_sizes = torch.tensor(m_sizes_list, dtype=torch.int64, device="cuda")
    # Ensure sum(m_sizes) == batch_size, critical for processing all input tokens.
    # Consider adding: assert torch.sum(m_sizes).item() == batch_size, "Sum of m_sizes must equal batch_size"

[gemini-code-assist]

Related to the m_sizes calculation: if m_sizes is updated to correctly distribute all tokens (especially when batch_size is not divisible by num_groups), the calculation for seg_indptr also needs to use these corrected m_sizes rather than the potentially inaccurate tokens_per_group.

SGLang's grouped GEMM relies on seg_indptr to define token segments for each expert. Using the actual m_sizes for its calculation is crucial for correctness.

    seg_indptr = torch.zeros(num_groups + 1, dtype=torch.int64, device="cuda")
    # seg_indptr[0] is 0. Calculate cumulative sum of actual m_sizes.
    # This assumes 'm_sizes' has been redefined to correctly sum to batch_size.
    torch.cumsum(m_sizes, dim=0, out=seg_indptr[1:])

[gemini-code-assist]

Similar to the create_test_data function, the m_sizes calculation here in create_fp8_test_data might not account for all tokens if batch_size is not divisible by num_groups. This could affect the FP8 benchmark path.

Would it be beneficial to apply the same robust m_sizes calculation here (distributing remainder tokens) to ensure all batch_size tokens are processed?

    base_tokens_per_group = batch_size // num_groups
    remainder_tokens = batch_size % num_groups
    m_sizes_list = [base_tokens_per_group] * num_groups
    for i in range(remainder_tokens):
        m_sizes_list[i] += 1
    m_sizes = torch.tensor(m_sizes_list, dtype=torch.int64, device="cuda")

[gemini-code-assist]

The fallback logic in get_model_config for unhandled model architectures defaults to using config.num_local_experts and config.intermediate_size. If a new, unhandled architecture doesn't conform to these attribute names, an AttributeError will occur. While this error is caught in main() and a hardcoded default configuration is used, this might lead to the benchmark running with parameters that don't match the user's intended model.

Could we enhance this fallback? For example, by printing a more specific warning when this else branch is taken for an architecture not explicitly listed in the if/elif conditions? This would alert the user that the benchmark might be using assumed (and potentially incorrect) parameters for their model.

    else:
        # Fallback for unrecognized architectures.
        print(f"Warning: Model architecture '{config.architectures[0]}' not explicitly handled. "
              f"Attempting to use default attributes 'num_local_experts' and 'intermediate_size'. "
              f"This may lead to incorrect benchmark parameters if these attributes are not applicable or missing.")
        num_groups = config.num_local_experts
        intermediate_size = config.intermediate_size

[gemini-code-assist]

This TODO comment correctly highlights a key maintainability aspect: there's significant code duplication between the different Triton kernel versions (e.g., standard vs. warp-specialized, FP16/BF16 vs. FP8).

While refactoring Triton kernels to minimize duplication without sacrificing performance is non-trivial, especially when adapting existing code, have you considered any strategies for this as a potential long-term improvement? For instance, could some common logic be encapsulated in tl.device_func utilities, or could Python-level pre-processing/code generation help manage variations if the core algorithmic structure is similar?

[zhyncs]

note: dead code

[jwfromm]

Hi @zhyncs, good to see you again! It's great that there's some value FBGEMM can bring to the SGLang community! I just wanted to note that FBGEMM's GenAI ops can now be pretty simply pip installed pip install fbgemm-gpu-genai and that would allow SGLang to not only use the triton kernels highlighted in this PR, but also the highly optimized CUTLASS grouped gemm operators that are comparably fast and often easier to deploy (https://github.com/pytorch/FBGEMM/blob/main/fbgemm_gpu/experimental/gen_ai/bench/quantize_ops.py#L1181). I'd love to see more integration between FBGEMM and the broader GenAI optimization community :)

[yiakwy-xpu-ml-framework-team]

@jwfromm do you have plan to work on fp8 (nv_fp8_e4m3, hip_fp8_e4m3_fnuz) ?

I just checked the source code, the hip is not supported yet :

https://github.com/pytorch/FBGEMM/blob/2f2a1ef555571022e36aa66f15b826cb3c55ba73/fbgemm_gpu/experimental/gen_ai/bench/quantize_ops.py#L1656

cc @HaiShaw

[jianyuh]

@yiakwy-xpu-ml-framework-team do you refer to fp8 support for grouped GEMM, or FP8 x INT4 GEMM?

For grouped GEMM (benchmarked in this PR), these two FP8 data types (on H100 / MI300x) are already supported:
(Triton version): https://github.com/pytorch/FBGEMM/blob/2f2a1ef555571022e36aa66f15b826cb3c55ba73/fbgemm_gpu/experimental/gemm/triton_gemm/grouped_gemm.py#L552

(CUTLASS version):
https://github.com/pytorch/FBGEMM/tree/2f2a1ef555571022e36aa66f15b826cb3c55ba73/fbgemm_gpu/experimental/gen_ai/src/quantize/cutlass_extensions/f8f8bf16_rowwise_grouped

(CK version):
https://github.com/pytorch/FBGEMM/tree/2f2a1ef555571022e36aa66f15b826cb3c55ba73/fbgemm_gpu/experimental/gen_ai/src/quantize/ck_extensions/fp8_rowwise_grouped

For FP8xINT4 GEMM (where you points to https://github.com/pytorch/FBGEMM/blob/2f2a1ef555571022e36aa66f15b826cb3c55ba73/fbgemm_gpu/experimental/gen_ai/bench/quantize_ops.py#L1656 ), it currently only has Nvidia H100 side support. We have efforts on enabling this on AMD MI300x.

[jianyuh]

When running benchmarks with triton==3.2.0, the following warning appears: we can't use warp-specialized features, but persistent kernels and TMA load/store remain available.

If we can run with triton 3.3, we might expect higher perf with warp specializations. cc @levendlee

[jwfromm]

@yiakwy-xpu-ml-framework-team Just to add a bit more to Jianyu's pointers, I actually think FBGEMM brings the most value in enabling AMD kernels since the options for Nvidia outside FBGEMM are a bit more fleshed out. We have full support across Nvidia and AMD for all kernels except the mixed dtype variety (FP8 activation INT4 weight) that you pointed out. Fortunately, AMD GPUs tend to have more RAM so the INT4 weight compression is less relevant. I definitely recommend trying out some of the FP8 activation FP8 weight kernels on AMD, I think you'll be really happy with the performance :)