Skip to content

Capture PyTorch profiling trace in HybridEPBuffer #583

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
shengf-nv wants to merge 23 commits into
base: hybrid-ep
Choose a base branch
from
Draft

Capture PyTorch profiling trace in HybridEPBuffer #583

Show file tree
Hide file tree
Changes from 22 commits
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
c2746f5
first implementation for collecting ET
shengf-nv Feb 18, 2026
9c13226
cleanup code
shengf-nv Feb 18, 2026
b0dc913
use c10::impl::GenericList
shengf-nv Feb 18, 2026
bb5bc8f
support all APIs
shengf-nv Feb 18, 2026
f8480e7
put new APIs into a new file
shengf-nv Feb 18, 2026
43db483
put new APIs into a new file
shengf-nv Feb 18, 2026
8fc1510
clean up head files
shengf-nv Feb 18, 2026
a09d4db
get pg id
shengf-nv Feb 19, 2026
ba9a0ab
get pg id
shengf-nv Feb 19, 2026
cc98dd8
add schema
shengf-nv Feb 22, 2026
d2456cb
remove schema
shengf-nv Feb 23, 2026
921f2a1
change innput/output from a single arrary to a list of inputs
shengf-nv Feb 24, 2026
68f1e92
fixed schema
shengf-nv Feb 24, 2026
830e126
add permute tests
shengf-nv Feb 27, 2026
3e706d0
add two more APIs
shengf-nv Feb 28, 2026
67a1d21
support four APIs
shengf-nv Mar 2, 2026
4853562
record HybridEPBuffer init when profiling is on
shengf-nv Mar 5, 2026
2fc7aa9
collect et with kineto trace
shengf-nv Mar 7, 2026
c450e13
check if profiler is enabled
Mar 17, 2026
72ca05f
updated with the latest version of hybrid-ep
May 30, 2026
e1d76c6
expose attributes
shengf-nv May 30, 2026
9616d6d
fix test
shengf-nv Jun 1, 2026
9b2a80a
dispatch combine worked
shengf-nv Jun 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 63 additions & 0 deletions csrc/hybrid_ep/config.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#include <stdexcept>
#include <optional>
#include "utils.cuh"
#include <ATen/core/ivalue.h>

// Now we support up to 72(GB200) ranks per node.
// This will be used to initialize the template param_t for communication kernel.
Expand Down Expand Up @@ -100,6 +101,26 @@ 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<c10::IValue> elements;
elements.reserve(13);
elements.push_back(c10::IValue(static_cast<int64_t>(hidden_dim)));
elements.push_back(c10::IValue(static_cast<int64_t>(max_num_of_tokens_per_rank)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_experts_per_rank)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_ranks_per_node)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_nodes)));
elements.push_back(c10::IValue(type_to_string(token_data_type)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_blocks_preprocessing_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_blocks_dispatch_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_blocks_combine_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_tokens_per_chunk_dispatch_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_tokens_per_chunk_combine_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_dispatch_chunks)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_combine_chunks)));
return c10::IValue(c10::ivalue::Tuple::create(std::move(elements)));
}
};

// Config used for hybrid-ep kernel.
Expand Down Expand Up @@ -189,6 +210,48 @@ struct HybridEpConfigInstance {
return valid;
}

/** Convert all attributes to a single IValue holding a tuple of IValues (ints/bools/string). */
c10::IValue to_ivalue_tuple() const {
std::vector<c10::IValue> elements;
elements.reserve(32);
// Hybrid-ep Config
elements.push_back(c10::IValue(static_cast<int64_t>(hidden_dim)));
elements.push_back(c10::IValue(static_cast<int64_t>(max_num_of_tokens_per_rank)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_experts_per_rank)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_ranks_per_node)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_nodes)));
elements.push_back(c10::IValue(static_cast<int64_t>(pad_multiple)));
// Metadata-preprocessing API Config
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_tokens_per_chunk_preprocessing_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_threads_per_block_preprocessing_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_blocks_preprocessing_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_blocks_permute)));
elements.push_back(c10::IValue(static_cast<int64_t>(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<int64_t>(num_of_stages_dispatch_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_stages_permute_block_dispatch_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_in_flight_s2g_dispatch_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_in_flight_s2g_permute_block_dispatch_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_additional_in_flight_s2g_dispatch_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_tokens_per_chunk_dispatch_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(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<int64_t>(num_of_stages_g2s_combine_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_stages_s2g_combine_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_stages_g2s_unpermute_block)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_stages_s2g_unpermute_block)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_tokens_per_chunk_combine_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_tokens_per_group_combine_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_blocks_combine_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(num_of_additional_in_flight_s2g_combine_api)));
elements.push_back(c10::IValue(static_cast<int64_t>(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)));
}
bool operator<(const HybridEpConfigInstance& other) const {
return std::memcmp(this, &other, sizeof(HybridEpConfigInstance)) < 0;
}
Expand Down
20 changes: 20 additions & 0 deletions csrc/hybrid_ep/executor/executor.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,26 @@ struct HandleImpl {
int64_t num_permuted_tokens = -1;
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). */
c10::IValue to_ivalue_tuple() const {
std::vector<c10::IValue> 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<int64_t>(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<int64_t>(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)));
}
};

class Executor {
Expand Down
33 changes: 26 additions & 7 deletions csrc/hybrid_ep/pybind_hybrid_ep.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
#include "hybrid_ep.cuh"
#include "utils.cuh"
#include "config.cuh"
#include "pybind_hybrid_ep_utils.cuh"

namespace py = pybind11;

Expand Down Expand Up @@ -109,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",
Expand All @@ -118,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",
Expand Down Expand Up @@ -170,7 +178,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
.def_readwrite("dense_to_expert_map", &HandleImpl::dense_to_expert_map);

pybind11::class_<HybridEPBuffer>(m, "HybridEPBuffer")
.def(py::init<py::object, BufferConfig, int, int, int, std::string, bool, bool, bool>(),
.def(py::init(&hybrid_ep_buffer_init),
py::arg("process_group"),
py::arg("config"),
py::arg("local_rank"),
Expand All @@ -180,8 +188,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"),
Expand All @@ -191,18 +202,24 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
py::arg("enable_permute") = false,
py::arg("fuse_permute_dispatch") = false,
py::arg("non_blocking") = false)
.def("dispatch", &HybridEPBuffer::dispatch, py::kw_only(),
.def("dispatch",
hybrid_ep_buffer_dispatch(&HybridEPBuffer::dispatch),
py::kw_only(),
py::arg("hidden"),
py::arg("probs") = c10::nullopt,
py::arg("scaling_factor") = c10::nullopt,
py::arg("handle"),
py::arg("with_probs"))
.def("combine", &HybridEPBuffer::combine, py::kw_only(),
.def("combine",
hybrid_ep_buffer_combine(&HybridEPBuffer::combine),
py::kw_only(),
py::arg("hidden"),
py::arg("probs") = c10::nullopt,
py::arg("handle"),
py::arg("with_probs"))
.def("dispatch_with_permute", &HybridEPBuffer::dispatch_with_permute, py::kw_only(),
.def("dispatch_with_permute",
hybrid_ep_buffer_dispatch_with_permute(&HybridEPBuffer::dispatch_with_permute),
py::kw_only(),
py::arg("hidden"),
py::arg("probs") = c10::nullopt,
py::arg("scaling_factor") = c10::nullopt,
Expand All @@ -211,7 +228,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
py::arg("fuse_permute_dispatch") = false,
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",
hybrid_ep_buffer_combine_with_unpermute(&HybridEPBuffer::combine_with_unpermute),
py::kw_only(),
py::arg("hidden"),
py::arg("probs") = c10::nullopt,
py::arg("handle"),
Expand Down
Loading