[Bug] HybridEP dispatcher passes incorrect max_num_of_tokens_per_rank to DeepEP, causing RDMA QP assertion failure

Summary

When using the HybridEP backend (--moe_flex_dispatcher_backend hybridep) for MoE expert-parallel training across multiple nodes, HybridEPDispatch.forward() in fused_a2a.py incorrectly uses the total number of tokens in a micro-batch (seq_length × micro_batch_size) as the max_num_of_tokens_per_rank parameter for HybridEPBuffer. This causes the RDMA Queue Pair send-queue depth (tx_depth) to exceed the hardware limit of 65535, triggering an assertion failure in DeepEP's internode communication initialization.

Environment

Megatron-LM: latest (hybrid-ep branch of DeepEP integrated)
DeepEP: v1.2.1 (hybrid-ep branch)
Model: Qwen3-30B-A3B (MoE, 128 experts)
Hardware: 2 nodes × 8 GPUs, InfiniBand RDMA interconnect
Training config: --max_length 8192 --micro_batch_size 8 --packing true --expert_model_parallel_size 16 --moe_flex_dispatcher_backend hybridep

Error

python: /path/to/DeepEP/csrc/hybrid_ep/buffer/internode.cu:167:
void setup_qp_init_attr(..., int): Assertion `tx_depth > 0 && tx_depth < 65536' failed.

All ranks crash with SIGABRT (signal 6) during HybridEP buffer initialization.

Root Cause

The call chain

MoEFlexTokenDispatcher.dispatch_preprocess() (token_dispatcher.py:1438) reshapes hidden_states from [seq_length, batch_size, hidden_size] to [seq_length * batch_size, hidden_size]:
```
hidden_states = hidden_states.view(-1, self.hidden_shape[-1])
```

HybridEPDispatch.forward() (fused_a2a.py:354-359) extracts the first dimension of the already-flattened tensor and uses it as seq_len:

if _hybrid_ep_buffer is None:
    seq_len, hidden_dim = x.shape[-2:]  # x is [seq_len * batch_size, hidden_dim]
    init_hybrid_ep_buffer(group, hidden_dim, seq_len, ...)  # seq_len is actually num_total_tokens

init_hybrid_ep_buffer() (fused_a2a.py:316) passes this value directly as max_num_of_tokens_per_rank:

_hybrid_ep_buffer = HybridEPBuffer(
    group=group,
    hidden_dim=hidden_dim,
    max_num_of_tokens_per_rank=seq_len,  # <-- This is seq_length * micro_batch_size, not per-rank tokens
    ...
)

DeepEP internode.cu uses max_num_of_tokens_per_rank to compute the RDMA QP send-queue depth:

// internode.cu:430 (dispatch)
setup_qp_init_attr(..., 3 * buffer_config.max_num_of_tokens_per_rank + 1);
// internode.cu:587 (combine)
setup_qp_init_attr(..., 2 * buffer_config.max_num_of_tokens_per_rank + 1);

The assertion enforces the IB hardware limit:

assert(tx_depth > 0 && tx_depth < 65536);

Concrete example

With max_length=8192, micro_batch_size=8, packing=true:

Parameter	Value
`x.shape[0]` (after flatten)	`8192 × 8 = 65536`
`max_num_of_tokens_per_rank` passed to DeepEP	`65536`
dispatch `tx_depth`	`3 × 65536 + 1 = 196609`
Hardware limit	`< 65536`

The dispatch tx_depth exceeds the limit by 3×.

Why single-node works but multi-node fails

The RDMA QP initialization (and thus the tx_depth assertion) only runs when num_of_nodes > 1. Single-node setups use NVLink-only communication and never hit this code path.

Steps to Reproduce

Configure a multi-node MoE training with HybridEP:

--expert_model_parallel_size 16 \
--moe_token_dispatcher_type flex \
--moe_flex_dispatcher_backend hybridep \
--micro_batch_size 8 \
--max_length 8192 \
--packing true

Run training across 2+ nodes with RDMA/InfiniBand.
Training crashes immediately during HybridEP buffer initialization with:
```
Assertion `tx_depth > 0 && tx_depth < 65536' failed.
```

Note: The issue is triggered when 3 × seq_length × micro_batch_size + 1 > 65535, i.e., seq_length × micro_batch_size > 21845. Common configurations like 8192 × 4 = 32768 or 4096 × 8 = 32768 will hit this.

Affected Code

megatron/core/transformer/moe/fused_a2a.py — HybridEPDispatch.forward() (line 354-359) and init_hybrid_ep_buffer() (line 316)

Additional context

Add any other context about the problem here.

Hey, i would like to take this issue could you assign me ?

@Shreyas-S-809

Hey @Shreyas-S-809, I assigned you to the issue. Feel free to just open a PR when you are done.

@janEbert

Hey @Shreyas-S-809, I assigned you to the issue. Feel free to just open a PR when you are done.

Hey @janEbert ,

Thanks for the assignment! I've just opened a PR with the proposed fix. #4094

As a quick heads-up, I don't have access to a multi-node InfiniBand cluster to test the RDMA queue initialization locally. I'll be relying on the CI/CD pipeline to validate the multi-node execution path.

Please let me know if the tests catch anything or if any adjustments are needed during the review process. Thanks!

Thank you also a lot for the initiative! Although I must say it's a bit unexpected that you'd wish to work on an issue that you can't test during development. 😅

When reviewing the PR, I got a bit curious and actually looked at the related DeepEP code myself. It seems like what's passed is actually correct. The seq_len variable in HybridEPDispatch.forward and init_hybrid_ep_buffer, which is badly named and should probably be called num_tokens, seems like it should be seq_len * batch_size like it currently is. DeepEP never explicitly names a batch size anywhere when using it (e.g., in its tests), but the code only makes sense assuming that batch_size is folded in, which does match what num_tokens usually refers to (seq_len * batch_size).

Maybe @NVIDIA/mixture-of-experts-adlr or @NVIDIA/mixture-of-experts-devtech could confirm (and take a look at the PR, too).

However, there's still the issue with the RDMA QP assertion. This could probably be worked around by setting a higher TP or CP degree to reduce the number of tokens per rank.

@janEbert

Thank you also a lot for the initiative! Although I must say it's a bit unexpected that you'd wish to work on an issue that you can't test during development. 😅

When reviewing the PR, I got a bit curious and actually looked at the related DeepEP code myself. It seems like what's passed is actually correct. The seq_len variable in HybridEPDispatch.forward and init_hybrid_ep_buffer, which is badly named and should probably be called num_tokens, seems like it should be seq_len * batch_size like it currently is. DeepEP never explicitly names a batch size anywhere when using it (e.g., in its tests), but the code only makes sense assuming that batch_size is folded in, which does match what num_tokens usually refers to (seq_len * batch_size).

Maybe @NVIDIA/mixture-of-experts-adlr or @NVIDIA/mixture-of-experts-devtech could confirm (and take a look at the PR, too).

However, there's still the issue with the RDMA QP assertion. This could probably be worked around by setting a higher TP or CP degree to reduce the number of tokens per rank.

Hey @janEbert thanks for looking into this so quickly!

That makes complete sense. If DeepEP inherently expects seq_len * batch_size folded together, then the original logic of passing x.shape[-2] (total tokens) was architecturally correct, and changing it to an expected per-rank capacity breaks that assumption.

If we agree that this is a hardware constraint rather than a routing bug, and the correct workaround is for users to increase their TP/CP degrees to reduce tokens per rank, do you think we should add a Python-side guardrail here?

Currently, hitting this limit triggers a hard SIGABRT from the C++ backend across all ranks, which is tough to debug for end users. Would a PR be welcome that simply checks (3 * num_tokens + 1) >= 65536 in HybridEPDispatch.forward and raises a clean ValueError? We could explicitly advise the user in the error message to increase their Tensor or Context Parallelism degrees to avoid the InfiniBand RDMA Queue Pair limit.

Let me know what you and the @NVIDIA/mixture-of-experts-adlr team think! Happy to either close my current PR or pivot it into adding that Python-side assertion + renaming seq_len to num_tokens for clarity.

Raising a cleaner error message with some tips on avoiding the error is a nice suggestion! But for now, let's wait for feedback from the MoE teams before continuing. :)

Raising a cleaner error message with some tips on avoiding the error is a nice suggestion! But for now, let's wait for feedback from the MoE teams before continuing. :)

Understood. While that's on hold, are there any other open issues you'd like me to tackle? Just let me know where I can be most helpful!. Thank you.

@Shreyas-S-809

Hey @Shreyas-S-809, thanks a lot for the understanding and continued motivation. Issues labeled as "community-request"s are usually a good place to look for open issues. Best pick something that you can test locally before PRing. ;)

@janEbert

Hey @Shreyas-S-809, sorry that we haven't followed up yet. I think you should just go ahead with the suggestion you made to raise a nice and clean ValueError.

Thanks a lot!

Hey @janEbert , I have updated the PR #4094 . I reverted the tensor shape modifications, renamed seq_len to num_tokens for clarity, and added the ValueError guardrail to gracefully catch the InfiniBand limit before the C++ backend crashes. Let me know if the error message wording looks good to the team!

@janEbert

Hey @janEbert , Just waiting for your response. Thank you.

Hey, let's keep the PR-related discussion in the PR, please. :)

@janEbert

Hey @janEbert , thanks for the ping!

That makes complete sense. If DeepEP is doing more dynamic buffer handling under the hood, hardcoding that strict 3x + 1 bound in Python was definitely risking false positives on valid configs. Great catch by the Megatron-Bridge tests!

I completely agree with your second approach. Filtering and wrapping the actual error raised by DeepEP seems much cleaner and more robust than trying to mirror their internal C++ math perfectly in our Python layer.

Thanks for linking my original issue and taking the time to explain the revert. Let me know if you need an extra set of eyes on the new implementation!

@janEbert

hello @janEbert @Shreyas-S-809
I noticed that a defensive check for cross-node communication was recently added to the hybrid EP dispatcher, but this check does not apply to pure NVL domains, especially NVL72. Moreover, it incorrectly blocks some larger-MBS intra-NVL-domain communications that should be allowed to run.
If we want to add this check in MCore, we need a concrete and reliable logic to determine whether the current EP communication will actually go through the NIC, which doesn't seem straightforward.
Therefore, I would suggest that this kind of check be proposed in a PR targeting the hybrid EP itself rather than MCore.
This PR is currently blocking some experiments from running. Do we have any way to resolve this?

@Autumn1998

Hey @Autumn1998, thanks for the detailed context!
That makes total sense—I completely overlooked how this would affect pure NVL domains like the NVL72 where the InfiniBand NIC limits wouldn't apply. I agree that trying to infer the underlying hardware topology (NVLink vs. NIC) from the Python layer in MCore is way too fragile and overly conservative.

@janEbert already went ahead and reverted the PR in #4718 to unblock those experiments, which was definitely the right call. Filtering it at the DeepEP level is the much better approach.

I really appreciate you taking the time to explain the hardware routing nuances here!

@Autumn1998

@Autumn1998 please try #4846

Navigation Menu

Repository navigation

[Bug] HybridEP dispatcher passes incorrect max_num_of_tokens_per_rank to DeepEP, causing RDMA QP assertion failure #3999

Summary

Environment

Error

Root Cause

The call chain

Concrete example

Why single-node works but multi-node fails

Steps to Reproduce

Affected Code

6 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[Bug] HybridEP dispatcher passes incorrect max_num_of_tokens_per_rank to DeepEP, causing RDMA QP assertion failure #3999

Description

Summary

Environment

Error

Root Cause

The call chain

Concrete example

Why single-node works but multi-node fails

Steps to Reproduce

Affected Code

Activity

Shreyas-S-809 commented on Mar 29, 2026

janEbert commented on Mar 31, 2026

Shreyas-S-809 commented on Apr 1, 2026

janEbert commented on Apr 2, 2026

Shreyas-S-809 commented on Apr 2, 2026

janEbert commented on Apr 2, 2026

Shreyas-S-809 commented on Apr 3, 2026

janEbert commented on Apr 7, 2026

6 remaining items

Shreyas-S-809 commented on Apr 27, 2026

Shreyas-S-809 commented on May 2, 2026

janEbert commented on May 4, 2026

Shreyas-S-809 commented on May 11, 2026

Autumn1998 commented on May 14, 2026

Shreyas-S-809 commented on May 14, 2026

janEbert commented on May 18, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions