From 9ec287c8668422e5b6349a1e971420d16eb056ee Mon Sep 17 00:00:00 2001 From: PanZezhong Date: Mon, 6 Jul 2026 02:03:05 +0000 Subject: [PATCH] fix: moe topk softmax test script and fix bias --- src/infinicore/pybind11/ops.hpp | 2 + .../pybind11/ops/moe_topk_softmax.hpp | 57 ++++ .../nvidia/moe_topk_softmax_nvidia.cu | 110 +++---- test/infinicore/ops/moe_topk_softmax.py | 284 ++++++++++++++++++ 4 files changed, 386 insertions(+), 67 deletions(-) create mode 100644 src/infinicore/pybind11/ops/moe_topk_softmax.hpp create mode 100644 test/infinicore/ops/moe_topk_softmax.py diff --git a/src/infinicore/pybind11/ops.hpp b/src/infinicore/pybind11/ops.hpp index 582008fcd..38aceb2fd 100644 --- a/src/infinicore/pybind11/ops.hpp +++ b/src/infinicore/pybind11/ops.hpp @@ -82,6 +82,7 @@ #include "ops/mha.hpp" #include "ops/mha_kvcache.hpp" #include "ops/mha_varlen.hpp" +#include "ops/moe_topk_softmax.hpp" #include "ops/mrope.hpp" #include "ops/mul.hpp" #include "ops/multi_margin_loss.hpp" @@ -189,6 +190,7 @@ inline void bind(py::module &m) { bind_mha_kvcache(m); bind_mha_varlen(m); bind_mha(m); + bind_moe_topk_softmax(m); bind_mrope(m); bind_hardswish(m); bind_hardtanh(m); diff --git a/src/infinicore/pybind11/ops/moe_topk_softmax.hpp b/src/infinicore/pybind11/ops/moe_topk_softmax.hpp new file mode 100644 index 000000000..5466fcd79 --- /dev/null +++ b/src/infinicore/pybind11/ops/moe_topk_softmax.hpp @@ -0,0 +1,57 @@ +#pragma once + +#include + +#include "infinicore/ops/moe_topk_softmax.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +std::tuple py_moe_topk_softmax(Tensor gating_output, + size_t topk, + bool renormalize, + float moe_softcapping, + py::object correction_bias) { + Tensor bias; + if (!correction_bias.is_none()) { + bias = correction_bias.cast(); + } + return op::moe_topk_softmax(gating_output, topk, renormalize, moe_softcapping, bias); +} + +void py_moe_topk_softmax_(Tensor topk_weights, + Tensor topk_indices, + Tensor gating_output, + py::object correction_bias, + bool renormalize, + float moe_softcapping) { + Tensor bias; + if (!correction_bias.is_none()) { + bias = correction_bias.cast(); + } + op::moe_topk_softmax_(topk_weights, topk_indices, gating_output, bias, renormalize, moe_softcapping); +} + +inline void bind_moe_topk_softmax(py::module &m) { + m.def("moe_topk_softmax", + &py_moe_topk_softmax, + py::arg("gating_output"), + py::arg("topk"), + py::arg("renormalize") = false, + py::arg("moe_softcapping") = 0.0f, + py::arg("correction_bias") = py::none(), + R"doc(MoE top-k softmax.)doc"); + + m.def("moe_topk_softmax_", + &py_moe_topk_softmax_, + py::arg("topk_weights"), + py::arg("topk_indices"), + py::arg("gating_output"), + py::arg("correction_bias") = py::none(), + py::arg("renormalize") = false, + py::arg("moe_softcapping") = 0.0f, + R"doc(In-place MoE top-k softmax.)doc"); +} + +} // namespace infinicore::ops diff --git a/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu b/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu index 2fd301c66..48842e144 100644 --- a/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu +++ b/src/infiniop/ops/moe_topk_softmax/nvidia/moe_topk_softmax_nvidia.cu @@ -62,20 +62,15 @@ struct MaxReduceOp { template __launch_bounds__(TPB) __global__ void moeSoftmax( const T *input, - const bool *finished, float *output, const int num_cols, - const float moe_softcapping, - const float *correction_bias) { + const float moe_softcapping) { using BlockReduce = cub::BlockReduce; __shared__ typename BlockReduce::TempStorage tmp_storage; __shared__ float normalizing_factor; __shared__ float float_max; const int thread_row_offset = blockIdx.x * num_cols; - if ((finished != nullptr) && finished[blockIdx.x]) { - return; - } float thread_data = -FLT_MAX; for (int ii = threadIdx.x; ii < num_cols; ii += TPB) { @@ -84,9 +79,6 @@ __launch_bounds__(TPB) __global__ void moeSoftmax( if (moe_softcapping != 0.0f) { val = tanhf(val / moe_softcapping) * moe_softcapping; } - if (correction_bias != nullptr) { - val += correction_bias[ii]; - } output[idx] = val; thread_data = fmaxf(val, thread_data); } @@ -151,21 +143,18 @@ struct TopKPairArgMax { template __launch_bounds__(TPB) __global__ void moeTopKFast( float *inputs_after_softmax, - const bool *finished, float *output, int *indices, const int num_experts, const int k, - const int start_expert, - const int end_expert, - const bool renormalize) { + const bool renormalize, + const float *correction_bias) { using namespace moe; using BlockReduce = cub::BlockReduce; __shared__ typename BlockReduce::TempStorage tmp_storage; TopKPair thread_pair; const int block_row = blockIdx.x; - const bool row_is_active = finished ? !finished[block_row] : true; const int thread_read_offset = blockIdx.x * num_experts; float row_sum_for_renormalize = 0.0f; @@ -179,7 +168,8 @@ __launch_bounds__(TPB) __global__ void moeTopKFast( for (int expert = threadIdx.x; expert < num_experts; expert += TPB) { const int idx = thread_read_offset + expert; inp_kvp.key = expert; - inp_kvp.value = inputs_after_softmax[idx]; + const float prob = inputs_after_softmax[idx]; + inp_kvp.value = correction_bias == nullptr ? prob : prob + correction_bias[expert]; if (inp_kvp.value > thread_pair.max.value) { thread_pair.secondMax = thread_pair.max; thread_pair.max = inp_kvp; @@ -198,13 +188,12 @@ __launch_bounds__(TPB) __global__ void moeTopKFast( } cub_kvp result = (i == TopKPair::MAX_INDEX) ? result_pair.max : result_pair.secondMax; int expert = result.key; - bool node_uses_expert = expert >= start_expert && expert < end_expert; - bool should_process_row = row_is_active && node_uses_expert; - inputs_after_softmax[thread_read_offset + expert] = -1.0f; + const float prob = inputs_after_softmax[thread_read_offset + expert]; + inputs_after_softmax[thread_read_offset + expert] = -FLT_MAX; int idx = k * block_row + k_idx * 2 + i; - output[idx] = result.value; - indices[idx] = should_process_row ? (expert - start_expert) : num_experts; - row_sum_for_renormalize += result.value; + output[idx] = prob; + indices[idx] = expert; + row_sum_for_renormalize += prob; } } __syncthreads(); @@ -222,14 +211,12 @@ __launch_bounds__(TPB) __global__ void moeTopKFast( template __launch_bounds__(TPB) __global__ void moeTopK( float *inputs_after_softmax, - const bool *finished, float *output, int *indices, const int num_experts, const int k, - const int start_expert, - const int end_expert, - const bool renormalize) { + const bool renormalize, + const float *correction_bias) { using cub_kvp = cub::KeyValuePair; using BlockReduce = cub::BlockReduce; __shared__ typename BlockReduce::TempStorage tmp_storage; @@ -237,7 +224,6 @@ __launch_bounds__(TPB) __global__ void moeTopK( cub::ArgMax arg_max; const int block_row = blockIdx.x; - const bool row_is_active = finished ? !finished[block_row] : true; const int thread_read_offset = blockIdx.x * num_experts; float row_sum_for_renormalize = 0.0f; @@ -248,20 +234,20 @@ __launch_bounds__(TPB) __global__ void moeTopK( for (int expert = threadIdx.x; expert < num_experts; expert += TPB) { const int idx = thread_read_offset + expert; inp_kvp.key = expert; - inp_kvp.value = inputs_after_softmax[idx]; + const float prob = inputs_after_softmax[idx]; + inp_kvp.value = correction_bias == nullptr ? prob : prob + correction_bias[expert]; thread_kvp = arg_max(inp_kvp, thread_kvp); } const cub_kvp result_kvp = BlockReduce(tmp_storage).Reduce(thread_kvp, arg_max); if (threadIdx.x == 0) { const int expert = result_kvp.key; - const bool node_uses_expert = expert >= start_expert && expert < end_expert; - const bool should_process_row = row_is_active && node_uses_expert; const int idx = k * block_row + k_idx; - output[idx] = result_kvp.value; - indices[idx] = should_process_row ? (expert - start_expert) : num_experts; - row_sum_for_renormalize += result_kvp.value; - inputs_after_softmax[thread_read_offset + expert] = -1.0f; + const float prob = inputs_after_softmax[thread_read_offset + expert]; + output[idx] = prob; + indices[idx] = expert; + row_sum_for_renormalize += prob; + inputs_after_softmax[thread_read_offset + expert] = -FLT_MAX; } __syncthreads(); } @@ -278,13 +264,10 @@ __launch_bounds__(TPB) __global__ void moeTopK( template __launch_bounds__(WARPS_PER_CTA *WARP_SIZE) __global__ void topkGatingSoftmax( const T *input, - const bool *finished, float *output, const int num_rows, int *indices, const int k, - const int start_expert, - const int end_expert, const bool renormalize, const float moe_softcapping, const float *correction_bias) { @@ -315,7 +298,6 @@ __launch_bounds__(WARPS_PER_CTA *WARP_SIZE) __global__ void topkGatingSoftmax( if (thread_row >= num_rows) { return; } - const bool row_is_active = finished ? !finished[thread_row] : true; const T *thread_row_ptr = input + thread_row * ELTS_PER_ROW; const int thread_group_idx = threadIdx.x % THREADS_PER_ROW; @@ -337,19 +319,13 @@ __launch_bounds__(WARPS_PER_CTA *WARP_SIZE) __global__ void topkGatingSoftmax( row_chunk[ii] = convert_to_float(row_chunk_temp[ii]); } - if (moe_softcapping != 0.0f || correction_bias != nullptr) { + if (moe_softcapping != 0.0f) { #pragma unroll for (int ii = 0; ii < VPT; ++ii) { float val = row_chunk[ii]; if (moe_softcapping != 0.0f) { val = tanhf(val / moe_softcapping) * moe_softcapping; } - if (correction_bias != nullptr) { - const int group_id = ii / ELTS_PER_LDG; - const int local_id = ii % ELTS_PER_LDG; - const int expert_idx = first_elt_read_by_thread + group_id * THREADS_PER_ROW * ELTS_PER_LDG + local_id; - val += correction_bias[expert_idx]; - } row_chunk[ii] = val; } } @@ -383,37 +359,41 @@ __launch_bounds__(WARPS_PER_CTA *WARP_SIZE) __global__ void topkGatingSoftmax( const int start_col = first_elt_read_by_thread; float row_sum_for_renormalize = 0.0f; for (int k_idx = 0; k_idx < k; ++k_idx) { - float max_val = row_chunk[0]; + float max_prob = row_chunk[0]; + float max_choice = correction_bias == nullptr ? max_prob : max_prob + correction_bias[start_col]; int expert = start_col; #pragma unroll for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; ++ldg, col += COLS_PER_GROUP_LDG) { #pragma unroll for (int ii = 0; ii < ELTS_PER_LDG; ++ii) { - float val = row_chunk[ldg * ELTS_PER_LDG + ii]; - if (val > max_val) { - max_val = val; - expert = col + ii; + const int expert_idx = col + ii; + float prob = row_chunk[ldg * ELTS_PER_LDG + ii]; + float choice = correction_bias == nullptr ? prob : prob + correction_bias[expert_idx]; + if (choice > max_choice) { + max_choice = choice; + max_prob = prob; + expert = expert_idx; } } } #pragma unroll for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { - float other_max = __shfl_xor_sync(0xffffffff, max_val, mask, THREADS_PER_ROW); + float other_choice = __shfl_xor_sync(0xffffffff, max_choice, mask, THREADS_PER_ROW); + float other_prob = __shfl_xor_sync(0xffffffff, max_prob, mask, THREADS_PER_ROW); int other_expert = __shfl_xor_sync(0xffffffff, expert, mask, THREADS_PER_ROW); - if (other_max > max_val || (other_max == max_val && other_expert < expert)) { - max_val = other_max; + if (other_choice > max_choice || (other_choice == max_choice && other_expert < expert)) { + max_choice = other_choice; + max_prob = other_prob; expert = other_expert; } } if (thread_group_idx == 0) { - const bool node_uses_expert = expert >= start_expert && expert < end_expert; - const bool should_process_row = row_is_active && node_uses_expert; const int idx = k * thread_row + k_idx; - output[idx] = max_val; - indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS; - row_sum_for_renormalize += max_val; + output[idx] = max_prob; + indices[idx] = expert; + row_sum_for_renormalize += max_prob; } if (k_idx + 1 < k) { @@ -421,7 +401,7 @@ __launch_bounds__(WARPS_PER_CTA *WARP_SIZE) __global__ void topkGatingSoftmax( const int thread_to_clear_in_group = (expert / ELTS_PER_LDG) % THREADS_PER_ROW; if (thread_group_idx == thread_to_clear_in_group) { const int offset_for_expert = expert % ELTS_PER_LDG; - row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = -10000.0f; + row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = -FLT_MAX; } } } @@ -453,13 +433,10 @@ struct TopkConstants { template void topkGatingSoftmaxLauncherHelper( const T *input, - const bool *finished, float *output, int *indices, const int num_rows, const int k, - const int start_expert, - const int end_expert, const bool renormalize, const float moe_softcapping, const float *correction_bias, @@ -475,13 +452,12 @@ void topkGatingSoftmaxLauncherHelper( const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB; dim3 block_dim(WARP_SIZE, WARPS_PER_TB); topkGatingSoftmax<<>>( - input, finished, output, num_rows, indices, k, start_expert, end_expert, - renormalize, moe_softcapping, correction_bias); + input, output, num_rows, indices, k, renormalize, moe_softcapping, correction_bias); } #define LAUNCH_SOFTMAX(TYPE, NUM_EXPERTS, WARPS_PER_TB) \ topkGatingSoftmaxLauncherHelper( \ - gating_output, nullptr, topk_weights, topk_indices, num_tokens, topk, 0, num_experts, renormalize, moe_softcapping, correction_bias, stream) + gating_output, topk_weights, topk_indices, num_tokens, topk, renormalize, moe_softcapping, correction_bias, stream) template infiniStatus_t topkGatingSoftmaxKernelLauncher( @@ -534,13 +510,13 @@ infiniStatus_t topkGatingSoftmaxKernelLauncher( } static constexpr int TPB = 256; moeSoftmax<<>>( - gating_output, nullptr, softmax_workspace, num_experts, moe_softcapping, correction_bias); + gating_output, softmax_workspace, num_experts, moe_softcapping); if (topk == 1) { moeTopK<<>>( - softmax_workspace, nullptr, topk_weights, topk_indices, num_experts, topk, 0, num_experts, renormalize); + softmax_workspace, topk_weights, topk_indices, num_experts, topk, renormalize, correction_bias); } else { moeTopKFast<<>>( - softmax_workspace, nullptr, topk_weights, topk_indices, num_experts, topk, 0, num_experts, renormalize); + softmax_workspace, topk_weights, topk_indices, num_experts, topk, renormalize, correction_bias); } break; } diff --git a/test/infinicore/ops/moe_topk_softmax.py b/test/infinicore/ops/moe_topk_softmax.py new file mode 100644 index 000000000..ab9434bc0 --- /dev/null +++ b/test/infinicore/ops/moe_topk_softmax.py @@ -0,0 +1,284 @@ +import os +import sys + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +import infinicore +import torch +import torch.nn.functional as F +from framework import ( + BaseOperatorTest, + GenericTestRunner, + TensorInitializer, + TensorSpec, + TestCase, +) +from infinicore.lib import _infinicore + + +# Test case format: +# (shape, topk, renormalize, moe_softcapping, bias_mode, description) +_TEST_CASES_DATA = [ + ((2, 4), 2, True, 0.0, "none", "ordinary softmax top-k without correction bias"), + ( + (2, 4), + 2, + True, + 0.0, + "zero", + "zero correction bias matches ordinary softmax top-k", + ), + ( + (3, 4), + 2, + True, + 0.0, + "force_last", + "bias changes selected indices, weights use original probs", + ), + ( + (2, 6), + 2, + True, + 0.0, + "force_middle", + "bias semantics on non-power-of-two expert count", + ), + ( + (4, 64), + 6, + True, + 0.0, + "small_random", + "power-of-two expert count with random correction bias", + ), + ( + (2, 8), + 3, + False, + 1.5, + "force_last", + "softcapping with unnormalized selected weights", + ), +] + +_TOLERANCE_MAP = { + infinicore.float16: {"atol": 2e-3, "rtol": 2e-3}, + infinicore.bfloat16: {"atol": 1e-2, "rtol": 1e-2}, + infinicore.float32: {"atol": 1e-5, "rtol": 1e-5}, +} + +_TENSOR_DTYPES = [infinicore.float16, infinicore.bfloat16, infinicore.float32] + + +def _make_logits(shape, seed, forced_expert=None): + generator = torch.Generator(device="cpu").manual_seed(seed) + logits = torch.randn(shape, generator=generator, dtype=torch.float32) * 1.25 + if forced_expert is not None: + logits[:, forced_expert] -= 4.0 + return logits + + +def _make_correction_bias(logits, mode, seed): + num_experts = logits.shape[-1] + if mode == "none": + return None + if mode == "zero": + return torch.zeros(num_experts, dtype=torch.float32) + if mode == "small_random": + generator = torch.Generator(device="cpu").manual_seed(seed) + return torch.randn(num_experts, generator=generator, dtype=torch.float32) * 0.05 + + if mode == "force_last": + target = num_experts - 1 + elif mode == "force_middle": + target = num_experts // 2 + else: + raise ValueError(f"Unsupported bias mode: {mode}") + + probs = F.softmax(logits, dim=-1, dtype=torch.float32) + other_probs = torch.cat([probs[:, :target], probs[:, target + 1 :]], dim=-1) + required_margin = (other_probs.max(dim=-1).values - probs[:, target]).max().item() + bias = torch.zeros(num_experts, dtype=torch.float32) + bias[target] = required_margin + 0.25 + return bias + + +def torch_moe_topk_softmax( + gating_output, topk, correction_bias=None, renormalize=True, moe_softcapping=0.0 +): + logits = gating_output.to(torch.float32) + if moe_softcapping != 0.0: + logits = torch.tanh(logits / moe_softcapping) * moe_softcapping + + probs = F.softmax(logits, dim=-1, dtype=torch.float32) + choice_scores = probs + if correction_bias is not None: + choice_scores = choice_scores + correction_bias.reshape(1, -1).to(torch.float32) + + _, selected_experts = torch.topk(choice_scores, topk, dim=-1) + routing_weights = torch.gather(probs, dim=-1, index=selected_experts) + if renormalize: + routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True) + + return routing_weights, selected_experts.to(torch.int32) + + +def parse_test_cases(): + test_cases = [] + for case_idx, ( + shape, + topk, + renormalize, + softcap, + bias_mode, + description, + ) in enumerate(_TEST_CASES_DATA): + num_experts = shape[-1] + forced_expert = None + if bias_mode == "force_last": + forced_expert = num_experts - 1 + elif bias_mode == "force_middle": + forced_expert = num_experts // 2 + + logits = _make_logits(shape, seed=2027 + case_idx, forced_expert=forced_expert) + bias = _make_correction_bias(logits, bias_mode, seed=4099 + case_idx) + + for dtype in _TENSOR_DTYPES: + input_spec = TensorSpec.from_tensor( + shape, + None, + dtype, + init_mode=TensorInitializer.MANUAL, + set_tensor=logits, + name="gating_output", + ) + inputs = [input_spec] + if bias is not None: + inputs.append( + TensorSpec.from_tensor( + (num_experts,), + None, + infinicore.float32, + init_mode=TensorInitializer.MANUAL, + set_tensor=bias, + name="correction_bias", + ) + ) + + kwargs = { + "topk": topk, + "renormalize": renormalize, + "moe_softcapping": softcap, + } + case_name = ( + f"moe_topk_softmax - {description} - dtype={dtype}, " + f"shape={shape}, topk={topk}, bias={bias_mode}" + ) + + test_cases.append( + TestCase( + inputs=inputs, + kwargs=kwargs, + output_spec=None, + comparison_target=None, + tolerance=_TOLERANCE_MAP[dtype], + description=f"{case_name} - OUT_OF_PLACE", + output_count=2, + ) + ) + + out_shape = (shape[0], topk) + test_cases.append( + TestCase( + inputs=inputs, + kwargs=kwargs.copy(), + output_specs=[ + TensorSpec.from_tensor( + out_shape, None, infinicore.float32, name="topk_weights" + ), + TensorSpec.from_tensor( + out_shape, None, infinicore.int32, name="topk_indices" + ), + ], + comparison_target="out", + tolerance=_TOLERANCE_MAP[dtype], + description=f"{case_name} - INPLACE(out)", + output_count=2, + ) + ) + return test_cases + + +class OpTest(BaseOperatorTest): + def __init__(self): + super().__init__("moe_topk_softmax") + + def get_test_cases(self): + return parse_test_cases() + + def torch_operator( + self, + gating_output, + correction_bias=None, + topk=2, + renormalize=True, + moe_softcapping=0.0, + out=None, + **kwargs, + ): + values, indices = torch_moe_topk_softmax( + gating_output, + topk, + correction_bias=correction_bias, + renormalize=renormalize, + moe_softcapping=moe_softcapping, + ) + if out is not None: + out_v, out_i = out + out_v.copy_(values) + out_i.copy_(indices) + return values, indices + + def infinicore_operator( + self, + gating_output, + correction_bias=None, + topk=2, + renormalize=True, + moe_softcapping=0.0, + out=None, + **kwargs, + ): + if out is None: + values = infinicore.empty( + (gating_output.shape[0], topk), + dtype=infinicore.float32, + device=gating_output.device, + ) + indices = infinicore.empty( + (gating_output.shape[0], topk), + dtype=infinicore.int32, + device=gating_output.device, + ) + else: + values, indices = out + + _infinicore.moe_topk_softmax_( + values._underlying, + indices._underlying, + gating_output._underlying, + correction_bias._underlying if correction_bias is not None else None, + renormalize, + moe_softcapping, + ) + return values, indices + + +def main(): + runner = GenericTestRunner(OpTest) + runner.run_and_exit() + + +if __name__ == "__main__": + main()