Add dsv3 router gemm benchmark on blackwell - #17707
Conversation
Summary of ChangesHello @harrisonlimh, 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! This pull request introduces a comprehensive benchmarking suite to evaluate the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe 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 by creating a comment using either
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 Limitations & Feedback Gemini Code Assist 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. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds a new benchmark script to compare the performance of dsv3_router_gemm from sglang and flashinfer on Blackwell architecture. The script includes both correctness checks and performance measurements.
My review focuses on improving the correctness, readability, and maintainability of the new benchmark script. I've identified a bug in the command-line argument parsing that makes a flag ineffective. I've also suggested several refactorings to clarify the logic, especially around handling benchmark timings, which is currently confusing and error-prone. Other suggestions include minor performance improvements and code simplification.
| if provider == "sglang": | ||
| ms, min_ms, max_ms = triton.testing.do_bench( | ||
| lambda: dsv3_router_gemm_sgl( | ||
| hidden_states.clone(memory_format=torch.contiguous_format), | ||
| router_weights.clone(memory_format=torch.contiguous_format), | ||
| ), | ||
| quantiles=quantiles, | ||
| ) | ||
| elif provider == "flashinfer": | ||
| ms, min_ms, max_ms = triton.testing.do_bench( | ||
| lambda: dsv3_router_gemm_flashinfer( | ||
| hidden_states.clone(memory_format=torch.contiguous_format), | ||
| router_weights.clone(memory_format=torch.contiguous_format), | ||
| launch_with_pdl, | ||
| ), | ||
| quantiles=quantiles, | ||
| ) | ||
|
|
||
| # Calculate TFLOPS | ||
| flops = 2 * m * n * k # multiply-adds | ||
| tflops = flops / (ms * 1e-3) / 1e12 | ||
|
|
||
| # Print shape-specific results with TFLOPS | ||
| print(f"Time: {ms*1000:.2f} us, TFLOPS: {tflops:.2f}") | ||
| return ms, max_ms, min_ms |
There was a problem hiding this comment.
The handling of timing results from triton.testing.do_bench is confusing and error-prone due to variable naming and value swapping. do_bench returns (median, min, max). The current implementation swaps min and max values between functions, which makes the code hard to follow and maintain.
I suggest refactoring to use clearer variable names and a more direct data flow. This makes the code easier to understand and less prone to bugs. The suggested change also removes the redundant memory_format=torch.contiguous_format from .clone() calls, as the tensors are already contiguous.
You can refactor _benchmark as suggested. Then, in get_benchmark_plot_friendly and get_benchmark, the inner benchmark function should be updated to:
def benchmark(cfg_id, provider):
m, n, k, tp_size, launch_with_pdl = all_configs[cfg_id]
median_ms, min_ms, max_ms = _benchmark(m, n, k, tp_size, launch_with_pdl, provider)
return median_ms * 1000, min_ms * 1000, max_ms * 1000 if provider == "sglang":
median_ms, min_ms, max_ms = triton.testing.do_bench(
lambda: dsv3_router_gemm_sgl(
hidden_states.clone(),
router_weights.clone(),
),
quantiles=quantiles,
)
elif provider == "flashinfer":
median_ms, min_ms, max_ms = triton.testing.do_bench(
lambda: dsv3_router_gemm_flashinfer(
hidden_states.clone(),
router_weights.clone(),
launch_with_pdl,
),
quantiles=quantiles,
)
# Calculate TFLOPS
flops = 2 * m * n * k # multiply-adds
tflops = flops / (median_ms * 1e-3) / 1e12
# Print shape-specific results with TFLOPS
print(f"Time: {median_ms*1000:.2f} us, TFLOPS: {tflops:.2f}")
return median_ms, min_ms, max_ms| "--run-correctness", | ||
| action="store_true", | ||
| default=True, | ||
| help="Whether to run correctness test", |
There was a problem hiding this comment.
The argument --run-correctness is defined with action="store_true" and default=True. This combination makes the value of args.run_correctness always True, regardless of whether the flag is provided on the command line. This prevents disabling the correctness tests.
To make this flag work as intended (i.e., run correctness tests only when the flag is present), you should remove default=True. The default for action="store_true" is False.
"--run-correctness",
action="store_true",
help="Whether to run correctness test",| import torch | ||
| import triton | ||
| from flashinfer.gemm.routergemm_dsv3 import mm_M1_16_K7168_N256 | ||
| from sgl_kernel import dsv3_router_gemm as dsv3_router_gemm |
There was a problem hiding this comment.
The alias dsv3_router_gemm is the same as the imported name. This is redundant and can be simplified for better code clarity.
| from sgl_kernel import dsv3_router_gemm as dsv3_router_gemm | |
| from sgl_kernel import dsv3_router_gemm |
| output = torch.randn( | ||
| hidden_states.shape[0], | ||
| router_weights.shape[0], | ||
| device="cuda", | ||
| dtype=torch.float32, | ||
| ).contiguous() |
There was a problem hiding this comment.
The output tensor is initialized with random values using torch.randn, but its content is immediately overwritten by the mm_M1_16_K7168_N256 kernel. It's more efficient to just allocate the memory without initializing it by using torch.empty.
| output = torch.randn( | |
| hidden_states.shape[0], | |
| router_weights.shape[0], | |
| device="cuda", | |
| dtype=torch.float32, | |
| ).contiguous() | |
| output = torch.empty( | |
| hidden_states.shape[0], | |
| router_weights.shape[0], | |
| device="cuda", | |
| dtype=torch.float32, | |
| ).contiguous() |
| mismatch_percent = 1.0 - match_ratio.item() | ||
| if mismatch_percent > 1 - percent: | ||
| print( | ||
| f"Mismatch percentage is {mismatch_percent:.4f} for rtol {rtol} " | ||
| f"(threshold: {1 - percent:.4f})" | ||
| ) | ||
| return False |
There was a problem hiding this comment.
The condition if mismatch_percent > 1 - percent: is redundant. If the code reaches this point, it means match_ratio < percent, which is equivalent to mismatch_percent > 1 - percent. This condition will always be true. You can simplify the logic by removing this if statement, making the code easier to understand.
mismatch_percent = 1.0 - match_ratio.item()
print(
f"Mismatch percentage is {mismatch_percent:.4f} for rtol {rtol} "
f"(threshold: {1 - percent:.4f})"
)
return False| benchmark = ( | ||
| get_benchmark_plot_friendly(args.tp_sizes) | ||
| if args.plot_friendly | ||
| else get_benchmark(args.tp_sizes) | ||
| ) |
There was a problem hiding this comment.
This line uses a ternary operator inside parentheses to select the benchmark function. While it works, a standard if/else block would be more readable and is generally preferred for this kind of logic.
if args.plot_friendly:
benchmark = get_benchmark_plot_friendly(args.tp_sizes)
else:
benchmark = get_benchmark(args.tp_sizes)
|
The result seems to suggest that the new flashinfer kernel only boosts the performance for m=6 and is limited to 0.365% gain. As I am new to the area, I will wait for the team's review on the test set up before adding the usage of the kernel. cc: @Fridge003 |
@harrisonlimh Your findings are correct. This comes from the fact that both the sglang kernel and the current flashinfer kernel derive from the same TRTLLM kernel with some differences. So the performance being similar is expected. We plan to upgrade/fix the flashinfer kernel to have performance parity (or better) than the sglang kernel. Once there is an issue and/or work started on that we will let you know. Thank you for your effort thus far on this. |
|
@harrisonlimh I think we should not compare the native implementation with |
There was a problem hiding this comment.
@harrisonlimh
Could you help fix the issues of tensor creation, pdl issue and rerun the benchmark? I think on kernel-level, two kernel should be identical, so no performance difference is expected.
Thanks!
| launch_with_pdl: bool, | ||
| ): | ||
| """Flashinfer implementation of dsv3 router gemm""" | ||
| output = torch.randn( |
There was a problem hiding this comment.
As benchmark script, output should not be created with randn, as it introduce addtional overhead. We should use empty just like how SGLang function works.
Moreover, contiguous() is not needed, and it might intoduce additional overhead
There was a problem hiding this comment.
Hi! Thank you for the feedback!
Please see below for follow ups!:
- Made the suggested change and attached the result that compares PDL enabled and disabled results for flashinfer and sgl-kernel. The performance is on par as suggested.
- The initial regression on flashinfer was caused by the use of
radnandempty. Removing those alone confirmed that both kernels perform similarly. - Added slight refactoring to compare those kernels with both pdl on and off;
launch_with_pdlparam seems to take precedent over the env var for flashinfer, so used the env var for sgl-kernel and thelaunch_with_pdlfor flashinfer kernel
|
Flashinfer vs. SGLang kernel performance comparison
Full benchmark data
|
|
@harrisonlimh I discussed with @nv-yunzheq and based on your current findings with PDL enabled showing performance parity, it is safe to proceed with integrating the flashinfer dsv3 router gemm. |
|
Hi @leejnau @harrisonlimh , this PR breaks SM103, it seems CC 103 was not added to FI (it should be compat anyway...) Could either of you help upd: https://github.com/flashinfer-ai/flashinfer/blob/19329d838236803c59c8f571d42910b1c0c2a18e/flashinfer/gemm/routergemm.py#L180 to add 103 to supported CC? Thanks~ |
|
Hi @b8zhong, thank you for the heads up on this! Would a simple allowlisting be sufficient to enable the kernal for SM 103? |
|
@leejnau Could you check if this is a FlashInfer bug/issue? That kernel should work for both SM100 and SM103 |
…2991) <!-- .github/pull_request_template.md --> add SM 103 support for mm_M1_16_K7168_N256. cc: @b8zhong, @Fridge003, @leejnau <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> ## 🔍 Related Issues * sgl-project/sglang#17707 <!-- Link any related issues here --> ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [ ] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [ ] I have installed the hooks with `pre-commit install`. - [ ] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Broadened GPU architecture support so optimized matrix-multiplication routines also target an additional modern compute capability. * **Tests** * Updated test gating so positive and negative tests for the optimized routine run on the newly supported compute capability as well. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Brian K. Ryu <bryu@nvidia.com>
6 participants
Motivation
Modifications
Accuracy Tests
Confirmed accuracy tests pass for m=[1...16], n=256, k=7168, tp=[1...8] with PDL=[False, True]
Benchmarking and Profiling
Benchmark result
Flashinfer kernel performed better on below configs:
Flashinfer kernel performed equally as the existing kernel on below configs:
In all other scenarios, the existing kernel performed better.
Full result
Checklist
Review Process
/tag-run-ci-label,/rerun-failed-ci,/tag-and-rerun-ci