Skip to content

[https://nvbugs/6430702][perf] Restore NVLink one-sided A2A fast path #16200

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

Merged
chienchunhung merged 2 commits into from
Jul 14, 2026
Merged

[https://nvbugs/6430702][perf] Restore NVLink one-sided A2A fast path #16200

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b76a9e3
[https://nvbugs/6430702][perf] Restore NVLink one-sided A2A fast path
chienchunhung Jul 9, 2026
79c9cbc
[https://nvbugs/6430702][fix] Address A2A rank-mask review feedback
chienchunhung Jul 10, 2026
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
119 changes: 71 additions & 48 deletions cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -392,7 +392,7 @@ __global__ void moeA2APrepareDispatchKernel(
// Dispatch Kernels
// ============================================================================

template <typename ThreadingPolicy, int TOP_K, bool ENABLE_EPLB>
template <typename ThreadingPolicy, int TOP_K, bool ENABLE_EPLB, bool ENABLE_RANK_MASK>
__global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [local_num_tokens, TOP_K]
const DispatchKernelPointers ptrs, // Struct containing all kernel pointers
int num_payloads, // Number of payloads
Expand Down Expand Up @@ -440,15 +440,17 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [
// Supports the non-divisible case where num_experts % ep_size != 0.
int target_rank = compute_target_rank_id(expert_id, ep_base, ep_remainder);

// Skip duplicates AND dead ranks: both produce the same -1 sentinel that combine
// checks via topk_send_indices[k] < 0. A token whose only target is dead is dropped
// from this collective; higher-layer logic (EPLB redistribution) is responsible
// for re-routing such tokens on subsequent iterations.
int const mask_word = target_rank >> 6;
uint64_t const mask_bit = 1ULL << (target_rank & 63);
bool const target_already_copied = already_copied[mask_word] & mask_bit;
bool const target_dead = !is_rank_active(ptrs.active_rank_mask, target_rank);
if (target_already_copied || target_dead)
bool const target_already_copied = (already_copied[mask_word] & mask_bit) != 0;
bool skip_target = target_already_copied;
if constexpr (ENABLE_RANK_MASK)
{
// This is a fail-closed safety guard until post-commit routing is enforced end to end.
// A masked route is not valid model output; the failed execution epoch must be discarded.
skip_target = skip_target || !is_rank_active(ptrs.active_rank_mask, target_rank);
}
if (skip_target)
{
if (thread_idx == 0)
{
Expand Down Expand Up @@ -532,8 +534,11 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [
#pragma unroll 1 // No unroll as one iter is typically enough
for (int target_rank = lane_id; target_rank < ep_size; target_rank += warpSize)
{
if (!is_rank_active(ptrs.active_rank_mask, target_rank))
continue;
if constexpr (ENABLE_RANK_MASK)
{
if (!is_rank_active(ptrs.active_rank_mask, target_rank))
continue;
}
int send_count = ptrs.send_counters[target_rank];
ptrs.recv_counters[target_rank][rank_id] = send_count;
}
Expand All @@ -545,8 +550,11 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [
#pragma unroll 1
for (int target_rank = 0; target_rank < ep_size; ++target_rank)
{
if (!is_rank_active(ptrs.active_rank_mask, target_rank))
continue;
if constexpr (ENABLE_RANK_MASK)
{
if (!is_rank_active(ptrs.active_rank_mask, target_rank))
continue;
}
int* target_stats = ptrs.eplb_gathered_stats[target_rank];
for (int expert_id = lane_id; expert_id < eplb_stats_num_experts; expert_id += warpSize)
{
Expand All @@ -570,8 +578,11 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [
#pragma unroll 1 // No unroll as one iter is typically enough
for (int target_rank = lane_id; target_rank < ep_size; target_rank += warpSize)
{
if (!is_rank_active(ptrs.active_rank_mask, target_rank))
continue;
if constexpr (ENABLE_RANK_MASK)
{
if (!is_rank_active(ptrs.active_rank_mask, target_rank))
continue;
}
uint32_t* flag_addr = &ptrs.completion_flags[target_rank][rank_id];
asm volatile("st.relaxed.sys.u32 [%0], %1;" ::"l"(flag_addr), "r"(expected_value));

Expand All @@ -586,8 +597,11 @@ __global__ void moeA2ADispatchKernel(int32_t const* token_selected_experts, // [
#pragma unroll 1 // No unroll
for (int peer_rank = lane_id; peer_rank < ep_size; peer_rank += warpSize)
{
if (!is_rank_active(ptrs.active_rank_mask, peer_rank))
continue;
if constexpr (ENABLE_RANK_MASK)
{
if (!is_rank_active(ptrs.active_rank_mask, peer_rank))
continue;
}
bool flag_set = false;
auto s = clock64();
do
Expand Down Expand Up @@ -636,10 +650,11 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params)
TLLM_CHECK(params.ep_rank >= 0 && params.ep_rank < params.ep_size);
TLLM_CHECK(params.local_num_tokens >= 0);
TLLM_CHECK(params.num_payloads > 0 && params.num_payloads <= kMaxPayloads);
// The local rank must always be marked active in its own view of the mask;
// otherwise the kernel itself would be running on a "dead" rank.
TLLM_CHECK_WITH_INFO((params.active_rank_mask[params.ep_rank >> 6] >> (params.ep_rank & 63)) & 1ULL,
"active_rank_mask must mark the local ep_rank (%d) as active", params.ep_rank);
if (params.enable_rank_mask)
{
TLLM_CHECK_WITH_INFO((params.active_rank_mask[params.ep_rank >> 6] >> (params.ep_rank & 63)) & 1ULL,
"active_rank_mask must mark the local ep_rank (%d) as active", params.ep_rank);
}

// Prepare kernel pointers struct
DispatchKernelPointers kernel_ptrs = {};
Expand Down Expand Up @@ -693,12 +708,15 @@ void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params)
grid_size = 1;
}
int shared_bytes = 2 * params.top_k * (int) sizeof(int);
SWITCH_BOOL(params.enable_eplb, EPLB_STATS, SWITCH_TOP_K(params.top_k, TOP_K, {
auto kernel_fn = moeA2ADispatchKernel<BlockPolicy, TOP_K, EPLB_STATS>;
launchWithPdlWhenEnabled("moeA2ADispatchKernel", kernel_fn, grid_size, kBlockSize, shared_bytes, params.stream,
params.token_selected_experts, kernel_ptrs, params.num_payloads, params.max_tokens_per_rank,
params.local_num_tokens, params.ep_rank, params.ep_size, params.num_experts, params.eplb_stats_num_experts);
}))
SWITCH_BOOL(params.enable_rank_mask, ENABLE_RANK_MASK, {SWITCH_BOOL(params.enable_eplb, EPLB_STATS, {
SWITCH_TOP_K(params.top_k, TOP_K, {
auto kernel_fn = moeA2ADispatchKernel<BlockPolicy, TOP_K, EPLB_STATS, ENABLE_RANK_MASK>;
launchWithPdlWhenEnabled("moeA2ADispatchKernel", kernel_fn, grid_size, kBlockSize, shared_bytes,
params.stream, params.token_selected_experts, kernel_ptrs, params.num_payloads,
params.max_tokens_per_rank, params.local_num_tokens, params.ep_rank, params.ep_size, params.num_experts,
params.eplb_stats_num_experts);
});
})})
}

// ============================================================================
Expand Down Expand Up @@ -741,7 +759,7 @@ __device__ void vectorized_combine_impl(T* dst_typed_base, int size_per_token, i
{
int target_rank = ptrs.topk_target_ranks[local_token_idx * TOP_K + k];
int dst_idx = ptrs.topk_send_indices[local_token_idx * TOP_K + k];
if (dst_idx < 0 || !is_rank_active(ptrs.active_rank_mask, target_rank))
if (dst_idx < 0)
{
acc[k].fill(0.0f);
continue;
Expand All @@ -766,12 +784,8 @@ __device__ void vectorized_combine_impl(T* dst_typed_base, int size_per_token, i
#pragma unroll
for (int k = 0; k < TOP_K; ++k)
{
int target_rank = ptrs.topk_target_ranks[local_token_idx * TOP_K + k];
int dst_idx = ptrs.topk_send_indices[local_token_idx * TOP_K + k];
if (dst_idx < 0 || !is_rank_active(ptrs.active_rank_mask, target_rank))
{
if (ptrs.topk_send_indices[local_token_idx * TOP_K + k] < 0)
continue; // acc[k] already holds 0.0f from fill() above
}
#pragma unroll
for (int j = elems_per_vec - 1; j >= 0; --j)
acc[k][j] = static_cast<float>(reinterpret_cast<InT const*>(&acc[k])[j]);
Expand Down Expand Up @@ -1158,7 +1172,7 @@ __global__ void moeA2APrepareCombineKernel(uint8_t* recv_buffer_bytes, void cons
// Generic Combine Kernel Implementation (Templated by data type)
// ============================================================================

template <typename T, typename ThreadingPolicy, int TOP_K>
template <typename T, typename ThreadingPolicy, int TOP_K, bool ENABLE_RANK_MASK>
__global__ void moeA2ACombineKernel(
const CombineKernelPointers ptrs, // Combine-specific struct, src_data_ptrs[0] is output
int max_tokens_per_rank, int elements_per_token, int local_num_tokens, int rank_id, int ep_size,
Expand Down Expand Up @@ -1203,8 +1217,11 @@ __global__ void moeA2ACombineKernel(
#pragma unroll 1 // No unroll
for (int peer_rank = lane_id; peer_rank < ep_size; peer_rank += warpSize)
{
if (!is_rank_active(ptrs.active_rank_mask, peer_rank))
continue;
if constexpr (ENABLE_RANK_MASK)
{
if (!is_rank_active(ptrs.active_rank_mask, peer_rank))
continue;
}
uint32_t* flag_addr = &ptrs.completion_flags[peer_rank][rank_id];
asm volatile("st.relaxed.sys.u32 [%0], %1;" ::"l"(flag_addr), "r"(expected_value));
#if ENABLE_DEBUG_PRINT
Expand All @@ -1219,8 +1236,11 @@ __global__ void moeA2ACombineKernel(
#pragma unroll 1 // No unroll
for (int peer_rank = lane_id; peer_rank < ep_size; peer_rank += warpSize)
{
if (!is_rank_active(ptrs.active_rank_mask, peer_rank))
continue;
if constexpr (ENABLE_RANK_MASK)
{
if (!is_rank_active(ptrs.active_rank_mask, peer_rank))
continue;
}
bool flag_set = false;
auto s = clock64();
do
Expand Down Expand Up @@ -1327,10 +1347,11 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params)
TLLM_CHECK(params.ep_rank >= 0 && params.ep_rank < params.ep_size);
TLLM_CHECK(params.local_num_tokens >= 0);
TLLM_CHECK(params.elements_per_token > 0);
// The local rank must always be marked active in its own view of the mask;
// otherwise the kernel itself would be running on a "dead" rank.
TLLM_CHECK_WITH_INFO((params.active_rank_mask[params.ep_rank >> 6] >> (params.ep_rank & 63)) & 1ULL,
"active_rank_mask must mark the local ep_rank (%d) as active", params.ep_rank);
if (params.enable_rank_mask)
{
TLLM_CHECK_WITH_INFO((params.active_rank_mask[params.ep_rank >> 6] >> (params.ep_rank & 63)) & 1ULL,
"active_rank_mask must mark the local ep_rank (%d) as active", params.ep_rank);
}

// Configure kernel launch (one block per token).
int const kBlockSize = tensorrt_llm::common::getEnvMoeA2ACombineBlockSize();
Expand Down Expand Up @@ -1385,14 +1406,16 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params)
auto const effective_dtype = params.use_low_precision ? nvinfer1::DataType::kFP8 : params.dtype;

// Launch appropriate kernel with compact macros
SWITCH_DTYPE(effective_dtype, TKernelType, {
SWITCH_TOP_K(params.top_k, TOP_K, {
auto kernel_fn = moeA2ACombineKernel<TKernelType, BlockPolicy, TOP_K>;
launchWithPdlWhenEnabled("moeA2ACombineKernel", kernel_fn, grid, kBlockSize, 0, params.stream, kernel_ptrs,
params.max_tokens_per_rank, params.elements_per_token, params.local_num_tokens, params.ep_rank,
params.ep_size, stride_per_token);
SWITCH_BOOL(params.enable_rank_mask, ENABLE_RANK_MASK, {
SWITCH_DTYPE(effective_dtype, TKernelType, {
SWITCH_TOP_K(params.top_k, TOP_K, {
auto kernel_fn = moeA2ACombineKernel<TKernelType, BlockPolicy, TOP_K, ENABLE_RANK_MASK>;
launchWithPdlWhenEnabled("moeA2ACombineKernel", kernel_fn, grid, kBlockSize, 0, params.stream,
kernel_ptrs, params.max_tokens_per_rank, params.elements_per_token, params.local_num_tokens,
params.ep_rank, params.ep_size, stride_per_token);
});
});
});
})
}

// Kernel to sanitize expert ids for invalid tokens
Expand Down
42 changes: 25 additions & 17 deletions cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,17 +62,16 @@ struct DispatchKernelPointers
int* local_token_counter; // Atomic counter for completed tokens

// Top-K compact routing info per local token (size: [local_num_tokens, top_k])
int* topk_target_ranks; // target rank per k, -1 for duplicates
int* topk_send_indices; // dst index per k, -1 for duplicates
int* topk_target_ranks; // target rank per k, -1 for invalid or duplicate routes
int* topk_send_indices; // dst index per k, -1 for invalid or duplicate routes

// Optional: Statistics for EPLB
int const* eplb_local_stats; // [eplb_stats_num_experts]
int* eplb_gathered_stats[kMaxRanks]; // [ep_size, eplb_stats_num_experts] per rank

// Active-rank bitmask: bit i set => rank i is alive and participates in this collective.
// Word 0 covers ranks 0..63; word 1 covers ranks 64..127. Tokens routed to a masked
// rank are dropped (topk_*[k] = -1); flag writes/waits to/from masked peers are skipped.
// The local rank's own bit must always be set; this is checked at launch time.
// Active-rank bitmask: bit i set => rank i participates in this collective.
// Word 0 covers ranks 0..63; word 1 covers ranks 64..127. The masked kernel
// rejects inactive route targets and skips their peer counters, stats, and flags.
uint64_t active_rank_mask[kRankMaskWords];
};

Expand All @@ -89,12 +88,11 @@ struct CombineKernelPointers
uint32_t* flag_val; // The value of the flag for this round (stored on the local rank)

// Top-K compact routing info per local token (size: [local_num_tokens, top_k])
int const* topk_target_ranks; // target rank per k, -1 for duplicates
int const* topk_send_indices; // dst index per k, -1 for duplicates
int const* topk_target_ranks; // target rank per k, -1 for invalid or duplicate routes
int const* topk_send_indices; // dst index per k, -1 for invalid or duplicate routes

// Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Combine skips flag
// writes/waits to/from masked peers and also skips per-token accumulation for ranks that
// become inactive between dispatch and combine.
// Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Combine skips
// completion flag writes/waits to/from inactive peers.
uint64_t active_rank_mask[kRankMaskWords];
};

Expand Down Expand Up @@ -139,9 +137,14 @@ struct MoeA2ADispatchParams
int const* eplb_local_stats; // [eplb_stats_num_experts]
int* eplb_gathered_stats[kMaxRanks]; // [ep_size, eplb_stats_num_experts] per rank

// Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. The launch function
// copies these words into the kernel pointers struct. Defaults to all-ones for
// backwards-compatible "no masking" behavior.
// Whether to instantiate a kernel with active-rank checks.
// This is a launch-lifetime mode, independent of future execution-abort handling.
bool enable_rank_mask{false};

// Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Used only when
Comment thread
chienchunhung marked this conversation as resolved.
// enable_rank_mask is true; defaults to all-ones for backwards-compatible behavior.
// The mask is copied by value into kernel arguments. Rank-mask mode must reject
// CUDA graph replay until generation-scoped invalidation and recapture are available.
uint64_t active_rank_mask[kRankMaskWords] = {~uint64_t{0}, ~uint64_t{0}};

// CUDA stream
Expand Down Expand Up @@ -189,9 +192,14 @@ struct MoeA2ACombineParams
// rank has signaled the target rank
void const* recv_buffers[kMaxRanks]; // Per-rank receive buffers (only for single payload)

// Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. The launch function
// copies these words into the kernel pointers struct. Defaults to all-ones for
// backwards-compatible "no masking" behavior.
// Whether to instantiate a kernel with active-rank checks in peer synchronization.
// This is a launch-lifetime mode, independent of future execution-abort handling.
bool enable_rank_mask{false};

// Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Used only when
// enable_rank_mask is true; defaults to all-ones for backwards-compatible behavior.
// The mask is copied by value into kernel arguments. Rank-mask mode must reject
// CUDA graph replay until generation-scoped invalidation and recapture are available.
uint64_t active_rank_mask[kRankMaskWords] = {~uint64_t{0}, ~uint64_t{0}};

// CUDA stream
Expand Down
Loading
Loading