microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-07-15 19:54:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge feature/ep benchmark routing updates

Browse Source 
Resolve the C++ benchmark conflict by combining NCCL-EP-style random top-k routing and masked selections with the current BF16/FP8 MoERuntime API and format-aware byte accounting.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: efbacae6-f679-430b-bc16-b45ae162fc76
This commit is contained in:
Binyang Li
2026-07-13 21:49:11 +00:00
parent ad0b63d21b 359e681e8e
commit 205556f8cb
1 changed files with 30 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
test/python/ep/mscclpp_ep_bench.cu
View File

@@ -29,6 +29,7 @@
#include <mscclpp/gpu_data_types.hpp>
#include <random>
#include <string>
#include <utility>
#include <vector>
#include "api.cuh"
@@ -292,33 +293,42 @@ int main(int argc, char** argv) {
CUDA_CHECK(cudaMalloc(&d_layout, (size_t)Elocal * W * sizeof(int64_t)));
CUDA_CHECK(cudaMalloc(&d_count, (size_t)Elocal * sizeof(int)));
// Inputs (content is immaterial to timing). Route each rank independently to
// random distinct experts so top-k selections span destination ranks.
// Inputs. Token payloads are immaterial to timing, but the top-k routing is
// generated with the SAME scheme as NCCL-EP's ep_bench (generateRandomTopkIndicesLL):
// per-token abs(randn)+1 scores, take the top-k experts by score, then mask 10
// random (token, slot) positions with -1 to simulate dropped tokens.
CUDA_CHECK(cudaMemset(d_x, 0, (size_t)T * H * sizeof(Bf16)));
std::vector<int64_t> h_topk((size_t)T * K);
std::vector<float> h_weights((size_t)T * K);
std::mt19937 rng(args.seed + rank);
std::uniform_int_distribution<int> expertDist(0, E - 1);
std::uniform_real_distribution<float> weightDist(0.5f, 1.5f);
for (int t = 0; t < T; ++t) {
for (int j = 0; j < K; ++j) {
int expert;
bool duplicate;
do {
expert = expertDist(rng);
duplicate = false;
for (int previous = 0; previous < j; ++previous) {
duplicate |= h_topk[(size_t)t * K + previous] == expert;
}
} while (duplicate);
h_topk[(size_t)t * K + j] = expert;
h_weights[(size_t)t * K + j] = weightDist(rng);
std::vector<float> h_weights((size_t)T * K, 1.0f);
{
std::mt19937 gen(args.seed + rank);
std::normal_distribution<float> dist(0.0f, 1.0f);
std::vector<std::pair<float, int>> scoreIdx(E);
for (int t = 0; t < T; ++t) {
for (int e = 0; e < E; ++e) scoreIdx[e] = {std::abs(dist(gen)) + 1.0f, e};
std::partial_sort(scoreIdx.begin(), scoreIdx.begin() + K, scoreIdx.end(),
[](const auto& a, const auto& b) { return a.first > b.first; });
for (int j = 0; j < K; ++j) h_topk[(size_t)t * K + j] = scoreIdx[j].second;
}
// Randomly mask 10 positions with -1 (simulates dropped tokens); mirrors
// ep_bench. Guarded on T > 0 so the distribution bound is valid.
if (T > 0) {
std::uniform_int_distribution<int> tokenDist(0, T - 1);
std::uniform_int_distribution<int> topkDist(0, K - 1);
for (int i = 0; i < 10; ++i) {
int ti = tokenDist(gen), ki = topkDist(gen);
h_topk[(size_t)ti * K + ki] = -1;
}
}
}
CUDA_CHECK(cudaMemcpy(d_topk, h_topk.data(), h_topk.size() * sizeof(int64_t), cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(d_weights, h_weights.data(), h_weights.size() * sizeof(float), cudaMemcpyHostToDevice));
const long long num_valid_selections = (long long)T * K;
// Byte accounting counts only valid selections (topk >= 0), matching ep_bench's
// calculateLowLatencyBytes after the -1 masking above.
long long num_valid_selections = 0;
for (size_t i = 0; i < h_topk.size(); ++i)
if (h_topk[i] >= 0) ++num_valid_selections;
const double dispatchBytesPerToken = fp8Dispatch ? H + (H / 128) * sizeof(float) : H * sizeof(Bf16);
const double disp_bytes = (double)num_valid_selections * dispatchBytesPerToken;
const double comb_bytes = (double)num_valid_selections * H * sizeof(Bf16);