support TRTLLM_DEEP_EP_TOKEN_LIMIT to allow run deep-ep on memory-con…

[ttyio]

none: support TRTLLM_DEEP_EP_TOKEN_LIMIT to allow run deep-ep on memory-constrained GPUs.

Description

DeepEP requires additional RDMA memory for communication, and on memory-constrained GPUs, we may not have enough memory to enable DeepEP for both the context and decoding phases. In disaggregated serving scenarios, it's straightforward to enable DeepEP only on the decoding server. However, for inflight batching, we need to apply a token limit so that DeepEP is only used during decoding.

[ttyio]

/bot run

[tensorrt-cicd]

PR_Github #10696 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #10696 [ run ] completed with state FAILURE
/LLM/main/L0_MergeRequest_PR pipeline #7906 completed with status: 'ABORTED'

[ttyio]

/bot run

[tensorrt-cicd]

PR_Github #10859 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #10859 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #8025 completed with status: 'FAILURE'

[ttyio]

/bot run

[tensorrt-cicd]

PR_Github #10865 [ run ] triggered by Bot

[ttyio]

/bot run

[tensorrt-cicd]

PR_Github #10868 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #10865 [ run ] completed with state ABORTED

[tensorrt-cicd]

PR_Github #10868 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #8034 completed with status: 'FAILURE'

[ttyio]

/bot run

[tensorrt-cicd]

PR_Github #11186 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #11186 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #8275 completed with status: 'FAILURE'

[ttyio]

/bot run

[tensorrt-cicd]

PR_Github #11325 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #11325 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #8378 completed with status: 'SUCCESS'

[yilin-void]

The code related to self.use_postquant_alltoall should include the use_all_to_all check as well. Otherwise, there will be a scenario where the AG/RS path is taken, but the post quant all2all logic is triggered, just like the following code snippet.

if not disable_fp4_allgather() or self.use_postquant_alltoall:
    if isinstance(x, Fp4QuantizedTensor):
        x, x_sf = x.fp4_tensor, x.scaling_factor
        x_row = x.shape[0]
        # note: we use uint8 to store 2 fp4 values
        x_col = x.shape[1] * 2
    else:
        sf_swizzle = not self.use_postquant_alltoall
        x_row = x.shape[0]
        x_col = x.shape[1]
        x, x_sf = torch.ops.trtllm.fp4_quantize(
            x, self.fc31_input_scale, self.scaling_vector_size,
            False, sf_swizzle)
        if self.use_postquant_alltoall:
            x_sf = x_sf.view((x_row, -1))

[yuantailing]

There are a lot variables related to communication method selection: enable_alltoall, use_all_to_all (from this PR), and use_allgather (will be introduced by #5684 ), which makes the logic hard to follow.

This comment is not to block this PR, but suggest a future refactor on communication method selection.

(Discuss in the following comment)

[ttyio]

/bot run

[tensorrt-cicd]

PR_Github #11682 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #11682 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #8649 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

[ttyio]

/bot reuse-pipeline

[tensorrt-cicd]

PR_Github #11825 [ reuse-pipeline ] triggered by Bot

[tensorrt-cicd]

PR_Github #11825 [ reuse-pipeline ] completed with state SUCCESS
Reusing PR_Github #11682 for commit c64b707

[ttyio]

/bot run

[tensorrt-cicd]

PR_Github #11828 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #11828 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #8764 completed with status: 'FAILURE'

[ttyio]

/bot run

[tensorrt-cicd]

PR_Github #11850 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #11850 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #8781 completed with status: 'FAILURE'

[ttyio]

/bot run

[tensorrt-cicd]

PR_Github #11859 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #11859 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #8788 completed with status: 'SUCCESS'

[yuantailing]

Hi @ttyio ,
Do you have a suggested value for TRTLLM_DEEP_EP_TOKEN_LIMIT? Should it be equal to max_batch_size * (1 + MTP) or something else?

[ttyio]

Hi @ttyio , Do you have a suggested value for TRTLLM_DEEP_EP_TOKEN_LIMIT? Should it be equal to max_batch_size * (1 + MTP) or something else?

Hi @yuantailing , I have not test MTP, for non-MTP, I used max_local_batch_size.

[yuantailing]

Thank you! @ttyio