perf: reuse HybridEP routing allgather buffer

[akoumpa]

Summary

Reuse HybridEP's preallocated allgather output buffer for the NCCL/inter-node routing-map allgather instead of allocating a temporary CUDA tensor during every dispatch.

The routing map is consumed by metadata preprocessing on the same stream before the next dispatch can reuse the buffer, so its lifetime does not overlap the following collective.

Memory impact

The removed temporary allocation contains the globally gathered routing payload on every rank.

For sparse Boolean routing, the saved allocation per rank is:

tokens_per_rank x world_size x num_experts bytes

For dense int16 routing, it is:

tokens_per_rank x world_size x topk x 2 bytes

For the stress configuration used during validation (8192 tokens/rank, 8 ranks, 256 experts, and top-k 36), this avoids:

• Sparse routing: 16 MiB per rank, or 128 MiB across the process group.
• Dense routing: 4.5 MiB per rank, or 36 MiB across the process group.

Implementation

• Wrap CustomAllgather::get_output_buffer() as the NCCL allgather output tensor.
• Preserve the existing NCCL behavior for inter-node execution, disabled custom allgather, and unaligned payloads.
• Preserve the upstream dense-routing behavior, including the dynamic bool/int16 dtype, routing shape, and byte views used for NCCL.
• Add a distributed output and gradient regression test.
• Assert the requested node topology so multi-node validation cannot silently execute as a single-node test.

Testing

Validation was performed on H100 GPUs using a multi-node DOCA/RDMA build in both of these topologies:

• One node with eight local ranks.
• Two nodes with four local ranks per node, for eight global ranks.

Output and gradient regression

Added tests/test_allgather_grad_validation.py, which compares the reused-buffer path with the existing allgather behavior.

The test covers the Cartesian product of:

• Dense expert indices and sparse Boolean routing maps.
• Routing with and without probabilities.
• Three independently executed dispatch/combine cycles.

This produces 12 cases per topology. Each case validates:

• Dispatched hidden states and routing probabilities.
• Per-expert token counts and dispatch-handle metadata.
• Forward combine hidden-state and probability outputs.
• Hidden-state and probability gradients for the backward-of-combine communication.
• Hidden-state and probability gradients for the backward-of-dispatch communication.
• Tensor shapes, dtypes, and optional None behavior.

All 12 cases passed on both the single-node and two-node topologies.

Full HybridEP reference validation

Ran tests/test_hybrid_ep.py with 512 tokens/rank, eight local experts, top-k 8, hidden dimension 1024, and padding multiple 32.

The following passed for both BF16 and FP8 on both topologies:

• Standard dispatch and combine.
• Sparse and dense routing.
• Routing with and without probabilities.
• Dispatch with permutation and combine with unpermutation.
• Non-fused and fused permute/dispatch execution.
• Torch reference comparisons.
• Intra-node NVLink and inter-node RDMA execution.

The API-level and kernel-level benchmark loops also completed successfully.

CUDA graph and nonblocking validation

Ran tests/test_graphed_hybrid_ep.py across:

• BF16 and FP8.
• With and without routing probabilities.
• Fused and non-fused execution.
• CUDA graph capture/replay and nonblocking dispatch/combine behavior.

All combinations passed on every rank in both topologies.

Harry Zhou's standalone allgather stress test

Thanks to Harry Zhou for providing test_allgather.py, which was used as an additional correctness and stress test.

It was run in both single-node and two-node configurations with its original stress dimensions:

tokens per rank  = 8192
hidden dimension = 7168
global experts   = 256
local experts    = 32
top-k            = 36
global ranks     = 8

For both dense-index and sparse-Boolean routing:

• Dispatched hidden tokens matched bitwise.
• Dispatched probabilities matched bitwise.
• The custom, NCCL, and legacy benchmark loops completed successfully.
• Repeated execution completed without exposing buffer-lifetime or reuse issues.

Thanks again to Harry for the focused test coverage and stress configuration.

Static checks

The new regression test also passed Ruff linting, Ruff formatting, Python bytecode compilation, and repository pre-commit checks.