feat: support flashinfer sample kernel on AMD

Author: CrimsonDump

rtp_llm/cpp/devices/rocm_impl/ROCmSampleOp.cc

@@ -4,6 +4,7 @@
 #include "rtp_llm/cpp/kernels/sampling_topp_kernels.h"
 #include "rtp_llm/cpp/kernels/sampling_penalty_kernels.h"
 #include "rtp_llm/cpp/core/torch_utils/BufferTorchUtils.h"
+#include "rtp_llm/cpp/kernels/rocm/sampling/sampling.h"
 
 using namespace std;
 
@@ -19,6 +20,7 @@ using SamplerT = float;
 // topk should has higher proirity than topp.
 
 GreedyOutput ROCmDevice::sampleGreedy(const GreedyParams& params) {
+    bool enable_flashinfer = init_params_.sampler_config.enable_flashinfer_sample_kernel;
     const auto& logits     = params.logits;
     const auto  batch_size = logits.shape()[0];
     const auto vocab_size_padded = logits.shape()[1];
@@ -219,6 +221,81 @@ GreedyOutput ROCmDevice::sampleGreedy(const GreedyParams& params) {
         return GreedyOutput{};
     }
 
+    if (enable_flashinfer) {
+        const auto batch_size = params.logits.shape()[0];
+        auto&      top_k      = params.top_k;
+        auto&      top_p      = params.top_p;
+
+        auto      logits_ref = params.logits.slice(0, params.logits.shape()[0]);
+        auto      probs   = softmax({logits_ref, std::nullopt, std::nullopt, 1.0f, DataType::TYPE_INVALID, std::nullopt});
+        auto      samples = transposed_tokens->view(transposed_tokens->shape()[0] - 1, 1);
+        torch::TensorOptions options =
+            torch::TensorOptions(dataTypeToTorchType(probs->type())).device(torch::Device(torch::kCUDA));
+        bool deterministic = false;
+        bool          need_output_all_probs = params.output_all_probs.has_value();
+        torch::Tensor probs_t               = Buffer2torchTensor(probs, false);
+        torch::Tensor samples_t             = Buffer2torchTensor(samples, false).flatten();
+        torch::Tensor top_k_t               = Buffer2torchTensor(top_k, false);
+        torch::Tensor top_p_t               = Buffer2torchTensor(top_p, false);
+        torch::Tensor output_all_probs_t;
+        if (need_output_all_probs) {
+            output_all_probs_t = Buffer2torchTensor(params.output_all_probs.value().get(), false);
+        }
+        std::transform(top_p.data<float>(), top_p.data<float>() + batch_size, top_p.data<float>(), [&](auto t) {
+            return std::abs(t) < 1e-7 ? 1.0 : t;
+        });
+        if (std::all_of(top_k.data<uint32_t>(), top_k.data<uint32_t>() + batch_size, [&](auto t) { return t == 1; })) {
+            torch::Tensor selected_tokens = torch::argmax(probs_t, -1, /*keepdim=*/false);
+            samples_t.copy_(selected_tokens);
+            if (need_output_all_probs) {
+                top_k_renorm_probs(probs_t, output_all_probs_t, top_k_t, 0, reinterpret_cast<uintptr_t>(stream_));
+            }
+        } else if (std::all_of(
+                       top_k.data<uint32_t>(), top_k.data<uint32_t>() + batch_size, [&](auto t) { return t <= 0; })) {
+            top_p_sampling_from_probs(probs_t, samples_t, std::nullopt, top_p_t, 1.0, deterministic, 0, 0, reinterpret_cast<uintptr_t>(stream_));
+            if (need_output_all_probs) {
+                top_p_renorm_probs(probs_t, output_all_probs_t, top_p_t, 1.0, reinterpret_cast<uintptr_t>(stream_));
+            }
+        } else if (std::all_of(top_p.data<float>(), top_p.data<float>() + batch_size, [&](auto t) {
+                       return std::abs(t - 1.0f) < 1e-7;
+                   })) {
+            std::transform(top_k.data<uint32_t>(),
+                           top_k.data<uint32_t>() + batch_size,
+                           top_k.data<uint32_t>(),
+                           [&](auto t) { return t <= 0 ? 1 << 30 : t; });
+            top_k_sampling_from_probs(
+                probs_t, samples_t, std::nullopt, top_k_t, 0, deterministic, 0, 0, reinterpret_cast<uintptr_t>(stream_));
+            if (need_output_all_probs) {
+                top_k_renorm_probs(probs_t, output_all_probs_t, top_k_t, 0, reinterpret_cast<uintptr_t>(stream_));
+            }
+        } else {
+            std::transform(top_k.data<uint32_t>(),
+                           top_k.data<uint32_t>() + batch_size,
+                           top_k.data<uint32_t>(),
+                           [&](auto t) { return t <= 0 ? 1 << 30 : t; });
+            top_k_top_p_sampling_from_probs(probs_t,
+                                            samples_t,
+                                            std::nullopt,
+                                            top_k_t,
+                                            0,
+                                            top_p_t,
+                                            1.0,
+                                            deterministic,
+                                            0,
+                                            0,
+                                            reinterpret_cast<uintptr_t>(stream_));
+            if (need_output_all_probs) {
+                torch::Tensor temp_t = torch::zeros_like(output_all_probs_t);
+                top_k_renorm_probs(probs_t, temp_t, top_k_t, 1.0, reinterpret_cast<uintptr_t>(stream_));
+                top_p_renorm_probs(temp_t, output_all_probs_t, top_p_t, 1.0, reinterpret_cast<uintptr_t>(stream_));
+            }
+        }
+        auto output_tokens = transpose({*transposed_tokens});
+        copy({params.token_ids, *output_tokens});
+        check_cuda_error();
+        return GreedyOutput{};
+    }
+
     // 4. run sampling
     // 4.1 run top_k
     invokeSetupTopKRuntimeArgs(batch_size,

rtp_llm/cpp/kernels/BUILD

@@ -235,6 +235,7 @@ cc_library(
     deps = [
         ":rocm_basic",
         ":rocm_mla",
+        "//rtp_llm/cpp/kernels/rocm/sampling:sampling",
     ],
     visibility = ["//visibility:public"],
 )

rtp_llm/cpp/kernels/rocm/sampling/BUILD

@@ -0,0 +1,41 @@
+load("//:def.bzl", "rocm_copts")
+load("//bazel:arch_select.bzl", "torch_deps")
+
+cc_library(
+    name = "sampling",
+    srcs = ["api.cc"],
+    hdrs = [
+        "kernel.cuh",
+        "sampling.h",
+        "utils.h",
+    ],
+    deps = [
+        "//rtp_llm/cpp/kernels:rocm_utils",
+        "@local_config_rocm//rocm:rocm_headers",
+    ] + torch_deps(),
+    copts = rocm_copts() + ["-DUSE_ROCM=1"],
+    visibility = ["//visibility:public"],
+)
+
+cc_binary(
+    name = "bind.so",
+    srcs = [
+        "bind.cc",
+    ],
+    deps = [
+        ":sampling",
+        "//rtp_llm/cpp/pybind:py_utils",
+    ],
+    linkshared = True,
+    linkstatic = False,
+)
+
+py_test(
+    name = "test",
+    srcs = ["test.py"],
+    data = [
+        ":bind.so",
+    ],
+    imports = ["."],
+    python_version = "PY3",
+)

rtp_llm/cpp/kernels/rocm/sampling/api.cc

@@ -0,0 +1,126 @@
+// based on flashinfer 0.4.1 https://github.com/flashinfer-ai/flashinfer/tree/a88349f9f43df74d31d1d52ad5aa20c28824a790
+/*
+ * Copyright (c) 2024 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "sampling.h"
+#include "utils.h"
+#include "kernel.cuh"
+
+namespace rtp_llm {
+
+void top_p_sampling_from_probs(torch::Tensor probs, torch::Tensor output,
+                               std::optional<torch::Tensor> maybe_indices,
+                               std::optional<torch::Tensor> maybe_top_p_arr, double top_p_val,
+                               bool deterministic, uint64_t philox_seed, uint64_t philox_offset, uintptr_t stream) {
+  CHECK_INPUT(probs);
+  CHECK_DIM(2, probs);  // probs: (batch_size, vocab_size)
+  unsigned int batch_size = output.sizes()[0];
+  unsigned int vocab_size = probs.sizes()[1];
+  bool has_top_p_arr = maybe_top_p_arr.has_value();
+
+  hipSetDevice(probs.get_device());
+  hipError_t status = sampling::TopPSamplingFromProb<float, int>(
+      static_cast<float*>(probs.data_ptr()), static_cast<int*>(output.data_ptr()),
+      maybe_indices.has_value() ? static_cast<int*>(maybe_indices->data_ptr()) : nullptr,
+      has_top_p_arr ? static_cast<float*>(maybe_top_p_arr->data_ptr()) : nullptr, batch_size,
+      top_p_val, vocab_size, deterministic, philox_seed, philox_offset, reinterpret_cast<hipStream_t>(stream));
+  TORCH_CHECK(status == hipSuccess, "TopPSamplingFromProbs failed with error code " + std::string(hipGetErrorString(status)));
+}
+
+void top_k_sampling_from_probs(torch::Tensor probs, torch::Tensor output,
+                               std::optional<torch::Tensor> maybe_indices,
+                               std::optional<torch::Tensor> maybe_top_k_arr, int64_t top_k_val,
+                               bool deterministic, uint64_t philox_seed, uint64_t philox_offset, uintptr_t stream) {
+  CHECK_INPUT(probs);
+  CHECK_INPUT(output);
+  CHECK_DEVICE(output, probs);
+  CHECK_DIM(2, probs);   // probs: (batch_size, vocab_size)
+  CHECK_DIM(1, output);  // output: (batch_size)
+  unsigned int batch_size = output.sizes()[0];
+  unsigned int vocab_size = probs.sizes()[1];
+  bool has_top_k_arr = maybe_top_k_arr.has_value();
+
+  hipSetDevice(probs.get_device());
+  hipError_t status = sampling::TopKSamplingFromProb<float, int>(
+      static_cast<float*>(probs.data_ptr()), static_cast<int*>(output.data_ptr()),
+      maybe_indices.has_value() ? static_cast<int*>(maybe_indices->data_ptr()) : nullptr,
+      has_top_k_arr ? static_cast<float*>(maybe_top_k_arr->data_ptr()) : nullptr, batch_size,
+      top_k_val, vocab_size, deterministic, philox_seed, philox_offset, reinterpret_cast<hipStream_t>(stream));
+  TORCH_CHECK(status == hipSuccess, "TopKSamplingFromProbs failed with error code " + std::string(hipGetErrorString(status)));
+}
+
+void top_k_top_p_sampling_from_probs(torch::Tensor probs, torch::Tensor output,
+                                     std::optional<torch::Tensor> maybe_indices,
+                                     std::optional<torch::Tensor> maybe_top_k_arr, double top_k_val,
+                                     std::optional<torch::Tensor> maybe_top_p_arr, double top_p_val,
+                                     bool deterministic, uint64_t philox_seed,
+                                     uint64_t philox_offset, uintptr_t stream) {
+  CHECK_INPUT(probs);
+  CHECK_INPUT(output);
+  CHECK_DEVICE(output, probs);
+  CHECK_DIM(2, probs);   // probs: (batch_size, vocab_size)
+  CHECK_DIM(1, output);  // output: (batch_size)
+  unsigned int batch_size = output.sizes()[0];
+  unsigned int vocab_size = probs.sizes()[1];
+  bool has_top_k_arr = maybe_top_k_arr.has_value();
+  bool has_top_p_arr = maybe_top_p_arr.has_value();
+
+  hipSetDevice(probs.get_device());
+  hipError_t status = sampling::TopKTopPSamplingFromProb<float, int>(
+      static_cast<float*>(probs.data_ptr()),
+      has_top_k_arr ? static_cast<int*>(maybe_top_k_arr->data_ptr()) : nullptr,
+      has_top_p_arr ? static_cast<float*>(maybe_top_p_arr->data_ptr()) : nullptr,
+      static_cast<int*>(output.data_ptr()),
+      maybe_indices.has_value() ? static_cast<int*>(maybe_indices->data_ptr()) : nullptr,
+      batch_size, top_k_val, top_p_val, vocab_size, deterministic, philox_seed, philox_offset,
+      reinterpret_cast<hipStream_t>(stream));
+  TORCH_CHECK(status == hipSuccess, "TopKTopPSamplingFromProb failed with error code " + std::string(hipGetErrorString(status)));
+}
+
+void top_p_renorm_probs(torch::Tensor probs, torch::Tensor renorm_probs,
+                        std::optional<torch::Tensor> maybe_top_p_arr, double top_p_val, uintptr_t stream) {
+  CHECK_INPUT(probs);
+  CHECK_DIM(2, probs);  // probs: (batch_size, vocab_size)
+  unsigned int batch_size = probs.sizes()[0];
+  unsigned int vocab_size = probs.sizes()[1];
+  bool has_top_p_arr = maybe_top_p_arr.has_value();
+
+  hipSetDevice(probs.get_device());
+  hipError_t status = sampling::TopPRenormProb<float>(
+      static_cast<float*>(probs.data_ptr()), static_cast<float*>(renorm_probs.data_ptr()),
+      has_top_p_arr ? static_cast<float*>(maybe_top_p_arr->data_ptr()) : nullptr, batch_size,
+      top_p_val, vocab_size, reinterpret_cast<hipStream_t>(stream));
+  
+  TORCH_CHECK(status == hipSuccess, "TopPRenormProb failed with error code " + std::string(hipGetErrorString(status)));
+}
+
+void top_k_renorm_probs(torch::Tensor probs, torch::Tensor renorm_probs,
+                        std::optional<torch::Tensor> maybe_top_k_arr, int64_t top_k_val, uintptr_t stream) {
+  CHECK_INPUT(probs);
+  CHECK_DIM(2, probs);  // probs: (batch_size, vocab_size)
+  unsigned int batch_size = probs.sizes()[0];
+  unsigned int vocab_size = probs.sizes()[1];
+  bool has_top_k_arr = maybe_top_k_arr.has_value();
+
+  hipSetDevice(probs.get_device());
+  hipError_t status = sampling::TopKRenormProb<float>(
+      static_cast<float*>(probs.data_ptr()), static_cast<float*>(renorm_probs.data_ptr()),
+      has_top_k_arr ? static_cast<int*>(maybe_top_k_arr->data_ptr()) : nullptr, batch_size,
+      top_k_val, vocab_size, reinterpret_cast<hipStream_t>(stream));
+
+  TORCH_CHECK(status == hipSuccess, "TopKRenormProb failed with error code " + std::string(hipGetErrorString(status)));
+}
+
+}

rtp_llm/cpp/kernels/rocm/sampling/bind.cc

@@ -0,0 +1,18 @@
+#include <pybind11/pybind11.h>
+
+#include "sampling.h"
+
+namespace py = pybind11;
+
+namespace rtp_llm {
+
+PYBIND11_MODULE(bind, m) {
+    m.doc() = "sampling c++ api for test";
+    m.def("top_p_renorm_probs", &top_p_renorm_probs, py::arg(), py::arg(), py::arg(), py::arg(), py::arg("stream") = 0, "top_p_renorm_probs");
+    m.def("top_k_renorm_probs", &top_k_renorm_probs, py::arg(), py::arg(), py::arg(), py::arg(), py::arg("stream") = 0, "top_k_renorm_probs");
+    m.def("top_p_sampling_from_probs", &top_p_sampling_from_probs, py::arg(), py::arg(), py::arg(), py::arg(), py::arg(), py::arg(), py::arg(), py::arg(), py::arg("stream") = 0, "top_p_sampling_from_probs");
+    m.def("top_k_sampling_from_probs", &top_k_sampling_from_probs, py::arg(), py::arg(), py::arg(), py::arg(), py::arg(), py::arg(), py::arg(), py::arg(), py::arg("stream") = 0, "top_k_sampling_from_probs");
+    m.def("top_k_top_p_sampling_from_probs", &top_k_top_p_sampling_from_probs, py::arg(), py::arg(), py::arg(), py::arg(), py::arg(), py::arg(), py::arg(), py::arg(), py::arg(), py::arg(), py::arg("stream") = 0, "top_k_top_p_sampling_from_probs");
+}
+
+}