wehub-resource-sync
ec436095dd
Book-CI / test (macos-latest) (push) Has been cancelled
Book-CI / test (ubuntu-latest) (push) Has been cancelled
Book-CI / test (windows-latest) (push) Has been cancelled
Release Fake Tag / publish (push) Has been cancelled
Deploy / deploy (macos-latest) (push) Has been cancelled
Deploy / deploy (ubuntu-latest) (push) Has been cancelled
Deploy / deploy (windows-latest) (push) Has been cancelled
Release to PyPI / Build & publish sglang-kt (push) Has been cancelled
Release to PyPI / Build kt-kernel (Python 3.11) (push) Has been cancelled
Release to PyPI / Build kt-kernel (Python 3.12) (push) Has been cancelled
Release to PyPI / Publish kt-kernel to PyPI (push) Has been cancelled
399 lines
19 KiB
C++
399 lines
19 KiB
C++
// #ifndef LLAMAFILE_MLA_HPP
|
|
// #define LLAMAFILE_MLA_HPP
|
|
|
|
// #include "../common.hpp"
|
|
// #include "../mla-tp.hpp"
|
|
// #include "../rms-norm.hpp"
|
|
// #include "../rope.hpp"
|
|
// #include "ggml-quants.h"
|
|
// #include "ggml.h"
|
|
// #include "llamafile/sgemm.h"
|
|
|
|
// #include <algorithm>
|
|
// #include <cstddef>
|
|
// #include <utility>
|
|
// #include <vector>
|
|
|
|
// #define DIRECT_OR_POOL_BY(what, threshold, var, fn) \
|
|
// do { \
|
|
// if ((what) < (threshold)) { \
|
|
// for (int i = 0; i < (var); i++) { \
|
|
// (fn)(i); \
|
|
// } \
|
|
// } else { \
|
|
// pool->do_work_stealing_job((var), nullptr, (fn), nullptr); \
|
|
// } \
|
|
// } while (0)
|
|
|
|
// #define VEC_DOT_TYPE(type) (ggml_internal_get_type_traits((ggml_type)(type)).vec_dot_type)
|
|
// #define QUANT_BLCK_COUNT(size, type) (((size_t)(size)) / (ggml_blck_size((ggml_type)(type))))
|
|
// #define QUANT_BLCK_SIZE(size, type) (QUANT_BLCK_COUNT(size, type) * (ggml_type_size((ggml_type)(type))))
|
|
// #define QUANT_OFFSET(ptr, type, n, n_elements) \
|
|
// (offset_pointer((ptr), (size_t)(n) * QUANT_BLCK_SIZE((n_elements), (type))))
|
|
|
|
// #define LLAMAFILE_SGEMM_QUANT_FULL_MATMUL(m, n, k, a, a_type, b, b_col, c, c_col) \
|
|
// do { \
|
|
// llamafile_sgemm((m), (n), QUANT_BLCK_COUNT((k), (a_type)), (a), QUANT_BLCK_COUNT((k), (a_type)), \
|
|
// QUANT_OFFSET((b), VEC_DOT_TYPE((a_type)), (b_col), (k)), \
|
|
// QUANT_BLCK_COUNT((k), VEC_DOT_TYPE((a_type))), offset_pointer((c), (c_col) * (m) *
|
|
// sizeof(float)), \
|
|
// (k), 0, 1, GGML_TASK_TYPE_COMPUTE, (a_type), VEC_DOT_TYPE((a_type)), GGML_TYPE_F32, \
|
|
// GGML_PREC_DEFAULT); \
|
|
// } while (0)
|
|
|
|
// #define LLAMAFILE_SGEMM_MATMUL_F32(m, n, k, a, lda, b, ldb, c, ldc) \
|
|
// do { \
|
|
// llamafile_sgemm((m), (n), (k), (a), (lda), (b), (ldb), (c), (ldc), 0, 1, GGML_TASK_TYPE_COMPUTE, GGML_TYPE_F32, \
|
|
// GGML_TYPE_F32, GGML_TYPE_F32, GGML_PREC_DEFAULT); \
|
|
// } while (0)
|
|
|
|
// // bool decide_absorb(size_t a,int a_type,size_t b,int b_type,size_t c,int c_type,size_t d,int d_type){
|
|
// // size_t flops1 = ;
|
|
|
|
// // }
|
|
|
|
// inline void transpose(void *start, size_t dim0, size_t stride, size_t dim1) {
|
|
// // static_assert(false, "TODO");
|
|
// }
|
|
|
|
// template <RMS_NORM T_RMSNorm = RMSNorm, ROPE_APPLIER T_RopeApplier = Rope, ROPE_ANGLE T_RopeAngle = Yarn>
|
|
// class LLAMA_MLA_TP {
|
|
// private:
|
|
// GeneralMLAConfig config;
|
|
// int tp_part_idx;
|
|
// std::vector<void *> nope_pages; // [page_count * page_token_count * nope]
|
|
// std::vector<void *> rope_pages; // [page_count * page_token_count * nope]
|
|
|
|
// // weights
|
|
// void *local_q_a_proj; // [hidden_size * q_lora_rank]
|
|
// void *local_q_a_norm; // [q_lora_rank]
|
|
// std::vector<void *> local_q_b_proj; // [num_heads * (nope_size + rope_size))]
|
|
// void *local_kv_a_proj_with_mqa; // [hidden_size * (kv_lora_rank + rope)]
|
|
// void *local_kv_a_norm_with_mqa;
|
|
// void *local_kv_b_proj; // [(num_heads * (nope_size + nope_size) * kv_lora_rank)],
|
|
// // q_absorb: [num_heads * nope_size * kv_lora_rank]
|
|
// // out_absorb: [num_heads * nope_size * kv_lora_rank]
|
|
// std::vector<void *> local_k_b_proj_nope; // [(num_heads * kv_lora_rank * nope)],
|
|
// void *local_w_o; // [(num_heads * nope_size) * hidden_size]
|
|
// T_RopeAngle rope_angle;
|
|
|
|
// // intermediate
|
|
|
|
// void *quant_input; // [qlen, hidden size(Q)]
|
|
// void *q_a_proj_output; // [qlen, q_lora_rank]
|
|
// void *quant_q_a_proj_output; // [qlen, q_lora_rank(Q)]
|
|
|
|
// // for each query
|
|
// std::vector<void *> q_pe; // [num_heads * max_qlen * rope_size]
|
|
// std::vector<void *> k_pe; // [num_threads * rope_size]
|
|
// std::vector<void *> q_nope; // [num_heads * max_qlen * nope_size]
|
|
// std::vector<void *> attention_weights; // [num_heads * max_qlen * max_klen];
|
|
// std::vector<void *> q_absorb; // [num_heads, max_qlen, kv_lora_rank], or [num_heads, kv_lora_rank,
|
|
// max_qlen] std::vector<void *> o_absorb; // [num_heads, max_qlen, kv_lora_rank], or [num_heads,
|
|
// kv_lora_rank, max_qlen] std::vector<void *> compressed_kv_tmp; // [num_threads * token_count_in_page *
|
|
// kv_lora_rank] std::vector<void *> quant_o_absorb; // [num_heads, max_qlen, kv_lora_rank], or [num_heads,
|
|
// kv_lora_rank, max_qlen] std::vector<void *> attention_output; // [num_threads * max_qlen * nope] std::vector<void
|
|
// *> quant_attention_output; // [num_threads * max_qlen * nope]
|
|
|
|
// public:
|
|
// using output_t = float;
|
|
|
|
// LLAMA_MLA_TP(GeneralMLAConfig config, int tp_part_idx) : config(config), tp_part_idx(tp_part_idx) {
|
|
// std::vector<std::pair<void **, uint64_t>> s_mem_requests;
|
|
// }
|
|
|
|
// void set_pages(std::vector<void *> cache_pages) { this->nope_pages = cache_pages; }
|
|
// void set_pages(std::vector<void *> cache_pages, std::vector<void *> pe_pages) {
|
|
// this->nope_pages = cache_pages;
|
|
// this->rope_pages = pe_pages;
|
|
// }
|
|
|
|
// void forward(std::vector<int> qlens, std::vector<std::vector<int>> page_tables, std::vector<int> kv_lens,
|
|
// const void *input, void *output) {}
|
|
|
|
// void forward_prefill(std::vector<int> qlens, std::vector<std::vector<int>> page_tables, std::vector<int> kvlens,
|
|
// const void *input_raw, void *output) {
|
|
// auto pool = config.pool->get_subpool(tp_part_idx);
|
|
|
|
// float *input = (float *)input_raw;
|
|
// std::vector<int> qlen_split, total_len_split;
|
|
// qlen_split.reserve(qlens.size() + 1);
|
|
// qlen_split.push_back(0);
|
|
// total_len_split.reserve(qlens.size() + 1);
|
|
// int qlen_sum = 0;
|
|
// int total_len_sum = 0;
|
|
// for (size_t i = 0; i < qlens.size(); i++) {
|
|
// qlen_sum += qlens[i];
|
|
// qlen_split.push_back(qlen_sum);
|
|
|
|
// total_len_sum += qlens[i] + kvlens[i];
|
|
// total_len_split.push_back(total_len_sum);
|
|
// }
|
|
|
|
// auto which_query_by_qlen_sum = [&](int token_nth) -> std::pair<size_t, size_t> {
|
|
// auto query_idx = std::upper_bound(qlen_split.begin(), qlen_split.end(), token_nth) - qlen_split.begin() - 1;
|
|
// auto token_nth_from_start = token_nth - qlen_split.at(query_idx) + kvlens.at(query_idx);
|
|
// return {query_idx, token_nth_from_start};
|
|
// };
|
|
// auto which_query_by_total_sum = [&](int token_nth) -> std::pair<size_t, size_t> {
|
|
// auto query_idx =
|
|
// std::upper_bound(total_len_split.begin(), total_len_split.end(), token_nth) - total_len_split.begin() - 1;
|
|
// auto token_nth_from_start = token_nth - total_len_split.at(query_idx);
|
|
// return {query_idx, token_nth_from_start};
|
|
// };
|
|
|
|
// auto which_page = [&](int query, int token_nth_from_start) -> std::pair<size_t, size_t> {
|
|
// size_t page_idx = page_tables.at(query).at(div_up((size_t)token_nth_from_start, config.token_count_in_page));
|
|
|
|
// size_t token_at_in_page = token_nth_from_start % config.token_count_in_page;
|
|
// return {page_idx, token_at_in_page};
|
|
// };
|
|
|
|
// ggml_type vec_dot_type = ggml_internal_get_type_traits((ggml_type)config.q_a_proj_type).vec_dot_type;
|
|
// size_t hidden_size_float_bytes = config.hidden_size * sizeof(float);
|
|
// size_t hidden_size_quant_blck_count = config.hidden_size / ggml_blck_size(vec_dot_type);
|
|
// size_t hidden_size_quant_bytes = hidden_size_quant_blck_count * ggml_type_size(vec_dot_type);
|
|
// // quant to q8 0
|
|
|
|
// DIRECT_OR_POOL_BY(qlen_sum, 10, qlen_sum, [&](int token_at_i) {
|
|
// size_t token_at = token_at_i;
|
|
// quantize_q8_0(offset_pointer(input, token_at * config.hidden_size * sizeof(float)),
|
|
// offset_pointer(quant_input,
|
|
// token_at * QUANT_BLCK_SIZE(config.hidden_size,
|
|
// VEC_DOT_TYPE(config.q_a_proj_type))),
|
|
// 1, config.hidden_size, nullptr);
|
|
// });
|
|
|
|
// {
|
|
// // q lora rank
|
|
// // maybe this should be up to non tp
|
|
// auto proj_lora_a = [&](int task_id) {
|
|
// size_t token_at = task_id % qlen_sum;
|
|
// bool do_q_or_kv = (task_id / qlen_sum) == 0;
|
|
// if (do_q_or_kv) {
|
|
// auto this_q_a_proj_output =
|
|
// (float *)offset_pointer(q_a_proj_output, token_at * config.hidden_size * sizeof(float));
|
|
|
|
// LLAMAFILE_SGEMM_QUANT_FULL_MATMUL(config.q_lora_rank, 1, config.hidden_size, local_q_a_proj,
|
|
// config.q_a_proj_type, quant_input, token_at, this_q_a_proj_output, 0);
|
|
|
|
// T_RMSNorm::rms_norm_single(config.q_lora_rank, (float *)local_q_a_norm, this_q_a_proj_output);
|
|
|
|
// quantize_q8_0(
|
|
// this_q_a_proj_output,
|
|
// offset_pointer(quant_q_a_proj_output,
|
|
// token_at * QUANT_BLCK_SIZE(config.q_lora_rank, VEC_DOT_TYPE(config.q_b_proj_type))),
|
|
// 1, config.q_lora_rank, nullptr);
|
|
|
|
// } else {
|
|
// auto [query, token_from_start] = which_query_by_qlen_sum(token_at);
|
|
// auto [page_idx, token_at_in_page] = which_page(query, token_from_start);
|
|
// LLAMAFILE_SGEMM_QUANT_FULL_MATMUL(config.kv_lora_rank, 1, config.hidden_size, local_kv_a_proj_with_mqa,
|
|
// config.kv_a_proj_with_mqa_type, quant_input, token_at,
|
|
// rope_pages.at(page_idx), token_at_in_page);
|
|
// T_RMSNorm::rms_norm_single(
|
|
// config.kv_lora_rank, (float *)local_kv_a_norm_with_mqa,
|
|
// (float *)offset_pointer(rope_pages.at(page_idx), token_at_in_page * config.kv_lora_rank *
|
|
// sizeof(float)));
|
|
// LLAMAFILE_SGEMM_QUANT_FULL_MATMUL(config.rope_size, 1, config.hidden_size,
|
|
// QUANT_OFFSET(local_kv_a_proj_with_mqa, config.kv_a_proj_with_mqa_type,
|
|
// config.kv_lora_rank, config.hidden_size),
|
|
// config.kv_a_proj_with_mqa_type, quant_input, token_at,
|
|
// nope_pages.at(page_idx), token_at_in_page);
|
|
// }
|
|
// };
|
|
// DIRECT_OR_POOL_BY(qlen_sum, 10, qlen_sum * 2, proj_lora_a);
|
|
// }
|
|
|
|
// {
|
|
// int task_count = config.num_heads * 2 * qlen_sum; // head, rope/nope, qlen
|
|
// auto q_proj_lora_b = [&](int task_id) {
|
|
// size_t head_idx = task_id / (2 * qlen_sum);
|
|
// task_id %= (2 * qlen_sum);
|
|
// bool nope_or_rope = (task_id / qlen_sum) == 0;
|
|
// task_id %= qlen_sum;
|
|
// size_t token_at = task_id;
|
|
// auto [query, token_from_start] = which_query_by_qlen_sum(token_at);
|
|
// auto [page_idx, token_at_in_page] = which_page(query, token_from_start);
|
|
|
|
// if (nope_or_rope) {
|
|
// LLAMAFILE_SGEMM_QUANT_FULL_MATMUL(
|
|
// config.nope_size, 1, config.q_lora_rank,
|
|
// QUANT_OFFSET(local_q_b_proj.at(head_idx), config.q_b_proj_type,
|
|
// head_idx * (config.nope_size + config.rope_size), config.q_lora_rank),
|
|
// config.q_b_proj_type, quant_q_a_proj_output, token_at, q_nope.at(head_idx), token_at);
|
|
// } else {
|
|
// LLAMAFILE_SGEMM_QUANT_FULL_MATMUL(
|
|
// config.rope_size, 1, config.q_lora_rank,
|
|
// QUANT_OFFSET(local_q_b_proj.at(head_idx), config.q_b_proj_type,
|
|
// head_idx * (config.nope_size + config.rope_size) + config.nope_size, config.q_lora_rank),
|
|
// config.q_b_proj_type, quant_q_a_proj_output, token_at, q_pe.at(head_idx), token_at);
|
|
// T_RopeApplier::apply_single(config.rope_size,
|
|
// offset_pointer(q_pe.at(head_idx), token_at * config.rope_size * sizeof(float)),
|
|
// rope_angle.cos(token_at), rope_angle.sin(token_at));
|
|
// }
|
|
// };
|
|
// pool->do_work_stealing_job(task_count, nullptr, q_proj_lora_b, nullptr);
|
|
// }
|
|
|
|
// for (int query = 0; query < qlens.size(); query++) {
|
|
// {
|
|
// // pe attention
|
|
// // apply k pe online
|
|
// int task_count = config.num_heads * (qlens[query] + kvlens[query]); // by kvlen
|
|
// auto pe_attn = [&](int task_id) {
|
|
// size_t head_idx = task_id / (qlens[query] + kvlens[query]);
|
|
// size_t token_from_start = task_id % (qlens[query] + kvlens[query]);
|
|
|
|
// // auto q_token_at = qlen_split[query] + qlens[query];
|
|
|
|
// auto [page_idx, token_at_in_page] = which_page(query, token_from_start);
|
|
// memcpy(k_pe[WorkerPool::thread_local_id],
|
|
// offset_pointer(rope_pages.at(page_idx), token_at_in_page * config.rope_size * sizeof(float)),
|
|
// sizeof(float) * config.rope_size);
|
|
// T_RopeApplier::apply_single(config.rope_size, k_pe[WorkerPool::thread_local_id],
|
|
// rope_angle.cos(token_from_start), rope_angle.sin(token_from_start));
|
|
|
|
// LLAMAFILE_SGEMM_MATMUL_F32(1, qlens[query], config.rope_size, k_pe[WorkerPool::thread_local_id],
|
|
// config.rope_size, q_pe.at(head_idx), config.rope_size,
|
|
// attention_weights[head_idx], config.max_kvlen);
|
|
// };
|
|
// pool->do_work_stealing_job(task_count, pe_attn);
|
|
// }
|
|
// {
|
|
// // clear q absorb
|
|
// pool->do_work_stealing_job(config.num_heads, [&](int task_id) {
|
|
// memset(q_absorb[task_id], 0, config.kv_lora_rank * config.max_qlen * sizeof(float));
|
|
// });
|
|
|
|
// // aborb W_uk
|
|
// int task_count = config.num_heads * qlens[query];
|
|
// auto task = [&](int task_id) {
|
|
// size_t head_idx = task_id / qlens[query];
|
|
// size_t token_at = task_id % qlens[query];
|
|
|
|
// // q_absorb now [kvrank, max_qlen]
|
|
// LLAMAFILE_SGEMM_MATMUL_F32(qlens[query], config.kv_lora_rank, config.nope_size, q_nope[head_idx],
|
|
// config.nope_size, local_k_b_proj_nope[head_idx], config.nope_size,
|
|
// q_absorb[head_idx], config.max_qlen);
|
|
// transpose(q_nope[head_idx], config.kv_lora_rank, config.max_qlen, qlens[query]);
|
|
// };
|
|
// pool->do_work_stealing_job(task_count, task);
|
|
// }
|
|
|
|
// {
|
|
// // nope attention weights
|
|
// size_t page_count = div_up((size_t)kvlens[query], config.token_count_in_page);
|
|
// int task_count = config.num_heads * page_count;
|
|
// auto task = [&](int task_id) {
|
|
// size_t head_idx = task_id / page_count;
|
|
// size_t page_idx = task_id % page_count;
|
|
// void *page_ptr = nope_pages[page_tables[query][page_idx]]; // mla no head
|
|
|
|
// size_t kvlen =
|
|
// page_idx == (page_count - 1) ? (kvlens[query] % config.token_count_in_page) :
|
|
// config.token_count_in_page;
|
|
|
|
// LLAMAFILE_SGEMM_MATMUL_F32(
|
|
// kvlen, qlens[query], config.kv_lora_rank, page_ptr, config.kv_lora_rank, q_absorb[head_idx],
|
|
// config.max_qlen,
|
|
// offset_pointer(attention_weights[head_idx], page_idx * config.token_count_in_page * sizeof(float)),
|
|
// config.max_kvlen);
|
|
// // static_assert(false, "soft max todo");
|
|
// };
|
|
// pool->do_work_stealing_job(task_count, task);
|
|
// }
|
|
|
|
// {
|
|
// // clear o absorb
|
|
// pool->do_work_stealing_job(config.num_heads, [&](int task_id) {
|
|
// memset(o_absorb[task_id], 0, config.kv_lora_rank * config.max_qlen * sizeof(float));
|
|
// });
|
|
|
|
// // o absorb
|
|
// size_t page_count = div_up((size_t)kvlens[query], config.token_count_in_page);
|
|
// int task_count = config.num_heads * page_count;
|
|
// auto task = [&](int task_id) {
|
|
// size_t head_idx = task_id / page_count;
|
|
// size_t page_idx = task_id % page_count;
|
|
// void *page_ptr = nope_pages[page_tables[query][page_idx]]; // mla no head
|
|
// size_t kvlen =
|
|
// page_idx == (page_count - 1) ? (kvlens[query] % config.token_count_in_page) :
|
|
// config.token_count_in_page;
|
|
|
|
// memcpy(compressed_kv_tmp[WorkerPool::thread_local_id], page_ptr,
|
|
// config.token_count_in_page * config.kv_lora_rank * sizeof(float));
|
|
// transpose(compressed_kv_tmp[WorkerPool::thread_local_id], config.token_count_in_page, config.kv_lora_rank,
|
|
// kvlen);
|
|
|
|
// LLAMAFILE_SGEMM_MATMUL_F32(
|
|
// config.kv_lora_rank, qlens[query], kvlen, compressed_kv_tmp[WorkerPool::thread_local_id],
|
|
// config.token_count_in_page,
|
|
// offset_pointer(attention_weights[head_idx], page_idx * config.token_count_in_page * sizeof(float)),
|
|
// config.max_kvlen, o_absorb[head_idx], config.kv_lora_rank);
|
|
// };
|
|
// pool->do_work_stealing_job(task_count, task);
|
|
// }
|
|
|
|
// {
|
|
|
|
// // clear
|
|
// pool->do_work_stealing_job(config.num_heads, [&](int task_id) {
|
|
// memset(attention_output[task_id], 0, config.nope_size * config.max_qlen * sizeof(float));
|
|
// });
|
|
|
|
// // attention output
|
|
// int task_count = config.num_heads * qlens[query];
|
|
// auto task = [&](int task_id) {
|
|
// size_t head_idx = task_id / qlens[query];
|
|
// size_t token_at = task_id % qlens[query];
|
|
|
|
// quantize_q8_0((float *)offset_pointer(o_absorb[head_idx], config.kv_lora_rank * token_at * sizeof(float)),
|
|
// offset_pointer(quant_o_absorb[head_idx],
|
|
// QUANT_BLCK_SIZE(config.kv_lora_rank, VEC_DOT_TYPE(config.kv_b_proj_type))),
|
|
// 1, config.kv_lora_rank, nullptr);
|
|
|
|
// auto kv_b_proj_ptr =
|
|
// offset_pointer(local_kv_b_proj, ((head_idx * 2 + 1) * config.nope_size) *
|
|
// QUANT_BLCK_SIZE(config.kv_lora_rank, config.kv_b_proj_type));
|
|
|
|
// LLAMAFILE_SGEMM_QUANT_FULL_MATMUL(config.nope_size, 1, config.kv_lora_rank, kv_b_proj_ptr,
|
|
// config.kv_b_proj_type, quant_o_absorb[head_idx], token_at,
|
|
// attention_output[head_idx], token_at);
|
|
// };
|
|
// pool->do_work_stealing_job(task_count, task);
|
|
// }
|
|
|
|
// {
|
|
// // quant attention output
|
|
// // static_assert(false,"TODO" );
|
|
// }
|
|
|
|
// {
|
|
// // get final output
|
|
// // static_assert(false,"TODO" );
|
|
// }
|
|
// }
|
|
// }
|
|
|
|
// void load_weights(int complete_num_heads, int offset) {}
|
|
// };
|
|
// template <typename Norm, typename Rope, typename RopeAngle>
|
|
// class TP_MLA<LLAMA_MLA_TP<Norm, Rope, RopeAngle>> : public TP_MLA_Common<LLAMA_MLA_TP<Norm, Rope, RopeAngle>> {
|
|
// public:
|
|
// using TP_MLA_Common<LLAMA_MLA_TP<Norm, Rope, RopeAngle>>::TP_MLA_Common;
|
|
|
|
// void load_weights() {
|
|
// auto pool = this->config.pool;
|
|
// auto tp_num_heads = this->config.num_heads / this->tp_count;
|
|
// pool->dispense_backend()->do_numa_job([this, pool, tp_num_heads](int tp_id) {
|
|
// this->tps[tp_id]->load_weights(this->config.num_heads, tp_id * tp_num_heads);
|
|
// });
|
|
// this->weights_loaded = true;
|
|
// }
|
|
|
|
// void merge_results(int qlen, void *output) {}
|
|
// };
|
|
|
|
// #endif
|