Fix MTP with deepep_low_latency

[MatthewBonanni]

Purpose

Using MTP with VLLM_ALL2ALL_BACKEND=deepep_low_latency causes an AssertionError due to incorrect handling of shared experts. This PR is a fix, carrying the logic from FusedMoE.forward_impl over to FusedMoE.forward_impl_chunked

Test Plan

VLLM_ALL2ALL_BACKEND=deepep_low_latency \
VLLM_ATTENTION_BACKEND=FLASH_ATTN_MLA \
vllm serve deepseek-ai/DeepSeek-R1 \
    --tensor-parallel-size 1 \
    --enable-expert-parallel \
    --data-parallel-size 8 \
    --trust-remote-code \
    --host 0.0.0.0 \
    --port 8000 \
    --speculative-config '{"method": "deepseek_mtp", "num_speculative_tokens": 1}'

Test Result

Previously hits AssertionError during startup, now starts up successfully.

---

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.

[gemini-code-assist]

Code Review

This pull request addresses an AssertionError that occurs when using speculative decoding with the deepep_low_latency backend on models featuring shared experts. The root cause was the forward_impl_chunked method in the FusedMoE layer lacking the logic to handle shared experts, which was already present in the non-chunked forward_impl method. The changes introduce this missing logic, ensuring consistent behavior and resolving the crash. The fix appears correct and aligns the chunked and non-chunked execution paths.

[MatthewBonanni]

cc @LucasWilkinson

[LucasWilkinson]

why does MTP mean we don't end up using the modular kernel?

[MatthewBonanni]

@LucasWilkinson the modular kernel is still used for the standard transformer blocks. With MTP, though, when running the mtp_block, self.quant_method.fused_experts is None and there are self.shared_experts (as a DeepseekV2MLP), so they need to be processed.

[LucasWilkinson]

Thanks!