[Kernel] Masked act_mul and fp8-quant Kernels for Batched MoE

[varun-sundar-rabindranath]

Purpose

Optimization to reduce unnecessary compute

For the batched MoEs we allocate tensors of shape [num_experts, max_tokens_per_expert, hidden_size]. On main we process all the elements (num_experts x max_tokens_per_expert x hidden_size) - but not all max_tokens_per_expert are valid and we can skip some of these. To this effect, add batched versions of silu_mul and per_token_quant fp8 kernels.

Test Plan

• Unit tests

• Local E2E testing
  commands:
  DeepSeek V2 lite:-

VLLM_ALL2ALL_BACKEND="pplx"  vllm serve deepseek-ai/DeepSeek-V2-Lite --trust-remote-code  --data-parallel-size 2 --tensor-parallel-size 1  --enable-expert-parallel --port 9020  --no-enable-prefix-caching 

 lm_eval --model local-completions --tasks gsm8k --model_args model=deepseek-ai/DeepSeek-V2-Lite,base_url=http://127.0.0.1:9020/v1/completions,num_concurrent=30,max_retries=1,tokenized_requests=False  --limit 100 --seed 42

Qwen FP8:-

VLLM_ALL2ALL_BACKEND="deepep_low_latency" VLLM_USE_DEEP_GEMM=1  vllm serve Qwen/Qwen3-30B-A3B-FP8 --trust-remote-code  --data-parallel-size 2 --tensor-parallel-size 1  --enable-expert-parallel --port 9020  --no-enable-prefix-caching

 lm_eval --model local-completions --tasks gsm8k --model_args model=Qwen/Qwen3-30B-A3B-FP8,base_url=http://127.0.0.1:9020/v1/completions,num_concurrent=1,max_retries=1,tokenized_requests=False  --limit 100 --seed 42

Test Result

DeepSeek v2 lite

PR

|Tasks|Version| 	Filter 	|n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|  	3|flexible-extract| 	5|exact_match|↑  | 0.27|±  |0.0446|
| 	|   	|strict-match	| 	5|exact_match|↑  | 0.27|±  |0.0446|

main
|Tasks|Version| 	Filter 	|n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|  	3|flexible-extract| 	5|exact_match|↑  | 0.27|±  |0.0446|
| 	|   	|strict-match	| 	5|exact_match|↑  | 0.27|±  |0.0446|

Qwen Fp8

PR

|Tasks|Version| 	Filter 	|n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|  	3|flexible-extract| 	5|exact_match|↑  | 0.84|±  |0.0368|
| 	|   	|strict-match	| 	5|exact_match|↑  | 0.88|±  |0.0327|

main
|Tasks|Version| 	Filter 	|n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|  	3|flexible-extract| 	5|exact_match|↑  | 0.88|±  |0.0327|
| 	|   	|strict-match	| 	5|exact_match|↑  | 0.92|±  |0.0273|

Note:

The lm_eval results become a bit finicky when I try to use big num_concurrent values. This happens also on main - I have set it to 1 here to produce output that are a bit consistent.

(Optional) Documentation Update

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[varun-sundar-rabindranath]

Marking this draft -- These kernels are not a priority at the moment given that a masked-fused-act-mul-quant exists in https://github.com/vllm-project/vllm/tree/ll_deepgemm_opt . We can revive this when needed.