diff --git a/src/native/cambricon/ops/topk_softmax/kernel.h b/src/native/cambricon/ops/topk_softmax/kernel.h new file mode 100644 index 000000000..a29233e4d --- /dev/null +++ b/src/native/cambricon/ops/topk_softmax/kernel.h @@ -0,0 +1,80 @@ +#ifndef INFINI_OPS_CAMBRICON_TOPK_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_CAMBRICON_TOPK_SOFTMAX_KERNEL_H_ + +#include +#include +#include + +#include "base/topk_softmax.h" +#include "dispatcher.h" +#include "native/cambricon/common.h" +#include "native/cambricon/data_type_.h" + +namespace infini::ops { + +template +void TopkSoftmaxUnion(cnrtQueue_t queue, int core_per_cluster, + int cluster_count, const T* gating_output, + const float* bias, const uint8_t* is_padding, + float* topk_weights, Index* topk_indices, + int32_t* token_expert_indices, int32_t num_tokens, + int32_t num_experts, int32_t topk, bool renormalize); + +template <> +class Operator : public TopkSoftmax { + public: + Operator(const Tensor gating_output, std::optional bias, + std::optional is_padding, const bool renormalize, + Tensor topk_weights, Tensor topk_indices, + Tensor token_expert_indices) + : TopkSoftmax{gating_output, bias, is_padding, + renormalize, topk_weights, topk_indices, + token_expert_indices} { + cnrt_utils::GetLaunchConfig(gating_output.device(), &core_per_cluster_, + &cluster_count_); + } + + void operator()(const Tensor gating_output, std::optional bias, + std::optional is_padding, const bool renormalize, + Tensor topk_weights, Tensor topk_indices, + Tensor token_expert_indices) const override { + ValidateCallMetadata(gating_output, bias, is_padding, renormalize, + topk_weights, topk_indices, token_expert_indices); + if (num_tokens_ == 0) { + return; + } + + const auto queue = static_cast(stream_ ? stream_ : 0); + using InputTypes = + List; + using IndexTypes = + List; + DispatchFunc( + {input_dtype_, index_dtype_}, + [&](auto input_tag, auto index_tag) { + using T = typename decltype(input_tag)::type; + using Index = typename decltype(index_tag)::type; + TopkSoftmaxUnion( + queue, core_per_cluster_, cluster_count_, + static_cast(gating_output.data()), + bias ? static_cast(bias->data()) : nullptr, + is_padding ? static_cast(is_padding->data()) + : nullptr, + static_cast(topk_weights.data()), + static_cast(topk_indices.data()), + static_cast(token_expert_indices.data()), + static_cast(num_tokens_), + static_cast(num_experts_), static_cast(topk_), + renormalize_); + }, + "CambriconTopkSoftmax::operator()"); + } + + private: + int core_per_cluster_{0}; + int cluster_count_{0}; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_CAMBRICON_TOPK_SOFTMAX_KERNEL_H_ diff --git a/src/native/cambricon/ops/topk_softmax/kernel.mlu b/src/native/cambricon/ops/topk_softmax/kernel.mlu new file mode 100644 index 000000000..86b8ce5a6 --- /dev/null +++ b/src/native/cambricon/ops/topk_softmax/kernel.mlu @@ -0,0 +1,182 @@ +#include +#include +#include +#include +#include + +#include "kernel.h" + +__nram__ char topk_softmax_nram[NRAM_MAX_SIZE] __attribute__((aligned(128))); + +namespace infini::ops { +namespace { + +constexpr std::size_t kAlignment = 128; + +__mlu_device__ std::size_t AlignUpDevice(std::size_t size) { + return (size + kAlignment - 1) / kAlignment * kAlignment; +} + +template +__mlu_device__ void LoadAsFloat(float* destination, T* cache, const T* source, + std::size_t count) { + __memcpy(cache, const_cast(source), count * sizeof(T), GDRAM2NRAM); + if constexpr (std::is_same_v) { + __bang_half2float(destination, reinterpret_cast(cache), count); + } else if constexpr (std::is_same_v) { + __bang_bfloat162float(destination, cache, count); + } else { + __memcpy(destination, cache, count * sizeof(float), NRAM2NRAM); + } +} + +template +__mlu_device__ void StoreIndex(Index* destination, bool padded, + int32_t selected_index) { + if (!padded) { + *destination = static_cast(selected_index); + } else if constexpr (std::is_same_v) { + auto* words = reinterpret_cast(destination); + words[0] = UINT32_MAX; + words[1] = UINT32_MAX; + } else { + *destination = static_cast(UINT32_MAX); + } +} + +template +__mlu_global__ void TopkSoftmaxKernel(const T* gating_output, const float* bias, + const uint8_t* is_padding, + float* topk_weights, Index* topk_indices, + int32_t* token_expert_indices, + int32_t num_tokens, int32_t num_experts, + int32_t topk, bool renormalize) { + char* cursor = topk_softmax_nram; + auto* input_cache = reinterpret_cast(cursor); + cursor += AlignUpDevice(static_cast(num_experts) * sizeof(T)); + auto* probabilities = reinterpret_cast(cursor); + cursor += + AlignUpDevice(static_cast(num_experts) * sizeof(float)); + auto* selection = reinterpret_cast(cursor); + cursor += + AlignUpDevice(static_cast(num_experts) * sizeof(float)); + auto* selected_weights = reinterpret_cast(cursor); + cursor += AlignUpDevice(static_cast(topk) * sizeof(float)); + auto* selected_indices = reinterpret_cast(cursor); + + for (int32_t token = taskId; token < num_tokens; token += taskDim) { + LoadAsFloat(probabilities, input_cache, + gating_output + static_cast(token) * num_experts, + num_experts); + + float max_value = -INFINITY; + for (int32_t expert = 0; expert < num_experts; ++expert) { + max_value = std::max(max_value, probabilities[expert]); + } + float sum = 0.0f; + for (int32_t expert = 0; expert < num_experts; ++expert) { + probabilities[expert] = expf(probabilities[expert] - max_value); + sum += probabilities[expert]; + } + const float inverse_sum = 1.0f / sum; + for (int32_t expert = 0; expert < num_experts; ++expert) { + probabilities[expert] *= inverse_sum; + } + + if (bias != nullptr) { + __memcpy(selection, const_cast(bias), + static_cast(num_experts) * sizeof(float), + GDRAM2NRAM); + for (int32_t expert = 0; expert < num_experts; ++expert) { + const float bias_value = selection[expert]; + selection[expert] = isnan(bias_value) || bias_value == -INFINITY + ? probabilities[expert] + : probabilities[expert] + bias_value; + } + } else { + __memcpy(selection, probabilities, + static_cast(num_experts) * sizeof(float), + NRAM2NRAM); + } + + float selected_sum = 0.0f; + for (int32_t rank = 0; rank < topk; ++rank) { + bool found = false; + float best_value = -INFINITY; + int32_t best_expert = 0; + for (int32_t expert = 0; expert < num_experts; ++expert) { + const float value = selection[expert]; + if (isnan(value)) { + continue; + } + if (!found || value > best_value || + (value == best_value && expert < best_expert)) { + found = true; + best_value = value; + best_expert = expert; + } + } + selected_indices[rank] = best_expert; + selected_weights[rank] = probabilities[best_expert]; + selected_sum += selected_weights[rank]; + selection[best_expert] = NAN; + } + + const bool padded = is_padding != nullptr && is_padding[token] != 0; + const float weight_scale = renormalize ? 1.0f / selected_sum : 1.0f; + for (int32_t rank = 0; rank < topk; ++rank) { + const std::size_t output_index = + static_cast(token) * topk + rank; + topk_weights[output_index] = selected_weights[rank] * weight_scale; + StoreIndex(topk_indices + output_index, padded, selected_indices[rank]); + token_expert_indices[output_index] = rank * num_tokens + token; + } + } +} + +} // namespace + +template +void TopkSoftmaxUnion(cnrtQueue_t queue, int core_per_cluster, + int cluster_count, const T* gating_output, + const float* bias, const uint8_t* is_padding, + float* topk_weights, Index* topk_indices, + int32_t* token_expert_indices, int32_t num_tokens, + int32_t num_experts, int32_t topk, bool renormalize) { + const auto align_up = [](std::size_t size) { + return (size + kAlignment - 1) / kAlignment * kAlignment; + }; + const std::size_t required_bytes = + align_up(static_cast(num_experts) * sizeof(T)) + + 2 * align_up(static_cast(num_experts) * sizeof(float)) + + align_up(static_cast(topk) * sizeof(float)) + + static_cast(topk) * sizeof(int32_t); + assert(required_bytes <= NRAM_MAX_SIZE && + "TopkSoftmax inputs do not fit in NRAM"); + + const cnrtDim3_t kernel_dim = {static_cast(core_per_cluster), + static_cast(cluster_count), 1}; + TopkSoftmaxKernel<<>>( + gating_output, bias, is_padding, topk_weights, topk_indices, + token_expert_indices, num_tokens, num_experts, topk, renormalize); + CNRT_CHECK(cnrtGetLastError()); +} + +#define INSTANTIATE_TOPK_SOFTMAX(T, INDEX) \ + template void TopkSoftmaxUnion( \ + cnrtQueue_t, int, int, const T*, const float*, const uint8_t*, float*, \ + INDEX*, int32_t*, int32_t, int32_t, int32_t, bool) + +INSTANTIATE_TOPK_SOFTMAX(__half, int32_t); +INSTANTIATE_TOPK_SOFTMAX(__half, uint32_t); +INSTANTIATE_TOPK_SOFTMAX(__half, int64_t); +INSTANTIATE_TOPK_SOFTMAX(__bang_bfloat16, int32_t); +INSTANTIATE_TOPK_SOFTMAX(__bang_bfloat16, uint32_t); +INSTANTIATE_TOPK_SOFTMAX(__bang_bfloat16, int64_t); +INSTANTIATE_TOPK_SOFTMAX(float, int32_t); +INSTANTIATE_TOPK_SOFTMAX(float, uint32_t); +INSTANTIATE_TOPK_SOFTMAX(float, int64_t); + +#undef INSTANTIATE_TOPK_SOFTMAX + +} // namespace infini::ops diff --git a/tests/test_topk_softmax.py b/tests/test_topk_softmax.py index 825731c99..730733e55 100644 --- a/tests/test_topk_softmax.py +++ b/tests/test_topk_softmax.py @@ -20,9 +20,9 @@ (torch.bfloat16, 1e-6, 1e-6), ), ) -def test_topk_softmax(dtype, index_dtype, has_bias, renormalize, rtol, atol): - if not torch.cuda.is_available(): - pytest.skip("`topk_softmax` requires the NVIDIA backend") +def test_topk_softmax(dtype, index_dtype, has_bias, renormalize, rtol, atol, device): + if device == "mlu" and index_dtype == torch.uint32: + pytest.skip("torch_mlu cannot construct uint32 test outputs") gating_output = torch.tensor( ( @@ -31,7 +31,7 @@ def test_topk_softmax(dtype, index_dtype, has_bias, renormalize, rtol, atol): (0.125, 0.75, 2.5, 1.0, -0.75), ), dtype=dtype, - device="cuda", + device=device, ) bias = None if has_bias: @@ -58,12 +58,10 @@ def test_topk_softmax(dtype, index_dtype, has_bias, renormalize, rtol, atol): torch.testing.assert_close(outputs[2], expected[2], rtol=0, atol=0) -def test_topk_softmax_bias_only_changes_selection(): - if not torch.cuda.is_available(): - pytest.skip("`topk_softmax` requires the NVIDIA backend") +def test_topk_softmax_bias_only_changes_selection(device): - gating_output = torch.tensor(((3.0, 2.0, 1.0),), dtype=torch.float32, device="cuda") - bias = torch.tensor((-4.0, 0.0, 3.0), dtype=torch.float32, device="cuda") + gating_output = torch.tensor(((3.0, 2.0, 1.0),), dtype=torch.float32, device=device) + bias = torch.tensor((-4.0, 0.0, 3.0), dtype=torch.float32, device=device) outputs = _make_outputs(gating_output, topk=1, index_dtype=torch.int32) infini.ops.topk_softmax( @@ -81,14 +79,12 @@ def test_topk_softmax_bias_only_changes_selection(): @pytest.mark.parametrize("bias_value", (float("nan"), float("-inf"))) -def test_topk_softmax_nonfinite_bias_selects_valid_experts(bias_value): - if not torch.cuda.is_available(): - pytest.skip("`topk_softmax` requires the NVIDIA backend") +def test_topk_softmax_nonfinite_bias_selects_valid_experts(bias_value, device): gating_output = torch.tensor( - ((1.0, 2.0, 3.0, 4.0),), dtype=torch.float32, device="cuda" + ((1.0, 2.0, 3.0, 4.0),), dtype=torch.float32, device=device ) - bias = torch.full((4,), bias_value, dtype=torch.float32, device="cuda") + bias = torch.full((4,), bias_value, dtype=torch.float32, device=device) first = _make_outputs(gating_output, topk=3, index_dtype=torch.int32) second = _make_outputs(gating_output, topk=3, index_dtype=torch.int32) @@ -108,9 +104,7 @@ def test_topk_softmax_nonfinite_bias_selects_valid_experts(bias_value): assert torch.unique(first[1]).numel() == 3 -def test_topk_softmax_padding_and_token_expert_indices(): - if not torch.cuda.is_available(): - pytest.skip("`topk_softmax` requires the NVIDIA backend") +def test_topk_softmax_padding_and_token_expert_indices(device): gating_output = torch.tensor( ( @@ -119,9 +113,9 @@ def test_topk_softmax_padding_and_token_expert_indices(): (-1.0, 0.5, 2.25, 1.25), ), dtype=torch.float16, - device="cuda", + device=device, ) - is_padding = torch.tensor((False, True, False), dtype=torch.bool, device="cuda") + is_padding = torch.tensor((False, True, False), dtype=torch.bool, device=device) outputs = _make_outputs(gating_output, topk=3, index_dtype=torch.int64) infini.ops.topk_softmax( @@ -135,18 +129,22 @@ def test_topk_softmax_padding_and_token_expert_indices(): expected = _reference(gating_output, None, is_padding, 3, False) torch.testing.assert_close(outputs[0], expected[0]) - torch.testing.assert_close(outputs[1], expected[1], rtol=0, atol=0) + torch.testing.assert_close(outputs[1].cpu(), expected[1].cpu(), rtol=0, atol=0) torch.testing.assert_close(outputs[2], expected[2], rtol=0, atol=0) -def test_topk_softmax_non_default_stream(): - if not torch.cuda.is_available(): - pytest.skip("non-default CUDA streams require the NVIDIA backend") +def test_topk_softmax_non_default_stream(device): + if device == "cuda": + accelerator = torch.cuda + stream_attribute = "cuda_stream" + else: + accelerator = torch.mlu + stream_attribute = "mlu_stream" - gating_output = torch.randn((7, 13), dtype=torch.bfloat16, device="cuda") + gating_output = torch.randn((7, 13), dtype=torch.bfloat16, device=device) outputs = _make_outputs(gating_output, topk=4, index_dtype=torch.int32) - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) + stream = accelerator.Stream() + stream.wait_stream(accelerator.current_stream()) infini.ops.topk_softmax( gating_output, @@ -154,7 +152,7 @@ def test_topk_softmax_non_default_stream(): None, True, *outputs, - stream=stream.cuda_stream, + stream=getattr(stream, stream_attribute), ) stream.synchronize() @@ -166,11 +164,9 @@ def test_topk_softmax_non_default_stream(): torch.testing.assert_close(outputs[2], expected[2], rtol=0, atol=0) -def test_topk_softmax_tie_selects_smaller_expert_index(): - if not torch.cuda.is_available(): - pytest.skip("`topk_softmax` requires the NVIDIA backend") +def test_topk_softmax_tie_selects_smaller_expert_index(device): - gating_output = torch.zeros((1, 4), dtype=torch.float32, device="cuda") + gating_output = torch.zeros((1, 4), dtype=torch.float32, device=device) outputs = _make_outputs(gating_output, topk=3, index_dtype=torch.int32) infini.ops.topk_softmax( @@ -183,16 +179,16 @@ def test_topk_softmax_tie_selects_smaller_expert_index(): ) torch.testing.assert_close( - outputs[1], torch.tensor(((0, 1, 2),), dtype=torch.int32, device="cuda") + outputs[1], torch.tensor(((0, 1, 2),), dtype=torch.int32, device=device) ) -def test_topk_softmax_padding_uses_uint32_max_sentinel(): - if not torch.cuda.is_available(): - pytest.skip("`topk_softmax` requires the NVIDIA backend") +def test_topk_softmax_padding_uses_uint32_max_sentinel(device): + if device == "mlu": + pytest.skip("torch_mlu cannot construct uint32 test outputs") - gating_output = torch.tensor(((1.0, 2.0, 3.0),), dtype=torch.float32, device="cuda") - is_padding = torch.tensor((True,), dtype=torch.bool, device="cuda") + gating_output = torch.tensor(((1.0, 2.0, 3.0),), dtype=torch.float32, device=device) + is_padding = torch.tensor((True,), dtype=torch.bool, device=device) outputs = _make_outputs(gating_output, topk=2, index_dtype=torch.uint32) infini.ops.topk_softmax( @@ -208,16 +204,14 @@ def test_topk_softmax_padding_uses_uint32_max_sentinel(): (1, 2), torch.iinfo(torch.uint32).max, dtype=torch.uint32, - device="cuda", + device=device, ) torch.testing.assert_close(outputs[1], expected, rtol=0, atol=0) -def test_topk_softmax_empty_tokens(): - if not torch.cuda.is_available(): - pytest.skip("`topk_softmax` requires the NVIDIA backend") +def test_topk_softmax_empty_tokens(device): - gating_output = torch.empty((0, 4), dtype=torch.float16, device="cuda") + gating_output = torch.empty((0, 4), dtype=torch.float16, device=device) outputs = _make_outputs(gating_output, topk=2, index_dtype=torch.int64) result = infini.ops.topk_softmax( @@ -233,13 +227,11 @@ def test_topk_softmax_empty_tokens(): assert all(output.shape == (0, 2) for output in outputs) -def test_topk_softmax_descriptor_reuses_matching_metadata(): - if not torch.cuda.is_available(): - pytest.skip("`topk_softmax` requires the NVIDIA backend") +def test_topk_softmax_descriptor_reuses_matching_metadata(device): - gating_output = torch.randn((3, 7), dtype=torch.float16, device="cuda") - bias = torch.randn((7,), dtype=torch.float32, device="cuda") - is_padding = torch.tensor((False, True, False), dtype=torch.bool, device="cuda") + gating_output = torch.randn((3, 7), dtype=torch.float16, device=device) + bias = torch.randn((7,), dtype=torch.float32, device=device) + is_padding = torch.tensor((False, True, False), dtype=torch.bool, device=device) outputs = _make_outputs(gating_output, topk=2, index_dtype=torch.int64) operator = infini.ops.TopkSoftmax(gating_output, bias, is_padding, True, *outputs) reused_input = torch.randn_like(gating_output) @@ -251,7 +243,9 @@ def test_topk_softmax_descriptor_reuses_matching_metadata(): expected = _reference(reused_input, reused_bias, reused_padding, 2, True) torch.testing.assert_close(reused_outputs[0], expected[0], rtol=1e-6, atol=1e-6) - torch.testing.assert_close(reused_outputs[1], expected[1], rtol=0, atol=0) + torch.testing.assert_close( + reused_outputs[1].cpu(), expected[1].cpu(), rtol=0, atol=0 + ) torch.testing.assert_close(reused_outputs[2], expected[2], rtol=0, atol=0)