From cf6f091408c57de6db93af3b66761f9343e318ae Mon Sep 17 00:00:00 2001 From: 01Chen Date: Fri, 17 Jul 2026 15:02:03 +0800 Subject: [PATCH] Enhance EP demo runtime configuration --- tests/ep/demo/run_tilexr_ep_dispatch_demo.sh | 100 +++- tests/ep/demo/tilexr_ep_dispatch_demo.cpp | 479 ++++++++++++++---- .../ep/unit/test_tilexr_ep_kernel_sources.cpp | 55 ++ 3 files changed, 545 insertions(+), 89 deletions(-) diff --git a/tests/ep/demo/run_tilexr_ep_dispatch_demo.sh b/tests/ep/demo/run_tilexr_ep_dispatch_demo.sh index f7baff33..3d458776 100755 --- a/tests/ep/demo/run_tilexr_ep_dispatch_demo.sh +++ b/tests/ep/demo/run_tilexr_ep_dispatch_demo.sh @@ -7,8 +7,17 @@ TILEXR_ROOT="$(cd "${EP_DIR}/../.." && pwd)" INSTALL_DIR="${EP_DIR}/install" rank_size="${1:-2}" -npu_count="${2:-${rank_size}}" -first_npu="${3:-0}" +if [[ $# -gt 0 ]]; then + shift +fi +npu_count="${1:-${rank_size}}" +if [[ $# -gt 0 ]]; then + shift +fi +first_npu="${1:-0}" +if [[ $# -gt 0 ]]; then + shift +fi : "${ASCEND_HOME_PATH:=}" : "${LD_LIBRARY_PATH:=}" @@ -19,6 +28,91 @@ export TILEXR_DEMO_NPUS="${npu_count}" export TILEXR_DEMO_FIRST_NPU="${first_npu}" export LD_LIBRARY_PATH="${TILEXR_ROOT}/install/lib64:${TILEXR_ROOT}/install/lib:${INSTALL_DIR}/lib64:${INSTALL_DIR}/lib:${LD_LIBRARY_PATH:-}" +demo_args=() +while [[ $# -gt 0 ]]; do + case "$1" in + --mode) + export TILEXR_EP_DEMO_MODE="${2:?missing value for --mode}" + demo_args+=("$1" "$2") + shift 2 + ;; + --bs) + export TILEXR_EP_DEMO_BS="${2:?missing value for --bs}" + demo_args+=("$1" "$2") + shift 2 + ;; + --h) + export TILEXR_EP_DEMO_H="${2:?missing value for --h}" + demo_args+=("$1" "$2") + shift 2 + ;; + --k) + export TILEXR_EP_DEMO_TOPK="${2:?missing value for --k}" + demo_args+=("$1" "$2") + shift 2 + ;; + --moe-expert-num) + export TILEXR_EP_DEMO_MOE_EXPERT_NUM="${2:?missing value for --moe-expert-num}" + demo_args+=("$1" "$2") + shift 2 + ;; + --expert-token-nums-type) + export TILEXR_EP_DEMO_EXPERT_TOKEN_NUMS_TYPE="${2:?missing value for --expert-token-nums-type}" + demo_args+=("$1" "$2") + shift 2 + ;; + --active-mask) + export TILEXR_EP_DEMO_ACTIVE_MASK=1 + demo_args+=("$1") + shift + ;; + --tp-world-size) + export TILEXR_EP_DEMO_TP_WORLD_SIZE="${2:?missing value for --tp-world-size}" + demo_args+=("$1" "$2") + shift 2 + ;; + --tp-rank-id) + export TILEXR_EP_DEMO_TP_RANK_ID="${2:?missing value for --tp-rank-id}" + demo_args+=("$1" "$2") + shift 2 + ;; + --tp-recv-counts) + export TILEXR_EP_DEMO_TP_RECV_COUNTS=1 + demo_args+=("$1") + shift + ;; + --shared-expert-num) + export TILEXR_EP_DEMO_SHARED_EXPERT_NUM="${2:?missing value for --shared-expert-num}" + demo_args+=("$1" "$2") + shift 2 + ;; + --shared-expert-rank-num) + export TILEXR_EP_DEMO_SHARED_EXPERT_RANK_NUM="${2:?missing value for --shared-expert-rank-num}" + demo_args+=("$1" "$2") + shift 2 + ;; + --quant-mode) + export TILEXR_EP_DEMO_QUANT_MODE="${2:?missing value for --quant-mode}" + demo_args+=("$1" "$2") + shift 2 + ;; + --static-quant-scale) + export TILEXR_EP_DEMO_STATIC_QUANT_SCALE="${2:?missing value for --static-quant-scale}" + demo_args+=("$1" "$2") + shift 2 + ;; + --dump-window) + export TILEXR_EP_DEMO_DUMP_WINDOW=1 + demo_args+=("$1") + shift + ;; + *) + demo_args+=("$1") + shift + ;; + esac +done + bin="${INSTALL_DIR}/bin/tilexr_ep_dispatch_demo" if [[ ! -x "${bin}" ]]; then echo "Missing demo binary: ${bin}" >&2 @@ -37,7 +131,7 @@ for ((rank = 0; rank < rank_size; ++rank)); do ( export RANK="${rank}" export RANK_SIZE="${rank_size}" - exec "${bin}" "${rank_size}" "${rank}" "${npu_count}" "${first_npu}" + exec "${bin}" "${rank_size}" "${rank}" "${npu_count}" "${first_npu}" "${demo_args[@]}" ) >"${log}" 2>&1 & pids+=("$!") done diff --git a/tests/ep/demo/tilexr_ep_dispatch_demo.cpp b/tests/ep/demo/tilexr_ep_dispatch_demo.cpp index e3f34b77..479088a1 100644 --- a/tests/ep/demo/tilexr_ep_dispatch_demo.cpp +++ b/tests/ep/demo/tilexr_ep_dispatch_demo.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -22,30 +23,43 @@ namespace { -constexpr int64_t kBs = 4; -constexpr int64_t kH = 8; -constexpr int64_t kTopK = 2; -constexpr int64_t kRoutes = kBs * kTopK; -constexpr int64_t kXElements = kBs * kH; constexpr int64_t kAssistInts = 4; constexpr uint16_t kFp16One = 0x3c00; -constexpr uint16_t kFp16Two = 0x4000; constexpr std::size_t kUdmaCacheLineBytes = 64; constexpr std::size_t kUdmaRegistrationAlignment = 2 * 1024 * 1024; TileXRUDMAMemHandle g_workspaceHandle = 0; bool g_workspaceRegistered = false; +enum class DemoMode { + kDispatchCombine, + kDispatchOnly, + kCombineOnly, +}; + struct DemoConfig { + int64_t bs = 4; + int64_t h = 8; + int64_t topK = 2; int64_t moeExpertNum = 8; int64_t sharedExpertNum = 0; int64_t sharedExpertRankNum = 0; int64_t tpWorldSize = 0; int64_t tpRankId = 0; + int64_t routes() const + { + return bs * topK; + } + + int64_t xElements() const + { + return bs * h; + } + int64_t maxRoutesPerRank() const { - return kBs * (kTopK + sharedExpertNum); + return bs * (topK + sharedExpertNum); } int64_t effectiveTpWorldSize() const @@ -55,7 +69,7 @@ struct DemoConfig { int64_t expandedElements() const { - return maxRoutesPerRank() * effectiveTpWorldSize() * kH; + return maxRoutesPerRank() * effectiveTpWorldSize() * h; } }; @@ -70,13 +84,66 @@ bool EnvEnabled(const char *name) return value != nullptr && value[0] != '\0' && std::string(value) != "0"; } -int GetEnvInt(const char *name, int fallback) +bool ParseInt64(const std::string &label, const char *text, int64_t *out) { - const char *value = std::getenv(name); - if (value == nullptr || value[0] == '\0') { - return fallback; + if (text == nullptr || text[0] == '\0' || out == nullptr) { + std::cerr << "missing integer value for " << label << std::endl; + return false; } - return std::atoi(value); + errno = 0; + char *end = nullptr; + const long long value = std::strtoll(text, &end, 10); + if (errno != 0 || end == text || *end != '\0') { + std::cerr << "invalid " << label << ": " << text << std::endl; + return false; + } + *out = static_cast(value); + return true; +} + +bool LoadEnvInt64(const char *name, int64_t *value) +{ + const char *text = std::getenv(name); + if (text == nullptr || text[0] == '\0') { + return true; + } + return ParseInt64(name, text, value); +} + +bool ParseDemoMode(const std::string &text, DemoMode *mode) +{ + if (mode == nullptr) { + return false; + } + if (text == "dispatch-combine") { + *mode = DemoMode::kDispatchCombine; + return true; + } + if (text == "dispatch-only") { + *mode = DemoMode::kDispatchOnly; + return true; + } + if (text == "combine-only") { + *mode = DemoMode::kCombineOnly; + return true; + } + return false; +} + +bool LoadEnvDemoMode(DemoMode *mode) +{ + const char *text = std::getenv("TILEXR_EP_DEMO_MODE"); + if (text == nullptr || text[0] == '\0') { + if (EnvEnabled("TILEXR_DEMO_DISPATCH_ONLY")) { + *mode = DemoMode::kDispatchOnly; + } + return true; + } + if (!ParseDemoMode(text, mode)) { + std::cerr << "invalid TILEXR_EP_DEMO_MODE: " << text << std::endl; + return false; + } + return true; } bool ParseHostPort(const std::string &text, HostPort *out) @@ -85,12 +152,13 @@ bool ParseHostPort(const std::string &text, HostPort *out) if (out == nullptr || pos == std::string::npos || pos == 0 || pos + 1 >= text.size()) { return false; } - const int port = std::atoi(text.substr(pos + 1).c_str()); - if (port <= 0 || port > 65535) { + int64_t portValue = 0; + const std::string portText = text.substr(pos + 1); + if (!ParseInt64("port", portText.c_str(), &portValue) || portValue <= 0 || portValue > 65535) { return false; } out->host = text.substr(0, pos); - out->port = port; + out->port = static_cast(portValue); return true; } @@ -394,36 +462,36 @@ int8_t QuantizedXValue(int rank, int64_t token, int64_t h, float scale) return static_cast(rounded); } -float DynamicScaleForXValue(int rank, int64_t token) +float DynamicScaleForXValue(const DemoConfig &config, int rank, int64_t token) { float maxAbs = 0.0f; - for (int64_t h = 0; h < kH; ++h) { + for (int64_t h = 0; h < config.h; ++h) { const float value = std::fabs(HalfBitsToFloat(XValue(rank, token, h))); maxAbs = std::max(maxAbs, value); } return maxAbs > 0.0f ? maxAbs / 127.0f : 1.0f; } -int8_t DynamicQuantizedXValue(int rank, int64_t token, int64_t h) +int8_t DynamicQuantizedXValue(const DemoConfig &config, int rank, int64_t token, int64_t h) { - const float scale = DynamicScaleForXValue(rank, token); + const float scale = DynamicScaleForXValue(config, rank, token); return QuantizedXValue(rank, token, h, scale > 0.0f ? 1.0f / scale : 1.0f); } std::vector ExpertIds(const DemoConfig &config) { - std::vector expertIds(kRoutes); - for (int64_t route = 0; route < kRoutes; ++route) { + std::vector expertIds(static_cast(config.routes())); + for (int64_t route = 0; route < config.routes(); ++route) { expertIds[route] = static_cast(route % config.moeExpertNum); } return expertIds; } -std::vector ActiveMask(bool enabled) +std::vector ActiveMask(const DemoConfig &config, bool enabled) { - std::vector mask(kBs, 1); - if (enabled && kBs > 0) { - mask[kBs - 1] = 0; + std::vector mask(static_cast(config.bs), 1); + if (enabled && config.bs > 0) { + mask[static_cast(config.bs - 1)] = 0; } return mask; } @@ -493,18 +561,18 @@ std::vector BuildExpectedRoutes( if (effectiveTpWorldSize > 1 && srcRank % effectiveTpWorldSize != targetTpRankId) { continue; } - for (int64_t token = 0; token < kBs; ++token) { + for (int64_t token = 0; token < config.bs; ++token) { if (!activeMask.empty() && activeMask[token] == 0) { continue; } for (int64_t sharedExpertId = 0; sharedExpertId < config.sharedExpertNum; ++sharedExpertId) { if (RouteBelongsToRank(static_cast(sharedExpertId), rank, rankSize, config)) { expected.push_back(ExpectedRoute {srcRank, static_cast(token), - static_cast(kTopK + sharedExpertId), static_cast(sharedExpertId)}); + static_cast(config.topK + sharedExpertId), static_cast(sharedExpertId)}); } } - for (int64_t topKId = 0; topKId < kTopK; ++topKId) { - const int64_t route = token * kTopK + topKId; + for (int64_t topKId = 0; topKId < config.topK; ++topKId) { + const int64_t route = token * config.topK + topKId; const int32_t expertId = static_cast(config.sharedExpertNum) + expertIds[route]; if (RouteBelongsToRank(expertId, rank, rankSize, config)) { expected.push_back(ExpectedRoute {srcRank, static_cast(token), @@ -592,7 +660,7 @@ bool ValidateOutputs(int rank, int rankSize, const DemoConfig &config, const std } if (usePerTokenDynamicQuant) { - const float expectedScale = DynamicScaleForXValue(route.srcRank, route.tokenId); + const float expectedScale = DynamicScaleForXValue(config, route.srcRank, route.tokenId); const float actualScale = dynamicScalesOut[row]; if (std::fabs(actualScale - expectedScale) > 1.0e-5f) { std::cerr << "rank " << rank << " dynamicScalesOut[" << row << "] expected " @@ -601,14 +669,15 @@ bool ValidateOutputs(int rank, int rankSize, const DemoConfig &config, const std } } - for (int64_t h = 0; h < kH; ++h) { + for (int64_t h = 0; h < config.h; ++h) { const bool useInt8Output = useStaticQuant || usePerTokenDynamicQuant; - const std::size_t byteOffset = row * kH * (useInt8Output ? sizeof(int8_t) : sizeof(uint16_t)) + + const std::size_t byteOffset = row * static_cast(config.h) * + (useInt8Output ? sizeof(int8_t) : sizeof(uint16_t)) + h * (useInt8Output ? sizeof(int8_t) : sizeof(uint16_t)); const int expectedValue = useStaticQuant ? static_cast(QuantizedXValue(route.srcRank, route.tokenId, h, staticQuantScale)) : (usePerTokenDynamicQuant ? - static_cast(DynamicQuantizedXValue(route.srcRank, route.tokenId, h)) : + static_cast(DynamicQuantizedXValue(config, route.srcRank, route.tokenId, h)) : static_cast(XValue(route.srcRank, route.tokenId, h))); const int actualValue = useInt8Output ? static_cast(*reinterpret_cast(&expandX[byteOffset])) : @@ -624,14 +693,69 @@ bool ValidateOutputs(int rank, int rankSize, const DemoConfig &config, const std return true; } -bool ValidateCombineOutputs(int rank, const std::vector &yOut) +uint16_t HalfBitsForSmallPositiveInteger(int64_t value) +{ + if (value <= 0 || value > 2048) { + return 0; + } + int exponent = 0; + int64_t base = 1; + while ((base << 1) <= value) { + base <<= 1; + ++exponent; + } + const int64_t mantissa = ((value - base) * 1024) / base; + return static_cast(((exponent + 15) << 10) | mantissa); +} + +bool SeedCombineInputs(int rank, int rankSize, const DemoConfig &config, const std::vector &activeMask, + std::vector *expertOut, std::vector *recvCounts, std::vector *assist) { - for (int64_t token = 0; token < kBs; ++token) { - for (int64_t h = 0; h < kH; ++h) { - const uint16_t actualValue = yOut[token * kH + h]; - if (actualValue != kFp16Two) { + if (expertOut == nullptr || recvCounts == nullptr || assist == nullptr) { + return false; + } + std::fill(expertOut->begin(), expertOut->end(), kFp16One); + std::fill(recvCounts->begin(), recvCounts->end(), 0); + std::fill(assist->begin(), assist->end(), 0); + + const std::vector expected = BuildExpectedTpRoutes(rank, rankSize, config, activeMask); + if (assist->size() < expected.size() * static_cast(kAssistInts)) { + std::cerr << "rank " << rank << " assist buffer is too small for combine-only seed" << std::endl; + return false; + } + for (std::size_t row = 0; row < expected.size(); ++row) { + const ExpectedRoute &route = expected[row]; + if (route.srcRank < 0 || route.srcRank >= rankSize) { + std::cerr << "rank " << rank << " invalid seeded src rank " << route.srcRank << std::endl; + return false; + } + (*recvCounts)[static_cast(route.srcRank)] += 1; + const std::size_t offset = row * static_cast(kAssistInts); + (*assist)[offset] = route.srcRank; + (*assist)[offset + 1] = route.tokenId; + (*assist)[offset + 2] = route.topKId; + (*assist)[offset + 3] = route.expertId; + } + return true; +} + +bool ValidateCombineOutputs(int rank, const DemoConfig &config, const std::vector &activeMask, + const std::vector &yOut) +{ + const uint16_t activeExpected = HalfBitsForSmallPositiveInteger(config.topK); + if (activeExpected == 0) { + std::cerr << "rank " << rank << " unsupported combine validation topK=" << config.topK + << " (supported range: 1..2048)" << std::endl; + return false; + } + for (int64_t token = 0; token < config.bs; ++token) { + const uint16_t expected = (!activeMask.empty() && activeMask[static_cast(token)] == 0) ? + 0 : activeExpected; + for (int64_t h = 0; h < config.h; ++h) { + const uint16_t actualValue = yOut[static_cast(token * config.h + h)]; + if (actualValue != expected) { std::cerr << "rank " << rank << " yOut[" << token << "][" << h << "] expected 0x" - << std::hex << kFp16Two << " got 0x" << actualValue << std::dec << std::endl; + << std::hex << expected << " got 0x" << actualValue << std::dec << std::endl; return false; } } @@ -706,30 +830,182 @@ void Cleanup(TileXRCommPtr comm, aclrtStream stream, int deviceId, bool deviceSe int main(int argc, char **argv) { - const int rankSize = argc > 1 ? std::atoi(argv[1]) : GetEnvInt("RANK_SIZE", 2); - const int rank = argc > 2 ? std::atoi(argv[2]) : GetEnvInt("RANK", 0); - const int npuCount = argc > 3 ? std::atoi(argv[3]) : GetEnvInt("TILEXR_DEMO_NPUS", rankSize); - const int firstNpu = argc > 4 ? std::atoi(argv[4]) : GetEnvInt("TILEXR_DEMO_FIRST_NPU", 0); - const bool dispatchOnly = argc > 5 ? std::atoi(argv[5]) != 0 : GetEnvInt("TILEXR_DEMO_DISPATCH_ONLY", 0) != 0; - const bool useActiveMask = EnvEnabled("TILEXR_EP_DEMO_ACTIVE_MASK"); - const bool requestedTpRecvCounts = EnvEnabled("TILEXR_EP_DEMO_TP_RECV_COUNTS"); - const int expertTokenNumsType = GetEnvInt("TILEXR_EP_DEMO_EXPERT_TOKEN_NUMS_TYPE", 1); - const int quantMode = GetEnvInt("TILEXR_EP_DEMO_QUANT_MODE", 0); + int64_t rankSizeValue = 2; + int64_t rankValue = 0; + int64_t npuCountValue = 0; + int64_t firstNpuValue = 0; + if (!LoadEnvInt64("RANK_SIZE", &rankSizeValue) || + !LoadEnvInt64("RANK", &rankValue) || + !LoadEnvInt64("TILEXR_DEMO_NPUS", &npuCountValue) || + !LoadEnvInt64("TILEXR_DEMO_FIRST_NPU", &firstNpuValue)) { + return 2; + } + if (argc > 1 && !ParseInt64("rank_size", argv[1], &rankSizeValue)) { + return 2; + } + if (argc > 2 && !ParseInt64("rank", argv[2], &rankValue)) { + return 2; + } + if (argc > 3 && !ParseInt64("npu_count", argv[3], &npuCountValue)) { + return 2; + } + if (argc > 4 && !ParseInt64("first_npu", argv[4], &firstNpuValue)) { + return 2; + } + + DemoMode mode = DemoMode::kDispatchCombine; + if (!LoadEnvDemoMode(&mode)) { + return 2; + } + bool useActiveMask = EnvEnabled("TILEXR_EP_DEMO_ACTIVE_MASK"); + bool requestedTpRecvCounts = EnvEnabled("TILEXR_EP_DEMO_TP_RECV_COUNTS"); + int64_t expertTokenNumsType = 1; + int64_t quantMode = 0; + int64_t staticQuantScaleValue = 1; + bool dumpWindow = EnvEnabled("TILEXR_EP_DEMO_DUMP_WINDOW"); + bool tpRankIdExplicit = std::getenv("TILEXR_EP_DEMO_TP_RANK_ID") != nullptr; + DemoConfig config {}; + config.tpRankId = 0; + if (!LoadEnvInt64("TILEXR_EP_DEMO_BS", &config.bs) || + !LoadEnvInt64("TILEXR_EP_DEMO_H", &config.h) || + !LoadEnvInt64("TILEXR_EP_DEMO_TOPK", &config.topK) || + !LoadEnvInt64("TILEXR_EP_DEMO_MOE_EXPERT_NUM", &config.moeExpertNum) || + !LoadEnvInt64("TILEXR_EP_DEMO_SHARED_EXPERT_NUM", &config.sharedExpertNum) || + !LoadEnvInt64("TILEXR_EP_DEMO_SHARED_EXPERT_RANK_NUM", &config.sharedExpertRankNum) || + !LoadEnvInt64("TILEXR_EP_DEMO_TP_WORLD_SIZE", &config.tpWorldSize) || + !LoadEnvInt64("TILEXR_EP_DEMO_TP_RANK_ID", &config.tpRankId) || + !LoadEnvInt64("TILEXR_EP_DEMO_EXPERT_TOKEN_NUMS_TYPE", &expertTokenNumsType) || + !LoadEnvInt64("TILEXR_EP_DEMO_QUANT_MODE", &quantMode) || + !LoadEnvInt64("TILEXR_EP_DEMO_STATIC_QUANT_SCALE", &staticQuantScaleValue)) { + return 2; + } + + int optionIndex = 5; + if (argc > optionIndex && std::string(argv[optionIndex]).find("--") != 0) { + int64_t legacyDispatchOnly = 0; + if (!ParseInt64("dispatch_only", argv[optionIndex], &legacyDispatchOnly)) { + return 2; + } + if (legacyDispatchOnly != 0) { + mode = DemoMode::kDispatchOnly; + } + ++optionIndex; + } + for (int index = optionIndex; index < argc; ++index) { + const std::string option(argv[index]); + auto requireValue = [&](const std::string &name) -> const char * { + if (index + 1 >= argc) { + std::cerr << "missing value for " << name << std::endl; + return nullptr; + } + return argv[++index]; + }; + if (option == "--mode") { + const char *value = requireValue(option); + if (value == nullptr || !ParseDemoMode(value, &mode)) { + std::cerr << "invalid --mode: " << (value == nullptr ? "" : value) << std::endl; + return 2; + } + } else if (option == "--bs") { + const char *value = requireValue(option); + if (value == nullptr || !ParseInt64(option, value, &config.bs)) { + return 2; + } + } else if (option == "--h") { + const char *value = requireValue(option); + if (value == nullptr || !ParseInt64(option, value, &config.h)) { + return 2; + } + } else if (option == "--k") { + const char *value = requireValue(option); + if (value == nullptr || !ParseInt64(option, value, &config.topK)) { + return 2; + } + } else if (option == "--moe-expert-num") { + const char *value = requireValue(option); + if (value == nullptr || !ParseInt64(option, value, &config.moeExpertNum)) { + return 2; + } + } else if (option == "--expert-token-nums-type") { + const char *value = requireValue(option); + if (value == nullptr || !ParseInt64(option, value, &expertTokenNumsType)) { + return 2; + } + } else if (option == "--active-mask") { + useActiveMask = true; + } else if (option == "--tp-world-size") { + const char *value = requireValue(option); + if (value == nullptr || !ParseInt64(option, value, &config.tpWorldSize)) { + return 2; + } + } else if (option == "--tp-rank-id") { + const char *value = requireValue(option); + if (value == nullptr || !ParseInt64(option, value, &config.tpRankId)) { + return 2; + } + tpRankIdExplicit = true; + } else if (option == "--tp-recv-counts") { + requestedTpRecvCounts = true; + } else if (option == "--shared-expert-num") { + const char *value = requireValue(option); + if (value == nullptr || !ParseInt64(option, value, &config.sharedExpertNum)) { + return 2; + } + } else if (option == "--shared-expert-rank-num") { + const char *value = requireValue(option); + if (value == nullptr || !ParseInt64(option, value, &config.sharedExpertRankNum)) { + return 2; + } + } else if (option == "--quant-mode") { + const char *value = requireValue(option); + if (value == nullptr || !ParseInt64(option, value, &quantMode)) { + return 2; + } + } else if (option == "--static-quant-scale") { + const char *value = requireValue(option); + if (value == nullptr || !ParseInt64(option, value, &staticQuantScaleValue)) { + return 2; + } + } else if (option == "--dump-window") { + dumpWindow = true; + } else { + std::cerr << "unknown option: " << option << std::endl; + return 2; + } + } + + if (config.effectiveTpWorldSize() > 1 && !tpRankIdExplicit) { + config.tpRankId = rankValue % config.effectiveTpWorldSize(); + } + if (npuCountValue == 0) { + npuCountValue = rankSizeValue; + } + if (rankSizeValue <= 0 || rankSizeValue > std::numeric_limits::max() || + rankValue < 0 || rankValue > std::numeric_limits::max() || + npuCountValue <= 0 || npuCountValue > std::numeric_limits::max() || + firstNpuValue < 0 || firstNpuValue > std::numeric_limits::max()) { + std::cerr << "invalid rank/device arguments: rankSize=" << rankSizeValue + << " rank=" << rankValue << " npuCount=" << npuCountValue + << " firstNpu=" << firstNpuValue << std::endl; + return 2; + } + const int rankSize = static_cast(rankSizeValue); + const int rank = static_cast(rankValue); + const int npuCount = static_cast(npuCountValue); + const int firstNpu = static_cast(firstNpuValue); + const bool dispatchOnly = mode == DemoMode::kDispatchOnly; + const bool combineOnly = mode == DemoMode::kCombineOnly; + const bool runDispatch = !combineOnly; + const bool runCombine = !dispatchOnly; const bool useStaticQuant = quantMode == 1; const bool usePerTokenDynamicQuant = quantMode == 2; - const float staticQuantScale = static_cast(GetEnvInt("TILEXR_EP_DEMO_STATIC_QUANT_SCALE", 1)); - DemoConfig config {}; - config.moeExpertNum = GetEnvInt("TILEXR_EP_DEMO_MOE_EXPERT_NUM", static_cast(config.moeExpertNum)); - config.sharedExpertNum = GetEnvInt("TILEXR_EP_DEMO_SHARED_EXPERT_NUM", 0); - config.sharedExpertRankNum = GetEnvInt("TILEXR_EP_DEMO_SHARED_EXPERT_RANK_NUM", 0); - config.tpWorldSize = GetEnvInt("TILEXR_EP_DEMO_TP_WORLD_SIZE", 0); - config.tpRankId = GetEnvInt("TILEXR_EP_DEMO_TP_RANK_ID", - config.effectiveTpWorldSize() > 1 ? rank % config.effectiveTpWorldSize() : 0); + const float staticQuantScale = static_cast(staticQuantScaleValue); const bool useTpRecvCounts = requestedTpRecvCounts || config.effectiveTpWorldSize() != 1; const int64_t expertRankSize = static_cast(rankSize) / config.effectiveTpWorldSize(); const int64_t moeRankNum = expertRankSize - config.sharedExpertRankNum; - if (rankSize <= 0 || rank < 0 || rank >= rankSize || config.effectiveTpWorldSize() <= 0 || + if (rankSize <= 0 || rank < 0 || rank >= rankSize || config.bs <= 0 || config.h <= 0 || config.topK <= 0 || + config.routes() <= 0 || config.xElements() <= 0 || config.effectiveTpWorldSize() <= 0 || rankSize % config.effectiveTpWorldSize() != 0 || moeRankNum <= 0 || config.moeExpertNum <= 0 || config.moeExpertNum % moeRankNum != 0 || config.sharedExpertNum < 0 || config.sharedExpertRankNum < 0 || @@ -740,6 +1016,9 @@ int main(int argc, char **argv) ((useStaticQuant || usePerTokenDynamicQuant) && !dispatchOnly)) { std::cerr << "This demo expects a valid rank and moeExpertNum divisible by MoE rank num, got moeExpertNum=" << config.moeExpertNum << " rankSize=" << rankSize + << " bs=" << config.bs + << " h=" << config.h + << " topK=" << config.topK << " sharedExpertNum=" << config.sharedExpertNum << " sharedExpertRankNum=" << config.sharedExpertRankNum << " tpWorldSize=" << config.tpWorldSize @@ -749,6 +1028,10 @@ int main(int argc, char **argv) << " quantMode=" << quantMode << std::endl; return 2; } + if (combineOnly && (useStaticQuant || usePerTokenDynamicQuant)) { + std::cerr << "combine-only does not support quantMode=" << quantMode << std::endl; + return 2; + } const int64_t localExpertNum = LocalExpertNum(rankSize, config); const int deviceId = GetDeviceIdFromEnv(rank, npuCount, firstNpu); @@ -777,14 +1060,14 @@ int main(int argc, char **argv) return 1; } - std::vector hostX(kXElements); - for (int64_t token = 0; token < kBs; ++token) { - for (int64_t h = 0; h < kH; ++h) { - hostX[token * kH + h] = XValue(rank, token, h); + std::vector hostX(static_cast(config.xElements())); + for (int64_t token = 0; token < config.bs; ++token) { + for (int64_t h = 0; h < config.h; ++h) { + hostX[static_cast(token * config.h + h)] = XValue(rank, token, h); } } const std::vector hostExpertIds = ExpertIds(config); - const std::vector hostActiveMask = ActiveMask(useActiveMask); + const std::vector hostActiveMask = ActiveMask(config, useActiveMask); const std::size_t expectedRouteCount = BuildExpectedTpRoutes(rank, rankSize, config, hostActiveMask).size(); @@ -809,19 +1092,19 @@ int main(int argc, char **argv) const std::size_t expertIdsBytes = hostExpertIds.size() * sizeof(int32_t); const std::size_t xActiveMaskBytes = hostActiveMask.size() * sizeof(uint8_t); const std::size_t expandedElements = std::max(static_cast(config.expandedElements()), - expectedRouteCount * static_cast(kH)); + expectedRouteCount * static_cast(config.h)); const std::size_t maxRoutesPerRank = static_cast(config.maxRoutesPerRank()); const std::size_t expandElementBytes = (useStaticQuant || usePerTokenDynamicQuant) ? sizeof(int8_t) : sizeof(uint16_t); const std::size_t expandXBytes = expandedElements * expandElementBytes; - const std::size_t expandedRows = expandedElements / kH; + const std::size_t expandedRows = expandedElements / static_cast(config.h); const std::size_t dynamicScalesBytes = expandedRows * sizeof(float); const std::size_t expertTokenNumsBytes = localExpertNum * sizeof(int64_t); const std::size_t recvCountsBytes = rankSize * sizeof(int32_t); const std::size_t tpRecvCountsBytes = recvCountsBytes; - const std::size_t assistBytes = (expandedElements / kH) * kAssistInts * sizeof(int32_t); - const std::size_t yOutBytes = kXElements * sizeof(uint16_t); - const std::size_t payloadRowBytes = kH * expandElementBytes; + const std::size_t assistBytes = expandedRows * kAssistInts * sizeof(int32_t); + const std::size_t yOutBytes = static_cast(config.xElements()) * sizeof(uint16_t); + const std::size_t payloadRowBytes = static_cast(config.h) * expandElementBytes; const std::size_t dispatchWindowBytes = EpWindowBytes(rankSize, config, payloadRowBytes, usePerTokenDynamicQuant); const std::size_t dispatchPayloadBytes = AlignSize(dispatchWindowBytes, 32) * static_cast(config.effectiveTpWorldSize() + 2); @@ -899,31 +1182,51 @@ int main(int argc, char **argv) g_workspaceRegistered = true; } - const std::vector hostExpertOut(expandedElements, kFp16One); + std::vector hostExpertOut(expandedElements, kFp16One); + if (combineOnly) { + std::vector seededRecvCounts(rankSize); + std::vector seededAssist(expandedRows * static_cast(kAssistInts)); + if (!SeedCombineInputs(rank, rankSize, config, hostActiveMask, &hostExpertOut, + &seededRecvCounts, &seededAssist)) { + Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); + return 1; + } + if (!CheckAcl(aclrtMemcpy(recvCountsDev, recvCountsBytes, seededRecvCounts.data(), recvCountsBytes, + ACL_MEMCPY_HOST_TO_DEVICE), "copy seeded recvCounts") || + !CheckAcl(aclrtMemcpy(assistDev, assistBytes, seededAssist.data(), assistBytes, + ACL_MEMCPY_HOST_TO_DEVICE), "copy seeded assist")) { + Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); + return 1; + } + } if (!CheckAcl(aclrtMemcpy(expertOutDev, expandXBytes, hostExpertOut.data(), expandXBytes, ACL_MEMCPY_HOST_TO_DEVICE), "copy expertOut")) { Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); return 1; } + if (runDispatch) { const bool useSharedExperts = config.sharedExpertNum != 0 || config.sharedExpertRankNum != 0; const bool useTp = config.effectiveTpWorldSize() != 1; const bool useDispatchV2 = crossNode || useActiveMask || useTpRecvCounts || expertTokenNumsType != 1 || useSharedExperts || useTp || useStaticQuant || usePerTokenDynamicQuant; const int dispatchRet = useDispatchV2 ? TileXRMoeEpDispatchV2(xDev, static_cast(expertIdsDev), scalesDev, - static_cast(xActiveMaskDev), nullptr, comm, kBs, kH, kTopK, config.moeExpertNum, + static_cast(xActiveMaskDev), nullptr, comm, config.bs, config.h, config.topK, + config.moeExpertNum, expertRankSize, ExpertRankForRank(rank, config), config.tpWorldSize, config.tpRankId, 0, - config.sharedExpertNum, config.sharedExpertRankNum, quantMode, kBs * rankSize, expertTokenNumsType, + config.sharedExpertNum, config.sharedExpertRankNum, quantMode, config.bs * rankSize, + expertTokenNumsType, expandXDev, dynamicScalesDev, static_cast(assistDev), static_cast(expertTokenNumsDev), static_cast(recvCountsDev), static_cast(tpRecvCountsDev), nullptr, workspaceDev, (useStaticQuant || usePerTokenDynamicQuant) ? TileXR::TILEXR_DATA_TYPE_INT8 : TileXR::TILEXR_DATA_TYPE_FP16, stream) : - TileXRMoeEpDispatch(xDev, static_cast(expertIdsDev), comm, kBs, kH, kTopK, config.moeExpertNum, - expandXDev, static_cast(expertTokenNumsDev), static_cast(recvCountsDev), - static_cast(assistDev), TileXR::TILEXR_DATA_TYPE_FP16, stream); + TileXRMoeEpDispatch(xDev, static_cast(expertIdsDev), comm, config.bs, config.h, config.topK, + config.moeExpertNum, expandXDev, static_cast(expertTokenNumsDev), + static_cast(recvCountsDev), static_cast(assistDev), + TileXR::TILEXR_DATA_TYPE_FP16, stream); if (!CheckTileXR(dispatchRet, "TileXRMoeEpDispatch") || !CheckAcl(aclrtSynchronizeStream(stream), "aclrtSynchronizeStream")) { Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); @@ -934,8 +1237,8 @@ int main(int argc, char **argv) return 1; } - if (EnvEnabled("TILEXR_EP_DEMO_DUMP_WINDOW") && commArgsHost != nullptr) { - const std::size_t rowBytes = kH * expandElementBytes; + if (dumpWindow && commArgsHost != nullptr) { + const std::size_t rowBytes = static_cast(config.h) * expandElementBytes; const std::size_t payloadBytes = AlignSize(maxRoutesPerRank * rowBytes + (usePerTokenDynamicQuant ? maxRoutesPerRank * sizeof(float) : 0), 32); const std::size_t assistWindowBytes = AlignSize(maxRoutesPerRank * kAssistInts * sizeof(int32_t), 32); @@ -981,7 +1284,7 @@ int main(int argc, char **argv) std::vector hostExpertTokenNums(localExpertNum); std::vector hostRecvCounts(rankSize); std::vector hostTpRecvCounts(rankSize); - std::vector hostAssist((expandedElements / kH) * kAssistInts); + std::vector hostAssist(expandedRows * static_cast(kAssistInts)); std::vector hostDynamicScales(expandedRows); if (!CheckAcl(aclrtMemcpy(hostExpandX.data(), expandXBytes, expandXDev, expandXBytes, @@ -1007,23 +1310,27 @@ int main(int argc, char **argv) } const bool dispatchOk = ValidateOutputs(rank, rankSize, config, hostExpandX, hostExpertTokenNums, hostRecvCounts, - hostAssist, hostDynamicScales, hostActiveMask, expertTokenNumsType, useStaticQuant, + hostAssist, hostDynamicScales, hostActiveMask, static_cast(expertTokenNumsType), useStaticQuant, usePerTokenDynamicQuant, staticQuantScale) && (!useTpRecvCounts || ValidateTpRecvCounts(rank, rankSize, config, hostActiveMask, hostRecvCounts, hostTpRecvCounts)); - std::cout << "rank " << rank << " validation " << (dispatchOk ? "PASS" : "FAIL") << std::endl; - if (!dispatchOk || dispatchOnly) { + std::cout << "rank " << rank << " dispatch validation " << (dispatchOk ? "PASS" : "FAIL") << std::endl; + if (!dispatchOk || !runCombine) { Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); return dispatchOk ? 0 : 1; } + } else if (!DemoBarrierAll(rank, rankSize, "combine inputs seeded")) { + Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); + return 1; + } const int combineRet = crossNode ? TileXRMoeEpCombineV2(expertOutDev, static_cast(assistDev), - static_cast(recvCountsDev), comm, kBs, kH, kTopK, config.moeExpertNum, yOutDev, workspaceDev, - TileXR::TILEXR_DATA_TYPE_FP16, stream) : + static_cast(recvCountsDev), comm, config.bs, config.h, config.topK, config.moeExpertNum, + yOutDev, workspaceDev, TileXR::TILEXR_DATA_TYPE_FP16, stream) : TileXRMoeEpCombine(expertOutDev, static_cast(assistDev), - static_cast(recvCountsDev), comm, kBs, kH, kTopK, config.moeExpertNum, yOutDev, - TileXR::TILEXR_DATA_TYPE_FP16, stream); + static_cast(recvCountsDev), comm, config.bs, config.h, config.topK, config.moeExpertNum, + yOutDev, TileXR::TILEXR_DATA_TYPE_FP16, stream); if (!CheckTileXR(combineRet, "TileXRMoeEpCombine") || !CheckAcl(aclrtSynchronizeStream(stream), "aclrtSynchronizeStream combine")) { Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); @@ -1034,14 +1341,14 @@ int main(int argc, char **argv) return 1; } - std::vector hostYOut(kXElements); + std::vector hostYOut(static_cast(config.xElements())); if (!CheckAcl(aclrtMemcpy(hostYOut.data(), yOutBytes, yOutDev, yOutBytes, ACL_MEMCPY_DEVICE_TO_HOST), "copy yOut")) { Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); return 1; } - const bool combineOk = ValidateCombineOutputs(rank, hostYOut); + const bool combineOk = ValidateCombineOutputs(rank, config, hostActiveMask, hostYOut); std::cout << "rank " << rank << " combine validation " << (combineOk ? "PASS" : "FAIL") << std::endl; Cleanup(comm, stream, deviceId, deviceSet, aclReady, buffers); return combineOk ? 0 : 1; diff --git a/tests/ep/unit/test_tilexr_ep_kernel_sources.cpp b/tests/ep/unit/test_tilexr_ep_kernel_sources.cpp index a8100561..46723ec3 100644 --- a/tests/ep/unit/test_tilexr_ep_kernel_sources.cpp +++ b/tests/ep/unit/test_tilexr_ep_kernel_sources.cpp @@ -341,6 +341,58 @@ void TestDispatchDemoExercisesV2OptionalInputs() CheckContains(path, contents, "expertTokenNumsType"); } +void TestDispatchDemoUsesRuntimeShapeConfig() +{ + const std::string path = "tests/ep/demo/tilexr_ep_dispatch_demo.cpp"; + std::string contents; + if (!ReadFile(path, &contents)) { + return; + } + + CheckContains(path, contents, "TILEXR_EP_DEMO_BS"); + CheckContains(path, contents, "TILEXR_EP_DEMO_H"); + CheckContains(path, contents, "TILEXR_EP_DEMO_TOPK"); + CheckContains(path, contents, "config.bs"); + CheckContains(path, contents, "config.h"); + CheckContains(path, contents, "config.topK"); + CheckNotContains(path, contents, "constexpr int64_t kBs"); + CheckNotContains(path, contents, "constexpr int64_t kH"); + CheckNotContains(path, contents, "constexpr int64_t kTopK"); +} + +void TestDispatchDemoSupportsCombineOnlyMode() +{ + const std::string path = "tests/ep/demo/tilexr_ep_dispatch_demo.cpp"; + std::string contents; + if (!ReadFile(path, &contents)) { + return; + } + + CheckContains(path, contents, "TILEXR_EP_DEMO_MODE"); + CheckContains(path, contents, "combine-only"); + CheckContains(path, contents, "SeedCombineInputs"); + CheckContains(path, contents, "ValidateCombineOutputs"); + CheckContains(path, contents, "dispatch validation"); +} + +void TestDemoRunnerAcceptsNamedShapeArgs() +{ + const std::string path = "tests/ep/demo/run_tilexr_ep_dispatch_demo.sh"; + std::string contents; + if (!ReadFile(path, &contents)) { + return; + } + + CheckContains(path, contents, "--bs"); + CheckContains(path, contents, "--h"); + CheckContains(path, contents, "--k"); + CheckContains(path, contents, "--mode"); + CheckContains(path, contents, "TILEXR_EP_DEMO_BS"); + CheckContains(path, contents, "TILEXR_EP_DEMO_H"); + CheckContains(path, contents, "TILEXR_EP_DEMO_TOPK"); + CheckContains(path, contents, "TILEXR_EP_DEMO_MODE"); +} + void TestKernelForwardsSharedExpertConfig() { std::string hostLaunch; @@ -464,6 +516,9 @@ int main() TestKernelForwardsExpertTokenNumsType(); TestKernelForwardsTpRecvCountsOut(); TestDispatchDemoExercisesV2OptionalInputs(); + TestDispatchDemoUsesRuntimeShapeConfig(); + TestDispatchDemoSupportsCombineOnlyMode(); + TestDemoRunnerAcceptsNamedShapeArgs(); TestKernelForwardsSharedExpertConfig(); TestKernelForwardsStaticQuantConfig(); TestKernelForwardsPerTokenDynamicQuantConfig();