[HybridEP] Add dense top-k routing scan path

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.

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<int>();
-    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<bool>(), 
+      config, global_routing_map.data_ptr(),
       preprocessing_tmp, preprocessing_local_experts_tmp,
       handle.sparse_to_dense_map.data_ptr<int32_t>(),
       handle.rdma_to_attn_map.data_ptr<bool>(), handle.attn_to_rdma_map.data_ptr<bool>(),
@@ -301,6 +316,19 @@ void Executor::dispatch_core(HybridEpConfigInstance config, DispatchArgs& args)
     param.multinode_aux_ptr = reinterpret_cast<void*>(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<int>();
+            size_t probs_sz = static_cast<size_t>(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<DType>(config, param, args.fuse_permute_dispatch, args.non_blocking, args.stream);
     nvtxRangePop();  // End of dispatch_core nvtx range

csrc/hybrid_ep/extension/allgather.cu

@@ -3,6 +3,9 @@
 
 #include "allgather.cuh"
 
+#include <algorithm>
+#include <cstdint>
+
 #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<int>(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();
-}
+}

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;
-};
+};

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<PreprocessingFuncPtr>(preprocessing_instance);
 
     // Run the kernel
@@ -455,4 +457,3 @@ void KernelCache::run_combine_kernel(
     func_ptr(param, stream);
 }
 
-

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,

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(