For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python function

calculate_ep_max_tokens_per_rank

calculate_ep_max_tokens_per_rank()

max.nn.comm.ep.calculate_ep_max_tokens_per_rank(*, max_batch_input_tokens, ep_size, data_parallel_degree, use_allreduce=False)

source

Calculate the maximum number of tokens per rank for EP communication.

Derives the tensor parallelism degree from ep_size and data_parallel_degree, then divides the batch tokens accordingly. When TP > 1, attention scatters tokens across ranks so each rank holds fewer tokens for the subsequent EP dispatch/combine phases.

Parameters:

• max_batch_input_tokens (int) – Maximum number of input tokens per batch.
• ep_size (int) – Expert parallelism size (total number of GPUs across nodes).
• data_parallel_degree (int) – Degree of data parallelism.
• use_allreduce (bool) – Is allreduce-backed expert parallelism enabled.

Returns:

Maximum tokens per rank for EP communication buffers.

Return type:

int