[Bugfix] [Performance] DeepEPHighThroughput + DeepSeek : Quant before Dispatch

[varun-sundar-rabindranath]

Purpose

DeepEPHighThroughput All2All kernel when used with DeepSeek models dispatches the tokens in 16bit datatype and quantizes after dispatch. This is inefficient for 2 reasons,

• More data in communication
• More data to quantize after dispatch

This PR introduces a fix to quantize to fp8 first and then dispatch the fp8 tensor.

Test Plan

canhazgpu run -g2 -- pytest -s tests/kernels/moe/test_modular_kernel_combinations.py

canhazgpu run -g2 -- pytest tests/kernels/moe/test_deepep_deepgemm_moe.py

VLLM_ALL2ALL_BACKEND="deepep_high_throughput" VLLM_USE_DEEP_GEMM=1  canhazgpu run -g 2 --  vllm serve Qwen/Qwen3-30B-A3B-FP8  --trust-remote-code --enable-expert-parallel --data-parallel-size 2 --port 9010 --no-enable-prefix-caching

Test Result

All tests pass for canhazgpu run -g2 -- pytest -s tests/kernels/moe/test_modular_kernel_combinations.py

All tests pass for canhazgpu run -g2 -- pytest tests/kernels/moe/test_deepep_deepgemm_moe.py

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  | 0.86|±  |0.0349|
|     |       |strict-match    |     5|exact_match|↑  | 0.94|±  |0.0239|

[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.

🚀

[gemini-code-assist]

Code Review

This pull request introduces a performance optimization for MoE layers using DeepEPHighThroughput with block quantization (e.g., for DeepSeek models). The change correctly modifies the logic to quantize the activations before dispatching them, which reduces communication overhead and is more efficient.

The implementation is clean and effective. The condition for pre-quantization is correctly expanded to include block-quantized cases, and the call to the quantization kernel is updated to pass the correct parameters, which also fixes a potential bug that the logical change would have otherwise introduced.

Overall, the changes look solid and align well with the stated purpose. I couldn't find any issues of high or critical severity.

[varun-sundar-rabindranath]

@tlrmchlsmth @bnellnm PTAL ! Thanks 🙌

[bnellnm]

So we still go down the "quantize after" codepath if the quantization is per-tensor? Is there some reason that quantization can't happen beforehand in that case also? Or does DeepEP not support that?

[varun-sundar-rabindranath]

So we still go down the "quantize after" codepath if the quantization is per-tensor? Is there some reason that quantization can't happen beforehand in that case also? Or does DeepEP not support that?

It is a DeepEP limitation. DeepEP doesn't support that.

[bnellnm]

So we still go down the "quantize after" codepath if the quantization is per-tensor? Is there some reason that quantization can't happen beforehand in that case also? Or does DeepEP not support that?

It is a DeepEP limitation. DeepEP doesn't support that.

Would it make sense to fake it out by replicating the scale and then resizing/truncating them after the dispatch?

[varun-sundar-rabindranath]

So we still go down the "quantize after" codepath if the quantization is per-tensor? Is there some reason that quantization can't happen beforehand in that case also? Or does DeepEP not support that?

It is a DeepEP limitation. DeepEP doesn't support that.

Would it make sense to fake it out by replicating the scale and then resizing/truncating them after the dispatch?

I went back and looked at the DeepEP documentation here
The documentation suggests that only block-quantization is supported. But the function seemingly also supports per-token quantization (We have unit test that have been passing - look here ).

However, it looks like we are an assert away in the DeepEP repo from crashing. To be safe, I have updated the code to support only block-quantization for the "Quant-then-Dispatch" block. For any other quantization we will "Dispatch-then-Quant"

cc @tlrmchlsmth

[tlrmchlsmth]

Thanks!