[Bugfix] Fix Shared Expert/Zero expert code in FusedMoE.process_chunk

[SageMoore]

Purpose

#23991 appears to break Deepseek V2 Lite, and presumably other models that use shared experts. This PR should fix this issue.

Test Results

Deepseek V2 Lite

VLLM_ALL2ALL_BACKEND=deepep_low_latency vllm serve --model="deepseek-ai/DeepSeek-V2-Lite" --data-parallel-size 2 --enable-expert-parallel --gpu-memory-utilization 0.75 -compilation_config '{"cudagraph_mode": "full_decode_only"}'

Before

(EngineCore_DP0 pid=3298771)     process_chunk(chunk_start,
(EngineCore_DP0 pid=3298771)   File "/home/sagemoore/git/nm-vllm/vllm/model_executor/layers/fused_moe/layer.py", line 1943, in process_chunk
(EngineCore_DP0 pid=3298771)     chunk_start:chunk_end, :].copy_(final_hidden_states[1],
(EngineCore_DP0 pid=3298771) IndexError: index 1 is out of bounds for dimension 0 with size 1

After

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.3733|±  |0.0280|
|     |       |strict-match    |     5|exact_match|↑  |0.3700|±  |0.0279|

Longcat Flash

[gemini-code-assist]

Code Review

This pull request aims to fix an issue with shared experts and zero experts in FusedMoE.process_chunk. The changes correctly separate the logic for handling outputs from zero experts and shared experts, which was previously conflated, leading to incorrect behavior. Specifically, the logic now correctly handles the tuple returned when using shared experts. However, a new assertion assert len(final_hidden_states) == 1 has been introduced which is incorrect and will likely lead to runtime assertion errors. I have provided a critical comment with a suggested fix for this issue.

[SageMoore]

CC @OftenDream. I'm trying to get setup on a server large enough to run meituan-longcat/LongCat-Flash-Chat-FP8 but can you take a look at this fix?