Add CUDA kernel size analysis tool for sgl-kernel optimization

[BBuf]

Motivation

This PR adds a kernel size analysis tool to help identify and optimize large CUDA kernels in sgl-kernel wheel packages.

• Extracts and analyzes all CUDA kernels from wheel files
• Groups kernels by name prefix (before <) to identify bloat from template instantiation
• Sorts kernels by size to find optimization targets
• Handles files without CUDA code gracefully
• Outputs both human-readable text reports and machine-readable JSON

Output In B200(SM100)

============================================================================================================================================
CUDA Kernel Size Analysis
============================================================================================================================================

Total kernels: 10262
Total size: 1631.80 MB (1,711,061,408 bytes)
Average kernel size: 162.83 KB

============================================================================================================================================
Kernel Groups (by name prefix)
============================================================================================================================================
Rank   Kernel Prefix                                                                    Count    Total (MB)   %       
--------------------------------------------------------------------------------------------------------------------------------------------
1      void cutlass::device_kernel                                                      3648     651.90       39.95   
2      void marlin_moe_wna16::Marlin                                                    720      370.32       22.69   
3      void marlin::Marlin                                                              1080     292.63       17.93   
4      void fast_hadamard_transform_kernel                                              294      89.20        5.47    
5      void flash::flash_fwd_sparse_kernel                                              32       11.43        0.70    
6      void cutlass::Kernel2                                                            96       11.24        0.69    
7      void per_token_group_quant_8bit_kernel                                           196      8.82         0.54    
8      void flashinfer::sampling::TopKTopPSamplingFromProbKernel                        40       7.54         0.46    
9      void flashinfer::sampling::TopKSamplingFromProbKernel                            40       7.50         0.46    
10     void router_gemm_kernel_bf16_output                                              64       6.89         0.42    
11     void router_gemm_kernel_float_output                                             64       6.83         0.42    
12     void mscclpp::executionKernel                                                    30       5.48         0.34    
13     void moe_fused_gate_kernel                                                       24       3.51         0.22    
14     void flashinfer::sampling::OnlineSoftmaxFusedKernel                              40       3.44         0.21    
15     void topkGatingSigmoid                                                           54       3.43         0.21    
16     void flashinfer::sampling::ChainSpeculativeSampling                              20       3.33         0.20    
17     void flashinfer::sampling::TopPSamplingFromProbKernel                            20       3.28         0.20    
18     void flashinfer::norm::FusedAddRMSNormKernel                                     30       3.15         0.19    
19     void topkGatingSoftmax                                                           54       3.03         0.19    
20     void flashinfer::BatchQKApplyRotaryPosIdsCosSinCacheEnhancedKernel               96       2.92         0.18    
21     void moe_vec_q                                                                   114      2.85         0.17    
22     void sglang::gptq::gemm_half_q_half_gptq_8bit_kernel                             16       2.83         0.17    
23     void flashinfer::sampling::MinPSamplingFromProbKernel                            20       2.77         0.17    
24     void mul_mat_vec_q                                                               114      2.56         0.16    
25     void flashinfer::sampling::TopKMaskLogitsKernel                                  20       2.56         0.16    
26     void marlin_moe_wna16::permute_cols_kernel                                       10       2.33         0.14    
27     void flashinfer::sampling::TopKRenormProbKernel                                  20       2.18         0.13    
28     void flashinfer::BatchQKApplyRotaryPosIdsCosSinCacheEnhancedHeadParallelis...    96       2.14         0.13    
29     void causal_conv1d_fwd_kernel                                                    36       2.13         0.13    
30     void flashinfer::sampling::SamplingFromProbKernel                                20       1.91         0.12    
31     void flashinfer::sampling::OnlineSoftmaxReduceKernel                             20       1.89         0.12    
32     void sglang::gptq::gemm_half_q_half_gptq_4bit_kernel                             16       1.86         0.11    
33     void flashinfer::sampling::TopPRenormProbKernel                                  10       1.73         0.11    
34     void sglang::gptq::gemm_half_q_half_gptq_3bit_kernel                             16       1.60         0.10    
35     void flashinfer::sampling::TreeSpeculativeSamplingTargetOnly                     20       1.57         0.10    
36     void flashinfer::sampling::SamplingFromLogitsKernel                              20       1.43         0.09    
37     void flashinfer::sampling::OnlineSoftmaxMapKernel                                20       1.41         0.09    
38     void moe_q2_K                                                                    12       1.39         0.08    
39     void dense_kernel0                                                               14       1.38         0.08    
40     void mul_mat_q2_K                                                                12       1.32         0.08    
41     void causal_conv1d_update_kernel                                                 36       1.31         0.08    
42     void moe_q3_K                                                                    12       1.29         0.08    
43     void flashinfer::BatchQKApplyRotaryPosIdsCosSinCacheKernel                       48       1.23         0.08    
44     void mul_mat_q3_K                                                                12       1.23         0.08    
45     void moe_q6_K                                                                    12       1.15         0.07    
46     void flashinfer::norm::RMSNormKernel                                             20       1.14         0.07    
47     void moe_fused_gate_kernel_dynamic                                               6        1.05         0.06    
48     void sglang::gptq::gemm_half_q_half_gptq_2bit_kernel                             16       1.04         0.06    
49     void per_token_quant_fp8_kernel                                                  18       1.02         0.06    
50     void moe_q5_K                                                                    12       1.00         0.06    
51     void moe_q4_K                                                                    12       0.93         0.06    
52     void per_token_quant_fp8_small_batch_kernel                                      18       0.89         0.05    
53     void moe_q5_0                                                                    12       0.88         0.05    
54     void flashinfer::BatchQKApplyRotaryPosIdsCosSinCacheHeadParallelismKernel        48       0.88         0.05    
55     void mul_mat_q6_K                                                                12       0.84         0.05    
56     void mul_mat_q5_0                                                                12       0.83         0.05    
57     void sglang::cross_device_reduce_2stage                                          24       0.79         0.05    
58     void moe_q5_1                                                                    12       0.73         0.04    
59     void mul_mat_q5_K                                                                12       0.73         0.04    
60     void moe_sum_reduce_kernel_warp_token_topk                                       32       0.73         0.04    
61     void moe_align_block_size_small_batch_expert_kernel                              10       0.71         0.04    
62     void tensorrt_llm::kernels::fusedQKNormRopeKernel                                12       0.70         0.04    
63     void mul_mat_q5_1                                                                12       0.69         0.04    
64     void mul_mat_q4_K                                                                12       0.67         0.04    
65     void flashinfer::activation::act_and_mul_kernel                                  18       0.65         0.04    
66     void moe_q4_0                                                                    12       0.63         0.04    
67     _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.63         0.04    
68     _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.62         0.04    
69     void moe_q4_1                                                                    12       0.60         0.04    
70     void sglang::allreduce_LL_2node                                                  6        0.60         0.04    
71     _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.59         0.04    
72     _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.59         0.04    
73     void mul_mat_q4_0                                                                12       0.59         0.04    
74     void mul_mat_q4_1                                                                12       0.57         0.03    
75     _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.55         0.03    
76     _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.55         0.03    
77     _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.55         0.03    
78     _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.55         0.03    
79     void marlin::gptq_marlin_repack_kernel                                           8        0.55         0.03    
80     void LogitsBitmaskKernel                                                         48       0.54         0.03    
81     _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.54         0.03    
82     _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.54         0.03    
83     (anonymous namespace)::topk_transform_prefill_kernel((anonymous namespace)...    2        0.51         0.03    
84     void apply_shuffle_mul_sum_kernel                                                6        0.50         0.03    
85     (anonymous namespace)::topk_kernel((anonymous namespace)::FastTopKParams)        2        0.49         0.03    
86     void moeSoftmax                                                                  6        0.49         0.03    
87     (anonymous namespace)::topk_transform_prefill_ragged_kernel((anonymous nam...    2        0.48         0.03    
88     (anonymous namespace)::topk_transform_decode_kernel((anonymous namespace):...    2        0.48         0.03    
89     _ZN7cutlass6KernelINS_4gemm6kernel23GemmWithEpilogueVisitorINS1_11threadbl...    2        0.43         0.03    
90     void moe_q8_0                                                                    12       0.41         0.03    
91     _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.39         0.02    
92     _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.39         0.02    
93     _ZN7cutlass6KernelINS_4gemm6kernel23GemmWithEpilogueVisitorINS1_11threadbl...    2        0.39         0.02    
94     _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.39         0.02    
95     _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.39         0.02    
96     void sglang::cross_device_reduce_1stage                                          24       0.38         0.02    
97     void moe_align_block_size_kernel                                                 10       0.38         0.02    
98     _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.38         0.02    
99     _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.38         0.02    
100    _ZN7cutlass6KernelINS_4gemm6kernel23GemmWithEpilogueVisitorINS1_11threadbl...    2        0.37         0.02    
101    void cvt_fp16_to_fp4                                                             24       0.37         0.02    
102    void allreduce8                                                                  8        0.36         0.02    
103    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.35         0.02    
104    void mul_mat_q8_0                                                                12       0.35         0.02    
105    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.35         0.02    
106    _ZN7cutlass6KernelINS_4gemm6kernel23GemmWithEpilogueVisitorINS1_11threadbl...    2        0.34         0.02    
107    _ZN7cutlass6KernelINS_4gemm6kernel23GemmWithEpilogueVisitorINS1_11threadbl...    2        0.34         0.02    
108    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.34         0.02    
109    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.34         0.02    
110    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.33         0.02    
111    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.33         0.02    
112    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.32         0.02    
113    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.32         0.02    
114    _ZN7cutlass6KernelINS_4gemm6kernel23GemmWithEpilogueVisitorINS1_11threadbl...    2        0.32         0.02    
115    _ZN7cutlass6KernelINS_4gemm6kernel23GemmWithEpilogueVisitorINS1_11threadbl...    2        0.32         0.02    
116    sglang::gptq::reconstruct_gptq_3bit_kernel(unsigned int const*, __half con...    2        0.32         0.02    
117    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.32         0.02    
118    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.32         0.02    
119    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.32         0.02    
120    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.31         0.02    
121    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.31         0.02    
122    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.31         0.02    
123    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.31         0.02    
124    _ZN7cutlass6KernelINS_4gemm6kernel23GemmWithEpilogueVisitorINS1_11threadbl...    2        0.31         0.02    
125    _ZN7cutlass6KernelINS_4gemm6kernel23GemmWithEpilogueVisitorINS1_11threadbl...    2        0.31         0.02    
126    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.31         0.02    
127    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.30         0.02    
128    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.30         0.02    
129    void sglang::allreduce_LL_1node                                                  6        0.30         0.02    
130    void moeTopK                                                                     4        0.30         0.02    
131    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.30         0.02    
132    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.29         0.02    
133    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.29         0.02    
134    sglang::gptq::make_sequential_3bit_kernel(unsigned int const*, unsigned in...    2        0.29         0.02    
135    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.29         0.02    
136    _ZN7cutlass6KernelINS_4gemm6kernel23GemmWithEpilogueVisitorINS1_11threadbl...    2        0.28         0.02    
137    _ZN7cutlass6KernelINS_4gemm6kernel23GemmWithEpilogueVisitorINS1_11threadbl...    2        0.28         0.02    
138    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.28         0.02    
139    void moe_sum_reduce_kernel                                                       32       0.28         0.02    
140    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.28         0.02    
141    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.28         0.02    
142    void flash::flash_fwd_splitkv_mla_kernel                                         2        0.28         0.02    
143    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.27         0.02    
144    void marlin::awq_marlin_repack_kernel                                            4        0.26         0.02    
145    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.26         0.02    
146    void moe_sum_kernel                                                              18       0.26         0.02    
147    void per_tensor_quant_fp8_kernel                                                 6        0.26         0.02    
148    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.26         0.02    
149    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalIN4cute5tupleIJii...    2        0.26         0.02    
150    _ZN7cutlass6KernelINS_4gemm6kernel23GemmWithEpilogueVisitorINS1_11threadbl...    2        0.25         0.02    
151    _ZN7cutlass6KernelINS_4gemm6kernel23GemmWithEpilogueVisitorINS1_11threadbl...    2        0.25         0.02    
152    void transfer_kernel_impl                                                        16       0.25         0.02    
153    _ZN7cutlass13device_kernelINS_4gemm6kernel13GemmUniversalINS1_17GroupProbl...    2        0.25         0.02    
154    _ZN7cutlass6KernelINS_4gemm6kernel23GemmWithEpilogueVisitorINS1_11threadbl...    2        0.24         0.01    
155    _ZN7cutlass6KernelINS_4gemm6kernel23GemmWithEpilogueVisitorINS1_11threadbl...    2        0.23         0.01    
.....

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.
• Work with maintainers to merge your PR. See the PR Merge Process

[gemini-code-assist]

Summary of Changes

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!

This pull request introduces a valuable diagnostic tool aimed at improving the performance and reducing the binary size of sgl-kernel packages. By providing a detailed breakdown of CUDA kernel sizes, developers can pinpoint specific kernels or template instantiations that contribute significantly to the overall binary footprint, guiding future optimization efforts. This enhancement empowers more informed decisions regarding kernel design and compilation strategies.

Highlights

• New CUDA Kernel Size Analysis Tool: Introduced a new Python script, analyze_whl_kernel_sizes.py, designed to analyze CUDA kernel sizes within sgl-kernel wheel packages.
• Optimization Identification: The tool helps identify large CUDA kernels and potential bloat from template instantiation by grouping kernels by name prefix and sorting them by size.
• Comprehensive Reporting: It generates both human-readable text reports and machine-readable JSON outputs, providing detailed insights into kernel counts, total size, average size, and individual kernel contributions.
• Dependency Update: The cubloaty tool, which is essential for the kernel analysis, has been added as a dependency to the Dockerfile.
• Documentation: The sgl-kernel/README.md has been updated with clear instructions on how to install and use the new kernel size analysis tool.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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 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

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 valuable tool for analyzing CUDA kernel sizes within wheel packages, which will be very helpful for optimization efforts. The script is well-structured and handles various cases, such as files without CUDA code. I've added a few suggestions to improve the script's robustness, maintainability, and adherence to Python best practices. These include making file path handling more robust, using constants for formatting, and improving code style.

[Fridge003]

Seems really useful

[Fridge003]

@BBuf Please fix lint

[BBuf]

@BBuf Please fix lint

Thanks! Done.