Skip to content

[Common] Default EP comm kernels to 32 SMs; keep shuffle/preprocess on all SMs #3422

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
phu0ngng wants to merge 1 commit into
base: main
Choose a base branch
from
+14 −3
Open

[Common] Default EP comm kernels to 32 SMs; keep shuffle/preprocess on all SMs #3422

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 12 additions & 1 deletion transformer_engine/common/ep/ep_backend.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -239,9 +239,12 @@ void EPBackend::init(ncclComm_t ep_comm, NVTEEpGroupConfig group_config) {
cfg.rdma_buffer_size = NCCL_EP_AUTO;
cfg.num_qp_per_rank = NCCL_EP_AUTO;
cfg.num_channels = NCCL_EP_AUTO;
// Default the dispatch/combine (comm) kernels to 32 SMs, clamped to the device SM count.
constexpr int kDefaultCommSms = 32;
const int device_sms = cuda::sm_count();
cfg.max_num_sms = group_config.num_comm_sms > 0
? static_cast<unsigned int>(group_config.num_comm_sms)
: NCCL_EP_AUTO;
: static_cast<unsigned int>(std::min(kDefaultCommSms, device_sms));
Comment on lines 245 to +247

@greptile-apps greptile-apps Bot Aug 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Stale zero-budget diagnostic

The implementation now maps num_comm_sms == 0 to min(32, device_sms), but the native validation message still describes zero as automatic sizing, misleading users who are diagnosing or selecting the communication-kernel SM budget.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

// 0 = NCCL_EP_AUTO, which enables eager mode (recv buffers sized per routing).
cfg.max_recv_tokens_per_rank = static_cast<unsigned int>(group_config.max_recv_tokens_per_rank);
cfg.zero_copy = group_config.zero_copy ? NCCL_EP_ZERO_COPY_ON : NCCL_EP_ZERO_COPY_OFF;
Expand All @@ -250,6 +253,14 @@ void EPBackend::init(ncclComm_t ep_comm, NVTEEpGroupConfig group_config) {
cfg.overflow_policy =
group_config.drop_on_overflow ? NCCL_EP_OVERFLOW_DROP : NCCL_EP_OVERFLOW_AUTO;

// Keep the local shuffle/preprocess kernels on all SMs by default (their cost scales inversely
// with SM count) so the comm-SM cap above does not throttle them. overwrite=0 respects a
// user-set value and only fills in the default when unset.
char sm_buf[16];
std::snprintf(sm_buf, sizeof(sm_buf), "%d", device_sms);
setenv("NCCL_EP_SHUFFLE_SMS", sm_buf, /*overwrite=*/0);
setenv("NCCL_EP_PREPROCESS_NUM_SMS", sm_buf, /*overwrite=*/0);

NVTE_CHECK_NCCL(ncclEpCreateGroup(&ep_group_, ep_comm, &cfg));

ep_comm_ = ep_comm;
Expand Down
2 changes: 1 addition & 1 deletion transformer_engine/common/include/transformer_engine/ep.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ typedef struct {
int max_recv_tokens_per_rank;
/*! Token hidden dimension. */
int hidden_dim;
/*! Max SMs for NCCL EP dispatch/combine kernels. 0 = auto. */
/*! Max SMs for NCCL EP dispatch/combine kernels. 0 = default (32, clamped to device SM count). */
int num_comm_sms;
/*! Widest token dtype the group will dispatch; sizes staging buffers.
* Required (no default): must be set to a real token dtype. Per-dispatch
Expand Down
2 changes: 1 addition & 1 deletion transformer_engine/jax/ep.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,7 @@ def ep_bootstrap(
at least ep_size * max_tokens_per_rank * top_k to avoid drops.
hidden_dim: Feature dimension of token tensors passed to ep_dispatch.
max_token_dtype: Widest dtype the group will dispatch (only bfloat16 supported).
max_num_sms: SM budget for EP kernels; 0 = auto.
max_num_sms: SM budget for the dispatch/combine kernels; 0 = default (32).
drop_on_overflow: Drop tokens exceeding recv_capacity_per_rank instead of
trapping on overflow. Dropped tokens are still counted in
total_recv_tokens, so callers can detect overflow from it.
Expand Down
Loading