142 lines
5.4 KiB
C++
142 lines
5.4 KiB
C++
/* ilu_ops_api.h — Standalone header for project_6 ex_engine.
|
|
*
|
|
* Adapted from xllm/core/kernels/ilu/ilu_ops_api.h.
|
|
* Removes xllm-internal deps (glog, kernels/kernels.h, framework/*).
|
|
* Only requires: torch, ixformer.h (ixformer::infer namespace).
|
|
*/
|
|
#pragma once
|
|
|
|
#include <torch/all.h>
|
|
// #include <optional> // use c10::optional instead
|
|
#include <iostream>
|
|
#include <stdexcept>
|
|
|
|
#include "ixformer.h"
|
|
|
|
using namespace ixformer;
|
|
|
|
/* ---- Minimal LOG(FATAL) replacement ------------------------------------ */
|
|
#ifndef LOG
|
|
struct FatalLogStream {
|
|
std::ostringstream ss;
|
|
[[noreturn]] ~FatalLogStream() noexcept(false) {
|
|
std::cerr << ss.str() << std::endl;
|
|
throw std::runtime_error(ss.str());
|
|
}
|
|
template <typename T> FatalLogStream& operator<<(const T& v) {
|
|
ss << v; return *this;
|
|
}
|
|
};
|
|
#define LOG(level) FatalLogStream()
|
|
#endif
|
|
|
|
namespace xllm::kernel::ilu {
|
|
|
|
void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query,
|
|
torch::Tensor& key,
|
|
torch::Tensor& cos_sin_cache,
|
|
torch::Tensor& positions,
|
|
bool interleave);
|
|
|
|
void act_and_mul(torch::Tensor out,
|
|
torch::Tensor input,
|
|
const std::string& act_mode);
|
|
|
|
void reshape_paged_cache(
|
|
torch::Tensor& key,
|
|
c10::optional<torch::Tensor>& value,
|
|
torch::Tensor& key_cache,
|
|
c10::optional<torch::Tensor>& value_cache,
|
|
torch::Tensor& slot_mapping);
|
|
|
|
void batch_prefill(torch::Tensor& query,
|
|
const torch::Tensor& key,
|
|
const c10::optional<torch::Tensor>& value,
|
|
torch::Tensor& output,
|
|
c10::optional<torch::Tensor>& output_lse,
|
|
const c10::optional<torch::Tensor>& q_cu_seq_lens,
|
|
const c10::optional<torch::Tensor>& kv_cu_seq_lens,
|
|
const c10::optional<torch::Tensor>& alibi_slope,
|
|
const c10::optional<torch::Tensor>& attn_bias,
|
|
const c10::optional<torch::Tensor>& q_quant_scale,
|
|
const c10::optional<torch::Tensor>& k_quant_scale,
|
|
const c10::optional<torch::Tensor>& v_quant_scale,
|
|
const torch::Tensor& block_tables,
|
|
int64_t max_query_len,
|
|
int64_t max_seq_len,
|
|
float scale,
|
|
bool is_causal,
|
|
int64_t window_size_left,
|
|
int64_t window_size_right,
|
|
const std::string& compute_dtype,
|
|
bool return_lse);
|
|
|
|
void batch_decode(torch::Tensor& query,
|
|
const torch::Tensor& k_cache,
|
|
torch::Tensor& output,
|
|
const torch::Tensor& block_table,
|
|
const torch::Tensor& seq_lens,
|
|
const c10::optional<torch::Tensor>& v_cache,
|
|
c10::optional<torch::Tensor>& output_lse,
|
|
const c10::optional<torch::Tensor>& q_quant_scale,
|
|
const c10::optional<torch::Tensor>& k_cache_quant_scale,
|
|
const c10::optional<torch::Tensor>& v_cache_quant_scale,
|
|
const c10::optional<torch::Tensor>& out_quant_scale,
|
|
const c10::optional<torch::Tensor>& alibi_slope,
|
|
const c10::optional<torch::Tensor>& mask,
|
|
const std::string& compute_dtype,
|
|
int64_t max_seq_len,
|
|
int64_t window_size_left,
|
|
int64_t window_size_right,
|
|
float scale,
|
|
bool return_lse,
|
|
bool is_causal,
|
|
int64_t kv_cache_quant_bit_size);
|
|
|
|
void residual_layer_norm(torch::Tensor& input,
|
|
torch::Tensor& output,
|
|
c10::optional<torch::Tensor>& residual,
|
|
torch::Tensor& weight,
|
|
c10::optional<torch::Tensor>& bias,
|
|
c10::optional<torch::Tensor>& residual_out,
|
|
double eps);
|
|
|
|
void rms_norm(torch::Tensor& output,
|
|
torch::Tensor& input,
|
|
torch::Tensor& weight,
|
|
double eps);
|
|
|
|
torch::Tensor matmul(torch::Tensor a,
|
|
torch::Tensor b,
|
|
c10::optional<torch::Tensor> bias);
|
|
|
|
std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
|
|
const torch::Tensor& input,
|
|
int64_t topk,
|
|
int64_t num_expert_group,
|
|
int64_t topk_group,
|
|
bool normalize,
|
|
const c10::optional<torch::Tensor>& mask,
|
|
const std::string& normed_by,
|
|
const std::string& scoring_func,
|
|
double route_scale,
|
|
const c10::optional<torch::Tensor>& e_score_correction_bias);
|
|
|
|
std::vector<torch::Tensor> moe_gen_idx(torch::Tensor& expert_id,
|
|
int64_t expert_num);
|
|
|
|
torch::Tensor moe_expand_input(const torch::Tensor& input,
|
|
const torch::Tensor& gather_index,
|
|
const torch::Tensor& combine_idx,
|
|
int64_t topk);
|
|
|
|
torch::Tensor group_gemm(torch::Tensor& input,
|
|
torch::Tensor& weight,
|
|
torch::Tensor& tokens_per_experts,
|
|
const c10::optional<torch::Tensor>& dst_to_src,
|
|
torch::Tensor& output);
|
|
|
|
torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight);
|
|
|
|
} // namespace xllm::kernel::ilu
|