Lora MoE Align Improvements

[gnovack]

Purpose

This PR includes the following changes to the MoE LoRA Align kernel:

• Adds a global memory variant of moe_lora_align_sum_kernel (following the moe_align_block_size_kernel implementation). Previously only the shared memory variant was implemented, which led to worse performance for models with larger num_experts.
• Uses two cuda streams to execute moe_lora_align_sum_kernel and moe_align_block_size_kernel in parallel
• Refactors moe_lora_align_sum_kernel and moe_align_block_size_kernel to reduce duplicate logic between LoRA and non-LoRA cases (e.g. moe_align_block_size_kernel and moe_lora_align_block_size_kernel now call a common _moe_align_block_size function which supports both LoRA and non-LoRA cases).

FIX #30026

Test Plan

• Ran existing LoRA and moe_align_sum test cases

Test Result

Benchmark results w/ gpt-oss-120b before vs. after this change

Before:

============ Serving Benchmark Result ============
Successful requests:                     400       
Failed requests:                         0         
Maximum request concurrency:             8         
Benchmark duration (s):                  500.09    
Total input tokens:                      641602    
Total generated tokens:                  240183    
Request throughput (req/s):              0.80      
Output token throughput (tok/s):         480.28    
Peak output token throughput (tok/s):    576.00    
Peak concurrent requests:                13.00     
Total Token throughput (tok/s):          1763.27   
---------------Time to First Token----------------
Mean TTFT (ms):                          155.95    
Median TTFT (ms):                        142.50    
P99 TTFT (ms):                           582.98    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          16.38     
Median TPOT (ms):                        16.42     
P99 TPOT (ms):                           17.02     
---------------Inter-token Latency----------------
Mean ITL (ms):                           16.38     
Median ITL (ms):                         15.15     
P99 ITL (ms):                            121.57    
==================================================

After:

============ Serving Benchmark Result ============
Successful requests:                     400       
Failed requests:                         0         
Maximum request concurrency:             8         
Benchmark duration (s):                  480.35    
Total input tokens:                      641602    
Total generated tokens:                  240183    
Request throughput (req/s):              0.83      
Output token throughput (tok/s):         500.01    
Peak output token throughput (tok/s):    648.00    
Peak concurrent requests:                13.00     
Total Token throughput (tok/s):          1835.70   
---------------Time to First Token----------------
Mean TTFT (ms):                          166.84    
Median TTFT (ms):                        149.03    
P99 TTFT (ms):                           631.58    
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          15.71     
Median TPOT (ms):                        15.73     
P99 TPOT (ms):                           16.71     
---------------Inter-token Latency----------------
Mean ITL (ms):                           15.71     
Median ITL (ms):                         14.38     
P99 ITL (ms):                            129.20    
==================================================

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @gnovack.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[gemini-code-assist]

Code Review

This pull request introduces several significant improvements to the MoE LoRA alignment kernels. The refactoring to unify LoRA and non-LoRA logic by using common __device__ functions is a great step towards reducing code duplication and improving maintainability. The introduction of a global memory variant for moe_lora_align_sum_kernel and the use of a separate CUDA stream for parallel execution are solid performance optimizations. However, I've found a critical issue in csrc/moe/moe_align_sum_kernels.cu where a data type mismatch for token_lora_mapping could lead to incorrect memory access and undefined behavior. This needs to be addressed.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[ApostaC]

cc @tylertitsworth

[jeejeelee]

Most tests in test_moe_lora_align_sum.py are now failing.

[gnovack]

Most tests in test_moe_lora_align_sum.py are now failing.

These should be fixed now

[jeejeelee]

This check should verify that the num_experts` is less than or equal to 1024.
We also need to add more number in https://github.com/vllm-project/vllm/blob/main/tests/lora/test_moe_lora_align_sum.py#L35

[gnovack]

good catch. I just updated this check and added test cases for larger num_experts

[jeejeelee]

LGTM, @yewentao256 could you please take a look?

[jeejeelee]

@gnovack All LoRA tests are failing

[yewentao256]

Thanks for the work! Please fix the unit tests, all related I think

[gnovack]

Thanks for the work! Please fix the unit tests, all related I think

No problem! I think I've found the issue which is causing these test failures, so should have a fix out in the next few hours

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @gnovack.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[jeejeelee]

Thank you for contribution adn paitence