From c2746f5db8ef54c6d6fce6e3087fcc4963f59ac1 Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Wed, 18 Feb 2026 02:37:34 +0200 Subject: [PATCH 01/22] first implementation for collecting ET --- csrc/hybrid_ep/config.cuh | 55 ++++++++++++++++++++++++++++++ csrc/hybrid_ep/pybind_hybrid_ep.cu | 52 +++++++++++++++++++++++++++- 2 files changed, 106 insertions(+), 1 deletion(-) diff --git a/csrc/hybrid_ep/config.cuh b/csrc/hybrid_ep/config.cuh index 8b1a95382..7955d0384 100644 --- a/csrc/hybrid_ep/config.cuh +++ b/csrc/hybrid_ep/config.cuh @@ -3,7 +3,9 @@ #pragma once #include +#include #include "utils.cuh" +#include // Now we support up to 72(GB200) ranks per node. // This will be used to initialize the template param_t for communication kernel. @@ -45,6 +47,25 @@ struct BufferConfig { } return valid; } + + /** Convert all attributes to a single IValue holding a tuple of IValues (ints/string). */ + c10::IValue to_ivalue_tuple() const { + std::vector elements; + elements.reserve(12); + elements.push_back(c10::IValue(static_cast(hidden_dim))); + elements.push_back(c10::IValue(static_cast(max_num_of_tokens_per_rank))); + elements.push_back(c10::IValue(static_cast(num_of_experts_per_rank))); + elements.push_back(c10::IValue(static_cast(num_of_ranks_per_node))); + elements.push_back(c10::IValue(static_cast(num_of_nodes))); + elements.push_back(c10::IValue(type_to_string(token_data_type))); + elements.push_back(c10::IValue(static_cast(num_of_blocks_preprocessing_api))); + elements.push_back(c10::IValue(static_cast(num_of_blocks_dispatch_api))); + elements.push_back(c10::IValue(static_cast(num_of_blocks_combine_api))); + elements.push_back(c10::IValue(static_cast(num_of_blocks_permute_api))); + elements.push_back(c10::IValue(static_cast(num_of_tokens_per_chunk_dispatch_api))); + elements.push_back(c10::IValue(static_cast(num_of_tokens_per_chunk_combine_api))); + return c10::IValue(c10::ivalue::Tuple::create(std::move(elements))); + } }; // Config used for hybrid-ep kernel. @@ -109,4 +130,38 @@ struct HybridEpConfigInstance { } return valid; } + + /** Convert all attributes to a single IValue holding a tuple of IValues (ints/bools). */ + c10::IValue to_ivalue_tuple() const { + std::vector elements; + elements.reserve(24); + // Hybrid-ep Config + elements.push_back(c10::IValue(static_cast(hidden_dim))); + elements.push_back(c10::IValue(static_cast(max_num_of_tokens_per_rank))); + elements.push_back(c10::IValue(static_cast(num_of_experts_per_rank))); + elements.push_back(c10::IValue(static_cast(num_of_ranks_per_node))); + elements.push_back(c10::IValue(static_cast(num_of_nodes))); + // Metadata-preprocessing API Config + elements.push_back(c10::IValue(static_cast(num_of_threads_per_block_preprocessing_api))); + elements.push_back(c10::IValue(static_cast(num_of_blocks_preprocessing_api))); + elements.push_back(c10::IValue(static_cast(num_of_blocks_permute_api))); + // Dispatch API Config + elements.push_back(c10::IValue(type_to_string(token_data_type))); + elements.push_back(c10::IValue(static_cast(num_of_stages_dispatch_api))); + elements.push_back(c10::IValue(static_cast(num_of_in_flight_s2g_dispatch_api))); + elements.push_back(c10::IValue(static_cast(num_of_tokens_per_chunk_dispatch_api))); + elements.push_back(c10::IValue(static_cast(num_of_blocks_dispatch_api))); + elements.push_back(c10::IValue(forward_dispatch_api)); + elements.push_back(c10::IValue(device_side_sync_dispatch_api)); + // Combine API Config + elements.push_back(c10::IValue(static_cast(num_of_stages_g2s_combine_api))); + elements.push_back(c10::IValue(static_cast(num_of_stages_s2g_combine_api))); + elements.push_back(c10::IValue(static_cast(num_of_tokens_per_chunk_combine_api))); + elements.push_back(c10::IValue(static_cast(num_of_tokens_per_group_combine_api))); + elements.push_back(c10::IValue(static_cast(num_of_blocks_combine_api))); + elements.push_back(c10::IValue(static_cast(num_of_additional_in_flight_s2g_combine_api))); + elements.push_back(c10::IValue(backward_combine_api)); + elements.push_back(c10::IValue(device_side_sync_combine_api)); + return c10::IValue(c10::ivalue::Tuple::create(std::move(elements))); + } }; diff --git a/csrc/hybrid_ep/pybind_hybrid_ep.cu b/csrc/hybrid_ep/pybind_hybrid_ep.cu index 812954f1f..49e83443f 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep.cu +++ b/csrc/hybrid_ep/pybind_hybrid_ep.cu @@ -9,9 +9,57 @@ #include "hybrid_ep.cuh" #include "utils.cuh" #include "config.cuh" +#include +#include +#include +#include +#include +#include +#include namespace py = pybind11; +// Wrapper for combine with logging +template +auto wrap_with_logging(Func func, const std::string& name) { + return [func, name](HybridEPBuffer& self, + HybridEpConfigInstance config, + torch::Tensor hidden, + c10::optional probs, + torch::Tensor sparse_to_dense_map, + torch::Tensor rdma_to_attn_map, + torch::Tensor attn_to_rdma_map, + int64_t num_of_tokens_per_rank, + bool with_probs) { + + auto result = (self.*func)(config, hidden, probs, sparse_to_dense_map, + rdma_to_attn_map, attn_to_rdma_map, + num_of_tokens_per_rank, with_probs); + + if (at::isRecordFunctionEnabled()) { + std::initializer_list inputList = { + config.to_ivalue_tuple(), + c10::IValue(hidden), + probs.has_value() ? c10::IValue(probs.value()) : c10::IValue(), + c10::IValue(sparse_to_dense_map), + c10::IValue(rdma_to_attn_map), + c10::IValue(attn_to_rdma_map), + c10::IValue(num_of_tokens_per_rank), + c10::IValue(with_probs), + }; + c10::IValue out0(std::get<0>(result)); + c10::IValue out1(std::get<1>(result)); + std::initializer_list outputList = { out0, out1 }; + + c10::ArrayRef inputsArray(inputList); + c10::ArrayRef outputsArray(outputList); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); + } + + return result; + }; +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "HybridEP, efficiently enable the expert-parallel communication in " "the Hopper+ architectures"; @@ -142,7 +190,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("attn_to_rdma_map"), py::arg("num_dispatched_tokens_tensor"), py::arg("num_dispatched_tokens") = std::nullopt, py::arg("num_of_tokens_per_rank"), py::arg("with_probs")) - .def("combine", &HybridEPBuffer::combine, py::kw_only(), + .def("combine", + wrap_with_logging(&HybridEPBuffer::combine, "combine"), + py::kw_only(), py::arg("config"), py::arg("hidden"), py::arg("probs") = c10::nullopt, py::arg("sparse_to_dense_map"), py::arg("rdma_to_attn_map"), py::arg("attn_to_rdma_map"), From 9c1322686ca19eb334a4d6387dec25792c7b6bd4 Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Wed, 18 Feb 2026 02:53:30 +0200 Subject: [PATCH 02/22] cleanup code --- csrc/hybrid_ep/pybind_hybrid_ep.cu | 6 +- tests/test_hybrid_ep_collect_et.py | 129 +++++++++++++++++++++++++++++ 2 files changed, 132 insertions(+), 3 deletions(-) create mode 100644 tests/test_hybrid_ep_collect_et.py diff --git a/csrc/hybrid_ep/pybind_hybrid_ep.cu b/csrc/hybrid_ep/pybind_hybrid_ep.cu index 49e83443f..fb72d2a66 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep.cu +++ b/csrc/hybrid_ep/pybind_hybrid_ep.cu @@ -19,9 +19,9 @@ namespace py = pybind11; -// Wrapper for combine with logging +// Wrapper for enable pytorch profiler tracing template -auto wrap_with_logging(Func func, const std::string& name) { +auto wrap_with_tracing(Func func, const std::string& name) { return [func, name](HybridEPBuffer& self, HybridEpConfigInstance config, torch::Tensor hidden, @@ -191,7 +191,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("num_dispatched_tokens") = std::nullopt, py::arg("num_of_tokens_per_rank"), py::arg("with_probs")) .def("combine", - wrap_with_logging(&HybridEPBuffer::combine, "combine"), + wrap_with_tracing(&HybridEPBuffer::combine, "HybridEPBuffer::combine"), py::kw_only(), py::arg("config"), py::arg("hidden"), py::arg("probs") = c10::nullopt, py::arg("sparse_to_dense_map"), diff --git a/tests/test_hybrid_ep_collect_et.py b/tests/test_hybrid_ep_collect_et.py new file mode 100644 index 000000000..6e2e83983 --- /dev/null +++ b/tests/test_hybrid_ep_collect_et.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: MIT +# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved +import argparse +from operator import truediv +import time +import torch +import torch.distributed as dist +import os +import deep_ep +from torch.profiler import ExecutionTraceObserver + +from utils import TorchRef, bench, bench_kineto, init_dist, count_rdma_send_from_routing_map + +HIDDEN_DIM = int(os.environ.get("HIDDEN_DIM", 7168)) +MAX_NUM_OF_TOKENS_PER_RANK = int(os.environ.get("MAX_NUM_OF_TOKENS_PER_RANK", 4096)) +# NUM_TOKENS_PER_RANK should equal or less than MAX_NUM_OF_TOKENS_PER_RANK +NUM_TOKENS_PER_RANK = int(os.environ.get("NUM_TOKENS_PER_RANK", 4096)) +NUM_LOCAL_EXPERTS = int(os.environ.get("NUM_LOCAL_EXPERTS", 8)) +TOPK = int(os.environ.get("TOPK", 8)) +PAD_MULTIPLE = int(os.environ.get("PAD_MULTIPLE", 32)) +ITERATIONS = int(os.environ.get("ITERATIONS", 100)) +SEED = int(os.environ.get("SEED", 42)) +USE_MNNVL = os.environ.get("USE_MNNVL", "0").strip().lower() in {"1", "true", "t", "yes", "y", "on"} +torch.manual_seed(SEED) +torch.cuda.manual_seed(SEED) +torch.cuda.manual_seed_all(SEED) +torch.backends.cudnn.deterministic = True +torch.backends.cudnn.benchmark = False +# Will be set after the process group is initialized +NUM_OF_RANKS_PER_NODE = None +NUM_OF_NODES = None +NUM_OF_EXPERTS = None + +def init_tensor( + hidden_dim: int, + seq_len: int, + topk: int, + num_of_experts: int, + use_fp8: bool = False, +): + if use_fp8: + hidden = torch.randint( + low=0, + high=256, + size=(seq_len, hidden_dim), + device="cuda", + dtype=torch.uint8, + ) + else: + hidden = torch.randn(seq_len, hidden_dim, device="cuda", dtype=torch.bfloat16) + probs = torch.zeros(seq_len, num_of_experts, device="cuda", dtype=torch.float32) + topk_idx = torch.zeros(seq_len, topk, device="cuda", dtype=torch.int64) + topk_weights = torch.zeros(seq_len, topk, device="cuda", dtype=torch.float32) + scaling_factor = torch.randn( + seq_len, hidden_dim // 128, device="cuda", dtype=torch.float32 + ) + + routing_map = torch.zeros(seq_len, num_of_experts, device="cuda", dtype=torch.bool) + + for i in range(seq_len): + # Force balanced routing for testing + # selected_experts = torch.tensor([ + # ((i * topk) % num_of_experts + val) % num_of_experts for val in range(topk) + # ], device="cuda") + selected_experts = torch.randperm(num_of_experts, device="cuda")[:topk] + topk_idx[i, :] = selected_experts.to(torch.int64) + topk_weights[i, :] = torch.ones(topk, device="cuda", dtype=torch.float32) + routing_map[i, selected_experts] = True + probs[i, selected_experts] = topk_weights[i, :] + + return hidden, probs, scaling_factor, routing_map, topk_idx, topk_weights + + +def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + _, _, group = init_dist(local_rank, num_local_ranks) + + # Set missing global vars + global NUM_OF_RANKS_PER_NODE, NUM_OF_NODES, NUM_OF_EXPERTS + if USE_MNNVL: + NUM_OF_RANKS_PER_NODE = group.size() + NUM_OF_NODES = 1 + NUM_OF_EXPERTS = NUM_LOCAL_EXPERTS * NUM_OF_RANKS_PER_NODE * NUM_OF_NODES + else: + NUM_OF_RANKS_PER_NODE = args.num_processes + NUM_OF_NODES = group.size() // NUM_OF_RANKS_PER_NODE + NUM_OF_EXPERTS = NUM_LOCAL_EXPERTS * NUM_OF_RANKS_PER_NODE * NUM_OF_NODES + + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + buffer = deep_ep.HybridEPBuffer( + group=group, + hidden_dim=HIDDEN_DIM, + max_num_of_tokens_per_rank=MAX_NUM_OF_TOKENS_PER_RANK, + num_local_experts=NUM_LOCAL_EXPERTS, + use_fp8=True + ) + + hidden, probs, scaling_factor, routing_map, topk_idx, topk_weights = init_tensor( + hidden_dim=HIDDEN_DIM, + seq_len=NUM_TOKENS_PER_RANK, + topk=TOPK, + num_of_experts=NUM_OF_EXPERTS, + use_fp8=truediv, + ) + + et = ExecutionTraceObserver().register_callback(f"et_{dist.get_rank()}.json") + et.start() + + 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) + ) + # The combine only support bf16 + dispatched_hidden_bf16 = dispatched_hidden.to(torch.bfloat16) + dispatched_probs = None + _, _ = buffer.combine(dispatched_hidden_bf16, dispatched_probs, handle) + + et.stop() + et.unregister_callback() + + time.sleep(10) + dist.barrier() + dist.destroy_process_group() + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description='Test intranode EP kernels') + parser.add_argument('--num-processes', type=int, default=4, + help='Number of processes to spawn (default: 4)') + args = parser.parse_args() + torch.multiprocessing.spawn(test_main, args=(args.num_processes, args), nprocs=args.num_processes) From b0dc913e97b1a1b5eb0655dd99df651c35a4ffc1 Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Wed, 18 Feb 2026 04:35:38 +0200 Subject: [PATCH 03/22] use c10::impl::GenericList --- csrc/hybrid_ep/pybind_hybrid_ep.cu | 28 +++++++++++++++------------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/csrc/hybrid_ep/pybind_hybrid_ep.cu b/csrc/hybrid_ep/pybind_hybrid_ep.cu index fb72d2a66..a03a46fa8 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep.cu +++ b/csrc/hybrid_ep/pybind_hybrid_ep.cu @@ -37,19 +37,21 @@ auto wrap_with_tracing(Func func, const std::string& name) { num_of_tokens_per_rank, with_probs); if (at::isRecordFunctionEnabled()) { - std::initializer_list inputList = { - config.to_ivalue_tuple(), - c10::IValue(hidden), - probs.has_value() ? c10::IValue(probs.value()) : c10::IValue(), - c10::IValue(sparse_to_dense_map), - c10::IValue(rdma_to_attn_map), - c10::IValue(attn_to_rdma_map), - c10::IValue(num_of_tokens_per_rank), - c10::IValue(with_probs), - }; - c10::IValue out0(std::get<0>(result)); - c10::IValue out1(std::get<1>(result)); - std::initializer_list outputList = { out0, out1 }; + c10::impl::GenericList inputList(c10::AnyType::get()); + inputList.reserve(8); + inputList.emplace_back(config.to_ivalue_tuple()); + inputList.emplace_back(c10::IValue(hidden)); + inputList.emplace_back(probs.has_value() ? c10::IValue(probs.value()) : c10::IValue()); + inputList.emplace_back(c10::IValue(sparse_to_dense_map)); + inputList.emplace_back(c10::IValue(rdma_to_attn_map)); + inputList.emplace_back(c10::IValue(attn_to_rdma_map)); + inputList.emplace_back(c10::IValue(num_of_tokens_per_rank)); + inputList.emplace_back(c10::IValue(with_probs)); + + c10::impl::GenericList outputList(c10::AnyType::get()); + outputList.reserve(2); + outputList.emplace_back(std::get<0>(result)); + outputList.emplace_back(std::get<1>(result)); c10::ArrayRef inputsArray(inputList); c10::ArrayRef outputsArray(outputList); From bb5bc8ff4199d35b038a067a3701d134277b2b76 Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Thu, 19 Feb 2026 01:03:28 +0200 Subject: [PATCH 04/22] support all APIs --- csrc/hybrid_ep/pybind_hybrid_ep.cu | 228 ++++++++++++++++++++++++++--- tests/test_hybrid_ep_collect_et.py | 6 +- 2 files changed, 207 insertions(+), 27 deletions(-) diff --git a/csrc/hybrid_ep/pybind_hybrid_ep.cu b/csrc/hybrid_ep/pybind_hybrid_ep.cu index a03a46fa8..70fc1c740 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep.cu +++ b/csrc/hybrid_ep/pybind_hybrid_ep.cu @@ -19,9 +19,32 @@ namespace py = pybind11; -// Wrapper for enable pytorch profiler tracing +// Convert a single argument to c10::IValue for tracing +template +c10::IValue to_ivalue(T&& arg) { + if constexpr (std::is_same_v, HybridEpConfigInstance>) { + return std::forward(arg).to_ivalue_tuple(); + } else if constexpr (std::is_same_v, BufferConfig>) { + return std::forward(arg).to_ivalue_tuple(); + } else if constexpr (std::is_same_v, c10::optional>) { + return arg.has_value() ? c10::IValue(arg.value()) : c10::IValue(); + } else if constexpr (std::is_same_v, c10::optional>) { + return arg.has_value() ? c10::IValue(arg.value()) : c10::IValue(); + } else { + return c10::IValue(std::forward(arg)); + } +} + +template +static void record_result(c10::impl::GenericList& outputList, Result&& result) { + std::apply([&outputList](const auto&... elem) { + (outputList.emplace_back(c10::IValue(elem)), ...); + }, result); +} + +// Wrapper for combine template -auto wrap_with_tracing(Func func, const std::string& name) { +auto wrap_with_tracing_combine(Func func, const std::string& name) { return [func, name](HybridEPBuffer& self, HybridEpConfigInstance config, torch::Tensor hidden, @@ -31,37 +54,188 @@ auto wrap_with_tracing(Func func, const std::string& name) { torch::Tensor attn_to_rdma_map, int64_t num_of_tokens_per_rank, bool with_probs) { - auto result = (self.*func)(config, hidden, probs, sparse_to_dense_map, rdma_to_attn_map, attn_to_rdma_map, num_of_tokens_per_rank, with_probs); - if (at::isRecordFunctionEnabled()) { c10::impl::GenericList inputList(c10::AnyType::get()); - inputList.reserve(8); - inputList.emplace_back(config.to_ivalue_tuple()); - inputList.emplace_back(c10::IValue(hidden)); - inputList.emplace_back(probs.has_value() ? c10::IValue(probs.value()) : c10::IValue()); - inputList.emplace_back(c10::IValue(sparse_to_dense_map)); - inputList.emplace_back(c10::IValue(rdma_to_attn_map)); - inputList.emplace_back(c10::IValue(attn_to_rdma_map)); - inputList.emplace_back(c10::IValue(num_of_tokens_per_rank)); - inputList.emplace_back(c10::IValue(with_probs)); + inputList.emplace_back(to_ivalue(config)); + inputList.emplace_back(to_ivalue(hidden)); + inputList.emplace_back(to_ivalue(probs)); + inputList.emplace_back(to_ivalue(sparse_to_dense_map)); + inputList.emplace_back(to_ivalue(rdma_to_attn_map)); + inputList.emplace_back(to_ivalue(attn_to_rdma_map)); + inputList.emplace_back(to_ivalue(num_of_tokens_per_rank)); + inputList.emplace_back(to_ivalue(with_probs)); + c10::impl::GenericList outputList(c10::AnyType::get()); + record_result(outputList, result); + c10::ArrayRef inputsArray(inputList); + c10::ArrayRef outputsArray(outputList); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); + } + return result; + }; +} +// Wrapper for dispatch (returns 3-tuple) +template +auto wrap_with_tracing_dispatch(Func func, const std::string& name) { + return [func, name](HybridEPBuffer& self, + HybridEpConfigInstance config, + torch::Tensor hidden, + c10::optional probs, + c10::optional scaling_factor, + torch::Tensor sparse_to_dense_map, + torch::Tensor rdma_to_attn_map, + torch::Tensor attn_to_rdma_map, + c10::optional num_dispatched_tokens_tensor, + c10::optional num_dispatched_tokens, + int64_t num_of_tokens_per_rank, + bool with_probs) { + auto result = (self.*func)(config, hidden, probs, scaling_factor, + sparse_to_dense_map, rdma_to_attn_map, attn_to_rdma_map, + num_dispatched_tokens_tensor, num_dispatched_tokens, + num_of_tokens_per_rank, with_probs); + if (at::isRecordFunctionEnabled()) { + c10::impl::GenericList inputList(c10::AnyType::get()); + inputList.emplace_back(to_ivalue(config)); + inputList.emplace_back(to_ivalue(hidden)); + inputList.emplace_back(to_ivalue(probs)); + inputList.emplace_back(to_ivalue(scaling_factor)); + inputList.emplace_back(to_ivalue(sparse_to_dense_map)); + inputList.emplace_back(to_ivalue(rdma_to_attn_map)); + inputList.emplace_back(to_ivalue(attn_to_rdma_map)); + inputList.emplace_back(to_ivalue(num_dispatched_tokens_tensor)); + inputList.emplace_back(to_ivalue(num_dispatched_tokens)); + inputList.emplace_back(to_ivalue(num_of_tokens_per_rank)); + inputList.emplace_back(to_ivalue(with_probs)); c10::impl::GenericList outputList(c10::AnyType::get()); - outputList.reserve(2); - outputList.emplace_back(std::get<0>(result)); - outputList.emplace_back(std::get<1>(result)); + record_result(outputList, result); + c10::ArrayRef inputsArray(inputList); + c10::ArrayRef outputsArray(outputList); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); + } + return result; + }; +} - c10::ArrayRef inputsArray(inputList); - c10::ArrayRef outputsArray(outputList); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); +// Wrapper for dispatch_with_permute (returns 6-tuple) +template +auto wrap_with_tracing_dispatch_with_permute(Func func, const std::string& name) { + return [func, name](HybridEPBuffer& self, + HybridEpConfigInstance config, + torch::Tensor hidden, + c10::optional probs, + c10::optional scaling_factor, + torch::Tensor sparse_to_dense_map, + torch::Tensor rdma_to_attn_map, + torch::Tensor attn_to_rdma_map, + c10::optional num_dispatched_tokens_tensor, + c10::optional local_expert_routing_map, + c10::optional row_id_map, + c10::optional num_permuted_tokens, + int64_t num_of_tokens_per_rank, + c10::optional pad_multiple, + bool non_blocking, + bool with_probs) { + auto result = (self.*func)(config, hidden, probs, scaling_factor, + sparse_to_dense_map, rdma_to_attn_map, attn_to_rdma_map, + num_dispatched_tokens_tensor, local_expert_routing_map, row_id_map, + num_permuted_tokens, num_of_tokens_per_rank, pad_multiple, + non_blocking, with_probs); + if (at::isRecordFunctionEnabled()) { + c10::impl::GenericList inputList(c10::AnyType::get()); + inputList.emplace_back(to_ivalue(config)); + inputList.emplace_back(to_ivalue(hidden)); + inputList.emplace_back(to_ivalue(probs)); + inputList.emplace_back(to_ivalue(scaling_factor)); + inputList.emplace_back(to_ivalue(sparse_to_dense_map)); + inputList.emplace_back(to_ivalue(rdma_to_attn_map)); + inputList.emplace_back(to_ivalue(attn_to_rdma_map)); + inputList.emplace_back(to_ivalue(num_dispatched_tokens_tensor)); + inputList.emplace_back(to_ivalue(local_expert_routing_map)); + inputList.emplace_back(to_ivalue(row_id_map)); + inputList.emplace_back(to_ivalue(num_permuted_tokens)); + inputList.emplace_back(to_ivalue(num_of_tokens_per_rank)); + inputList.emplace_back(to_ivalue(pad_multiple)); + inputList.emplace_back(to_ivalue(non_blocking)); + inputList.emplace_back(to_ivalue(with_probs)); + c10::impl::GenericList outputList(c10::AnyType::get()); + record_result(outputList, result); + c10::ArrayRef inputsArray(inputList); + c10::ArrayRef outputsArray(outputList); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); } + return result; + }; +} +// Wrapper for combine_with_unpermute (returns 2-tuple) +template +auto wrap_with_tracing_combine_with_unpermute(Func func, const std::string& name) { + return [func, name](HybridEPBuffer& self, + HybridEpConfigInstance config, + torch::Tensor hidden, + c10::optional probs, + torch::Tensor sparse_to_dense_map, + torch::Tensor rdma_to_attn_map, + torch::Tensor attn_to_rdma_map, + c10::optional num_dispatched_tokens_tensor, + c10::optional row_id_map, + int64_t num_of_tokens_per_rank, + c10::optional pad_multiple, + bool with_probs) { + auto result = (self.*func)(config, hidden, probs, sparse_to_dense_map, + rdma_to_attn_map, attn_to_rdma_map, + num_dispatched_tokens_tensor, row_id_map, + num_of_tokens_per_rank, pad_multiple, with_probs); + if (at::isRecordFunctionEnabled()) { + c10::impl::GenericList inputList(c10::AnyType::get()); + inputList.emplace_back(to_ivalue(config)); + inputList.emplace_back(to_ivalue(hidden)); + inputList.emplace_back(to_ivalue(probs)); + inputList.emplace_back(to_ivalue(sparse_to_dense_map)); + inputList.emplace_back(to_ivalue(rdma_to_attn_map)); + inputList.emplace_back(to_ivalue(attn_to_rdma_map)); + inputList.emplace_back(to_ivalue(num_dispatched_tokens_tensor)); + inputList.emplace_back(to_ivalue(row_id_map)); + inputList.emplace_back(to_ivalue(num_of_tokens_per_rank)); + inputList.emplace_back(to_ivalue(pad_multiple)); + inputList.emplace_back(to_ivalue(with_probs)); + c10::impl::GenericList outputList(c10::AnyType::get()); + record_result(outputList, result); + c10::ArrayRef inputsArray(inputList); + c10::ArrayRef outputsArray(outputList); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); + } return result; }; } +// Wrapper for HybridEPBuffer constructor with tracing +static HybridEPBuffer* wrap_hybrid_ep_buffer_init( + py::object process_group, BufferConfig config, int local_rank, int node_rank, + int group_size, std::string base_path, bool load_cached_kernels, + bool use_shared_buffer, bool enable_custom_allgather) { + if (at::isRecordFunctionEnabled()) { + c10::impl::GenericList inputList(c10::AnyType::get()); + inputList.emplace_back(to_ivalue(config)); + inputList.emplace_back(to_ivalue(local_rank)); + inputList.emplace_back(to_ivalue(node_rank)); + inputList.emplace_back(to_ivalue(group_size)); + inputList.emplace_back(to_ivalue(base_path)); + inputList.emplace_back(to_ivalue(load_cached_kernels)); + inputList.emplace_back(to_ivalue(use_shared_buffer)); + inputList.emplace_back(to_ivalue(enable_custom_allgather)); + c10::impl::GenericList outputList(c10::AnyType::get()); + c10::ArrayRef inputsArray(inputList); + c10::ArrayRef outputsArray(outputList); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::__init__", inputsArray, outputsArray); + } + return new HybridEPBuffer(process_group, config, local_rank, node_rank, group_size, + base_path, load_cached_kernels, use_shared_buffer, enable_custom_allgather); +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "HybridEP, efficiently enable the expert-parallel communication in " "the Hopper+ architectures"; @@ -171,7 +345,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { }); pybind11::class_(m, "HybridEPBuffer") - .def(py::init(), + .def(py::init(&wrap_hybrid_ep_buffer_init), py::arg("process_group"), py::arg("config"), py::arg("local_rank"), @@ -184,7 +358,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("update_buffer", &HybridEPBuffer::update_buffer, py::arg("config")) .def("metadata_preprocessing", &HybridEPBuffer::metadata_preprocessing, py::kw_only(), py::arg("config"), py::arg("routing_map"), py::arg("num_of_tokens_per_rank"), py::arg("non_blocking") = false) - .def("dispatch", &HybridEPBuffer::dispatch, py::kw_only(), + .def("dispatch", + wrap_with_tracing_dispatch(&HybridEPBuffer::dispatch, "HybridEPBuffer::dispatch"), + py::kw_only(), py::arg("config"), py::arg("hidden"), py::arg("probs") = c10::nullopt, py::arg("scaling_factor") = c10::nullopt, @@ -193,14 +369,16 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("num_dispatched_tokens") = std::nullopt, py::arg("num_of_tokens_per_rank"), py::arg("with_probs")) .def("combine", - wrap_with_tracing(&HybridEPBuffer::combine, "HybridEPBuffer::combine"), + wrap_with_tracing_combine(&HybridEPBuffer::combine, "HybridEPBuffer::combine"), py::kw_only(), py::arg("config"), py::arg("hidden"), py::arg("probs") = c10::nullopt, py::arg("sparse_to_dense_map"), py::arg("rdma_to_attn_map"), py::arg("attn_to_rdma_map"), py::arg("num_of_tokens_per_rank"), py::arg("with_probs")) - .def("dispatch_with_permute", &HybridEPBuffer::dispatch_with_permute, py::kw_only(), + .def("dispatch_with_permute", + wrap_with_tracing_dispatch_with_permute(&HybridEPBuffer::dispatch_with_permute, "HybridEPBuffer::dispatch_with_permute"), + py::kw_only(), py::arg("config"), py::arg("hidden"), py::arg("probs") = c10::nullopt, py::arg("scaling_factor") = c10::nullopt, @@ -210,7 +388,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("num_permuted_tokens") = std::nullopt, py::arg("num_of_tokens_per_rank"), py::arg("pad_multiple") = std::nullopt, py::arg("non_blocking") = false, py::arg("with_probs") = false) - .def("combine_with_unpermute", &HybridEPBuffer::combine_with_unpermute, py::kw_only(), + .def("combine_with_unpermute", + wrap_with_tracing_combine_with_unpermute(&HybridEPBuffer::combine_with_unpermute, "HybridEPBuffer::combine_with_unpermute"), + py::kw_only(), py::arg("config"), py::arg("hidden"), py::arg("probs") = c10::nullopt, py::arg("sparse_to_dense_map"), py::arg("rdma_to_attn_map"), diff --git a/tests/test_hybrid_ep_collect_et.py b/tests/test_hybrid_ep_collect_et.py index 6e2e83983..e5eb8a2a6 100644 --- a/tests/test_hybrid_ep_collect_et.py +++ b/tests/test_hybrid_ep_collect_et.py @@ -87,6 +87,9 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): stream = torch.cuda.Stream() with torch.cuda.stream(stream): + et = ExecutionTraceObserver().register_callback(f"et_{dist.get_rank()}.json") + et.start() + buffer = deep_ep.HybridEPBuffer( group=group, hidden_dim=HIDDEN_DIM, @@ -103,9 +106,6 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): use_fp8=truediv, ) - et = ExecutionTraceObserver().register_callback(f"et_{dist.get_rank()}.json") - et.start() - 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 f8480e7a2cd87997a06aa0fcdf2d2045794313d2 Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Thu, 19 Feb 2026 01:20:52 +0200 Subject: [PATCH 05/22] put new APIs into a new file --- csrc/hybrid_ep/pybind_hybrid_ep.cu | 230 +---------------------------- 1 file changed, 6 insertions(+), 224 deletions(-) diff --git a/csrc/hybrid_ep/pybind_hybrid_ep.cu b/csrc/hybrid_ep/pybind_hybrid_ep.cu index 70fc1c740..3bd4f7752 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep.cu +++ b/csrc/hybrid_ep/pybind_hybrid_ep.cu @@ -9,233 +9,15 @@ #include "hybrid_ep.cuh" #include "utils.cuh" #include "config.cuh" +#include "pybind_hybrid_ep_utils.cuh" #include #include #include #include #include -#include -#include namespace py = pybind11; -// Convert a single argument to c10::IValue for tracing -template -c10::IValue to_ivalue(T&& arg) { - if constexpr (std::is_same_v, HybridEpConfigInstance>) { - return std::forward(arg).to_ivalue_tuple(); - } else if constexpr (std::is_same_v, BufferConfig>) { - return std::forward(arg).to_ivalue_tuple(); - } else if constexpr (std::is_same_v, c10::optional>) { - return arg.has_value() ? c10::IValue(arg.value()) : c10::IValue(); - } else if constexpr (std::is_same_v, c10::optional>) { - return arg.has_value() ? c10::IValue(arg.value()) : c10::IValue(); - } else { - return c10::IValue(std::forward(arg)); - } -} - -template -static void record_result(c10::impl::GenericList& outputList, Result&& result) { - std::apply([&outputList](const auto&... elem) { - (outputList.emplace_back(c10::IValue(elem)), ...); - }, result); -} - -// Wrapper for combine -template -auto wrap_with_tracing_combine(Func func, const std::string& name) { - return [func, name](HybridEPBuffer& self, - HybridEpConfigInstance config, - torch::Tensor hidden, - c10::optional probs, - torch::Tensor sparse_to_dense_map, - torch::Tensor rdma_to_attn_map, - torch::Tensor attn_to_rdma_map, - int64_t num_of_tokens_per_rank, - bool with_probs) { - auto result = (self.*func)(config, hidden, probs, sparse_to_dense_map, - rdma_to_attn_map, attn_to_rdma_map, - num_of_tokens_per_rank, with_probs); - if (at::isRecordFunctionEnabled()) { - c10::impl::GenericList inputList(c10::AnyType::get()); - inputList.emplace_back(to_ivalue(config)); - inputList.emplace_back(to_ivalue(hidden)); - inputList.emplace_back(to_ivalue(probs)); - inputList.emplace_back(to_ivalue(sparse_to_dense_map)); - inputList.emplace_back(to_ivalue(rdma_to_attn_map)); - inputList.emplace_back(to_ivalue(attn_to_rdma_map)); - inputList.emplace_back(to_ivalue(num_of_tokens_per_rank)); - inputList.emplace_back(to_ivalue(with_probs)); - c10::impl::GenericList outputList(c10::AnyType::get()); - record_result(outputList, result); - c10::ArrayRef inputsArray(inputList); - c10::ArrayRef outputsArray(outputList); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); - } - return result; - }; -} - -// Wrapper for dispatch (returns 3-tuple) -template -auto wrap_with_tracing_dispatch(Func func, const std::string& name) { - return [func, name](HybridEPBuffer& self, - HybridEpConfigInstance config, - torch::Tensor hidden, - c10::optional probs, - c10::optional scaling_factor, - torch::Tensor sparse_to_dense_map, - torch::Tensor rdma_to_attn_map, - torch::Tensor attn_to_rdma_map, - c10::optional num_dispatched_tokens_tensor, - c10::optional num_dispatched_tokens, - int64_t num_of_tokens_per_rank, - bool with_probs) { - auto result = (self.*func)(config, hidden, probs, scaling_factor, - sparse_to_dense_map, rdma_to_attn_map, attn_to_rdma_map, - num_dispatched_tokens_tensor, num_dispatched_tokens, - num_of_tokens_per_rank, with_probs); - if (at::isRecordFunctionEnabled()) { - c10::impl::GenericList inputList(c10::AnyType::get()); - inputList.emplace_back(to_ivalue(config)); - inputList.emplace_back(to_ivalue(hidden)); - inputList.emplace_back(to_ivalue(probs)); - inputList.emplace_back(to_ivalue(scaling_factor)); - inputList.emplace_back(to_ivalue(sparse_to_dense_map)); - inputList.emplace_back(to_ivalue(rdma_to_attn_map)); - inputList.emplace_back(to_ivalue(attn_to_rdma_map)); - inputList.emplace_back(to_ivalue(num_dispatched_tokens_tensor)); - inputList.emplace_back(to_ivalue(num_dispatched_tokens)); - inputList.emplace_back(to_ivalue(num_of_tokens_per_rank)); - inputList.emplace_back(to_ivalue(with_probs)); - c10::impl::GenericList outputList(c10::AnyType::get()); - record_result(outputList, result); - c10::ArrayRef inputsArray(inputList); - c10::ArrayRef outputsArray(outputList); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); - } - return result; - }; -} - -// Wrapper for dispatch_with_permute (returns 6-tuple) -template -auto wrap_with_tracing_dispatch_with_permute(Func func, const std::string& name) { - return [func, name](HybridEPBuffer& self, - HybridEpConfigInstance config, - torch::Tensor hidden, - c10::optional probs, - c10::optional scaling_factor, - torch::Tensor sparse_to_dense_map, - torch::Tensor rdma_to_attn_map, - torch::Tensor attn_to_rdma_map, - c10::optional num_dispatched_tokens_tensor, - c10::optional local_expert_routing_map, - c10::optional row_id_map, - c10::optional num_permuted_tokens, - int64_t num_of_tokens_per_rank, - c10::optional pad_multiple, - bool non_blocking, - bool with_probs) { - auto result = (self.*func)(config, hidden, probs, scaling_factor, - sparse_to_dense_map, rdma_to_attn_map, attn_to_rdma_map, - num_dispatched_tokens_tensor, local_expert_routing_map, row_id_map, - num_permuted_tokens, num_of_tokens_per_rank, pad_multiple, - non_blocking, with_probs); - if (at::isRecordFunctionEnabled()) { - c10::impl::GenericList inputList(c10::AnyType::get()); - inputList.emplace_back(to_ivalue(config)); - inputList.emplace_back(to_ivalue(hidden)); - inputList.emplace_back(to_ivalue(probs)); - inputList.emplace_back(to_ivalue(scaling_factor)); - inputList.emplace_back(to_ivalue(sparse_to_dense_map)); - inputList.emplace_back(to_ivalue(rdma_to_attn_map)); - inputList.emplace_back(to_ivalue(attn_to_rdma_map)); - inputList.emplace_back(to_ivalue(num_dispatched_tokens_tensor)); - inputList.emplace_back(to_ivalue(local_expert_routing_map)); - inputList.emplace_back(to_ivalue(row_id_map)); - inputList.emplace_back(to_ivalue(num_permuted_tokens)); - inputList.emplace_back(to_ivalue(num_of_tokens_per_rank)); - inputList.emplace_back(to_ivalue(pad_multiple)); - inputList.emplace_back(to_ivalue(non_blocking)); - inputList.emplace_back(to_ivalue(with_probs)); - c10::impl::GenericList outputList(c10::AnyType::get()); - record_result(outputList, result); - c10::ArrayRef inputsArray(inputList); - c10::ArrayRef outputsArray(outputList); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); - } - return result; - }; -} - -// Wrapper for combine_with_unpermute (returns 2-tuple) -template -auto wrap_with_tracing_combine_with_unpermute(Func func, const std::string& name) { - return [func, name](HybridEPBuffer& self, - HybridEpConfigInstance config, - torch::Tensor hidden, - c10::optional probs, - torch::Tensor sparse_to_dense_map, - torch::Tensor rdma_to_attn_map, - torch::Tensor attn_to_rdma_map, - c10::optional num_dispatched_tokens_tensor, - c10::optional row_id_map, - int64_t num_of_tokens_per_rank, - c10::optional pad_multiple, - bool with_probs) { - auto result = (self.*func)(config, hidden, probs, sparse_to_dense_map, - rdma_to_attn_map, attn_to_rdma_map, - num_dispatched_tokens_tensor, row_id_map, - num_of_tokens_per_rank, pad_multiple, with_probs); - if (at::isRecordFunctionEnabled()) { - c10::impl::GenericList inputList(c10::AnyType::get()); - inputList.emplace_back(to_ivalue(config)); - inputList.emplace_back(to_ivalue(hidden)); - inputList.emplace_back(to_ivalue(probs)); - inputList.emplace_back(to_ivalue(sparse_to_dense_map)); - inputList.emplace_back(to_ivalue(rdma_to_attn_map)); - inputList.emplace_back(to_ivalue(attn_to_rdma_map)); - inputList.emplace_back(to_ivalue(num_dispatched_tokens_tensor)); - inputList.emplace_back(to_ivalue(row_id_map)); - inputList.emplace_back(to_ivalue(num_of_tokens_per_rank)); - inputList.emplace_back(to_ivalue(pad_multiple)); - inputList.emplace_back(to_ivalue(with_probs)); - c10::impl::GenericList outputList(c10::AnyType::get()); - record_result(outputList, result); - c10::ArrayRef inputsArray(inputList); - c10::ArrayRef outputsArray(outputList); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); - } - return result; - }; -} - -// Wrapper for HybridEPBuffer constructor with tracing -static HybridEPBuffer* wrap_hybrid_ep_buffer_init( - py::object process_group, BufferConfig config, int local_rank, int node_rank, - int group_size, std::string base_path, bool load_cached_kernels, - bool use_shared_buffer, bool enable_custom_allgather) { - if (at::isRecordFunctionEnabled()) { - c10::impl::GenericList inputList(c10::AnyType::get()); - inputList.emplace_back(to_ivalue(config)); - inputList.emplace_back(to_ivalue(local_rank)); - inputList.emplace_back(to_ivalue(node_rank)); - inputList.emplace_back(to_ivalue(group_size)); - inputList.emplace_back(to_ivalue(base_path)); - inputList.emplace_back(to_ivalue(load_cached_kernels)); - inputList.emplace_back(to_ivalue(use_shared_buffer)); - inputList.emplace_back(to_ivalue(enable_custom_allgather)); - c10::impl::GenericList outputList(c10::AnyType::get()); - c10::ArrayRef inputsArray(inputList); - c10::ArrayRef outputsArray(outputList); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::__init__", inputsArray, outputsArray); - } - return new HybridEPBuffer(process_group, config, local_rank, node_rank, group_size, - base_path, load_cached_kernels, use_shared_buffer, enable_custom_allgather); -} - PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "HybridEP, efficiently enable the expert-parallel communication in " "the Hopper+ architectures"; @@ -345,7 +127,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { }); pybind11::class_(m, "HybridEPBuffer") - .def(py::init(&wrap_hybrid_ep_buffer_init), + .def(py::init(&hybrid_ep_buffer_init), py::arg("process_group"), py::arg("config"), py::arg("local_rank"), @@ -359,7 +141,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("metadata_preprocessing", &HybridEPBuffer::metadata_preprocessing, py::kw_only(), py::arg("config"), py::arg("routing_map"), py::arg("num_of_tokens_per_rank"), py::arg("non_blocking") = false) .def("dispatch", - wrap_with_tracing_dispatch(&HybridEPBuffer::dispatch, "HybridEPBuffer::dispatch"), + hybrid_ep_buffer_dispatch(&HybridEPBuffer::dispatch, "HybridEPBuffer::dispatch"), py::kw_only(), py::arg("config"), py::arg("hidden"), py::arg("probs") = c10::nullopt, @@ -369,7 +151,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("num_dispatched_tokens") = std::nullopt, py::arg("num_of_tokens_per_rank"), py::arg("with_probs")) .def("combine", - wrap_with_tracing_combine(&HybridEPBuffer::combine, "HybridEPBuffer::combine"), + hybrid_ep_buffer_combine(&HybridEPBuffer::combine, "HybridEPBuffer::combine"), py::kw_only(), py::arg("config"), py::arg("hidden"), py::arg("probs") = c10::nullopt, py::arg("sparse_to_dense_map"), @@ -377,7 +159,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("num_of_tokens_per_rank"), py::arg("with_probs")) .def("dispatch_with_permute", - wrap_with_tracing_dispatch_with_permute(&HybridEPBuffer::dispatch_with_permute, "HybridEPBuffer::dispatch_with_permute"), + hybrid_ep_buffer_dispatch_with_permute(&HybridEPBuffer::dispatch_with_permute, "HybridEPBuffer::dispatch_with_permute"), py::kw_only(), py::arg("config"), py::arg("hidden"), py::arg("probs") = c10::nullopt, @@ -389,7 +171,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("num_of_tokens_per_rank"), py::arg("pad_multiple") = std::nullopt, py::arg("non_blocking") = false, py::arg("with_probs") = false) .def("combine_with_unpermute", - wrap_with_tracing_combine_with_unpermute(&HybridEPBuffer::combine_with_unpermute, "HybridEPBuffer::combine_with_unpermute"), + hybrid_ep_buffer_combine_with_unpermute(&HybridEPBuffer::combine_with_unpermute, "HybridEPBuffer::combine_with_unpermute"), py::kw_only(), py::arg("config"), py::arg("hidden"), py::arg("probs") = c10::nullopt, From 43db4837847e62c8d7902d74f76ff3376ef00e09 Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Thu, 19 Feb 2026 01:21:35 +0200 Subject: [PATCH 06/22] put new APIs into a new file --- csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh | 232 ++++++++++++++++++++++ 1 file changed, 232 insertions(+) create mode 100644 csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh diff --git a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh new file mode 100644 index 000000000..8b43df308 --- /dev/null +++ b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh @@ -0,0 +1,232 @@ +// SPDX-License-Identifier: MIT +// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +#pragma once + +#include +#include "hybrid_ep.cuh" +#include "config.cuh" +#include +#include +#include +#include +#include + +namespace py = pybind11; + +// Convert a single argument to c10::IValue for tracing +template +inline c10::IValue to_ivalue(T&& arg) { + if constexpr (std::is_same_v, HybridEpConfigInstance>) { + return std::forward(arg).to_ivalue_tuple(); + } else if constexpr (std::is_same_v, BufferConfig>) { + return std::forward(arg).to_ivalue_tuple(); + } else if constexpr (std::is_same_v, c10::optional>) { + return arg.has_value() ? c10::IValue(arg.value()) : c10::IValue(); + } else if constexpr (std::is_same_v, c10::optional>) { + return arg.has_value() ? c10::IValue(arg.value()) : c10::IValue(); + } else { + return c10::IValue(std::forward(arg)); + } +} + +template +inline void record_result(c10::impl::GenericList& outputList, Result&& result) { + std::apply([&outputList](const auto&... elem) { + (outputList.emplace_back(c10::IValue(elem)), ...); + }, result); +} + +// Wrapper for combine +template +auto hybrid_ep_buffer_combine(Func func, const std::string& name) { + return [func, name](HybridEPBuffer& self, + HybridEpConfigInstance config, + torch::Tensor hidden, + c10::optional probs, + torch::Tensor sparse_to_dense_map, + torch::Tensor rdma_to_attn_map, + torch::Tensor attn_to_rdma_map, + int64_t num_of_tokens_per_rank, + bool with_probs) { + auto result = (self.*func)(config, hidden, probs, sparse_to_dense_map, + rdma_to_attn_map, attn_to_rdma_map, + num_of_tokens_per_rank, with_probs); + if (at::isRecordFunctionEnabled()) { + c10::impl::GenericList inputList(c10::AnyType::get()); + inputList.emplace_back(to_ivalue(config)); + inputList.emplace_back(to_ivalue(hidden)); + inputList.emplace_back(to_ivalue(probs)); + inputList.emplace_back(to_ivalue(sparse_to_dense_map)); + inputList.emplace_back(to_ivalue(rdma_to_attn_map)); + inputList.emplace_back(to_ivalue(attn_to_rdma_map)); + inputList.emplace_back(to_ivalue(num_of_tokens_per_rank)); + inputList.emplace_back(to_ivalue(with_probs)); + c10::impl::GenericList outputList(c10::AnyType::get()); + record_result(outputList, result); + c10::ArrayRef inputsArray(inputList); + c10::ArrayRef outputsArray(outputList); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); + } + return result; + }; +} + +// Wrapper for dispatch (returns 3-tuple) +template +auto hybrid_ep_buffer_dispatch(Func func, const std::string& name) { + return [func, name](HybridEPBuffer& self, + HybridEpConfigInstance config, + torch::Tensor hidden, + c10::optional probs, + c10::optional scaling_factor, + torch::Tensor sparse_to_dense_map, + torch::Tensor rdma_to_attn_map, + torch::Tensor attn_to_rdma_map, + c10::optional num_dispatched_tokens_tensor, + c10::optional num_dispatched_tokens, + int64_t num_of_tokens_per_rank, + bool with_probs) { + auto result = (self.*func)(config, hidden, probs, scaling_factor, + sparse_to_dense_map, rdma_to_attn_map, attn_to_rdma_map, + num_dispatched_tokens_tensor, num_dispatched_tokens, + num_of_tokens_per_rank, with_probs); + if (at::isRecordFunctionEnabled()) { + c10::impl::GenericList inputList(c10::AnyType::get()); + inputList.emplace_back(to_ivalue(config)); + inputList.emplace_back(to_ivalue(hidden)); + inputList.emplace_back(to_ivalue(probs)); + inputList.emplace_back(to_ivalue(scaling_factor)); + inputList.emplace_back(to_ivalue(sparse_to_dense_map)); + inputList.emplace_back(to_ivalue(rdma_to_attn_map)); + inputList.emplace_back(to_ivalue(attn_to_rdma_map)); + inputList.emplace_back(to_ivalue(num_dispatched_tokens_tensor)); + inputList.emplace_back(to_ivalue(num_dispatched_tokens)); + inputList.emplace_back(to_ivalue(num_of_tokens_per_rank)); + inputList.emplace_back(to_ivalue(with_probs)); + c10::impl::GenericList outputList(c10::AnyType::get()); + record_result(outputList, result); + c10::ArrayRef inputsArray(inputList); + c10::ArrayRef outputsArray(outputList); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); + } + return result; + }; +} + +// Wrapper for dispatch_with_permute (returns 6-tuple) +template +auto hybrid_ep_buffer_dispatch_with_permute(Func func, const std::string& name) { + return [func, name](HybridEPBuffer& self, + HybridEpConfigInstance config, + torch::Tensor hidden, + c10::optional probs, + c10::optional scaling_factor, + torch::Tensor sparse_to_dense_map, + torch::Tensor rdma_to_attn_map, + torch::Tensor attn_to_rdma_map, + c10::optional num_dispatched_tokens_tensor, + c10::optional local_expert_routing_map, + c10::optional row_id_map, + c10::optional num_permuted_tokens, + int64_t num_of_tokens_per_rank, + c10::optional pad_multiple, + bool non_blocking, + bool with_probs) { + auto result = (self.*func)(config, hidden, probs, scaling_factor, + sparse_to_dense_map, rdma_to_attn_map, attn_to_rdma_map, + num_dispatched_tokens_tensor, local_expert_routing_map, row_id_map, + num_permuted_tokens, num_of_tokens_per_rank, pad_multiple, + non_blocking, with_probs); + if (at::isRecordFunctionEnabled()) { + c10::impl::GenericList inputList(c10::AnyType::get()); + inputList.emplace_back(to_ivalue(config)); + inputList.emplace_back(to_ivalue(hidden)); + inputList.emplace_back(to_ivalue(probs)); + inputList.emplace_back(to_ivalue(scaling_factor)); + inputList.emplace_back(to_ivalue(sparse_to_dense_map)); + inputList.emplace_back(to_ivalue(rdma_to_attn_map)); + inputList.emplace_back(to_ivalue(attn_to_rdma_map)); + inputList.emplace_back(to_ivalue(num_dispatched_tokens_tensor)); + inputList.emplace_back(to_ivalue(local_expert_routing_map)); + inputList.emplace_back(to_ivalue(row_id_map)); + inputList.emplace_back(to_ivalue(num_permuted_tokens)); + inputList.emplace_back(to_ivalue(num_of_tokens_per_rank)); + inputList.emplace_back(to_ivalue(pad_multiple)); + inputList.emplace_back(to_ivalue(non_blocking)); + inputList.emplace_back(to_ivalue(with_probs)); + c10::impl::GenericList outputList(c10::AnyType::get()); + record_result(outputList, result); + c10::ArrayRef inputsArray(inputList); + c10::ArrayRef outputsArray(outputList); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); + } + return result; + }; +} + +// Wrapper for combine_with_unpermute (returns 2-tuple) +template +auto hybrid_ep_buffer_combine_with_unpermute(Func func, const std::string& name) { + return [func, name](HybridEPBuffer& self, + HybridEpConfigInstance config, + torch::Tensor hidden, + c10::optional probs, + torch::Tensor sparse_to_dense_map, + torch::Tensor rdma_to_attn_map, + torch::Tensor attn_to_rdma_map, + c10::optional num_dispatched_tokens_tensor, + c10::optional row_id_map, + int64_t num_of_tokens_per_rank, + c10::optional pad_multiple, + bool with_probs) { + auto result = (self.*func)(config, hidden, probs, sparse_to_dense_map, + rdma_to_attn_map, attn_to_rdma_map, + num_dispatched_tokens_tensor, row_id_map, + num_of_tokens_per_rank, pad_multiple, with_probs); + if (at::isRecordFunctionEnabled()) { + c10::impl::GenericList inputList(c10::AnyType::get()); + inputList.emplace_back(to_ivalue(config)); + inputList.emplace_back(to_ivalue(hidden)); + inputList.emplace_back(to_ivalue(probs)); + inputList.emplace_back(to_ivalue(sparse_to_dense_map)); + inputList.emplace_back(to_ivalue(rdma_to_attn_map)); + inputList.emplace_back(to_ivalue(attn_to_rdma_map)); + inputList.emplace_back(to_ivalue(num_dispatched_tokens_tensor)); + inputList.emplace_back(to_ivalue(row_id_map)); + inputList.emplace_back(to_ivalue(num_of_tokens_per_rank)); + inputList.emplace_back(to_ivalue(pad_multiple)); + inputList.emplace_back(to_ivalue(with_probs)); + c10::impl::GenericList outputList(c10::AnyType::get()); + record_result(outputList, result); + c10::ArrayRef inputsArray(inputList); + c10::ArrayRef outputsArray(outputList); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); + } + return result; + }; +} + +// HybridEPBuffer constructor with tracing +inline HybridEPBuffer* hybrid_ep_buffer_init( + py::object process_group, BufferConfig config, int local_rank, int node_rank, + int group_size, std::string base_path, bool load_cached_kernels, + bool use_shared_buffer, bool enable_custom_allgather) { + if (at::isRecordFunctionEnabled()) { + c10::impl::GenericList inputList(c10::AnyType::get()); + inputList.emplace_back(to_ivalue(config)); + inputList.emplace_back(to_ivalue(local_rank)); + inputList.emplace_back(to_ivalue(node_rank)); + inputList.emplace_back(to_ivalue(group_size)); + inputList.emplace_back(to_ivalue(base_path)); + inputList.emplace_back(to_ivalue(load_cached_kernels)); + inputList.emplace_back(to_ivalue(use_shared_buffer)); + inputList.emplace_back(to_ivalue(enable_custom_allgather)); + c10::impl::GenericList outputList(c10::AnyType::get()); + c10::ArrayRef inputsArray(inputList); + c10::ArrayRef outputsArray(outputList); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::__init__", inputsArray, outputsArray); + } + return new HybridEPBuffer(process_group, config, local_rank, node_rank, group_size, + base_path, load_cached_kernels, use_shared_buffer, enable_custom_allgather); +} From 8fc15108af11250be721821d9ecbd670a8764361 Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Thu, 19 Feb 2026 01:26:21 +0200 Subject: [PATCH 07/22] clean up head files --- csrc/hybrid_ep/pybind_hybrid_ep.cu | 5 ----- 1 file changed, 5 deletions(-) diff --git a/csrc/hybrid_ep/pybind_hybrid_ep.cu b/csrc/hybrid_ep/pybind_hybrid_ep.cu index 3bd4f7752..dac53f471 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep.cu +++ b/csrc/hybrid_ep/pybind_hybrid_ep.cu @@ -10,11 +10,6 @@ #include "utils.cuh" #include "config.cuh" #include "pybind_hybrid_ep_utils.cuh" -#include -#include -#include -#include -#include namespace py = pybind11; From a09d4dba959ae6cf38e4a046954dd55e241eb312 Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Thu, 19 Feb 2026 03:10:49 +0200 Subject: [PATCH 08/22] get pg id --- csrc/hybrid_ep/hybrid_ep.cuh | 2 + csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh | 49 +++++++++++------------ 2 files changed, 26 insertions(+), 25 deletions(-) diff --git a/csrc/hybrid_ep/hybrid_ep.cuh b/csrc/hybrid_ep/hybrid_ep.cuh index 3463e51c5..8ed54b3e0 100644 --- a/csrc/hybrid_ep/hybrid_ep.cuh +++ b/csrc/hybrid_ep/hybrid_ep.cuh @@ -19,6 +19,8 @@ #include "buffer/internode.cuh" #endif +std::string get_comm_id(pybind11::object process_group); + class HybridEPBuffer { public: HybridEPBuffer(pybind11::object process_group, BufferConfig config, int local_rank, int node_rank, int group_size, std::string base_path, bool load_cached_kernels, bool use_shared_buffer, bool enable_custom_allgather); diff --git a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh index 8b43df308..3ee19a790 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh +++ b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh @@ -37,7 +37,30 @@ inline void record_result(c10::impl::GenericList& outputList, Result&& result) { }, result); } -// Wrapper for combine +inline HybridEPBuffer* hybrid_ep_buffer_init( + py::object process_group, BufferConfig config, int local_rank, int node_rank, + int group_size, std::string base_path, bool load_cached_kernels, + bool use_shared_buffer, bool enable_custom_allgather) { + if (at::isRecordFunctionEnabled()) { + c10::impl::GenericList inputList(c10::AnyType::get()); + inputList.emplace_back(c10::IValue(get_comm_id(process_group))); + inputList.emplace_back(to_ivalue(config)); + inputList.emplace_back(to_ivalue(local_rank)); + inputList.emplace_back(to_ivalue(node_rank)); + inputList.emplace_back(to_ivalue(group_size)); + inputList.emplace_back(to_ivalue(base_path)); + inputList.emplace_back(to_ivalue(load_cached_kernels)); + inputList.emplace_back(to_ivalue(use_shared_buffer)); + inputList.emplace_back(to_ivalue(enable_custom_allgather)); + c10::impl::GenericList outputList(c10::AnyType::get()); + c10::ArrayRef inputsArray(inputList); + c10::ArrayRef outputsArray(outputList); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::__init__", inputsArray, outputsArray); + } + return new HybridEPBuffer(process_group, config, local_rank, node_rank, group_size, + base_path, load_cached_kernels, use_shared_buffer, enable_custom_allgather); +} + template auto hybrid_ep_buffer_combine(Func func, const std::string& name) { return [func, name](HybridEPBuffer& self, @@ -206,27 +229,3 @@ auto hybrid_ep_buffer_combine_with_unpermute(Func func, const std::string& name) return result; }; } - -// HybridEPBuffer constructor with tracing -inline HybridEPBuffer* hybrid_ep_buffer_init( - py::object process_group, BufferConfig config, int local_rank, int node_rank, - int group_size, std::string base_path, bool load_cached_kernels, - bool use_shared_buffer, bool enable_custom_allgather) { - if (at::isRecordFunctionEnabled()) { - c10::impl::GenericList inputList(c10::AnyType::get()); - inputList.emplace_back(to_ivalue(config)); - inputList.emplace_back(to_ivalue(local_rank)); - inputList.emplace_back(to_ivalue(node_rank)); - inputList.emplace_back(to_ivalue(group_size)); - inputList.emplace_back(to_ivalue(base_path)); - inputList.emplace_back(to_ivalue(load_cached_kernels)); - inputList.emplace_back(to_ivalue(use_shared_buffer)); - inputList.emplace_back(to_ivalue(enable_custom_allgather)); - c10::impl::GenericList outputList(c10::AnyType::get()); - c10::ArrayRef inputsArray(inputList); - c10::ArrayRef outputsArray(outputList); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::__init__", inputsArray, outputsArray); - } - return new HybridEPBuffer(process_group, config, local_rank, node_rank, group_size, - base_path, load_cached_kernels, use_shared_buffer, enable_custom_allgather); -} From ba9a0ab62ba1e87017d02fddc38ae38c04e1910b Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Thu, 19 Feb 2026 03:48:03 +0200 Subject: [PATCH 09/22] get pg id --- csrc/hybrid_ep/hybrid_ep.cuh | 2 -- csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh | 2 +- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/csrc/hybrid_ep/hybrid_ep.cuh b/csrc/hybrid_ep/hybrid_ep.cuh index 8ed54b3e0..3463e51c5 100644 --- a/csrc/hybrid_ep/hybrid_ep.cuh +++ b/csrc/hybrid_ep/hybrid_ep.cuh @@ -19,8 +19,6 @@ #include "buffer/internode.cuh" #endif -std::string get_comm_id(pybind11::object process_group); - class HybridEPBuffer { public: HybridEPBuffer(pybind11::object process_group, BufferConfig config, int local_rank, int node_rank, int group_size, std::string base_path, bool load_cached_kernels, bool use_shared_buffer, bool enable_custom_allgather); diff --git a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh index 3ee19a790..9983874c1 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh +++ b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh @@ -43,7 +43,7 @@ inline HybridEPBuffer* hybrid_ep_buffer_init( bool use_shared_buffer, bool enable_custom_allgather) { if (at::isRecordFunctionEnabled()) { c10::impl::GenericList inputList(c10::AnyType::get()); - inputList.emplace_back(c10::IValue(get_comm_id(process_group))); + inputList.emplace_back(c10::IValue(process_group.attr("group_name").cast())); inputList.emplace_back(to_ivalue(config)); inputList.emplace_back(to_ivalue(local_rank)); inputList.emplace_back(to_ivalue(node_rank)); From cc98dd8e7d47ba55bb9011c8718b3114f054ebe1 Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Sun, 22 Feb 2026 02:09:21 +0200 Subject: [PATCH 10/22] add schema --- csrc/hybrid_ep/pybind_hybrid_ep.cu | 8 ++-- csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh | 53 +++++++++++++++++------ tests/test_hybrid_ep_collect_et.py | 2 +- 3 files changed, 45 insertions(+), 18 deletions(-) diff --git a/csrc/hybrid_ep/pybind_hybrid_ep.cu b/csrc/hybrid_ep/pybind_hybrid_ep.cu index dac53f471..92570ee9e 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep.cu +++ b/csrc/hybrid_ep/pybind_hybrid_ep.cu @@ -136,7 +136,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("metadata_preprocessing", &HybridEPBuffer::metadata_preprocessing, py::kw_only(), py::arg("config"), py::arg("routing_map"), py::arg("num_of_tokens_per_rank"), py::arg("non_blocking") = false) .def("dispatch", - hybrid_ep_buffer_dispatch(&HybridEPBuffer::dispatch, "HybridEPBuffer::dispatch"), + hybrid_ep_buffer_dispatch(&HybridEPBuffer::dispatch), py::kw_only(), py::arg("config"), py::arg("hidden"), py::arg("probs") = c10::nullopt, @@ -146,7 +146,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("num_dispatched_tokens") = std::nullopt, py::arg("num_of_tokens_per_rank"), py::arg("with_probs")) .def("combine", - hybrid_ep_buffer_combine(&HybridEPBuffer::combine, "HybridEPBuffer::combine"), + hybrid_ep_buffer_combine(&HybridEPBuffer::combine), py::kw_only(), py::arg("config"), py::arg("hidden"), py::arg("probs") = c10::nullopt, py::arg("sparse_to_dense_map"), @@ -154,7 +154,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("num_of_tokens_per_rank"), py::arg("with_probs")) .def("dispatch_with_permute", - hybrid_ep_buffer_dispatch_with_permute(&HybridEPBuffer::dispatch_with_permute, "HybridEPBuffer::dispatch_with_permute"), + hybrid_ep_buffer_dispatch_with_permute(&HybridEPBuffer::dispatch_with_permute), py::kw_only(), py::arg("config"), py::arg("hidden"), py::arg("probs") = c10::nullopt, @@ -166,7 +166,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("num_of_tokens_per_rank"), py::arg("pad_multiple") = std::nullopt, py::arg("non_blocking") = false, py::arg("with_probs") = false) .def("combine_with_unpermute", - hybrid_ep_buffer_combine_with_unpermute(&HybridEPBuffer::combine_with_unpermute, "HybridEPBuffer::combine_with_unpermute"), + hybrid_ep_buffer_combine_with_unpermute(&HybridEPBuffer::combine_with_unpermute), py::kw_only(), py::arg("config"), py::arg("hidden"), py::arg("probs") = c10::nullopt, diff --git a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh index 9983874c1..8e26e17d1 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh +++ b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh @@ -10,6 +10,7 @@ #include #include #include +#include #include namespace py = pybind11; @@ -55,15 +56,20 @@ inline HybridEPBuffer* hybrid_ep_buffer_init( c10::impl::GenericList outputList(c10::AnyType::get()); c10::ArrayRef inputsArray(inputList); c10::ArrayRef outputsArray(outputList); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::__init__", inputsArray, outputsArray); + std::vector args; + std::vector returns; + c10::FunctionSchema schema( + "HybridEPBuffer::__init__", "", + std::move(args), std::move(returns), false, false); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, inputsArray, outputsArray); } return new HybridEPBuffer(process_group, config, local_rank, node_rank, group_size, base_path, load_cached_kernels, use_shared_buffer, enable_custom_allgather); } template -auto hybrid_ep_buffer_combine(Func func, const std::string& name) { - return [func, name](HybridEPBuffer& self, +auto hybrid_ep_buffer_combine(Func func) { + return [func](HybridEPBuffer& self, HybridEpConfigInstance config, torch::Tensor hidden, c10::optional probs, @@ -89,7 +95,12 @@ auto hybrid_ep_buffer_combine(Func func, const std::string& name) { record_result(outputList, result); c10::ArrayRef inputsArray(inputList); c10::ArrayRef outputsArray(outputList); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); + std::vector args; + std::vector returns; + c10::FunctionSchema schema( + "HybridEPBuffer::combine", "", + std::move(args), std::move(returns), false, false); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, inputsArray, outputsArray); } return result; }; @@ -97,8 +108,8 @@ auto hybrid_ep_buffer_combine(Func func, const std::string& name) { // Wrapper for dispatch (returns 3-tuple) template -auto hybrid_ep_buffer_dispatch(Func func, const std::string& name) { - return [func, name](HybridEPBuffer& self, +auto hybrid_ep_buffer_dispatch(Func func) { + return [func](HybridEPBuffer& self, HybridEpConfigInstance config, torch::Tensor hidden, c10::optional probs, @@ -131,7 +142,13 @@ auto hybrid_ep_buffer_dispatch(Func func, const std::string& name) { record_result(outputList, result); c10::ArrayRef inputsArray(inputList); c10::ArrayRef outputsArray(outputList); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); + + std::vector args; + std::vector returns; + c10::FunctionSchema schema( + "HybridEPBuffer::dispatch", "", + std::move(args), std::move(returns), false, false); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, inputsArray, outputsArray); } return result; }; @@ -139,8 +156,8 @@ auto hybrid_ep_buffer_dispatch(Func func, const std::string& name) { // Wrapper for dispatch_with_permute (returns 6-tuple) template -auto hybrid_ep_buffer_dispatch_with_permute(Func func, const std::string& name) { - return [func, name](HybridEPBuffer& self, +auto hybrid_ep_buffer_dispatch_with_permute(Func func) { + return [func](HybridEPBuffer& self, HybridEpConfigInstance config, torch::Tensor hidden, c10::optional probs, @@ -182,7 +199,12 @@ auto hybrid_ep_buffer_dispatch_with_permute(Func func, const std::string& name) record_result(outputList, result); c10::ArrayRef inputsArray(inputList); c10::ArrayRef outputsArray(outputList); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); + std::vector args; + std::vector returns; + c10::FunctionSchema schema( + "HybridEPBuffer::dispatch_with_permute", "", + std::move(args), std::move(returns), false, false); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, inputsArray, outputsArray); } return result; }; @@ -190,8 +212,8 @@ auto hybrid_ep_buffer_dispatch_with_permute(Func func, const std::string& name) // Wrapper for combine_with_unpermute (returns 2-tuple) template -auto hybrid_ep_buffer_combine_with_unpermute(Func func, const std::string& name) { - return [func, name](HybridEPBuffer& self, +auto hybrid_ep_buffer_combine_with_unpermute(Func func) { + return [func](HybridEPBuffer& self, HybridEpConfigInstance config, torch::Tensor hidden, c10::optional probs, @@ -224,7 +246,12 @@ auto hybrid_ep_buffer_combine_with_unpermute(Func func, const std::string& name) record_result(outputList, result); c10::ArrayRef inputsArray(inputList); c10::ArrayRef outputsArray(outputList); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(name, inputsArray, outputsArray); + std::vector args; + std::vector returns; + c10::FunctionSchema schema( + "HybridEPBuffer::combine_with_unpermute", "", + std::move(args), std::move(returns), false, false); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, inputsArray, outputsArray); } return result; }; diff --git a/tests/test_hybrid_ep_collect_et.py b/tests/test_hybrid_ep_collect_et.py index e5eb8a2a6..7a7da3ea0 100644 --- a/tests/test_hybrid_ep_collect_et.py +++ b/tests/test_hybrid_ep_collect_et.py @@ -87,7 +87,7 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): stream = torch.cuda.Stream() with torch.cuda.stream(stream): - et = ExecutionTraceObserver().register_callback(f"et_{dist.get_rank()}.json") + et = ExecutionTraceObserver().register_callback(f"rank-{dist.get_rank()}.json") et.start() buffer = deep_ep.HybridEPBuffer( From d2456cb64ba6430a8baaafec66e6691232a17dea Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Mon, 23 Feb 2026 04:57:26 +0200 Subject: [PATCH 11/22] remove schema --- csrc/hybrid_ep/config.cuh | 2 +- csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh | 91 ++++++++++++++++++++--- 2 files changed, 83 insertions(+), 10 deletions(-) diff --git a/csrc/hybrid_ep/config.cuh b/csrc/hybrid_ep/config.cuh index 7955d0384..7d262d920 100644 --- a/csrc/hybrid_ep/config.cuh +++ b/csrc/hybrid_ep/config.cuh @@ -134,7 +134,7 @@ struct HybridEpConfigInstance { /** Convert all attributes to a single IValue holding a tuple of IValues (ints/bools). */ c10::IValue to_ivalue_tuple() const { std::vector elements; - elements.reserve(24); + elements.reserve(23); // Hybrid-ep Config elements.push_back(c10::IValue(static_cast(hidden_dim))); elements.push_back(c10::IValue(static_cast(max_num_of_tokens_per_rank))); diff --git a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh index 8e26e17d1..fb7c9ac04 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh +++ b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh @@ -6,11 +6,13 @@ #include #include "hybrid_ep.cuh" #include "config.cuh" +#include #include #include #include #include #include +#include #include namespace py = pybind11; @@ -57,11 +59,21 @@ inline HybridEPBuffer* hybrid_ep_buffer_init( c10::ArrayRef inputsArray(inputList); c10::ArrayRef outputsArray(outputList); std::vector args; + args.emplace_back("group_name", c10::StringType::get()); + args.emplace_back("config", c10::AnyType::get()); + args.emplace_back("local_rank", c10::IntType::get()); + args.emplace_back("node_rank", c10::IntType::get()); + args.emplace_back("group_size", c10::IntType::get()); + args.emplace_back("base_path", c10::StringType::get()); + args.emplace_back("load_cached_kernels", c10::BoolType::get()); + args.emplace_back("use_shared_buffer", c10::BoolType::get()); + args.emplace_back("enable_custom_allgather", c10::BoolType::get()); std::vector returns; + // outputList is empty for __init__ (buffer pointer not recorded as IValue) c10::FunctionSchema schema( "HybridEPBuffer::__init__", "", std::move(args), std::move(returns), false, false); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, inputsArray, outputsArray); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::__init__", inputsArray, outputsArray); } return new HybridEPBuffer(process_group, config, local_rank, node_rank, group_size, base_path, load_cached_kernels, use_shared_buffer, enable_custom_allgather); @@ -95,12 +107,26 @@ auto hybrid_ep_buffer_combine(Func func) { record_result(outputList, result); c10::ArrayRef inputsArray(inputList); c10::ArrayRef outputsArray(outputList); + /* + std::cout << "outputList size: " << outputList.size() << std::endl; std::vector args; + args.emplace_back("config", c10::AnyType::get()); + args.emplace_back("hidden", c10::TensorType::get()); + args.emplace_back("probs", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("sparse_to_dense_map", c10::TensorType::get()); + args.emplace_back("rdma_to_attn_map", c10::TensorType::get()); + args.emplace_back("attn_to_rdma_map", c10::TensorType::get()); + args.emplace_back("num_of_tokens_per_rank", c10::IntType::get()); + args.emplace_back("with_probs", c10::BoolType::get()); std::vector returns; + returns.emplace_back("", c10::TensorType::get()); + returns.emplace_back("", c10::TensorType::get()); c10::FunctionSchema schema( "HybridEPBuffer::combine", "", std::move(args), std::move(returns), false, false); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, inputsArray, outputsArray); + */ + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::combine", inputsArray, outputsArray); + //RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, inputsArray, outputsArray); } return result; }; @@ -142,13 +168,26 @@ auto hybrid_ep_buffer_dispatch(Func func) { record_result(outputList, result); c10::ArrayRef inputsArray(inputList); c10::ArrayRef outputsArray(outputList); - - std::vector args; - std::vector returns; - c10::FunctionSchema schema( - "HybridEPBuffer::dispatch", "", - std::move(args), std::move(returns), false, false); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, inputsArray, outputsArray); + // std::vector args; + // args.emplace_back("config", c10::AnyType::get()); + // args.emplace_back("hidden", c10::TensorType::get()); + // args.emplace_back("probs", c10::OptionalType::create(c10::TensorType::get())); + // args.emplace_back("scaling_factor", c10::OptionalType::create(c10::TensorType::get())); + // args.emplace_back("sparse_to_dense_map", c10::TensorType::get()); + // args.emplace_back("rdma_to_attn_map", c10::TensorType::get()); + // args.emplace_back("attn_to_rdma_map", c10::TensorType::get()); + // args.emplace_back("num_dispatched_tokens_tensor", c10::OptionalType::create(c10::TensorType::get())); + // args.emplace_back("num_dispatched_tokens", c10::OptionalType::create(c10::IntType::get())); + // args.emplace_back("num_of_tokens_per_rank", c10::IntType::get()); + // args.emplace_back("with_probs", c10::BoolType::get()); + // std::vector returns; + // returns.emplace_back("", c10::TensorType::get()); + // returns.emplace_back("", c10::TensorType::get()); + // returns.emplace_back("", c10::TensorType::get()); + // c10::FunctionSchema schema( + // "HybridEPBuffer::dispatch", "", + // std::move(args), std::move(returns), false, false); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::dispatch", inputsArray, outputsArray); } return result; }; @@ -200,7 +239,28 @@ auto hybrid_ep_buffer_dispatch_with_permute(Func func) { c10::ArrayRef inputsArray(inputList); c10::ArrayRef outputsArray(outputList); std::vector args; + args.emplace_back("config", c10::AnyType::get()); + args.emplace_back("hidden", c10::TensorType::get()); + args.emplace_back("probs", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("scaling_factor", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("sparse_to_dense_map", c10::TensorType::get()); + args.emplace_back("rdma_to_attn_map", c10::TensorType::get()); + args.emplace_back("attn_to_rdma_map", c10::TensorType::get()); + args.emplace_back("num_dispatched_tokens_tensor", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("local_expert_routing_map", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("row_id_map", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("num_permuted_tokens", c10::OptionalType::create(c10::IntType::get())); + args.emplace_back("num_of_tokens_per_rank", c10::IntType::get()); + args.emplace_back("pad_multiple", c10::OptionalType::create(c10::IntType::get())); + args.emplace_back("non_blocking", c10::BoolType::get()); + args.emplace_back("with_probs", c10::BoolType::get()); std::vector returns; + returns.emplace_back("", c10::TensorType::get()); + returns.emplace_back("", c10::TensorType::get()); + returns.emplace_back("", c10::TensorType::get()); + returns.emplace_back("", c10::TensorType::get()); + returns.emplace_back("", c10::TensorType::get()); + returns.emplace_back("", c10::TensorType::get()); c10::FunctionSchema schema( "HybridEPBuffer::dispatch_with_permute", "", std::move(args), std::move(returns), false, false); @@ -247,7 +307,20 @@ auto hybrid_ep_buffer_combine_with_unpermute(Func func) { c10::ArrayRef inputsArray(inputList); c10::ArrayRef outputsArray(outputList); std::vector args; + args.emplace_back("config", c10::AnyType::get()); + args.emplace_back("hidden", c10::TensorType::get()); + args.emplace_back("probs", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("sparse_to_dense_map", c10::TensorType::get()); + args.emplace_back("rdma_to_attn_map", c10::TensorType::get()); + args.emplace_back("attn_to_rdma_map", c10::TensorType::get()); + args.emplace_back("num_dispatched_tokens_tensor", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("row_id_map", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("num_of_tokens_per_rank", c10::IntType::get()); + args.emplace_back("pad_multiple", c10::OptionalType::create(c10::IntType::get())); + args.emplace_back("with_probs", c10::BoolType::get()); std::vector returns; + returns.emplace_back("", c10::TensorType::get()); + returns.emplace_back("", c10::TensorType::get()); c10::FunctionSchema schema( "HybridEPBuffer::combine_with_unpermute", "", std::move(args), std::move(returns), false, false); From 921f2a1845ffc8b6acda087f417b883842b76666 Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Tue, 24 Feb 2026 03:48:17 +0200 Subject: [PATCH 12/22] change innput/output from a single arrary to a list of inputs --- csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh | 290 +++++++++++----------- 1 file changed, 146 insertions(+), 144 deletions(-) diff --git a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh index fb7c9ac04..54f3d4c73 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh +++ b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh @@ -34,7 +34,7 @@ inline c10::IValue to_ivalue(T&& arg) { } template -inline void record_result(c10::impl::GenericList& outputList, Result&& result) { +inline void record_result(std::vector& outputList, Result&& result) { std::apply([&outputList](const auto&... elem) { (outputList.emplace_back(c10::IValue(elem)), ...); }, result); @@ -45,35 +45,35 @@ inline HybridEPBuffer* hybrid_ep_buffer_init( int group_size, std::string base_path, bool load_cached_kernels, bool use_shared_buffer, bool enable_custom_allgather) { if (at::isRecordFunctionEnabled()) { - c10::impl::GenericList inputList(c10::AnyType::get()); - inputList.emplace_back(c10::IValue(process_group.attr("group_name").cast())); - inputList.emplace_back(to_ivalue(config)); - inputList.emplace_back(to_ivalue(local_rank)); - inputList.emplace_back(to_ivalue(node_rank)); - inputList.emplace_back(to_ivalue(group_size)); - inputList.emplace_back(to_ivalue(base_path)); - inputList.emplace_back(to_ivalue(load_cached_kernels)); - inputList.emplace_back(to_ivalue(use_shared_buffer)); - inputList.emplace_back(to_ivalue(enable_custom_allgather)); - c10::impl::GenericList outputList(c10::AnyType::get()); - c10::ArrayRef inputsArray(inputList); - c10::ArrayRef outputsArray(outputList); - std::vector args; - args.emplace_back("group_name", c10::StringType::get()); - args.emplace_back("config", c10::AnyType::get()); - args.emplace_back("local_rank", c10::IntType::get()); - args.emplace_back("node_rank", c10::IntType::get()); - args.emplace_back("group_size", c10::IntType::get()); - args.emplace_back("base_path", c10::StringType::get()); - args.emplace_back("load_cached_kernels", c10::BoolType::get()); - args.emplace_back("use_shared_buffer", c10::BoolType::get()); - args.emplace_back("enable_custom_allgather", c10::BoolType::get()); - std::vector returns; - // outputList is empty for __init__ (buffer pointer not recorded as IValue) - c10::FunctionSchema schema( - "HybridEPBuffer::__init__", "", - std::move(args), std::move(returns), false, false); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::__init__", inputsArray, outputsArray); + std::vector inputValues; + inputValues.reserve(9); + inputValues.emplace_back(c10::IValue(process_group.attr("group_name").cast())); + inputValues.emplace_back(to_ivalue(config)); + inputValues.emplace_back(to_ivalue(local_rank)); + inputValues.emplace_back(to_ivalue(node_rank)); + inputValues.emplace_back(to_ivalue(group_size)); + inputValues.emplace_back(to_ivalue(base_path)); + inputValues.emplace_back(to_ivalue(load_cached_kernels)); + inputValues.emplace_back(to_ivalue(use_shared_buffer)); + inputValues.emplace_back(to_ivalue(enable_custom_allgather)); + + std::vector outputValues; + // std::vector args; + // args.emplace_back("group_name", c10::StringType::get()); + // args.emplace_back("config", c10::AnyType::get()); + // args.emplace_back("local_rank", c10::IntType::get()); + // args.emplace_back("node_rank", c10::IntType::get()); + // args.emplace_back("group_size", c10::IntType::get()); + // args.emplace_back("base_path", c10::StringType::get()); + // args.emplace_back("load_cached_kernels", c10::BoolType::get()); + // args.emplace_back("use_shared_buffer", c10::BoolType::get()); + // args.emplace_back("enable_custom_allgather", c10::BoolType::get()); + // std::vector returns; + // // outputList is empty for __init__ (buffer pointer not recorded as IValue) + // c10::FunctionSchema schema( + // "HybridEPBuffer::__init__", "", + // std::move(args), std::move(returns), false, false); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::__init__", &inputValues, outputValues); } return new HybridEPBuffer(process_group, config, local_rank, node_rank, group_size, base_path, load_cached_kernels, use_shared_buffer, enable_custom_allgather); @@ -94,19 +94,20 @@ auto hybrid_ep_buffer_combine(Func func) { rdma_to_attn_map, attn_to_rdma_map, num_of_tokens_per_rank, with_probs); if (at::isRecordFunctionEnabled()) { - c10::impl::GenericList inputList(c10::AnyType::get()); - inputList.emplace_back(to_ivalue(config)); - inputList.emplace_back(to_ivalue(hidden)); - inputList.emplace_back(to_ivalue(probs)); - inputList.emplace_back(to_ivalue(sparse_to_dense_map)); - inputList.emplace_back(to_ivalue(rdma_to_attn_map)); - inputList.emplace_back(to_ivalue(attn_to_rdma_map)); - inputList.emplace_back(to_ivalue(num_of_tokens_per_rank)); - inputList.emplace_back(to_ivalue(with_probs)); - c10::impl::GenericList outputList(c10::AnyType::get()); - record_result(outputList, result); - c10::ArrayRef inputsArray(inputList); - c10::ArrayRef outputsArray(outputList); + std::vector inputValues; + inputValues.reserve(8); + inputValues.push_back(to_ivalue(config)); + inputValues.push_back(to_ivalue(hidden)); + inputValues.push_back(to_ivalue(probs)); + inputValues.push_back(to_ivalue(sparse_to_dense_map)); + inputValues.push_back(to_ivalue(rdma_to_attn_map)); + inputValues.push_back(to_ivalue(attn_to_rdma_map)); + inputValues.push_back(to_ivalue(num_of_tokens_per_rank)); + inputValues.push_back(to_ivalue(with_probs)); + + std::vector outputValues; + record_result(outputValues, result); + // c10::ArrayRef outputsArray(outputList); /* std::cout << "outputList size: " << outputList.size() << std::endl; std::vector args; @@ -125,8 +126,7 @@ auto hybrid_ep_buffer_combine(Func func) { "HybridEPBuffer::combine", "", std::move(args), std::move(returns), false, false); */ - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::combine", inputsArray, outputsArray); - //RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, inputsArray, outputsArray); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::combine", &inputValues, outputValues); } return result; }; @@ -152,22 +152,22 @@ auto hybrid_ep_buffer_dispatch(Func func) { num_dispatched_tokens_tensor, num_dispatched_tokens, num_of_tokens_per_rank, with_probs); if (at::isRecordFunctionEnabled()) { - c10::impl::GenericList inputList(c10::AnyType::get()); - inputList.emplace_back(to_ivalue(config)); - inputList.emplace_back(to_ivalue(hidden)); - inputList.emplace_back(to_ivalue(probs)); - inputList.emplace_back(to_ivalue(scaling_factor)); - inputList.emplace_back(to_ivalue(sparse_to_dense_map)); - inputList.emplace_back(to_ivalue(rdma_to_attn_map)); - inputList.emplace_back(to_ivalue(attn_to_rdma_map)); - inputList.emplace_back(to_ivalue(num_dispatched_tokens_tensor)); - inputList.emplace_back(to_ivalue(num_dispatched_tokens)); - inputList.emplace_back(to_ivalue(num_of_tokens_per_rank)); - inputList.emplace_back(to_ivalue(with_probs)); - c10::impl::GenericList outputList(c10::AnyType::get()); - record_result(outputList, result); - c10::ArrayRef inputsArray(inputList); - c10::ArrayRef outputsArray(outputList); + std::vector inputValues; + inputValues.reserve(11); + inputValues.push_back(to_ivalue(config)); + inputValues.push_back(to_ivalue(hidden)); + inputValues.push_back(to_ivalue(probs)); + inputValues.push_back(to_ivalue(scaling_factor)); + inputValues.push_back(to_ivalue(sparse_to_dense_map)); + inputValues.push_back(to_ivalue(rdma_to_attn_map)); + inputValues.push_back(to_ivalue(attn_to_rdma_map)); + inputValues.push_back(to_ivalue(num_dispatched_tokens_tensor)); + inputValues.push_back(to_ivalue(num_dispatched_tokens)); + inputValues.push_back(to_ivalue(num_of_tokens_per_rank)); + inputValues.push_back(to_ivalue(with_probs)); + + std::vector outputValues; + record_result(outputValues, result); // std::vector args; // args.emplace_back("config", c10::AnyType::get()); // args.emplace_back("hidden", c10::TensorType::get()); @@ -187,7 +187,7 @@ auto hybrid_ep_buffer_dispatch(Func func) { // c10::FunctionSchema schema( // "HybridEPBuffer::dispatch", "", // std::move(args), std::move(returns), false, false); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::dispatch", inputsArray, outputsArray); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::dispatch", &inputValues, outputValues); } return result; }; @@ -218,53 +218,54 @@ auto hybrid_ep_buffer_dispatch_with_permute(Func func) { num_permuted_tokens, num_of_tokens_per_rank, pad_multiple, non_blocking, with_probs); if (at::isRecordFunctionEnabled()) { - c10::impl::GenericList inputList(c10::AnyType::get()); - inputList.emplace_back(to_ivalue(config)); - inputList.emplace_back(to_ivalue(hidden)); - inputList.emplace_back(to_ivalue(probs)); - inputList.emplace_back(to_ivalue(scaling_factor)); - inputList.emplace_back(to_ivalue(sparse_to_dense_map)); - inputList.emplace_back(to_ivalue(rdma_to_attn_map)); - inputList.emplace_back(to_ivalue(attn_to_rdma_map)); - inputList.emplace_back(to_ivalue(num_dispatched_tokens_tensor)); - inputList.emplace_back(to_ivalue(local_expert_routing_map)); - inputList.emplace_back(to_ivalue(row_id_map)); - inputList.emplace_back(to_ivalue(num_permuted_tokens)); - inputList.emplace_back(to_ivalue(num_of_tokens_per_rank)); - inputList.emplace_back(to_ivalue(pad_multiple)); - inputList.emplace_back(to_ivalue(non_blocking)); - inputList.emplace_back(to_ivalue(with_probs)); - c10::impl::GenericList outputList(c10::AnyType::get()); - record_result(outputList, result); - c10::ArrayRef inputsArray(inputList); - c10::ArrayRef outputsArray(outputList); - std::vector args; - args.emplace_back("config", c10::AnyType::get()); - args.emplace_back("hidden", c10::TensorType::get()); - args.emplace_back("probs", c10::OptionalType::create(c10::TensorType::get())); - args.emplace_back("scaling_factor", c10::OptionalType::create(c10::TensorType::get())); - args.emplace_back("sparse_to_dense_map", c10::TensorType::get()); - args.emplace_back("rdma_to_attn_map", c10::TensorType::get()); - args.emplace_back("attn_to_rdma_map", c10::TensorType::get()); - args.emplace_back("num_dispatched_tokens_tensor", c10::OptionalType::create(c10::TensorType::get())); - args.emplace_back("local_expert_routing_map", c10::OptionalType::create(c10::TensorType::get())); - args.emplace_back("row_id_map", c10::OptionalType::create(c10::TensorType::get())); - args.emplace_back("num_permuted_tokens", c10::OptionalType::create(c10::IntType::get())); - args.emplace_back("num_of_tokens_per_rank", c10::IntType::get()); - args.emplace_back("pad_multiple", c10::OptionalType::create(c10::IntType::get())); - args.emplace_back("non_blocking", c10::BoolType::get()); - args.emplace_back("with_probs", c10::BoolType::get()); - std::vector returns; - returns.emplace_back("", c10::TensorType::get()); - returns.emplace_back("", c10::TensorType::get()); - returns.emplace_back("", c10::TensorType::get()); - returns.emplace_back("", c10::TensorType::get()); - returns.emplace_back("", c10::TensorType::get()); - returns.emplace_back("", c10::TensorType::get()); - c10::FunctionSchema schema( - "HybridEPBuffer::dispatch_with_permute", "", - std::move(args), std::move(returns), false, false); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, inputsArray, outputsArray); + std::vector inputValues; + inputValues.reserve(15); + inputValues.push_back(to_ivalue(config)); + inputValues.push_back(to_ivalue(hidden)); + inputValues.push_back(to_ivalue(probs)); + inputValues.push_back(to_ivalue(scaling_factor)); + inputValues.push_back(to_ivalue(sparse_to_dense_map)); + inputValues.push_back(to_ivalue(rdma_to_attn_map)); + inputValues.push_back(to_ivalue(attn_to_rdma_map)); + inputValues.push_back(to_ivalue(num_dispatched_tokens_tensor)); + inputValues.push_back(to_ivalue(local_expert_routing_map)); + inputValues.push_back(to_ivalue(row_id_map)); + inputValues.push_back(to_ivalue(num_permuted_tokens)); + inputValues.push_back(to_ivalue(num_of_tokens_per_rank)); + inputValues.push_back(to_ivalue(pad_multiple)); + inputValues.push_back(to_ivalue(non_blocking)); + inputValues.push_back(to_ivalue(with_probs)); + + std::vector outputValues; + record_result(outputValues, result); + + // std::vector args; + // args.emplace_back("config", c10::AnyType::get()); + // args.emplace_back("hidden", c10::TensorType::get()); + // args.emplace_back("probs", c10::OptionalType::create(c10::TensorType::get())); + // args.emplace_back("scaling_factor", c10::OptionalType::create(c10::TensorType::get())); + // args.emplace_back("sparse_to_dense_map", c10::TensorType::get()); + // args.emplace_back("rdma_to_attn_map", c10::TensorType::get()); + // args.emplace_back("attn_to_rdma_map", c10::TensorType::get()); + // args.emplace_back("num_dispatched_tokens_tensor", c10::OptionalType::create(c10::TensorType::get())); + // args.emplace_back("local_expert_routing_map", c10::OptionalType::create(c10::TensorType::get())); + // args.emplace_back("row_id_map", c10::OptionalType::create(c10::TensorType::get())); + // args.emplace_back("num_permuted_tokens", c10::OptionalType::create(c10::IntType::get())); + // args.emplace_back("num_of_tokens_per_rank", c10::IntType::get()); + // args.emplace_back("pad_multiple", c10::OptionalType::create(c10::IntType::get())); + // args.emplace_back("non_blocking", c10::BoolType::get()); + // args.emplace_back("with_probs", c10::BoolType::get()); + // std::vector returns; + // returns.emplace_back("", c10::TensorType::get()); + // returns.emplace_back("", c10::TensorType::get()); + // returns.emplace_back("", c10::TensorType::get()); + // returns.emplace_back("", c10::TensorType::get()); + // returns.emplace_back("", c10::TensorType::get()); + // returns.emplace_back("", c10::TensorType::get()); + // c10::FunctionSchema schema( + // "HybridEPBuffer::dispatch_with_permute", "", + // std::move(args), std::move(returns), false, false); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::dispatch_with_permute", &inputValues, outputValues); } return result; }; @@ -290,41 +291,42 @@ auto hybrid_ep_buffer_combine_with_unpermute(Func func) { num_dispatched_tokens_tensor, row_id_map, num_of_tokens_per_rank, pad_multiple, with_probs); if (at::isRecordFunctionEnabled()) { - c10::impl::GenericList inputList(c10::AnyType::get()); - inputList.emplace_back(to_ivalue(config)); - inputList.emplace_back(to_ivalue(hidden)); - inputList.emplace_back(to_ivalue(probs)); - inputList.emplace_back(to_ivalue(sparse_to_dense_map)); - inputList.emplace_back(to_ivalue(rdma_to_attn_map)); - inputList.emplace_back(to_ivalue(attn_to_rdma_map)); - inputList.emplace_back(to_ivalue(num_dispatched_tokens_tensor)); - inputList.emplace_back(to_ivalue(row_id_map)); - inputList.emplace_back(to_ivalue(num_of_tokens_per_rank)); - inputList.emplace_back(to_ivalue(pad_multiple)); - inputList.emplace_back(to_ivalue(with_probs)); - c10::impl::GenericList outputList(c10::AnyType::get()); - record_result(outputList, result); - c10::ArrayRef inputsArray(inputList); - c10::ArrayRef outputsArray(outputList); - std::vector args; - args.emplace_back("config", c10::AnyType::get()); - args.emplace_back("hidden", c10::TensorType::get()); - args.emplace_back("probs", c10::OptionalType::create(c10::TensorType::get())); - args.emplace_back("sparse_to_dense_map", c10::TensorType::get()); - args.emplace_back("rdma_to_attn_map", c10::TensorType::get()); - args.emplace_back("attn_to_rdma_map", c10::TensorType::get()); - args.emplace_back("num_dispatched_tokens_tensor", c10::OptionalType::create(c10::TensorType::get())); - args.emplace_back("row_id_map", c10::OptionalType::create(c10::TensorType::get())); - args.emplace_back("num_of_tokens_per_rank", c10::IntType::get()); - args.emplace_back("pad_multiple", c10::OptionalType::create(c10::IntType::get())); - args.emplace_back("with_probs", c10::BoolType::get()); - std::vector returns; - returns.emplace_back("", c10::TensorType::get()); - returns.emplace_back("", c10::TensorType::get()); - c10::FunctionSchema schema( - "HybridEPBuffer::combine_with_unpermute", "", - std::move(args), std::move(returns), false, false); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, inputsArray, outputsArray); + std::vector inputValues; + inputValues.reserve(10); + inputValues.emplace_back(to_ivalue(config)); + inputValues.emplace_back(to_ivalue(hidden)); + inputValues.emplace_back(to_ivalue(probs)); + inputValues.emplace_back(to_ivalue(sparse_to_dense_map)); + inputValues.emplace_back(to_ivalue(rdma_to_attn_map)); + inputValues.emplace_back(to_ivalue(attn_to_rdma_map)); + inputValues.emplace_back(to_ivalue(num_dispatched_tokens_tensor)); + inputValues.emplace_back(to_ivalue(row_id_map)); + inputValues.emplace_back(to_ivalue(num_of_tokens_per_rank)); + inputValues.emplace_back(to_ivalue(pad_multiple)); + inputValues.emplace_back(to_ivalue(with_probs)); + + std::vector outputValues; + record_result(outputValues, result); + + // std::vector args; + // args.emplace_back("config", c10::AnyType::get()); + // args.emplace_back("hidden", c10::TensorType::get()); + // args.emplace_back("probs", c10::OptionalType::create(c10::TensorType::get())); + // args.emplace_back("sparse_to_dense_map", c10::TensorType::get()); + // args.emplace_back("rdma_to_attn_map", c10::TensorType::get()); + // args.emplace_back("attn_to_rdma_map", c10::TensorType::get()); + // args.emplace_back("num_dispatched_tokens_tensor", c10::OptionalType::create(c10::TensorType::get())); + // args.emplace_back("row_id_map", c10::OptionalType::create(c10::TensorType::get())); + // args.emplace_back("num_of_tokens_per_rank", c10::IntType::get()); + // args.emplace_back("pad_multiple", c10::OptionalType::create(c10::IntType::get())); + // args.emplace_back("with_probs", c10::BoolType::get()); + // std::vector returns; + // returns.emplace_back("", c10::TensorType::get()); + // returns.emplace_back("", c10::TensorType::get()); + // c10::FunctionSchema schema( + // "HybridEPBuffer::combine_with_unpermute", "", + // std::move(args), std::move(returns), false, false); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::combine_with_unpermute", &inputValues, outputValues); } return result; }; From 68f1e92beacf15003c3c6477a80dfba52717589e Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Tue, 24 Feb 2026 03:59:50 +0200 Subject: [PATCH 13/22] fixed schema --- csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh | 172 +++++++++++----------- 1 file changed, 85 insertions(+), 87 deletions(-) diff --git a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh index 54f3d4c73..a77149f8e 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh +++ b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh @@ -58,22 +58,23 @@ inline HybridEPBuffer* hybrid_ep_buffer_init( inputValues.emplace_back(to_ivalue(enable_custom_allgather)); std::vector outputValues; - // std::vector args; - // args.emplace_back("group_name", c10::StringType::get()); - // args.emplace_back("config", c10::AnyType::get()); - // args.emplace_back("local_rank", c10::IntType::get()); - // args.emplace_back("node_rank", c10::IntType::get()); - // args.emplace_back("group_size", c10::IntType::get()); - // args.emplace_back("base_path", c10::StringType::get()); - // args.emplace_back("load_cached_kernels", c10::BoolType::get()); - // args.emplace_back("use_shared_buffer", c10::BoolType::get()); - // args.emplace_back("enable_custom_allgather", c10::BoolType::get()); - // std::vector returns; - // // outputList is empty for __init__ (buffer pointer not recorded as IValue) - // c10::FunctionSchema schema( - // "HybridEPBuffer::__init__", "", - // std::move(args), std::move(returns), false, false); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::__init__", &inputValues, outputValues); + + std::vector args; + args.emplace_back("group_name", c10::StringType::get()); + args.emplace_back("config", c10::AnyType::get()); + args.emplace_back("local_rank", c10::IntType::get()); + args.emplace_back("node_rank", c10::IntType::get()); + args.emplace_back("group_size", c10::IntType::get()); + args.emplace_back("base_path", c10::StringType::get()); + args.emplace_back("load_cached_kernels", c10::BoolType::get()); + args.emplace_back("use_shared_buffer", c10::BoolType::get()); + args.emplace_back("enable_custom_allgather", c10::BoolType::get()); + std::vector returns; + c10::FunctionSchema schema( + "HybridEPBuffer::__init__", "", + std::move(args), std::move(returns), false, false); + //RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::__init__", &inputValues, outputValues); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, &inputValues, outputValues); } return new HybridEPBuffer(process_group, config, local_rank, node_rank, group_size, base_path, load_cached_kernels, use_shared_buffer, enable_custom_allgather); @@ -107,9 +108,7 @@ auto hybrid_ep_buffer_combine(Func func) { std::vector outputValues; record_result(outputValues, result); - // c10::ArrayRef outputsArray(outputList); - /* - std::cout << "outputList size: " << outputList.size() << std::endl; + std::vector args; args.emplace_back("config", c10::AnyType::get()); args.emplace_back("hidden", c10::TensorType::get()); @@ -125,8 +124,7 @@ auto hybrid_ep_buffer_combine(Func func) { c10::FunctionSchema schema( "HybridEPBuffer::combine", "", std::move(args), std::move(returns), false, false); - */ - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::combine", &inputValues, outputValues); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, &inputValues, outputValues); } return result; }; @@ -168,26 +166,26 @@ auto hybrid_ep_buffer_dispatch(Func func) { std::vector outputValues; record_result(outputValues, result); - // std::vector args; - // args.emplace_back("config", c10::AnyType::get()); - // args.emplace_back("hidden", c10::TensorType::get()); - // args.emplace_back("probs", c10::OptionalType::create(c10::TensorType::get())); - // args.emplace_back("scaling_factor", c10::OptionalType::create(c10::TensorType::get())); - // args.emplace_back("sparse_to_dense_map", c10::TensorType::get()); - // args.emplace_back("rdma_to_attn_map", c10::TensorType::get()); - // args.emplace_back("attn_to_rdma_map", c10::TensorType::get()); - // args.emplace_back("num_dispatched_tokens_tensor", c10::OptionalType::create(c10::TensorType::get())); - // args.emplace_back("num_dispatched_tokens", c10::OptionalType::create(c10::IntType::get())); - // args.emplace_back("num_of_tokens_per_rank", c10::IntType::get()); - // args.emplace_back("with_probs", c10::BoolType::get()); - // std::vector returns; - // returns.emplace_back("", c10::TensorType::get()); - // returns.emplace_back("", c10::TensorType::get()); - // returns.emplace_back("", c10::TensorType::get()); - // c10::FunctionSchema schema( - // "HybridEPBuffer::dispatch", "", - // std::move(args), std::move(returns), false, false); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::dispatch", &inputValues, outputValues); + std::vector args; + args.emplace_back("config", c10::AnyType::get()); + args.emplace_back("hidden", c10::TensorType::get()); + args.emplace_back("probs", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("scaling_factor", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("sparse_to_dense_map", c10::TensorType::get()); + args.emplace_back("rdma_to_attn_map", c10::TensorType::get()); + args.emplace_back("attn_to_rdma_map", c10::TensorType::get()); + args.emplace_back("num_dispatched_tokens_tensor", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("num_dispatched_tokens", c10::OptionalType::create(c10::IntType::get())); + args.emplace_back("num_of_tokens_per_rank", c10::IntType::get()); + args.emplace_back("with_probs", c10::BoolType::get()); + std::vector returns; + returns.emplace_back("", c10::TensorType::get()); + returns.emplace_back("", c10::TensorType::get()); + returns.emplace_back("", c10::TensorType::get()); + c10::FunctionSchema schema( + "HybridEPBuffer::dispatch", "", + std::move(args), std::move(returns), false, false); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, &inputValues, outputValues); } return result; }; @@ -239,33 +237,33 @@ auto hybrid_ep_buffer_dispatch_with_permute(Func func) { std::vector outputValues; record_result(outputValues, result); - // std::vector args; - // args.emplace_back("config", c10::AnyType::get()); - // args.emplace_back("hidden", c10::TensorType::get()); - // args.emplace_back("probs", c10::OptionalType::create(c10::TensorType::get())); - // args.emplace_back("scaling_factor", c10::OptionalType::create(c10::TensorType::get())); - // args.emplace_back("sparse_to_dense_map", c10::TensorType::get()); - // args.emplace_back("rdma_to_attn_map", c10::TensorType::get()); - // args.emplace_back("attn_to_rdma_map", c10::TensorType::get()); - // args.emplace_back("num_dispatched_tokens_tensor", c10::OptionalType::create(c10::TensorType::get())); - // args.emplace_back("local_expert_routing_map", c10::OptionalType::create(c10::TensorType::get())); - // args.emplace_back("row_id_map", c10::OptionalType::create(c10::TensorType::get())); - // args.emplace_back("num_permuted_tokens", c10::OptionalType::create(c10::IntType::get())); - // args.emplace_back("num_of_tokens_per_rank", c10::IntType::get()); - // args.emplace_back("pad_multiple", c10::OptionalType::create(c10::IntType::get())); - // args.emplace_back("non_blocking", c10::BoolType::get()); - // args.emplace_back("with_probs", c10::BoolType::get()); - // std::vector returns; - // returns.emplace_back("", c10::TensorType::get()); - // returns.emplace_back("", c10::TensorType::get()); - // returns.emplace_back("", c10::TensorType::get()); - // returns.emplace_back("", c10::TensorType::get()); - // returns.emplace_back("", c10::TensorType::get()); - // returns.emplace_back("", c10::TensorType::get()); - // c10::FunctionSchema schema( - // "HybridEPBuffer::dispatch_with_permute", "", - // std::move(args), std::move(returns), false, false); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::dispatch_with_permute", &inputValues, outputValues); + std::vector args; + args.emplace_back("config", c10::AnyType::get()); + args.emplace_back("hidden", c10::TensorType::get()); + args.emplace_back("probs", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("scaling_factor", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("sparse_to_dense_map", c10::TensorType::get()); + args.emplace_back("rdma_to_attn_map", c10::TensorType::get()); + args.emplace_back("attn_to_rdma_map", c10::TensorType::get()); + args.emplace_back("num_dispatched_tokens_tensor", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("local_expert_routing_map", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("row_id_map", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("num_permuted_tokens", c10::OptionalType::create(c10::IntType::get())); + args.emplace_back("num_of_tokens_per_rank", c10::IntType::get()); + args.emplace_back("pad_multiple", c10::OptionalType::create(c10::IntType::get())); + args.emplace_back("non_blocking", c10::BoolType::get()); + args.emplace_back("with_probs", c10::BoolType::get()); + std::vector returns; + returns.emplace_back("", c10::TensorType::get()); + returns.emplace_back("", c10::TensorType::get()); + returns.emplace_back("", c10::TensorType::get()); + returns.emplace_back("", c10::TensorType::get()); + returns.emplace_back("", c10::TensorType::get()); + returns.emplace_back("", c10::TensorType::get()); + c10::FunctionSchema schema( + "HybridEPBuffer::dispatch_with_permute", "", + std::move(args), std::move(returns), false, false); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, &inputValues, outputValues); } return result; }; @@ -308,25 +306,25 @@ auto hybrid_ep_buffer_combine_with_unpermute(Func func) { std::vector outputValues; record_result(outputValues, result); - // std::vector args; - // args.emplace_back("config", c10::AnyType::get()); - // args.emplace_back("hidden", c10::TensorType::get()); - // args.emplace_back("probs", c10::OptionalType::create(c10::TensorType::get())); - // args.emplace_back("sparse_to_dense_map", c10::TensorType::get()); - // args.emplace_back("rdma_to_attn_map", c10::TensorType::get()); - // args.emplace_back("attn_to_rdma_map", c10::TensorType::get()); - // args.emplace_back("num_dispatched_tokens_tensor", c10::OptionalType::create(c10::TensorType::get())); - // args.emplace_back("row_id_map", c10::OptionalType::create(c10::TensorType::get())); - // args.emplace_back("num_of_tokens_per_rank", c10::IntType::get()); - // args.emplace_back("pad_multiple", c10::OptionalType::create(c10::IntType::get())); - // args.emplace_back("with_probs", c10::BoolType::get()); - // std::vector returns; - // returns.emplace_back("", c10::TensorType::get()); - // returns.emplace_back("", c10::TensorType::get()); - // c10::FunctionSchema schema( - // "HybridEPBuffer::combine_with_unpermute", "", - // std::move(args), std::move(returns), false, false); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::combine_with_unpermute", &inputValues, outputValues); + std::vector args; + args.emplace_back("config", c10::AnyType::get()); + args.emplace_back("hidden", c10::TensorType::get()); + args.emplace_back("probs", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("sparse_to_dense_map", c10::TensorType::get()); + args.emplace_back("rdma_to_attn_map", c10::TensorType::get()); + args.emplace_back("attn_to_rdma_map", c10::TensorType::get()); + args.emplace_back("num_dispatched_tokens_tensor", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("row_id_map", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("num_of_tokens_per_rank", c10::IntType::get()); + args.emplace_back("pad_multiple", c10::OptionalType::create(c10::IntType::get())); + args.emplace_back("with_probs", c10::BoolType::get()); + std::vector returns; + returns.emplace_back("", c10::TensorType::get()); + returns.emplace_back("", c10::TensorType::get()); + c10::FunctionSchema schema( + "HybridEPBuffer::combine_with_unpermute", "", + std::move(args), std::move(returns), false, false); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, &inputValues, outputValues); } return result; }; From 830e126d80779d4e9db6b06d9a03e909906c3233 Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Sat, 28 Feb 2026 00:24:29 +0200 Subject: [PATCH 14/22] add permute tests --- tests/test_hybrid_ep_collect_et.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/test_hybrid_ep_collect_et.py b/tests/test_hybrid_ep_collect_et.py index 7a7da3ea0..ab7dd8745 100644 --- a/tests/test_hybrid_ep_collect_et.py +++ b/tests/test_hybrid_ep_collect_et.py @@ -87,7 +87,8 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): stream = torch.cuda.Stream() with torch.cuda.stream(stream): - et = ExecutionTraceObserver().register_callback(f"rank-{dist.get_rank()}.json") + et = ExecutionTraceObserver().register_callback(f"./et/rank-{dist.get_rank()}.json") + et.set_extra_resource_collection(True) et.start() buffer = deep_ep.HybridEPBuffer( @@ -114,6 +115,15 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): dispatched_probs = None _, _ = buffer.combine(dispatched_hidden_bf16, dispatched_probs, handle) + + dispatched_hidden_with_permute, dispatched_probs_with_permute, _, _, handle_with_permute= ( + buffer.dispatch_with_permute(hidden=hidden, scaling_factor=scaling_factor, routing_map=routing_map, probs=probs, pad_multiple=PAD_MULTIPLE) + ) + dispatched_hidden_bf16_with_permute = dispatched_hidden_with_permute.to(torch.bfloat16) + + combine_with_unpermute_args = {'hidden': dispatched_hidden_bf16_with_permute, 'probs': dispatched_probs_with_permute, 'handle': handle_with_permute, 'pad_multiple': PAD_MULTIPLE} + buffer.combine_with_unpermute(**combine_with_unpermute_args) + et.stop() et.unregister_callback() From 3e706d0b0edad211a035f202ca19e954622022fe Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Sat, 28 Feb 2026 02:16:39 +0200 Subject: [PATCH 15/22] add two more APIs --- csrc/hybrid_ep/pybind_hybrid_ep.cu | 7 ++- csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh | 64 +++++++++++++++++++++++ tests/test_hybrid_ep_collect_et.py | 38 +++++++------- 3 files changed, 88 insertions(+), 21 deletions(-) diff --git a/csrc/hybrid_ep/pybind_hybrid_ep.cu b/csrc/hybrid_ep/pybind_hybrid_ep.cu index 92570ee9e..f36464ee4 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep.cu +++ b/csrc/hybrid_ep/pybind_hybrid_ep.cu @@ -132,8 +132,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("load_cached_kernels") = false, py::arg("use_shared_buffer") = true, py::arg("enable_custom_allgather") = true) - .def("update_buffer", &HybridEPBuffer::update_buffer, py::arg("config")) - .def("metadata_preprocessing", &HybridEPBuffer::metadata_preprocessing, + .def("update_buffer", + hybrid_ep_buffer_update_buffer(&HybridEPBuffer::update_buffer), + py::arg("config")) + .def("metadata_preprocessing", + hybrid_ep_buffer_metadata_preprocessing(&HybridEPBuffer::metadata_preprocessing), py::kw_only(), py::arg("config"), py::arg("routing_map"), py::arg("num_of_tokens_per_rank"), py::arg("non_blocking") = false) .def("dispatch", hybrid_ep_buffer_dispatch(&HybridEPBuffer::dispatch), diff --git a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh index a77149f8e..4b3e9d1df 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh +++ b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh @@ -130,6 +130,70 @@ auto hybrid_ep_buffer_combine(Func func) { }; } +// Wrapper for update_buffer (returns bool) +template +auto hybrid_ep_buffer_update_buffer(Func func) { + return [func](HybridEPBuffer& self, HybridEpConfigInstance config) { + auto result = (self.*func)(config); + if (at::isRecordFunctionEnabled()) { + std::vector inputValues; + inputValues.reserve(1); + inputValues.push_back(to_ivalue(config)); + + std::vector outputValues; + outputValues.push_back(c10::IValue(result)); + + std::vector args; + args.emplace_back("config", c10::AnyType::get()); + std::vector returns; + returns.emplace_back("", c10::BoolType::get()); + c10::FunctionSchema schema( + "HybridEPBuffer::update_buffer", "", + std::move(args), std::move(returns), false, false); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, &inputValues, outputValues); + } + return result; + }; +} + +// Wrapper for metadata_preprocessing (returns 5-tuple of tensors) +template +auto hybrid_ep_buffer_metadata_preprocessing(Func func) { + return [func](HybridEPBuffer& self, + HybridEpConfigInstance config, + torch::Tensor routing_map, + int64_t num_of_tokens_per_rank, + bool non_blocking) { + auto result = (self.*func)(config, routing_map, num_of_tokens_per_rank, non_blocking); + if (at::isRecordFunctionEnabled()) { + std::vector inputValues; + inputValues.reserve(4); + inputValues.push_back(to_ivalue(config)); + inputValues.push_back(to_ivalue(routing_map)); + inputValues.push_back(to_ivalue(num_of_tokens_per_rank)); + inputValues.push_back(to_ivalue(non_blocking)); + + std::vector outputValues; + record_result(outputValues, result); + + std::vector args; + args.emplace_back("config", c10::AnyType::get()); + args.emplace_back("routing_map", c10::TensorType::get()); + args.emplace_back("num_of_tokens_per_rank", c10::IntType::get()); + args.emplace_back("non_blocking", c10::BoolType::get()); + std::vector returns; + for (int i = 0; i < 5; ++i) { + returns.emplace_back("", c10::TensorType::get()); + } + c10::FunctionSchema schema( + "HybridEPBuffer::metadata_preprocessing", "", + std::move(args), std::move(returns), false, false); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, &inputValues, outputValues); + } + return result; + }; +} + // Wrapper for dispatch (returns 3-tuple) template auto hybrid_ep_buffer_dispatch(Func func) { diff --git a/tests/test_hybrid_ep_collect_et.py b/tests/test_hybrid_ep_collect_et.py index ab7dd8745..eec3e67f3 100644 --- a/tests/test_hybrid_ep_collect_et.py +++ b/tests/test_hybrid_ep_collect_et.py @@ -87,7 +87,15 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): stream = torch.cuda.Stream() with torch.cuda.stream(stream): - et = ExecutionTraceObserver().register_callback(f"./et/rank-{dist.get_rank()}.json") + hidden, probs, scaling_factor, routing_map, topk_idx, topk_weights = init_tensor( + hidden_dim=HIDDEN_DIM, + seq_len=NUM_TOKENS_PER_RANK, + topk=TOPK, + num_of_experts=NUM_OF_EXPERTS, + use_fp8=truediv, + ) + + et = ExecutionTraceObserver().register_callback(f"./et-permute/rank-{dist.get_rank()}.json") et.set_extra_resource_collection(True) et.start() @@ -99,30 +107,22 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): use_fp8=True ) - hidden, probs, scaling_factor, routing_map, topk_idx, topk_weights = init_tensor( - hidden_dim=HIDDEN_DIM, - seq_len=NUM_TOKENS_PER_RANK, - topk=TOPK, - num_of_experts=NUM_OF_EXPERTS, - use_fp8=truediv, - ) - - 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) - ) - # The combine only support bf16 - dispatched_hidden_bf16 = dispatched_hidden.to(torch.bfloat16) - dispatched_probs = None - _, _ = buffer.combine(dispatched_hidden_bf16, dispatched_probs, handle) + # 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) + # ) + # # The combine only support bf16 + # dispatched_hidden_bf16 = dispatched_hidden.to(torch.bfloat16) + # dispatched_probs = None + # _, _ = buffer.combine(dispatched_hidden_bf16, dispatched_probs, handle) dispatched_hidden_with_permute, dispatched_probs_with_permute, _, _, handle_with_permute= ( buffer.dispatch_with_permute(hidden=hidden, scaling_factor=scaling_factor, routing_map=routing_map, probs=probs, pad_multiple=PAD_MULTIPLE) ) - dispatched_hidden_bf16_with_permute = dispatched_hidden_with_permute.to(torch.bfloat16) + # dispatched_hidden_bf16_with_permute = dispatched_hidden_with_permute.to(torch.bfloat16) - combine_with_unpermute_args = {'hidden': dispatched_hidden_bf16_with_permute, 'probs': dispatched_probs_with_permute, 'handle': handle_with_permute, 'pad_multiple': PAD_MULTIPLE} - buffer.combine_with_unpermute(**combine_with_unpermute_args) + # combine_with_unpermute_args = {'hidden': dispatched_hidden_bf16_with_permute, 'probs': dispatched_probs_with_permute, 'handle': handle_with_permute, 'pad_multiple': PAD_MULTIPLE} + # buffer.combine_with_unpermute(**combine_with_unpermute_args) et.stop() et.unregister_callback() From 67a1d21b6e5dff8adbe15a29149cae473dfdd89d Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Tue, 3 Mar 2026 00:27:15 +0200 Subject: [PATCH 16/22] support four APIs --- tests/test_hybrid_ep_collect_et.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/test_hybrid_ep_collect_et.py b/tests/test_hybrid_ep_collect_et.py index eec3e67f3..acf01abbf 100644 --- a/tests/test_hybrid_ep_collect_et.py +++ b/tests/test_hybrid_ep_collect_et.py @@ -107,22 +107,22 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): use_fp8=True ) - # 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) - # ) - # # The combine only support bf16 - # dispatched_hidden_bf16 = dispatched_hidden.to(torch.bfloat16) - # dispatched_probs = None - # _, _ = buffer.combine(dispatched_hidden_bf16, dispatched_probs, handle) + 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) + ) + # The combine only support bf16 + dispatched_hidden_bf16 = dispatched_hidden.to(torch.bfloat16) + dispatched_probs = None + _, _ = buffer.combine(dispatched_hidden_bf16, dispatched_probs, handle) dispatched_hidden_with_permute, dispatched_probs_with_permute, _, _, handle_with_permute= ( buffer.dispatch_with_permute(hidden=hidden, scaling_factor=scaling_factor, routing_map=routing_map, probs=probs, pad_multiple=PAD_MULTIPLE) ) - # dispatched_hidden_bf16_with_permute = dispatched_hidden_with_permute.to(torch.bfloat16) + dispatched_hidden_bf16_with_permute = dispatched_hidden_with_permute.to(torch.bfloat16) - # combine_with_unpermute_args = {'hidden': dispatched_hidden_bf16_with_permute, 'probs': dispatched_probs_with_permute, 'handle': handle_with_permute, 'pad_multiple': PAD_MULTIPLE} - # buffer.combine_with_unpermute(**combine_with_unpermute_args) + combine_with_unpermute_args = {'hidden': dispatched_hidden_bf16_with_permute, 'probs': dispatched_probs_with_permute, 'handle': handle_with_permute, 'pad_multiple': PAD_MULTIPLE} + buffer.combine_with_unpermute(**combine_with_unpermute_args) et.stop() et.unregister_callback() From 48535628291f7c8385297a44b36eb11df2f46280 Mon Sep 17 00:00:00 2001 From: shengfu-nv Date: Fri, 6 Mar 2026 00:52:20 +0200 Subject: [PATCH 17/22] record HybridEPBuffer init when profiling is on --- csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh | 82 ++++++++++++++++++----- 1 file changed, 65 insertions(+), 17 deletions(-) diff --git a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh index 4b3e9d1df..47c188149 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh +++ b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh @@ -9,7 +9,10 @@ #include #include #include +#include +#include #include +#include #include #include #include @@ -40,24 +43,56 @@ inline void record_result(std::vector& outputList, Result&& result) }, result); } +// Struct to hold recording data for HybridEPBuffer::__init__ (can be consumed by record_hybrid_ep_buffer_init). +struct HybridEpBufferInitInfo { + std::shared_mutex mutex; + std::vector inputValues; + std::optional schema; + bool recorded = false; +}; + +HybridEpBufferInitInfo hybrid_ep_buffer_init_info; + +// Call this to record HybridEPBuffer::__init__ using data stored in hybrid_ep_buffer_init_info. +inline void record_hybrid_ep_buffer_init() { + { + std::shared_lock lock(hybrid_ep_buffer_init_info.mutex); + if (hybrid_ep_buffer_init_info.recorded) return; + } + std::unique_lock lock(hybrid_ep_buffer_init_info.mutex); + if (hybrid_ep_buffer_init_info.recorded) return; + + std::vector outputValues; + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS( + hybrid_ep_buffer_init_info.schema.value(), + &hybrid_ep_buffer_init_info.inputValues, + outputValues); + hybrid_ep_buffer_init_info.recorded = true; +} + inline HybridEPBuffer* hybrid_ep_buffer_init( py::object process_group, BufferConfig config, int local_rank, int node_rank, int group_size, std::string base_path, bool load_cached_kernels, bool use_shared_buffer, bool enable_custom_allgather) { - if (at::isRecordFunctionEnabled()) { - std::vector inputValues; - inputValues.reserve(9); - inputValues.emplace_back(c10::IValue(process_group.attr("group_name").cast())); - inputValues.emplace_back(to_ivalue(config)); - inputValues.emplace_back(to_ivalue(local_rank)); - inputValues.emplace_back(to_ivalue(node_rank)); - inputValues.emplace_back(to_ivalue(group_size)); - inputValues.emplace_back(to_ivalue(base_path)); - inputValues.emplace_back(to_ivalue(load_cached_kernels)); - inputValues.emplace_back(to_ivalue(use_shared_buffer)); - inputValues.emplace_back(to_ivalue(enable_custom_allgather)); - - std::vector outputValues; + + // Initializing HybridEPBuffer may happens before pytorch profiling starts + // save the information to record it later when other operators are called. + // It is based on the assumption there is only one HybridEPBuffer instance + // in one process + { + std::unique_lock lock(hybrid_ep_buffer_init_info.mutex); + + hybrid_ep_buffer_init_info.inputValues.clear(); + hybrid_ep_buffer_init_info.inputValues.reserve(9); + hybrid_ep_buffer_init_info.inputValues.emplace_back(c10::IValue(process_group.attr("group_name").cast())); + hybrid_ep_buffer_init_info.inputValues.emplace_back(to_ivalue(config)); + hybrid_ep_buffer_init_info.inputValues.emplace_back(to_ivalue(local_rank)); + hybrid_ep_buffer_init_info.inputValues.emplace_back(to_ivalue(node_rank)); + hybrid_ep_buffer_init_info.inputValues.emplace_back(to_ivalue(group_size)); + hybrid_ep_buffer_init_info.inputValues.emplace_back(to_ivalue(base_path)); + hybrid_ep_buffer_init_info.inputValues.emplace_back(to_ivalue(load_cached_kernels)); + hybrid_ep_buffer_init_info.inputValues.emplace_back(to_ivalue(use_shared_buffer)); + hybrid_ep_buffer_init_info.inputValues.emplace_back(to_ivalue(enable_custom_allgather)); std::vector args; args.emplace_back("group_name", c10::StringType::get()); @@ -70,12 +105,13 @@ inline HybridEPBuffer* hybrid_ep_buffer_init( args.emplace_back("use_shared_buffer", c10::BoolType::get()); args.emplace_back("enable_custom_allgather", c10::BoolType::get()); std::vector returns; - c10::FunctionSchema schema( + hybrid_ep_buffer_init_info.schema.emplace( "HybridEPBuffer::__init__", "", std::move(args), std::move(returns), false, false); - //RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::__init__", &inputValues, outputValues); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, &inputValues, outputValues); + + hybrid_ep_buffer_init_info.recorded = false; } + return new HybridEPBuffer(process_group, config, local_rank, node_rank, group_size, base_path, load_cached_kernels, use_shared_buffer, enable_custom_allgather); } @@ -95,6 +131,8 @@ auto hybrid_ep_buffer_combine(Func func) { rdma_to_attn_map, attn_to_rdma_map, num_of_tokens_per_rank, with_probs); if (at::isRecordFunctionEnabled()) { + record_hybrid_ep_buffer_init(); + std::vector inputValues; inputValues.reserve(8); inputValues.push_back(to_ivalue(config)); @@ -136,6 +174,8 @@ auto hybrid_ep_buffer_update_buffer(Func func) { return [func](HybridEPBuffer& self, HybridEpConfigInstance config) { auto result = (self.*func)(config); if (at::isRecordFunctionEnabled()) { + record_hybrid_ep_buffer_init(); + std::vector inputValues; inputValues.reserve(1); inputValues.push_back(to_ivalue(config)); @@ -166,6 +206,8 @@ auto hybrid_ep_buffer_metadata_preprocessing(Func func) { bool non_blocking) { auto result = (self.*func)(config, routing_map, num_of_tokens_per_rank, non_blocking); if (at::isRecordFunctionEnabled()) { + record_hybrid_ep_buffer_init(); + std::vector inputValues; inputValues.reserve(4); inputValues.push_back(to_ivalue(config)); @@ -214,6 +256,8 @@ auto hybrid_ep_buffer_dispatch(Func func) { num_dispatched_tokens_tensor, num_dispatched_tokens, num_of_tokens_per_rank, with_probs); if (at::isRecordFunctionEnabled()) { + record_hybrid_ep_buffer_init(); + std::vector inputValues; inputValues.reserve(11); inputValues.push_back(to_ivalue(config)); @@ -280,6 +324,8 @@ auto hybrid_ep_buffer_dispatch_with_permute(Func func) { num_permuted_tokens, num_of_tokens_per_rank, pad_multiple, non_blocking, with_probs); if (at::isRecordFunctionEnabled()) { + record_hybrid_ep_buffer_init(); + std::vector inputValues; inputValues.reserve(15); inputValues.push_back(to_ivalue(config)); @@ -353,6 +399,8 @@ auto hybrid_ep_buffer_combine_with_unpermute(Func func) { num_dispatched_tokens_tensor, row_id_map, num_of_tokens_per_rank, pad_multiple, with_probs); if (at::isRecordFunctionEnabled()) { + record_hybrid_ep_buffer_init(); + std::vector inputValues; inputValues.reserve(10); inputValues.emplace_back(to_ivalue(config)); From 2fc7aa9d557d20bc52dfaf4390688ba7beb10628 Mon Sep 17 00:00:00 2001 From: shengf Date: Fri, 6 Mar 2026 18:00:52 -0800 Subject: [PATCH 18/22] collect et with kineto trace --- tests/test_hybrid_ep_collect_et.py | 53 ++++++++++++++++++++---------- 1 file changed, 35 insertions(+), 18 deletions(-) diff --git a/tests/test_hybrid_ep_collect_et.py b/tests/test_hybrid_ep_collect_et.py index acf01abbf..17a834365 100644 --- a/tests/test_hybrid_ep_collect_et.py +++ b/tests/test_hybrid_ep_collect_et.py @@ -95,10 +95,6 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): use_fp8=truediv, ) - et = ExecutionTraceObserver().register_callback(f"./et-permute/rank-{dist.get_rank()}.json") - et.set_extra_resource_collection(True) - et.start() - buffer = deep_ep.HybridEPBuffer( group=group, hidden_dim=HIDDEN_DIM, @@ -107,26 +103,47 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): use_fp8=True ) - 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) + et = ExecutionTraceObserver().register_callback(f"./et-permute/rank-{dist.get_rank()}.json") + et.set_extra_resource_collection(True) + + prof = torch.profiler.profile( + schedule=torch.profiler.schedule( + wait=0, + warmup=0, + active=1, + repeat=1, + ), + record_shapes=False, + with_stack=False, + execution_trace_observer=et, ) - # The combine only support bf16 - dispatched_hidden_bf16 = dispatched_hidden.to(torch.bfloat16) - dispatched_probs = None - _, _ = buffer.combine(dispatched_hidden_bf16, dispatched_probs, handle) + prof.start() + for i in range(5): - dispatched_hidden_with_permute, dispatched_probs_with_permute, _, _, handle_with_permute= ( - buffer.dispatch_with_permute(hidden=hidden, scaling_factor=scaling_factor, routing_map=routing_map, probs=probs, pad_multiple=PAD_MULTIPLE) - ) - dispatched_hidden_bf16_with_permute = dispatched_hidden_with_permute.to(torch.bfloat16) - combine_with_unpermute_args = {'hidden': dispatched_hidden_bf16_with_permute, 'probs': dispatched_probs_with_permute, 'handle': handle_with_permute, 'pad_multiple': PAD_MULTIPLE} - buffer.combine_with_unpermute(**combine_with_unpermute_args) + 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) + ) + # The combine only support bf16 + dispatched_hidden_bf16 = dispatched_hidden.to(torch.bfloat16) + dispatched_probs = None + _, _ = buffer.combine(dispatched_hidden_bf16, dispatched_probs, handle) - et.stop() + + dispatched_hidden_with_permute, dispatched_probs_with_permute, _, _, handle_with_permute= ( + buffer.dispatch_with_permute(hidden=hidden, scaling_factor=scaling_factor, routing_map=routing_map, probs=probs, pad_multiple=PAD_MULTIPLE) + ) + dispatched_hidden_bf16_with_permute = dispatched_hidden_with_permute.to(torch.bfloat16) + + combine_with_unpermute_args = {'hidden': dispatched_hidden_bf16_with_permute, 'probs': dispatched_probs_with_permute, 'handle': handle_with_permute, 'pad_multiple': PAD_MULTIPLE} + buffer.combine_with_unpermute(**combine_with_unpermute_args) + + prof.step() + + prof.stop() et.unregister_callback() - + time.sleep(10) dist.barrier() dist.destroy_process_group() From c450e130126bf0d8ee0df9c287e755bded28c7b1 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 16 Mar 2026 17:12:55 -0700 Subject: [PATCH 19/22] check if profiler is enabled --- csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh | 16 +++++++++------- tests/test_hybrid_ep_collect_et.py | 6 ++---- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh index 47c188149..55956c388 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh +++ b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh @@ -16,7 +16,9 @@ #include #include #include -#include +#include + +using namespace torch::autograd::profiler; namespace py = pybind11; @@ -130,7 +132,7 @@ auto hybrid_ep_buffer_combine(Func func) { auto result = (self.*func)(config, hidden, probs, sparse_to_dense_map, rdma_to_attn_map, attn_to_rdma_map, num_of_tokens_per_rank, with_probs); - if (at::isRecordFunctionEnabled()) { + if (isProfilerEnabledInMainThread()) { record_hybrid_ep_buffer_init(); std::vector inputValues; @@ -173,7 +175,7 @@ template auto hybrid_ep_buffer_update_buffer(Func func) { return [func](HybridEPBuffer& self, HybridEpConfigInstance config) { auto result = (self.*func)(config); - if (at::isRecordFunctionEnabled()) { + if (isProfilerEnabledInMainThread()) { record_hybrid_ep_buffer_init(); std::vector inputValues; @@ -205,7 +207,7 @@ auto hybrid_ep_buffer_metadata_preprocessing(Func func) { int64_t num_of_tokens_per_rank, bool non_blocking) { auto result = (self.*func)(config, routing_map, num_of_tokens_per_rank, non_blocking); - if (at::isRecordFunctionEnabled()) { + if (isProfilerEnabledInMainThread()) { record_hybrid_ep_buffer_init(); std::vector inputValues; @@ -255,7 +257,7 @@ auto hybrid_ep_buffer_dispatch(Func func) { sparse_to_dense_map, rdma_to_attn_map, attn_to_rdma_map, num_dispatched_tokens_tensor, num_dispatched_tokens, num_of_tokens_per_rank, with_probs); - if (at::isRecordFunctionEnabled()) { + if (isProfilerEnabledInMainThread()) { record_hybrid_ep_buffer_init(); std::vector inputValues; @@ -323,7 +325,7 @@ auto hybrid_ep_buffer_dispatch_with_permute(Func func) { num_dispatched_tokens_tensor, local_expert_routing_map, row_id_map, num_permuted_tokens, num_of_tokens_per_rank, pad_multiple, non_blocking, with_probs); - if (at::isRecordFunctionEnabled()) { + if (isProfilerEnabledInMainThread()) { record_hybrid_ep_buffer_init(); std::vector inputValues; @@ -398,7 +400,7 @@ auto hybrid_ep_buffer_combine_with_unpermute(Func func) { rdma_to_attn_map, attn_to_rdma_map, num_dispatched_tokens_tensor, row_id_map, num_of_tokens_per_rank, pad_multiple, with_probs); - if (at::isRecordFunctionEnabled()) { + if (isProfilerEnabledInMainThread()) { record_hybrid_ep_buffer_init(); std::vector inputValues; diff --git a/tests/test_hybrid_ep_collect_et.py b/tests/test_hybrid_ep_collect_et.py index 17a834365..a36a9fa70 100644 --- a/tests/test_hybrid_ep_collect_et.py +++ b/tests/test_hybrid_ep_collect_et.py @@ -108,8 +108,8 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): prof = torch.profiler.profile( schedule=torch.profiler.schedule( - wait=0, - warmup=0, + wait=1, + warmup=2, active=1, repeat=1, ), @@ -120,8 +120,6 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): prof.start() for i in range(5): - - 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 e1d76c649c16d2962f8e0f7248654497bd96e91e Mon Sep 17 00:00:00 2001 From: shengf Date: Fri, 29 May 2026 20:38:37 -0700 Subject: [PATCH 20/22] expose attributes --- csrc/hybrid_ep/pybind_hybrid_ep.cu | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/csrc/hybrid_ep/pybind_hybrid_ep.cu b/csrc/hybrid_ep/pybind_hybrid_ep.cu index 61124f445..b362c95f2 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep.cu +++ b/csrc/hybrid_ep/pybind_hybrid_ep.cu @@ -110,6 +110,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { &HybridEpConfigInstance::num_of_stages_g2s_combine_api) .def_readwrite("num_of_stages_s2g_combine_api", &HybridEpConfigInstance::num_of_stages_s2g_combine_api) + .def_readwrite("num_of_stages_g2s_unpermute_block", + &HybridEpConfigInstance::num_of_stages_g2s_unpermute_block) + .def_readwrite("num_of_stages_s2g_unpermute_block", + &HybridEpConfigInstance::num_of_stages_s2g_unpermute_block) .def_readwrite("num_of_tokens_per_chunk_combine_api", &HybridEpConfigInstance::num_of_tokens_per_chunk_combine_api) .def_readwrite("num_of_tokens_per_group_combine_api", @@ -119,6 +123,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def_readwrite( "num_of_additional_in_flight_s2g_combine_api", &HybridEpConfigInstance::num_of_additional_in_flight_s2g_combine_api) + .def_readwrite( + "num_of_additional_in_flight_s2g_unpermute_block_combine_api", + &HybridEpConfigInstance::num_of_additional_in_flight_s2g_unpermute_block_combine_api) .def_readwrite("backward_combine_api", &HybridEpConfigInstance::backward_combine_api) .def_readwrite("device_side_sync_combine_api", From 9616d6d2437895a3193e0da8a2d5f5d5112e6369 Mon Sep 17 00:00:00 2001 From: shengf Date: Mon, 1 Jun 2026 15:59:06 -0700 Subject: [PATCH 21/22] fix test --- tests/test_hybrid_ep_collect_et.py | 50 ++++++++++++++++++------------ 1 file changed, 31 insertions(+), 19 deletions(-) diff --git a/tests/test_hybrid_ep_collect_et.py b/tests/test_hybrid_ep_collect_et.py index a36a9fa70..cd4a4e99c 100644 --- a/tests/test_hybrid_ep_collect_et.py +++ b/tests/test_hybrid_ep_collect_et.py @@ -108,8 +108,8 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): prof = torch.profiler.profile( schedule=torch.profiler.schedule( - wait=1, - warmup=2, + wait=0, + warmup=0, active=1, repeat=1, ), @@ -119,26 +119,38 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): ) prof.start() - for i in range(5): - 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) - ) - # The combine only support bf16 - dispatched_hidden_bf16 = dispatched_hidden.to(torch.bfloat16) - dispatched_probs = None - _, _ = buffer.combine(dispatched_hidden_bf16, dispatched_probs, handle) - - - dispatched_hidden_with_permute, dispatched_probs_with_permute, _, _, handle_with_permute= ( - buffer.dispatch_with_permute(hidden=hidden, scaling_factor=scaling_factor, routing_map=routing_map, probs=probs, pad_multiple=PAD_MULTIPLE) - ) - dispatched_hidden_bf16_with_permute = dispatched_hidden_with_permute.to(torch.bfloat16) - - combine_with_unpermute_args = {'hidden': dispatched_hidden_bf16_with_permute, 'probs': dispatched_probs_with_permute, 'handle': handle_with_permute, 'pad_multiple': PAD_MULTIPLE} - buffer.combine_with_unpermute(**combine_with_unpermute_args) + for i in range(1): + # Non-permute + 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)) + # dispatched_hidden_bf16 = dispatched_hidden.to(torch.bfloat16) + # combine_args = {'hidden': dispatched_hidden_bf16, 'probs': dispatched_probs, 'handle': handle} + # buffer.combine(**combine_args) + + # # Permute (non-fused) + # dispatched_hidden_wp, dispatched_probs_wp, _, tpe_wp, handle_wp = ( + # buffer.dispatch_with_permute(hidden=hidden, scaling_factor=scaling_factor, + # routing_map=routing_map, probs=probs, pad_multiple=PAD_MULTIPLE)) + # dispatched_hidden_bf16_wp = dispatched_hidden_wp.to(torch.bfloat16) + # combine_wp_args = {'hidden': dispatched_hidden_bf16_wp, 'probs': dispatched_probs_wp, + # 'handle': handle_wp, 'pad_multiple': PAD_MULTIPLE} + # buffer.combine_with_unpermute(**combine_wp_args) + + # # Fused permute-dispatch + # dispatched_hidden_fused, dispatched_probs_fused, _, tpe_fused, handle_fused = ( + # buffer.dispatch_with_permute(hidden=hidden, scaling_factor=scaling_factor, + # routing_map=routing_map, probs=probs, pad_multiple=PAD_MULTIPLE, fuse_permute_dispatch=True)) + # dispatched_hidden_bf16_fused = dispatched_hidden_fused.to(torch.bfloat16) + + # combine_fused_args = {'hidden': dispatched_hidden_bf16_fused, 'probs': dispatched_probs_fused, + # 'handle': handle_fused, 'pad_multiple': PAD_MULTIPLE, 'fuse_unpermute_combine': True} + # buffer.combine_with_unpermute(**combine_fused_args) + prof.step() + prof.stop() et.unregister_callback() From 9b2a80ac20b02935412e5484312f868a5641a74d Mon Sep 17 00:00:00 2001 From: shengf Date: Tue, 2 Jun 2026 21:34:34 -0700 Subject: [PATCH 22/22] dispatch combine worked --- csrc/hybrid_ep/config.cuh | 107 +++++++++++----------- csrc/hybrid_ep/executor/executor.cuh | 34 +++---- csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh | 95 +++++++++---------- tests/test_hybrid_ep_collect_et.py | 31 ++++--- 4 files changed, 130 insertions(+), 137 deletions(-) diff --git a/csrc/hybrid_ep/config.cuh b/csrc/hybrid_ep/config.cuh index 514caab0b..7c6046e07 100644 --- a/csrc/hybrid_ep/config.cuh +++ b/csrc/hybrid_ep/config.cuh @@ -11,6 +11,7 @@ #include #include "utils.cuh" #include +#include // Now we support up to 72(GB200) ranks per node. // This will be used to initialize the template param_t for communication kernel. @@ -102,24 +103,24 @@ struct BufferConfig { return valid; } - /** Convert all attributes to a single IValue holding a tuple of IValues (ints/string). */ + /** Convert all attributes to a single IValue holding a GenericList of IValues (ints/string). */ c10::IValue to_ivalue_tuple() const { - std::vector elements; - elements.reserve(13); - elements.push_back(c10::IValue(static_cast(hidden_dim))); - elements.push_back(c10::IValue(static_cast(max_num_of_tokens_per_rank))); - elements.push_back(c10::IValue(static_cast(num_of_experts_per_rank))); - elements.push_back(c10::IValue(static_cast(num_of_ranks_per_node))); - elements.push_back(c10::IValue(static_cast(num_of_nodes))); - elements.push_back(c10::IValue(type_to_string(token_data_type))); - elements.push_back(c10::IValue(static_cast(num_of_blocks_preprocessing_api))); - elements.push_back(c10::IValue(static_cast(num_of_blocks_dispatch_api))); - elements.push_back(c10::IValue(static_cast(num_of_blocks_combine_api))); - elements.push_back(c10::IValue(static_cast(num_of_tokens_per_chunk_dispatch_api))); - elements.push_back(c10::IValue(static_cast(num_of_tokens_per_chunk_combine_api))); - elements.push_back(c10::IValue(static_cast(num_of_dispatch_chunks))); - elements.push_back(c10::IValue(static_cast(num_of_combine_chunks))); - return c10::IValue(c10::ivalue::Tuple::create(std::move(elements))); + c10::impl::GenericList list(c10::AnyType::get()); + list.reserve(13); + list.push_back(c10::IValue(static_cast(hidden_dim))); + list.push_back(c10::IValue(static_cast(max_num_of_tokens_per_rank))); + list.push_back(c10::IValue(static_cast(num_of_experts_per_rank))); + list.push_back(c10::IValue(static_cast(num_of_ranks_per_node))); + list.push_back(c10::IValue(static_cast(num_of_nodes))); + list.push_back(c10::IValue(type_to_string(token_data_type))); + list.push_back(c10::IValue(static_cast(num_of_blocks_preprocessing_api))); + list.push_back(c10::IValue(static_cast(num_of_blocks_dispatch_api))); + list.push_back(c10::IValue(static_cast(num_of_blocks_combine_api))); + list.push_back(c10::IValue(static_cast(num_of_tokens_per_chunk_dispatch_api))); + list.push_back(c10::IValue(static_cast(num_of_tokens_per_chunk_combine_api))); + list.push_back(c10::IValue(static_cast(num_of_dispatch_chunks))); + list.push_back(c10::IValue(static_cast(num_of_combine_chunks))); + return c10::IValue(std::move(list)); } }; @@ -210,47 +211,47 @@ struct HybridEpConfigInstance { return valid; } - /** Convert all attributes to a single IValue holding a tuple of IValues (ints/bools/string). */ + /** Convert all attributes to a single IValue holding a GenericList of IValues (ints/bools/string). */ c10::IValue to_ivalue_tuple() const { - std::vector elements; - elements.reserve(32); + c10::impl::GenericList list(c10::AnyType::get()); + list.reserve(32); // Hybrid-ep Config - elements.push_back(c10::IValue(static_cast(hidden_dim))); - elements.push_back(c10::IValue(static_cast(max_num_of_tokens_per_rank))); - elements.push_back(c10::IValue(static_cast(num_of_experts_per_rank))); - elements.push_back(c10::IValue(static_cast(num_of_ranks_per_node))); - elements.push_back(c10::IValue(static_cast(num_of_nodes))); - elements.push_back(c10::IValue(static_cast(pad_multiple))); + list.push_back(c10::IValue(static_cast(hidden_dim))); + list.push_back(c10::IValue(static_cast(max_num_of_tokens_per_rank))); + list.push_back(c10::IValue(static_cast(num_of_experts_per_rank))); + list.push_back(c10::IValue(static_cast(num_of_ranks_per_node))); + list.push_back(c10::IValue(static_cast(num_of_nodes))); + list.push_back(c10::IValue(static_cast(pad_multiple))); // Metadata-preprocessing API Config - elements.push_back(c10::IValue(static_cast(num_of_tokens_per_chunk_preprocessing_api))); - elements.push_back(c10::IValue(static_cast(num_of_threads_per_block_preprocessing_api))); - elements.push_back(c10::IValue(static_cast(num_of_blocks_preprocessing_api))); - elements.push_back(c10::IValue(static_cast(num_of_blocks_permute))); - elements.push_back(c10::IValue(static_cast(num_of_blocks_unpermute))); + list.push_back(c10::IValue(static_cast(num_of_tokens_per_chunk_preprocessing_api))); + list.push_back(c10::IValue(static_cast(num_of_threads_per_block_preprocessing_api))); + list.push_back(c10::IValue(static_cast(num_of_blocks_preprocessing_api))); + list.push_back(c10::IValue(static_cast(num_of_blocks_permute))); + list.push_back(c10::IValue(static_cast(num_of_blocks_unpermute))); // Dispatch API Config - elements.push_back(c10::IValue(type_to_string(token_data_type))); - elements.push_back(c10::IValue(static_cast(num_of_stages_dispatch_api))); - elements.push_back(c10::IValue(static_cast(num_of_stages_permute_block_dispatch_api))); - elements.push_back(c10::IValue(static_cast(num_of_in_flight_s2g_dispatch_api))); - elements.push_back(c10::IValue(static_cast(num_of_in_flight_s2g_permute_block_dispatch_api))); - elements.push_back(c10::IValue(static_cast(num_of_additional_in_flight_s2g_dispatch_api))); - elements.push_back(c10::IValue(static_cast(num_of_tokens_per_chunk_dispatch_api))); - elements.push_back(c10::IValue(static_cast(num_of_blocks_dispatch_api))); - elements.push_back(c10::IValue(forward_dispatch_api)); - elements.push_back(c10::IValue(device_side_sync_dispatch_api)); + list.push_back(c10::IValue(type_to_string(token_data_type))); + list.push_back(c10::IValue(static_cast(num_of_stages_dispatch_api))); + list.push_back(c10::IValue(static_cast(num_of_stages_permute_block_dispatch_api))); + list.push_back(c10::IValue(static_cast(num_of_in_flight_s2g_dispatch_api))); + list.push_back(c10::IValue(static_cast(num_of_in_flight_s2g_permute_block_dispatch_api))); + list.push_back(c10::IValue(static_cast(num_of_additional_in_flight_s2g_dispatch_api))); + list.push_back(c10::IValue(static_cast(num_of_tokens_per_chunk_dispatch_api))); + list.push_back(c10::IValue(static_cast(num_of_blocks_dispatch_api))); + list.push_back(c10::IValue(forward_dispatch_api)); + list.push_back(c10::IValue(device_side_sync_dispatch_api)); // Combine API Config - elements.push_back(c10::IValue(static_cast(num_of_stages_g2s_combine_api))); - elements.push_back(c10::IValue(static_cast(num_of_stages_s2g_combine_api))); - elements.push_back(c10::IValue(static_cast(num_of_stages_g2s_unpermute_block))); - elements.push_back(c10::IValue(static_cast(num_of_stages_s2g_unpermute_block))); - elements.push_back(c10::IValue(static_cast(num_of_tokens_per_chunk_combine_api))); - elements.push_back(c10::IValue(static_cast(num_of_tokens_per_group_combine_api))); - elements.push_back(c10::IValue(static_cast(num_of_blocks_combine_api))); - elements.push_back(c10::IValue(static_cast(num_of_additional_in_flight_s2g_combine_api))); - elements.push_back(c10::IValue(static_cast(num_of_additional_in_flight_s2g_unpermute_block_combine_api))); - elements.push_back(c10::IValue(backward_combine_api)); - elements.push_back(c10::IValue(device_side_sync_combine_api)); - return c10::IValue(c10::ivalue::Tuple::create(std::move(elements))); + list.push_back(c10::IValue(static_cast(num_of_stages_g2s_combine_api))); + list.push_back(c10::IValue(static_cast(num_of_stages_s2g_combine_api))); + list.push_back(c10::IValue(static_cast(num_of_stages_g2s_unpermute_block))); + list.push_back(c10::IValue(static_cast(num_of_stages_s2g_unpermute_block))); + list.push_back(c10::IValue(static_cast(num_of_tokens_per_chunk_combine_api))); + list.push_back(c10::IValue(static_cast(num_of_tokens_per_group_combine_api))); + list.push_back(c10::IValue(static_cast(num_of_blocks_combine_api))); + list.push_back(c10::IValue(static_cast(num_of_additional_in_flight_s2g_combine_api))); + list.push_back(c10::IValue(static_cast(num_of_additional_in_flight_s2g_unpermute_block_combine_api))); + list.push_back(c10::IValue(backward_combine_api)); + list.push_back(c10::IValue(device_side_sync_combine_api)); + return c10::IValue(std::move(list)); } bool operator<(const HybridEpConfigInstance& other) const { return std::memcmp(this, &other, sizeof(HybridEpConfigInstance)) < 0; diff --git a/csrc/hybrid_ep/executor/executor.cuh b/csrc/hybrid_ep/executor/executor.cuh index 8087160b0..d4e23d9d3 100644 --- a/csrc/hybrid_ep/executor/executor.cuh +++ b/csrc/hybrid_ep/executor/executor.cuh @@ -34,24 +34,24 @@ struct HandleImpl { torch::Tensor dense_chunk_layout; torch::Tensor dense_to_expert_map; - /** Convert all attributes to a single IValue holding a tuple of IValues (tensors/ints/nested config tuple). */ + /** Convert all attributes to a single IValue holding a GenericList of IValues (tensors/ints/nested config list). */ c10::IValue to_ivalue_tuple() const { - std::vector elements; - elements.reserve(13); - elements.push_back(c10::IValue(sparse_to_dense_map)); - elements.push_back(c10::IValue(rdma_to_attn_map)); - elements.push_back(c10::IValue(attn_to_rdma_map)); - elements.push_back(c10::IValue(num_dispatched_tokens_tensor)); - elements.push_back(c10::IValue(local_expert_routing_map)); - elements.push_back(c10::IValue(static_cast(num_of_tokens_per_rank))); - elements.push_back(config.to_ivalue_tuple()); - elements.push_back(c10::IValue(tokens_per_expert)); - elements.push_back(c10::IValue(padded_tokens_per_expert)); - elements.push_back(c10::IValue(overflow_flag)); - elements.push_back(c10::IValue(static_cast(num_permuted_tokens))); - elements.push_back(c10::IValue(dense_chunk_layout)); - elements.push_back(c10::IValue(dense_to_expert_map)); - return c10::IValue(c10::ivalue::Tuple::create(std::move(elements))); + c10::impl::GenericList list(c10::AnyType::get()); + list.reserve(13); + list.push_back(c10::IValue(sparse_to_dense_map)); + list.push_back(c10::IValue(rdma_to_attn_map)); + list.push_back(c10::IValue(attn_to_rdma_map)); + list.push_back(c10::IValue(num_dispatched_tokens_tensor)); + list.push_back(c10::IValue(local_expert_routing_map)); + list.push_back(c10::IValue(static_cast(num_of_tokens_per_rank))); + list.push_back(config.to_ivalue_tuple()); + list.push_back(c10::IValue(tokens_per_expert)); + list.push_back(c10::IValue(padded_tokens_per_expert)); + list.push_back(c10::IValue(overflow_flag)); + list.push_back(c10::IValue(static_cast(num_permuted_tokens))); + list.push_back(c10::IValue(dense_chunk_layout)); + list.push_back(c10::IValue(dense_to_expert_map)); + return c10::IValue(std::move(list)); } }; diff --git a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh index a49eff3d5..74f2b06ef 100644 --- a/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh +++ b/csrc/hybrid_ep/pybind_hybrid_ep_utils.cuh @@ -120,44 +120,6 @@ inline HybridEPBuffer* hybrid_ep_buffer_init( base_path, load_cached_kernels, use_shared_buffer, enable_custom_allgather); } -template -auto hybrid_ep_buffer_combine(Func func) { - return [func](HybridEPBuffer& self, - torch::Tensor hidden, - c10::optional probs, - HandleImpl handle, - bool with_probs) { - auto result = (self.*func)(hidden, probs, handle, with_probs); - if (isProfilerEnabledInMainThread()) { - record_hybrid_ep_buffer_init(); - - std::vector inputValues; - inputValues.reserve(4); - inputValues.push_back(to_ivalue(hidden)); - inputValues.push_back(to_ivalue(probs)); - inputValues.push_back(to_ivalue(handle)); - inputValues.push_back(to_ivalue(with_probs)); - - std::vector outputValues; - record_result(outputValues, result); - - std::vector args; - args.emplace_back("hidden", c10::TensorType::get()); - args.emplace_back("probs", c10::OptionalType::create(c10::TensorType::get())); - args.emplace_back("handle", c10::AnyType::get()); - args.emplace_back("with_probs", c10::BoolType::get()); - std::vector returns; - returns.emplace_back("", c10::TensorType::get()); - returns.emplace_back("", c10::TensorType::get()); - c10::FunctionSchema schema( - "HybridEPBuffer::combine", "", - std::move(args), std::move(returns), false, false); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, &inputValues, outputValues); - } - return result; - }; -} - // Wrapper for update_buffer (returns bool) template auto hybrid_ep_buffer_update_buffer(Func func) { @@ -203,7 +165,6 @@ auto hybrid_ep_buffer_metadata_preprocessing(Func func) { fuse_permute_dispatch, non_blocking); if (isProfilerEnabledInMainThread()) { record_hybrid_ep_buffer_init(); - std::vector inputValues; inputValues.reserve(8); inputValues.push_back(to_ivalue(config)); @@ -228,24 +189,12 @@ auto hybrid_ep_buffer_metadata_preprocessing(Func func) { args.emplace_back("fuse_permute_dispatch", c10::BoolType::get()); args.emplace_back("non_blocking", c10::BoolType::get()); std::vector returns; - returns.reserve(13); - returns.emplace_back("", c10::TensorType::get()); // sparse_to_dense_map - returns.emplace_back("", c10::TensorType::get()); // rdma_to_attn_map - returns.emplace_back("", c10::TensorType::get()); // attn_to_rdma_map - returns.emplace_back("", c10::TensorType::get()); // num_dispatched_tokens_tensor - returns.emplace_back("", c10::TensorType::get()); // local_expert_routing_map - returns.emplace_back("", c10::IntType::get()); // num_of_tokens_per_rank - returns.emplace_back("", c10::AnyType::get()); // config - returns.emplace_back("", c10::TensorType::get()); // tokens_per_expert - returns.emplace_back("", c10::TensorType::get()); // padded_tokens_per_expert - returns.emplace_back("", c10::TensorType::get()); // overflow_flag - returns.emplace_back("", c10::IntType::get()); // num_permuted_tokens - returns.emplace_back("", c10::TensorType::get()); // dense_chunk_layout - returns.emplace_back("", c10::TensorType::get()); // dense_to_expert_map + returns.emplace_back("HandleImpl", c10::AnyType::get()); c10::FunctionSchema schema( "HybridEPBuffer::metadata_preprocessing", "", std::move(args), std::move(returns), false, false); - RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, &inputValues, outputValues); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, &inputValues, outputValues); + // RECORD_FUNCTION_WITH_INPUTS_OUTPUTS("HybridEPBuffer::metadata_preprocessing", &inputValues, outputValues); } return result; }; @@ -347,6 +296,44 @@ auto hybrid_ep_buffer_dispatch_with_permute(Func func) { }; } +template +auto hybrid_ep_buffer_combine(Func func) { + return [func](HybridEPBuffer& self, + torch::Tensor hidden, + c10::optional probs, + HandleImpl handle, + bool with_probs) { + auto result = (self.*func)(hidden, probs, handle, with_probs); + if (isProfilerEnabledInMainThread()) { + record_hybrid_ep_buffer_init(); + + std::vector inputValues; + inputValues.reserve(4); + inputValues.push_back(to_ivalue(hidden)); + inputValues.push_back(to_ivalue(probs)); + inputValues.push_back(to_ivalue(handle)); + inputValues.push_back(to_ivalue(with_probs)); + + std::vector outputValues; + record_result(outputValues, result); + + std::vector args; + args.emplace_back("hidden", c10::TensorType::get()); + args.emplace_back("probs", c10::OptionalType::create(c10::TensorType::get())); + args.emplace_back("handle", c10::AnyType::get()); + args.emplace_back("with_probs", c10::BoolType::get()); + std::vector returns; + returns.emplace_back("", c10::TensorType::get()); + returns.emplace_back("", c10::TensorType::get()); + c10::FunctionSchema schema( + "HybridEPBuffer::combine", "", + std::move(args), std::move(returns), false, false); + RECORD_FUNCTION_WITH_INPUTS_OUTPUTS(schema, &inputValues, outputValues); + } + return result; + }; +} + // Wrapper for combine_with_unpermute (returns 2-tuple) template auto hybrid_ep_buffer_combine_with_unpermute(Func func) { diff --git a/tests/test_hybrid_ep_collect_et.py b/tests/test_hybrid_ep_collect_et.py index cd4a4e99c..d30c505e2 100644 --- a/tests/test_hybrid_ep_collect_et.py +++ b/tests/test_hybrid_ep_collect_et.py @@ -121,22 +121,27 @@ def test_main(local_rank: int, num_local_ranks: int, args: argparse.Namespace): for i in range(1): + hidden = torch.relu(hidden) # Non-permute - 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)) - # dispatched_hidden_bf16 = dispatched_hidden.to(torch.bfloat16) - # combine_args = {'hidden': dispatched_hidden_bf16, 'probs': dispatched_probs, 'handle': handle} - # buffer.combine(**combine_args) + 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, + ) + dispatched_hidden_bf16 = dispatched_hidden.to(torch.bfloat16) + combine_args = {'hidden': dispatched_hidden_bf16, 'probs': dispatched_probs, 'handle': handle} + buffer.combine(**combine_args) # # Permute (non-fused) - # dispatched_hidden_wp, dispatched_probs_wp, _, tpe_wp, handle_wp = ( - # buffer.dispatch_with_permute(hidden=hidden, scaling_factor=scaling_factor, - # routing_map=routing_map, probs=probs, pad_multiple=PAD_MULTIPLE)) - # dispatched_hidden_bf16_wp = dispatched_hidden_wp.to(torch.bfloat16) - # combine_wp_args = {'hidden': dispatched_hidden_bf16_wp, 'probs': dispatched_probs_wp, - # 'handle': handle_wp, 'pad_multiple': PAD_MULTIPLE} - # buffer.combine_with_unpermute(**combine_wp_args) + dispatched_hidden_wp, dispatched_probs_wp, _, tpe_wp, handle_wp = ( + buffer.dispatch_with_permute(hidden=hidden, scaling_factor=scaling_factor, + routing_map=routing_map, probs=probs, pad_multiple=PAD_MULTIPLE)) + dispatched_hidden_bf16_wp = dispatched_hidden_wp.to(torch.bfloat16) + combine_wp_args = {'hidden': dispatched_hidden_bf16_wp, 'probs': dispatched_probs_wp, + 'handle': handle_wp, 'pad_multiple': PAD_MULTIPLE} + buffer.combine_with_unpermute(**combine_wp_args) # # Fused permute-dispatch # dispatched_hidden_fused, dispatched_probs_fused, _, tpe_fused, handle_fused = (