Skip to content

feat(normal): add optional diagnostic probes for DeepEP normal mode #720

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

Open
Marsssqqq wants to merge 16 commits into
base: antgroup-opt
Choose a base branch
from
Open

feat(normal): add optional diagnostic probes for DeepEP normal mode #720

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
699d17b
fix(build): locate CUDA 13 CCCL headers
Jul 28, 2026
5867e09
feat(deepxtrace): validate and bind normal stats contract
Aug 4, 2026
4365d77
feat(normal): add notify full-kernel duration probes
Jul 27, 2026
f46f3ff
feat(normal): add dispatch final completion probes
Jul 27, 2026
b10f01d
feat(normal): add dispatch RDMA receive completion probes
Aug 4, 2026
7ab8aa8
feat(normal): add combine logical receive completion probes
Jul 27, 2026
9572707
feat(deepxtrace): make normal diagnosis optional and configurable
Jul 28, 2026
fd9182a
fix(buffer): preserve backward-compatible optional defaults
Aug 10, 2026
5caf04e
fix(normal): harden optional DeepXTrace integration
Aug 10, 2026
77977ed
fix(normal): isolate diagnostics from low-latency APIs
Aug 10, 2026
fc64467
fix(normal): harden probe contracts and collection
Aug 10, 2026
7d672b0
test(normal): cover optional DeepXTrace integration
Aug 10, 2026
21a2499
style(normal): apply repository formatters
Aug 10, 2026
d66d866
refactor(deepxtrace): decouple diagnosis lifecycle from DeepEP
Aug 11, 2026
7064154
test(normal): add optional DeepXTrace normal-mode test
Aug 12, 2026
8a2d98f
refactor(normal): group diagnostic probe arguments
Aug 26, 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
177 changes: 170 additions & 7 deletions csrc/deep_ep.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,49 @@ void SharedMemoryAllocator::close_mem_handle(void* ptr) {

namespace deep_ep {

namespace {

internode::NormalNotifyStats prepare_normal_notify_stats(const NormalNotifyStats& stats, int64_t* timer_state) {
const bool enabled = stats.duration_ns.has_value();
EP_HOST_ASSERT(stats.count.has_value() == enabled);
if (not enabled)
return {};

EP_HOST_ASSERT(stats.duration_ns->is_cuda() and stats.count->is_cuda());
EP_HOST_ASSERT(stats.duration_ns->scalar_type() == torch::kInt64 and stats.count->scalar_type() == torch::kInt64);
EP_HOST_ASSERT(stats.duration_ns->dim() == 1 and stats.duration_ns->numel() == 1 and stats.duration_ns->is_contiguous());
EP_HOST_ASSERT(stats.count->dim() == 1 and stats.count->numel() == 1 and stats.count->is_contiguous());

return {
stats.duration_ns->data_ptr<int64_t>(),
stats.count->data_ptr<int64_t>(),
timer_state,
};
}

internode::NormalCompletionStats prepare_normal_completion_stats(const NormalCompletionStats& stats, int num_ranks) {
const bool enabled = stats.cost.has_value();
EP_HOST_ASSERT(stats.sample_count.has_value() == enabled);
EP_HOST_ASSERT(stats.token_count.has_value() == enabled);
if (not enabled)
return {};

for (const auto& tensor : {stats.cost, stats.sample_count, stats.token_count}) {
EP_HOST_ASSERT(tensor->is_cuda());
EP_HOST_ASSERT(tensor->scalar_type() == torch::kInt64);
EP_HOST_ASSERT(tensor->dim() == 1 and tensor->is_contiguous());
EP_HOST_ASSERT(tensor->numel() == num_ranks);
}

return {
stats.cost->data_ptr<int64_t>(),
stats.sample_count->data_ptr<int64_t>(),
stats.token_count->data_ptr<int64_t>(),
};
}

} // namespace

Buffer::Buffer(int rank,
int num_ranks,
int64_t num_nvl_bytes,
Expand Down Expand Up @@ -198,6 +241,7 @@ Buffer::Buffer(int rank,
// Create 32 MiB workspace
CUDA_CHECK(cudaMalloc(&workspace, NUM_WORKSPACE_BYTES));
CUDA_CHECK(cudaMemsetAsync(workspace, 0, NUM_WORKSPACE_BYTES, comm_stream));
CUDA_CHECK(cudaMalloc(&normal_notify_full_kernel_timer_states, 3 * 2 * sizeof(int64_t)));

// MoE counter
CUDA_CHECK(cudaMallocHost(&moe_recv_counter, sizeof(int64_t), cudaHostAllocMapped));
Expand Down Expand Up @@ -316,6 +360,7 @@ void Buffer::destroy() {

// Free workspace and MoE counter
CUDA_CHECK(cudaFree(workspace));
CUDA_CHECK(cudaFree(normal_notify_full_kernel_timer_states));
CUDA_CHECK(cudaFreeHost(const_cast<int*>(moe_recv_counter)));

// Free chunked mode staffs
Expand Down Expand Up @@ -946,7 +991,8 @@ Buffer::internode_dispatch(const torch::Tensor& x,
const Config& config,
std::optional<EventHandle>& previous_event,
bool async,
bool allocate_on_comm_stream) {
bool allocate_on_comm_stream,
const NormalDispatchStats& normal_stats) {
#ifndef DISABLE_NVSHMEM
// In dispatch, CPU will busy-wait until GPU receive tensor size metadata from other ranks, which can be quite long.
// If users of DeepEP need to execute other Python code on other threads, such as KV transfer, their code will get stuck due to GIL
Expand Down Expand Up @@ -1004,6 +1050,10 @@ Buffer::internode_dispatch(const torch::Tensor& x,
EP_HOST_ASSERT(num_tokens_per_expert->size(0) % num_ranks == 0);
EP_HOST_ASSERT(num_tokens_per_expert->size(0) / num_ranks <= NUM_MAX_LOCAL_EXPERTS);
}
const auto notify_stats = prepare_normal_notify_stats(normal_stats.notify, normal_notify_full_kernel_timer_states);
const auto cached_notify_stats = prepare_normal_notify_stats(normal_stats.cached_notify, normal_notify_full_kernel_timer_states + 2);
const auto final_completion_stats = prepare_normal_completion_stats(normal_stats.final_completion, num_ranks);
const auto rdma_recv_completion_stats = prepare_normal_completion_stats(normal_stats.rdma_recv_completion, num_ranks);

auto num_tokens = static_cast<int>(x.size(0)), hidden = static_cast<int>(x.size(1)),
hidden_int4 = static_cast<int>(x.size(1) * x.element_size() / sizeof(int4));
Expand Down Expand Up @@ -1094,7 +1144,8 @@ Buffer::internode_dispatch(const torch::Tensor& x,
config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks),
num_nvl_bytes,
true,
low_latency_mode);
low_latency_mode,
cached_notify_stats);
} else {
rdma_channel_prefix_matrix = torch::empty({num_rdma_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA));
recv_rdma_rank_prefix_sum = torch::empty({num_rdma_ranks}, dtype(torch::kInt32).device(torch::kCUDA));
Expand Down Expand Up @@ -1134,7 +1185,8 @@ Buffer::internode_dispatch(const torch::Tensor& x,
comm_stream,
config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks),
num_nvl_bytes,
low_latency_mode);
low_latency_mode,
notify_stats);

// Synchronize total received tokens and tokens per expert
if (num_worst_tokens > 0) {
Expand Down Expand Up @@ -1235,6 +1287,8 @@ Buffer::internode_dispatch(const torch::Tensor& x,
buffer_ptrs_gpu,
config.num_max_nvl_chunked_send_tokens,
config.num_max_nvl_chunked_recv_tokens,
final_completion_stats,
rdma_recv_completion_stats,
rank,
num_ranks,
cached_mode,
Expand Down Expand Up @@ -1324,7 +1378,8 @@ std::tuple<torch::Tensor, std::optional<torch::Tensor>, std::optional<EventHandl
const Config& config,
std::optional<EventHandle>& previous_event,
bool async,
bool allocate_on_comm_stream) {
bool allocate_on_comm_stream,
const NormalCombineStats& normal_stats) {
#ifndef DISABLE_NVSHMEM
const int num_channels = config.num_sms / 2;
EP_HOST_ASSERT(config.num_sms % 2 == 0);
Expand Down Expand Up @@ -1356,6 +1411,8 @@ std::tuple<torch::Tensor, std::optional<torch::Tensor>, std::optional<EventHandl
EP_HOST_ASSERT(combined_rdma_head.dim() == 2 and combined_rdma_head.size(0) == num_combined_tokens and
combined_rdma_head.size(1) == num_rdma_ranks);
EP_HOST_ASSERT(combined_nvl_head.dim() == 2 and combined_nvl_head.size(1) == NUM_MAX_NVL_PEERS);
const auto cached_notify_stats = prepare_normal_notify_stats(normal_stats.cached_notify, normal_notify_full_kernel_timer_states + 4);
const auto logical_recv_completion_stats = prepare_normal_completion_stats(normal_stats.logical_recv_completion, num_ranks);

// Allocate all tensors on comm stream if set
// NOTES: do not allocate tensors upfront!
Expand Down Expand Up @@ -1413,7 +1470,8 @@ std::tuple<torch::Tensor, std::optional<torch::Tensor>, std::optional<EventHandl
config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks),
num_nvl_bytes,
false,
low_latency_mode);
low_latency_mode,
cached_notify_stats);

// Assign bias pointers
auto bias_opts = std::vector<std::optional<torch::Tensor>>({bias_0, bias_1});
Expand Down Expand Up @@ -1453,6 +1511,7 @@ std::tuple<torch::Tensor, std::optional<torch::Tensor>, std::optional<EventHandl
buffer_ptrs_gpu,
config.num_max_nvl_chunked_send_tokens,
config.num_max_nvl_chunked_recv_tokens,
logical_recv_completion_stats,
rank,
num_ranks,
comm_stream,
Comment thread
Marsssqqq marked this conversation as resolved.
Expand Down Expand Up @@ -1866,6 +1925,71 @@ void Buffer::low_latency_clean_mask_buffer() {

} // namespace deep_ep

namespace {

std::optional<torch::Tensor> cast_optional_tensor(pybind11::handle value) {
return value.is_none() ? std::nullopt : std::make_optional(pybind11::cast<torch::Tensor>(value));
}

pybind11::sequence require_normal_stats_sequence(pybind11::handle source, size_t expected_size, const char* argument_name) {
if (not pybind11::isinstance<pybind11::tuple>(source) and not pybind11::isinstance<pybind11::list>(source))
throw pybind11::value_error(std::string("`") + argument_name + "` must be a tuple or list");

auto sequence = pybind11::reinterpret_borrow<pybind11::sequence>(source);
if (static_cast<size_t>(sequence.size()) != expected_size)
throw pybind11::value_error(std::string("`") + argument_name + "` must contain " + std::to_string(expected_size) +
" tensors, but got " + std::to_string(sequence.size()));
return sequence;
}

} // namespace

namespace pybind11::detail {

template <>
struct type_caster<deep_ep::NormalDispatchStats> {
public:
PYBIND11_TYPE_CASTER(deep_ep::NormalDispatchStats, const_name("NormalDispatchStats"));

bool load(handle source, bool) {
if (source.is_none()) {
value = {};
return true;
}

auto stats = require_normal_stats_sequence(source, 10, "normal_dispatch_stats");
value = {
{cast_optional_tensor(stats[0]), cast_optional_tensor(stats[1])},
{cast_optional_tensor(stats[2]), cast_optional_tensor(stats[3])},
{cast_optional_tensor(stats[4]), cast_optional_tensor(stats[5]), cast_optional_tensor(stats[6])},
{cast_optional_tensor(stats[7]), cast_optional_tensor(stats[8]), cast_optional_tensor(stats[9])},
};
return true;
}
};

template <>
struct type_caster<deep_ep::NormalCombineStats> {
public:
PYBIND11_TYPE_CASTER(deep_ep::NormalCombineStats, const_name("NormalCombineStats"));

bool load(handle source, bool) {
if (source.is_none()) {
value = {};
return true;
}

auto stats = require_normal_stats_sequence(source, 5, "normal_combine_stats");
value = {
{cast_optional_tensor(stats[0]), cast_optional_tensor(stats[1])},
{cast_optional_tensor(stats[2]), cast_optional_tensor(stats[3]), cast_optional_tensor(stats[4])},
};
return true;
}
};

} // namespace pybind11::detail

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.doc() = "DeepEP: an efficient expert-parallel communication library";

Expand Down Expand Up @@ -1900,8 +2024,47 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
.def("get_dispatch_layout", &deep_ep::Buffer::get_dispatch_layout)
.def("intranode_dispatch", &deep_ep::Buffer::intranode_dispatch)
.def("intranode_combine", &deep_ep::Buffer::intranode_combine)
.def("internode_dispatch", &deep_ep::Buffer::internode_dispatch)
.def("internode_combine", &deep_ep::Buffer::internode_combine)
.def("internode_dispatch",
&deep_ep::Buffer::internode_dispatch,
py::arg("x"),
py::arg("x_scales"),
py::arg("topk_idx"),
py::arg("topk_weights"),
py::arg("num_tokens_per_rank"),
py::arg("num_tokens_per_rdma_rank"),
py::arg("is_token_in_rank"),
py::arg("num_tokens_per_expert"),
py::arg("cached_num_recv_tokens"),
py::arg("cached_num_rdma_recv_tokens"),
py::arg("cached_rdma_channel_prefix_matrix"),
py::arg("cached_recv_rdma_rank_prefix_sum"),
py::arg("cached_gbl_channel_prefix_matrix"),
py::arg("cached_recv_gbl_rank_prefix_sum"),
py::arg("expert_alignment"),
py::arg("num_worst_tokens"),
py::arg("config"),
py::arg("previous_event"),
py::arg("async"),
py::arg("allocate_on_comm_stream"),
py::arg("normal_dispatch_stats") = py::none())
.def("internode_combine",
&deep_ep::Buffer::internode_combine,
py::arg("x"),
py::arg("topk_weights"),
py::arg("bias_0"),
py::arg("bias_1"),
py::arg("src_meta"),
py::arg("is_combined_token_in_rank"),
py::arg("rdma_channel_prefix_matrix"),
py::arg("rdma_rank_prefix_sum"),
py::arg("gbl_channel_prefix_matrix"),
py::arg("combined_rdma_head"),
py::arg("combined_nvl_head"),
py::arg("config"),
py::arg("previous_event"),
py::arg("async"),
py::arg("allocate_on_comm_stream"),
py::arg("normal_combine_stats") = py::none())
.def("clean_low_latency_buffer", &deep_ep::Buffer::clean_low_latency_buffer)
.def("low_latency_dispatch", &deep_ep::Buffer::low_latency_dispatch)
.def("low_latency_combine", &deep_ep::Buffer::low_latency_combine)
Expand Down
32 changes: 30 additions & 2 deletions csrc/deep_ep.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,29 @@ class SharedMemoryAllocator {

namespace deep_ep {

struct NormalNotifyStats {
std::optional<torch::Tensor> duration_ns;
std::optional<torch::Tensor> count;
};

struct NormalCompletionStats {
std::optional<torch::Tensor> cost;
std::optional<torch::Tensor> sample_count;
std::optional<torch::Tensor> token_count;
};

struct NormalDispatchStats {
NormalNotifyStats notify;
NormalNotifyStats cached_notify;
NormalCompletionStats final_completion;
NormalCompletionStats rdma_recv_completion;
};

struct NormalCombineStats {
NormalNotifyStats cached_notify;
NormalCompletionStats logical_recv_completion;
};

struct Buffer {
EP_STATIC_ASSERT(NUM_MAX_NVL_PEERS == 8, "The number of maximum NVLink peers must be 8");

Expand Down Expand Up @@ -99,6 +122,9 @@ struct Buffer {
// Workspace
void* workspace = nullptr;

// Per-launch scratch for the three normal-mode Notify timers
int64_t* normal_notify_full_kernel_timer_states = nullptr;

// Host-side MoE info
volatile int* moe_recv_counter = nullptr;
int* moe_recv_counter_mapped = nullptr;
Expand Down Expand Up @@ -235,7 +261,8 @@ struct Buffer {
const Config& config,
std::optional<EventHandle>& previous_event,
bool async,
bool allocate_on_comm_stream);
bool allocate_on_comm_stream,
const NormalDispatchStats& normal_stats = {});

std::tuple<torch::Tensor, std::optional<torch::Tensor>, std::optional<EventHandle>> internode_combine(
const torch::Tensor& x,
Expand All @@ -252,7 +279,8 @@ struct Buffer {
const Config& config,
std::optional<EventHandle>& previous_event,
bool async,
bool allocate_on_comm_stream);
bool allocate_on_comm_stream,
const NormalCombineStats& normal_stats = {});

void clean_low_latency_buffer(int num_max_dispatch_tokens_per_rank, int hidden, int num_experts);

Expand Down
25 changes: 23 additions & 2 deletions csrc/kernels/api.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,18 @@ namespace internode {

int get_source_meta_bytes();

struct NormalNotifyStats {
int64_t* duration_ns = nullptr;
int64_t* count = nullptr;
int64_t* timer_state = nullptr;
};

struct NormalCompletionStats {
int64_t* cost = nullptr;
int64_t* sample_count = nullptr;
int64_t* token_count = nullptr;
};

void notify_dispatch(const int* num_tokens_per_rank,
int* moe_recv_counter_mapped,
int num_ranks,
Expand Down Expand Up @@ -173,7 +185,8 @@ void notify_dispatch(const int* num_tokens_per_rank,
cudaStream_t stream,
int64_t num_rdma_bytes,
int64_t num_nvl_bytes,
bool low_latency_mode);
bool low_latency_mode,
NormalNotifyStats normal_notify_stats);

void dispatch(void* recv_x,
float* recv_x_scales,
Expand Down Expand Up @@ -207,6 +220,10 @@ void dispatch(void* recv_x,
void** buffer_ptrs,
int num_max_nvl_chunked_send_tokens,
int num_max_nvl_chunked_recv_tokens,
// Completion cost tensors accumulate clock64() SM cycles;
// Notify duration tensors accumulate %globaltimer nanoseconds.
NormalCompletionStats normal_final_completion_stats,
NormalCompletionStats normal_rdma_recv_completion_stats,
int rank,
int num_ranks,
bool is_cached_dispatch,
Expand Down Expand Up @@ -235,7 +252,8 @@ void cached_notify(int hidden_int4,
int64_t num_rdma_bytes,
int64_t num_nvl_bytes,
bool is_cached_dispatch,
bool low_latency_mode);
bool low_latency_mode,
NormalNotifyStats normal_notify_stats);

void combine(cudaDataType_t type,
void* combined_x,
Expand All @@ -261,6 +279,9 @@ void combine(cudaDataType_t type,
void** buffer_ptrs,
int num_max_nvl_chunked_send_tokens,
int num_max_nvl_chunked_recv_tokens,
// Completion cost tensors accumulate clock64() SM cycles;
// Notify duration tensors accumulate %globaltimer nanoseconds.
NormalCompletionStats normal_logical_recv_completion_stats,
int rank,
int num_ranks,
cudaStream_t stream,
Expand Down
Loading