From 8af54cabcf04d38773eee3ef8992694922690e6b Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 24 Aug 2026 13:42:32 -0700 Subject: [PATCH] [Common] Default EP comm kernels to 32 SMs and keep shuffle/preprocess on all SMs Signed-off-by: Phuong Nguyen --- transformer_engine/common/ep/ep_backend.cpp | 13 ++++++++++++- .../common/include/transformer_engine/ep.h | 2 +- transformer_engine/jax/ep.py | 2 +- 3 files changed, 14 insertions(+), 3 deletions(-) diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index b726784d03..fb608639a8 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -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(group_config.num_comm_sms) - : NCCL_EP_AUTO; + : static_cast(std::min(kDefaultCommSms, device_sms)); // 0 = NCCL_EP_AUTO, which enables eager mode (recv buffers sized per routing). cfg.max_recv_tokens_per_rank = static_cast(group_config.max_recv_tokens_per_rank); cfg.zero_copy = group_config.zero_copy ? NCCL_EP_ZERO_COPY_ON : NCCL_EP_ZERO_COPY_OFF; @@ -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; diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h index 31417fbff1..9c4b23912a 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -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 diff --git a/transformer_engine/jax/ep.py b/transformer_engine/jax/ep.py index 32074b49e4..2222a41e48 100644 --- a/transformer_engine/jax/ep.py +++ b/transformer_engine/jax/ep.py @@ -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.