// Copyright © 2025 Advanced Micro Devices, Inc.
// SPDX-License-Identifier: MIT

// clang-format off
#pragma once

#include <aotriton/config.h>
#include <aotriton/dtypes.h>
#include <aotriton/runtime.h>
#include <aotriton/util.h>
#include <aotriton/_internal/lazy_tensor_internal.h>
#include <functional>
#include <string>
#include <vector>
#include "aotriton/_internal/flash/aiter.h"

namespace AOTRITON_NS::v3::flash {
    struct OpAttnFwdParams;
}

namespace AOTRITON_NS::v3::flash {

using AOTRITON_NS::v3::flash::OpAttnFwdParams;

struct AiterFmhaV3FwdContext {
    const OpAttnFwdParams *params = nullptr;
    const char* check_inputs_are_supported(Gpu);
#if 1
    mutable aiter::mha_fwd_args cookie;
#endif
    // As a thin wrapper, SlimAffine kernels does not understand operator
    // (ParentContext) or Conditional Kernel (condition).
    template <typename ParentContext>
    AiterFmhaV3FwdContext(const ParentContext& pcontext, bool condition) {
        params = pcontext.params;
    }

    hipError_t lookup_optimal(Gpu gpu);
    hipError_t launch(hipStream_t stream) const;

    static std::tuple<int, int> get_archmod_number(Gpu gpu);
};

}

// vim: set fileencoding=utf-8


