Skip to content

Commit 1954e0e

Browse files
committed
[TRTLLM-9493][feat] Custom AllToAll for helix parallelism
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
1 parent 9a1750c commit 1954e0e

File tree

9 files changed

+1522
-38
lines changed

9 files changed

+1522
-38
lines changed

cpp/tensorrt_llm/common/envUtils.cpp

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -531,6 +531,12 @@ bool getEnvEplbForceGdrcopy()
531531
return getBoolEnv("TRTLLM_EPLB_FORCE_GDRCOPY");
532532
}
533533

534+
bool getEnvUseNcclForHelix()
535+
{
536+
static bool const useNcclForHelix = getBoolEnv("TRTLLM_USE_NCCL_FOR_HELIX");
537+
return useNcclForHelix;
538+
}
539+
534540
} // namespace common
535541

536542
TRTLLM_NAMESPACE_END

cpp/tensorrt_llm/common/envUtils.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,9 @@ bool getEnvKVCacheTransferAllBlocksForWindow();
156156

157157
bool getEnvEplbForceGdrcopy();
158158

159+
// Whether to use NCCL-based all-to-all for Helix parallelism.
160+
bool getEnvUseNcclForHelix();
161+
159162
} // namespace common
160163

161164
TRTLLM_NAMESPACE_END

0 commit comments

Comments
 (0)