Skip to content

fix ep_moe_reorder kernel bugs#6858

Merged
BBuf merged 7 commits intomainfrom
fix_ep_moe_reorder_kernel_bug
Jun 4, 2025
Merged

fix ep_moe_reorder kernel bugs#6858
BBuf merged 7 commits intomainfrom
fix_ep_moe_reorder_kernel_bug

Conversation

@BBuf
Copy link
Copy Markdown
Collaborator

@BBuf BBuf commented Jun 4, 2025

Motivation

  • The original cuda kernel implementation only supported float32, while inference inputs typically use bfloat16 or float16 dtypes, so I fixed it.
  • The CUDA kernel currently lacks type-specific vectorization length handling and doesn't account for non-divisible cases.
  • The scale computation could be moved outside topk when a1_scales_ptr != nullptr and use_per_token_if_dynamic is true, reducing redundant global memory reads during topk iterations
  • In benchmark script, the input tensor allocation was mistakenly included in the timing function.
  • Unit tests were not added.

benchmark in h100

ep-moe-pre-reorder-performance:
   batch_size  CUDA Kernel  Triton Kernel
0        64.0     9.952000      15.584000
1       128.0    10.144000      15.712000
2       256.0    12.864000      17.440001
3       512.0    17.824000      22.528000
4       640.0    23.712000      24.831999
5       768.0    24.896000      27.456000
6      1024.0    29.247999      33.920001
7      2048.0    55.039998      66.399999
8      4096.0   108.287998     147.551998

Modifications

Checklist

Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

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

Hello @BBuf, I'm Gemini Code Assist1! 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 here, providing a summary of this pull request to help everyone quickly understand the changes. This PR, authored by BBuf, focuses on fixing several issues within the ep_moe_reorder CUDA kernel and its associated tooling. The primary goals are to add support for different data types (specifically bfloat16 and float16), improve the kernel's handling of vectorization and non-divisible sizes, correct a benchmark timing issue, and add unit tests to verify the kernel's correctness.

Highlights

  • CUDA Kernel Dtype Support: The CUDA kernel ep_pre_reorder_cuda_kernel has been updated to be templated, allowing it to support different floating-point data types beyond just float32, including float16 and bfloat16. A PyTorch dtype dispatch mechanism is added to handle this.
  • Improved Vectorization and Remainder Handling: The CUDA kernel's vectorization logic has been refined. It now correctly calculates the vector size based on the data type and includes a separate loop to process the remaining elements in the hidden dimension that are not perfectly divisible by the vector size, ensuring all data is processed.
  • Benchmark Timing Fix: The benchmark script bench_moe_ep_pre_reorder.py was updated to move the tensor allocation step outside of the timed section, providing a more accurate measurement of the kernel execution time itself.
  • Unit Tests Added: A new test file test_ep_moe_reorder_kernel.py has been added. This file includes parameterized tests that compare the output of the CUDA kernel against the existing Triton kernel for various batch sizes, hidden sizes, topk values, data types (float16, bfloat16, float32), and scaling configurations.

Changelog

Click here to see the changelog
  • sgl-kernel/benchmark/bench_moe_ep_pre_reorder.py
    • Removed unused imports (itertools, triton.language).
    • Changed the default dtype for benchmarking from float32 to bfloat16.
    • Moved tensor allocation (alloc_tensors()) outside the triton.testing.do_bench calls for both CUDA and Triton providers to fix timing.
  • sgl-kernel/csrc/moe/ep_moe_reorder_kernel.cu
    • Templated the ep_pre_reorder_cuda_kernel on scalar_t to support multiple data types.
    • Updated input and output pointers to use the templated scalar_t type.
    • Adjusted the calculation of vec_size based on sizeof(scalar_t).
    • Added a second loop to handle elements in the hidden dimension that are not processed by the main vectorized loop (i.e., hidden_size % vec_size != 0).
    • Implemented DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16 in ep_moe_pre_reorder to dispatch the kernel based on the input tensor's dtype.
  • sgl-kernel/tests/test_ep_moe_reorder_kernel.py
    • Added a new test file for the ep_moe_reorder kernel.
    • Included helper functions create_test_tensors, run_cuda_kernel, and run_triton_kernel for test setup and execution.
    • Added a parameterized test test_ep_moe_pre_reorder_vs_triton covering various dimensions, topk, dtypes (float16, bfloat16, float32), and scaling types.
    • Used torch.testing.assert_close to compare CUDA and Triton kernel outputs, casting to float for comparison with specified tolerances.
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.


Kernel code, a tricky art,
Fixing bugs, playing a part.
Float16, bfloat too,
Now the kernel sees them through.
Tests confirm, the work is done.

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.

Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

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

Code Review

This pull request effectively addresses several important bugs and limitations in the ep_moe_reorder CUDA kernel and its associated benchmark and tests. The changes are well-implemented and significantly improve the kernel's functionality and robustness.

Key improvements include:

  • Multi-Dtype Support: The CUDA kernel is now templated to support float32, float16, and bfloat16, which is crucial for inference workloads. This is handled correctly using PyTorch's type dispatch macros.
  • Corrected Vectorization and Tail Handling: The kernel now correctly calculates vectorization length based on data type (vec_size = 16 / sizeof(scalar_t)) and includes a loop to process remaining elements when hidden_size is not perfectly divisible by vec_size. This fixes a significant correctness issue.
  • Benchmark Accuracy: The tensor allocation in the benchmark script has been moved outside the timed section, leading to more accurate performance measurements.
  • Unit Tests: A comprehensive suite of unit tests has been added, comparing the CUDA kernel's output against a Triton implementation across various configurations. This is a valuable addition for ensuring correctness and preventing regressions.

The code quality is high, and the changes align well with the stated motivations in the pull request description. Excellent work!

Summary of Findings

  • CUDA Kernel Block Size: In sgl-kernel/csrc/moe/ep_moe_reorder_kernel.cu, the block_size for the CUDA kernel launch is hardcoded to 512 in the host wrapper function ep_moe_pre_reorder. While 512 is a common and often effective block size, making this parameter configurable or dynamically tuned based on input shapes or device properties could offer further optimization opportunities in diverse scenarios. This is a minor point and more of a consideration for future enhancements. (Severity: low - not commented due to review settings)

Merge Readiness

The pull request is in excellent shape and addresses all its stated goals effectively. The code changes are correct, well-tested, and improve the kernel significantly. I have no major concerns and believe the PR is ready for merging after the standard internal checks and approvals. As an AI reviewer, I am not authorized to approve the pull request.

Comment thread sgl-kernel/tests/test_ep_moe_reorder_kernel.py Outdated
BBuf and others added 3 commits June 4, 2025 12:30
Co-authored-by: JieXin Liang <Alcanderian@users.noreply.github.com>
@FlamingoPg
Copy link
Copy Markdown
Collaborator

looks good

Copy link
Copy Markdown
Collaborator

@yuan-luo yuan-luo left a comment

Choose a reason for hiding this comment

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

Please rename test_ep_moe_reorder_kernel.py to test_ep_moe_pre_reorder_kernel.py

Copy link
Copy Markdown
Collaborator

@Alcanderian Alcanderian left a comment

Choose a reason for hiding this comment

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

LGTM

@BBuf
Copy link
Copy Markdown
Collaborator Author

BBuf commented Jun 4, 2025

Please rename test_ep_moe_reorder_kernel.py to test_ep_moe_pre_reorder_kernel.py

done

@BBuf BBuf requested a review from yuan-luo June 4, 2025 07:06
@BBuf BBuf merged commit bd75690 into main Jun 4, 2025
78 of 93 checks passed
@BBuf BBuf deleted the fix_ep_moe_reorder_kernel_bug branch June 4, 2025 11:14
ch-wan pushed a commit that referenced this pull request Jun 4, 2025
Co-authored-by: JieXin Liang <Alcanderian@users.noreply.github.com>
jianan-gu pushed a commit to jianan-gu/sglang that referenced this pull request Jun 12, 2025
Co-authored-by: JieXin Liang <Alcanderian@users.noreply.github.com>
xwu-intel pushed a commit to xwu-intel/sglang that referenced this pull request Jun 17, 2025
Co-authored-by: JieXin Liang <Alcanderian@users.noreply.github.com>
walker-ai pushed a commit to walker-ai/sglang that referenced this pull request Jul 8, 2025
Merge branch 'sgl_20250610_sync_tag047 of git@code.alipay.com:Theta/SGLang.git into main

https://code.alipay.com/Theta/SGLang/pull_requests/52


Reviewed-by: 剑川 <jianchuan.gys@antgroup.com>


* [Bugfix] Fix slice operation when chunk size mismatch (sgl-project#6697)
* [Bugfix] Fix ChatCompletion endpoint of mini_lb when stream is set (sgl-project#6703)
* [CI] Fix setup of disaggregation with different tp (sgl-project#6706)
* [PD] Remove Unnecessary Exception Handling for FastQueue.get() (sgl-project#6712)
* Fuse routed_scaling_factor in DeepSeek (sgl-project#6710)
* Overlap two kernels in DeepSeek with communication (sgl-project#6711)
* Minor refactor two-batch overlap (sgl-project#6682)
* Speed up when having padding tokens two-batch overlap (sgl-project#6668)
* [Feature] Support Flashinfer fp8 blockwise GEMM kernel on Blackwell (sgl-project#6479)
* Fix LoRA bench (sgl-project#6719)
* temp
* Fix PP for Qwen3 MoE (sgl-project#6709)
* [feat] triton kernel for get_last_loc (sgl-project#6676)
* [fix] more mem for draft_extend cuda_graph (sgl-project#6726)
* [PD] bug fix:  Update status if nixl receiver send a a dummy req. (sgl-project#6720)
* Tune memory arguments on B200 (sgl-project#6718)
* Add DeepSeek-R1-0528 function call chat template (sgl-project#6725)
* refactor(tool call): Fix BaseFormatDetector tool_index issue and refactor `parse_streaming_increment` (sgl-project#6715)
* Add draft extend CUDA graph for Triton backend (sgl-project#6705)
* refactor apply_w8a8_block_fp8_linear in fp (sgl-project#6545)
* [PD] Support completion endpoint (sgl-project#6729)
* PD Rust LB (PO2) (sgl-project#6437)
* Super tiny enable sole usage of expert distribution metrics and update doc (sgl-project#6680)
* Support picking variants of EPLB algorithms (sgl-project#6728)
* Support tuning DeepEP configs (sgl-project#6742)
* [test] add ut and bm for get_last_loc (sgl-project#6746)
* Fix mem_fraction_static for AMD CI (sgl-project#6748)
* [fix][RL] Fix DeepSeekV3ForCausalLM.post_load_weights for multiple update weight (sgl-project#6265)
* Improve EPLB logical to physical dispatch map (sgl-project#6727)
* Update DeepSeek-R1-0528 function call chat template (sgl-project#6765)
* [PD] Optimize time out logic and add env var doc for mooncake (sgl-project#6761)
* Fix aiohttp 'Chunk too big' in bench_serving (sgl-project#6737)
* Support sliding window in triton backend (sgl-project#6509)
* Fix shared experts fusion error (sgl-project#6289)
* Fix one bug in the grouped-gemm triton kernel (sgl-project#6772)
* update llama4 chat template and pythonic parser (sgl-project#6679)
* feat(tool call): Enhance Llama32Detector for improved JSON parsing in non-stream (sgl-project#6784)
* Support token-level quantization for EP MoE (sgl-project#6782)
* Temporarily lower mmlu threshold for triton sliding window backend (sgl-project#6785)
* ci: relax test_function_call_required (sgl-project#6786)
* Add intel_amx backend for Radix Attention for CPU (sgl-project#6408)
* Fix incorrect LoRA weight loading for fused gate_up_proj (sgl-project#6734)
* fix(PD-disaggregation): Can not get local ip (sgl-project#6792)
* [FIX] mmmu bench serving result display error (sgl-project#6525) (sgl-project#6791)
* Bump torch to 2.7.0 (sgl-project#6788)
* chore: bump sgl-kernel v0.1.5 (sgl-project#6794)
* Improve profiler and integrate profiler in bench_one_batch_server (sgl-project#6787)
* chore: upgrade sgl-kernel v0.1.5 (sgl-project#6795)
* [Minor] Always append newline after image token when parsing chat message (sgl-project#6797)
* Update CI tests for Llama4 models (sgl-project#6421)
* [Feat] Enable PDL automatically on Hopper architecture (sgl-project#5981)
* chore: update blackwell docker (sgl-project#6800)
* misc: cache is_hopper_arch (sgl-project#6799)
* Remove contiguous before Flashinfer groupwise fp8 gemm (sgl-project#6804)
* Correctly abort the failed grammar requests & Improve the handling of abort (sgl-project#6803)
* [EP] Add cuda kernel for moe_ep_pre_reorder (sgl-project#6699)
* Add draft extend CUDA graph for flashinfer backend  (sgl-project#6805)
* Refactor CustomOp to avoid confusing bugs (sgl-project#5382)
* Tiny log prefill time (sgl-project#6780)
* Tiny fix EPLB assertion about rebalancing period and recorder window size (sgl-project#6813)
* Add simple utility to dump tensors for debugging (sgl-project#6815)
* Fix profiles do not have consistent names (sgl-project#6811)
* Speed up rebalancing when using non-static dispatch algorithms (sgl-project#6812)
* [1/2] Add Kernel support for Cutlass based Fused FP4 MoE (sgl-project#6093)
* [Router] Fix k8s Service Discovery (sgl-project#6766)
* Add CPU optimized kernels for topk and rope fusions  (sgl-project#6456)
* fix new_page_count_next_decode (sgl-project#6671)
* Fix wrong weight reference in dynamic EPLB (sgl-project#6818)
* Minor add metrics to expert location updater (sgl-project#6816)
* [Refactor] Rename `n_share_experts_fusion` as `num_fused_shared_experts` (sgl-project#6735)
* [FEAT] Add transformers backend support  (sgl-project#5929)
* [fix] recover auto-dispatch for rmsnorm and rope (sgl-project#6745)
* fix ep_moe_reorder kernel bugs (sgl-project#6858)
* [Refactor] Multimodal data processing for VLM (sgl-project#6659)
* Decoder-only Scoring API (sgl-project#6460)
* feat: add dp-rank to KV events (sgl-project#6852)
* Set `num_fused_shared_experts` as `num_shared_experts` when shared_experts fusion is not disabled (sgl-project#6736)
* Fix one missing arg in DeepEP (sgl-project#6878)
* Support LoRA in TestOpenAIVisionServer and fix fused kv_proj loading bug. (sgl-project#6861)
* support 1 shot allreduce  in 1-node and 2-node using mscclpp (sgl-project#6277)
* Fix Qwen3MoE missing token padding optimization (sgl-project#6820)
* Tiny update error hints (sgl-project#6846)
* Support layerwise rebalancing experts (sgl-project#6851)
* Tiny allow profiler API to auto create directory (sgl-project#6865)
* Support Blackwell DeepEP docker images (sgl-project#6868)
* [EP] Add cuda kernel for moe_ep_post_reorder (sgl-project#6837)
* [theta]merge 0605
* oai: fix openAI client error with single request via batch api (sgl-project#6170)
* [PD] Fix potential perf spike caused by tracker gc and optimize doc (sgl-project#6764)
* Use deepgemm instead of triton for fused_qkv_a_proj_with_mqa (sgl-project#6890)
* [CUTLASS-FP4-MOE]  Introduce CutlassMoEParams class for easy initialization of Cutlass Grouped Gems Metadata (sgl-project#6887)
* bugfix(OAI): Fix image_data processing for jinja chat templates (sgl-project#6877)
* [CPU] enable CI for PRs, add Dockerfile and auto build task (sgl-project#6458)
* AITER backend extension and workload optimizations (sgl-project#6838)
* [theta]merge
* [theta]merge
* [Feature] Support Flashinfer fmha on Blackwell (sgl-project#6930)
* Fix a bug in abort & Improve docstrings for abort (sgl-project#6931)
* Tiny support customize DeepEP max dispatch tokens per rank (sgl-project#6934)
* Sync the changes on cuda graph runners (sgl-project#6932)
* [PD] Optimize transfer queue forward logic for dummy rank (sgl-project#6922)
* [Refactor] image data process in bench_serving (sgl-project#6879)
* [fix] logical_to_all_physical_map index 256 is out of bounds in EP parallel. (sgl-project#6767)
* Add triton fused moe kernel config for E=257 on B200 (sgl-project#6939)
* [sgl-kernel] update deepgemm (sgl-project#6942)
* chore: bump sgl-kernel v0.1.6 (sgl-project#6943)
* Minor compile fused topk (sgl-project#6944)
* [Bugfix] pipeline parallelism and Eagle Qwen2 (sgl-project#6910)
* Tiny re-introduce profile id logging (sgl-project#6912)
* Add triton version as a fused_moe_triton config search key to avoid performace decrease in different Triton version (sgl-project#5955)
* reduce torch.zeros overhead in moe align block size kernel (sgl-project#6369)
* chore: upgrade sgl-kernel v0.1.6 (sgl-project#6945)
* add fbgemm moe grouped gemm kernel benchmark (sgl-project#6924)
* [Docker] Add docker file for SGL Router (sgl-project#6915)
* Disabling mixed chunked prefill when eagle is enabled (sgl-project#6874)
* Add canary for EPLB rebalancing (sgl-project#6895)
* Refactor global_server_args_dict (sgl-project#6866)
* Fuse routed scaling factor in topk_reduce kernel (sgl-project#6220)
* Update server timeout time in AMD CI. (sgl-project#6953)
* [misc] add is_cpu() (sgl-project#6950)
* Add H20 fused MoE kernel tuning configs for DeepSeek-R1/V3 (sgl-project#6885)
* Add a CUDA kernel for fusing mapping and weighted sum for MoE. (sgl-project#6916)
* chore: bump sgl-kernel v0.1.6.post1 (sgl-project#6955)
* chore: upgrade sgl-kernel v0.1.6.post1 (sgl-project#6957)
* [DeepseekR1-FP4] Add Support for nvidia/DeepSeekR1-FP4 model (sgl-project#6853)
* Revert "Fuse routed scaling factor in topk_reduce kernel (sgl-project#6220)" (sgl-project#6968)
* [AMD] Add more tests to per-commit-amd (sgl-project#6926)
* chore: bump sgl-kernel v0.1.7 (sgl-project#6963)
* Slightly improve the sampler to skip unnecessary steps (sgl-project#6956)
* rebase h20 fused_moe config (sgl-project#6966)
* Fix CI and triton moe Configs (sgl-project#6974)
* Remove unnecessary kernels of num_token_non_padded (sgl-project#6965)
* Extend cuda graph capture bs for B200 (sgl-project#6937)
* Fuse routed scaling factor in deepseek (sgl-project#6970)
* Sync cuda graph runners (sgl-project#6976)
* Fix draft extend ut stability with flush cache (sgl-project#6979)
* Fix triton sliding window test case (sgl-project#6981)
* Fix expert distribution dumping causes OOM (sgl-project#6967)
* Minor remove one kernel for DeepSeek (sgl-project#6977)
* [perf][sgl-kernel] extend cutlass_mla_decode to support num_head < 128 (sgl-project#6929)
* Enable more unit tests for AMD CI. (sgl-project#6983)
* Use torch.compile to fuse flash attention decode metadata preparation (sgl-project#6973)
* Eliminate stream sync to speed up LoRA batch init  (sgl-project#6960)
* support qwen3 emebedding (sgl-project#6990)
* Fix torch profiler bugs for bench_offline_throughput.py (sgl-project#6557)
* chore: upgrade flashinfer v0.2.6.post1 jit (sgl-project#6958)
* cleanup tmp dir (sgl-project#7007)
* chore: update pr test xeon (sgl-project#7008)
* Fix cutlass MLA gets almost zero accuracy (sgl-project#6998)
* Update amd nightly models CI. (sgl-project#6992)
* feat: add direct routing strategy to DP worker (sgl-project#6884)
* Fallback to lower triton version for unfound fused moe configs (sgl-project#7013)
* Fix torchvision version for Blackwell (sgl-project#7015)
* Simplify prepare_extend_after_decode (sgl-project#6987)
* Migrate to assertEqual (sgl-project#6741)
* Fix torch version in blackwell dockerfile (sgl-project#7017)
* chore: update pr test xeon (sgl-project#7018)
* Update default settings for blackwell (sgl-project#7023)
* Support both approximate and exact expert distribution collection (sgl-project#6964)
* Add decode req pool (sgl-project#6980)
* [theta]merge 0610
* [theta]merge 0610
* [CI] Add CI workflow for sgl-router docker build (sgl-project#7027)
* Fix fused_moe triton configs (sgl-project#7029)
* CPU: map changes from developing branch in sgl-kernel (sgl-project#6833)
* chore: bump v0.4.7 (sgl-project#7038)
* Update README.md (sgl-project#7040)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants