From 149ff76f7079d8c38c784e2e8361258476e4c32d Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Tue, 12 May 2026 14:08:00 +0800 Subject: [PATCH 1/8] Improve test_hybrid_ep.py: check_bitwise diagnostics, probs=False benchmarks, --only-bf16, permute/unpermute kernel profiling --- tests/test_hybrid_ep.py | 89 ++++++++++++++++++++++++++++++++--------- 1 file changed, 71 insertions(+), 18 deletions(-) diff --git a/tests/test_hybrid_ep.py b/tests/test_hybrid_ep.py index 0fd7334a2..90f91a3ce 100644 --- a/tests/test_hybrid_ep.py +++ b/tests/test_hybrid_ep.py @@ -200,18 +200,30 @@ def test_hybrid_ep_correctness(buffer: deep_ep.HybridEPBuffer, ref: TorchRef, us enable_permute=True, ) - assert bitwise_equal(dispatched_hidden_ref, dispatched_hidden), \ - f"Dispatch hidden mismatch (with_probs={with_probs}, fuse={fuse_permute_dispatch})" - if dispatched_probs is not None and dispatched_probs_ref is not None: - assert bitwise_equal(dispatched_probs_ref, dispatched_probs), \ - f"Dispatch probs mismatch (with_probs={with_probs}, fuse={fuse_permute_dispatch})" - if ( - dispatched_scaling_factor is not None - and dispatched_scaling_factor_ref is not None - ): - assert bitwise_equal( - dispatched_scaling_factor_ref, dispatched_scaling_factor - ), f"Dispatch scaling_factor mismatch (with_probs={with_probs}, fuse={fuse_permute_dispatch})" + def check_bitwise(name, ref, test): + if ref is None or test is None: + return + if not bitwise_equal(ref, test): + total = ref.numel() + mismatch = (ref.contiguous().view(torch.uint8) != test.contiguous().view(torch.uint8)).sum().item() + # Count element-level mismatches (not byte-level) + elem_mismatch = (ref != test).sum().item() + pct = 100.0 * elem_mismatch / max(ref.numel(), 1) + msg = (f"{name} mismatch (with_probs={with_probs}, fuse={fuse_permute_dispatch}): " + f"{elem_mismatch}/{ref.numel()} elements ({pct:.2f}%), " + f"{mismatch} bytes differ, shape={list(ref.shape)}") + # Print first few mismatching positions + flat_ref = ref.contiguous().view(-1) + flat_test = test.contiguous().view(-1) + diff_idx = (flat_ref != flat_test).nonzero(as_tuple=True)[0][:5] + for idx in diff_idx: + i = idx.item() + msg += f"\n [{i}]: ref={flat_ref[i].item()}, got={flat_test[i].item()}" + assert False, msg + + check_bitwise("Dispatch hidden", dispatched_hidden_ref, dispatched_hidden) + check_bitwise("Dispatch probs", dispatched_probs_ref, dispatched_probs) + check_bitwise("Dispatch scaling_factor", dispatched_scaling_factor_ref, dispatched_scaling_factor) # The combine only support bf16 dispatched_hidden = dispatched_hidden.to(torch.bfloat16) @@ -318,6 +330,11 @@ def test_hybrid_ep_benchmark(buffer: deep_ep.HybridEPBuffer, group: dist.Process 'topk_weights': topk_weights, 'num_of_experts': NUM_OF_EXPERTS, 'handle': handle} combine_args = {'hidden': dispatched_hidden_bf16, 'probs': dispatched_probs, 'handle': handle} + # Dispatch/combine with probs=False variants + dispatch_noprob_args = {'hidden': hidden, 'scaling_factor': scaling_factor, 'topk_idx': topk_idx, + 'topk_weights': None, 'num_of_experts': NUM_OF_EXPERTS, 'handle': handle} + combine_noprob_args = {'hidden': dispatched_hidden_bf16, 'probs': None, 'handle': handle} + # Permute (non-fused) dispatched_hidden_wp, dispatched_probs_wp, _, tpe_wp, handle_wp = ( buffer.dispatch_with_permute(hidden=hidden, scaling_factor=scaling_factor, @@ -365,11 +382,17 @@ def test_hybrid_ep_benchmark(buffer: deep_ep.HybridEPBuffer, group: dist.Process print(f' combine = fused_combine_unpermute_kernel + misc', flush=True) print(f' (misc = device_sync, update_flag, etc.)', flush=True) - # Non-permute + # Non-permute (probs=True) t = bench(lambda: buffer.dispatch(**dispatch_args))[0] - _report_bw(f'dispatch ({dtype_str})', t, nvl_dispatch_actual, 'nvl_recv_bytes', rdma_dispatch, 'rdma_send_bytes') + _report_bw(f'dispatch ({dtype_str}, probs=True)', t, nvl_dispatch_actual, 'nvl_recv_bytes', rdma_dispatch, 'rdma_send_bytes') t = bench(lambda: buffer.combine(**combine_args))[0] - _report_bw('combine', t, nvl_combine, 'combine_send_bytes', rdma_combine, 'rdma_recv_bytes') + _report_bw('combine (probs=True)', t, nvl_combine, 'combine_send_bytes', rdma_combine, 'rdma_recv_bytes') + + # Non-permute (probs=False) + t = bench(lambda: buffer.dispatch(**dispatch_noprob_args))[0] + _report_bw(f'dispatch ({dtype_str}, probs=False)', t, nvl_dispatch_actual, 'nvl_recv_bytes', rdma_dispatch, 'rdma_send_bytes') + t = bench(lambda: buffer.combine(**combine_noprob_args))[0] + _report_bw('combine (probs=False)', t, nvl_combine, 'combine_send_bytes', rdma_combine, 'rdma_recv_bytes') # Permute (non-fused) t = bench(lambda: buffer.dispatch_with_permute(**dispatch_wp_args))[0] @@ -391,12 +414,20 @@ def test_hybrid_ep_benchmark(buffer: deep_ep.HybridEPBuffer, group: dist.Process print(f' Non-fused: dispatch_kernel only | combine_kernel only', flush=True) print(f' Fused: fused_permute_dispatch_kernel only | fused_combine_unpermute_kernel only', flush=True) - # Non-fused kernel profiling + # Non-fused kernel profiling (probs=True) group.barrier() dispatch_t, combine_t = bench_kineto( lambda: (buffer.dispatch(**dispatch_args), buffer.combine(**combine_args)), kernel_names=('dispatch_kernel', 'combine_kernel'), barrier_comm_profiling=True, suppress_kineto_output=True) - _report_kineto(f'dispatch kernel ({dtype_str})', 'combine kernel', + _report_kineto(f'dispatch kernel ({dtype_str}, probs=True)', 'combine kernel (probs=True)', + dispatch_t, nvl_dispatch_actual, combine_t, nvl_combine, rdma_dispatch, rdma_combine) + + # Non-fused kernel profiling (probs=False) + group.barrier() + dispatch_t, combine_t = bench_kineto( + lambda: (buffer.dispatch(**dispatch_noprob_args), buffer.combine(**combine_noprob_args)), + kernel_names=('dispatch_kernel', 'combine_kernel'), barrier_comm_profiling=True, suppress_kineto_output=True) + _report_kineto(f'dispatch kernel ({dtype_str}, probs=False)', 'combine kernel (probs=False)', dispatch_t, nvl_dispatch_actual, combine_t, nvl_combine, rdma_dispatch, rdma_combine) # Fused kernel profiling @@ -406,6 +437,25 @@ def test_hybrid_ep_benchmark(buffer: deep_ep.HybridEPBuffer, group: dist.Process kernel_names=('dispatch_kernel', 'combine_kernel'), barrier_comm_profiling=True, suppress_kineto_output=True) _report_kineto(f'fused dispatch+permute kernel ({dtype_str})', 'fused combine+unpermute kernel', dispatch_t, nvl_dispatch_actual, combine_t, nvl_combine, rdma_dispatch, rdma_combine) + + # Non-fused permute/unpermute kernel profiling (isolate permute_kernel and unpermute_kernel times) + group.barrier() + dispatch_t, permute_t = bench_kineto( + lambda: buffer.dispatch_with_permute(**dispatch_wp_args), + kernel_names=('dispatch_kernel', 'permute_kernel'), barrier_comm_profiling=True, suppress_kineto_output=True) + combine_t, unpermute_t = bench_kineto( + lambda: buffer.combine_with_unpermute(**combine_wp_args), + kernel_names=('combine_kernel', 'unpermute_kernel'), barrier_comm_profiling=True, suppress_kineto_output=True) + d_times = _gather_times(dispatch_t) + p_times = _gather_times(permute_t) + c_times = _gather_times(combine_t) + u_times = _gather_times(unpermute_t) + if rank == 0: + fmt = lambda ts: f'avg={sum(ts)/len(ts)*1e6:.1f} us [min={min(ts)*1e6:.1f}, max={max(ts)*1e6:.1f}]' + print(f' {"dispatch_kernel (in dispatch+permute):":<{LOG_LABEL_WIDTH}} {fmt(d_times)}', flush=True) + print(f' {"permute_kernel:":<{LOG_LABEL_WIDTH}} {fmt(p_times)}', flush=True) + print(f' {"unpermute_kernel:":<{LOG_LABEL_WIDTH}} {fmt(u_times)}', flush=True) + print(f' {"combine_kernel (in combine+unpermute):":<{LOG_LABEL_WIDTH}} {fmt(c_times)}', flush=True) else: if torch.distributed.get_rank() == 0: torch.cuda.profiler.start() @@ -446,7 +496,8 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): stream = torch.cuda.Stream() with torch.cuda.stream(stream): - for use_fp8 in [False, True]: + fp8_modes = [False] if args.only_bf16 else [False, True] + for use_fp8 in fp8_modes: buffer = deep_ep.HybridEPBuffer( group=group, hidden_dim=HIDDEN_DIM, @@ -487,5 +538,7 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): help='Number of processes to spawn (default: 4)') parser.add_argument('--nsys-profile', action='store_true', default=False, help='benchmark with nsys profile or not (default: False)') + parser.add_argument('--only-bf16', action='store_true', default=False, + help='Skip FP8 tests, only run BF16 (default: False)') args = parser.parse_args() torch.multiprocessing.spawn(test_main, args=(args.num_processes, args), nprocs=args.num_processes) From 62f7232f39035199de59ed1f6744f076b1507397 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Tue, 12 May 2026 15:07:28 +0800 Subject: [PATCH 2/8] Optimize sparse prob communication --- csrc/hybrid_ep/backend/hybrid_ep_backend.cuh | 75 ++++++++++++++------ csrc/hybrid_ep/config.cuh | 14 +++- tests/test_hybrid_ep.py | 2 +- 3 files changed, 66 insertions(+), 25 deletions(-) diff --git a/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh b/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh index ad7f94f8a..d4b1c0698 100644 --- a/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh +++ b/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh @@ -399,18 +399,21 @@ struct combine_kernel_dynamic_shared_memory_buffer_tconsumer, 2nd for consumer->producer. Should be 8B alignment(natural alignment). alignas(8) uint64_t intra_node_mbarrier_G2S_buffer[NUM_OF_STAGES_G2S][2]; // Shared memory mbarrier that protect inter node red warp group G2S token entry. 1st for producer->consumer, 2nd for consumer->producer. Should be 8B alignment(natural alignment). @@ -446,12 +449,15 @@ struct combine_kernel_dynamic_shared_memory_buffer_tconsumer, 2nd for consumer->producer. Should be 8B alignment(natural alignment). alignas(8) uint64_t inter_node_mbarrier_G2S_buffer[NUM_OF_STAGES_G2S][2]; @@ -1637,13 +1643,22 @@ inline __device__ void S2G_warp_group_device_function(const int local_rank, (uint32_t)(HIDDEN_DIM * sizeof(TOKEN_DATA_TYPE))); // Store the prob from shared to remote global for FW dispatch. + // Only send the destination rank's E_per_rank slice (not the full E*R vector). + // The source SMEM prob buffer contains the full E*R probs (mostly zeros for sparse routing). + // Each rank's E_per_rank slice already has zeros at non-active expert positions. if constexpr(FORWARD_DISPATCH){ - float* remote_prob_addr = remote_expert_output_prob[remote_rank_id] + (output_buffer_index * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE)); + static_assert(NUM_OF_EXPERTS_PER_RANK * sizeof(float) >= 16, + "NUM_OF_EXPERTS_PER_RANK * sizeof(float) must be >= 16 for TMA minimum transfer size."); + static_assert((NUM_OF_EXPERTS_PER_RANK * sizeof(float)) % 16 == 0, + "NUM_OF_EXPERTS_PER_RANK * sizeof(float) must be 16B-aligned for TMA."); + float* remote_prob_addr = remote_expert_output_prob[remote_rank_id] + + output_buffer_index * static_cast(NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE) + + remote_rank_id * NUM_OF_EXPERTS_PER_RANK; cuda::ptx::cp_async_bulk(cuda::ptx::space_global, cuda::ptx::space_shared, reinterpret_cast(remote_prob_addr), - reinterpret_cast(&smem_buffer_ptr->intra_node_prob_buffer[stage][0]), - (uint32_t)((NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE) * sizeof(float))); + reinterpret_cast(&smem_buffer_ptr->intra_node_prob_buffer[stage][remote_rank_id * NUM_OF_EXPERTS_PER_RANK]), + (uint32_t)(NUM_OF_EXPERTS_PER_RANK * sizeof(float))); } @@ -2359,14 +2374,19 @@ inline __device__ void intra_node_G2S_warp_group_device_function(const int node_ total_tx_size += (uint32_t)(HIDDEN_DIM * sizeof(uint16_t)); if constexpr(BACKWARD_COMBINE){ + // Sparse prob optimization: only read the source rank's E_per_rank slice. + // The source rank (current_src_token_id) wrote its probs at offset + // current_src_token_id * E_per_rank within the E*R_per_node prob vector. cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, cuda::ptx::space_global, reinterpret_cast(&smem_buffer_ptr->intra_node_prob_G2S_buffer[token_stage][0]), - reinterpret_cast(remote_expert_input_prob[current_src_token_id] + (sparse_to_dense_map_value * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE))), - (uint32_t)((NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE) * sizeof(float)), + reinterpret_cast(remote_expert_input_prob[current_src_token_id] + (sparse_to_dense_map_value * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE)) + current_src_token_id * NUM_OF_EXPERTS_PER_RANK), + (uint32_t)(NUM_OF_EXPERTS_PER_RANK * sizeof(float)), &smem_buffer_ptr->intra_node_mbarrier_G2S_buffer[token_stage][0]); - total_tx_size += (uint32_t)((NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE) * sizeof(float)); + total_tx_size += (uint32_t)(NUM_OF_EXPERTS_PER_RANK * sizeof(float)); + // Store the source rank ID so the reduction warp group can place probs at the correct offset. + smem_buffer_ptr->intra_node_prob_src_rank_G2S_buffer[token_stage] = current_src_token_id; } if(current_src_token_id == last_src_token_id){ @@ -2578,12 +2598,16 @@ inline __device__ void intra_node_red_warp_group_device_function(const int node_ } if constexpr(BACKWARD_COMBINE){ + // Sparse prob optimization: SMEM only has E_per_rank elements for the source rank's slice. + // Read the source rank ID and accumulate into the correct position in acc_prob. + int src_rank = smem_buffer_ptr->intra_node_prob_src_rank_G2S_buffer[token_stage]; + int prob_offset = src_rank * NUM_OF_EXPERTS_PER_RANK; #pragma unroll for(int n = 0; n < NUM_OF_PROB_VEC_ELEMENT_PER_THREAD; n++){ - int element_id = INTRA_NODE_RED_GROUP::thread_rank() + n * INTRA_NODE_RED_GROUP::size(); - if(element_id < NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE){ - float src_data = load_prob_base_ptr[element_id]; - acc_prob[n] += src_data; + int global_element_id = INTRA_NODE_RED_GROUP::thread_rank() + n * INTRA_NODE_RED_GROUP::size(); + int local_element_id = global_element_id - prob_offset; + if(local_element_id >= 0 && local_element_id < NUM_OF_EXPERTS_PER_RANK){ + acc_prob[n] += load_prob_base_ptr[local_element_id]; } } } @@ -3082,14 +3106,17 @@ inline __device__ void inter_node_G2S_warp_group_device_function(const int node_ total_tx_size += (uint32_t)(HIDDEN_DIM * sizeof(uint16_t)); if constexpr(BACKWARD_COMBINE){ + // Sparse prob optimization: only read the source rank's E_per_rank slice. cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, cuda::ptx::space_global, reinterpret_cast(&smem_buffer_ptr->inter_node_prob_G2S_buffer[token_stage][0]), - reinterpret_cast(remote_expert_input_prob[current_src_token_id] + (sparse_to_dense_map_value * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE))), - (uint32_t)((NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE) * sizeof(float)), + reinterpret_cast(remote_expert_input_prob[current_src_token_id] + (sparse_to_dense_map_value * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE)) + current_src_token_id * NUM_OF_EXPERTS_PER_RANK), + (uint32_t)(NUM_OF_EXPERTS_PER_RANK * sizeof(float)), &smem_buffer_ptr->inter_node_mbarrier_G2S_buffer[token_stage][0]); - total_tx_size += (uint32_t)((NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE) * sizeof(float)); + total_tx_size += (uint32_t)(NUM_OF_EXPERTS_PER_RANK * sizeof(float)); + // Store the source rank ID so the reduction warp group can place probs at the correct offset. + smem_buffer_ptr->inter_node_prob_src_rank_G2S_buffer[token_stage] = current_src_token_id; } if(current_src_token_id == last_src_token_id){ @@ -3376,12 +3403,16 @@ inline __device__ void inter_node_red_warp_group_device_function(const int node_ } if constexpr(BACKWARD_COMBINE){ + // Sparse prob optimization: SMEM only has E_per_rank elements for the source rank's slice. + // Read the source rank ID and accumulate into the correct position in acc_prob[0]. + int src_rank = smem_buffer_ptr->inter_node_prob_src_rank_G2S_buffer[token_stage]; + int prob_offset = src_rank * NUM_OF_EXPERTS_PER_RANK; #pragma unroll for(int n = 0; n < NUM_OF_PROB_VEC_ELEMENT_PER_THREAD; n++){ - int element_id = thread_rank_within_pipeline + n * NUM_OF_THREADS_PER_PIPELINE; - if(element_id < NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE){ - float src_data = load_prob_base_ptr[element_id]; - acc_prob[0][n] += src_data; + int global_element_id = thread_rank_within_pipeline + n * NUM_OF_THREADS_PER_PIPELINE; + int local_element_id = global_element_id - prob_offset; + if(local_element_id >= 0 && local_element_id < NUM_OF_EXPERTS_PER_RANK){ + acc_prob[0][n] += load_prob_base_ptr[local_element_id]; } } } diff --git a/csrc/hybrid_ep/config.cuh b/csrc/hybrid_ep/config.cuh index 230d61f0c..769a9a737 100644 --- a/csrc/hybrid_ep/config.cuh +++ b/csrc/hybrid_ep/config.cuh @@ -287,13 +287,23 @@ static SmemSizes compute_smem_sizes(const HybridEpConfigInstance& c) { b.add((int64_t)c.num_of_stages_s2g_combine_api * c.hidden_dim * 2, 128); if (c.backward_combine_api) { if (multinode) { - b.add((int64_t)c.num_of_stages_g2s_combine_api * c.num_of_experts_per_rank * c.num_of_ranks_per_node * 4, 16); + // intra_node_prob_G2S: E_per_rank per stage (sparse prob optimization) + b.add((int64_t)c.num_of_stages_g2s_combine_api * c.num_of_experts_per_rank * 4, 16); + // intra_node_prob_S2G: E*R per stage (full vector for RDMA) b.add((int64_t)c.num_of_stages_s2g_combine_api * c.num_of_experts_per_rank * c.num_of_ranks_per_node * 4, 16); + // inter_node_prob_G2S: E*R per stage (reads from RDMA landing buffer, already reduced) b.add((int64_t)c.num_of_stages_g2s_combine_api * c.num_of_experts_per_rank * c.num_of_ranks_per_node * 4, 16); + // inter_node_prob_S2G: E*R*N per stage b.add((int64_t)c.num_of_stages_s2g_combine_api * c.num_of_experts_per_rank * c.num_of_ranks_per_node * c.num_of_nodes * 4, 16); + // intra_node_prob_src_rank: int per stage + b.add((int64_t)c.num_of_stages_g2s_combine_api * 4, 4); } else { - b.add((int64_t)c.num_of_stages_g2s_combine_api * c.num_of_experts_per_rank * c.num_of_ranks_per_node * 4, 16); + // inter_node_prob_G2S: E_per_rank per stage (sparse prob optimization) + b.add((int64_t)c.num_of_stages_g2s_combine_api * c.num_of_experts_per_rank * 4, 16); + // inter_node_prob_S2G: E*R per stage b.add((int64_t)c.num_of_stages_s2g_combine_api * c.num_of_experts_per_rank * c.num_of_ranks_per_node * 4, 16); + // inter_node_prob_src_rank: int per stage + b.add((int64_t)c.num_of_stages_g2s_combine_api * 4, 4); } } if (multinode) diff --git a/tests/test_hybrid_ep.py b/tests/test_hybrid_ep.py index 90f91a3ce..791277fdb 100644 --- a/tests/test_hybrid_ep.py +++ b/tests/test_hybrid_ep.py @@ -321,7 +321,7 @@ def test_hybrid_ep_benchmark(buffer: deep_ep.HybridEPBuffer, group: dist.Process multinode = (NUM_OF_NODES > 1) # ---- Setup: collect handles, build args dicts (also serves as warmup) ---- - # Non-permute + # Non-permute (forward dispatch with probs) dispatched_hidden, dispatched_probs, _, handle = ( buffer.dispatch(hidden=hidden, scaling_factor=scaling_factor, topk_idx=topk_idx, topk_weights=topk_weights, num_of_experts=NUM_OF_EXPERTS)) From db9880bd510d1a6647c0c1770e48159c327c4e37 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Tue, 12 May 2026 15:36:12 +0800 Subject: [PATCH 3/8] Fix sparse prob copy edge cases --- csrc/hybrid_ep/backend/hybrid_ep_backend.cuh | 82 +++++++++++++------- csrc/hybrid_ep/config.cuh | 2 + csrc/hybrid_ep/executor/executor.cu | 13 ++++ 3 files changed, 71 insertions(+), 26 deletions(-) diff --git a/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh b/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh index d4b1c0698..22c5edee0 100644 --- a/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh +++ b/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh @@ -76,6 +76,14 @@ using Copy_t = >::type >::type; +template +inline __device__ void scalar_copy_float_slice(float* dst, const float* src) { + #pragma unroll + for(int i = 0; i < NUM_OF_FLOATS; i++){ + dst[i] = src[i]; + } +} + enum scan_state{ EMPTY = 0, PRIV_SUM = 1 @@ -405,7 +413,7 @@ struct combine_kernel_dynamic_shared_memory_buffer_tconsumer, 2nd for consumer->producer. Should be 8B alignment(natural alignment). alignas(8) uint64_t intra_node_mbarrier_G2S_buffer[NUM_OF_STAGES_G2S][2]; @@ -1647,18 +1657,20 @@ inline __device__ void S2G_warp_group_device_function(const int local_rank, // The source SMEM prob buffer contains the full E*R probs (mostly zeros for sparse routing). // Each rank's E_per_rank slice already has zeros at non-active expert positions. if constexpr(FORWARD_DISPATCH){ - static_assert(NUM_OF_EXPERTS_PER_RANK * sizeof(float) >= 16, - "NUM_OF_EXPERTS_PER_RANK * sizeof(float) must be >= 16 for TMA minimum transfer size."); - static_assert((NUM_OF_EXPERTS_PER_RANK * sizeof(float)) % 16 == 0, - "NUM_OF_EXPERTS_PER_RANK * sizeof(float) must be 16B-aligned for TMA."); float* remote_prob_addr = remote_expert_output_prob[remote_rank_id] + output_buffer_index * static_cast(NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE) + remote_rank_id * NUM_OF_EXPERTS_PER_RANK; - cuda::ptx::cp_async_bulk(cuda::ptx::space_global, - cuda::ptx::space_shared, - reinterpret_cast(remote_prob_addr), - reinterpret_cast(&smem_buffer_ptr->intra_node_prob_buffer[stage][remote_rank_id * NUM_OF_EXPERTS_PER_RANK]), - (uint32_t)(NUM_OF_EXPERTS_PER_RANK * sizeof(float))); + const float* local_prob_addr = &smem_buffer_ptr->intra_node_prob_buffer[stage][remote_rank_id * NUM_OF_EXPERTS_PER_RANK]; + if constexpr(NUM_OF_EXPERTS_PER_RANK * sizeof(float) >= 16 && + (NUM_OF_EXPERTS_PER_RANK * sizeof(float)) % 16 == 0){ + cuda::ptx::cp_async_bulk(cuda::ptx::space_global, + cuda::ptx::space_shared, + reinterpret_cast(remote_prob_addr), + reinterpret_cast(local_prob_addr), + (uint32_t)(NUM_OF_EXPERTS_PER_RANK * sizeof(float))); + } else { + scalar_copy_float_slice(remote_prob_addr, local_prob_addr); + } } @@ -2377,14 +2389,23 @@ inline __device__ void intra_node_G2S_warp_group_device_function(const int node_ // Sparse prob optimization: only read the source rank's E_per_rank slice. // The source rank (current_src_token_id) wrote its probs at offset // current_src_token_id * E_per_rank within the E*R_per_node prob vector. - cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, - cuda::ptx::space_global, - reinterpret_cast(&smem_buffer_ptr->intra_node_prob_G2S_buffer[token_stage][0]), - reinterpret_cast(remote_expert_input_prob[current_src_token_id] + (sparse_to_dense_map_value * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE)) + current_src_token_id * NUM_OF_EXPERTS_PER_RANK), - (uint32_t)(NUM_OF_EXPERTS_PER_RANK * sizeof(float)), - &smem_buffer_ptr->intra_node_mbarrier_G2S_buffer[token_stage][0]); - - total_tx_size += (uint32_t)(NUM_OF_EXPERTS_PER_RANK * sizeof(float)); + float* local_prob_addr = &smem_buffer_ptr->intra_node_prob_G2S_buffer[token_stage][0]; + const float* remote_prob_addr = remote_expert_input_prob[current_src_token_id] + + sparse_to_dense_map_value * static_cast(NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE) + + current_src_token_id * NUM_OF_EXPERTS_PER_RANK; + if constexpr(NUM_OF_EXPERTS_PER_RANK * sizeof(float) >= 16 && + (NUM_OF_EXPERTS_PER_RANK * sizeof(float)) % 16 == 0){ + cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, + cuda::ptx::space_global, + reinterpret_cast(local_prob_addr), + reinterpret_cast(remote_prob_addr), + (uint32_t)(NUM_OF_EXPERTS_PER_RANK * sizeof(float)), + &smem_buffer_ptr->intra_node_mbarrier_G2S_buffer[token_stage][0]); + + total_tx_size += (uint32_t)(NUM_OF_EXPERTS_PER_RANK * sizeof(float)); + } else { + scalar_copy_float_slice(local_prob_addr, remote_prob_addr); + } // Store the source rank ID so the reduction warp group can place probs at the correct offset. smem_buffer_ptr->intra_node_prob_src_rank_G2S_buffer[token_stage] = current_src_token_id; } @@ -3107,14 +3128,23 @@ inline __device__ void inter_node_G2S_warp_group_device_function(const int node_ if constexpr(BACKWARD_COMBINE){ // Sparse prob optimization: only read the source rank's E_per_rank slice. - cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, - cuda::ptx::space_global, - reinterpret_cast(&smem_buffer_ptr->inter_node_prob_G2S_buffer[token_stage][0]), - reinterpret_cast(remote_expert_input_prob[current_src_token_id] + (sparse_to_dense_map_value * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE)) + current_src_token_id * NUM_OF_EXPERTS_PER_RANK), - (uint32_t)(NUM_OF_EXPERTS_PER_RANK * sizeof(float)), - &smem_buffer_ptr->inter_node_mbarrier_G2S_buffer[token_stage][0]); - - total_tx_size += (uint32_t)(NUM_OF_EXPERTS_PER_RANK * sizeof(float)); + float* local_prob_addr = &smem_buffer_ptr->inter_node_prob_G2S_buffer[token_stage][0]; + const float* remote_prob_addr = remote_expert_input_prob[current_src_token_id] + + sparse_to_dense_map_value * static_cast(NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE) + + current_src_token_id * NUM_OF_EXPERTS_PER_RANK; + if constexpr(NUM_OF_EXPERTS_PER_RANK * sizeof(float) >= 16 && + (NUM_OF_EXPERTS_PER_RANK * sizeof(float)) % 16 == 0){ + cuda::ptx::cp_async_bulk(cuda::ptx::space_shared, + cuda::ptx::space_global, + reinterpret_cast(local_prob_addr), + reinterpret_cast(remote_prob_addr), + (uint32_t)(NUM_OF_EXPERTS_PER_RANK * sizeof(float)), + &smem_buffer_ptr->inter_node_mbarrier_G2S_buffer[token_stage][0]); + + total_tx_size += (uint32_t)(NUM_OF_EXPERTS_PER_RANK * sizeof(float)); + } else { + scalar_copy_float_slice(local_prob_addr, remote_prob_addr); + } // Store the source rank ID so the reduction warp group can place probs at the correct offset. smem_buffer_ptr->inter_node_prob_src_rank_G2S_buffer[token_stage] = current_src_token_id; } diff --git a/csrc/hybrid_ep/config.cuh b/csrc/hybrid_ep/config.cuh index 769a9a737..a1bbdab3d 100644 --- a/csrc/hybrid_ep/config.cuh +++ b/csrc/hybrid_ep/config.cuh @@ -297,6 +297,8 @@ static SmemSizes compute_smem_sizes(const HybridEpConfigInstance& c) { b.add((int64_t)c.num_of_stages_s2g_combine_api * c.num_of_experts_per_rank * c.num_of_ranks_per_node * c.num_of_nodes * 4, 16); // intra_node_prob_src_rank: int per stage b.add((int64_t)c.num_of_stages_g2s_combine_api * 4, 4); + // inter_node_prob_src_rank: int per stage for local-source G2S entries + b.add((int64_t)c.num_of_stages_g2s_combine_api * 4, 4); } else { // inter_node_prob_G2S: E_per_rank per stage (sparse prob optimization) b.add((int64_t)c.num_of_stages_g2s_combine_api * c.num_of_experts_per_rank * 4, 16); diff --git a/csrc/hybrid_ep/executor/executor.cu b/csrc/hybrid_ep/executor/executor.cu index d21da301f..6a9e49314 100644 --- a/csrc/hybrid_ep/executor/executor.cu +++ b/csrc/hybrid_ep/executor/executor.cu @@ -301,6 +301,19 @@ void Executor::dispatch_core(HybridEpConfigInstance config, DispatchArgs& args) param.multinode_aux_ptr = reinterpret_cast(inter_node_dispatch_buffers->mr_info); #endif #endif + static constexpr bool kZeroDispatchProbBufferBeforeSparseWrites = false; + if constexpr (kZeroDispatchProbBufferBeforeSparseWrites) { + if (config.forward_dispatch_api && args.num_dispatched_tokens_tensor.has_value()) { + // Sparse prob dispatch writes only the destination rank's E_per_rank slice. + // A dense memset would preserve the old full-probs output contract, but it can be + // redundant and expensive when TOPK is large, so keep this path disabled by default. + int num_dispatched_tokens = args.num_dispatched_tokens_tensor.value().item(); + size_t probs_sz = static_cast(num_dispatched_tokens) * + config.num_of_experts_per_rank * config.num_of_ranks_per_node * sizeof(float); + CUDA_CHECK(cudaMemsetAsync(intra_node_dispatch_buffers->expert_output_prob, 0, probs_sz, args.stream)); + } + } + // Launch kernel kernel_cache.run_dispatch_kernel(config, param, args.fuse_permute_dispatch, args.non_blocking, args.stream); nvtxRangePop(); // End of dispatch_core nvtx range From 7e6544479a8c974425a778811ea9c463ed3091a8 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Tue, 12 May 2026 17:08:50 +0800 Subject: [PATCH 4/8] Add dense routing map for scan input Teach the Hybrid-EP metadata preprocessing path to accept dense int16 topk_idx rows in addition to the existing sparse bool routing map. The scan kernel now templates on TOPK, derives per-rank/per-node routing by range-checking dense expert IDs, reconstructs local expert routing data where needed, and includes TOPK in the JIT cache key so sparse and dense preprocessing kernels do not alias. Plumb the dense routing mode through HybridEpConfigInstance, pybind, the JIT compiler, Executor::allgather_routing_map, dispatch(), and dispatch_with_permute(). Dense routing converts topk_idx to contiguous int16, keeps handle reuse working without requiring topk_idx again, and masks -1 dropped-token sentinels when reconstructing probs on the Python side. The API rejects expert counts beyond the int16 representable range to avoid silent wraparound. Handle collectives for the dense layout: NCCL allgather views int16 tensors as int8 because NCCL does not directly support int16, and the custom NVLink allgather now sizes its backing buffer for the larger of sparse bool rows and dense topk rows. If a routing row is not 16B-aligned, the executor falls back to NCCL instead of launching the custom uint4 allgather kernel. Extend test_hybrid_ep.py to compare dense-routing dispatch against the existing sparse reference for with-probs and no-probs cases, then run combine with the dense handle to validate the metadata maps are usable beyond dispatch. --- csrc/hybrid_ep/backend/hybrid_ep_backend.cuh | 362 ++++++++++++------- csrc/hybrid_ep/config.cuh | 2 + csrc/hybrid_ep/executor/executor.cu | 33 +- csrc/hybrid_ep/extension/allgather.cu | 13 +- csrc/hybrid_ep/extension/allgather.cuh | 3 +- csrc/hybrid_ep/jit/compiler.cu | 9 +- csrc/hybrid_ep/jit/compiler.cuh | 2 +- csrc/hybrid_ep/pybind_hybrid_ep.cu | 1 + deep_ep/hybrid_ep_buffer.py | 83 ++++- tests/test_hybrid_ep.py | 46 +++ 10 files changed, 405 insertions(+), 149 deletions(-) diff --git a/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh b/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh index 22c5edee0..806e1def9 100644 --- a/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh +++ b/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh @@ -4858,12 +4858,14 @@ __global__ void combine_kernel(const __grid_constant__ combine_kernel_param_t pa template + int NUM_OF_EXPERTS_PER_RANK, + int TOPK = 0> // 0 = bool routing map mode; >0 = dense topk_idx mode with this K value __launch_bounds__(NUM_THREADS_PER_BLOCK, 1) -__global__ void scan(const bool* input_routing_map, +__global__ void scan(const void* input_routing_data, // bool* (TOPK==0) or int16_t* (TOPK>0) tmp_state_t* tmp, tmp_state_t* local_experts_tmp, int32_t* sparse_to_dense_map, @@ -4880,6 +4882,12 @@ __global__ void scan(const bool* input_routing_map, const int local_experts_tokens_limit, // This size MUST be multiple of LOCAL_EXPERTS_PADDING_SIZE! const int num_of_tokens_per_rank) { + // Dense vs sparse routing mode. + constexpr bool DENSE_ROUTING = (TOPK > 0); + // Cast input pointer to the correct type. + const bool* input_routing_map = DENSE_ROUTING ? nullptr : reinterpret_cast(input_routing_data); + const int16_t* input_topk_idx = DENSE_ROUTING ? reinterpret_cast(input_routing_data) : nullptr; + // Calculate the warps per block. constexpr int WARP_SIZE = 32; constexpr int NUM_OF_WARPS_PER_BLOCK = NUM_THREADS_PER_BLOCK / WARP_SIZE; @@ -4907,7 +4915,7 @@ __global__ void scan(const bool* input_routing_map, #endif // For each token(row in routing map), calculate how many bytes need to be loaded from the routing map and how to load them. static_assert(sizeof(bool) == 1, "Bool is not 1 byte???"); - constexpr int NUM_OF_BYTES_TO_LOAD_FOR_EACH_TOKEN = NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE; + constexpr int NUM_OF_BYTES_TO_LOAD_FOR_EACH_TOKEN = DENSE_ROUTING ? (TOPK * sizeof(int16_t)) : (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE); using copy_t = Copy_t; static_assert(NUM_OF_BYTES_TO_LOAD_FOR_EACH_TOKEN % sizeof(copy_t) == 0, "NUM_OF_BYTES_TO_LOAD_FOR_EACH_TOKEN and copy_t mismatch"); constexpr int ROUTING_MAP_LOAD_ITER = NUM_OF_BYTES_TO_LOAD_FOR_EACH_TOKEN / sizeof(copy_t); @@ -5010,47 +5018,99 @@ __global__ void scan(const bool* input_routing_map, // We need to calculate the per-node routing info and save back to rdma_to_attn_map. bool per_node_routing_info = (current_token_local_rank == local_rank); int current_token_rdma_to_attn_map_id = current_token_node_rank * rdma_to_attn_map_size_per_node + current_token_local_id; - // Global routing map load base addr for current token. - const copy_t* routing_map_load_base_addr = reinterpret_cast(input_routing_map + - current_token_id * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE * NUM_OF_NODES) + - node_rank * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE)); - - // Load the routing map for current token. - bool token_routing_map[NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE]; - #pragma unroll - for(int j = 0; j < ROUTING_MAP_LOAD_ITER; j++){ - *(reinterpret_cast(token_routing_map) + j) = routing_map_load_base_addr[j]; - } - - // Convert the routing map to per rank routing info and accumulate to accumulator. - // Also convert the per rank routing info to per node routing info. - // When permute fusion is enabled, also accumulate local experts to accumulator. + // Load routing data and compute per-rank routing info. bool token_needed_by_this_node = false; + // Per-rank routing flag array (reused in both modes). + bool token_needed_by_rank[NUM_OF_RANKS_PER_NODE]; #pragma unroll for(int j = 0; j < NUM_OF_RANKS_PER_NODE; j++){ - bool token_needed_by_this_rank = false; + token_needed_by_rank[j] = false; + } + + if constexpr(DENSE_ROUTING){ + // Dense mode: load TOPK int16 expert indices and check against rank ranges. + // Input layout: [num_total_tokens, TOPK] int16_t (NOT per-node sliced — global expert IDs). + // Dropped tokens use -1 as sentinel (negative, so naturally excluded by range checks). + const copy_t* topk_load_base_addr = reinterpret_cast(input_topk_idx + current_token_id * TOPK); + int16_t topk_experts[TOPK]; #pragma unroll - for(int k = 0; k < EXPERTS_TO_RANK_REDUCE_ITER; k++){ - int current_expert_to_rank_t_id = j * EXPERTS_TO_RANK_REDUCE_ITER + k; - expert_to_rank_t reduction_data = *(reinterpret_cast(token_routing_map) + current_expert_to_rank_t_id); - if(reduction_data != (expert_to_rank_t)0){ - token_needed_by_this_rank = true; - break; + for(int j = 0; j < ROUTING_MAP_LOAD_ITER; j++){ + *(reinterpret_cast(topk_experts) + j) = topk_load_base_addr[j]; + } + // Compute per-rank and per-node routing from topk expert indices. + // Expert global_id maps to node = global_id / (E_per_rank * R_per_node), rank_within_node = (global_id / E_per_rank) % R_per_node. + // We only care about experts on the local node (node_rank). + constexpr int EXPERTS_PER_NODE = NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE; + int node_expert_start = node_rank * EXPERTS_PER_NODE; + int node_expert_end = node_expert_start + EXPERTS_PER_NODE; + #pragma unroll + for(int k = 0; k < TOPK; k++){ + int expert_id = (int)topk_experts[k]; + if(expert_id >= node_expert_start && expert_id < node_expert_end){ + int local_expert_id = expert_id - node_expert_start; + int rank_within_node = local_expert_id / NUM_OF_EXPERTS_PER_RANK; + token_needed_by_rank[rank_within_node] = true; + token_needed_by_this_node = true; } } - if(token_needed_by_this_rank){ - token_routing_map_sum[j] += 1; - token_needed_by_this_node = true; + // Accumulate per-rank sums. + #pragma unroll + for(int j = 0; j < NUM_OF_RANKS_PER_NODE; j++){ + if(token_needed_by_rank[j]){ + token_routing_map_sum[j] += 1; + } } #ifdef HYBRID_EP_BUILD_PERMUTE_FUSION_ENABLE - if(j == local_rank){ - int current_local_expert_id = j * NUM_OF_EXPERTS_PER_RANK; - #pragma unroll - for(int k = 0; k < NUM_OF_EXPERTS_PER_RANK; k++){ - token_local_experts_routing_map_sum[k] += (int32_t)(token_routing_map[current_local_expert_id + k]); + // For permute fusion: compute per-local-expert routing. + #pragma unroll + for(int k = 0; k < TOPK; k++){ + int expert_id = (int)topk_experts[k]; + int local_expert_start = node_expert_start + local_rank * NUM_OF_EXPERTS_PER_RANK; + int local_expert_end = local_expert_start + NUM_OF_EXPERTS_PER_RANK; + if(expert_id >= local_expert_start && expert_id < local_expert_end){ + int local_expert_idx = expert_id - local_expert_start; + token_local_experts_routing_map_sum[local_expert_idx] += 1; } } #endif + } else { + // Sparse bool mode: load E_per_rank * R_per_node bools for local node slice. + const copy_t* routing_map_load_base_addr = reinterpret_cast(input_routing_map + + current_token_id * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE * NUM_OF_NODES) + + node_rank * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE)); + bool token_routing_map[NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE]; + #pragma unroll + for(int j = 0; j < ROUTING_MAP_LOAD_ITER; j++){ + *(reinterpret_cast(token_routing_map) + j) = routing_map_load_base_addr[j]; + } + // Convert per-expert bools to per-rank bools. + #pragma unroll + for(int j = 0; j < NUM_OF_RANKS_PER_NODE; j++){ + bool token_needed_by_this_rank = false; + #pragma unroll + for(int k = 0; k < EXPERTS_TO_RANK_REDUCE_ITER; k++){ + int current_expert_to_rank_t_id = j * EXPERTS_TO_RANK_REDUCE_ITER + k; + expert_to_rank_t reduction_data = *(reinterpret_cast(token_routing_map) + current_expert_to_rank_t_id); + if(reduction_data != (expert_to_rank_t)0){ + token_needed_by_this_rank = true; + break; + } + } + token_needed_by_rank[j] = token_needed_by_this_rank; + if(token_needed_by_this_rank){ + token_routing_map_sum[j] += 1; + token_needed_by_this_node = true; + } +#ifdef HYBRID_EP_BUILD_PERMUTE_FUSION_ENABLE + if(j == local_rank){ + int current_local_expert_id = j * NUM_OF_EXPERTS_PER_RANK; + #pragma unroll + for(int k = 0; k < NUM_OF_EXPERTS_PER_RANK; k++){ + token_local_experts_routing_map_sum[k] += (int32_t)(token_routing_map[current_local_expert_id + k]); + } + } +#endif + } } // Save the per node routing info back to rdma_to_attn_map if needed. @@ -5392,75 +5452,72 @@ __global__ void scan(const bool* input_routing_map, bool token_needed_by_dense_chunk_layout = first_token_of_a_chunk && current_token_global_chunk_id > 0 && current_token_global_chunk_id < num_of_total_attn_chunks; #endif - // Global routing map load base addr for current token. - const copy_t* routing_map_load_base_addr = reinterpret_cast(input_routing_map + - current_token_id * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE * NUM_OF_NODES) + - node_rank * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE)); - - // Load the routing map for current token. Only load when the token is not out-of-bound. - bool token_routing_map[NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE]; - if(token_out_of_bound == 0){ - #pragma unroll - for(int j = 0; j < ROUTING_MAP_LOAD_ITER; j++){ - *(reinterpret_cast(token_routing_map) + j) = routing_map_load_base_addr[j]; - } - } - - // Convert the routing map to per rank routing info for current token, - // then produce the per-rank final exclusive scan within the warp for this tile. - int32_t final_ex_scan[NUM_OF_RANKS_PER_NODE]; + // Load routing data for current token (reuse for all ranks). + // In dense mode: load TOPK int16 expert indices. + // In sparse mode: load E_per_rank * R_per_node bools for local node slice. + // Also compute per-rank routing flags. + bool token_needed_by_rank_step2[NUM_OF_RANKS_PER_NODE]; + // Sparse mode needs the full routing map for local_expert_routing_map writes later. + bool token_routing_map[DENSE_ROUTING ? 1 : (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE)]; #pragma unroll for(int j = 0; j < NUM_OF_RANKS_PER_NODE; j++){ - int32_t temp_scan = 0; - bool token_needed_by_this_rank = false; - // Old warp-level scan implementation, using warp shuffle, suitable for general data type, but not fast enough for bool type. - // If the token is not out-of-bound, check whether this rank need this token. - /*if(token_out_of_bound == 0){ + token_needed_by_rank_step2[j] = false; + } + + if constexpr(DENSE_ROUTING){ + if(token_out_of_bound == 0){ + const copy_t* topk_load_base_addr = reinterpret_cast(input_topk_idx + current_token_id * TOPK); + int16_t topk_experts[TOPK]; #pragma unroll - for(int k = 0; k < EXPERTS_TO_RANK_REDUCE_ITER; k++){ - int current_expert_to_rank_t_id = j * EXPERTS_TO_RANK_REDUCE_ITER + k; - expert_to_rank_t reduction_data = *(reinterpret_cast(token_routing_map) + current_expert_to_rank_t_id); - if(reduction_data != (expert_to_rank_t)0){ - token_needed_by_this_rank = true; - break; - } + for(int j = 0; j < ROUTING_MAP_LOAD_ITER; j++){ + *(reinterpret_cast(topk_experts) + j) = topk_load_base_addr[j]; } - if(token_needed_by_this_rank){ - temp_scan = 1; - }else{ - temp_scan = 0; + constexpr int EXPERTS_PER_NODE = NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE; + int node_expert_start = node_rank * EXPERTS_PER_NODE; + int node_expert_end = node_expert_start + EXPERTS_PER_NODE; + #pragma unroll + for(int k = 0; k < TOPK; k++){ + int expert_id = (int)topk_experts[k]; + if(expert_id >= node_expert_start && expert_id < node_expert_end){ + int local_expert_id = expert_id - node_expert_start; + int rank_within_node = local_expert_id / NUM_OF_EXPERTS_PER_RANK; + token_needed_by_rank_step2[rank_within_node] = true; + } } } - - // Each warp perform a inclusive scan from all threads(lanes). - #pragma unroll - for(int k = 1; k < WARP_SIZE; k *= 2){ - int32_t temp = __shfl_up_sync(~0, temp_scan, k); - if(lane_id >= k){ - temp_scan += temp; + } else { + // Sparse bool mode: load and reduce as before. + const copy_t* routing_map_load_base_addr = reinterpret_cast(input_routing_map + + current_token_id * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE * NUM_OF_NODES) + + node_rank * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE)); + if(token_out_of_bound == 0){ + #pragma unroll + for(int j = 0; j < ROUTING_MAP_LOAD_ITER; j++){ + *(reinterpret_cast(token_routing_map) + j) = routing_map_load_base_addr[j]; } } - - // The inclusive scan from last lane is the sum of this rank of this tile. Need to accumulate that for later tiles. - int32_t temp_sum = __shfl_sync(~0, temp_scan, WARP_SIZE - 1); - - // Make scan exclusive. - int32_t exclusive_scan = __shfl_up_sync(~0, temp_scan, 1); - temp_scan = (lane_id >= 1) ? exclusive_scan : 0;*/ - - // New warp-level scan implementation for bool value, using warp vote instead of warp shuffle. Warp vote is way faster than warp shuffle. - // If the token is not out-of-bound, check whether this rank need this token. if(token_out_of_bound == 0){ #pragma unroll - for(int k = 0; k < EXPERTS_TO_RANK_REDUCE_ITER; k++){ - int current_expert_to_rank_t_id = j * EXPERTS_TO_RANK_REDUCE_ITER + k; - expert_to_rank_t reduction_data = *(reinterpret_cast(token_routing_map) + current_expert_to_rank_t_id); - if(reduction_data != (expert_to_rank_t)0){ - token_needed_by_this_rank = true; - break; + for(int j = 0; j < NUM_OF_RANKS_PER_NODE; j++){ + #pragma unroll + for(int k = 0; k < EXPERTS_TO_RANK_REDUCE_ITER; k++){ + int current_expert_to_rank_t_id = j * EXPERTS_TO_RANK_REDUCE_ITER + k; + expert_to_rank_t reduction_data = *(reinterpret_cast(token_routing_map) + current_expert_to_rank_t_id); + if(reduction_data != (expert_to_rank_t)0){ + token_needed_by_rank_step2[j] = true; + break; + } } } } + } + + // Convert the per-rank routing flags to per-rank final exclusive scan within the warp for this tile. + int32_t final_ex_scan[NUM_OF_RANKS_PER_NODE]; + #pragma unroll + for(int j = 0; j < NUM_OF_RANKS_PER_NODE; j++){ + int32_t temp_scan = 0; + bool token_needed_by_this_rank = token_needed_by_rank_step2[j]; // Each warp vote to create a bit mask indicating which token is needed by this rank within this tile. unsigned vote_result = __ballot_sync(~0, token_needed_by_this_rank); @@ -5481,7 +5538,21 @@ __global__ void scan(const bool* input_routing_map, for(int k = 0; k < NUM_OF_EXPERTS_PER_RANK; k++){ bool token_needed_by_this_local_expert = false; if(token_out_of_bound == 0){ - token_needed_by_this_local_expert = token_routing_map[j * NUM_OF_EXPERTS_PER_RANK + k]; + if constexpr(DENSE_ROUTING){ + // Check if any topk index matches this local expert. + int target_expert = node_rank * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE) + j * NUM_OF_EXPERTS_PER_RANK + k; + // Re-load topk indices. The compiler should hoist this out of the k loop. + const int16_t* topk_base = input_topk_idx + current_token_id * TOPK; + #pragma unroll + for(int m = 0; m < TOPK; m++){ + if((int)topk_base[m] == target_expert){ + token_needed_by_this_local_expert = true; + break; + } + } + } else { + token_needed_by_this_local_expert = token_routing_map[j * NUM_OF_EXPERTS_PER_RANK + k]; + } } unsigned local_expert_vote_result = __ballot_sync(~0, token_needed_by_this_local_expert); int32_t local_expert_temp_sum = __popc(local_expert_vote_result); @@ -5510,11 +5581,37 @@ __global__ void scan(const bool* input_routing_map, #else // Each thread save local routing map for this token of the local rank to local_expert_routing_map if this token is needed by the local rank. if(j == local_rank && token_needed_by_this_rank){ - expert_to_rank_t* local_expert_routing_map_store_base_addr = reinterpret_cast(local_expert_routing_map + (final_ex_scan[j] * NUM_OF_EXPERTS_PER_RANK)); - #pragma unroll - for(int k = 0; k < EXPERTS_TO_RANK_REDUCE_ITER; k++){ - int current_expert_to_rank_t_id = j * EXPERTS_TO_RANK_REDUCE_ITER + k; - local_expert_routing_map_store_base_addr[k] = *(reinterpret_cast(token_routing_map) + current_expert_to_rank_t_id); + if constexpr(DENSE_ROUTING){ + // Reconstruct per-expert bool routing map from topk indices for the local rank. + bool local_expert_bools[NUM_OF_EXPERTS_PER_RANK]; + #pragma unroll + for(int k = 0; k < NUM_OF_EXPERTS_PER_RANK; k++){ + local_expert_bools[k] = false; + } + // Re-load topk indices (or reuse from above if we cached them). + const int16_t* topk_base = input_topk_idx + current_token_id * TOPK; + int local_expert_start = node_rank * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE) + local_rank * NUM_OF_EXPERTS_PER_RANK; + #pragma unroll + for(int k = 0; k < TOPK; k++){ + int expert_id = (int)topk_base[k]; + int local_idx = expert_id - local_expert_start; + if(local_idx >= 0 && local_idx < NUM_OF_EXPERTS_PER_RANK){ + local_expert_bools[local_idx] = true; + } + } + // Store as expert_to_rank_t chunks. + expert_to_rank_t* local_expert_routing_map_store_base_addr = reinterpret_cast(local_expert_routing_map + (final_ex_scan[j] * NUM_OF_EXPERTS_PER_RANK)); + #pragma unroll + for(int k = 0; k < EXPERTS_TO_RANK_REDUCE_ITER; k++){ + local_expert_routing_map_store_base_addr[k] = *(reinterpret_cast(local_expert_bools) + k); + } + } else { + expert_to_rank_t* local_expert_routing_map_store_base_addr = reinterpret_cast(local_expert_routing_map + (final_ex_scan[j] * NUM_OF_EXPERTS_PER_RANK)); + #pragma unroll + for(int k = 0; k < EXPERTS_TO_RANK_REDUCE_ITER; k++){ + int current_expert_to_rank_t_id = j * EXPERTS_TO_RANK_REDUCE_ITER + k; + local_expert_routing_map_store_base_addr[k] = *(reinterpret_cast(token_routing_map) + current_expert_to_rank_t_id); + } } } #endif @@ -5563,37 +5660,54 @@ __global__ void scan(const bool* input_routing_map, int current_token_node_id = current_token_attn_to_rdma_map_node_id < node_rank ? current_token_attn_to_rdma_map_node_id : current_token_attn_to_rdma_map_node_id + 1; int current_token_local_id = current_token_id / (NUM_OF_NODES - 1); - const copy_t* routing_map_load_base_addr = reinterpret_cast(input_routing_map + - ((node_rank * NUM_OF_RANKS_PER_NODE + local_rank) * num_of_tokens_per_rank + current_token_local_id) * - (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE * NUM_OF_NODES) + - (current_token_node_id * NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE)); - bool* attn_to_rdma_map_base_addr = attn_to_rdma_map + (current_token_local_id * (NUM_OF_NODES - 1) + current_token_attn_to_rdma_map_node_id); - // Load the routing map for current token row. - bool token_routing_map[NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE]; - #pragma unroll - for(int j = 0; j < ROUTING_MAP_LOAD_ITER; j++){ - *(reinterpret_cast(token_routing_map) + j) = routing_map_load_base_addr[j]; - } - - // Convert the routing map to per rank routing info and then to per node routing info. + // Check if any expert on the target node is needed by this token. bool token_needed_by_this_node = false; - #pragma unroll - for(int j = 0; j < NUM_OF_RANKS_PER_NODE; j++){ - bool token_needed_by_this_rank = false; + + if constexpr(DENSE_ROUTING){ + // Dense mode: check if any topk expert falls on current_token_node_id. + // The global token id for this local token on the local rank. + int global_token_id = (node_rank * NUM_OF_RANKS_PER_NODE + local_rank) * num_of_tokens_per_rank + current_token_local_id; + const int16_t* topk_base = input_topk_idx + global_token_id * TOPK; + constexpr int EXPERTS_PER_NODE = NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE; + int target_node_expert_start = current_token_node_id * EXPERTS_PER_NODE; + int target_node_expert_end = target_node_expert_start + EXPERTS_PER_NODE; #pragma unroll - for(int k = 0; k < EXPERTS_TO_RANK_REDUCE_ITER; k++){ - int current_expert_to_rank_t_id = j * EXPERTS_TO_RANK_REDUCE_ITER + k; - expert_to_rank_t reduction_data = *(reinterpret_cast(token_routing_map) + current_expert_to_rank_t_id); - if(reduction_data != (expert_to_rank_t)0){ - token_needed_by_this_rank = true; + for(int k = 0; k < TOPK; k++){ + int expert_id = (int)topk_base[k]; + if(expert_id >= target_node_expert_start && expert_id < target_node_expert_end){ + token_needed_by_this_node = true; break; } } - if(token_needed_by_this_rank){ - token_needed_by_this_node = true; - break; + } else { + // Sparse bool mode: load routing map for target node and reduce. + const copy_t* routing_map_load_base_addr = reinterpret_cast(input_routing_map + + ((node_rank * NUM_OF_RANKS_PER_NODE + local_rank) * num_of_tokens_per_rank + current_token_local_id) * + (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE * NUM_OF_NODES) + + (current_token_node_id * NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE)); + bool token_routing_map_step3[NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE]; + #pragma unroll + for(int j = 0; j < ROUTING_MAP_LOAD_ITER; j++){ + *(reinterpret_cast(token_routing_map_step3) + j) = routing_map_load_base_addr[j]; + } + #pragma unroll + for(int j = 0; j < NUM_OF_RANKS_PER_NODE; j++){ + bool token_needed_by_this_rank = false; + #pragma unroll + for(int k = 0; k < EXPERTS_TO_RANK_REDUCE_ITER; k++){ + int current_expert_to_rank_t_id = j * EXPERTS_TO_RANK_REDUCE_ITER + k; + expert_to_rank_t reduction_data = *(reinterpret_cast(token_routing_map_step3) + current_expert_to_rank_t_id); + if(reduction_data != (expert_to_rank_t)0){ + token_needed_by_this_rank = true; + break; + } + } + if(token_needed_by_this_rank){ + token_needed_by_this_node = true; + break; + } } } @@ -5660,8 +5774,10 @@ public: // Block size for preprocessing kernel. int NUM_THREADS_PER_BLOCK, // Grid size for preprocessing kernel(1:1 block:SM mapping). - int NUM_OF_BLOCKS> - static void metadata_preprocessing(const bool* input_routing_map, + int NUM_OF_BLOCKS, + // 0 = bool routing map mode; >0 = dense topk_idx mode with this K value + int TOPK = 0> + static void metadata_preprocessing(const void* input_routing_data, // bool* (TOPK==0) or int16_t* (TOPK>0) tmp_state_t* preprocessing_tmp, tmp_state_t* preprocessing_local_experts_tmp, int32_t* sparse_to_dense_map, @@ -5700,9 +5816,9 @@ public: #endif // Launch the preprocessing kernel to process the global routing map. - scan + scan <<>> - (input_routing_map, preprocessing_tmp, preprocessing_local_experts_tmp, sparse_to_dense_map, rdma_to_attn_map, attn_to_rdma_map, num_of_tokens_for_experts, local_expert_routing_map, + (input_routing_data, preprocessing_tmp, preprocessing_local_experts_tmp, sparse_to_dense_map, rdma_to_attn_map, attn_to_rdma_map, num_of_tokens_for_experts, local_expert_routing_map, dense_chunk_layout, dense_to_expert_map, num_of_local_experts_tokens, token_drop_triggered, node_rank, local_rank, local_experts_tokens_limit, num_of_tokens_per_rank); // Check if there is any CUDA error. diff --git a/csrc/hybrid_ep/config.cuh b/csrc/hybrid_ep/config.cuh index a1bbdab3d..9309650ce 100644 --- a/csrc/hybrid_ep/config.cuh +++ b/csrc/hybrid_ep/config.cuh @@ -113,6 +113,7 @@ struct HybridEpConfigInstance { int num_of_ranks_per_node; int num_of_nodes; int pad_multiple; + int topk; // 0 = sparse bool routing map; >0 = dense int16 topk_idx mode /* * Metadata-preprocessing API Config @@ -452,6 +453,7 @@ public: config.num_of_additional_in_flight_s2g_unpermute_block_combine_api = get_env_int("NUM_OF_ADDITIONAL_IN_FLIGHT_S2G_UNPERMUTE_BLOCK_COMBINE_API", 2); config.pad_multiple = 1; + config.topk = 0; // default: sparse bool routing map // If we use the fused permute-dispatch kernel, the number of blocks // for the permute part is the same as the number of blocks for the dispatch part. diff --git a/csrc/hybrid_ep/executor/executor.cu b/csrc/hybrid_ep/executor/executor.cu index 6a9e49314..ceed02f64 100644 --- a/csrc/hybrid_ep/executor/executor.cu +++ b/csrc/hybrid_ep/executor/executor.cu @@ -29,25 +29,40 @@ torch::Tensor Executor::allgather_routing_map( nvtxRangePushA("allgather_routing_map in hybrid-ep"); auto torch_distributed = py::module_::import("torch.distributed"); - auto num_of_expert = local_routing_map.size(-1); + auto num_cols = local_routing_map.size(-1); // E_total (sparse) or TOPK (dense) auto num_of_tokens_per_rank = local_routing_map.size(-2); auto group_size = process_group.attr("size")().cast(); - assert(num_of_expert == config.num_of_experts_per_rank * config.num_of_ranks_per_node * config.num_of_nodes); + bool dense_routing = (config.topk > 0); + bool custom_allgather_aligned = (local_routing_map.numel() * local_routing_map.element_size()) % 16 == 0; + if (!dense_routing) { + assert(num_cols == config.num_of_experts_per_rank * config.num_of_ranks_per_node * config.num_of_nodes); + } else { + assert(num_cols == config.topk); + } + auto dtype = dense_routing ? torch::kInt16 : torch::kBool; torch::Tensor global_routing_map; // At inter-node case, we will use NCCL allgather - if(config.num_of_nodes > 1 || !enable_custom_allgather) { + if(config.num_of_nodes > 1 || !enable_custom_allgather || !custom_allgather_aligned) { global_routing_map = torch::empty( - {num_of_tokens_per_rank * group_size, num_of_expert}, - torch::TensorOptions().dtype(torch::kBool).device(torch::kCUDA) + {num_of_tokens_per_rank * group_size, num_cols}, + torch::TensorOptions().dtype(dtype).device(torch::kCUDA) ); - torch_distributed.attr("all_gather_into_tensor")(global_routing_map, local_routing_map, process_group); + if (dense_routing) { + // NCCL does not support int16 directly. View as int8 for the collective, + // then reinterpret the gathered bytes through the int16 output tensor. + auto local_as_bytes = local_routing_map.view(torch::kInt8); + auto global_as_bytes = global_routing_map.view(torch::kInt8); + torch_distributed.attr("all_gather_into_tensor")(global_as_bytes, local_as_bytes, process_group); + } else { + torch_distributed.attr("all_gather_into_tensor")(global_routing_map, local_routing_map, process_group); + } } else { // At intra-node case, we will use custom allgather allgather_obj.launch(local_routing_map, /*NUM_OF_SMS=*/32, at::cuda::getCurrentCUDAStream()); global_routing_map = torch::from_blob( allgather_obj.get_output_buffer(), - {num_of_tokens_per_rank * group_size, num_of_expert}, - torch::TensorOptions().dtype(torch::kBool).device(torch::kCUDA) + {num_of_tokens_per_rank * group_size, num_cols}, + torch::TensorOptions().dtype(dtype).device(torch::kCUDA) ); } @@ -133,7 +148,7 @@ HandleImpl Executor::metadata_preprocess_core( } kernel_cache.run_preprocess_kernel( - config, global_routing_map.data_ptr(), + config, global_routing_map.data_ptr(), preprocessing_tmp, preprocessing_local_experts_tmp, handle.sparse_to_dense_map.data_ptr(), handle.rdma_to_attn_map.data_ptr(), handle.attn_to_rdma_map.data_ptr(), diff --git a/csrc/hybrid_ep/extension/allgather.cu b/csrc/hybrid_ep/extension/allgather.cu index 481699329..dd8cb997f 100644 --- a/csrc/hybrid_ep/extension/allgather.cu +++ b/csrc/hybrid_ep/extension/allgather.cu @@ -3,6 +3,9 @@ #include "allgather.cuh" +#include +#include + #define MAX_BLOCKS 256 #define TIMEOUT 20000000000ull @@ -115,6 +118,7 @@ void CustomAllgather::init(pybind11::object process_group, int rank_idx, BufferC this->num_of_experts_per_rank = buffer_config.num_of_experts_per_rank; this->num_of_tokens_per_rank = buffer_config.max_num_of_tokens_per_rank; this->num_of_nodes = buffer_config.num_of_nodes; + this->routing_bytes_per_token = num_of_experts_per_rank * num_of_ranks_per_node * num_of_nodes * sizeof(bool); this->allocator = allocator; this->process_group = process_group; } @@ -125,6 +129,9 @@ bool CustomAllgather::grow_buffer_config(const HybridEpConfigInstance& config, B changed |= grow_to(buf_config.num_of_experts_per_rank, config.num_of_experts_per_rank); changed |= grow_to(buf_config.max_num_of_tokens_per_rank, config.max_num_of_tokens_per_rank); changed |= grow_to(buf_config.num_of_nodes, config.num_of_nodes); + int sparse_bytes_per_token = config.num_of_experts_per_rank * config.num_of_ranks_per_node * config.num_of_nodes * sizeof(bool); + int dense_bytes_per_token = config.topk > 0 ? config.topk * static_cast(sizeof(int16_t)) : 0; + changed |= grow_to(routing_bytes_per_token, std::max(sparse_bytes_per_token, dense_bytes_per_token)); return changed; } @@ -141,9 +148,7 @@ void CustomAllgather::allocate_buffers() { void CustomAllgather::allocate_ag_buffer() { // Allocate the output buffer - auto num_of_expert = num_of_experts_per_rank * num_of_ranks_per_node * num_of_nodes; - auto gathered_elets = num_of_expert * num_of_tokens_per_rank * num_of_ranks_per_node * num_of_nodes; - auto gathered_bytes = gathered_elets * sizeof(bool); + auto gathered_bytes = routing_bytes_per_token * num_of_tokens_per_rank * num_of_ranks_per_node * num_of_nodes; allocator->allocate(&dst_buffer, gathered_bytes); if(num_of_nodes == 1) { @@ -278,4 +283,4 @@ void * CustomAllgather::get_output_buffer() { CustomAllgather::~CustomAllgather() { destroy(); -} \ No newline at end of file +} diff --git a/csrc/hybrid_ep/extension/allgather.cuh b/csrc/hybrid_ep/extension/allgather.cuh index 1ab35d5c3..a1c0d253e 100644 --- a/csrc/hybrid_ep/extension/allgather.cuh +++ b/csrc/hybrid_ep/extension/allgather.cuh @@ -38,6 +38,7 @@ private: int num_of_experts_per_rank; int num_of_tokens_per_rank; int num_of_nodes; + int routing_bytes_per_token; ExtendedMemoryAllocator *allocator; pybind11::object process_group; -}; \ No newline at end of file +}; diff --git a/csrc/hybrid_ep/jit/compiler.cu b/csrc/hybrid_ep/jit/compiler.cu index 128aba4d6..ec0f3eaad 100644 --- a/csrc/hybrid_ep/jit/compiler.cu +++ b/csrc/hybrid_ep/jit/compiler.cu @@ -208,7 +208,8 @@ std::string NVCCCompiler::get_metadata_preprocessing_code(HybridEpConfigInstance std::to_string(config.num_of_ranks_per_node) + ", " + std::to_string(config.num_of_nodes) + ", " + std::to_string(config.num_of_experts_per_rank) + ">::metadata_preprocessing<" + std::to_string(config.pad_multiple) + ", " + std::to_string(config.num_of_tokens_per_chunk_preprocessing_api) + ", " + - std::to_string(config.num_of_threads_per_block_preprocessing_api) + ", " + std::to_string(config.num_of_blocks_preprocessing_api) + R"(>; + std::to_string(config.num_of_threads_per_block_preprocessing_api) + ", " + std::to_string(config.num_of_blocks_preprocessing_api) + ", " + + std::to_string(config.topk) + R"(>; return func_ptr; } } @@ -280,7 +281,7 @@ node_rank(node_rank), local_rank(local_rank), nvcc_compiler(base_path, comm_id) void KernelCache::run_preprocess_kernel( HybridEpConfigInstance config, - const bool* input_routing_map, + const void* input_routing_map, hybrid_ep::tmp_state_t* preprocessing_tmp, hybrid_ep::tmp_state_t* preprocessing_local_experts_tmp, int32_t* sparse_to_dense_map, @@ -312,6 +313,7 @@ void KernelCache::run_preprocess_kernel( config.num_of_tokens_per_chunk_preprocessing_api, config.num_of_threads_per_block_preprocessing_api, config.num_of_blocks_preprocessing_api, + config.topk, fuse_permute_dispatch, non_blocking ); @@ -325,7 +327,7 @@ void KernelCache::run_preprocess_kernel( auto preprocessing_instance = kernel_cache[preprocess_kernel_key]; // Cast the function pointer to the correct type - using PreprocessingFuncPtr = void (*)(const bool*, hybrid_ep::tmp_state_t*, hybrid_ep::tmp_state_t*, int32_t*, bool*, bool*, int32_t*, bool*, int32_t*, int32_t*, int32_t*, int*, const int, const int, const int, const int, cudaStream_t); + using PreprocessingFuncPtr = void (*)(const void*, hybrid_ep::tmp_state_t*, hybrid_ep::tmp_state_t*, int32_t*, bool*, bool*, int32_t*, bool*, int32_t*, int32_t*, int32_t*, int*, const int, const int, const int, const int, cudaStream_t); auto func_ptr = std::any_cast(preprocessing_instance); // Run the kernel @@ -455,4 +457,3 @@ void KernelCache::run_combine_kernel( func_ptr(param, stream); } - diff --git a/csrc/hybrid_ep/jit/compiler.cuh b/csrc/hybrid_ep/jit/compiler.cuh index 36c290da6..206aa037d 100644 --- a/csrc/hybrid_ep/jit/compiler.cuh +++ b/csrc/hybrid_ep/jit/compiler.cuh @@ -70,7 +70,7 @@ public: void run_preprocess_kernel( HybridEpConfigInstance config, - const bool* input_routing_map, + const void* input_routing_map, hybrid_ep::tmp_state_t* preprocessing_tmp, hybrid_ep::tmp_state_t* preprocessing_local_experts_tmp, int32_t* sparse_to_dense_map, diff --git a/csrc/hybrid_ep/pybind_hybrid_ep.cu b/csrc/hybrid_ep/pybind_hybrid_ep.cu index 3a3186867..598d2d525 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep.cu +++ b/csrc/hybrid_ep/pybind_hybrid_ep.cu @@ -73,6 +73,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { &HybridEpConfigInstance::num_of_ranks_per_node) .def_readwrite("num_of_nodes", &HybridEpConfigInstance::num_of_nodes) .def_readwrite("pad_multiple", &HybridEpConfigInstance::pad_multiple) + .def_readwrite("topk", &HybridEpConfigInstance::topk) // Metadata-preprocessing API Config .def_readwrite("num_of_tokens_per_chunk_preprocessing_api", &HybridEpConfigInstance::num_of_tokens_per_chunk_preprocessing_api) .def_readwrite( diff --git a/deep_ep/hybrid_ep_buffer.py b/deep_ep/hybrid_ep_buffer.py index 1800e0e7c..94a43ccb0 100644 --- a/deep_ep/hybrid_ep_buffer.py +++ b/deep_ep/hybrid_ep_buffer.py @@ -6,6 +6,8 @@ import hybrid_ep_cpp import warnings +INT16_EXPERT_LIMIT = torch.iinfo(torch.int16).max + 1 + def indices_to_map( topk_idx: torch.Tensor, topk_weights: torch.Tensor, @@ -31,6 +33,22 @@ def indices_to_map( return routing_map, probs +def dense_indices_to_probs( + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + num_of_tokens: int, + num_of_experts: int, +): + probs = torch.zeros( + num_of_tokens, num_of_experts, device=topk_idx.device, dtype=torch.float32 + ) + valid = (topk_idx >= 0) & (topk_idx < num_of_experts) + safe_idx = torch.where(valid, topk_idx, torch.zeros_like(topk_idx)).long() + safe_weights = torch.where(valid, topk_weights, torch.zeros_like(topk_weights)).to(probs.dtype) + probs.scatter_add_(1, safe_idx, safe_weights) + return probs + + class HybridEPBuffer: def __init__( self, @@ -177,6 +195,7 @@ def dispatch( num_dispatched_tokens_tensor: torch.Tensor = None, num_dispatched_tokens: int = None, handle: tuple = None, + dense_routing: bool = False, ): """ Dispatch the data to the experts. @@ -186,33 +205,59 @@ def dispatch( Backward direction: combine_in_backward <- local_unpermute -> expert_mlp -> local_permute -> dispatch_in_backward + + When dense_routing=True, topk_idx is passed directly as int16 (skipping indices_to_map). + This reduces allgather size from T*E_total to T*K*2 bytes. + Dropped tokens should use -1 as sentinel (naturally ignored by range checks in the kernel). """ num_of_tokens, hidden_dim = hidden.shape - if routing_map is not None: + if dense_routing: + assert topk_idx is not None or handle is not None, \ + "topk_idx is required for dense_routing mode when handle is None" + assert num_of_experts is not None or topk_idx is None, \ + "num_of_experts is required for dense_routing mode" + assert num_of_experts is None or num_of_experts <= INT16_EXPERT_LIMIT, \ + "dense_routing uses int16 topk_idx, so num_of_experts must be <= 32768" + routing_data = None + if topk_idx is not None: + topk = topk_idx.size(-1) + routing_data = topk_idx.to(torch.int16).contiguous() + if probs is None and topk_weights is not None: + probs = dense_indices_to_probs( + topk_idx, topk_weights, num_of_tokens, num_of_experts + ) + else: + topk = 0 + elif routing_map is not None: assert routing_map.dtype == torch.bool num_of_experts = routing_map.size(-1) + topk = 0 + routing_data = routing_map else: # Generate the routing map and the probs according to the topk_idx and topk_weights. assert ( num_of_experts is not None ), "The number of experts should be provided on index-based routing." + topk = 0 if topk_idx is not None: routing_map, probs = indices_to_map( topk_idx, topk_weights, num_of_tokens, num_of_experts ) + routing_data = routing_map assert ( - handle is not None or routing_map is not None + handle is not None or routing_data is not None ), "The handle and routing_map should not be both None" if handle is None: config = self.update_template_config( hidden_dim=hidden_dim, num_of_tokens_per_rank=num_of_tokens, + topk=topk, ) handle_impl = self.runtime.metadata_preprocessing( config=config, - routing_map=routing_map, + routing_map=routing_data, num_of_tokens_per_rank=num_of_tokens, enable_permute=False, non_blocking=False, @@ -321,12 +366,14 @@ def dispatch_with_permute( # Otherwise, tokens_per_expert is copied through pinned memory so Python can derive num_permuted_tokens. non_blocking: bool = False, fuse_permute_dispatch: bool = False, + dense_routing: bool = False, # Deprecated parameters num_dispatched_tokens: int = None, use_host_meta: bool = None, ): """ Dispatch the data to the experts with permute. + When dense_routing=True, topk_idx is passed directly as int16 (skipping indices_to_map). """ if num_dispatched_tokens is not None: warnings.warn("The num_dispatched_tokens is deprecated, it will be removed in the future.") @@ -336,9 +383,28 @@ def dispatch_with_permute( with torch.cuda.nvtx.range("hybrid-ep dispatch with permute phase"): num_of_tokens_per_rank, hidden_dim = hidden.shape - if routing_map is not None: + if dense_routing: + assert topk_idx is not None or handle is not None, \ + "topk_idx is required for dense_routing mode when handle is None" + assert num_of_experts is not None or topk_idx is None, \ + "num_of_experts is required for dense_routing mode" + assert num_of_experts is None or num_of_experts <= INT16_EXPERT_LIMIT, \ + "dense_routing uses int16 topk_idx, so num_of_experts must be <= 32768" + routing_data = None + if topk_idx is not None: + topk = topk_idx.size(-1) + routing_data = topk_idx.to(torch.int16).contiguous() + if probs is None and topk_weights is not None: + probs = dense_indices_to_probs( + topk_idx, topk_weights, num_of_tokens_per_rank, num_of_experts + ) + else: + topk = 0 + elif routing_map is not None: assert routing_map.dtype == torch.bool num_of_experts = routing_map.size(-1) + topk = 0 + routing_data = routing_map else: # Generate the routing map and the probs according to the topk_idx and topk_weights. if topk_idx is not None: @@ -348,6 +414,8 @@ def dispatch_with_permute( routing_map, probs = indices_to_map( topk_idx, topk_weights, num_of_tokens_per_rank, num_of_experts ) + topk = 0 + routing_data = routing_map if non_blocking: assert num_permuted_tokens is not None and num_permuted_tokens >= 0, \ "The num_permuted_tokens is required for non-blocking mode." @@ -356,9 +424,9 @@ def dispatch_with_permute( f"num_permuted_tokens ({num_permuted_tokens}) must be a multiple of pad_multiple ({pad_multiple}) in non-blocking mode." if handle is None: - assert hidden.size(0) == routing_map.size( + assert hidden.size(0) == routing_data.size( 0 - ), "The hidden and the routing_map should have the same row number." + ), "The hidden and the routing data should have the same row number." config = self.update_template_config( hidden_dim=hidden_dim, num_of_tokens_per_rank=num_of_tokens_per_rank, @@ -366,10 +434,11 @@ def dispatch_with_permute( pad_multiple=pad_multiple, use_fp8=use_fp8, fuse_permute_dispatch=fuse_permute_dispatch, + topk=topk, ) handle_impl = self.runtime.metadata_preprocessing( config=config, - routing_map=routing_map, + routing_map=routing_data, num_of_tokens_per_rank=num_of_tokens_per_rank, num_permuted_tokens=num_permuted_tokens, pad_multiple=pad_multiple, diff --git a/tests/test_hybrid_ep.py b/tests/test_hybrid_ep.py index 791277fdb..c6b21ab76 100644 --- a/tests/test_hybrid_ep.py +++ b/tests/test_hybrid_ep.py @@ -168,6 +168,52 @@ def test_hybrid_ep_correctness(buffer: deep_ep.HybridEPBuffer, ref: TorchRef, us if dist.get_rank() == 0: print(' dispatch+combine API: PASS', flush=True) + # Dense routing dispatch correctness check + for with_probs in [True, False]: + dispatched_hidden_ref, dispatched_probs_ref, dispatched_scaling_factor_ref = ( + ref.dispatch( + hidden, routing_map, probs if with_probs else None, scaling_factor + ) + ) + ( + dispatched_hidden_dense, + dispatched_probs_dense, + dispatched_scaling_factor_dense, + handle_dense, + ) = buffer.dispatch( + hidden=hidden, scaling_factor=scaling_factor, topk_idx=topk_idx, + topk_weights=topk_weights if with_probs else None, num_of_experts=NUM_OF_EXPERTS, + dense_routing=True, + ) + + assert bitwise_equal(dispatched_hidden_ref, dispatched_hidden_dense) + if dispatched_probs_dense is not None and dispatched_probs_ref is not None: + start, end = ref._local_expert_range_per_node() + assert bitwise_equal(dispatched_probs_ref, dispatched_probs_dense[:, start:end]) + masked_probs = torch.zeros_like(dispatched_probs_dense) + masked_probs[:, start:end] = dispatched_probs_dense[:, start:end] + dispatched_probs_dense = masked_probs + + _, _, _, num_dispatched_tokens, local_expert_routing_map, _, _ = handle_dense + num_dispatched_tokens = num_dispatched_tokens.cpu() + local_expert_routing_map = local_expert_routing_map[ + : num_dispatched_tokens.item() + ] + copy_times = local_expert_routing_map.sum(dim=1) + hidden_to_combine = dispatched_hidden_dense.to(torch.bfloat16) * copy_times.unsqueeze(1) + combined_hidden, combined_probs = buffer.combine( + hidden_to_combine, dispatched_probs_dense, handle_dense + ) + combined_hidden = combined_hidden / TOPK + + assert torch.allclose(combined_hidden, hidden.to(torch.bfloat16), atol=2e-5, rtol=1e-2) + if combined_probs is not None and probs is not None: + assert bitwise_equal(combined_probs, probs) + + dist.barrier() + if dist.get_rank() == 0: + print(' dispatch+combine API (dense routing): PASS', flush=True) + # Dispatch with permute correctness check for fuse_permute_dispatch in [False, True]: for with_probs in [True, False]: From 482a958c1b1d6c422b70631c657194399b8d2c13 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Tue, 26 May 2026 14:02:25 +0800 Subject: [PATCH 5/8] [HybridEP] Optimize dense scan local expert lookup Signed-off-by: Harry Zhou --- csrc/hybrid_ep/backend/hybrid_ep_backend.cuh | 42 +++++++++----------- 1 file changed, 19 insertions(+), 23 deletions(-) diff --git a/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh b/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh index 806e1def9..6758030b2 100644 --- a/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh +++ b/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh @@ -5459,6 +5459,16 @@ __global__ void scan(const void* input_routing_data, // bool* (TOPK==0) or int1 bool token_needed_by_rank_step2[NUM_OF_RANKS_PER_NODE]; // Sparse mode needs the full routing map for local_expert_routing_map writes later. bool token_routing_map[DENSE_ROUTING ? 1 : (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE)]; + constexpr int NUM_LOCAL_EXPERT_MASK_WORDS = (NUM_OF_EXPERTS_PER_RANK + 31) / 32; + static_assert(!DENSE_ROUTING || NUM_LOCAL_EXPERT_MASK_WORDS <= 16, + "Dense scan supports up to 512 local experts per rank."); + uint32_t local_expert_mask[DENSE_ROUTING ? NUM_LOCAL_EXPERT_MASK_WORDS : 1]; + if constexpr(DENSE_ROUTING){ + #pragma unroll + for(int k = 0; k < NUM_LOCAL_EXPERT_MASK_WORDS; k++){ + local_expert_mask[k] = 0; + } + } #pragma unroll for(int j = 0; j < NUM_OF_RANKS_PER_NODE; j++){ token_needed_by_rank_step2[j] = false; @@ -5482,6 +5492,12 @@ __global__ void scan(const void* input_routing_data, // bool* (TOPK==0) or int1 int local_expert_id = expert_id - node_expert_start; int rank_within_node = local_expert_id / NUM_OF_EXPERTS_PER_RANK; token_needed_by_rank_step2[rank_within_node] = true; + int local_expert_start = local_rank * NUM_OF_EXPERTS_PER_RANK; + int local_expert_end = local_expert_start + NUM_OF_EXPERTS_PER_RANK; + if(local_expert_id >= local_expert_start && local_expert_id < local_expert_end){ + int local_expert_idx = local_expert_id - local_expert_start; + local_expert_mask[local_expert_idx / 32] |= 1u << (local_expert_idx % 32); + } } } } @@ -5539,17 +5555,8 @@ __global__ void scan(const void* input_routing_data, // bool* (TOPK==0) or int1 bool token_needed_by_this_local_expert = false; if(token_out_of_bound == 0){ if constexpr(DENSE_ROUTING){ - // Check if any topk index matches this local expert. - int target_expert = node_rank * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE) + j * NUM_OF_EXPERTS_PER_RANK + k; - // Re-load topk indices. The compiler should hoist this out of the k loop. - const int16_t* topk_base = input_topk_idx + current_token_id * TOPK; - #pragma unroll - for(int m = 0; m < TOPK; m++){ - if((int)topk_base[m] == target_expert){ - token_needed_by_this_local_expert = true; - break; - } - } + token_needed_by_this_local_expert = + ((local_expert_mask[k / 32] >> (k % 32)) & 1u) != 0; } else { token_needed_by_this_local_expert = token_routing_map[j * NUM_OF_EXPERTS_PER_RANK + k]; } @@ -5586,18 +5593,7 @@ __global__ void scan(const void* input_routing_data, // bool* (TOPK==0) or int1 bool local_expert_bools[NUM_OF_EXPERTS_PER_RANK]; #pragma unroll for(int k = 0; k < NUM_OF_EXPERTS_PER_RANK; k++){ - local_expert_bools[k] = false; - } - // Re-load topk indices (or reuse from above if we cached them). - const int16_t* topk_base = input_topk_idx + current_token_id * TOPK; - int local_expert_start = node_rank * (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE) + local_rank * NUM_OF_EXPERTS_PER_RANK; - #pragma unroll - for(int k = 0; k < TOPK; k++){ - int expert_id = (int)topk_base[k]; - int local_idx = expert_id - local_expert_start; - if(local_idx >= 0 && local_idx < NUM_OF_EXPERTS_PER_RANK){ - local_expert_bools[local_idx] = true; - } + local_expert_bools[k] = ((local_expert_mask[k / 32] >> (k % 32)) & 1u) != 0; } // Store as expert_to_rank_t chunks. expert_to_rank_t* local_expert_routing_map_store_base_addr = reinterpret_cast(local_expert_routing_map + (final_ex_scan[j] * NUM_OF_EXPERTS_PER_RANK)); From 10358abe0172427ba944c891bae47940e6494028 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Mon, 1 Jun 2026 14:02:54 +0800 Subject: [PATCH 6/8] [HybridEP] Use rank bitsets in dense scan Signed-off-by: Harry Zhou --- csrc/hybrid_ep/backend/hybrid_ep_backend.cuh | 52 ++++++++++++++------ 1 file changed, 38 insertions(+), 14 deletions(-) diff --git a/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh b/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh index 6758030b2..06b51a415 100644 --- a/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh +++ b/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh @@ -4943,6 +4943,10 @@ __global__ void scan(const void* input_routing_data, // bool* (TOPK==0) or int1 //static_assert(NUM_OF_RANKS_PER_NODE % sizeof(rank_to_node_t) == 0, "NUM_OF_RANKS_PER_NODE and rank_to_node_t mismatch"); //constexpr int RANKS_TO_NODE_REDUCE_ITER = NUM_OF_RANKS_PER_NODE / sizeof(rank_to_node_t); + constexpr int NUM_RANK_MASK_WORDS = (NUM_OF_RANKS_PER_NODE + 31) / 32; + static_assert(!DENSE_ROUTING || NUM_RANK_MASK_WORDS <= 16, + "Dense scan supports up to 512 ranks per node."); + // How do a warp save per-rank routing info back to shared memory. What's the max number of elements does each thread save back. constexpr int NUM_OF_RANKS_PER_THREAD = ((NUM_OF_RANKS_PER_NODE - 1) / WARP_SIZE) + 1; #ifdef HYBRID_EP_BUILD_PERMUTE_FUSION_ENABLE @@ -5020,11 +5024,20 @@ __global__ void scan(const void* input_routing_data, // bool* (TOPK==0) or int1 int current_token_rdma_to_attn_map_id = current_token_node_rank * rdma_to_attn_map_size_per_node + current_token_local_id; // Load routing data and compute per-rank routing info. bool token_needed_by_this_node = false; - // Per-rank routing flag array (reused in both modes). - bool token_needed_by_rank[NUM_OF_RANKS_PER_NODE]; - #pragma unroll - for(int j = 0; j < NUM_OF_RANKS_PER_NODE; j++){ - token_needed_by_rank[j] = false; + // Per-rank routing flags. Dense mode uses a compact bitset to avoid a large + // per-thread bool array when the NVL domain has many ranks. + bool token_needed_by_rank[DENSE_ROUTING ? 1 : NUM_OF_RANKS_PER_NODE]; + uint32_t rank_mask[DENSE_ROUTING ? NUM_RANK_MASK_WORDS : 1]; + if constexpr(DENSE_ROUTING){ + #pragma unroll + for(int j = 0; j < NUM_RANK_MASK_WORDS; j++){ + rank_mask[j] = 0; + } + } else { + #pragma unroll + for(int j = 0; j < NUM_OF_RANKS_PER_NODE; j++){ + token_needed_by_rank[j] = false; + } } if constexpr(DENSE_ROUTING){ @@ -5049,14 +5062,14 @@ __global__ void scan(const void* input_routing_data, // bool* (TOPK==0) or int1 if(expert_id >= node_expert_start && expert_id < node_expert_end){ int local_expert_id = expert_id - node_expert_start; int rank_within_node = local_expert_id / NUM_OF_EXPERTS_PER_RANK; - token_needed_by_rank[rank_within_node] = true; + rank_mask[rank_within_node / 32] |= 1u << (rank_within_node % 32); token_needed_by_this_node = true; } } // Accumulate per-rank sums. #pragma unroll for(int j = 0; j < NUM_OF_RANKS_PER_NODE; j++){ - if(token_needed_by_rank[j]){ + if(((rank_mask[j / 32] >> (j % 32)) & 1u) != 0){ token_routing_map_sum[j] += 1; } } @@ -5456,22 +5469,28 @@ __global__ void scan(const void* input_routing_data, // bool* (TOPK==0) or int1 // In dense mode: load TOPK int16 expert indices. // In sparse mode: load E_per_rank * R_per_node bools for local node slice. // Also compute per-rank routing flags. - bool token_needed_by_rank_step2[NUM_OF_RANKS_PER_NODE]; + bool token_needed_by_rank_step2[DENSE_ROUTING ? 1 : NUM_OF_RANKS_PER_NODE]; // Sparse mode needs the full routing map for local_expert_routing_map writes later. bool token_routing_map[DENSE_ROUTING ? 1 : (NUM_OF_EXPERTS_PER_RANK * NUM_OF_RANKS_PER_NODE)]; constexpr int NUM_LOCAL_EXPERT_MASK_WORDS = (NUM_OF_EXPERTS_PER_RANK + 31) / 32; static_assert(!DENSE_ROUTING || NUM_LOCAL_EXPERT_MASK_WORDS <= 16, "Dense scan supports up to 512 local experts per rank."); uint32_t local_expert_mask[DENSE_ROUTING ? NUM_LOCAL_EXPERT_MASK_WORDS : 1]; + uint32_t rank_mask[DENSE_ROUTING ? NUM_RANK_MASK_WORDS : 1]; if constexpr(DENSE_ROUTING){ #pragma unroll for(int k = 0; k < NUM_LOCAL_EXPERT_MASK_WORDS; k++){ local_expert_mask[k] = 0; } - } - #pragma unroll - for(int j = 0; j < NUM_OF_RANKS_PER_NODE; j++){ - token_needed_by_rank_step2[j] = false; + #pragma unroll + for(int k = 0; k < NUM_RANK_MASK_WORDS; k++){ + rank_mask[k] = 0; + } + } else { + #pragma unroll + for(int j = 0; j < NUM_OF_RANKS_PER_NODE; j++){ + token_needed_by_rank_step2[j] = false; + } } if constexpr(DENSE_ROUTING){ @@ -5491,7 +5510,7 @@ __global__ void scan(const void* input_routing_data, // bool* (TOPK==0) or int1 if(expert_id >= node_expert_start && expert_id < node_expert_end){ int local_expert_id = expert_id - node_expert_start; int rank_within_node = local_expert_id / NUM_OF_EXPERTS_PER_RANK; - token_needed_by_rank_step2[rank_within_node] = true; + rank_mask[rank_within_node / 32] |= 1u << (rank_within_node % 32); int local_expert_start = local_rank * NUM_OF_EXPERTS_PER_RANK; int local_expert_end = local_expert_start + NUM_OF_EXPERTS_PER_RANK; if(local_expert_id >= local_expert_start && local_expert_id < local_expert_end){ @@ -5533,7 +5552,12 @@ __global__ void scan(const void* input_routing_data, // bool* (TOPK==0) or int1 #pragma unroll for(int j = 0; j < NUM_OF_RANKS_PER_NODE; j++){ int32_t temp_scan = 0; - bool token_needed_by_this_rank = token_needed_by_rank_step2[j]; + bool token_needed_by_this_rank; + if constexpr(DENSE_ROUTING){ + token_needed_by_this_rank = ((rank_mask[j / 32] >> (j % 32)) & 1u) != 0; + } else { + token_needed_by_this_rank = token_needed_by_rank_step2[j]; + } // Each warp vote to create a bit mask indicating which token is needed by this rank within this tile. unsigned vote_result = __ballot_sync(~0, token_needed_by_this_rank); From 88f1b6a47a477878ea6e12f366d33f66fffcd99b Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Sun, 28 Jun 2026 21:09:54 -0700 Subject: [PATCH 7/8] [HybridEP] Standardize benchmark probability labels Signed-off-by: Harry Zhou --- tests/test_hybrid_ep.py | 42 ++++++++++++++++++++--------------------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/tests/test_hybrid_ep.py b/tests/test_hybrid_ep.py index c6b21ab76..8b4d0d1f2 100644 --- a/tests/test_hybrid_ep.py +++ b/tests/test_hybrid_ep.py @@ -432,25 +432,25 @@ def test_hybrid_ep_benchmark(buffer: deep_ep.HybridEPBuffer, group: dist.Process t = bench(lambda: buffer.dispatch(**dispatch_args))[0] _report_bw(f'dispatch ({dtype_str}, probs=True)', t, nvl_dispatch_actual, 'nvl_recv_bytes', rdma_dispatch, 'rdma_send_bytes') t = bench(lambda: buffer.combine(**combine_args))[0] - _report_bw('combine (probs=True)', t, nvl_combine, 'combine_send_bytes', rdma_combine, 'rdma_recv_bytes') + _report_bw(f'combine ({dtype_str}, probs=True)', t, nvl_combine, 'combine_send_bytes', rdma_combine, 'rdma_recv_bytes') # Non-permute (probs=False) t = bench(lambda: buffer.dispatch(**dispatch_noprob_args))[0] _report_bw(f'dispatch ({dtype_str}, probs=False)', t, nvl_dispatch_actual, 'nvl_recv_bytes', rdma_dispatch, 'rdma_send_bytes') t = bench(lambda: buffer.combine(**combine_noprob_args))[0] - _report_bw('combine (probs=False)', t, nvl_combine, 'combine_send_bytes', rdma_combine, 'rdma_recv_bytes') + _report_bw(f'combine ({dtype_str}, probs=False)', t, nvl_combine, 'combine_send_bytes', rdma_combine, 'rdma_recv_bytes') # Permute (non-fused) t = bench(lambda: buffer.dispatch_with_permute(**dispatch_wp_args))[0] - _report_bw(f'dispatch+permute ({dtype_str})', t, nvl_dispatch_actual, 'nvl_recv_bytes', rdma_dispatch, 'rdma_send_bytes') + _report_bw(f'dispatch+permute ({dtype_str}, probs=True)', t, nvl_dispatch_actual, 'nvl_recv_bytes', rdma_dispatch, 'rdma_send_bytes') t = bench(lambda: buffer.combine_with_unpermute(**combine_wp_args))[0] - _report_bw('combine+unpermute', t, nvl_combine, 'combine_send_bytes', rdma_combine, 'rdma_recv_bytes') + _report_bw(f'combine+unpermute ({dtype_str}, probs=True)', t, nvl_combine, 'combine_send_bytes', rdma_combine, 'rdma_recv_bytes') # Fused t = bench(lambda: buffer.dispatch_with_permute(**dispatch_fused_args))[0] - _report_bw(f'fused dispatch+permute ({dtype_str})', t, nvl_dispatch_actual, 'nvl_recv_bytes', rdma_dispatch, 'rdma_send_bytes') + _report_bw(f'fused dispatch+permute ({dtype_str}, probs=True)', t, nvl_dispatch_actual, 'nvl_recv_bytes', rdma_dispatch, 'rdma_send_bytes') t = bench(lambda: buffer.combine_with_unpermute(**combine_fused_args))[0] - _report_bw('fused combine+unpermute', t, nvl_combine, 'combine_send_bytes', rdma_combine, 'rdma_recv_bytes') + _report_bw(f'fused combine+unpermute ({dtype_str}, probs=True)', t, nvl_combine, 'combine_send_bytes', rdma_combine, 'rdma_recv_bytes') # ---- Kineto / nsys profiling ---- # Kineto measures pure GPU kernel time only (no CPU overhead, no d2d, no device_sync) @@ -465,7 +465,7 @@ def test_hybrid_ep_benchmark(buffer: deep_ep.HybridEPBuffer, group: dist.Process dispatch_t, combine_t = bench_kineto( lambda: (buffer.dispatch(**dispatch_args), buffer.combine(**combine_args)), kernel_names=('dispatch_kernel', 'combine_kernel'), barrier_comm_profiling=True, suppress_kineto_output=True) - _report_kineto(f'dispatch kernel ({dtype_str}, probs=True)', 'combine kernel (probs=True)', + _report_kineto(f'dispatch kernel ({dtype_str}, probs=True)', f'combine kernel ({dtype_str}, probs=True)', dispatch_t, nvl_dispatch_actual, combine_t, nvl_combine, rdma_dispatch, rdma_combine) # Non-fused kernel profiling (probs=False) @@ -473,7 +473,7 @@ def test_hybrid_ep_benchmark(buffer: deep_ep.HybridEPBuffer, group: dist.Process dispatch_t, combine_t = bench_kineto( lambda: (buffer.dispatch(**dispatch_noprob_args), buffer.combine(**combine_noprob_args)), kernel_names=('dispatch_kernel', 'combine_kernel'), barrier_comm_profiling=True, suppress_kineto_output=True) - _report_kineto(f'dispatch kernel ({dtype_str}, probs=False)', 'combine kernel (probs=False)', + _report_kineto(f'dispatch kernel ({dtype_str}, probs=False)', f'combine kernel ({dtype_str}, probs=False)', dispatch_t, nvl_dispatch_actual, combine_t, nvl_combine, rdma_dispatch, rdma_combine) # Fused kernel profiling @@ -481,7 +481,7 @@ def test_hybrid_ep_benchmark(buffer: deep_ep.HybridEPBuffer, group: dist.Process dispatch_t, combine_t = bench_kineto( lambda: (buffer.dispatch_with_permute(**dispatch_fused_args), buffer.combine_with_unpermute(**combine_fused_args)), kernel_names=('dispatch_kernel', 'combine_kernel'), barrier_comm_profiling=True, suppress_kineto_output=True) - _report_kineto(f'fused dispatch+permute kernel ({dtype_str})', 'fused combine+unpermute kernel', + _report_kineto(f'fused dispatch+permute kernel ({dtype_str}, probs=True)', f'fused combine+unpermute kernel ({dtype_str}, probs=True)', dispatch_t, nvl_dispatch_actual, combine_t, nvl_combine, rdma_dispatch, rdma_combine) # Non-fused permute/unpermute kernel profiling (isolate permute_kernel and unpermute_kernel times) @@ -505,32 +505,32 @@ def test_hybrid_ep_benchmark(buffer: deep_ep.HybridEPBuffer, group: dist.Process else: if torch.distributed.get_rank() == 0: torch.cuda.profiler.start() - with torch.cuda.nvtx.range(f"hybrid-ep dispatch ({dtype_str})"): + with torch.cuda.nvtx.range(f"hybrid-ep dispatch ({dtype_str}, probs=True)"): if rank == 0: - print(f"profile hybrid-ep dispatch ({dtype_str})", flush=True) + print(f"profile hybrid-ep dispatch ({dtype_str}, probs=True)", flush=True) nsys_dispatch_args = {'hidden': hidden, 'scaling_factor': scaling_factor, 'topk_idx': topk_idx, 'topk_weights': topk_weights, 'num_of_experts': NUM_OF_EXPERTS} bench(lambda: buffer.dispatch(**nsys_dispatch_args)) - with torch.cuda.nvtx.range("hybrid-ep combine"): + with torch.cuda.nvtx.range(f"hybrid-ep combine ({dtype_str}, probs=True)"): if rank == 0: - print(f"profile hybrid-ep combine", flush=True) + print(f"profile hybrid-ep combine ({dtype_str}, probs=True)", flush=True) bench(lambda: buffer.combine(**combine_args)) - with torch.cuda.nvtx.range(f"hybrid-ep dispatch+permute ({dtype_str})"): + with torch.cuda.nvtx.range(f"hybrid-ep dispatch+permute ({dtype_str}, probs=True)"): if rank == 0: - print(f"profile hybrid-ep dispatch+permute ({dtype_str})", flush=True) + print(f"profile hybrid-ep dispatch+permute ({dtype_str}, probs=True)", flush=True) nsys_dispatch_wp_args = {'hidden': hidden, 'scaling_factor': scaling_factor, 'routing_map': routing_map, 'probs': probs, 'pad_multiple': PAD_MULTIPLE} bench(lambda: buffer.dispatch_with_permute(**nsys_dispatch_wp_args)) - with torch.cuda.nvtx.range("hybrid-ep combine+unpermute"): + with torch.cuda.nvtx.range(f"hybrid-ep combine+unpermute ({dtype_str}, probs=True)"): if rank == 0: - print(f"profile hybrid-ep combine+unpermute", flush=True) + print(f"profile hybrid-ep combine+unpermute ({dtype_str}, probs=True)", flush=True) bench(lambda: buffer.combine_with_unpermute(**combine_wp_args)) - with torch.cuda.nvtx.range(f"hybrid-ep dispatch+permute fused"): + with torch.cuda.nvtx.range(f"hybrid-ep dispatch+permute fused ({dtype_str}, probs=True)"): if rank == 0: - print(f"profile hybrid-ep dispatch+permute fused", flush=True) + print(f"profile hybrid-ep dispatch+permute fused ({dtype_str}, probs=True)", flush=True) nsys_dispatch_fused_args = {'hidden': hidden, 'scaling_factor': scaling_factor, 'routing_map': routing_map, 'probs': probs, 'pad_multiple': PAD_MULTIPLE, 'fuse_permute_dispatch': True} bench(lambda: buffer.dispatch_with_permute(**nsys_dispatch_fused_args)) - with torch.cuda.nvtx.range("hybrid-ep combine+unpermute fused"): + with torch.cuda.nvtx.range(f"hybrid-ep combine+unpermute fused ({dtype_str}, probs=True)"): if rank == 0: - print(f"profile hybrid-ep combine+unpermute", flush=True) + print(f"profile hybrid-ep combine+unpermute fused ({dtype_str}, probs=True)", flush=True) bench(lambda: buffer.combine_with_unpermute(**combine_fused_args)) time.sleep(1) if torch.distributed.get_rank() == 0: From 18d401a0d047b8028e379f9ba54efa17d479a300 Mon Sep 17 00:00:00 2001 From: Harry Zhou Date: Sun, 28 Jun 2026 22:54:52 -0700 Subject: [PATCH 8/8] [HybridEP] Test dense topk permute routing Signed-off-by: Harry Zhou --- tests/test_hybrid_ep.py | 192 +++++++++++++++++++++++++++++++--------- 1 file changed, 151 insertions(+), 41 deletions(-) diff --git a/tests/test_hybrid_ep.py b/tests/test_hybrid_ep.py index 8b4d0d1f2..348e3c84b 100644 --- a/tests/test_hybrid_ep.py +++ b/tests/test_hybrid_ep.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: MIT # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved import argparse +import inspect import time import torch import torch.distributed as dist @@ -55,6 +56,33 @@ def bitwise_equal(a: torch.Tensor, b: torch.Tensor) -> bool: b_bytes = b.contiguous().view(torch.uint8) return torch.equal(a_bytes, b_bytes) + +def assert_bitwise_equal(name: str, ref: torch.Tensor, test: torch.Tensor, context: str = ""): + if ref is None or test is None: + return + if bitwise_equal(ref, test): + return + mismatch = (ref.contiguous().view(torch.uint8) != test.contiguous().view(torch.uint8)).sum().item() + elem_mismatch = (ref != test).sum().item() + pct = 100.0 * elem_mismatch / max(ref.numel(), 1) + msg = (f"{name} mismatch{context}: {elem_mismatch}/{ref.numel()} elements " + f"({pct:.2f}%), {mismatch} bytes differ, shape={list(ref.shape)}") + flat_ref = ref.contiguous().view(-1) + flat_test = test.contiguous().view(-1) + diff_idx = (flat_ref != flat_test).nonzero(as_tuple=True)[0][:5] + for idx in diff_idx: + i = idx.item() + msg += f"\n [{i}]: ref={flat_ref[i].item()}, got={flat_test[i].item()}" + assert False, msg + + +def supports_kwarg(fn, name: str) -> bool: + """Return whether a bound Python method accepts a keyword argument.""" + try: + return name in inspect.signature(fn).parameters + except (TypeError, ValueError): + return False + def init_tensor( hidden_dim: int, seq_len: int, @@ -104,6 +132,8 @@ def test_hybrid_ep_correctness(buffer: deep_ep.HybridEPBuffer, ref: TorchRef, us use_fp8=use_fp8, ) dtype_str = "FP8" if hidden.dtype == torch.uint8 else "BF16" + supports_dense_dispatch = supports_kwarg(buffer.dispatch, "dense_routing") + supports_dense_permute = supports_kwarg(buffer.dispatch_with_permute, "dense_routing") dist.barrier() if dist.get_rank() == 0: print(f'\n=== Correctness Check ({dtype_str}, {dist.get_world_size()} ranks) ===', flush=True) @@ -168,51 +198,61 @@ def test_hybrid_ep_correctness(buffer: deep_ep.HybridEPBuffer, ref: TorchRef, us if dist.get_rank() == 0: print(' dispatch+combine API: PASS', flush=True) - # Dense routing dispatch correctness check - for with_probs in [True, False]: - dispatched_hidden_ref, dispatched_probs_ref, dispatched_scaling_factor_ref = ( - ref.dispatch( - hidden, routing_map, probs if with_probs else None, scaling_factor + # Dense top-k routing dispatch correctness check. Older hybrid-ep branches do + # not expose dense_routing yet, so skip this block when the installed package + # lacks the keyword. + if supports_dense_dispatch: + for with_probs in [True, False]: + context = f" (dense_routing=True, with_probs={with_probs})" + dispatched_hidden_ref, dispatched_probs_ref, dispatched_scaling_factor_ref = ( + ref.dispatch( + hidden, routing_map, probs if with_probs else None, scaling_factor + ) + ) + ( + dispatched_hidden_dense, + dispatched_probs_dense, + dispatched_scaling_factor_dense, + handle_dense, + ) = buffer.dispatch( + hidden=hidden, scaling_factor=scaling_factor, topk_idx=topk_idx, + topk_weights=topk_weights if with_probs else None, num_of_experts=NUM_OF_EXPERTS, + dense_routing=True, ) - ) - ( - dispatched_hidden_dense, - dispatched_probs_dense, - dispatched_scaling_factor_dense, - handle_dense, - ) = buffer.dispatch( - hidden=hidden, scaling_factor=scaling_factor, topk_idx=topk_idx, - topk_weights=topk_weights if with_probs else None, num_of_experts=NUM_OF_EXPERTS, - dense_routing=True, - ) - - assert bitwise_equal(dispatched_hidden_ref, dispatched_hidden_dense) - if dispatched_probs_dense is not None and dispatched_probs_ref is not None: - start, end = ref._local_expert_range_per_node() - assert bitwise_equal(dispatched_probs_ref, dispatched_probs_dense[:, start:end]) - masked_probs = torch.zeros_like(dispatched_probs_dense) - masked_probs[:, start:end] = dispatched_probs_dense[:, start:end] - dispatched_probs_dense = masked_probs - _, _, _, num_dispatched_tokens, local_expert_routing_map, _, _ = handle_dense - num_dispatched_tokens = num_dispatched_tokens.cpu() - local_expert_routing_map = local_expert_routing_map[ - : num_dispatched_tokens.item() - ] - copy_times = local_expert_routing_map.sum(dim=1) - hidden_to_combine = dispatched_hidden_dense.to(torch.bfloat16) * copy_times.unsqueeze(1) - combined_hidden, combined_probs = buffer.combine( - hidden_to_combine, dispatched_probs_dense, handle_dense - ) - combined_hidden = combined_hidden / TOPK + assert_bitwise_equal("Dense dispatch hidden", dispatched_hidden_ref, dispatched_hidden_dense, context) + assert_bitwise_equal("Dense dispatch scaling_factor", dispatched_scaling_factor_ref, dispatched_scaling_factor_dense, context) + if dispatched_probs_dense is not None and dispatched_probs_ref is not None: + start, end = ref._local_expert_range_per_node() + assert_bitwise_equal("Dense dispatch probs", dispatched_probs_ref, dispatched_probs_dense[:, start:end], context) + masked_probs = torch.zeros_like(dispatched_probs_dense) + masked_probs[:, start:end] = dispatched_probs_dense[:, start:end] + dispatched_probs_dense = masked_probs + + _, _, _, num_dispatched_tokens, local_expert_routing_map, _, _ = handle_dense + num_dispatched_tokens = num_dispatched_tokens.cpu() + local_expert_routing_map = local_expert_routing_map[ + : num_dispatched_tokens.item() + ] + copy_times = local_expert_routing_map.sum(dim=1) + hidden_to_combine = dispatched_hidden_dense.to(torch.bfloat16) * copy_times.unsqueeze(1) + combined_hidden, combined_probs = buffer.combine( + hidden_to_combine, dispatched_probs_dense, handle_dense + ) + combined_hidden = combined_hidden / TOPK - assert torch.allclose(combined_hidden, hidden.to(torch.bfloat16), atol=2e-5, rtol=1e-2) - if combined_probs is not None and probs is not None: - assert bitwise_equal(combined_probs, probs) + assert torch.allclose(combined_hidden, hidden.to(torch.bfloat16), atol=2e-5, rtol=1e-2), \ + f"Dense combine hidden mismatch{context}" + if combined_probs is not None and probs is not None: + assert_bitwise_equal("Dense combine probs", probs, combined_probs, context) - dist.barrier() - if dist.get_rank() == 0: - print(' dispatch+combine API (dense routing): PASS', flush=True) + dist.barrier() + if dist.get_rank() == 0: + print(' dispatch+combine API (dense routing): PASS', flush=True) + else: + dist.barrier() + if dist.get_rank() == 0: + print(' dispatch+combine API (dense routing): SKIP (unsupported)', flush=True) # Dispatch with permute correctness check for fuse_permute_dispatch in [False, True]: @@ -304,6 +344,76 @@ def check_bitwise(name, ref, test): ) print(f' {api_name}: PASS', flush=True) + # Dense top-k routing with permute correctness check. This exercises scan + # with dense topk_idx input and enable_permute=True, i.e. the path that + # produces dense_chunk_layout and dense_to_expert_map. + if supports_dense_permute: + for fuse_permute_dispatch in [False, True]: + for with_probs in [True, False]: + context = (f" (dense_routing=True, with_probs={with_probs}, " + f"fuse_permute_dispatch={fuse_permute_dispatch})") + ( + dispatched_hidden_dense, + dispatched_probs_dense, + dispatched_scaling_factor_dense, + _tokens_per_expert_dense, + handle_dense, + ) = buffer.dispatch_with_permute( + hidden=hidden, + topk_idx=topk_idx, + topk_weights=topk_weights if with_probs else None, + num_of_experts=NUM_OF_EXPERTS, + scaling_factor=scaling_factor, + pad_multiple=PAD_MULTIPLE, + fuse_permute_dispatch=fuse_permute_dispatch, + dense_routing=True, + ) + + ( + dispatched_hidden_ref, + dispatched_probs_ref, + dispatched_scaling_factor_ref, + ) = ref.dispatch( + hidden, + routing_map, + probs if with_probs else None, + scaling_factor, + pad_multiple=PAD_MULTIPLE, + enable_permute=True, + ) + + assert_bitwise_equal("Dense dispatch+permute hidden", dispatched_hidden_ref, dispatched_hidden_dense, context) + assert_bitwise_equal("Dense dispatch+permute probs", dispatched_probs_ref, dispatched_probs_dense, context) + assert_bitwise_equal("Dense dispatch+permute scaling_factor", dispatched_scaling_factor_ref, dispatched_scaling_factor_dense, context) + + combined_hidden, combined_probs = buffer.combine_with_unpermute( + hidden=dispatched_hidden_dense.to(torch.bfloat16), + probs=dispatched_probs_dense, + handle=handle_dense, + pad_multiple=PAD_MULTIPLE, + fuse_unpermute_combine=fuse_permute_dispatch, + ) + combined_hidden = combined_hidden / TOPK + + assert torch.allclose( + combined_hidden, hidden.to(torch.bfloat16), atol=2e-5, rtol=1e-2 + ), f"Dense combine+unpermute hidden mismatch{context}" + if combined_probs is not None and probs is not None: + assert_bitwise_equal("Dense combine+unpermute probs", probs, combined_probs, context) + + dist.barrier() + if dist.get_rank() == 0: + api_name = ( + 'dispatch_with_permute + combine_with_unpermute API (dense routing, non-fused)' + if not fuse_permute_dispatch + else 'dispatch_with_permute + combine_with_unpermute API (dense routing, fused)' + ) + print(f' {api_name}: PASS', flush=True) + else: + dist.barrier() + if dist.get_rank() == 0: + print(' dispatch_with_permute + combine_with_unpermute API (dense routing): SKIP (unsupported)', flush=True) + def _gather_times(t): """Gather a scalar time from all ranks, return list on rank 0."""