diff --git a/lib/Runtime/Kernels/hip/qmoe_kernel.hip b/lib/Runtime/Kernels/hip/qmoe_kernel.hip index 7b8e0e37cd..8be255cfc7 100644 --- a/lib/Runtime/Kernels/hip/qmoe_kernel.hip +++ b/lib/Runtime/Kernels/hip/qmoe_kernel.hip @@ -17,6 +17,7 @@ #include #include #include +#include static constexpr int kBlock = 256; static constexpr int kMaxK = 16; @@ -1239,17 +1240,21 @@ static hipError_t dispatch_qmoe_fc2( n_blk_mid, block_size); } -// Online per-shape autotune for the decode fc1/fc2 GEMV kernels. The optimal -// (BLOCK_SIZE, TILE_N) depends on K/32 (work units per block) and N (grid +// Per-shape choice of (BLOCK_SIZE, TILE_N) for the decode fc1/fc2 GEMV +// kernels. The best one depends on K/32 (work units per block) and N (grid // width): the hand-picked BS=64/TN=8 suits gpt-oss-20b (K=2880) but is slower -// than BS=32 on Qwen3.5 (K=2048/512, fewer work units → BS=64 leaves threads -// idle). A synthetic tight-loop benchmark mis-ranks TILE_N: it reuses one -// expert set across all timed iters (hot cache), while in steady state each -// layer reads cold weights, flipping the memory-bound ranking. So tune ONLINE -// — sample each candidate on real decode calls (true cache regime), accumulate -// GPU time, lock in the fastest after enough samples. Keyed by -// (N, K, block_size, has_zp); dims are model-static so one tuning window per -// shape. State/lock have process lifetime (kernels are stateless). +// than BS=32 on Qwen3.5 (K=2048/512, fewer work units -> BS=64 leaves threads +// idle). Keyed by (N, K, block_size, has_zp); dims are model-static. +// +// The choice is made once per shape, on its first call and before that call +// computes anything: every candidate is timed on the call's own input and +// weights, each launch on a different set of k experts, so the weights are +// read cold as in steady-state decode, where each layer reads other weights +// (a loop over one expert set reads them hot and mis-ranks TILE_N). The +// call's real launches, and every later call's, then use the chosen +// configuration. The block size sets the reduction order, so sampling +// candidates on real calls instead - as this used to - gave the first ~100 +// calls of a process slightly different results for the same inputs. struct QmoeCfg { int bs; int tn; @@ -1264,23 +1269,15 @@ static constexpr int kNumQmoeCfgs = static_cast(sizeof(kQmoeCfgs) / sizeof(kQmoeCfgs[0])); static constexpr int kQmoeDefaultCfg = 4; // index of {64, 8} -// One round = one full sweep of all candidates (one sample each, round-robin, -// one per decode call). First kWarmupRounds sweeps are discarded (cold start); -// the next kSampleRounds accumulate; then best = min mean and done. +// One discarded round of every candidate (cold start), then kSampleRounds +// timed ones; the least total time wins. static constexpr int kQmoeWarmupRounds = 1; static constexpr int kQmoeSampleRounds = 8; -struct QmoeTuneState { - bool done = false; - QmoeCfg best = kQmoeCfgs[kQmoeDefaultCfg]; - double ms[kNumQmoeCfgs] = {}; - int n[kNumQmoeCfgs] = {}; - int cursor = 0; - int rounds = 0; -}; - -static std::unordered_map g_qmoe_fc1_state; -static std::unordered_map g_qmoe_fc2_state; +static std::unordered_map g_qmoe_fc1_choice; +static std::unordered_map g_qmoe_fc2_choice; +static std::unordered_map g_qmoe_fc1_dp4a_choice; +static std::unordered_map g_qmoe_fc2_dp4a_choice; static std::mutex g_qmoe_tune_mutex; static uint64_t qmoe_cfg_key(int n, int k, int block_size, bool has_zp) { @@ -1290,60 +1287,96 @@ static uint64_t qmoe_cfg_key(int n, int k, int block_size, bool has_zp) { (has_zp ? 1ull : 0ull); } - -// Pick the config for this call. While tuning, returns the candidate being -// sampled and sets *sample_idx >= 0 (caller times this launch and reports via -// qmoe_tune_record). Once done, returns best and *sample_idx = -1. -static QmoeCfg qmoe_tune_pick( - std::unordered_map& states, uint64_t key, - int* sample_idx) { - std::lock_guard lk(g_qmoe_tune_mutex); - QmoeTuneState& st = states[key]; - if (st.done) { - *sample_idx = -1; - return st.best; +// Device expert-index sets for timing: set c holds experts +// (c*k + j) % num_experts, j < k. Allocated once per (k, num_experts). +static const int32_t* qmoe_tune_indices(int64_t k, int64_t num_experts, + int* sets) { + static std::unordered_map> cache; + const uint64_t key = (static_cast(k) << 32) ^ + static_cast(num_experts); + auto it = cache.find(key); + if (it == cache.end()) { + const int n = static_cast( + num_experts >= k ? num_experts / k : 1); + std::vector host(static_cast(n * k)); + for (int c = 0; c < n; ++c) + for (int64_t j = 0; j < k; ++j) + host[static_cast(c * k + j)] = + static_cast((c * k + j) % num_experts); + int32_t* device = nullptr; + if (hipMalloc(&device, host.size() * sizeof(int32_t)) != hipSuccess || + hipMemcpy(device, host.data(), host.size() * sizeof(int32_t), + hipMemcpyHostToDevice) != hipSuccess) { + if (device) + (void)hipFree(device); + return nullptr; + } + it = cache.emplace(key, std::make_pair(device, n)).first; } - *sample_idx = st.cursor; - return kQmoeCfgs[st.cursor]; + *sets = it->second.second; + return it->second.first; } -// Record one measured sample (real-call GPU time) for the candidate at -// sample_idx, advance the round-robin cursor, and finalize once enough sample -// rounds have completed. Warmup rounds are discarded (cold-start noise). -static void qmoe_tune_record( - std::unordered_map& states, uint64_t key, - int sample_idx, float ms, const char* tag, int n, int kdim, - int block_size) { +// The configuration of shape `key`, timing the candidates first when it has +// none yet. `launch(cfg, expert_indices)` runs one candidate; its output is +// scratch the call overwrites afterwards. While the stream is being captured +// into a HIP graph nothing can be timed: the default is used for this call +// and the choice is left for a later one. +template +static QmoeCfg qmoe_choose(std::unordered_map& choices, + uint64_t key, hipStream_t s, int64_t k, + int64_t num_experts, const char* tag, + Launch launch) { std::lock_guard lk(g_qmoe_tune_mutex); - QmoeTuneState& st = states[key]; - if (st.done || sample_idx != st.cursor) - return; // stale sample (shape raced to done) - if (st.rounds >= kQmoeWarmupRounds) { - st.ms[sample_idx] += ms; - st.n[sample_idx] += 1; - } - if (++st.cursor >= kNumQmoeCfgs) { - st.cursor = 0; - if (++st.rounds >= kQmoeWarmupRounds + kQmoeSampleRounds) { - double best_mean = 1e30; - int best = kQmoeDefaultCfg; - for (int i = 0; i < kNumQmoeCfgs; i++) { - if (st.n[i] == 0) - continue; - double mean = st.ms[i] / st.n[i]; - if (mean < best_mean) { - best_mean = mean; - best = i; - } - } - st.best = kQmoeCfgs[best]; - st.done = true; - CUSTOM_KERNELS_DEBUG_LOG( - "[custom_kernels] qmoe %s online-autotune N=%d K=%d block=%d -> " - "BS=%d TN=%d (%.1f us)\n", - tag, n, kdim, block_size, st.best.bs, st.best.tn, best_mean * 1000.0); + auto found = choices.find(key); + if (found != choices.end()) + return found->second; + hipStreamCaptureStatus capture = hipStreamCaptureStatusNone; + if (hipStreamIsCapturing(s, &capture) != hipSuccess || + capture != hipStreamCaptureStatusNone || num_experts <= 0) + return kQmoeCfgs[kQmoeDefaultCfg]; + int sets = 0; + const int32_t* indices = qmoe_tune_indices(k, num_experts, &sets); + if (!indices) + return kQmoeCfgs[kQmoeDefaultCfg]; + constexpr int kRounds = kQmoeWarmupRounds + kQmoeSampleRounds; + std::vector events(static_cast(2 * kRounds * kNumQmoeCfgs)); + for (hipEvent_t& e : events) + if (hipEventCreate(&e) != hipSuccess) + e = nullptr; + int launches = 0; + bool ok = true; + for (int r = 0; r < kRounds && ok; ++r) + for (int c = 0; c < kNumQmoeCfgs && ok; ++c) { + const int32_t* set = indices + (launches++ % sets) * k; + hipEvent_t* pair = &events[static_cast(2 * (r * kNumQmoeCfgs + c))]; + ok = pair[0] && pair[1] && hipEventRecord(pair[0], s) == hipSuccess && + launch(kQmoeCfgs[c], set) == hipSuccess && + hipEventRecord(pair[1], s) == hipSuccess; } + int best = kQmoeDefaultCfg; + if (ok && hipEventSynchronize(events.back()) == hipSuccess) { + double total[kNumQmoeCfgs] = {}; + for (int r = kQmoeWarmupRounds; r < kRounds; ++r) + for (int c = 0; c < kNumQmoeCfgs; ++c) { + float ms = 0; + const size_t at = static_cast(2 * (r * kNumQmoeCfgs + c)); + hipEventElapsedTime(&ms, events[at], events[at + 1]); + total[c] += ms; + } + for (int c = 0; c < kNumQmoeCfgs; ++c) + if (total[c] < total[best]) + best = c; + CUSTOM_KERNELS_DEBUG_LOG( + "[custom_kernels] qmoe %s autotune -> BS=%d TN=%d (%.1f us)\n", tag, + kQmoeCfgs[best].bs, kQmoeCfgs[best].tn, + total[best] * 1000.0 / kQmoeSampleRounds); } + for (hipEvent_t e : events) + if (e) + (void)hipEventDestroy(e); + choices.emplace(key, kQmoeCfgs[best]); + return kQmoeCfgs[best]; } extern "C" int hip_qmoe_decode_fused( @@ -1361,7 +1394,7 @@ extern "C" int hip_qmoe_decode_fused( int64_t hidden_size, int64_t inter_size, int64_t k, int64_t block_size, float swiglu_alpha, float swiglu_beta, float swiglu_limit, - int64_t element_size_bytes) { + int64_t element_size_bytes, int64_t num_experts) { if (element_size_bytes != 2) { fprintf(stderr, "[custom_kernels] hip_qmoe_decode_fused: only fp16 " "supported (got element_size %lld)\n", @@ -1387,89 +1420,51 @@ extern "C" int hip_qmoe_decode_fused( } hipStream_t s = static_cast(stream); + // Each launch below reports hipGetLastError(); clear an error left by an + // earlier call so it is not taken for theirs, including the tiling timing. + (void)hipGetLastError(); const int fusion_inter = static_cast(2 * inter_size); const int n_blk_in = static_cast((hidden_size + block_size - 1) / block_size); const int n_blk_mid = static_cast((inter_size + block_size - 1) / block_size); - // FC1 config: N=fusion_inter, K=hidden. Online per-shape autotune. - int s1 = -1; - const uint64_t key1 = - qmoe_cfg_key(fusion_inter, static_cast(hidden_size), - static_cast(block_size), fc1_zero_points != nullptr); - const QmoeCfg c1 = qmoe_tune_pick(g_qmoe_fc1_state, key1, &s1); - // FC1 + SwiGLU fused: grid (fusion_inter/TN, k). Each block produces TN // matmul outputs internally and pair-wise SwiGLUs them down to TN/2 // act_out elements. No dynamic LDS (input read straight from L2/MALL). - // While s1 >= 0 the shape is still tuning: bracket this real launch with - // events and feed its GPU time back to the online tuner. - { - hipEvent_t me0 = nullptr, me1 = nullptr; - if (s1 >= 0) { - hipEventCreate(&me0); - hipEventCreate(&me1); - hipEventRecord(me0, s); - } - hipError_t e = dispatch_qmoe_fc1( - c1.bs, c1.tn, s, static_cast(k), fusion_inter, input, - expert_indices, fc1_weights, fc1_scales, fc1_zero_points, fc1_bias, - act_out, static_cast(hidden_size), n_blk_in, - static_cast(block_size), swiglu_alpha, swiglu_beta, swiglu_limit); - if (e != hipSuccess) { - fprintf(stderr, "[custom_kernels] qmoe_decode_fc1 failed: %s\n", - hipGetErrorString(e)); - return static_cast(e); - } - if (s1 >= 0) { - hipEventRecord(me1, s); - hipEventSynchronize(me1); - float ms = 0; - hipEventElapsedTime(&ms, me0, me1); - hipEventDestroy(me0); - hipEventDestroy(me1); - qmoe_tune_record(g_qmoe_fc1_state, key1, s1, ms, "fc1", fusion_inter, - static_cast(hidden_size), - static_cast(block_size)); - } + auto fc1 = [&](QmoeCfg c, const void* indices) { + return dispatch_qmoe_fc1( + c.bs, c.tn, s, static_cast(k), fusion_inter, input, indices, + fc1_weights, fc1_scales, fc1_zero_points, fc1_bias, act_out, + static_cast(hidden_size), n_blk_in, static_cast(block_size), + swiglu_alpha, swiglu_beta, swiglu_limit); + }; + const QmoeCfg c1 = qmoe_choose( + g_qmoe_fc1_choice, + qmoe_cfg_key(fusion_inter, static_cast(hidden_size), + static_cast(block_size), fc1_zero_points != nullptr), + s, k, num_experts, "fc1", fc1); + if (hipError_t e = fc1(c1, expert_indices); e != hipSuccess) { + fprintf(stderr, "[custom_kernels] qmoe_decode_fc1 failed: %s\n", + hipGetErrorString(e)); + return static_cast(e); } - // FC2 config: N=hidden, K=inter. Online per-shape autotune. - int s2 = -1; - const uint64_t key2 = - qmoe_cfg_key(static_cast(hidden_size), static_cast(inter_size), - static_cast(block_size), fc2_zero_points != nullptr); - const QmoeCfg c2 = qmoe_tune_pick(g_qmoe_fc2_state, key2, &s2); - // FC2 + weighted write to slot_buf: grid (hidden/TN, k). No LDS preload. - { - hipEvent_t me0 = nullptr, me1 = nullptr; - if (s2 >= 0) { - hipEventCreate(&me0); - hipEventCreate(&me1); - hipEventRecord(me0, s); - } - hipError_t e = dispatch_qmoe_fc2( - c2.bs, c2.tn, s, static_cast(k), act_out, expert_indices, - expert_weights, fc2_weights, fc2_scales, fc2_zero_points, fc2_bias, - slot_buf, static_cast(inter_size), static_cast(hidden_size), - n_blk_mid, static_cast(block_size)); - if (e != hipSuccess) { - fprintf(stderr, "[custom_kernels] qmoe_decode_fc2 failed: %s\n", - hipGetErrorString(e)); - return static_cast(e); - } - if (s2 >= 0) { - hipEventRecord(me1, s); - hipEventSynchronize(me1); - float ms = 0; - hipEventElapsedTime(&ms, me0, me1); - hipEventDestroy(me0); - hipEventDestroy(me1); - qmoe_tune_record(g_qmoe_fc2_state, key2, s2, ms, "fc2", - static_cast(hidden_size), - static_cast(inter_size), - static_cast(block_size)); - } + auto fc2 = [&](QmoeCfg c, const void* indices) { + return dispatch_qmoe_fc2( + c.bs, c.tn, s, static_cast(k), act_out, indices, expert_weights, + fc2_weights, fc2_scales, fc2_zero_points, fc2_bias, slot_buf, + static_cast(inter_size), static_cast(hidden_size), n_blk_mid, + static_cast(block_size)); + }; + const QmoeCfg c2 = qmoe_choose( + g_qmoe_fc2_choice, + qmoe_cfg_key(static_cast(hidden_size), static_cast(inter_size), + static_cast(block_size), fc2_zero_points != nullptr), + s, k, num_experts, "fc2", fc2); + if (hipError_t e = fc2(c2, expert_indices); e != hipSuccess) { + fprintf(stderr, "[custom_kernels] qmoe_decode_fc2 failed: %s\n", + hipGetErrorString(e)); + return static_cast(e); } // Reduce: grid (hidden/kBlock). @@ -1966,8 +1961,6 @@ static hipError_t dispatch_qmoe_fc2_dp4a( // Separate online-autotune state for the dp4a kernels (their per-config // ranking differs from the fp path; sharing keys would cross-contaminate). -static std::unordered_map g_qmoe_fc1_dp4a_state; -static std::unordered_map g_qmoe_fc2_dp4a_state; extern "C" int hip_qmoe_decode_fused_dp4a( void* stream, @@ -1986,7 +1979,7 @@ extern "C" int hip_qmoe_decode_fused_dp4a( int64_t hidden_size, int64_t inter_size, int64_t k, int64_t block_size, float swiglu_alpha, float swiglu_beta, float swiglu_limit, - int64_t element_size_bytes) { + int64_t element_size_bytes, int64_t num_experts) { if (element_size_bytes != 2) { fprintf(stderr, "[custom_kernels] hip_qmoe_decode_fused_dp4a: only fp16 " "supported (got element_size %lld)\n", @@ -2008,6 +2001,9 @@ extern "C" int hip_qmoe_decode_fused_dp4a( } hipStream_t s = static_cast(stream); + // Each launch below reports hipGetLastError(); clear an error left by an + // earlier call so it is not taken for theirs, including the tiling timing. + (void)hipGetLastError(); const int fusion_inter = static_cast(2 * inter_size); const int n_blk_in = static_cast((hidden_size + block_size - 1) / block_size); const int n_blk_mid = static_cast((inter_size + block_size - 1) / block_size); @@ -2029,33 +2025,23 @@ extern "C" int hip_qmoe_decode_fused_dp4a( } } - // FC1 + SwiGLU (dp4a), online-autotuned. N=fusion_inter, K=hidden. - int s1 = -1; - const uint64_t key1 = + // FC1 + SwiGLU (dp4a). N=fusion_inter, K=hidden. + auto fc1 = [&](QmoeCfg c, const void* indices) { + return dispatch_qmoe_fc1_dp4a( + c.bs, c.tn, s, static_cast(k), fusion_inter, a_qb_in, a_scale_in, + indices, fc1_weights, fc1_scales, fc1_zero_points, fc1_bias, act_out, + static_cast(hidden_size), n_blk_in, static_cast(block_size), + swiglu_alpha, swiglu_beta, swiglu_limit); + }; + const QmoeCfg c1 = qmoe_choose( + g_qmoe_fc1_dp4a_choice, qmoe_cfg_key(fusion_inter, static_cast(hidden_size), - static_cast(block_size), fc1_zero_points != nullptr); - const QmoeCfg c1 = qmoe_tune_pick(g_qmoe_fc1_dp4a_state, key1, &s1); - { - hipEvent_t me0 = nullptr, me1 = nullptr; - if (s1 >= 0) { hipEventCreate(&me0); hipEventCreate(&me1); hipEventRecord(me0, s); } - hipError_t e = dispatch_qmoe_fc1_dp4a( - c1.bs, c1.tn, s, static_cast(k), fusion_inter, a_qb_in, - a_scale_in, expert_indices, fc1_weights, fc1_scales, fc1_zero_points, - fc1_bias, act_out, static_cast(hidden_size), n_blk_in, - static_cast(block_size), swiglu_alpha, swiglu_beta, swiglu_limit); - if (e != hipSuccess) { - fprintf(stderr, "[custom_kernels] qmoe_decode_fc1_dp4a failed: %s\n", - hipGetErrorString(e)); - return static_cast(e); - } - if (s1 >= 0) { - hipEventRecord(me1, s); hipEventSynchronize(me1); - float ms = 0; hipEventElapsedTime(&ms, me0, me1); - hipEventDestroy(me0); hipEventDestroy(me1); - qmoe_tune_record(g_qmoe_fc1_dp4a_state, key1, s1, ms, "fc1-dp4a", - fusion_inter, static_cast(hidden_size), - static_cast(block_size)); - } + static_cast(block_size), fc1_zero_points != nullptr), + s, k, num_experts, "fc1-dp4a", fc1); + if (hipError_t e = fc1(c1, expert_indices); e != hipSuccess) { + fprintf(stderr, "[custom_kernels] qmoe_decode_fc1_dp4a failed: %s\n", + hipGetErrorString(e)); + return static_cast(e); } // Quantize the k [inter] expert-slot activations (grid = n_blk_mid x k). @@ -2075,34 +2061,23 @@ extern "C" int hip_qmoe_decode_fused_dp4a( } } - // FC2 (dp4a) + routing-weighted write, online-autotuned. N=hidden, K=inter. - int s2 = -1; - const uint64_t key2 = + // FC2 (dp4a) + routing-weighted write. N=hidden, K=inter. + auto fc2 = [&](QmoeCfg c, const void* indices) { + return dispatch_qmoe_fc2_dp4a( + c.bs, c.tn, s, static_cast(k), a_qb_mid, a_scale_mid, indices, + expert_weights, fc2_weights, fc2_scales, fc2_zero_points, fc2_bias, + slot_buf, static_cast(inter_size), static_cast(hidden_size), + n_blk_mid, static_cast(block_size)); + }; + const QmoeCfg c2 = qmoe_choose( + g_qmoe_fc2_dp4a_choice, qmoe_cfg_key(static_cast(hidden_size), static_cast(inter_size), - static_cast(block_size), fc2_zero_points != nullptr); - const QmoeCfg c2 = qmoe_tune_pick(g_qmoe_fc2_dp4a_state, key2, &s2); - { - hipEvent_t me0 = nullptr, me1 = nullptr; - if (s2 >= 0) { hipEventCreate(&me0); hipEventCreate(&me1); hipEventRecord(me0, s); } - hipError_t e = dispatch_qmoe_fc2_dp4a( - c2.bs, c2.tn, s, static_cast(k), a_qb_mid, a_scale_mid, - expert_indices, expert_weights, fc2_weights, fc2_scales, - fc2_zero_points, fc2_bias, slot_buf, static_cast(inter_size), - static_cast(hidden_size), n_blk_mid, static_cast(block_size)); - if (e != hipSuccess) { - fprintf(stderr, "[custom_kernels] qmoe_decode_fc2_dp4a failed: %s\n", - hipGetErrorString(e)); - return static_cast(e); - } - if (s2 >= 0) { - hipEventRecord(me1, s); hipEventSynchronize(me1); - float ms = 0; hipEventElapsedTime(&ms, me0, me1); - hipEventDestroy(me0); hipEventDestroy(me1); - qmoe_tune_record(g_qmoe_fc2_dp4a_state, key2, s2, ms, "fc2-dp4a", - static_cast(hidden_size), - static_cast(inter_size), - static_cast(block_size)); - } + static_cast(block_size), fc2_zero_points != nullptr), + s, k, num_experts, "fc2-dp4a", fc2); + if (hipError_t e = fc2(c2, expert_indices); e != hipSuccess) { + fprintf(stderr, "[custom_kernels] qmoe_decode_fc2_dp4a failed: %s\n", + hipGetErrorString(e)); + return static_cast(e); } // Reduce: output[n] = sum over slot of slot_buf[slot, n]. diff --git a/lib/Runtime/Kernels/include/hip_custom_kernels.h b/lib/Runtime/Kernels/include/hip_custom_kernels.h index b5d5639581..bbcdcc1828 100644 --- a/lib/Runtime/Kernels/include/hip_custom_kernels.h +++ b/lib/Runtime/Kernels/include/hip_custom_kernels.h @@ -3137,6 +3137,10 @@ HIP_KERNEL_API int hip_qmoe_reduce_sorted_pairs( * * Constraints: fp16 only (element_size_bytes == 2); hidden_size and * inter_size both multiples of 32; block_size > 0 and even. + * + * num_experts (E) lets the first call of a shape time the GEMV tile + * configurations on other experts' weights before computing; every call + * then uses the configuration chosen, so equal inputs give equal results. */ HIP_KERNEL_API int hip_qmoe_decode_fused( void* stream, @@ -3153,7 +3157,7 @@ HIP_KERNEL_API int hip_qmoe_decode_fused( int64_t hidden_size, int64_t inter_size, int64_t k, int64_t block_size, float swiglu_alpha, float swiglu_beta, float swiglu_limit, - int64_t element_size_bytes); + int64_t element_size_bytes, int64_t num_experts); // W4A8 dp4a decode variant of hip_qmoe_decode_fused (env-gated via // HIPDNN_EP_MATMUL_DP4A). Additionally takes 4 scratch buffers: quantized @@ -3176,7 +3180,7 @@ HIP_KERNEL_API int hip_qmoe_decode_fused_dp4a( int64_t hidden_size, int64_t inter_size, int64_t k, int64_t block_size, float swiglu_alpha, float swiglu_beta, float swiglu_limit, - int64_t element_size_bytes); + int64_t element_size_bytes, int64_t num_experts); /* ========================================================================= * com.amd QMoE Sub-Kernels (LatentMoE) diff --git a/lib/Runtime/Kernels/unit_test/CMakeLists.txt b/lib/Runtime/Kernels/unit_test/CMakeLists.txt index 891cfddfd8..b4183e685c 100644 --- a/lib/Runtime/Kernels/unit_test/CMakeLists.txt +++ b/lib/Runtime/Kernels/unit_test/CMakeLists.txt @@ -150,6 +150,9 @@ foreach(_arch IN LISTS HIP_ARCHITECTURES) add_kernel_package_ut(${_arch} hipdnn-kernel-ut-gather-block-quantized-xi64-wfp8-yfp16 gather_block_quantized/Xi64_Wfp8_Yfp16/test_gather_block_quantized_fp8.cpp) + add_kernel_package_ut(${_arch} + hipdnn-kernel-ut-qmoe-decode-xfp16-wu4-yfp16 + qmoe_decode/Xfp16_Wu4_Yfp16/test_qmoe_decode.cpp) endforeach() # One package-facing entry point finds the live GPU arch, invokes every leaf diff --git a/lib/Runtime/Kernels/unit_test/Makefile b/lib/Runtime/Kernels/unit_test/Makefile index 60129475fe..db6d0e7c0d 100644 --- a/lib/Runtime/Kernels/unit_test/Makefile +++ b/lib/Runtime/Kernels/unit_test/Makefile @@ -43,7 +43,7 @@ LEAVES := gather_block_quantized/Xi64_Wfp8_Yfp16 \ matmul_nbits/Xfp16_Wu4_Yfp16 matmul_nbits/Xfp16_Wi8_Yfp16 \ matmul_nbits/Xfp32_Wu2_Yfp32 matmul_nbits/Xfp32_Wu3_Yfp32 \ matmul_nbits/Xfp32_Wu4_Yfp32 matmul_nbits/Xfp32_Wi8_Yfp32 \ - qmoe_ragged_mm/Xfp16_Wu4_Yfp16 + qmoe_decode/Xfp16_Wu4_Yfp16 qmoe_ragged_mm/Xfp16_Wu4_Yfp16 ARCH := $(patsubst --offload-arch=%,%,$(OFFLOAD)) HIPCC := $(HIP_SDK)/bin/hipcc.exe diff --git a/lib/Runtime/Kernels/unit_test/README.md b/lib/Runtime/Kernels/unit_test/README.md index fbec98a0ef..d3cfc41715 100644 --- a/lib/Runtime/Kernels/unit_test/README.md +++ b/lib/Runtime/Kernels/unit_test/README.md @@ -90,7 +90,7 @@ controller from that `bin` directory: | `--mode` | `lookup`, `autotune` | `lookup` | `lookup` uses the tables embedded in `custom_kernels_.dll` and falls through to autotune on a miss, like production. `autotune` bypasses the tables (production's `online` mode). | | `--coverage` | `1`, `2`, `3` | `3` | Shape tier; see "Shape coverage". `1` suits a per-PR job, `3` a nightly one. | -The controller runs the 16 leaves serially, with `bin` as the working +The controller runs the 17 leaves serially, with `bin` as the working directory so every leaf finds `custom_kernels_.dll`. It prints one `[/] PASS|FAIL ( ms)` line per leaf and finally `Combined report: `. @@ -193,7 +193,7 @@ Concretely: never drops a categorical situation. A tier3 (default) run of one leaf is on the order of tens to ~150 cases, and -the whole 16-leaf suite at tier3 is budgeted at **up to ~30 min total** +the whole 17-leaf suite at tier3 is budgeted at **up to ~30 min total** (measured ~20-25 min in practice) — this is intentionally wider than a "finishes in a minute" smoke test, because the shape axis (M=1 vs M>1, multiple K/N/seq scenarios) needs real coverage too, not just the categorical @@ -256,6 +256,7 @@ unit_test/ gqa/decode/Xfp16_Wi8_Yfp16/ gqa/prefill/Xfp16_Wfp16_Yfp16/ gqa/prefill/Xfp16_Wi8_Yfp16/ + qmoe_decode/Xfp16_Wu4_Yfp16/ qmoe_ragged_mm/Xfp16_Wu4_Yfp16/ gather_block_quantized/Xi64_W{u4,u8,fp8}_Yfp16/ ``` @@ -362,6 +363,12 @@ COVERAGE=N` (any leaf, or the top-level aggregator) builds and runs it. (`genCase()`'s `O(M*K*N)`, no BLAS) is now multithreaded over `N` to keep this affordable; see the comment above `genCase()` in any `matmul_nbits/*/test_matmul_nbits*.cpp`. +- **qmoe_decode** (the fused single-token decode): `{fp, dp4a}` x block + size `{32,64,128}` x zero points `{off,on}` x bias `{off,on}` (24 combos) + on a shape list COVERAGE thins (tier1 = 1 shape, tier2 = 2 with + GPT-OSS-20B, tier3 = 3), each checked against a CPU reference and then + for 128 more bit-identical calls. Every tier also runs a first call made + over a pending HIP error and the rejected arguments. - **qmoe_ragged_mm**: routing `{balanced, empty experts, skewed}` x input `{contiguous, FC1 gather}` x bias `{off,on}` x zero points `{off,on}` (24 combos) on a shape list COVERAGE thins (tier1 = 1 shape, tier2 = 2, tier3 = diff --git a/lib/Runtime/Kernels/unit_test/kernel_ut_main.cpp b/lib/Runtime/Kernels/unit_test/kernel_ut_main.cpp index d320add7e5..1412117f1a 100644 --- a/lib/Runtime/Kernels/unit_test/kernel_ut_main.cpp +++ b/lib/Runtime/Kernels/unit_test/kernel_ut_main.cpp @@ -62,6 +62,8 @@ static const Leaf kLeaves[] = { "hipdnn-kernel-ut-gather-block-quantized-xi64-wu8-yfp16.exe"}, {"gather_block_quantized", "Xi64_Wfp8_Yfp16", "hipdnn-kernel-ut-gather-block-quantized-xi64-wfp8-yfp16.exe"}, + {"qmoe_decode", "Xfp16_Wu4_Yfp16", + "hipdnn-kernel-ut-qmoe-decode-xfp16-wu4-yfp16.exe"}, }; static std::string executableDirectory() { diff --git a/lib/Runtime/Kernels/unit_test/qmoe_decode/Xfp16_Wu4_Yfp16/Makefile b/lib/Runtime/Kernels/unit_test/qmoe_decode/Xfp16_Wu4_Yfp16/Makefile new file mode 100644 index 0000000000..4bbdce7215 --- /dev/null +++ b/lib/Runtime/Kernels/unit_test/qmoe_decode/Xfp16_Wu4_Yfp16/Makefile @@ -0,0 +1,129 @@ +# ============================================================ +# custom_kernels qmoe_decode (fused single-token QMoE decode: fp16 input, +# 4-bit weights, fp16 output; fp and dp4a variants) Test +# +# make test OFFLOAD=--offload-arch=gfx1151 HIP_SDK= [COVERAGE=1|2|3] +# make test_custom OFFLOAD=... HIP_SDK=... ARGS="--coverage 1" +# make clean +# +# MODE is accepted for the uniform API but only labels out/results.csv: the +# decode picks its GEMV tiling itself, on the first call of each shape. +# +# `test` runs every categorical situation (fp / dp4a x block size x zero +# points x bias, a pending HIP error, rejected arguments) on the shapes +# COVERAGE selects (default 3; see unit_test/README.md "Shape coverage"). +# Inputs and the CPU reference are generated in-process -- no python, no data +# files. +# ============================================================ + +ifeq ($(OS),Windows_NT) + DETECTED_OS := Windows + SHELL := cmd.exe + .SHELLFLAGS := /c +else + DETECTED_OS := WSL +endif + +# ---- Required from CI / CLI ---- + +OFFLOAD ?= --offload-arch=gfx1151 +HIP_SDK ?= $(HIP_PATH) +MODE ?= lookup +COVERAGE ?= 3 +ARGS ?= + +ifeq ($(HIP_SDK),) + $(error pass HIP_SDK=) +endif + +# ---- Paths ---- + +HIPCC := $(HIP_SDK)/bin/hipcc.exe +INCLUDE_DIR := ../../../include +KERNEL_HIP := ../../../hip/qmoe_kernel.hip + +BUILD_DIR ?= out +TARGET_DIRECT := $(BUILD_DIR)/test_direct.exe +RESULTS_CSV := $(BUILD_DIR)/results.csv + +ARCH := $(patsubst --offload-arch=%,%,$(OFFLOAD)) + +# ---- Build settings ---- + +KERNEL_FLAGS := -O3 -std=c++17 -x hip -ffp-contract=fast +CXXFLAGS := -O3 -std=c++17 +INCLUDES := -I$(INCLUDE_DIR) + +# ---- OS-specific: run command + clean ---- + +ifeq ($(DETECTED_OS),Windows) + HIP_SDK_WIN := $(subst /,\,$(HIP_SDK)) + RUN_CMD = set "PATH=$(HIP_SDK_WIN)\bin;%PATH%" && set "HIPDNN_RESULTS_CSV=$(subst /,\,$(RESULTS_CSV))" && set "HIPDNN_RESULTS_OP=qmoe_decode" && set "HIPDNN_RESULTS_LEAF=Xfp16_Wu4_Yfp16" && set "HIPDNN_RESULTS_ARCH=$(ARCH)" && set "HIPDNN_RESULTS_MODE=$(MODE)" && $(subst /,\,.\$(1)) + define DO_CLEAN + if exist $(BUILD_DIR) rmdir /s /q $(BUILD_DIR) + endef +else + HIP_SDK_WIN := $(shell wslpath -w "$(HIP_SDK)") + RUN_CMD = cmd.exe /c "set PATH=$(HIP_SDK_WIN)\bin;%PATH% && set HIPDNN_RESULTS_CSV=$(RESULTS_CSV) && set HIPDNN_RESULTS_OP=qmoe_decode && set HIPDNN_RESULTS_LEAF=Xfp16_Wu4_Yfp16 && set HIPDNN_RESULTS_ARCH=$(ARCH) && set HIPDNN_RESULTS_MODE=$(MODE) && %CD%\$(1)" + define DO_CLEAN + rm -rf $(BUILD_DIR) + endef +endif + +CSV_FRESH ?= 1 +RESET_RESULTS_CSV = if exist $(subst /,\,$(RESULTS_CSV)) del $(subst /,\,$(RESULTS_CSV)) + +# ============================================================ +# Build +# ============================================================ + +ifeq ($(DETECTED_OS),Windows) +$(BUILD_DIR): + @if not exist $(BUILD_DIR) mkdir $(BUILD_DIR) +else +$(BUILD_DIR): + @mkdir -p $@ +endif + +DIRECT_TEST_OBJ := $(BUILD_DIR)/direct_test.obj +DIRECT_KERNEL_OBJ := $(BUILD_DIR)/direct_kernel.obj + +direct: FORCE | $(BUILD_DIR) + $(HIPCC) -c $(CXXFLAGS) $(OFFLOAD) $(INCLUDES) -x hip test_qmoe_decode.cpp -o $(DIRECT_TEST_OBJ) + $(HIPCC) -c $(KERNEL_FLAGS) $(OFFLOAD) $(INCLUDES) $(KERNEL_HIP) -o $(DIRECT_KERNEL_OBJ) + $(HIPCC) $(OFFLOAD) $(DIRECT_TEST_OBJ) $(DIRECT_KERNEL_OBJ) -o $(TARGET_DIRECT) + +# ============================================================ +# Run / test targets (uniform API) +# ============================================================ + +clean: + $(DO_CLEAN) + +# make test = build + run every categorical situation at $(COVERAGE) +# (default 3). No data files -- everything is in-process. +test: direct + @$(RESET_RESULTS_CSV) + $(call RUN_CMD,$(TARGET_DIRECT)) --coverage $(COVERAGE) + +# Pass the exe's own flags through ARGS= (e.g. --coverage 1). +test_custom: direct +ifeq ($(CSV_FRESH),1) + @$(RESET_RESULTS_CSV) +endif + $(call RUN_CMD,$(TARGET_DIRECT)) $(ARGS) + +# Build only. The unit_test/ aggregator passes PKG_DIR and PKG_NAME to collect +# the exe into its hipdnn-kernel-ut package layout. +build: direct +ifneq ($(PKG_DIR),) +ifeq ($(DETECTED_OS),Windows) + @copy /Y $(subst /,\,$(TARGET_DIRECT)) "$(subst /,\,$(PKG_DIR))\$(PKG_NAME)" >nul +else + @cp $(TARGET_DIRECT) $(PKG_DIR)/$(PKG_NAME) +endif +endif + +FORCE: + +.PHONY: test test_custom clean build direct FORCE diff --git a/lib/Runtime/Kernels/unit_test/qmoe_decode/Xfp16_Wu4_Yfp16/README.md b/lib/Runtime/Kernels/unit_test/qmoe_decode/Xfp16_Wu4_Yfp16/README.md new file mode 100644 index 0000000000..b9bde89944 --- /dev/null +++ b/lib/Runtime/Kernels/unit_test/qmoe_decode/Xfp16_Wu4_Yfp16/README.md @@ -0,0 +1,74 @@ +# qmoe_decode Xfp16_Wu4_Yfp16 standalone test + +Standalone `hipcc`-only test for the fused single-token QMoE decode, +`hip_qmoe_decode_fused` and its W4A8 variant `hip_qmoe_decode_fused_dp4a`: +fp16 input, 4-bit expert weights with fp16 scales, fp16 output. It does not +configure or build hip-ep, ONNX Runtime, LLVM, or MLIR. `make` compiles +`qmoe_kernel.hip` into the leaf; the CMake package build links +`custom_kernels_.dll` instead. Inputs and the CPU reference are +generated in-process -- no python, no data files. + +``` +make test OFFLOAD=--offload-arch=%GFX% HIP_SDK=%HIP_SDK% [COVERAGE=1|2|3] +make test_custom OFFLOAD=... HIP_SDK=... ARGS="--coverage 1" +make clean +``` + +| Target | Meaning | +|---|---| +| `test` | Coverage-tiered sweep -- see "Shape coverage" below. | +| `test_custom` | Runs the exe with `ARGS=` (e.g. `--coverage 1`). | +| `clean` | Removes `out/`. | + +## What it checks + +Every case runs FC1 + SwiGLU, FC2 and the weighted sum over k distinct +experts, and checks: + +- the first call against a double-precision reference that rounds to fp16 + where the kernels store (SwiGLU output, weighted slot, output) and, for + dp4a, quantizes the activations to per-group int8 the way the kernels do. + Each element must be within `4e-3` of itself plus `2e-3` of the largest + output (the k slots are fp16 before they are summed, so an element where + they cancel may be off by an ulp of the slots), and relL2 within `1e-3`; +- 128 more calls on the same inputs, each bit-identical to the first. The + first call of a shape picks the GEMV tiling for the process, and the block + size sets the reduction order, so a tiling that changed after the first + call would show here; +- output guard regions, and a NaN fill that fails any element left unwritten. + +Categorical situations, at every tier: fp and dp4a, block size 32, 64 and +128, zero points absent (8) and present, bias absent and present. Only block +sizes that are multiples of 32 are run: both GEMVs read one group's scale and +zero point per 32 values. + +Also at every tier: + +- the first call of an untimed shape made while an earlier failed HIP call's + error is pending, for fp and dp4a, must return 0 with a correct, repeatable + result. Without the decode clearing the error on entry, dp4a reports it as + its own; fp's first tiling launch consumes it, so fp still succeeds but + keeps the default tiling, which this test cannot see; +- an fp32 element size, an odd block size, a dp4a block size of 48, a hidden + size of 48 and k = 0 must return -1 and write nothing. + +## Shape coverage (3 tiers) + +`COVERAGE` only thins the shapes (hidden x inter, experts, k): tier1 = +`64 x 64`, 4 experts, k 2; tier2 adds GPT-OSS-20B (`2880 x 2880`, 32, 4); +tier3 (default) adds `2048 x 768`, 128, 8. Each shape runs all 24 categories. +At startup the exe prints `qmoe_decode coverage=N cases=` (32 / 56 / 80); +a successful run prints one `PASS` line per case and ends with `ALL PASS`. + +## `MODE` and `out/results.csv` + +The decode has no autotune table: it times its GEMV tiling candidates on the +first call of each shape. `MODE` is accepted for consistency with other +leaves and only fills the `mode` column. + +Each case appends one row to `out/results.csv` (see `unit_test/README.md`), +`shape` like `H2880_I2880_E32_k4`, `config` like `dp4a_bs32_zp_bias`, `relL2` +against the reference, and `time_ms` the mean of 50 back-to-back calls after +the repeat check, on warm caches (the same experts every call). + +CI passes `OFFLOAD`/`HIP_SDK` explicitly; there is no personal default. diff --git a/lib/Runtime/Kernels/unit_test/qmoe_decode/Xfp16_Wu4_Yfp16/test_qmoe_decode.cpp b/lib/Runtime/Kernels/unit_test/qmoe_decode/Xfp16_Wu4_Yfp16/test_qmoe_decode.cpp new file mode 100644 index 0000000000..46cfa68b1b --- /dev/null +++ b/lib/Runtime/Kernels/unit_test/qmoe_decode/Xfp16_Wu4_Yfp16/test_qmoe_decode.cpp @@ -0,0 +1,680 @@ +/* + * Copyright (C) 2026 Advanced Micro Devices, Inc. All rights reserved. + * Licensed under the MIT License. + * + * Standalone unit test for the fused single-token QMoE decode + * (hip_qmoe_decode_fused and its dp4a variant): fp16 activations, 4-bit + * weights, fp16 output. + */ + +#include "hip_custom_kernels.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +#define HIP_CHECK(call) \ + do { \ + const hipError_t error_ = (call); \ + if (error_ != hipSuccess) { \ + std::cerr << "HIP failure " << hipGetErrorString(error_) << " at " \ + << __FILE__ << ':' << __LINE__ << '\n'; \ + std::exit(2); \ + } \ + } while (false) + +constexpr float kSwigluAlpha = 1.702f; +constexpr float kSwigluBeta = 1.0f; +constexpr float kSwigluLimit = 7.0f; +constexpr int kDefaultZeroPoint = 8; +constexpr int kGuardElements = 32; +// The first call of a shape picks its GEMV tiling. If the choice were made by +// sampling candidates on real calls, it would take 12 candidates x 9 rounds of +// calls, with results that differ in the low bits; these many calls must all +// match the first bit for bit. +constexpr int kRepeatCalls = 129; +constexpr int kTimedCalls = 50; +// The kernel accumulates in fp32 in its own order; the reference in double, +// rounding to fp16 where the kernel stores (SwiGLU output, weighted slot, +// output). Each element may then differ by a few fp16 ulps of itself, and, +// since the k slots are stored as fp16 before they are summed, by an ulp of +// the slots where they cancel: hence a term scaled by the largest output. +constexpr double kRelTolerance = 4e-3; +constexpr double kScaleTolerance = 2e-3; +constexpr double kRelL2Tolerance = 1e-3; + +struct Shape { + int64_t hidden; + int64_t inter; + int64_t experts; + int64_t k; + int min_coverage; +}; + +constexpr Shape kShapes[] = { + {64, 64, 4, 2, 1}, + {2880, 2880, 32, 4, 2}, // GPT-OSS-20B + {2048, 768, 128, 8, 3}, +}; + +constexpr int64_t kBlockSizes[] = {32, 64, 128}; + +struct Category { + bool dp4a; + bool zero_points; + bool bias; + int64_t block_size; +}; + +class Rng { +public: + explicit Rng(uint64_t seed) : state_(seed) {} + uint64_t next() { + uint64_t z = (state_ += 0x9e3779b97f4a7c15ull); + z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ull; + z = (z ^ (z >> 27)) * 0x94d049bb133111ebull; + return z ^ (z >> 31); + } + float uniform(float lo, float hi) { + return lo + (hi - lo) * static_cast(next() >> 40) / + static_cast(1ull << 24); + } + +private: + uint64_t state_; +}; + +float to_float(__half h) { return __half2float(h); } +__half to_half(double v) { return __float2half(static_cast(v)); } + +// One expert matrix [E, N, K] of 4-bit weights with fp16 group scales, packed +// 4-bit zero points and an fp16 bias. +struct Matrix { + int64_t n = 0, k = 0, blocks = 0, zp_stride = 0; + std::vector weights; // [E, N, K/2], element k in nibble k&1 + std::vector<__half> scales; // [E, N, blocks] + std::vector zero_points; // [E, N, ceil(blocks/2)], group g in nibble g&1 + std::vector<__half> bias; // [E, N] + + int q(int64_t e, int64_t row, int64_t col) const { + const uint8_t b = weights[((e * n + row) * k + col) / 2]; + return (col & 1) ? (b >> 4) : (b & 0x0f); + } + int zp(int64_t e, int64_t row, int64_t group, bool on) const { + if (!on) { + return kDefaultZeroPoint; + } + const uint8_t b = zero_points[(e * n + row) * zp_stride + group / 2]; + return (group & 1) ? (b >> 4) : (b & 0x0f); + } + float scale(int64_t e, int64_t row, int64_t group) const { + return to_float(scales[(e * n + row) * blocks + group]); + } +}; + +Matrix make_matrix(int64_t experts, int64_t n, int64_t k, int64_t block_size, + float scale_lo, float scale_hi, Rng& rng) { + Matrix m; + m.n = n; + m.k = k; + m.blocks = (k + block_size - 1) / block_size; + m.zp_stride = (m.blocks + 1) / 2; + m.weights.resize(static_cast(experts * n * k / 2)); + for (uint8_t& b : m.weights) { + b = static_cast(rng.next()); + } + m.scales.resize(static_cast(experts * n * m.blocks)); + for (__half& s : m.scales) { + s = to_half(rng.uniform(scale_lo, scale_hi)); + } + m.zero_points.resize(static_cast(experts * n * m.zp_stride)); + for (uint8_t& b : m.zero_points) { + b = static_cast(rng.next()); + } + m.bias.resize(static_cast(experts * n)); + for (__half& b : m.bias) { + b = to_half(rng.uniform(-0.5f, 0.5f)); + } + return m; +} + +// Per-group int8 quantization of an activation row, as the dp4a path does it: +// scale = max|a| / 127 per group, q = round-to-nearest-even(a * 127 / max|a|) +// clamped to +-127. +void quantize_row(const std::vector& a, int64_t block_size, + std::vector& q, std::vector& scale) { + const int64_t blocks = (static_cast(a.size()) + block_size - 1) / block_size; + q.assign(a.size(), 0); + scale.assign(static_cast(blocks), 1.0f); + for (int64_t g = 0; g < blocks; ++g) { + const int64_t end = std::min((g + 1) * block_size, a.size()); + float gmax = 0.0f; + for (int64_t i = g * block_size; i < end; ++i) { + gmax = std::max(gmax, std::fabs(a[i])); + } + scale[g] = gmax > 0.0f ? gmax * (1.0f / 127.0f) : 1.0f; + const float inv = gmax > 0.0f ? 127.0f / gmax : 0.0f; + for (int64_t i = g * block_size; i < end; ++i) { + q[i] = std::max(-127, std::min(127, static_cast(std::nearbyint(a[i] * inv)))); + } + } +} + +// One output row of an expert GEMV: sum over k of a[k] * (w - zp) * scale. +double gemv(const Matrix& m, int64_t e, int64_t row, + const std::vector& a, const std::vector& aq, + const std::vector& a_scale, const Category& c) { + double acc = 0.0; + for (int64_t g = 0; g < m.blocks; ++g) { + const int zp = m.zp(e, row, g, c.zero_points); + const float s = m.scale(e, row, g); + const int64_t end = std::min((g + 1) * c.block_size, m.k); + if (c.dp4a) { + int64_t dot = 0; + for (int64_t i = g * c.block_size; i < end; ++i) { + dot += static_cast(aq[i]) * (m.q(e, row, i) - zp); + } + acc += static_cast(a_scale[g] * s) * static_cast(dot); + } else { + double dot = 0.0; + for (int64_t i = g * c.block_size; i < end; ++i) { + dot += static_cast(a[i]) * (m.q(e, row, i) - zp); + } + acc += dot * s; + } + } + return acc; +} + +struct Problem { + Shape shape; + Matrix fc1, fc2; + std::vector<__half> input; + std::vector indices; + std::vector<__half> weights; +}; + +Problem make_problem(const Shape& shape, int64_t block_size, uint64_t seed) { + Rng rng(seed); + Problem p; + p.shape = shape; + p.fc1 = make_matrix(shape.experts, 2 * shape.inter, shape.hidden, block_size, + 0.002f, 0.02f, rng); + p.fc2 = make_matrix(shape.experts, shape.hidden, shape.inter, block_size, + 0.002f, 0.02f, rng); + p.input.resize(static_cast(shape.hidden)); + for (__half& v : p.input) { + v = to_half(rng.uniform(-1.0f, 1.0f)); + } + // k distinct experts, as top-k routing gives. + std::vector pool(static_cast(shape.experts)); + for (int64_t e = 0; e < shape.experts; ++e) { + pool[e] = static_cast(e); + } + for (int64_t s = 0; s < shape.k; ++s) { + const size_t pick = s + rng.next() % (pool.size() - s); + std::swap(pool[s], pool[pick]); + p.indices.push_back(pool[s]); + p.weights.push_back(to_half(rng.uniform(0.1f, 0.5f))); + } + return p; +} + +std::vector<__half> reference(const Problem& p, const Category& c) { + const Shape& sh = p.shape; + std::vector x(static_cast(sh.hidden)); + for (int64_t i = 0; i < sh.hidden; ++i) { + x[i] = to_float(p.input[i]); + } + std::vector xq; + std::vector x_scale; + if (c.dp4a) { + quantize_row(x, c.block_size, xq, x_scale); + } + std::vector out(static_cast(sh.hidden), 0.0); + for (int64_t s = 0; s < sh.k; ++s) { + const int64_t e = p.indices[s]; + // FC1 outputs pair (gate, linear) at (2j, 2j+1); SwiGLU, stored as fp16. + std::vector act(static_cast(sh.inter)); + for (int64_t j = 0; j < sh.inter; ++j) { + double g = gemv(p.fc1, e, 2 * j, x, xq, x_scale, c); + double l = gemv(p.fc1, e, 2 * j + 1, x, xq, x_scale, c); + if (c.bias) { + g += to_float(p.fc1.bias[e * p.fc1.n + 2 * j]); + l += to_float(p.fc1.bias[e * p.fc1.n + 2 * j + 1]); + } + const double G = std::min(g, static_cast(kSwigluLimit)); + const double L = std::max(std::min(l, static_cast(kSwigluLimit)), + -static_cast(kSwigluLimit)); + const double sig = 1.0 / (1.0 + std::exp(-kSwigluAlpha * G)); + act[j] = to_float(to_half(G * sig * (L + kSwigluBeta))); + } + std::vector aq; + std::vector a_scale; + if (c.dp4a) { + quantize_row(act, c.block_size, aq, a_scale); + } + const double w = to_float(p.weights[s]); + // FC2, scaled by the routing weight and stored as fp16 per slot. + for (int64_t n = 0; n < sh.hidden; ++n) { + double v = gemv(p.fc2, e, n, act, aq, a_scale, c); + if (c.bias) { + v += to_float(p.fc2.bias[e * p.fc2.n + n]); + } + out[n] += to_float(to_half(w * v)); + } + } + std::vector<__half> result(out.size()); + for (size_t i = 0; i < out.size(); ++i) { + result[i] = to_half(out[i]); + } + return result; +} + +template class DeviceBuffer { +public: + explicit DeviceBuffer(size_t count) { + if (count != 0) { + HIP_CHECK(hipMalloc(reinterpret_cast(&data_), count * sizeof(T))); + } + } + ~DeviceBuffer() { + if (data_) { + (void)hipFree(data_); + } + } + DeviceBuffer(const DeviceBuffer&) = delete; + DeviceBuffer& operator=(const DeviceBuffer&) = delete; + T* get() { return data_; } + +private: + T* data_ = nullptr; +}; + +template void upload(DeviceBuffer& d, const std::vector& h) { + HIP_CHECK(hipMemcpy(d.get(), h.data(), h.size() * sizeof(T), + hipMemcpyHostToDevice)); +} + +// Device copies of a problem plus the call's scratch and a guarded output. +struct Device { + explicit Device(const Problem& p, int64_t block_size) + : input(p.input.size()), indices(p.indices.size()), + weights(p.weights.size()), fc1_w(p.fc1.weights.size()), + fc1_s(p.fc1.scales.size()), fc1_z(p.fc1.zero_points.size()), + fc1_b(p.fc1.bias.size()), fc2_w(p.fc2.weights.size()), + fc2_s(p.fc2.scales.size()), fc2_z(p.fc2.zero_points.size()), + fc2_b(p.fc2.bias.size()), + slot(static_cast(p.shape.k * p.shape.hidden)), + act(static_cast(p.shape.k * p.shape.inter)), + output(static_cast(p.shape.hidden + 2 * kGuardElements)), + a_qb_in(static_cast(p.shape.hidden)), + a_scale_in(static_cast(p.shape.hidden / block_size + 1)), + a_qb_mid(static_cast(p.shape.k * p.shape.inter)), + a_scale_mid(static_cast(p.shape.k * (p.shape.inter / block_size + 1))) { + upload(input, p.input); + upload(indices, p.indices); + upload(weights, p.weights); + upload(fc1_w, p.fc1.weights); + upload(fc1_s, p.fc1.scales); + upload(fc1_z, p.fc1.zero_points); + upload(fc1_b, p.fc1.bias); + upload(fc2_w, p.fc2.weights); + upload(fc2_s, p.fc2.scales); + upload(fc2_z, p.fc2.zero_points); + upload(fc2_b, p.fc2.bias); + } + DeviceBuffer<__half> input; + DeviceBuffer indices; + DeviceBuffer<__half> weights; + DeviceBuffer fc1_w; + DeviceBuffer<__half> fc1_s; + DeviceBuffer fc1_z; + DeviceBuffer<__half> fc1_b; + DeviceBuffer fc2_w; + DeviceBuffer<__half> fc2_s; + DeviceBuffer fc2_z; + DeviceBuffer<__half> fc2_b; + DeviceBuffer<__half> slot; + DeviceBuffer<__half> act; + DeviceBuffer<__half> output; + DeviceBuffer a_qb_in; + DeviceBuffer a_scale_in; + DeviceBuffer a_qb_mid; + DeviceBuffer a_scale_mid; +}; + +int launch(Device& d, const Shape& sh, const Category& c, hipStream_t stream, + int64_t element_size = 2, int64_t hidden = -1, int64_t k = -1) { + hidden = hidden < 0 ? sh.hidden : hidden; + k = k < 0 ? sh.k : k; + const void* z1 = c.zero_points ? d.fc1_z.get() : nullptr; + const void* z2 = c.zero_points ? d.fc2_z.get() : nullptr; + const void* b1 = c.bias ? d.fc1_b.get() : nullptr; + const void* b2 = c.bias ? d.fc2_b.get() : nullptr; + __half* out = d.output.get() + kGuardElements; + if (c.dp4a) { + return hip_qmoe_decode_fused_dp4a( + reinterpret_cast(stream), d.input.get(), d.indices.get(), + d.weights.get(), d.fc1_w.get(), d.fc1_s.get(), z1, b1, d.fc2_w.get(), + d.fc2_s.get(), z2, b2, d.slot.get(), d.act.get(), out, d.a_qb_in.get(), + d.a_scale_in.get(), d.a_qb_mid.get(), d.a_scale_mid.get(), hidden, + sh.inter, k, c.block_size, kSwigluAlpha, kSwigluBeta, kSwigluLimit, + element_size, sh.experts); + } + return hip_qmoe_decode_fused( + reinterpret_cast(stream), d.input.get(), d.indices.get(), + d.weights.get(), d.fc1_w.get(), d.fc1_s.get(), z1, b1, d.fc2_w.get(), + d.fc2_s.get(), z2, b2, d.slot.get(), d.act.get(), out, hidden, sh.inter, + k, c.block_size, kSwigluAlpha, kSwigluBeta, kSwigluLimit, element_size, + sh.experts); +} + +std::vector read_output(Device& d, size_t count) { + std::vector bits(count + 2 * kGuardElements); + HIP_CHECK(hipMemcpy(bits.data(), d.output.get(), bits.size() * sizeof(uint16_t), + hipMemcpyDeviceToHost)); + return bits; +} + +// A NaN no case produces, so an element left unwritten fails the compare. +constexpr uint16_t kFillBits = 0x7fff; + +void fill_output(Device& d, size_t count) { + const std::vector fill(count + 2 * kGuardElements, kFillBits); + HIP_CHECK(hipMemcpy(d.output.get(), fill.data(), fill.size() * sizeof(uint16_t), + hipMemcpyHostToDevice)); +} + +bool guards_intact(const std::vector& bits, size_t count) { + for (int i = 0; i < kGuardElements; ++i) { + if (bits[i] != kFillBits || bits[kGuardElements + count + i] != kFillBits) { + return false; + } + } + return true; +} + +// Tiny out/results.csv writer (op,leaf,arch,mode,shape,config,time_ms,relL2, +// verdict), opt-in via HIPDNN_RESULTS_CSV (set by the Makefile). +struct CsvRow { + std::string shape, config; + double time_ms = 0.0; + double rel_l2 = 0.0; + std::string verdict; +}; +class CsvWriter { +public: + CsvWriter() { + const char* path = std::getenv("HIPDNN_RESULTS_CSV"); + if (!path || !path[0]) { + return; + } + path_ = path; + bool need_header = true; + if (FILE* probe = std::fopen(path_.c_str(), "rb")) { + std::fseek(probe, 0, SEEK_END); + need_header = std::ftell(probe) == 0; + std::fclose(probe); + } + f_ = std::fopen(path_.c_str(), "a"); + if (f_ && need_header) { + std::fprintf(f_, "op,leaf,arch,mode,shape,config,time_ms,relL2,verdict\n"); + std::fflush(f_); + } + } + ~CsvWriter() { + if (f_) { + std::fclose(f_); + } + } + void write(const CsvRow& r) { + if (!f_) { + return; + } + std::fprintf(f_, "%s,%s,%s,%s,%s,%s,%.6f,%.6e,%s\n", + env_or("HIPDNN_RESULTS_OP", "qmoe_decode"), + env_or("HIPDNN_RESULTS_LEAF", "unknown"), + env_or("HIPDNN_RESULTS_ARCH", "unknown"), + env_or("HIPDNN_RESULTS_MODE", "unknown"), r.shape.c_str(), + r.config.c_str(), r.time_ms, r.rel_l2, r.verdict.c_str()); + std::fflush(f_); + } + +private: + static const char* env_or(const char* name, const char* dflt) { + const char* v = std::getenv(name); + return (v && v[0]) ? v : dflt; + } + std::string path_; + FILE* f_ = nullptr; +}; + +// One case's results.csv row, written when it goes out of scope so that a +// case returning early still reports FAIL. +struct ResultRow { + CsvRow row; + ResultRow(std::string shape, std::string config) { + row.shape = std::move(shape); + row.config = std::move(config); + row.rel_l2 = std::numeric_limits::infinity(); + row.verdict = "FAIL"; + } + ResultRow(const ResultRow&) = delete; + ResultRow& operator=(const ResultRow&) = delete; + ~ResultRow() { CsvWriter().write(row); } + void finish(double time_ms, double rel_l2, bool pass) { + row.time_ms = time_ms; + row.rel_l2 = rel_l2; + row.verdict = pass ? "PASS" : "FAIL"; + } +}; + +std::string shape_name(const Shape& sh) { + return "H" + std::to_string(sh.hidden) + "_I" + std::to_string(sh.inter) + + "_E" + std::to_string(sh.experts) + "_k" + std::to_string(sh.k); +} + +std::string category_name(const Category& c) { + return std::string(c.dp4a ? "dp4a" : "fp") + "_bs" + + std::to_string(c.block_size) + (c.zero_points ? "_zp" : "") + + (c.bias ? "_bias" : ""); +} + +double mean_call_ms(Device& d, const Shape& sh, const Category& c, + hipStream_t stream) { + hipEvent_t start = nullptr, stop = nullptr; + HIP_CHECK(hipEventCreate(&start)); + HIP_CHECK(hipEventCreate(&stop)); + HIP_CHECK(hipEventRecord(start, stream)); + for (int i = 0; i < kTimedCalls; ++i) { + (void)launch(d, sh, c, stream); + } + HIP_CHECK(hipEventRecord(stop, stream)); + HIP_CHECK(hipEventSynchronize(stop)); + float ms = 0.0f; + HIP_CHECK(hipEventElapsedTime(&ms, start, stop)); + HIP_CHECK(hipEventDestroy(start)); + HIP_CHECK(hipEventDestroy(stop)); + return static_cast(ms) / kTimedCalls; +} + +// First call against the reference, then kRepeatCalls - 1 more that must match +// it bit for bit, then the mean steady-state call time. +bool run_case(const Shape& sh, const Category& c, uint64_t seed, + hipStream_t stream, const char* label = nullptr) { + const std::string name = (label ? std::string(label) + "_" : std::string()) + + shape_name(sh); + const std::string config = category_name(c); + ResultRow result(name, config); + const Problem p = make_problem(sh, c.block_size, seed); + const std::vector<__half> expected = reference(p, c); + Device d(p, c.block_size); + const size_t count = static_cast(sh.hidden); + fill_output(d, count); + + const int status = launch(d, sh, c, stream); + HIP_CHECK(hipStreamSynchronize(stream)); + const std::vector first = read_output(d, count); + const bool guards = guards_intact(first, count); + + double err_sq = 0.0, ref_sq = 0.0, max_abs = 0.0, ref_max = 0.0; + int64_t mismatches = 0; + for (const __half& h : expected) { + ref_max = std::max(ref_max, std::abs(static_cast(to_float(h)))); + } + for (size_t i = 0; i < count; ++i) { + uint16_t bits = first[kGuardElements + i]; + __half h; + std::memcpy(&h, &bits, sizeof(h)); + const double got = to_float(h); + const double want = to_float(expected[i]); + const double delta = std::isnan(got) ? std::numeric_limits::infinity() + : std::abs(got - want); + max_abs = std::max(max_abs, delta); + mismatches += + delta <= kRelTolerance * std::abs(want) + kScaleTolerance * ref_max ? 0 : 1; + err_sq += delta * delta; + ref_sq += want * want; + } + const double rel_l2 = std::sqrt(err_sq / std::max(ref_sq, 1e-30)); + const bool accurate = rel_l2 <= kRelL2Tolerance && mismatches == 0; + + int differing_call = -1; + for (int call = 1; call < kRepeatCalls && status == 0; ++call) { + (void)launch(d, sh, c, stream); + HIP_CHECK(hipStreamSynchronize(stream)); + if (read_output(d, count) != first) { + differing_call = call; + break; + } + } + const bool pass = status == 0 && guards && accurate && differing_call < 0; + const double ms = pass ? mean_call_ms(d, sh, c, stream) : 0.0; + + std::cout << (pass ? "PASS" : "FAIL") << " shape=" << name + << " config=" << config << " status=" << status + << " relL2=" << rel_l2 << " maxAbs=" << max_abs + << " mismatches=" << mismatches << " guards=" << (guards ? "PASS" : "FAIL") << " repeat=" + << (differing_call < 0 ? std::string("identical") + : "call " + std::to_string(differing_call) + + " differs") + << " time_ms=" << ms << '\n'; + result.finish(ms, rel_l2, pass); + return pass; +} + +// A failed HIP call earlier on this thread leaves its error pending. The first +// call of a shape times its tiling candidates before computing; a stale error +// must neither fail that timing nor be reported as the call's own. +bool run_pending_error(bool dp4a, uint64_t seed, hipStream_t stream) { + // A shape no other case uses, so this call is the one that picks the tiling. + const Shape sh{96, 160, 6, 2, 1}; + (void)hipSetDevice(-1); + if (hipPeekAtLastError() == hipSuccess) { + std::cout << "FAIL pending_error: hipSetDevice(-1) left no pending error\n"; + ResultRow result("pending_error_" + shape_name(sh), dp4a ? "dp4a" : "fp"); + return false; + } + return run_case(sh, {dp4a, true, true, 32}, seed, stream, "pending_error"); +} + +// Arguments the launchers must refuse with -1 and without writing anything. +bool run_rejected(const char* label, const Category& c, int64_t element_size, + int64_t hidden, int64_t k, hipStream_t stream) { + ResultRow result(std::string("rejected_") + label, category_name(c)); + const Shape sh{64, 64, 4, 2, 1}; + const Problem p = make_problem(sh, 32, 1); + Device d(p, 32); + fill_output(d, static_cast(sh.hidden)); + const int status = launch(d, sh, c, stream, element_size, hidden, k); + HIP_CHECK(hipStreamSynchronize(stream)); + const std::vector bits = read_output(d, static_cast(sh.hidden)); + const bool untouched = + std::all_of(bits.begin(), bits.end(), [](uint16_t b) { return b == kFillBits; }); + const bool pass = status == -1 && untouched; + std::cout << (pass ? "PASS" : "FAIL") << " rejected=" << label + << " config=" << category_name(c) << " status=" << status + << " output=" << (untouched ? "untouched" : "WRITTEN") << '\n'; + result.finish(0.0, pass ? 0.0 : 1.0, pass); + return pass; +} + +bool parse_options(int argc, char** argv, int* coverage) { + for (int i = 1; i < argc; ++i) { + const std::string argument = argv[i]; + if (argument == "--coverage" && i + 1 < argc) { + *coverage = std::atoi(argv[++i]); + } else { + std::cerr << "usage: " << argv[0] << " [--coverage 1|2|3]\n"; + return false; + } + } + return *coverage >= 1 && *coverage <= 3; +} + +} // namespace + +int main(int argc, char** argv) { + int coverage = 3; + if (!parse_options(argc, argv, &coverage)) { + return 2; + } + + int shape_count = 0; + for (const Shape& sh : kShapes) { + shape_count += sh.min_coverage <= coverage ? 1 : 0; + } + constexpr int kCategoryCount = 2 * 2 * 2 * 3; + constexpr int kExtraCount = 2 + 6; + std::cout << "qmoe_decode coverage=" << coverage + << " cases=" << shape_count * kCategoryCount + kExtraCount << '\n'; + + hipStream_t stream = nullptr; + HIP_CHECK(hipStreamCreate(&stream)); + bool all_pass = true; + uint64_t seed = 0; + // The pending-error cases first, while their shape is still untimed. + all_pass = run_pending_error(false, ++seed, stream) && all_pass; + all_pass = run_pending_error(true, ++seed, stream) && all_pass; + for (const Shape& sh : kShapes) { + ++seed; + if (sh.min_coverage > coverage) { + continue; + } + for (bool dp4a : {false, true}) { + for (int64_t block_size : kBlockSizes) { + for (bool zero_points : {false, true}) { + for (bool bias : {false, true}) { + all_pass = run_case(sh, {dp4a, zero_points, bias, block_size}, seed, + stream) && + all_pass; + } + } + } + } + } + all_pass = run_rejected("fp32", {false, false, false, 32}, 4, -1, -1, stream) && all_pass; + all_pass = run_rejected("odd_block", {false, false, false, 15}, 2, -1, -1, stream) && all_pass; + all_pass = run_rejected("dp4a_block48", {true, false, false, 48}, 2, -1, -1, stream) && all_pass; + all_pass = run_rejected("hidden48", {false, false, false, 32}, 2, 48, -1, stream) && all_pass; + all_pass = run_rejected("k0", {false, false, false, 32}, 2, -1, 0, stream) && all_pass; + all_pass = run_rejected("dp4a_k0", {true, false, false, 32}, 2, -1, 0, stream) && all_pass; + HIP_CHECK(hipStreamDestroy(stream)); + + std::cout << (all_pass ? "ALL PASS" : "SOME FAIL") << '\n'; + return all_pass ? 0 : 1; +} diff --git a/lib/Runtime/real/qmoe.cpp b/lib/Runtime/real/qmoe.cpp index daa806fadd..3e3f3d2802 100644 --- a/lib/Runtime/real/qmoe.cpp +++ b/lib/Runtime/real/qmoe.cpp @@ -273,7 +273,7 @@ int wrap_qmoe(RuntimeState *state, const void *input, const void *router_probs, fc2_zero_points, fc2_bias, d_fc2_buf, d_act_buf, output, d_a_qb_in, d_a_scale_in, d_a_qb_mid, d_a_scale_mid, hidden_size, inter_size, k, block_size, activation_alpha, activation_beta, swiglu_limit, - elem_size)); + elem_size, num_experts)); return 0; } RUNTIME_DEBUG_LOG("[REAL] wrap_qmoe: fused decode path (k=%lld)\n", @@ -283,7 +283,7 @@ int wrap_qmoe(RuntimeState *state, const void *input, const void *router_probs, fc1_scales, fc1_zero_points, fc1_bias, fc2_weights, fc2_scales, fc2_zero_points, fc2_bias, d_fc2_buf, d_act_buf, output, hidden_size, inter_size, k, block_size, activation_alpha, activation_beta, - swiglu_limit, elem_size)); + swiglu_limit, elem_size, num_experts)); return 0; } diff --git a/test/numeric/tests/test_qmoe.py b/test/numeric/tests/test_qmoe.py index 9cffee69e2..8fe623bc9b 100644 --- a/test/numeric/tests/test_qmoe.py +++ b/test/numeric/tests/test_qmoe.py @@ -343,6 +343,67 @@ def test_qmoe_gpt_oss_shape_mix_weights(self, model_runner, seq_len, normalize): compare_outputs(actual, expected, atol=1e-1, rtol=1e-2, cos_threshold=0.99) +# More decode calls than the fused decode would spend choosing a tiling if it +# chose on real calls (12 candidates x 9 rounds). +DETERMINISM_CALLS = 129 + + +class TestQMoEDecodeDeterminism: + """The fused decode must give the same output for the same input on every + call, including the first ones of a shape. Its GEMV tiling sets the + reduction order, so a tiling chosen by sampling candidates on real calls + would make those calls differ in the low bits. + + The tiling choice is cached per shape for the process, so the case uses a + shape no other test here uses. The reduction order only shows in the + fp16 output for some shapes; on gfx1151 a tiling sampled on real calls + made about 80 of the first 129 calls differ from the first for this one. + + Only the dp4a decode (the default) is reachable here: with 32-aligned + blocks and sizes the EP takes it, and the fp decode is the fallback for + other block sizes. The kernel unit test qmoe_decode/Xfp16_Wu4_Yfp16 runs + both. + """ + + @pytest.mark.parametrize( + "hidden,intermediate,num_experts,top_k,block_size", + [ + # 32-aligned blocks and sizes: the dp4a fused decode (the default). + (HIDDEN, HIDDEN, 16, 2, 32), + ], + ) + def test_qmoe_decode_repeatable( + self, model_runner, hidden, intermediate, num_experts, top_k, block_size + ): + backend = model_runner.backend + if not hasattr(backend, "create_session"): + pytest.skip("backend does not expose an ORT session") + model = _make_qmoe_model( + 1, 1, hidden, intermediate, num_experts, top_k, block_size + ) + sub = model_runner._next_subdir() + path = sub / "model.onnx" + path.write_bytes(model.SerializeToString()) + sess = backend.create_session(str(path)) + + rng = np.random.default_rng(7) + x = rng.uniform(-1, 1, [1, 1, hidden]).astype(np.float16) + router = (rng.standard_normal([1, num_experts]) * 1.2).astype(np.float16) + feeds = {i.name: v for i, v in zip(sess.get_inputs(), [x, router])} + + first = sess.run(None, feeds)[0] + assert np.isfinite(first).all() + for call in range(1, DETERMINISM_CALLS): + out = sess.run(None, feeds)[0] + diff = np.flatnonzero(out.view(np.uint16) != first.view(np.uint16)) + assert diff.size == 0, ( + f"call {call} differs from call 0 at {diff.size} of {out.size} " + f"values (max |diff| " + f"{np.abs(out.astype(np.float32) - first.astype(np.float32)).max()})" + ) + del sess + + @dataclass(frozen=True) class Relu2Config: seq_len: int