[RFC]: Parallel strategy for Deepseek v3.2 like models

[ganyi1996ppo]

Suggestion Description

Current Deepseek v3.2 is inefficient with both tp or dp alone, based on deepseek v3.2's attribute, we found deploy this model with sequence parallel with specialized metadata preparing will be much more efficient than other parallel strategy.

Old Impl

Problem with tp:

• Redundant calculation in indexer path where wq_b and w_k is data parallel
• mla runs with num_head 16 and each token requires separate cache calculation, the compute intensity is quite low during prefill especially for long prompt

Problem with dp:

• Risk of high latency
• potential load imbalance over request

New proposal

Sequence parallel over token dim for both prefill and decode over tensor parallel group. No shard on tensor parallel dim for the weight. All gather over cache before the kv cache update step

Pros :

• full head on mla part, high tflops
• no redundant calculation on indexer or mla part
• gathered kv cache is quite small and fast, which can also been overlapped by two stream
• communication data amount cut to half to hidden_states
• support both piecewise and fullgraph

Cons :

• complicated metadata calculation
• potential execution time divergence over ranks, especially in piecewise graph

Operating System

No response

GPU

No response

ROCm Component

No response

[wuhuikx]

I add the data flow in yellow boxes, can you help check is that aligned with your proposal? @ganyi1996ppo

For your proposal, my understanding is summarized as follow.
The key optimization is to reduce the performance bottleneck of "Indexer" module, which is the bottleneck for DS-V3.2. For the Sparse MLA module, each Q is computed the MLA with the previous top-k numbers of KV instead of all KV as in DS-V3. The indexer is used to selected the history top-K KV by for_loop iteration on sequence_length of Q (for-loop iterations Q times). So the sequence parallel on the Q will help reduce the for_loop iterations from sequence_length to sequence_length/N across all the cards.

[ganyi1996ppo]

This graph can illustrate this better:

The bottleneck include 2 parts:

• Doing sparse_mla with sharded head (inefficient since its only tiling on num_heads)
• Doring indexer with data parallel (each card execute same indexer operation while what's necessary is only 1/sp)

[wuhuikx]

So all the attention module uses SP while the MOE use TP. We can also use EP for MoE right?

1. Attention SP + MoE TP
2. Attention SP + MoE EP