Refactor communication logic of DeepSeek for extensibility and understandability

[fzyzcjy]

Motivation

Make code clean, not error-prune, extensible

Modifications

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[ch-wan]

forward_post_ffn

[fzyzcjy]

wondering why that is better than layer_end

Indeed it does the job of handling layer end transformations (e.g. make output shape correct for last layer). For non-last layers, it does not cooperate with pre_mlp, but indeed cooperate with next layer's pre_attn. So a verbose name (don't use it - only for illustration) should be forward_layer_end_transformation_or_cooperate_with_next_pre_attn

[ch-wan]

Checking enable_moe_dense_fully_dp() and (not self.is_layer_sparse) should be handled by the communicator. We can simplify this statement like this:

if hidden_states.shape[0] > 0 or self.layer_communicator.require_ffn_sync():

[fzyzcjy]

I am a bit confused about this, why is it related to communicators?

I originally wanted to move it inside DeepseekMLP with logic like "if tp_size=1 + hidden states is empty, then I do not compute" though. -> This is moved now to make code clearer (though it does not fit the PR title and I originally want to separately handle it).

[ch-wan]

It can be handled internally in LayerCommunicator. With this change, LayerScatterModes becomes a private class of communicator.py.

[fzyzcjy]

I make it exposed to allow in the future the LayerScatterModes to be used by modules other than LayerCommunicator, since LayerScatterModes tells us some facts about how layer input / attn / mlp / layer out shape looks like.