diff --git a/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh b/csrc/hybrid_ep/backend/hybrid_ep_backend.cuh index ad7f94f8a..06b51a415 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 @@ -399,18 +407,23 @@ 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 +459,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 +1653,24 @@ 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)); - 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))); + 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; + 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); + } } @@ -2359,14 +2386,28 @@ 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){ - 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)), - &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)); + // 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. + 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; } if(current_src_token_id == last_src_token_id){ @@ -2578,12 +2619,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 +3127,26 @@ 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){ - 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)), - &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)); + // Sparse prob optimization: only read the source rank's E_per_rank slice. + 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; } if(current_src_token_id == last_src_token_id){ @@ -3376,12 +3433,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]; } } } @@ -4797,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, @@ -4819,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; @@ -4846,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); @@ -4874,6 +4943,10 @@ __global__ void scan(const bool* input_routing_map, //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 @@ -4949,47 +5022,108 @@ __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]; + // Load routing data and compute per-rank routing info. + bool token_needed_by_this_node = 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; + } } - // 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. - 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: 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; + rank_mask[rank_within_node / 32] |= 1u << (rank_within_node % 32); + 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(((rank_mask[j / 32] >> (j % 32)) & 1u) != 0){ + 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. @@ -5331,75 +5465,99 @@ __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){ + // 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[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 j = 0; j < ROUTING_MAP_LOAD_ITER; j++){ - *(reinterpret_cast(token_routing_map) + j) = routing_map_load_base_addr[j]; + for(int k = 0; k < NUM_LOCAL_EXPERT_MASK_WORDS; k++){ + local_expert_mask[k] = 0; + } + #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; } } - - // 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]; - #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){ + + 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; + 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){ + int local_expert_idx = local_expert_id - local_expert_start; + local_expert_mask[local_expert_idx / 32] |= 1u << (local_expert_idx % 32); + } + } } } - - // 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; + 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); @@ -5420,7 +5578,12 @@ __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){ + 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]; + } } 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); @@ -5449,11 +5612,26 @@ __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] = ((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)); + #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 @@ -5502,37 +5680,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; + } } } @@ -5599,8 +5794,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, @@ -5639,9 +5836,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 230d61f0c..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 @@ -287,13 +288,25 @@ 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); + // 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 { - 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) @@ -440,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 d21da301f..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(), @@ -301,6 +316,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 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 0fd7334a2..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,6 +198,62 @@ 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 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, + ) + + 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), \ + 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) + 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]: for with_probs in [True, False]: @@ -200,18 +286,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) @@ -246,6 +344,76 @@ def test_hybrid_ep_correctness(buffer: deep_ep.HybridEPBuffer, ref: TorchRef, us ) 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.""" @@ -309,7 +477,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)) @@ -318,6 +486,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,23 +538,29 @@ 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(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(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) @@ -391,12 +570,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)', 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) + 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)', f'combine kernel ({dtype_str}, probs=False)', dispatch_t, nvl_dispatch_actual, combine_t, nvl_combine, rdma_dispatch, rdma_combine) # Fused kernel profiling @@ -404,37 +591,56 @@ 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) + 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() - 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: @@ -446,7 +652,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 +694,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)