/*
 * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * 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.
 */
#pragma once

#include <benchmark/benchmark.h>

#include <nlohmann/json.hpp>

#ifdef USING_OSS_CUTLASS_MOE_GEMM
#include "tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h"
#else
#include "moe_kernels.h"
#endif

#include "tensorrt_llm/kernels/cutlass_kernels/include/cutlass_kernel_selector.h"

#include "tensorrt_llm/common/cudaUtils.h"
#include "tensorrt_llm/common/memoryUtils.h"
#include "tensorrt_llm/common/nvtxUtils.h"
#include "tensorrt_llm/kernels/cutlass_kernels/cutlass_preprocessors.h"
#include "tensorrt_llm/runtime/bufferManager.h"
#include "tensorrt_llm/runtime/cudaStream.h"

#include <algorithm>
#include <cuda.h>
#include <memory>
#include <numeric>
#include <random>
#include <type_traits>
#include <unordered_map>
#include <vector>

using namespace tensorrt_llm::kernels;
using namespace tensorrt_llm::common;
using namespace tensorrt_llm::runtime;
using namespace tensorrt_llm::cutlass_extensions;

using namespace CUTLASS_MOE_GEMM_KERNELS_NAMESPACE;
using CUTLASS_MOE_GEMM_NAMESPACE::TmaWarpSpecializedGroupedGemmInput;
using CUTLASS_MOE_GEMM_KERNELS_NAMESPACE::CutlassMoeFCRunner;
using CUTLASS_MOE_GEMM_NAMESPACE::ActivationType;
using CUTLASS_MOE_GEMM_KERNELS_NAMESPACE::ActivationParams;
using CUTLASS_MOE_GEMM_NAMESPACE::isGatedActivation;

static BufferManager::CudaStreamPtr streamPtr;
static std::unique_ptr<BufferManager> bufferManager;
static int deviceCount;
static char* workloadFile = nullptr;
static bool useCudaGraph = true;

enum VERBOSE_LEVEL
{
    SILENT = 0,
    ERROR = 1,
    INFO = 2,
    VERBOSE = 3
};

constexpr int LOG_LEVEL = ERROR;

enum class GemmToProfile : int
{
    GEMM_1 = static_cast<int>(GemmProfilerBackend::GemmToProfile::GEMM_1),
    GEMM_2 = static_cast<int>(GemmProfilerBackend::GemmToProfile::GEMM_2),
    LAYER = static_cast<int>(3),
};

namespace
{
// Abstract class for routing config
struct RoutingConfig
{
    virtual void start(){};
    virtual void setRouting(int* selected_experts, int64_t num_experts, int64_t k, int64_t num_tokens) = 0;
    virtual std::string getName() = 0;
    virtual bool isDeterministic() const = 0;
    virtual bool supportsConfig(int64_t num_experts, int64_t k, int64_t num_tokens) const = 0;
};

/**
 * Generates a perfectly balanced routing configuration
 */
struct LoadBalancedRoutingConfig : public RoutingConfig
{
    std::string getName() override
    {
        return "balanced";
    }

    bool isDeterministic() const override
    {
        return true;
    }

    void setRouting(int* selected_experts, int64_t num_experts, int64_t k, int64_t num_tokens) override
    {
        std::vector<int> h_selected_experts(k * num_tokens, 0);
        int stride = tensorrt_llm::common::ceilDiv(num_experts, k);
        for (int token = 0; token < num_tokens; token++)
        {
            for (int i = 0; i < k; i++)
            {
                h_selected_experts[token * k + i] = (token + i * stride) % num_experts;
            }
        }

        check_cuda_error(cudaMemcpyAsync(selected_experts, h_selected_experts.data(),
            h_selected_experts.size() * sizeof(int), cudaMemcpyHostToDevice, streamPtr->get()));
        check_cuda_error(cudaStreamSynchronize(streamPtr->get()));
    }

    bool supportsConfig(int64_t, int64_t, int64_t) const override
    {
        return true;
    }
};

/**
 * Selects experts according to given random distribution
 */
struct RandomDistributionRoutingConfig : public RoutingConfig
{
    using ElementType = float;
    std::mt19937_64 twister{0xD5};
    std::vector<float> probabilities;
    int64_t num_experts;
    int64_t k;
    std::string name;

    RandomDistributionRoutingConfig(std::vector<ElementType> const& in_probabilities, int64_t num_experts, int64_t k,
        std::string name = "random_distribution")
        : probabilities(std::move(in_probabilities))
        , num_experts(num_experts)
        , k(k)
        , name(std::move(name))
    {
        TLLM_CHECK_WITH_INFO(num_experts == probabilities.size(),
            "Cannot create random routing distribution. Number of experts does not match the number of weights");
    }

    void start()
    {
        twister.seed(0xD5);
    }

    std::string getName() override
    {
        return name;
    }

    bool isDeterministic() const override
    {
        return false;
    }

    void doSample(float& curr_max, std::vector<int>& selected)
    {
        std::uniform_real_distribution<float> dist(0, curr_max);
        float value = dist(twister);
        float running_sum = 0;
        for (int expert = 0; expert < probabilities.size(); expert++)
        {
            if (std::find(selected.begin(), selected.end(), expert) != selected.end())
                continue; // Already picked
            float prob = probabilities[expert];
            running_sum += prob;
            if (value < running_sum)
            {
                curr_max -= prob;
                selected.push_back(expert);
                return;
            }
        }
    }

    void setRouting(int* selected_experts, int64_t num_experts, int64_t k, int64_t num_tokens) override
    {
        TLLM_CHECK(num_experts == probabilities.size());
        std::vector<int> h_selected_experts(k * num_tokens, 0);
        std::vector<int> selected;
        float max = std::accumulate(probabilities.begin(), probabilities.end(), 0.0f);
        // TODO Put this on the GPU
        for (int token = 0; token < num_tokens; token++)
        {
            selected.clear();
            float curr_max = max;
            for (int choice = 0; choice < k; choice++)
            {
                doSample(curr_max, selected);
            }
            std::copy(selected.begin(), selected.end(), h_selected_experts.begin() + token * k);
        }
        check_cuda_error(cudaMemcpyAsync(selected_experts, h_selected_experts.data(),
            h_selected_experts.size() * sizeof(int), cudaMemcpyHostToDevice, streamPtr->get()));
        check_cuda_error(cudaStreamSynchronize(streamPtr->get()));
    }

    bool supportsConfig(int64_t num_experts_, int64_t, int64_t) const override
    {
        return num_experts == num_experts_;
    }
};

/**
 * Generates routing values by sampling a uniform distribution [-1,1)
 */
struct UniformRoutingConfig : public RoutingConfig
{
    std::mt19937_64 twister{0xD5};

    void start()
    {
        twister.seed(0xD5);
    }

    std::string getName() override
    {
        return "uniform";
    }

    bool isDeterministic() const override
    {
        return false;
    }

    void setRouting(int* selected_experts, int64_t num_experts, int64_t k, int64_t num_tokens) override
    {
        std::uniform_int_distribution<int> dist(0, num_experts - 1);
        std::vector<int> input(k * num_tokens);
        for (int i = 0; i < num_tokens; i++)
        {
            for (int j = 0; j < k; j++)
            {
                while (true)
                {
                    int expert_id = dist(twister);
                    bool valid = true;
                    for (int prev_j = 0; prev_j < j; prev_j++)
                    {
                        if (expert_id == input[i * k + prev_j])
                        {
                            valid = false;
                            break;
                        }
                    }
                    if (valid)
                    {
                        input[i * k + j] = expert_id;
                        break;
                    }
                }
            }
        }
        check_cuda_error(cudaMemcpyAsync(
            selected_experts, input.data(), input.size() * sizeof(int), cudaMemcpyHostToDevice, streamPtr->get()));
        check_cuda_error(cudaStreamSynchronize(streamPtr->get()));
    }

    bool supportsConfig(int64_t, int64_t, int64_t) const override
    {
        return true;
    }
};

/**
 * Stores a specific routing configuration
 */
struct VectoredRoutingConfig : public RoutingConfig
{
    using ElementType = int;
    std::vector<int> config;
    int64_t num_experts;
    int64_t k;
    std::string name;

    VectoredRoutingConfig(
        std::vector<ElementType> config, int64_t num_experts, int64_t k, std::string name = "vectored")
        : config(config)
        , num_experts(num_experts)
        , k(k)
        , name(name)
    {
    }

    std::string getName() override
    {
        return name;
    }

    bool isDeterministic() const override
    {
        return true;
    }

    void setRouting(int* selected_experts, int64_t num_experts, int64_t k, int64_t num_tokens) override
    {
        for (int64_t i = 0; i < num_tokens; i += config.size())
        {
            int64_t num_to_copy = std::min(num_tokens - i, (int64_t) config.size());
            check_cuda_error(cudaMemcpyAsync(
                selected_experts, config.data(), num_to_copy * sizeof(int), cudaMemcpyHostToDevice, streamPtr->get()));
            selected_experts += num_to_copy;
        }
        check_cuda_error(cudaStreamSynchronize(streamPtr->get()));
    }

    bool supportsConfig(int64_t num_experts_, int64_t k_, int64_t) const override
    {
        return num_experts_ == num_experts && k_ == k;
    }
};

}; // namespace

constexpr int LOAD_BALANCED_ROUTING_CONFIG = 0;
constexpr int UNIFORM_ROUTING_CONFIG = 1;
std::vector<std::shared_ptr<RoutingConfig>> routingConfigCache{
    std::static_pointer_cast<RoutingConfig>(std::make_shared<LoadBalancedRoutingConfig>()),
    std::static_pointer_cast<RoutingConfig>(std::make_shared<UniformRoutingConfig>()),
};

#ifdef ENABLE_FP8
using SafeFP8 = __nv_fp8_e4m3;
#else
using SafeFP8 = void;
#endif

#ifdef ENABLE_FP4
using SafeFP4 = __nv_fp4_e2m1;
#else
using SafeFP4 = void;
#endif

template <class TypeTuple_>
class MixtureOfExpertsBenchmark : public ::benchmark::Fixture
{
public:
    using DataType = typename TypeTuple_::DataType;
    using WeightType = typename TypeTuple_::WeightType;
    using OutputType = typename TypeTuple_::OutputType;
    constexpr static bool INT4 = std::is_same_v<WeightType, cutlass::uint4b_t>;
    constexpr static bool NVFP4 = std::is_same_v<DataType, SafeFP4> && std::is_same_v<WeightType, SafeFP4>;
    constexpr static bool FP8 = std::is_same_v<DataType, SafeFP8> && std::is_same_v<WeightType, SafeFP8>;
    constexpr static bool WFP4AFP8 = std::is_same_v<WeightType, SafeFP4> && std::is_same_v<DataType, SafeFP8>;
    constexpr static bool INT_QUANT = !std::is_same_v<DataType, WeightType>
        && (std::is_same_v<WeightType, cutlass::uint4b_t> || std::is_same_v<WeightType, uint8_t>);
    constexpr static bool ANY_FP4 = NVFP4 || WFP4AFP8;
    using InputType = std::conditional_t<NVFP4, OutputType, DataType>;
    using WeightStorage = std::conditional_t<INT_QUANT || ANY_FP4, uint8_t, WeightType>;
    constexpr static int WEIGHT_ELEM_PER_BYTE = (INT4 || ANY_FP4) ? 2 : 1;
    int const BASE_HIDDEN_SIZE = 64 / sizeof(WeightType) * WEIGHT_ELEM_PER_BYTE;

    constexpr static int64_t FP4_VECTOR_SIZE = NVFP4 ? TmaWarpSpecializedGroupedGemmInput::NVFP4BlockScaleVectorSize
                                                     : TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaleVectorSize;
    constexpr static int64_t MinNDimAlignment = NVFP4 ? TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentNVFP4
                                                      : TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX;
    constexpr static int64_t MinKDimAlignment = NVFP4 ? TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentNVFP4
                                                      : TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentMXFPX;

    std::vector<BufferManager::IBufferPtr> managed_buffers;
    int* mSelectedExperts{};
    DataType* mInputTensor{};

    int64_t mHiddenSize{};
    int64_t mNumExperts{};
    int64_t mNumExpertsPerNode{};
    int64_t mK{};

    constexpr static nvinfer1::DataType toDTypeID()
    {
        if (FP8 || WFP4AFP8)
            return nvinfer1::DataType::kFP8;
        if (NVFP4)
            return nvinfer1::DataType::kFP4;
        if (INT_QUANT && INT4)
            return nvinfer1::DataType::kINT4;
        if (INT_QUANT)
            return nvinfer1::DataType::kINT8;
        if (std::is_same_v<DataType, float>)
            return nvinfer1::DataType::kFLOAT;
        if (std::is_same_v<DataType, half>)
            return nvinfer1::DataType::kHALF;
#ifdef ENABLE_BF16
        if (std::is_same_v<DataType, nv_bfloat16>)
            return nvinfer1::DataType::kBF16;
#endif
        TLLM_THROW("Unrecognised format");
    };

    constexpr static nvinfer1::DataType toWTypeID()
    {
        if (FP8)
            return nvinfer1::DataType::kFP8;
        if (NVFP4 || WFP4AFP8)
            return nvinfer1::DataType::kFP4;
        if (INT_QUANT && INT4)
            return nvinfer1::DataType::kINT4;
        if (INT_QUANT)
            return nvinfer1::DataType::kINT8;
        if (std::is_same_v<DataType, float>)
            return nvinfer1::DataType::kFLOAT;
        if (std::is_same_v<DataType, half>)
            return nvinfer1::DataType::kHALF;
#ifdef ENABLE_BF16
        if (std::is_same_v<DataType, nv_bfloat16>)
            return nvinfer1::DataType::kBF16;
#endif
        TLLM_THROW("Unrecognised format");
    };

    template <class T>
    constexpr static auto typeToDtypeID()
    {
        if constexpr (std::is_same_v<T, SafeFP8>)
        {
            return nvinfer1::DataType::kFP8;
        }
        else if constexpr (std::is_same_v<T, SafeFP4>)
        {
            return nvinfer1::DataType::kFP4;
        }
        else if constexpr (std::is_same_v<T, uint8_t>)
        {
            return nvinfer1::DataType::kINT8;
        }
        else if constexpr (std::is_same_v<T, cutlass::uint4b_t>)
        {
            return nvinfer1::DataType::kINT4;
        }
        else if constexpr (std::is_same_v<T, nv_bfloat16>)
        {
            return nvinfer1::DataType::kBF16;
        }
        else if constexpr (std::is_same_v<T, half>)
        {
            return nvinfer1::DataType::kHALF;
        }
        else if constexpr (std::is_same_v<T, float>)
        {
            return nvinfer1::DataType::kFLOAT;
        }
        else
        {
            // sizeof(T) to make the static assert dependent on the template
            static_assert(sizeof(T) == 0, "Unrecognised data type");
        }
    }

    static bool shouldSkip()
    {
#ifndef ENABLE_FP8
        static_assert(!FP8, "FP8 Tests enabled on unsupported CUDA version");
#endif
#ifndef ENABLE_FP4
        static_assert(!ANY_FP4, "FP4 Tests enabled on unsupported CUDA version");
#endif
        bool should_skip_unsupported_fp8 = getSMVersion() < 89 && FP8;
        bool should_skip_unsupported_fp4 = (getSMVersion() < 100 || getSMVersion() >= 120) && ANY_FP4;
        return should_skip_unsupported_fp8 || should_skip_unsupported_fp4;
    }

    // Deprecated, just here to suppress warnings
    void SetUp(benchmark::State const& s) override
    {
        abort();
    }

    void TearDown(benchmark::State const& s) override
    {
        abort();
    }

    cudaEvent_t mStartEvent, mEndEvent;

    void SetUp(benchmark::State& s) override
    {
        assert(bufferManager);
        if (shouldSkip())
        {
            s.SkipWithMessage("GPU does not support dtype");
        }

        // Makes sure nothing from a previous iteration hangs around
        check_cuda_error(cudaDeviceSynchronize());
        check_cuda_error(cudaEventCreate(&mStartEvent));
        check_cuda_error(cudaEventCreate(&mEndEvent));
    }

    void TearDown(benchmark::State& s) override
    {
        managed_buffers.clear();

        check_cuda_error(cudaEventDestroy(mStartEvent));
        check_cuda_error(cudaEventDestroy(mEndEvent));
        check_cuda_error(cudaDeviceSynchronize());
    }

    CutlassMoeFCRunner<DataType, WeightType, OutputType, InputType> mMoERunner{};
    GemmProfilerBackend mGemmProfilerBackend{};
    char* mGemmProfilerWorkspace{};
    char* mWorkspace{};
    float* mScaleProbs{};
    WeightStorage* mExpertWeight1{};
    WeightStorage* mExpertWeight2{};
    DataType* mExpertIntScale1{};
    DataType* mExpertIntScale2{};

    float* mExpertFP8Scale1{};
    float* mExpertFP8Scale2{};
    float* mExpertFP8Scale3{};

    float* mExpertFP4ActScale1{};
    using ElementSF = TmaWarpSpecializedGroupedGemmInput::ElementSF;
    ElementSF* mExpertFP4WeightSf1{};
    float* mExpertFP4GlobalScale1{};
    float* mExpertFP4ActScale2{};
    ElementSF* mExpertFP4WeightSf2{};
    float* mExpertFP4GlobalScale2{};

    DataType* mExpertBias1{};
    DataType* mExpertBias2{};

    OutputType* mFinalOutput{};
    int* mSourceToExpandedMap;
    int64_t mInterSize{};
    int64_t mTotalTokens{};

    int mRoutingConfigIndex = 0;

    bool mUseBias = true;
    bool mUseFinalScale = true;
    bool mIsGated = false;
    int mGatedMultiplier = 1;

    ActivationType mActType = ActivationType::Relu;

    constexpr static int64_t NUM_BUFFERS = 32;
    int64_t mNumWorkspaceBuffers = NUM_BUFFERS;
    int64_t mNumInputBuffers = NUM_BUFFERS;
    int64_t mNumGemmProfilerBuffers = NUM_BUFFERS;

    std::array<QuantParams, NUM_BUFFERS> mQuantParams{};
    bool mUseLora = false;
    bool mUsePrequantScale = false;
    int mGroupSize = -1;
    std::array<LoraParams, NUM_BUFFERS> mLoraParams{};

    std::optional<tensorrt_llm::cutlass_extensions::CutlassGemmConfig> mSelectedConfig = std::nullopt;

    int64_t mBufferIndex = 0;
    size_t mGemmProfilerWorkspaceSize = 0;
    size_t mWorkspaceSize = 0;
    size_t mExpertWeight1Size = 0;
    size_t mExpertWeight2Size = 0;
    size_t mExpertBias1Size = 0;
    size_t mExpertBias2Size = 0;
    size_t mInputTensorSize = 0;
    size_t mFinalOutputSize = 0;
    size_t mSourceToExpandedMapSize = 0;
    size_t mScaleProbsSize = 0;
    size_t mSelectedExpertsSize = 0;
    size_t mExpertFP4WeightSf1Size = 0;
    size_t mExpertFP4WeightSf2Size = 0;
    size_t mExpertIntScale1Size = 0;
    size_t mExpertIntScale2Size = 0;

    size_t padSize(size_t size)
    {
        return ceilDiv(size, 128) * 128;
    }

    template <class T>
    T* allocBuffer(size_t size)
    {
        size_t size_padded = padSize(size) * sizeof(T);
        auto i_buffer = bufferManager->gpu(size_padded);
        check_cuda_error(cudaGetLastError());
        managed_buffers.emplace_back(std::move(i_buffer));
        T* ptr = static_cast<T*>(managed_buffers.back()->data());
        populateRandomBuffer(ptr, size_padded, streamPtr->get());
        return ptr;
    }

    void initBuffersPermute(int64_t num_tokens, int64_t hidden_size, int64_t inter_size, int64_t num_experts, int64_t k,
        int64_t routing_config, MOEParallelismConfig parallelism_config, GemmToProfile gemm_to_profile)
    {
        assert(hidden_size % BASE_HIDDEN_SIZE == 0);

        managed_buffers.clear();

        mTotalTokens = num_tokens;
        mHiddenSize = hidden_size;
        mInterSize = inter_size / parallelism_config.tp_size;
        mNumExperts = num_experts;
        mNumExpertsPerNode = num_experts / parallelism_config.ep_size;
        mK = k;
        mIsGated = isGatedActivation(mActType);
        mGatedMultiplier = mIsGated ? 2 : 1;
        auto const gated_inter = mInterSize * mGatedMultiplier;
        size_t const expert_matrix_size = padSize(mNumExpertsPerNode * mHiddenSize * mInterSize);

        bool need_weight_1 = gemm_to_profile == GemmToProfile::GEMM_1 || gemm_to_profile == GemmToProfile::LAYER;
        bool need_weight_2 = gemm_to_profile == GemmToProfile::GEMM_2 || gemm_to_profile == GemmToProfile::LAYER;
        mExpertWeight1Size = need_weight_1 ? expert_matrix_size * mGatedMultiplier / WEIGHT_ELEM_PER_BYTE : 0;
        mExpertWeight2Size = need_weight_2 ? expert_matrix_size / WEIGHT_ELEM_PER_BYTE : 0;
        mExpertWeight1 = need_weight_1 ? allocBuffer<WeightStorage>(mExpertWeight1Size * NUM_BUFFERS) : nullptr;
        mExpertWeight2 = need_weight_2 ? allocBuffer<WeightStorage>(mExpertWeight2Size * NUM_BUFFERS) : nullptr;

        if (gemm_to_profile == GemmToProfile::LAYER)
        {
            mWorkspaceSize = mMoERunner.getWorkspaceSize(mTotalTokens, mHiddenSize, mInterSize, mNumExperts, mK,
                mActType, parallelism_config, mUseLora, /*use_deepseek_fp8_block_scale=*/false,
                /*min_latency_mode=*/false, mUsePrequantScale);

            mNumWorkspaceBuffers = mWorkspaceSize > 1024 * 1024 * 1024 ? 2 : NUM_BUFFERS;
            mWorkspace = allocBuffer<char>(mWorkspaceSize * mNumWorkspaceBuffers);

            mExpertBias1 = nullptr;
            mExpertBias2 = nullptr;
            if (mUseBias)
            {
                mExpertBias1Size = padSize(mNumExpertsPerNode * gated_inter);
                mExpertBias2Size = padSize(mNumExpertsPerNode * mHiddenSize);
                mExpertBias1 = allocBuffer<DataType>(mExpertBias1Size * NUM_BUFFERS);
                mExpertBias2 = allocBuffer<DataType>(mExpertBias2Size * NUM_BUFFERS);
            }

            if constexpr (INT_QUANT)
            {
                mExpertIntScale1Size = padSize(mNumExpertsPerNode * gated_inter);
                mExpertIntScale2Size = padSize(mNumExpertsPerNode * mHiddenSize);
                mExpertIntScale1 = allocBuffer<DataType>(mExpertIntScale1Size * NUM_BUFFERS);
                mExpertIntScale2 = allocBuffer<DataType>(mExpertIntScale2Size * NUM_BUFFERS);

                for (int i = 0; i < NUM_BUFFERS; i++)
                {
                    mQuantParams[i] = QuantParams::Int(
                        mExpertIntScale1 + mExpertIntScale1Size * i, mExpertIntScale2 + mExpertIntScale2Size * i);
                }
            }
            else if constexpr (FP8)
            {
                mExpertFP8Scale1 = allocBuffer<float>(mNumExpertsPerNode);
                mExpertFP8Scale2 = allocBuffer<float>(1);
                mExpertFP8Scale3 = allocBuffer<float>(mNumExpertsPerNode);

                for (int i = 0; i < NUM_BUFFERS; i++)
                {
                    mQuantParams[i] = QuantParams::FP8(mExpertFP8Scale1, mExpertFP8Scale2, mExpertFP8Scale3);
                }
            }
            else if constexpr (ANY_FP4)
            {
                mExpertFP4ActScale1 = allocBuffer<float>(mNumExpertsPerNode);
                mExpertFP4WeightSf1Size = mNumExpertsPerNode
                    * TmaWarpSpecializedGroupedGemmInput::alignToSfDim(gated_inter, MinNDimAlignment)
                    * TmaWarpSpecializedGroupedGemmInput::alignToSfDim(mHiddenSize, MinKDimAlignment) / FP4_VECTOR_SIZE;
                mExpertFP4WeightSf1 = allocBuffer<ElementSF>(mExpertFP4WeightSf1Size * NUM_BUFFERS);
                mExpertFP4GlobalScale1 = allocBuffer<float>(mNumExpertsPerNode);

                mExpertFP4ActScale2 = allocBuffer<float>(mNumExpertsPerNode);
                mExpertFP4WeightSf2Size = mNumExpertsPerNode
                    * TmaWarpSpecializedGroupedGemmInput::alignToSfDim(mInterSize, MinNDimAlignment)
                    * TmaWarpSpecializedGroupedGemmInput::alignToSfDim(mHiddenSize, MinKDimAlignment) / FP4_VECTOR_SIZE;
                mExpertFP4WeightSf2 = allocBuffer<ElementSF>(mExpertFP4WeightSf2Size * NUM_BUFFERS);
                mExpertFP4GlobalScale2 = allocBuffer<float>(mNumExpertsPerNode);

                auto func = NVFP4 ? QuantParams::FP4 : QuantParams::FP8MXFP4;
                for (int i = 0; i < NUM_BUFFERS; i++)
                {
                    mQuantParams[i] = func(mExpertFP4ActScale1, mExpertFP4WeightSf1 + mExpertFP4WeightSf1Size * i,
                        mExpertFP4GlobalScale1, mExpertFP4ActScale2, mExpertFP4WeightSf2 + mExpertFP4WeightSf2Size * i,
                        mExpertFP4GlobalScale2, false, false);
                }
            }

            mSelectedExpertsSize = padSize(mTotalTokens * mK);
            mSelectedExperts = allocBuffer<int>(mSelectedExpertsSize * NUM_BUFFERS);
            mScaleProbsSize = padSize(mTotalTokens * mK);
            mScaleProbs = allocBuffer<float>(mScaleProbsSize * NUM_BUFFERS);
            mInputTensorSize = padSize(mTotalTokens * mHiddenSize);
            mNumInputBuffers = mInputTensorSize > 1024 * 1024 * 1024 ? 2 : NUM_BUFFERS;
            mInputTensor = allocBuffer<DataType>(mInputTensorSize * mNumInputBuffers);
            mFinalOutputSize = padSize(mTotalTokens * mHiddenSize);
            mFinalOutput = allocBuffer<OutputType>(mFinalOutputSize * mNumInputBuffers);

            mSourceToExpandedMapSize = padSize(mTotalTokens * mK);
            mSourceToExpandedMap = allocBuffer<int>(mSourceToExpandedMapSize * NUM_BUFFERS);
            mRoutingConfigIndex = routing_config;
            auto tactic = routingConfigCache.at(routing_config);
            tactic->start();
            for (int i = 0; i < NUM_BUFFERS; i++)
            {
                tactic->setRouting(mSelectedExperts + mSelectedExpertsSize * i, mNumExperts, mK, mTotalTokens);
            }
        }

#ifdef USING_OSS_CUTLASS_MOE_GEMM
        mGemmProfilerBackend.init(mMoERunner, GemmProfilerBackend::GemmToProfile::Undefined, typeToDtypeID<DataType>(),
            typeToDtypeID<WeightType>(), typeToDtypeID<OutputType>(), mNumExperts, mK, mHiddenSize, mHiddenSize,
            mInterSize, mGroupSize, mActType, mUseBias, mUseLora, /*min_latency_mode=*/false,
            /*need_weights=*/false, parallelism_config, /*enable_alltoall=*/false);
#else
        mGemmProfilerBackend.init(mMoERunner, GemmProfilerBackend::GemmToProfile::Undefined, typeToDtypeID<DataType>(),
            typeToDtypeID<WeightType>(), typeToDtypeID<OutputType>(), mNumExperts, mK, mHiddenSize, mHiddenSize,
            mInterSize, mGroupSize, mActType, mUseBias, mUseLora, /*min_latency_mode=*/false,
            /*need_weights=*/false, parallelism_config);
#endif

        mGemmProfilerWorkspaceSize = 0;
        if (gemm_to_profile == GemmToProfile::GEMM_1 || gemm_to_profile == GemmToProfile::LAYER)
        {
            mGemmProfilerBackend.mGemmToProfile = GemmProfilerBackend::GemmToProfile::GEMM_1;
            mGemmProfilerWorkspaceSize
                = std::max(mGemmProfilerWorkspaceSize, mGemmProfilerBackend.getWorkspaceSize(mTotalTokens));
        }

        if (gemm_to_profile == GemmToProfile::GEMM_2 || gemm_to_profile == GemmToProfile::LAYER)
        {
            mGemmProfilerBackend.mGemmToProfile = GemmProfilerBackend::GemmToProfile::GEMM_2;
            mGemmProfilerWorkspaceSize
                = std::max(mGemmProfilerWorkspaceSize, mGemmProfilerBackend.getWorkspaceSize(mTotalTokens));
        }

        mGemmProfilerWorkspaceSize = padSize(mGemmProfilerWorkspaceSize);
        mNumGemmProfilerBuffers = mGemmProfilerWorkspaceSize > 1024 * 1024 * 1024 ? 2 : NUM_BUFFERS;
        mNumGemmProfilerBuffers = gemm_to_profile == GemmToProfile::LAYER ? 1 : mNumGemmProfilerBuffers;
        mGemmProfilerWorkspace = mGemmProfilerWorkspaceSize > 0
            ? allocBuffer<char>(mGemmProfilerWorkspaceSize * mNumGemmProfilerBuffers)
            : nullptr;

        check_cuda_error(cudaStreamSynchronize(streamPtr->get()));
    }

    void prepareGemmProfiler(GemmToProfile gemm_to_profile)
    {
        if (gemm_to_profile == GemmToProfile::LAYER)
            return;
        mGemmProfilerBackend.mGemmToProfile = static_cast<GemmProfilerBackend::GemmToProfile>(gemm_to_profile);
        auto* expert_weights = gemm_to_profile == GemmToProfile::GEMM_1 ? mExpertWeight1 : mExpertWeight2;
        auto expert_weights_size = gemm_to_profile == GemmToProfile::GEMM_1 ? mExpertWeight1Size : mExpertWeight2Size;
        mGemmProfilerBackend.prepare(mTotalTokens,
            mGemmProfilerWorkspace + mGemmProfilerWorkspaceSize * (mBufferIndex % mNumGemmProfilerBuffers),
            /*expert_weights=*/expert_weights + expert_weights_size * mBufferIndex, streamPtr->get());
    }

    std::array<cudaGraph_t, NUM_BUFFERS> mGraph{};

    std::array<cudaGraphExec_t, NUM_BUFFERS> mGraphInstance{};

    void createGraph(MOEParallelismConfig parallelism_config, GemmToProfile gemm_to_profile)
    {
        if (!useCudaGraph)
            return;

        NVTX3_SCOPED_RANGE(BuildGraph);

        for (int i = 0; i < NUM_BUFFERS; i++)
        {
            mBufferIndex = i;
            // Each buffer will have a different routing config for the gemm profiler
            prepareGemmProfiler(gemm_to_profile);
            check_cuda_error(cudaGraphCreate(&mGraph[i], 0));
            check_cuda_error(cudaStreamBeginCapture(streamPtr->get(), cudaStreamCaptureModeThreadLocal));
            runMoEPermute(parallelism_config, gemm_to_profile);
            check_cuda_error(cudaStreamEndCapture(streamPtr->get(), &mGraph[i]));
            check_cuda_error(cudaGraphInstantiate(&mGraphInstance[i], mGraph[i], nullptr, nullptr, 0));
        }
    }

    void destroyGraph()
    {
        if (!useCudaGraph)
            return;

        NVTX3_SCOPED_RANGE(DestroyGraph);

        for (int i = 0; i < NUM_BUFFERS; i++)
        {
            check_cuda_error(cudaGraphExecDestroy(mGraphInstance[i]));
            check_cuda_error(cudaGraphDestroy(mGraph[i]));
        }
    }

    float benchmarkLoop(MOEParallelismConfig parallelism_config, GemmToProfile gemm_to_profile)
    {
        mBufferIndex = (mBufferIndex + 1) % NUM_BUFFERS;

        // Setup the profiler state for this iteration. CUDA Graphs will do this when it captures the graph.
        if (gemm_to_profile != GemmToProfile::LAYER && !useCudaGraph)
        {
            prepareGemmProfiler(gemm_to_profile);
        }
        else if (gemm_to_profile == GemmToProfile::LAYER)
        {
            auto tactic = routingConfigCache.at(mRoutingConfigIndex);
            if (!tactic->isDeterministic())
            {
                tactic->setRouting(
                    mSelectedExperts + mSelectedExpertsSize * mBufferIndex, mNumExperts, mK, mTotalTokens);
            }
        }

        {
            NVTX3_SCOPED_RANGE(BenchmarkLoopIteration);
            check_cuda_error(cudaEventRecord(mStartEvent, streamPtr->get()));
            if (useCudaGraph)
            {
                cudaGraphLaunch(mGraphInstance[mBufferIndex], streamPtr->get());
            }
            else
            {
                runMoEPermute(parallelism_config, gemm_to_profile);
            }
            check_cuda_error(cudaEventRecord(mEndEvent, streamPtr->get()));
            check_cuda_error(cudaStreamSynchronize(streamPtr->get()));
        }

        float ms;
        check_cuda_error(cudaEventElapsedTime(&ms, mStartEvent, mEndEvent));
        return ms;
    }

    // An imprecise benchmark pass for picking the best tactic.
    // Runs for 3 iterations or 1 second and picks the best option
    int pickBestTactic(MOEParallelismConfig parallelism_config, GemmToProfile gemm_to_profile)
    {
        auto tactics = mMoERunner.getTactics(static_cast<MoeGemmId>(gemm_to_profile));
        ::nvtx3::scoped_range nvtx(tensorrt_llm::common::nvtx::nextColor(),
            "Tactic Profiling GEMM " + std::to_string(static_cast<int>(gemm_to_profile)));
        // We save space by reusing the same workspace buffer for all tactics when doing full layer profiling. So we
        // need to hardcode the buffer index to 0.
        auto old_buffer_index = mBufferIndex;
        mBufferIndex = 0;
        prepareGemmProfiler(gemm_to_profile);
        mBufferIndex = old_buffer_index;

        auto* mGemmProfilerExpertWeights = gemm_to_profile == GemmToProfile::GEMM_1 ? mExpertWeight1 : mExpertWeight2;

        float best_time = INFINITY;
        int best_idx = -1;
        for (int tidx = 0; tidx < tactics.size(); tidx++)
        {
            ::nvtx3::scoped_range nvtx(
                tensorrt_llm::common::nvtx::nextColor(), "Tactic Profiling Tactic Index: " + std::to_string(tidx));
            try
            {
                // Set the tactic
                auto const& t = tactics[tidx];

                {
                    ::nvtx3::scoped_range nvtx(tensorrt_llm::common::nvtx::nextColor(), "Tactic Profiling Warm-Up");
                    // Warm-Up run
                    mGemmProfilerBackend.runProfiler(mTotalTokens, t, mGemmProfilerWorkspace,
                        /*expert_weights=*/mGemmProfilerExpertWeights, streamPtr->get());
                    check_cuda_error(cudaStreamSynchronize(streamPtr->get()));
                }

                // Profile all samples or for 1 sec
                int const max_iters = mGemmProfilerBackend.NUM_ROUTING_SAMPLES * 2;
                float const max_time_ms = 1000.f;

                float time = 0.f;
                int iter = 0;
                while (iter < max_iters && time < max_time_ms)
                {
                    {
                        ::nvtx3::scoped_range nvtx(tensorrt_llm::common::nvtx::nextColor(),
                            "Tactic Profiling Iteration " + std::to_string(iter));

                        check_cuda_error(cudaEventRecord(mStartEvent, streamPtr->get()));
                        mGemmProfilerBackend.runProfiler(mTotalTokens, t, mGemmProfilerWorkspace,
                            /*expert_weights=*/mGemmProfilerExpertWeights, streamPtr->get());
                        check_cuda_error(cudaEventRecord(mEndEvent, streamPtr->get()));
                        check_cuda_error(cudaStreamSynchronize(streamPtr->get()));
                    }

                    float ms;
                    check_cuda_error(cudaEventElapsedTime(&ms, mStartEvent, mEndEvent));
                    time += ms;

                    iter++;
                }
                // Get average time per iteration
                time /= static_cast<float>(iter);

                if (LOG_LEVEL >= VERBOSE)
                {
                    std::cout << "Tactic " << tidx << " for GEMM" << (int) gemm_to_profile << ":\n"
                              << t.toString() << "\ntook: " << time << "ms\n";
                }

                // Update the best
                if (time < best_time)
                {
                    best_idx = tidx;
                    best_time = time;
                }
            }
            catch (std::exception const& e)
            {
                // Sync to tidy up
                if (LOG_LEVEL >= ERROR)
                    std::cout << "Tactic failed to run with: " << e.what() << std::endl;
                check_cuda_error(cudaDeviceSynchronize());
                // skip invalid tactic
                continue;
            }
        }

        return best_idx;
    }

    int mBestTacticGemm1 = -1;
    int mBestTacticGemm2 = -1;

    std::pair<int, int> setTactic(
        int tactic_idx1, int tactic_idx2, MOEParallelismConfig parallelism_config, GemmToProfile gemm_to_profile)
    {
        auto tactics1 = mMoERunner.getTactics(MoeGemmId::GEMM_1);
        auto tactics2 = mMoERunner.getTactics(MoeGemmId::GEMM_2);
        std::vector<std::pair<std::reference_wrapper<int>, GemmToProfile>> tactics_to_profile{
            {tactic_idx1, GemmToProfile::GEMM_1}, {tactic_idx2, GemmToProfile::GEMM_2}};
        for (auto& combo : tactics_to_profile)
        {
            auto& t = combo.first.get();
            auto& tactics = combo.second == GemmToProfile::GEMM_1 ? tactics1 : tactics2;
            if (combo.second != gemm_to_profile && gemm_to_profile != GemmToProfile::LAYER)
            {
                t = 0; // Unneeded tactic, set to 0
                continue;
            }
            if (t == -1)
            {
                t = pickBestTactic(parallelism_config, combo.second);
            }

            if (t < 0 || t >= tactics.size())
            {
                return {-1, -1};
            }
        }

        mMoERunner.setTactic(tactics1[tactic_idx1], tactics2[tactic_idx2]);
        mBestTacticGemm1 = tactic_idx1;
        mBestTacticGemm2 = tactic_idx2;
        return {tactic_idx1, tactic_idx2};
    }

    void runMoEPermute(MOEParallelismConfig parallelism_config, GemmToProfile gemm_to_profile)
    {
        switch (gemm_to_profile)
        {
        case GemmToProfile::GEMM_1:
        case GemmToProfile::GEMM_2:
        {
            auto tactic_idx = gemm_to_profile == GemmToProfile::GEMM_1 ? mBestTacticGemm1 : mBestTacticGemm2;
            auto* expert_weights = gemm_to_profile == GemmToProfile::GEMM_1 ? mExpertWeight1 : mExpertWeight2;
            auto expert_weights_size
                = gemm_to_profile == GemmToProfile::GEMM_1 ? mExpertWeight1Size : mExpertWeight2Size;

            auto tactics = mMoERunner.getTactics(static_cast<MoeGemmId>(gemm_to_profile))[tactic_idx];
            if (static_cast<int>(gemm_to_profile) != static_cast<int>(mGemmProfilerBackend.mGemmToProfile))
            {
                throw std::runtime_error("Configuration mismatch between mGemmProfilerBackend and runMoEPermute");
            }
            mGemmProfilerBackend.mSampleIndex = mBufferIndex % mGemmProfilerBackend.NUM_ROUTING_SAMPLES;
            mGemmProfilerBackend.runProfiler(mTotalTokens, tactics,
                mGemmProfilerWorkspace + mGemmProfilerWorkspaceSize * (mBufferIndex % mNumGemmProfilerBuffers),
                /*expert_weights=*/expert_weights + expert_weights_size * mBufferIndex, streamPtr->get());
            break;
        }
        case GemmToProfile::LAYER:
        {
            auto stream = streamPtr->get();
            MoeMinLatencyParams min_latency_params;
#ifdef USING_OSS_CUTLASS_MOE_GEMM
            mMoERunner.runMoe(mInputTensor + mInputTensorSize * (mBufferIndex % mNumInputBuffers), nullptr, true,
                mSelectedExperts + mSelectedExpertsSize * mBufferIndex,
                mUseFinalScale ? mScaleProbs + mScaleProbsSize * mBufferIndex : nullptr,
                mExpertWeight1 + mExpertWeight1Size * mBufferIndex, mExpertBias1 + mExpertBias1Size * mBufferIndex,
                ActivationParams(mActType), mExpertWeight2 + mExpertWeight2Size * mBufferIndex,
                mExpertBias2 + mExpertBias2Size * mBufferIndex, mQuantParams[mBufferIndex], mTotalTokens, mTotalTokens,
                mHiddenSize, mHiddenSize, mInterSize, mNumExperts, mK,
                mWorkspace + mWorkspaceSize * (mBufferIndex % mNumWorkspaceBuffers),
                mFinalOutput + mFinalOutputSize * (mBufferIndex % mNumInputBuffers),
                mSourceToExpandedMap + mSourceToExpandedMapSize * mBufferIndex, parallelism_config,
                /*enable_alltoall=*/false, mUseLora, mLoraParams[mBufferIndex],
                /*use_fp8_block_scaling=*/false, /*min_latency_mode=*/false, min_latency_params, stream);
#else
            mMoERunner.runMoe(mInputTensor + mInputTensorSize * (mBufferIndex % mNumInputBuffers), nullptr, true,
                mSelectedExperts + mSelectedExpertsSize * mBufferIndex,
                mUseFinalScale ? mScaleProbs + mScaleProbsSize * mBufferIndex : nullptr,
                mExpertWeight1 + mExpertWeight1Size * mBufferIndex, mExpertBias1 + mExpertBias1Size * mBufferIndex,
                ActivationParams(mActType), mExpertWeight2 + mExpertWeight2Size * mBufferIndex,
                mExpertBias2 + mExpertBias2Size * mBufferIndex, mQuantParams[mBufferIndex], mTotalTokens, mTotalTokens,
                mHiddenSize, mHiddenSize, mInterSize, mNumExperts, mK,
                mWorkspace + mWorkspaceSize * (mBufferIndex % mNumWorkspaceBuffers),
                mFinalOutput + mFinalOutputSize * (mBufferIndex % mNumInputBuffers),
                mSourceToExpandedMap + mSourceToExpandedMapSize * mBufferIndex, parallelism_config,
                /*enable_alltoall=*/false, mUseLora, mLoraParams[mBufferIndex],
                /*use_fp8_block_scaling=*/false, /*min_latency_mode=*/false, min_latency_params, stream);
#endif
            break;
        }
        }
    }

    void runBenchmark(benchmark::State& state);
};

template <class TypeTuple_>
void MixtureOfExpertsBenchmark<TypeTuple_>::runBenchmark(benchmark::State& state)
{
    NVTX3_SCOPED_RANGE(FullBenchmark);
    int const num_experts = state.range(0);
    int const top_k = state.range(1);
    int const hidden_size = state.range(2);
    int const inter_size = state.range(3);
    int const tp_size = state.range(4);
    int const ep_size = state.range(5);
    int const world_rank = state.range(6);
    int const num_tokens = state.range(7);
    mUseBias = state.range(8);
    mUseFinalScale = state.range(9);
    mActType = static_cast<ActivationType>(state.range(10));
    int tactic_idx1 = state.range(11);
    int tactic_idx2 = state.range(12);
    int const routing_config = state.range(13);
    GemmToProfile const gemm_to_profile = static_cast<GemmToProfile>(state.range(14));

    state.counters["num_experts"] = num_experts;
    state.counters["top_k"] = top_k;
    state.counters["hidden_size"] = hidden_size;
    state.counters["inter_size"] = inter_size;
    state.counters["tp_size"] = tp_size;
    state.counters["ep_size"] = ep_size;
    state.counters["world_rank"] = world_rank;
    state.counters["num_tokens"] = num_tokens;
    state.counters["use_bias"] = (int) mUseBias;
    state.counters["use_final_scale"] = (int) mUseFinalScale;
    state.counters["act_fn"] = (int) mActType;
    state.counters["routing_config"] = (int) routing_config;
    state.counters["dtype"] = (int) toDTypeID();
    state.counters["wtype"] = (int) toWTypeID();
    state.counters["gemm_to_profile"] = (int) gemm_to_profile;

    std::stringstream ss;
    ss << "Experts,K,Hidden,Inter,TP,EP,Rank,Tokens,Bias,Scale,Actfn,Tactic1,Tactic2,Gemm,Routing=";
    for (auto v : {num_experts, top_k, hidden_size, inter_size, tp_size, ep_size, world_rank, num_tokens,
             (int) mUseBias, (int) mUseFinalScale, (int) mActType, tactic_idx1, tactic_idx2, (int) gemm_to_profile})
    {
        ss << v << ",";
    }
    ss << routingConfigCache.at(routing_config)->getName();
    // state.SetLabel(ss.str());
    state.SetLabel(routingConfigCache.at(routing_config)->getName());

    // Always use EP size for moe config until we support TP+EP, we just divide the inter size for TP
    MOEParallelismConfig parallelism_config{tp_size, world_rank / ep_size, ep_size, world_rank % ep_size};
    initBuffersPermute(
        num_tokens, hidden_size, inter_size, num_experts, top_k, routing_config, parallelism_config, gemm_to_profile);

    // Parse the tactic, does checks for "auto" mode and out of range
    std::tie(tactic_idx1, tactic_idx2) = setTactic(tactic_idx1, tactic_idx2, parallelism_config, gemm_to_profile);
    if (tactic_idx1 < 0 || tactic_idx2 < 0)
    {
        state.SkipWithMessage("Out of range tactic");
        return;
    }
    auto tactics1 = mMoERunner.getTactics(MoeGemmId::GEMM_1);
    auto tactics2 = mMoERunner.getTactics(MoeGemmId::GEMM_2);
    if (LOG_LEVEL >= INFO)
    {
        std::cout << "Selected tactic #1: " << tactic_idx1 << "/" << tactics1.size() << "\n"
                  << tactics1[tactic_idx1].toString() << std::endl;
        std::cout << "Selected tactic #2: " << tactic_idx2 << "/" << tactics2.size() << "\n"
                  << tactics2[tactic_idx2].toString() << std::endl;
    }
    state.counters["tactic_idx1"] = tactic_idx1;
    state.counters["tactic_idx2"] = tactic_idx2;

    state.counters["t1_sm_version"] = tactics1[tactic_idx1].sm_version;
    state.counters["t1_tile_shape"] = tactics1[tactic_idx1].getTileConfigAsInt();
    state.counters["t1_cluster_shape"] = (int) tactics1[tactic_idx1].cluster_shape;
    state.counters["t1_dynamic_cluster_shape"] = (int) tactics1[tactic_idx1].dynamic_cluster_shape;
    state.counters["t1_fallback_cluster_shape"] = (int) tactics1[tactic_idx1].fallback_cluster_shape;
    state.counters["t1_epilogue_schedule"] = (int) tactics1[tactic_idx1].epilogue_schedule;

    state.counters["t2_sm_version"] = tactics2[tactic_idx2].sm_version;
    state.counters["t2_tile_shape"] = tactics2[tactic_idx2].getTileConfigAsInt();
    state.counters["t2_cluster_shape"] = (int) tactics2[tactic_idx2].cluster_shape;
    state.counters["t2_dynamic_cluster_shape"] = (int) tactics2[tactic_idx2].dynamic_cluster_shape;
    state.counters["t2_fallback_cluster_shape"] = (int) tactics2[tactic_idx2].fallback_cluster_shape;
    state.counters["t2_epilogue_schedule"] = (int) tactics2[tactic_idx2].epilogue_schedule;

    createGraph(parallelism_config, gemm_to_profile);

    {
        ::nvtx3::scoped_range nvtx(tensorrt_llm::common::nvtx::nextColor(), "BenchmarkRun " + ss.str());
        for (auto _ : state)
        {
            float ms = benchmarkLoop(parallelism_config, gemm_to_profile);
            state.SetIterationTime(ms / 1000.f);
        }
    }

    destroyGraph();

    state.SetItemsProcessed(state.iterations() * num_tokens);

    // Cleanup all the benchmark state
    managed_buffers.clear();
    check_cuda_error(cudaDeviceSynchronize());
}