#pragma once
// SPDX-License-Identifier: MIT
// Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
#include <torch/extension.h>

namespace aiter {
namespace torch_itfs {
std::vector<at::Tensor>
mha_batch_prefill(at::Tensor& q,                  // [total_q, hq, d]
                  const at::Tensor& k,            // [num_blocks, hk, d/8, block_size, 8]
                  const at::Tensor& v,            // [num_blocks, hk, block_size/8, d, 8]
                  const at::Tensor& cu_seqlens_q, // [b+1]
                  const at::Tensor& kv_indptr,    // [b+1]
                  const at::Tensor& kv_page_indices,
                  int max_seqlen_q,
                  int max_seqlen_k,
                  float p_dropout,
                  float softmax_scale,
                  float logits_soft_cap,
                  bool zero_tensors,
                  bool is_causal,
                  int window_size_left,
                  int window_size_right,
                  bool return_softmax_lse,
                  bool return_dropout_randval,
                  std::optional<at::Tensor> out_,                // [total_q, hq, d]
                  std::optional<const at::Tensor> bias_,         // [total_q, max_seqlen_k]
                  std::optional<const at::Tensor> alibi_slopes_, // [hq] or [b, hq]
                  // Per-tensor descale for PERTENSOR mode (Q/K/V each have one scale value)
                  std::optional<const at::Tensor> q_descale, // [1] per-tensor Q descale
                  std::optional<const at::Tensor> k_descale, // [1] per-tensor K descale
                  std::optional<const at::Tensor> v_descale, // [1] per-tensor V descale
                  // Per-page descale for KV_BLOCKSCALE mode (Q per-tensor, K/V per-page)
                  // Mutually exclusive with k_descale/v_descale
                  std::optional<const at::Tensor> kv_block_descale, // [num_block, num_kv_head, 2]
                  std::optional<const at::Tensor> kv_last_page_lens,
                  std::optional<const at::Tensor> block_table,
                  std::optional<const at::Tensor> seqlen_k,
                  std::optional<const at::Tensor> sink_ptr_, // [hq];
                  std::optional<at::Generator> gen_);

} // namespace torch_itfs
} // namespace aiter
