From f7fbe84ea06d9078a537d9236a0fec59904ecbcc Mon Sep 17 00:00:00 2001 From: Shreyas S Date: Tue, 28 Apr 2026 00:18:37 +0530 Subject: [PATCH 1/3] fix(moe): add HybridEP hardware limit guardrail and rename seq_len - Renames seq_len to um_tokens in init_hybrid_ep_buffer and HybridEPDispatch for clarity, as the variable actually represents the flattened micro-batch (seq_len * batch_size). - Adds a Python-side ValueError guardrail before DeepEP buffer initialization to catch RDMA Queue Pair depths that exceed the InfiniBand hardware limit (65535). This prevents ungraceful C++ SIGABRT crashes and instructs users to increase their Tensor/Context Parallelism degrees. --- megatron/core/transformer/moe/fused_a2a.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py index 39f50a4a670..592900684e0 100644 --- a/megatron/core/transformer/moe/fused_a2a.py +++ b/megatron/core/transformer/moe/fused_a2a.py @@ -278,7 +278,7 @@ def set_deepep_num_sms(num_sms): def init_hybrid_ep_buffer( group: torch.distributed.ProcessGroup, hidden_dim: int, - seq_len: int, + num_tokens: int, num_local_experts: int, num_sms_dispatch_api: int, num_sms_combine_api: int, @@ -313,7 +313,7 @@ def init_hybrid_ep_buffer( _hybrid_ep_buffer = HybridEPBuffer( group=group, hidden_dim=hidden_dim, - max_num_of_tokens_per_rank=seq_len, + max_num_of_tokens_per_rank=num_tokens, num_local_experts=num_local_experts, use_fp8=fp8_dispatch, num_sms_dispatch_api=num_sms_dispatch_api, From ff29c60d9000fbb9ac233c6742635e581686ec81 Mon Sep 17 00:00:00 2001 From: Shreyas S Date: Tue, 28 Apr 2026 00:27:27 +0530 Subject: [PATCH 2/3] fix(moe): add HybridEP hardware limit guardrail - Adds a Python-side ValueError guardrail before DeepEP buffer initialization to catch RDMA Queue Pair depths that exceed the InfiniBand hardware limit (65535). This prevents ungraceful C++ SIGABRT crashes and instructs users to increase their Tensor/Context Parallelism degrees. --- megatron/core/transformer/moe/fused_a2a.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py index 592900684e0..541035fd24e 100644 --- a/megatron/core/transformer/moe/fused_a2a.py +++ b/megatron/core/transformer/moe/fused_a2a.py @@ -351,12 +351,26 @@ def forward( Forward pass of fused dispatch of the HybridEP backend ''' if _hybrid_ep_buffer is None: - seq_len, hidden_dim = x.shape[-2:] + num_tokens, hidden_dim = x.shape[-2:] + + # --- Hardware Limit Gaurdrail --- + # DeepEP calculates tx_depth = 3 * num_tokens + 1. + # # InfiniBand strictly asserts tx_depth < 65536. + tx_depth = 3 * num_tokens + 1 + if tx_depth >= 65536: + raise ValueError( + f"HybridEP RDMA Queue Pair depth ({tx_depth}) exceeds the InfiniBand " + f"hardware limit of 65535. This occurs because the total tokens per rank " + f"({num_tokens}) is too high. Please reduce sequence length or micro-batch size, " + f"or increase Tensor Parallelism (TP) / Context Parallelism (CP) to reduce " + f"the number of tokens processed per rank." + ) + fp8_dispatch = False # Currently, we do not support fp8 dispatch init_hybrid_ep_buffer( group, hidden_dim, - seq_len, + num_tokens, num_local_experts, num_sms_dispatch_api, num_sms_combine_api, From 7548fb8654d5f59861a3ab24ef2aca0cca4e370f Mon Sep 17 00:00:00 2001 From: Shreyas S Date: Tue, 28 Apr 2026 18:18:41 +0530 Subject: [PATCH 3/3] fix lexical issues --- megatron/core/transformer/moe/fused_a2a.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/megatron/core/transformer/moe/fused_a2a.py b/megatron/core/transformer/moe/fused_a2a.py index 541035fd24e..e68c60d45ee 100644 --- a/megatron/core/transformer/moe/fused_a2a.py +++ b/megatron/core/transformer/moe/fused_a2a.py @@ -353,19 +353,18 @@ def forward( if _hybrid_ep_buffer is None: num_tokens, hidden_dim = x.shape[-2:] - # --- Hardware Limit Gaurdrail --- + # --- Hardware Limit Guardrail --- # DeepEP calculates tx_depth = 3 * num_tokens + 1. - # # InfiniBand strictly asserts tx_depth < 65536. + # InfiniBand strictly asserts tx_depth < 65536. tx_depth = 3 * num_tokens + 1 if tx_depth >= 65536: raise ValueError( f"HybridEP RDMA Queue Pair depth ({tx_depth}) exceeds the InfiniBand " f"hardware limit of 65535. This occurs because the total tokens per rank " - f"({num_tokens}) is too high. Please reduce sequence length or micro-batch size, " + f"({num_tokens}) too high. Reduce sequence length or micro-batch size, " f"or increase Tensor Parallelism (TP) / Context Parallelism (CP) to reduce " f"the number of tokens processed per rank." ) - fp8_dispatch = False # Currently, we do not support fp8 dispatch init_hybrid_ep_buffer( group,