mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-16 09:37:25 +08:00
188b7f9b8c
Signed-off-by: charlifu <charlifu@amd.com> Co-authored-by: Tyler Michael Smith <tysmith@redhat.com>
25 lines
1.1 KiB
C++
25 lines
1.1 KiB
C++
#pragma once
|
|
|
|
#include <torch/all.h>
|
|
|
|
torch::Tensor LLMM1(at::Tensor& in_a, at::Tensor& in_b,
|
|
const int64_t rows_per_block);
|
|
|
|
torch::Tensor wvSplitK(at::Tensor& in_a, at::Tensor& in_b,
|
|
const int64_t CuCount);
|
|
|
|
void wvSplitKQ(at::Tensor& in_a, at::Tensor& in_b, at::Tensor& out_c,
|
|
at::Tensor& scale_a, at::Tensor& scale_b, const int64_t CuCount);
|
|
|
|
void paged_attention(torch::Tensor& out, torch::Tensor& exp_sums,
|
|
torch::Tensor& max_logits, torch::Tensor& tmp_out,
|
|
torch::Tensor& query, torch::Tensor& key_cache,
|
|
torch::Tensor& value_cache, int64_t num_kv_heads,
|
|
double scale, torch::Tensor& block_tables,
|
|
torch::Tensor& context_lens,
|
|
const std::optional<torch::Tensor>& query_start_loc,
|
|
int64_t block_size, int64_t max_context_len,
|
|
const std::optional<torch::Tensor>& alibi_slopes,
|
|
const std::string& kv_cache_dtype, torch::Tensor& k_scale,
|
|
torch::Tensor& v_scale);
|