@xpmemeda I observe that you are reverting partial changes in #4390. However, using sum token is necessary to keep all intermediate tensor static. As an extreme case, if we have 8 DP workers and their processing batch sizes are [120, 1, 1, 1, 1, 1, 1, 1], the input of GroupedGeMM is of size [127, hidden_dim]. After this PR, it becomes [120*8, hidden_dim], most of the input tensors are padded with 0. The correct way to resolve this issue is to set a large batch size for cuda graph, or to use DP ffn and LM head. We are fixing our PRs to support the latter feature.

if we have 8 DP workers and their processing batch sizes are [120, 1, 1, 1, 1, 1, 1, 1], the input of GroupedGeMM is of size [127, hidden_dim].

@ch-wan 现状是 FFN 的输入大小会变成 [127 * 8, hidden_dim]，而不是 [127, hidden_dim]。

gently ping. @ch-wan @zhyncs @merrymercy

if we have 8 DP workers and their processing batch sizes are [120, 1, 1, 1, 1, 1, 1, 1], the input of GroupedGeMM is of size [127, hidden_dim].

@ch-wan 现状是 FFN 的输入大小会变成 [127 * 8, hidden_dim]，而不是 [127, hidden_dim]。

Sorry, I don't understand your response. Could you please add more details?

The current implementation was checked internally. It is correct.

To answer your question, bs in this line represents the global batch size so that it can reuse the input buffer of FFN. This design does not incur redundant computation because the sequence length of padded queries are 0s.

You concern regarding excessive tensor shape after all-gather is reasonable but not valid. We only copy the first several effective tokens to the communication buffer. See this.

The current implementation was checked internally. It is correct.

To answer your question, bs in this line represents the global batch size so that it can reuse the input buffer of FFN. This design does not incur redundant computation because the sequence length of padded queries are 0s.

You concern regarding excessive tensor shape after all-gather is reasonable but not valid. We only copy the first several effective tokens to the communication buffer. See this.

@ch-wan 好的，我理解了，谢谢回复。

The current implementation was checked internally. It is correct.

To answer your question, bs in this line represents the global batch size so that it can reuse the input buffer of FFN. This design does not incur redundant computation because the sequence length of padded queries are 0s.

You concern regarding excessive tensor shape after all-gather is reasonable but not valid. We only copy the first several effective tokens to the communication buffer. See this.

@ch-wan 这里引出了另一个问题，既然没有引入额外的计算，为什么改成 sum 之后，dp 性能大幅下降。也就是 issue中提到的内容，我前面的 benchmark 也验证了这点。

@xpmemeda Probably because the slow run did not activate CUDA graph. Please consider to increase --cuda-run-max-bs.

@xpmemeda Probably because the slow run did not activate CUDA graph. Please consider to increase --cuda-run-max-bs.

@ch-wan 我确认过了 CUDA graph 是生效的，nsight 分析了下，看起来用 sum 会导致大多数 kernel 耗时增加，可能要再确认一下 sum 是不是真的没有产生额外计算。

用 sum（before this pr）：

用 max（after this pr）：