Skip to content

GitLab

Unverified Commit 527cc978 authored by Jeffrey Morgan's avatar Jeffrey Morgan Committed by GitHub
Browse files

llama: update vendored code to commit 40c6d79f (#7875)

parent a37f4a86
Hide whitespace changes
Inline Side-by-side
Showing with 7228 additions and 3115 deletions
llama/llama-grammar.cpp
View file @ 527cc978
/**
* llama.cpp - commit 3f1ae2e32cde00c39b96be6d01c2997c29bae555 - do not edit this file
* llama.cpp - commit 40c6d79fb52f995f47507fedfeaae2ac05d9b35c - do not edit this file
*
* MIT License
*
......
llama/llama-grammar.h
View file @ 527cc978
/**
* llama.cpp - commit 3f1ae2e32cde00c39b96be6d01c2997c29bae555 - do not edit this file
* llama.cpp - commit 40c6d79fb52f995f47507fedfeaae2ac05d9b35c - do not edit this file
*
* MIT License
*
......
llama/llama-impl.h
View file @ 527cc978
/**
* llama.cpp - commit 3f1ae2e32cde00c39b96be6d01c2997c29bae555 - do not edit this file
* llama.cpp - commit 40c6d79fb52f995f47507fedfeaae2ac05d9b35c - do not edit this file
*
* MIT License
*
......
llama/llama-sampling.cpp
View file @ 527cc978
/**
* llama.cpp - commit 3f1ae2e32cde00c39b96be6d01c2997c29bae555 - do not edit this file
* llama.cpp - commit 40c6d79fb52f995f47507fedfeaae2ac05d9b35c - do not edit this file
*
* MIT License
*
......@@ -89,6 +89,30 @@ static void llama_log_softmax(float * array, size_t size) {
}
*/
static void llama_sampler_temp_impl(llama_token_data_array * cur_p, float temp) {
if (temp <= 0.0f) {
// find the token with the highest logit and set the rest to -inf
size_t max_i = 0;
float max_l = cur_p->data[0].logit;
for (size_t i = 1; i < cur_p->size; ++i) {
if (cur_p->data[i ].logit > max_l) {
cur_p->data[max_i].logit = -INFINITY;
max_i = i;
max_l = cur_p->data[i].logit;
} else {
cur_p->data[i].logit = -INFINITY;
}
}
return;
}
for (size_t i = 0; i < cur_p->size; ++i) {
cur_p->data[i].logit /= temp;
}
}
static void llama_sampler_softmax_impl(llama_token_data_array * cur_p) {
GGML_ASSERT(cur_p->size > 0);
......@@ -115,7 +139,7 @@ static void llama_sampler_softmax_impl(llama_token_data_array * cur_p) {
}
static void llama_sampler_top_k_impl(llama_token_data_array * cur_p, int32_t k) {
// TODO: move bucket sort to separate function so that top_p/tail_free/typical/softmax first is equally fast
// TODO: move bucket sort to separate function so that top_p/typical/softmax first is equally fast
// if (k >= (int32_t)cur_p->size) {
// return;
// }
......@@ -453,6 +477,9 @@ static const char * llama_sampler_dist_name(const struct llama_sampler * /*smpl*
static void llama_sampler_dist_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
auto * ctx = (llama_sampler_dist *) smpl->ctx;
llama_sampler_softmax_impl(cur_p);
cur_p->selected = llama_sample_dist(cur_p, ctx->rng);
}
......@@ -732,101 +759,6 @@ struct llama_sampler * llama_sampler_init_min_p(float p, size_t min_keep) {
};
}
// tail-free
struct llama_sampler_tail_free {
const float z;
const size_t min_keep;
};
static const char * llama_sampler_tail_free_name(const struct llama_sampler * /*smpl*/) {
return "tail-free";
}
static void llama_sampler_tail_free_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
const auto * ctx = (llama_sampler_tail_free *) smpl->ctx;
if (ctx->z >= 1.0f || cur_p->size <= 2) {
return;
}
llama_sampler_softmax_impl(cur_p);
// Compute the first and second derivatives
std::vector<float> first_derivatives(cur_p->size - 1);
std::vector<float> second_derivatives(cur_p->size - 2);
for (size_t i = 0; i < first_derivatives.size(); ++i) {
first_derivatives[i] = cur_p->data[i].p - cur_p->data[i + 1].p;
}
for (size_t i = 0; i < second_derivatives.size(); ++i) {
second_derivatives[i] = first_derivatives[i] - first_derivatives[i + 1];
}
// Calculate absolute value of second derivatives
for (size_t i = 0; i < second_derivatives.size(); ++i) {
second_derivatives[i] = std::abs(second_derivatives[i]);
}
// Normalize the second derivatives
{
const float second_derivatives_sum = std::accumulate(second_derivatives.begin(), second_derivatives.end(), 0.0f);
if (second_derivatives_sum > 1e-6f) {
for (float & value : second_derivatives) {
value /= second_derivatives_sum;
}
} else {
for (float & value : second_derivatives) {
value = 1.0f / second_derivatives.size();
}
}
}
float cum_sum = 0.0f;
size_t last_idx = cur_p->size;
for (size_t i = 0; i < second_derivatives.size(); ++i) {
cum_sum += second_derivatives[i];
// Check if the running sum is greater than z or if we have kept at least min_keep tokens
if (cum_sum > ctx->z && i >= ctx->min_keep) {
last_idx = i;
break;
}
}
// Resize the output vector to keep only the tokens above the tail location
cur_p->size = last_idx;
}
static struct llama_sampler * llama_sampler_tail_free_clone(const struct llama_sampler * smpl) {
const auto * ctx = (const llama_sampler_tail_free *) smpl->ctx;
return llama_sampler_init_tail_free(ctx->z, ctx->min_keep);
}
static void llama_sampler_tail_free_free(struct llama_sampler * smpl) {
delete (llama_sampler_tail_free *) smpl->ctx;
}
static struct llama_sampler_i llama_sampler_tail_free_i = {
/* .name = */ llama_sampler_tail_free_name,
/* .accept = */ nullptr,
/* .apply = */ llama_sampler_tail_free_apply,
/* .reset = */ nullptr,
/* .clone = */ llama_sampler_tail_free_clone,
/* .free = */ llama_sampler_tail_free_free,
};
struct llama_sampler * llama_sampler_init_tail_free(float z, size_t min_keep) {
return new llama_sampler {
/* .iface = */ &llama_sampler_tail_free_i,
/* .ctx = */ new llama_sampler_tail_free {
/* .z = */ z,
/*. min_keep = */ min_keep,
},
};
}
// typical
struct llama_sampler_typical {
......@@ -938,9 +870,8 @@ static const char * llama_sampler_temp_name(const struct llama_sampler * /*smpl*
static void llama_sampler_temp_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
const auto * ctx = (llama_sampler_temp *) smpl->ctx;
for (size_t i = 0; i < cur_p->size; ++i) {
cur_p->data[i].logit /= ctx->temp;
}
llama_sampler_temp_impl(cur_p, ctx->temp);
}
static struct llama_sampler * llama_sampler_temp_clone(const struct llama_sampler * smpl) {
......@@ -987,6 +918,7 @@ static void llama_sampler_temp_ext_apply(struct llama_sampler * smpl, llama_toke
if (ctx->delta > 0) {
const float min_temp = std::max(0.0f, ctx->temp - ctx->delta);
const float max_temp = ctx->temp + ctx->delta;
float exponent_val = ctx->exponent;
// no need to do anything if there is only one (or zero) candidates
......@@ -1024,9 +956,7 @@ static void llama_sampler_temp_ext_apply(struct llama_sampler * smpl, llama_toke
#endif
// Apply the dynamically calculated temperature scaling
for (size_t i = 0; i < cur_p->size; ++i) {
cur_p->data[i].logit /= dyn_temp;
}
llama_sampler_temp_impl(cur_p, dyn_temp);
// Re-compute softmax probabilities after scaling logits with dynamic temperature
const double max_l_double = cur_p->data[0].logit;
......@@ -1050,9 +980,7 @@ static void llama_sampler_temp_ext_apply(struct llama_sampler * smpl, llama_toke
}
#endif
} else {
for (size_t i = 0; i < cur_p->size; ++i) {
cur_p->data[i].logit /= ctx->temp;
}
llama_sampler_temp_impl(cur_p, ctx->temp);
}
}
......@@ -1085,6 +1013,101 @@ struct llama_sampler * llama_sampler_init_temp_ext(float temp, float delta, floa
};
}
// xtc
struct llama_sampler_xtc {
const float probability;
const float threshold;
const size_t min_keep;
const uint32_t seed;
uint32_t seed_cur;
std::mt19937 rng;
};
static const char * llama_sampler_xtc_name(const struct llama_sampler * /*smpl*/) {
return "xtc";
}
static void llama_sample_xtc_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
auto * ctx = (llama_sampler_xtc *) smpl->ctx;
if (ctx->probability <= 0.0f
|| ctx->threshold > 0.5f
|| cur_p->size < 2) {
return;
}
std::uniform_real_distribution<float> distribution(0.0f, 1.0f);
float chance = distribution(ctx->rng);
if (chance > ctx->probability) return;
// in case it's not sorted/recalculated yet
llama_sampler_softmax_impl(cur_p);
int pos_last = 0;
for (size_t i = 0; i < cur_p->size; ++i) {
if (cur_p->data[i].p >= ctx->threshold) {
pos_last = i;
} else break;
}
if (cur_p->size - pos_last >= ctx->min_keep && pos_last > 0) {
cur_p->data += pos_last;
cur_p->size -= pos_last;
}
}
static struct llama_sampler * llama_sampler_xtc_clone(const struct llama_sampler * smpl) {
const auto * ctx = (const llama_sampler_xtc *) smpl->ctx;
auto * result = llama_sampler_init_xtc(ctx->probability, ctx->threshold, ctx->min_keep, ctx->seed);
// copy the state
{
auto * result_ctx = (llama_sampler_xtc *) result->ctx;
result_ctx->rng = ctx->rng;
}
return result;
}
static void llama_sampler_xtc_free(struct llama_sampler * smpl) {
delete (llama_sampler_xtc *) smpl->ctx;
}
static void llama_sampler_xtc_reset(struct llama_sampler * smpl) {
auto * ctx = (llama_sampler_xtc *) smpl->ctx;
ctx->seed_cur = get_rng_seed(ctx->seed);
ctx->rng.seed(ctx->seed_cur);
}
static struct llama_sampler_i llama_sampler_xtc_i = {
/* .name = */ llama_sampler_xtc_name,
/* .accept = */ nullptr,
/* .apply = */ llama_sample_xtc_apply,
/* .reset = */ llama_sampler_xtc_reset,
/* .clone = */ llama_sampler_xtc_clone,
/* .free = */ llama_sampler_xtc_free,
};
struct llama_sampler * llama_sampler_init_xtc(float p, float t, size_t min_keep, uint32_t seed) {
auto seed_cur = get_rng_seed(seed);
return new llama_sampler {
/* .iface = */ &llama_sampler_xtc_i,
/* .ctx = */ new llama_sampler_xtc {
/* .probability = */ p,
/* .threshold = */ t,
/* .min_keep = */ min_keep,
/* .seed = */ seed,
/* .seed_cur = */ seed_cur,
/* .rng = */ std::mt19937(seed_cur),
},
};
}
// mirostat
struct llama_sampler_mirostat {
......@@ -1591,6 +1614,400 @@ struct llama_sampler * llama_sampler_init_penalties(
};
}
// DRY
struct llama_sampler_dry {
int32_t total_context_size;
const float dry_multiplier;
const float dry_base;
const int32_t dry_allowed_length;
const int32_t dry_penalty_last_n;
std::unordered_multimap<llama_token, std::vector<llama_token>> dry_processed_breakers;
std::vector<int> dry_repeat_count;
std::unordered_map<llama_token, int> dry_max_token_repeat;
ring_buffer<llama_token> last_tokens;
};
// Ported from Koboldcpp, original PR: https://github.com/LostRuins/koboldcpp/pull/982 (Original author: pi6am)
static void get_overlapping_token_sequences(const llama_vocab & vocab, const std::string& str, std::unordered_multimap<llama_token, std::vector<llama_token>>& token_sequences, int max_tail_len = -1) {
for (llama_token token_id = 0; token_id < (llama_token)vocab.n_vocab; token_id++) {
std::string word = llama_detokenize(vocab, {token_id}, true);
if (word.find(str) != std::string::npos) {
token_sequences.emplace(token_id, std::vector<llama_token>());
} else {
size_t word_len = word.size(), str_len = str.size();
size_t pos = -1;
while ((pos = word.find(str[0], pos + 1)) != std::string::npos) {
bool match = true;
size_t i;
for (i = 1; i < str_len && i + pos < word_len; ++i) {
if (word[pos + i] != str[i]) {
match = false;
break;
}
}
if (match) {
std::vector<llama_token> tokenization = llama_tokenize_internal(vocab, str.substr(i), false, false);
if (max_tail_len >= 0 && tokenization.size() > (size_t)max_tail_len) {
tokenization.resize(max_tail_len);
}
// Ensure we don't already have a duplicate matching tokenization
auto its = token_sequences.equal_range(token_id);
bool found = false;
for (auto it = its.first; it != its.second; ++it) {
if (tokenization == it->second) {
found = true;
break;
}
}
if (!found) {
token_sequences.emplace(token_id, tokenization);
}
}
}
}
}
}
static const char * llama_sampler_dry_name(const struct llama_sampler * /*smpl*/) {
return "dry";
}
static void llama_sampler_dry_accept(struct llama_sampler * smpl, llama_token token) {
auto * ctx = (llama_sampler_dry *) smpl->ctx;
if (ctx->dry_multiplier == 0.0f || ctx->dry_base < 1.0f || ctx->dry_penalty_last_n == 0) {
return;
}
ctx->last_tokens.push_back(token);
}
// Ported from Koboldcpp, original PR: https://github.com/LostRuins/koboldcpp/pull/982 (Original author: pi6am)
static void llama_sampler_dry_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
auto * ctx = (llama_sampler_dry *) smpl->ctx;
if (ctx->dry_multiplier == 0.0f || ctx->dry_base < 1.0f || ctx->dry_penalty_last_n == 0) {
return;
}
int32_t effective_dry_penalty_last_n = (ctx->dry_penalty_last_n == -1) ? ctx->total_context_size : std::max(ctx->dry_penalty_last_n, 0);
int last_n_repeat = std::min(std::min((int)ctx->last_tokens.size(), effective_dry_penalty_last_n), ctx->total_context_size);
if (last_n_repeat <= ctx->dry_allowed_length) {
return;
}
ctx->dry_repeat_count.assign(last_n_repeat, 0);
ctx->dry_max_token_repeat.clear();
// Step 1: Look for restart sequences to limit the maximum repetition length.
// Work backwards through the context looking for any token that begins a restart sequence.
//
// The collection `restart_sequences` is a mapping from a "head" token to all "tail"
// sequences that together comprise a restart sequence. This allows us to quickly check
// whether each token is the head of a complete sequence. Most restart sequences are actually
// a single token, and for these the "tail" is an empty vector.
//
// If the token is a "head", test all restart sequences that begin with this token
// (there will often only be one sequence for each token, but if sequences like 'aaaq1' and
// 'aaa1' are used as restart strings, both could start with 'aaa' when tokenized). The
// longest matching sequence (if any) is used to limit the maximum repetition length.
//
// Note that in the case case of a short sequence contained in a longer one, this might fail to
// find the smallest value for `rep_limit`. For example, if 'amniotic' and 'ni' are both used as
// restart sequences, 'ni' will be found first, and since it's shorter it will fail to suppress
// 'otic'. This is a minor issue since fully contained restart sequences are likely to be rare.
//
// This is theoretically worst-case O(N^2) for arbitrary restart sequences, which is why we
// have already clamped the maximum tail sequence length when generating `restart_sequences`.
// With clamping, this scan is O(N) in the context length.
int rep_limit = last_n_repeat;
for (int i = 0; i < last_n_repeat; ++i) {
llama_token token = ctx->last_tokens.rat(i);
auto its = ctx->dry_processed_breakers.equal_range(token);
if (its.first == ctx->dry_processed_breakers.end()) {
continue;
}
int longest_match = -1;
for (auto it = its.first; it != its.second; ++it) {
// Note that (*it) does not contain the head character, so seq_len will be
// the restart sequence length minus 1.
// In the common case of a single-token restart sequence, (*it) will be empty
// and we will trivially match.
int seq_len = (int)it->second.size();
if (seq_len > longest_match && seq_len <= (int)i) {
bool match = true;
for (int offset = 0; offset < seq_len; ++offset) {
// The -1 when indexing `last_tokens` is because we already matched the head.
if (it->second[offset] != ctx->last_tokens.rat(i - offset - 1)) {
match = false;
break;
}
}
if (match) {
longest_match = seq_len;
}
}
}
if (longest_match >= 0) {
// We found a restart sequence starting `i` tokens from the end and continuing for
// `longest_match` tokens.
rep_limit = i - longest_match;
break;
}
}
if (rep_limit < ctx->dry_allowed_length) {
return;
}
// Step 2: Iterate in reverse over the last N tokens of the context, using the "Z-algorithm" (in
// the reverse direction) to efficiently compute the positions and lengths of suffixes appearing
// elsewhere in the context. We limit the suffix length to `rep_limit` to respect restart sequences.
//
// This algorithm is not currently documented on Wikipedia, but there is a clear description here:
// https://ivanyu.me/blog/2014/10/15/z-algorithm/
//
// The code below is adapted from the public domain implementation by the same author here:
// https://github.com/ivanyu/string-algorithms/blob/master/z_algorithm.py
//
// Example:
// Last N tokens: a b c c b c y a b c
// Repeat counts: 0 0 3 1 0 2 0 0 0 0
// ^
// This `3` means that the last three tokens of the context (a b c) also appear here.
//
// This step is worst case O(N) since the Z-algorithm is linear, despite the appearance of nested
// for/while loops. This can be seen by observing that the `lt` and `rt` bounds are set after each
// repeated suffix is detected (i.e. after each while loop when n > 0). These bound variables
// ensure that the inner while loops only examine each token in the context once as the outer
// for loop iterates over the context.
{
const int last = last_n_repeat - 1;
int rt = 0, lt = 0;
for (int k = 1; k < last_n_repeat; ++k) {
if (k > rt) {
// If k is outside the current Z-box, do naive computation.
int n = 0;
while (n + k < last_n_repeat && ctx->last_tokens.rat(n) == ctx->last_tokens.rat(n+k)) {
++n;
}
ctx->dry_repeat_count[last - k] = std::min(n, rep_limit);
if (n > 0) {
lt = k;
rt = k+n-1;
}
} else {
// If k is inside the current Z-box, consider two cases.
int p = k - lt; // Pair index.
int right_part_len = rt - k + 1;
if (ctx->dry_repeat_count[last - p] < right_part_len) {
int n = std::min(ctx->dry_repeat_count[last - p], rep_limit);
ctx->dry_repeat_count[last - k] = n;
} else {
int i = rt + 1;
while (i < last_n_repeat && ctx->last_tokens.rat(i) == ctx->last_tokens.rat(i - k)) {
i += 1;
}
int n = std::min(i - k, rep_limit);
ctx->dry_repeat_count[last - k] = n;
lt = k;
rt = i - 1;
}
}
}
}
// Step 3: Iterate over dry_repeat_count and last_tokens, examining the maximum repeat length
// that would be generated by emitting each new token that would extend a sequence.
//
// Following the same example as above:
// Last N tokens: a b c c b c y a b c
// Repeat counts: 0 0 3 1 0 2 0 0 0 0
//
// For each non-zero, look ahead one token. This token, if emitted, would extend the repetition.
// c: 3 -> 4 (from `a b c` to `a b c c`)
// b: 1 -> 2 (from `c` to `c b`)
// y: 2 -> 3 (from `b c` to `b c y`)
for (int i = 0; i < last_n_repeat - 1; ++i) {
int repeat_len = ctx->dry_repeat_count[i];
if (repeat_len >= ctx->dry_allowed_length) {
// This token ends a repeat, so the next token would continue one.
// By convention, the value of `repeat_len` only includes the tokens currently
// in the context, not the new token that would be added.
llama_token token = ctx->last_tokens.rat(last_n_repeat - 2 - i);
// Track the maximum sequence ending in this token.
const auto& it = ctx->dry_max_token_repeat.find(token);
if (it == ctx->dry_max_token_repeat.end() || it->second < repeat_len) {
ctx->dry_max_token_repeat[token] = repeat_len;
}
}
}
// Step 4: Apply logit penalties based on the maximum repeat length for relevant tokens.
// Prevent floating point overflow in `pow(penalty_base, exponent)` by clamping to `max_exponent`.
// Compute it from `penalty_base` and the approximate log of `std::numeric_limits<float>::max()`
const float FLOAT_MAX_LOG = 88.7228391f;
int max_exponent = 0;
if (ctx->dry_base > 1.000001f) {
max_exponent = FLOAT_MAX_LOG / std::log(ctx->dry_base);
}
for (size_t i = 0; i < cur_p->size; ++i) {
const auto& af_kvp = ctx->dry_max_token_repeat.find(cur_p->data[i].id);
if (af_kvp != ctx->dry_max_token_repeat.end()) {
// Check all sequence breakers starting with this token
auto range = ctx->dry_processed_breakers.equal_range(cur_p->data[i].id);
bool is_single_token_breaker = false;
for (auto it = range.first; it != range.second; ++it) {
if (it->second.empty()) {
is_single_token_breaker = true;
break;
}
}
// Apply penalty only if it's not a single-token sequence breaker
if (!is_single_token_breaker) {
int repeat_exp = af_kvp->second - ctx->dry_allowed_length;
if (max_exponent > 0 && repeat_exp > max_exponent) {
repeat_exp = max_exponent;
}
float penalty = ctx->dry_multiplier * std::pow(ctx->dry_base, repeat_exp);
cur_p->data[i].logit -= penalty;
}
}
}
cur_p->sorted = false;
}
static void llama_sampler_dry_reset(struct llama_sampler * smpl) {
auto * ctx = (llama_sampler_dry *) smpl->ctx;
ctx->last_tokens.clear();
ctx->dry_repeat_count.clear();
ctx->dry_max_token_repeat.clear();
}
static struct llama_sampler * llama_sampler_dry_clone(const struct llama_sampler * smpl) {
const auto * ctx = (llama_sampler_dry *) smpl->ctx;
llama_vocab dummy_vocab;
// dummy vocab is passed because it is only needed for raw sequence breaker processing, which we have already done and will simply be copying
auto * result = llama_sampler_init_dry_impl(dummy_vocab, ctx->total_context_size, ctx->dry_multiplier, ctx->dry_base, ctx->dry_allowed_length, ctx->dry_penalty_last_n, NULL, 0);
// Copy the state, including the processed breakers
{
auto * result_ctx = (llama_sampler_dry *) result->ctx;
result_ctx->dry_processed_breakers = ctx->dry_processed_breakers;
result_ctx->dry_repeat_count = ctx->dry_repeat_count;
result_ctx->dry_max_token_repeat = ctx->dry_max_token_repeat;
result_ctx->last_tokens = ctx->last_tokens;
}
return result;
}
static void llama_sampler_dry_free(struct llama_sampler * smpl) {
delete (llama_sampler_dry *) smpl->ctx;
}
static struct llama_sampler_i llama_sampler_dry_i = {
/* .name = */ llama_sampler_dry_name,
/* .accept = */ llama_sampler_dry_accept,
/* .apply = */ llama_sampler_dry_apply,
/* .reset = */ llama_sampler_dry_reset,
/* .clone = */ llama_sampler_dry_clone,
/* .free = */ llama_sampler_dry_free,
};
struct llama_sampler * llama_sampler_init_dry_impl(const struct llama_vocab & vocab, int32_t context_size, float dry_multiplier, float dry_base, int32_t dry_allowed_length, int32_t dry_penalty_last_n, const char** seq_breakers, size_t num_breakers) {
int32_t effective_dry_penalty_last_n = (dry_penalty_last_n == -1) ? context_size : std::max(dry_penalty_last_n, 0);
std::unordered_multimap<llama_token, std::vector<llama_token>> processed_breakers;
const int MAX_CHAR_LEN = 40;
const int MAX_SEQ_LEN = 20;
const bool dry_enabled = (dry_multiplier != 0.0f && dry_base >= 1.0f && dry_penalty_last_n != 0);
if (dry_enabled && seq_breakers != nullptr && num_breakers > 0) {
// Process sequence breakers
for (size_t i = 0; i < num_breakers; ++i) {
if (seq_breakers[i] == nullptr || std::strlen(seq_breakers[i]) == 0) {
LLAMA_LOG_WARN("skipping null or empty DRY sequence breaker at index %zu\n", i);
continue;
}
std::string sequence_break(seq_breakers[i]);
if (sequence_break.empty()) {
LLAMA_LOG_WARN("skipping empty DRY sequence breaker\n");
continue;
}
if (sequence_break.size() > MAX_CHAR_LEN) {
LLAMA_LOG_WARN("truncating DRY sequence breaker to %d characters\n", MAX_CHAR_LEN);
sequence_break.resize(MAX_CHAR_LEN);
}
get_overlapping_token_sequences(vocab, sequence_break, processed_breakers, MAX_SEQ_LEN);
}
}
return new llama_sampler {
/* .iface = */ &llama_sampler_dry_i,
/* .ctx = */ new llama_sampler_dry {
/* .total_context_size = */ context_size,
/* .dry_multiplier = */ dry_multiplier,
/* .dry_base = */ dry_base,
/* .dry_allowed_length = */ dry_allowed_length,
/* .dry_penalty_last_n = */ dry_penalty_last_n,
/* .dry_processed_breakers = */ std::move(processed_breakers),
/* .dry_repeat_count = */ dry_enabled ? std::vector<int>(effective_dry_penalty_last_n, 0) : std::vector<int>{},
/* .dry_max_token_repeat = */ {},
/* .last_tokens = */ dry_enabled ? ring_buffer<llama_token>(effective_dry_penalty_last_n) : ring_buffer<llama_token>(0),
},
};
}
// wrapper for test-sampling.cpp
struct llama_sampler * llama_sampler_init_dry_testing(int32_t context_size, float dry_multiplier, float dry_base, int32_t dry_allowed_length, int32_t dry_penalty_last_n, const std::vector<std::vector<llama_token>>& seq_breakers) {
llama_vocab dummy_vocab;
auto * result = llama_sampler_init_dry_impl(dummy_vocab, context_size, dry_multiplier, dry_base, dry_allowed_length, dry_penalty_last_n, NULL, 0);
auto * ctx = (llama_sampler_dry *) result->ctx;
// Process the token-based sequence breakers
ctx->dry_processed_breakers.clear();
if (seq_breakers.empty()) {
LLAMA_LOG_WARN("empty DRY sequence breakers list in llama_sampler_init_dry_testing\n");
} else {
for (const auto& breaker : seq_breakers) {
if (breaker.empty()) {
LLAMA_LOG_WARN("skipping DRY empty sequence breaker\n");
continue;
}
llama_token head_token = breaker[0];
std::vector<llama_token> tail_tokens(breaker.begin() + 1, breaker.end());
ctx->dry_processed_breakers.emplace(head_token, std::move(tail_tokens));
}
if (ctx->dry_processed_breakers.empty()) {
LLAMA_LOG_WARN("no valid DRY sequence breakers processed in llama_sampler_init_dry_testing\n");
}
}
return result;
}
// logit-bias
struct llama_sampler_logit_bias {
......@@ -1670,6 +2087,229 @@ struct llama_sampler * llama_sampler_init_logit_bias(
};
}
// infill
//#define GGML_DEBUG_SAMPLER_INFILL
struct llama_sampler_infill {
const struct llama_vocab * vocab;
std::vector<char> buf0;
std::vector<char> buf1;
};
static const char * llama_sampler_infill_name(const struct llama_sampler * /*smpl*/) {
return "infill";
}
static void llama_sampler_infill_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
auto * ctx = (llama_sampler_infill *) smpl->ctx;
llama_sampler_softmax_impl(cur_p);
#if defined(GGML_DEBUG_SAMPLER_INFILL)
#define LOG_DBG_CUR LLAMA_LOG_DEBUG
#else
#define LOG_DBG_CUR(...)
#endif
for (size_t i = 0; i < cur_p->size; ++i) {
LOG_DBG_CUR("%s: cur_p[%3zu] = { id: %6d, p: %.6f, logit: %6.3f }\n", __func__, i, cur_p->data[i].id, cur_p->data[i].p, cur_p->data[i].logit);
}
float p_txt_sum = 0.0f;
float p_eog_sum = 0.0f;
for (size_t i = 0; i < cur_p->size; ++i) {
if (llama_token_is_eog_impl(*ctx->vocab, cur_p->data[i].id)) {
p_eog_sum += cur_p->data[i].p;
} else {
p_txt_sum += cur_p->data[i].p;
}
}
const float rat = p_eog_sum == 0.0 ? INFINITY : p_txt_sum / p_eog_sum; GGML_UNUSED(rat);
LOG_DBG_CUR("%s: p_txt_sum = %.2f, p_eog_sum = %.2f, rat = %.2f, n = %zu\n", __func__, p_txt_sum, p_eog_sum, rat, cur_p->size);
if (3*p_eog_sum*cur_p->size > p_txt_sum) {
LOG_DBG_CUR("%s: the ratio p_txt/p_eog = %.2f is too low -> sampling EOG\n", __func__, p_txt_sum/p_eog_sum);
// keep just the EOG tokens
const auto size_org = cur_p->size;
cur_p->size = 0;
float p_sum = 0.0f;
for (size_t i = 0; i < size_org; ++i) {
if (llama_token_is_eog_impl(*ctx->vocab, cur_p->data[i].id)) {
p_sum += cur_p->data[i].p;
cur_p->data[cur_p->size++] = cur_p->data[i];
}
}
// normalize probs
for (size_t i = 0; i < cur_p->size; ++i) {
cur_p->data[i].p /= p_sum;
}
return;
}
size_t n_combined = 0; GGML_UNUSED(n_combined);
// combine tokens with common prefix
for (size_t i0 = 0; i0 < cur_p->size; ++i0) {
for (size_t i1 = 0; i1 < cur_p->size; ++i1) {
if (cur_p->data[i0].logit == -INFINITY) {
break;
}
if (i0 == i1 || cur_p->data[i1].logit == -INFINITY) {
continue;
}
int len0 = llama_token_to_piece_impl(*ctx->vocab, cur_p->data[i0].id, ctx->buf0.data(), ctx->buf0.size(), 0, false);
if (len0 < 0) {
ctx->buf0.resize(len0);
len0 = llama_token_to_piece_impl(*ctx->vocab, cur_p->data[i0].id, ctx->buf0.data(), ctx->buf0.size(), 0, false);
assert(len0 > 0);
}
int len1 = llama_token_to_piece_impl(*ctx->vocab, cur_p->data[i1].id, ctx->buf1.data(), ctx->buf1.size(), 0, false);
if (len1 < 0) {
ctx->buf1.resize(len1);
len1 = llama_token_to_piece_impl(*ctx->vocab, cur_p->data[i1].id, ctx->buf1.data(), ctx->buf1.size(), 0, false);
assert(len1 > 0);
}
// token i0 is a prefix of token i1
if (len0 > 0 && len0 <= len1 && memcmp(ctx->buf0.data(), ctx->buf1.data(), len0) == 0) {
int dst = i0;
int src = i1;
// merge into the token with higher probability
if (cur_p->data[i1].p > cur_p->data[i0].p) {
std::swap(dst, src);
}
cur_p->data[dst].p += cur_p->data[src].p;
cur_p->data[src].logit = -INFINITY;
cur_p->data[src].p = 0.0f;
n_combined++;
}
}
}
size_t n_non_eog = 0;
size_t size_org = cur_p->size;
float p_sum = 0.0f;
float thold = 0.2f;
cur_p->size = 0;
LOG_DBG_CUR("%s: n_combined = %zu, applying thold = %.3f\n", __func__, n_combined, thold);
for (size_t i = 0; i < size_org; ++i) {
const bool is_eog = llama_token_is_eog_impl(*ctx->vocab, cur_p->data[i].id);
if (cur_p->data[i].p < thold && !is_eog) {
continue;
}
if (!is_eog) {
++n_non_eog;
}
p_sum += cur_p->data[i].p;
// keep this token
cur_p->data[cur_p->size++] = cur_p->data[i];
}
LOG_DBG_CUR("%s: n_non_eog = %zu\n", __func__, n_non_eog);
// if no non-EOG tokens are left -> reduce cur_p to single EOT token
if (n_non_eog == 0) {
cur_p->size = 1;
cur_p->data[0].id = llama_token_eot_impl(*ctx->vocab);
cur_p->data[0].logit = 1.0f;
return;
}
// normalize probs
for (size_t i = 0; i < cur_p->size; ++i) {
cur_p->data[i].p /= p_sum;
LOG_DBG_CUR("%s: cur_p[%3zu] = { id: %6d, p: %.6f, logit: %6.3f }\n", __func__, i, cur_p->data[i].id, cur_p->data[i].p, cur_p->data[i].logit);
}
size_org = cur_p->size;
p_sum = 0.0f;
thold = 1.0/(n_non_eog + 1);
cur_p->size = 0;
LOG_DBG_CUR("%s: applying thold = %.3f\n", __func__, thold);
for (size_t i = 0; i < size_org; ++i) {
const bool is_eog = llama_token_is_eog_impl(*ctx->vocab, cur_p->data[i].id);
if (cur_p->data[i].p < thold && !is_eog) {
continue;
}
p_sum += cur_p->data[i].p;
cur_p->data[cur_p->size++] = cur_p->data[i];
}
// normalize probs
for (size_t i = 0; i < cur_p->size; ++i) {
cur_p->data[i].p /= p_sum;
LOG_DBG_CUR("%s: cur_p[%3zu] = { id: %6d, p: %.6f, logit: %6.3f }\n", __func__, i, cur_p->data[i].id, cur_p->data[i].p, cur_p->data[i].logit);
}
#undef LOG_DBG_CUR
}
static struct llama_sampler * llama_sampler_infill_clone(const struct llama_sampler * smpl) {
const auto * ctx = (const llama_sampler_infill *) smpl->ctx;
return llama_sampler_init_infill_impl(*ctx->vocab);
}
static void llama_sampler_infill_free(struct llama_sampler * smpl) {
delete (llama_sampler_infill *) smpl->ctx;
}
static struct llama_sampler_i llama_sampler_infill_i = {
/* .name = */ llama_sampler_infill_name,
/* .accept = */ nullptr,
/* .apply = */ llama_sampler_infill_apply,
/* .reset = */ nullptr,
/* .clone = */ llama_sampler_infill_clone,
/* .free = */ llama_sampler_infill_free,
};
struct llama_sampler * llama_sampler_init_infill_impl(
const struct llama_vocab & vocab) {
return new llama_sampler {
/* .iface = */ &llama_sampler_infill_i,
/* .ctx = */ new llama_sampler_infill {
/* .vocab = */ &vocab,
/* .buf0 = */ std::vector<char>(512),
/* .buf1 = */ std::vector<char>(512),
},
};
}
// utils
uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl) {
......
llama/llama-sampling.h
View file @ 527cc978
/**
* llama.cpp - commit 3f1ae2e32cde00c39b96be6d01c2997c29bae555 - do not edit this file
* llama.cpp - commit 40c6d79fb52f995f47507fedfeaae2ac05d9b35c - do not edit this file
*
* MIT License
*
......@@ -30,8 +30,6 @@
#include "llama-grammar.h"
#include <unordered_map>
struct llama_vocab;
struct llama_grammar;
......@@ -53,3 +51,24 @@ struct llama_sampler * llama_sampler_init_grammar_impl(
const struct llama_vocab & vocab,
const char * grammar_str,
const char * grammar_root);
struct llama_sampler * llama_sampler_init_infill_impl(
const struct llama_vocab & vocab);
struct llama_sampler * llama_sampler_init_dry_impl(
const struct llama_vocab & vocab,
int32_t context_size,
float dry_multiplier,
float dry_base,
int32_t dry_allowed_length,
int32_t dry_penalty_last_n,
const char ** seq_breakers,
size_t num_breakers);
struct llama_sampler * llama_sampler_init_dry_testing(
int32_t context_size,
float dry_multiplier,
float dry_base,
int32_t dry_allowed_length,
int32_t dry_penalty_last_n,
const std::vector<std::vector<llama_token>>& seq_breakers);
llama/llama-vocab.cpp
View file @ 527cc978
/**
* llama.cpp - commit 3f1ae2e32cde00c39b96be6d01c2997c29bae555 - do not edit this file
* llama.cpp - commit 40c6d79fb52f995f47507fedfeaae2ac05d9b35c - do not edit this file
*
* MIT License
*
......@@ -247,7 +247,7 @@ struct llm_tokenizer_spm_session {
}
// seed the work queue with all possible 2-character tokens.
for (size_t i = 1; i < symbols.size(); ++i) {
for (int i = 1; i < (int) symbols.size(); ++i) {
try_add_bigram(i - 1, i);
}
......@@ -589,7 +589,7 @@ struct llm_tokenizer_bpe_session {
index++;
symbols.emplace_back(sym);
}
for (size_t i = 1; i < symbols.size(); ++i) {
for (int i = 1; i < (int) symbols.size(); ++i) {
add_new_bigram(i - 1, i);
}
......@@ -1689,6 +1689,14 @@ llama_token llama_token_eos_impl(const struct llama_vocab & vocab) {
return vocab.special_eos_id;
}
llama_token llama_token_eot_impl(const struct llama_vocab & vocab) {
return vocab.special_eot_id;
}
llama_token llama_token_eom_impl(const struct llama_vocab & vocab) {
return vocab.special_eom_id;
}
llama_token llama_token_cls_impl(const struct llama_vocab & vocab) {
return vocab.special_cls_id;
}
......@@ -1714,23 +1722,39 @@ bool llama_add_eos_token_impl(const struct llama_vocab & vocab) {
}
llama_token llama_token_prefix_impl(const struct llama_vocab & vocab) {
return vocab.special_prefix_id;
return vocab.special_fim_pre_id;
}
llama_token llama_token_middle_impl(const struct llama_vocab & vocab) {
return vocab.special_middle_id;
return vocab.special_fim_mid_id;
}
llama_token llama_token_suffix_impl(const struct llama_vocab & vocab) {
return vocab.special_suffix_id;
return vocab.special_fim_suf_id;
}
llama_token llama_token_eot_impl(const struct llama_vocab & vocab) {
return vocab.special_eot_id;
llama_token llama_token_fim_pre_impl(const struct llama_vocab & vocab) {
return vocab.special_fim_pre_id;
}
llama_token llama_token_eom_impl(const struct llama_vocab & vocab) {
return vocab.special_eom_id;
llama_token llama_token_fim_suf_impl(const struct llama_vocab & vocab) {
return vocab.special_fim_suf_id;
}
llama_token llama_token_fim_mid_impl(const struct llama_vocab & vocab) {
return vocab.special_fim_mid_id;
}
llama_token llama_token_fim_pad_impl(const struct llama_vocab & vocab) {
return vocab.special_fim_pad_id;
}
llama_token llama_token_fim_rep_impl(const struct llama_vocab & vocab) {
return vocab.special_fim_rep_id;
}
llama_token llama_token_fim_sep_impl(const struct llama_vocab & vocab) {
return vocab.special_fim_sep_id;
}
int32_t llama_tokenize_impl(
......@@ -1968,3 +1992,19 @@ int32_t llama_detokenize_impl(
return total <= text_len_max ? total : -total;
}
std::string llama_detokenize(const struct llama_vocab & vocab, const std::vector<llama_token> & tokens, bool special) {
std::string text;
text.resize(std::max(text.capacity(), tokens.size()));
int32_t n_chars = llama_detokenize_impl(vocab, tokens.data(), (int32_t)tokens.size(), &text[0], (int32_t)text.size(), false, special);
if (n_chars < 0) {
text.resize(-n_chars);
n_chars = llama_detokenize_impl(vocab, tokens.data(), (int32_t)tokens.size(), &text[0], (int32_t)text.size(), false, special);
GGML_ASSERT(n_chars <= (int32_t)text.size()); // whitespace trimming is performed after per-token detokenization
}
text.resize(n_chars);
// NOTE: the original tokenizer decodes bytes after collecting the pieces.
return text;
}
llama/llama-vocab.h
View file @ 527cc978
/**
* llama.cpp - commit 3f1ae2e32cde00c39b96be6d01c2997c29bae555 - do not edit this file
* llama.cpp - commit 40c6d79fb52f995f47507fedfeaae2ac05d9b35c - do not edit this file
*
* MIT License
*
......@@ -63,20 +63,26 @@ struct llama_vocab {
std::map<std::pair<std::string, std::string>, int> bpe_ranks;
// default LLaMA special tokens
// TODO: should we set all of these to LLAMA_TOKEN_NULL?
id special_bos_id = 1;
id special_eos_id = 2;
id special_eot_id = LLAMA_TOKEN_NULL;
id special_eom_id = LLAMA_TOKEN_NULL;
id special_unk_id = 0;
id special_sep_id = -1;
id special_pad_id = -1;
id special_cls_id = -1;
id special_mask_id = -1;
id linefeed_id = 13;
id special_prefix_id = -1;
id special_suffix_id = -1;
id special_middle_id = -1;
id special_eot_id = -1; // TODO: move above after "eos_id", and here add "file separator" token
id special_eom_id = -1;
id special_sep_id = LLAMA_TOKEN_NULL;
id special_pad_id = LLAMA_TOKEN_NULL;
id special_cls_id = LLAMA_TOKEN_NULL;
id special_mask_id = LLAMA_TOKEN_NULL;
id linefeed_id = 13;
// fim tokens
id special_fim_pre_id = LLAMA_TOKEN_NULL;
id special_fim_suf_id = LLAMA_TOKEN_NULL;
id special_fim_mid_id = LLAMA_TOKEN_NULL;
id special_fim_pad_id = LLAMA_TOKEN_NULL;
id special_fim_rep_id = LLAMA_TOKEN_NULL; // repo
id special_fim_sep_id = LLAMA_TOKEN_NULL; // file separator
// set of all tokens that cause "end of generation"
std::set<id> special_eog_ids;
......@@ -130,19 +136,26 @@ bool llama_token_is_control_impl(const struct llama_vocab & vocab, llama_token t
llama_token llama_token_bos_impl(const struct llama_vocab & vocab);
llama_token llama_token_eos_impl(const struct llama_vocab & vocab);
llama_token llama_token_eot_impl(const struct llama_vocab & vocab);
llama_token llama_token_eom_impl(const struct llama_vocab & vocab);
llama_token llama_token_cls_impl(const struct llama_vocab & vocab);
llama_token llama_token_sep_impl(const struct llama_vocab & vocab);
llama_token llama_token_nl_impl (const struct llama_vocab & vocab);
llama_token llama_token_pad_impl(const struct llama_vocab & vocab);
bool llama_add_bos_token_impl(const struct llama_vocab & vocab);
bool llama_add_eos_token_impl(const struct llama_vocab & vocab);
llama_token llama_token_prefix_impl(const struct llama_vocab & vocab);
llama_token llama_token_middle_impl(const struct llama_vocab & vocab);
llama_token llama_token_suffix_impl(const struct llama_vocab & vocab);
llama_token llama_token_eot_impl (const struct llama_vocab & vocab);
llama_token llama_token_eom_impl (const struct llama_vocab & vocab);
llama_token llama_token_fim_pre_impl(const struct llama_vocab & vocab);
llama_token llama_token_fim_suf_impl(const struct llama_vocab & vocab);
llama_token llama_token_fim_mid_impl(const struct llama_vocab & vocab);
llama_token llama_token_fim_pad_impl(const struct llama_vocab & vocab);
llama_token llama_token_fim_rep_impl(const struct llama_vocab & vocab);
llama_token llama_token_fim_sep_impl(const struct llama_vocab & vocab);
bool llama_add_bos_token_impl(const struct llama_vocab & vocab);
bool llama_add_eos_token_impl(const struct llama_vocab & vocab);
int32_t llama_tokenize_impl(
const struct llama_vocab & vocab,
......@@ -162,6 +175,12 @@ int32_t llama_token_to_piece_impl(
int32_t lstrip,
bool special);
// check if token0 is contained as a prefix in token1
bool llama_token_is_prefix_impl(
const struct llama_vocab & vocab,
llama_token token0,
llama_token token1);
int32_t llama_detokenize_impl(
const struct llama_vocab & vocab,
const llama_token * tokens,
......@@ -170,3 +189,8 @@ int32_t llama_detokenize_impl(
int32_t text_len_max,
bool remove_special,
bool unparse_special);
std::string llama_detokenize(
const struct llama_vocab & vocab,
const std::vector<llama_token> & tokens,
bool special);
llama/llama.cpp
View file @ 527cc978
This source diff could not be displayed because it is too large. You can view the blob instead.
llama/llama.go
View file @ 527cc978
......@@ -3,12 +3,12 @@ package llama
//go:generate make -j 8
/*
#cgo CFLAGS: -O2 -std=c11 -DGGML_BUILD=1 -DNDEBUG -DLOG_DISABLE_LOGS -DGGML_USE_LLAMAFILE
#cgo CXXFLAGS: -O2 -std=c++11 -DGGML_BUILD=1 -DNDEBUG -DLOG_DISABLE_LOGS -DGGML_USE_LLAMAFILE
#cgo CFLAGS: -O3 -std=c17 -DGGML_BUILD=1 -DNDEBUG -DLOG_DISABLE_LOGS -DGGML_USE_LLAMAFILE -DGGML_USE_CPU -DGGML_USE_CPU_AARCH64
#cgo CXXFLAGS: -O3 -std=c++17 -DGGML_BUILD=1 -DNDEBUG -DLOG_DISABLE_LOGS -DGGML_USE_LLAMAFILE -DGGML_USE_CPU -DGGML_USE_CPU_AARCH64
#cgo amd64,avx CFLAGS: -mavx
#cgo amd64,avx CXXFLAGS: -mavx
#cgo amd64,avx2 CFLAGS: -mavx2 -mfma
#cgo amd64,avx2 CXXFLAGS: -mavx2 -mfma
#cgo amd64,avx2 CFLAGS: -mavx2 -mfma -mf16c
#cgo amd64,avx2 CXXFLAGS: -mavx2 -mfma -mf16c
#cgo amd64,avx512 CFLAGS: -mavx512f -mavx512dq -mavx512bw
#cgo amd64,avx512 CXXFLAGS: -mavx512f -mavx512dq -mavx512bw
#cgo amd64,avx512bf16 CFLAGS: -mavx512bf16 -D__AVX512BF16__
......@@ -33,21 +33,22 @@ package llama
#cgo darwin,amd64,avx2 CFLAGS: -DGGML_USE_ACCELERATE -DACCELERATE_NEW_LAPACK -DACCELERATE_LAPACK_ILP64
#cgo darwin,amd64,avx2 CXXFLAGS: -DGGML_USE_ACCELERATE -DACCELERATE_NEW_LAPACK -DACCELERATE_LAPACK_ILP64
#cgo darwin,amd64,avx2 LDFLAGS: -framework Accelerate
#cgo darwin,arm64 CFLAGS: -DGGML_USE_METAL -DGGML_USE_ACCELERATE -DGGML_METAL_EMBED_LIBRARY -DACCELERATE_NEW_LAPACK -DACCELERATE_LAPACK_ILP64 -DGGML_USE_BLAS
#cgo darwin,arm64 CXXFLAGS: -DGGML_USE_METAL -DGGML_USE_ACCELERATE -DGGML_METAL_EMBED_LIBRARY -DACCELERATE_NEW_LAPACK -DACCELERATE_LAPACK_ILP64 -DGGML_USE_BLAS
#cgo darwin,arm64 CFLAGS: -DGGML_USE_METAL -DGGML_USE_ACCELERATE -DGGML_METAL_EMBED_LIBRARY -DACCELERATE_NEW_LAPACK -DACCELERATE_LAPACK_ILP64 -DGGML_USE_BLAS -DGGML_BLAS_USE_ACCELERATE
#cgo darwin,arm64 CXXFLAGS: -DGGML_USE_METAL -DGGML_USE_ACCELERATE -DGGML_METAL_EMBED_LIBRARY -DACCELERATE_NEW_LAPACK -DACCELERATE_LAPACK_ILP64 -DGGML_USE_BLAS -DGGML_BLAS_USE_ACCELERATE
#cgo darwin,arm64 LDFLAGS: -framework Foundation -framework Metal -framework MetalKit -framework Accelerate
#cgo linux CFLAGS: -D_GNU_SOURCE
#cgo linux CXXFLAGS: -D_GNU_SOURCE
#cgo linux LDFLAGS: -ldl
#cgo linux,amd64 LDFLAGS: -L${SRCDIR}/build/linux-amd64
#cgo linux,arm64 CFLAGS: -D__aarch64__ -D__ARM_NEON -D__ARM_FEATURE_FMA
#cgo linux,arm64 CXXFLAGS: -D__aarch64__ -D__ARM_NEON -D__ARM_FEATURE_FMA
#cgo linux,arm64 LDFLAGS: -L${SRCDIR}/build/linux-arm64
#cgo linux,arm64,sve CFLAGS: -march=armv8.6-a+sve
#cgo linux,arm64,sve CXXFLAGS: -march=armv8.6-a+sve
#cgo linux,cuda LDFLAGS: -lcuda -lcudart -lcublas -lcublasLt -lpthread -ldl -lrt -lresolv
#cgo linux,rocm LDFLAGS: -lpthread -ldl -lrt -lresolv
#cgo rocm CFLAGS: -DGGML_USE_CUDA -DGGML_USE_HIPBLAS -DGGML_CUDA_DMMV_X=32 -DGGML_CUDA_PEER_MAX_BATCH_SIZE=128 -DGGML_CUDA_MMV_Y=1 -DGGML_BUILD=1
#cgo rocm CXXFLAGS: -DGGML_USE_CUDA -DGGML_USE_HIPBLAS -DGGML_CUDA_DMMV_X=32 -DGGML_CUDA_PEER_MAX_BATCH_SIZE=128 -DGGML_CUDA_MMV_Y=1 -DGGML_BUILD=1
#cgo linux,cuda LDFLAGS: -lcuda -lcudart -lcublas -lcublasLt -lpthread -lrt -lresolv
#cgo linux,rocm LDFLAGS: -lpthread -lrt -lresolv
#cgo rocm CFLAGS: -DGGML_USE_CUDA -DGGML_USE_HIP -DGGML_CUDA_DMMV_X=32 -DGGML_CUDA_PEER_MAX_BATCH_SIZE=128 -DGGML_CUDA_MMV_Y=1 -DGGML_BUILD=1
#cgo rocm CXXFLAGS: -DGGML_USE_CUDA -DGGML_USE_HIP -DGGML_CUDA_DMMV_X=32 -DGGML_CUDA_PEER_MAX_BATCH_SIZE=128 -DGGML_CUDA_MMV_Y=1 -DGGML_BUILD=1
#cgo rocm LDFLAGS: -L${SRCDIR} -lggml_rocm -lhipblas -lamdhip64 -lrocblas
#cgo windows CFLAGS: -Wno-discarded-qualifiers -D_WIN32_WINNT=0x602
#cgo windows CXXFLAGS: -D_WIN32_WINNT=0x602
......@@ -67,7 +68,8 @@ package llama
#include "mllama.h"
#include "sampling_ext.h"
bool llamaProgressCallback(float progress, void *user_data);
extern bool llamaProgressCallback(float progress, void *user_data);
extern void llamaLog(int level, char* text, void* user_data);
typedef enum {COMP_UNKNOWN,COMP_GCC,COMP_CLANG} COMPILER;
COMPILER inline get_compiler() {
......@@ -79,6 +81,7 @@ COMPILER inline get_compiler() {
return UNKNOWN_COMPILER;
#endif
}
*/
import "C"
......@@ -93,6 +96,7 @@ import (
"runtime/cgo"
"slices"
"strings"
"sync/atomic"
"unsafe"
)
......@@ -113,6 +117,26 @@ func PrintSystemInfo() string {
return C.GoString(C.llama_print_system_info()) + compiler
}
var logLevel atomic.Int32
func init() {
logLevel.Store(int32(C.GGML_LOG_LEVEL_INFO))
C.llama_log_set((C.ggml_log_callback)(C.llamaLog), nil)
}
func EnableDebug() {
logLevel.Store(int32(C.GGML_LOG_LEVEL_DEBUG))
}
//export llamaLog
func llamaLog(level int32, text *C.char, _ unsafe.Pointer) {
if level < logLevel.Load() {
return
}
fmt.Print(C.GoString(text))
}
func GetModelArch(modelPath string) (string, error) {
mp := C.CString(modelPath)
defer C.free(unsafe.Pointer(mp))
......@@ -633,14 +657,13 @@ func (c *Context) Synchronize() {
// sampling
// TODO: this is a temporary wrapper to allow calling C++ code from CGo
type SamplingContext struct {
c *C.struct_gpt_sampler
c *C.struct_common_sampler
}
type SamplingParams struct {
TopK int
TopP float32
MinP float32
TfsZ float32
TypicalP float32
Temp float32
RepeatLastN int
......@@ -656,11 +679,10 @@ type SamplingParams struct {
}
func NewSamplingContext(model *Model, params SamplingParams) (*SamplingContext, error) {
var cparams C.struct_gpt_sampler_cparams
var cparams C.struct_common_sampler_cparams
cparams.top_k = C.int32_t(params.TopK)
cparams.top_p = C.float(params.TopP)
cparams.min_p = C.float(params.MinP)
cparams.tfs_z = C.float(params.TfsZ)
cparams.typical_p = C.float(params.TypicalP)
cparams.temp = C.float(params.Temp)
cparams.penalty_last_n = C.int32_t(params.RepeatLastN)
......@@ -677,26 +699,26 @@ func NewSamplingContext(model *Model, params SamplingParams) (*SamplingContext,
defer C.free(unsafe.Pointer(grammar))
cparams.grammar = grammar
context := &SamplingContext{c: C.gpt_sampler_cinit(model.c, &cparams)}
context := &SamplingContext{c: C.common_sampler_cinit(model.c, &cparams)}
if context.c == nil {
return nil, errors.New("unable to create sampling context")
}
runtime.SetFinalizer(context, func(s *SamplingContext) { C.gpt_sampler_cfree(s.c) })
runtime.SetFinalizer(context, func(s *SamplingContext) { C.common_sampler_cfree(s.c) })
return context, nil
}
func (s *SamplingContext) Reset() {
C.gpt_sampler_creset(s.c)
C.common_sampler_creset(s.c)
}
func (s *SamplingContext) Sample(llamaContext *Context, idx int) int {
return int(C.gpt_sampler_csample(s.c, llamaContext.c, C.int(idx)))
return int(C.common_sampler_csample(s.c, llamaContext.c, C.int(idx)))
}
func (s *SamplingContext) Accept(id int, applyGrammar bool) {
C.gpt_sampler_caccept(s.c, C.llama_token(id), C.bool(applyGrammar))
C.common_sampler_caccept(s.c, C.llama_token(id), C.bool(applyGrammar))
}
type JsonSchema struct {
......
llama/llama.h
View file @ 527cc978
/**
* llama.cpp - commit 3f1ae2e32cde00c39b96be6d01c2997c29bae555 - do not edit this file
* llama.cpp - commit 40c6d79fb52f995f47507fedfeaae2ac05d9b35c - do not edit this file
*
* MIT License
*
......@@ -28,6 +28,7 @@
#define LLAMA_H
#include "ggml.h"
#include "ggml-cpu.h"
#include "ggml-backend.h"
#include <stddef.h>
......@@ -231,7 +232,7 @@ extern "C" {
enum llama_split_mode {
LLAMA_SPLIT_MODE_NONE = 0, // single GPU
LLAMA_SPLIT_MODE_LAYER = 1, // split layers and KV across GPUs
LLAMA_SPLIT_MODE_ROW = 2, // split rows across GPUs
LLAMA_SPLIT_MODE_ROW = 2, // split layers and KV across GPUs, use tensor parallelism if supported
};
// TODO: simplify (https://github.com/ggerganov/llama.cpp/pull/9294#pullrequestreview-2286561979)
......@@ -243,6 +244,7 @@ extern "C" {
typedef struct llama_token_data_array {
// TODO: consider SoA
// NOTE: this pointer can be modified by the samplers
llama_token_data * data;
size_t size;
int64_t selected; // this is the index in the data array (i.e. not the token id)
......@@ -258,8 +260,11 @@ extern "C" {
// - token : the token ids of the input (used when embd is NULL)
// - embd : token embeddings (i.e. float vector of size n_embd) (used when token is NULL)
// - pos : the positions of the respective token in the sequence
// (if set to NULL, the token position will be tracked automatically by llama_decode)
// - seq_id : the sequence to which the respective token belongs
// (if set to NULL, the sequence ID will be assumed to be 0)
// - logits : if zero, the logits (and/or the embeddings) for the respective token will not be output
// (if set to NULL, only the logits for last token will be returned)
//
typedef struct llama_batch {
int32_t n_tokens;
......@@ -271,15 +276,6 @@ extern "C" {
int32_t * n_seq_id;
llama_seq_id ** seq_id;
int8_t * logits; // TODO: rename this to "output"
// NOTE: helpers for smooth API transition - can be deprecated in the future
// for future-proof code, use the above fields instead and ignore everything below
//
// pos[i] = all_pos_0 + i*all_pos_1
//
llama_pos all_pos_0; // used if pos == NULL
llama_pos all_pos_1; // used if pos == NULL
llama_seq_id all_seq_id; // used if seq_id == NULL
} llama_batch;
enum llama_model_kv_override_type {
......@@ -303,13 +299,13 @@ extern "C" {
};
struct llama_model_params {
// NULL-terminated list of devices to use for offloading (if NULL, all available devices are used)
ggml_backend_dev_t * devices;
int32_t n_gpu_layers; // number of layers to store in VRAM
enum llama_split_mode split_mode; // how to split the model across multiple GPUs
// main_gpu interpretation depends on split_mode:
// LLAMA_SPLIT_MODE_NONE: the GPU that is used for the entire model
// LLAMA_SPLIT_MODE_ROW: the GPU that is used for small tensors and intermediate results
// LLAMA_SPLIT_MODE_LAYER: ignored
// the GPU that is used for the entire model when split_mode is LLAMA_SPLIT_MODE_NONE
int32_t main_gpu;
// proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices()
......@@ -464,6 +460,7 @@ extern "C" {
LLAMA_API bool llama_supports_mmap (void);
LLAMA_API bool llama_supports_mlock (void);
LLAMA_API bool llama_supports_gpu_offload(void);
LLAMA_API bool llama_supports_rpc (void);
LLAMA_API uint32_t llama_n_ctx (const struct llama_context * ctx);
LLAMA_API uint32_t llama_n_batch (const struct llama_context * ctx);
......@@ -704,6 +701,9 @@ extern "C" {
// Apply the KV cache updates (such as K-shifts, defragmentation, etc.)
LLAMA_API void llama_kv_cache_update(struct llama_context * ctx);
// Check if the context supports KV cache shifting
LLAMA_API bool llama_kv_cache_can_shift(struct llama_context * ctx);
//
// State / sessions
//
......@@ -806,15 +806,15 @@ extern "C" {
// Decoding
//
// Return batch for single sequence of tokens starting at pos_0
// Return batch for single sequence of tokens
// The sequence ID will be fixed to 0
// The position of the tokens will be tracked automatically by llama_decode
//
// NOTE: this is a helper function to facilitate transition to the new batch API - avoid using it
//
LLAMA_API struct llama_batch llama_batch_get_one(
llama_token * tokens,
int32_t n_tokens,
llama_pos pos_0,
llama_seq_id seq_id);
int32_t n_tokens);
// Allocates a batch of tokens on the heap that can hold a maximum of n_tokens
// Each token can be assigned up to n_seq_max sequence ids
......@@ -834,7 +834,7 @@ extern "C" {
// Processes a batch of tokens with the ecoder part of the encoder-decoder model.
// Stores the encoder output internally for later use by the decoder cross-attention layers.
// 0 - success
// < 0 - error
// < 0 - error. the KV cache state is restored to the state before this call
LLAMA_API int32_t llama_encode(
struct llama_context * ctx,
struct llama_batch batch);
......@@ -842,7 +842,7 @@ extern "C" {
// Positive return values does not mean a fatal error, but rather a warning.
// 0 - success
// 1 - could not find a KV slot for the batch (try reducing the size of the batch or increase the context)
// < 0 - error
// < 0 - error. the KV cache state is restored to the state before this call
LLAMA_API int32_t llama_decode(
struct llama_context * ctx,
struct llama_batch batch);
......@@ -927,6 +927,7 @@ extern "C" {
// Special tokens
LLAMA_API llama_token llama_token_bos(const struct llama_model * model); // beginning-of-sentence
LLAMA_API llama_token llama_token_eos(const struct llama_model * model); // end-of-sentence
LLAMA_API llama_token llama_token_eot(const struct llama_model * model); // end-of-turn
LLAMA_API llama_token llama_token_cls(const struct llama_model * model); // classification
LLAMA_API llama_token llama_token_sep(const struct llama_model * model); // sentence separator
LLAMA_API llama_token llama_token_nl (const struct llama_model * model); // next-line
......@@ -935,11 +936,17 @@ extern "C" {
LLAMA_API bool llama_add_bos_token(const struct llama_model * model);
LLAMA_API bool llama_add_eos_token(const struct llama_model * model);
// Codellama infill tokens
LLAMA_API llama_token llama_token_prefix(const struct llama_model * model); // Beginning of infill prefix
LLAMA_API llama_token llama_token_middle(const struct llama_model * model); // Beginning of infill middle
LLAMA_API llama_token llama_token_suffix(const struct llama_model * model); // Beginning of infill suffix
LLAMA_API llama_token llama_token_eot (const struct llama_model * model); // End of infill middle
// infill tokens
DEPRECATED(LLAMA_API llama_token llama_token_prefix(const struct llama_model * model), "use llama_token_fim_pre instead");
DEPRECATED(LLAMA_API llama_token llama_token_middle(const struct llama_model * model), "use llama_token_fim_mid instead");
DEPRECATED(LLAMA_API llama_token llama_token_suffix(const struct llama_model * model), "use llama_token_fim_suf instead");
LLAMA_API llama_token llama_token_fim_pre(const struct llama_model * model);
LLAMA_API llama_token llama_token_fim_suf(const struct llama_model * model);
LLAMA_API llama_token llama_token_fim_mid(const struct llama_model * model);
LLAMA_API llama_token llama_token_fim_pad(const struct llama_model * model);
LLAMA_API llama_token llama_token_fim_rep(const struct llama_model * model);
LLAMA_API llama_token llama_token_fim_sep(const struct llama_model * model);
//
// Tokenization
......@@ -1014,6 +1021,9 @@ extern "C" {
char * buf,
int32_t length);
// Get list of built-in chat templates
LLAMA_API int32_t llama_chat_builtin_templates(const char ** output, size_t len);
//
// Sampling API
//
......@@ -1098,12 +1108,13 @@ extern "C" {
// available samplers:
LLAMA_API struct llama_sampler * llama_sampler_init_greedy (void);
LLAMA_API struct llama_sampler * llama_sampler_init_dist (uint32_t seed);
LLAMA_API struct llama_sampler * llama_sampler_init_greedy(void);
LLAMA_API struct llama_sampler * llama_sampler_init_dist (uint32_t seed);
/// @details Sorts candidate tokens by their logits in descending order and calculate probabilities based on logits.
/// NOTE: Avoid using on the full vocabulary as the sorting can become slow. For example, apply top-k or top-p sampling first.
LLAMA_API struct llama_sampler * llama_sampler_init_softmax (void);
DEPRECATED(LLAMA_API struct llama_sampler * llama_sampler_init_softmax (void),
"will be removed in the future (see https://github.com/ggerganov/llama.cpp/pull/9896#discussion_r1800920915)");
/// @details Top-K sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751
LLAMA_API struct llama_sampler * llama_sampler_init_top_k (int32_t k);
......@@ -1114,16 +1125,18 @@ extern "C" {
/// @details Minimum P sampling as described in https://github.com/ggerganov/llama.cpp/pull/3841
LLAMA_API struct llama_sampler * llama_sampler_init_min_p (float p, size_t min_keep);
/// @details Tail Free Sampling described in https://www.trentonbricken.com/Tail-Free-Sampling/.
LLAMA_API struct llama_sampler * llama_sampler_init_tail_free (float z, size_t min_keep);
/// @details Locally Typical Sampling implementation described in the paper https://arxiv.org/abs/2202.00666.
LLAMA_API struct llama_sampler * llama_sampler_init_typical (float p, size_t min_keep);
/// #details Updates the logits l_i` = l_i/t. When t <= 0.0f, the maximum logit is kept at it's original value, the rest are set to -inf
LLAMA_API struct llama_sampler * llama_sampler_init_temp (float t);
/// @details Dynamic temperature implementation (a.k.a. entropy) described in the paper https://arxiv.org/abs/2309.02772.
LLAMA_API struct llama_sampler * llama_sampler_init_temp_ext (float t, float delta, float exponent);
/// @details XTC sampler as described in https://github.com/oobabooga/text-generation-webui/pull/6335
LLAMA_API struct llama_sampler * llama_sampler_init_xtc (float p, float t, size_t min_keep, uint32_t seed);
/// @details Mirostat 1.0 algorithm described in the paper https://arxiv.org/abs/2007.14966. Uses tokens instead of words.
/// @param candidates A vector of `llama_token_data` containing the candidate tokens, their probabilities (p), and log-odds (logit) for the current position in the generated text.
/// @param tau The target cross-entropy (or surprise) value you want to achieve for the generated text. A higher value corresponds to more surprising or less predictable text, while a lower value corresponds to less surprising or more predictable text.
......@@ -1163,11 +1176,43 @@ extern "C" {
bool penalize_nl, // consider newlines as a repeatable token
bool ignore_eos); // ignore the end-of-sequence token
/// @details DRY sampler, designed by p-e-w, as described in: https://github.com/oobabooga/text-generation-webui/pull/5677, porting Koboldcpp implementation authored by pi6am: https://github.com/LostRuins/koboldcpp/pull/982
LLAMA_API struct llama_sampler * llama_sampler_init_dry(
const struct llama_model * model,
float dry_multiplier,
float dry_base,
int32_t dry_allowed_length,
int32_t dry_penalty_last_n,
const char ** seq_breakers,
size_t num_breakers);
LLAMA_API struct llama_sampler * llama_sampler_init_logit_bias(
int32_t n_vocab,
int32_t n_logit_bias,
const llama_logit_bias * logit_bias);
// this sampler is meant to be used for fill-in-the-middle infilling
// it's supposed to be used after top_k + top_p sampling
//
// 1. if the sum of the EOG probs times the number of candidates is higher than the sum of the other probs -> pick EOG
// 2. combine probs of tokens that have the same prefix
//
// example:
//
// - before:
// "hel": 0.5
// "hell": 0.2
// "hello": 0.1
// "dummy": 0.1
//
// - after:
// "hel": 0.8
// "dummy": 0.1
//
// 3. discard non-EOG tokens with low prob
// 4. if no tokens are left -> pick EOT
//
LLAMA_API struct llama_sampler * llama_sampler_init_infill(const struct llama_model * model);
// Returns the seed used by the sampler if applicable, LLAMA_DEFAULT_SEED otherwise
LLAMA_API uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl);
......@@ -1239,8 +1284,6 @@ extern "C" {
LLAMA_API void llama_perf_sampler_print(const struct llama_sampler * chain);
LLAMA_API void llama_perf_sampler_reset( struct llama_sampler * chain);
LLAMA_API void llama_perf_dump_yaml(FILE * stream, const struct llama_context * ctx);
#ifdef __cplusplus
}
#endif
......
llama/llama_darwin.c deleted 100644 → 0
View file @ a37f4a86
const char *ggml_metallib_start;
const char *ggml_metallib_end;
llama/llama_darwin.go deleted 100644 → 0
View file @ a37f4a86
package llama
// extern const char *ggml_metallib_start;
// extern const char *ggml_metallib_end;
import "C"
import (
_ "embed"
"strings"
"unsafe"
)
//go:embed ggml-common.h
var ggmlCommon string
//go:embed ggml-metal.metal
var ggmlMetal string
func init() {
metal := strings.ReplaceAll(ggmlMetal, `#include "ggml-common.h"`, ggmlCommon)
cMetal := C.CString(metal)
C.ggml_metallib_start = cMetal
C.ggml_metallib_end = (*C.char)(unsafe.Pointer(uintptr(unsafe.Pointer(cMetal)) + uintptr(len(metal))))
}
llama/llava.cpp
View file @ 527cc978
/**
* llama.cpp - commit 3f1ae2e32cde00c39b96be6d01c2997c29bae555 - do not edit this file
* llama.cpp - commit 40c6d79fb52f995f47507fedfeaae2ac05d9b35c - do not edit this file
*
* MIT License
*
......@@ -37,13 +37,17 @@
#include <limits>
#include <vector>
#define die(msg) do { fputs("error: " msg "\n", stderr); exit(1); } while (0)
#define die_fmt(fmt, ...) do { fprintf(stderr, "error: " fmt "\n", __VA_ARGS__); exit(1); } while (0)
#define LOG_INF(...) do { fprintf(stdout, __VA_ARGS__); } while (0)
#define LOG_WRN(...) do { fprintf(stderr, __VA_ARGS__); } while (0)
#define LOG_ERR(...) do { fprintf(stderr, __VA_ARGS__); } while (0)
#define LOG_DBG(...) do { fprintf(stdout, __VA_ARGS__); } while (0)
#if defined(LLAVA_LOG_OFF)
# define LOG_INF(...)
# define LOG_WRN(...)
# define LOG_ERR(...)
# define LOG_DBG(...)
#else // defined(LLAVA_LOG_OFF)
# define LOG_INF(...) do { fprintf(stdout, __VA_ARGS__); } while (0)
# define LOG_WRN(...) do { fprintf(stderr, __VA_ARGS__); } while (0)
# define LOG_ERR(...) do { fprintf(stderr, __VA_ARGS__); } while (0)
# define LOG_DBG(...) do { fprintf(stdout, __VA_ARGS__); } while (0)
#endif // defined(LLAVA_LOG_OFF)
// RGB uint8 image
struct clip_image_u8 {
......@@ -427,6 +431,40 @@ bool llava_image_embed_make_with_clip_img(clip_ctx * ctx_clip, int n_threads, co
return true;
}
struct llava_embd_batch {
std::vector<llama_pos> pos;
std::vector<int32_t> n_seq_id;
std::vector<llama_seq_id> seq_id_0;
std::vector<llama_seq_id *> seq_ids;
std::vector<int8_t> logits;
llama_batch batch;
llava_embd_batch(float * embd, int32_t n_embd, int32_t n_tokens, llama_pos pos_0, llama_seq_id seq_id) {
pos .resize(n_tokens);
n_seq_id.resize(n_tokens);
seq_ids .resize(n_tokens + 1);
logits .resize(n_tokens);
seq_id_0.resize(1);
seq_id_0[0] = seq_id;
seq_ids [n_tokens] = nullptr;
batch = {
/*n_tokens =*/ n_tokens,
/*tokens =*/ nullptr,
/*embd =*/ embd,
/*n_embd =*/ n_embd,
/*pos =*/ pos.data(),
/*n_seq_id =*/ n_seq_id.data(),
/*seq_id =*/ seq_ids.data(),
/*logits =*/ logits.data(),
};
for (int i = 0; i < n_tokens; i++) {
batch.pos [i] = pos_0 + i;
batch.n_seq_id[i] = 1;
batch.seq_id [i] = seq_id_0.data();
batch.logits [i] = false;
}
}
};
bool llava_eval_image_embed(llama_context * ctx_llama, const struct llava_image_embed * image_embed, int n_batch, int * n_past) {
int n_embd = llama_n_embd(llama_get_model(ctx_llama));
......@@ -435,8 +473,9 @@ bool llava_eval_image_embed(llama_context * ctx_llama, const struct llava_image_
if (n_eval > n_batch) {
n_eval = n_batch;
}
llama_batch batch = {int32_t(n_eval), nullptr, (image_embed->embed+i*n_embd), n_embd, nullptr, nullptr, nullptr, nullptr, *n_past, 1, 0, };
if (llama_decode(ctx_llama, batch)) {
float * embd = image_embed->embed+i*n_embd;
llava_embd_batch llava_batch = llava_embd_batch(embd, n_embd, n_eval, *n_past, 0);
if (llama_decode(ctx_llama, llava_batch.batch)) {
LOG_ERR("%s : failed to eval\n", __func__);
return false;
}
......@@ -458,7 +497,7 @@ struct llava_image_embed * llava_image_embed_make_with_bytes(struct clip_ctx * c
bool image_embed_result = llava_image_embed_make_with_clip_img(ctx_clip, n_threads, img, &image_embed, &n_image_pos);
if (!image_embed_result) {
clip_image_u8_free(img);
LOG_ERR("%s: coulnd't embed the image\n", __func__);
LOG_ERR("%s: couldn't embed the image\n", __func__);
return NULL;
}
......@@ -490,10 +529,16 @@ static bool load_file_to_bytes(const char* path, unsigned char** bytesOut, long
errno = 0;
size_t ret = fread(buffer, 1, fileSize, file); // Read the file into the buffer
if (ferror(file)) {
die_fmt("read error: %s", strerror(errno));
LOG_ERR("read error: %s", strerror(errno));
free(buffer);
fclose(file);
return false;
}
if (ret != (size_t) fileSize) {
die("unexpectedly reached end of file");
LOG_ERR("unexpectedly reached end of file");
free(buffer);
fclose(file);
return false;
}
fclose(file); // Close the file
......
llama/llava.h
View file @ 527cc978
/**
* llama.cpp - commit 3f1ae2e32cde00c39b96be6d01c2997c29bae555 - do not edit this file
* llama.cpp - commit 40c6d79fb52f995f47507fedfeaae2ac05d9b35c - do not edit this file
*
* MIT License
*
......
llama/log.cpp
View file @ 527cc978
/**
* llama.cpp - commit 3f1ae2e32cde00c39b96be6d01c2997c29bae555 - do not edit this file
* llama.cpp - commit 40c6d79fb52f995f47507fedfeaae2ac05d9b35c - do not edit this file
*
* MIT License
*
......@@ -34,10 +34,10 @@
#include <thread>
#include <vector>
int gpt_log_verbosity_thold = LOG_DEFAULT_LLAMA;
int common_log_verbosity_thold = LOG_DEFAULT_LLAMA;
void gpt_log_set_verbosity_thold(int verbosity) {
gpt_log_verbosity_thold = verbosity;
void common_log_set_verbosity_thold(int verbosity) {
common_log_verbosity_thold = verbosity;
}
#define LOG_COL_DEFAULT "\033[0m"
......@@ -55,16 +55,16 @@ static int64_t t_us() {
}
// colors
enum gpt_log_col : int {
GPT_LOG_COL_DEFAULT = 0,
GPT_LOG_COL_BOLD,
GPT_LOG_COL_RED,
GPT_LOG_COL_GREEN,
GPT_LOG_COL_YELLOW,
GPT_LOG_COL_BLUE,
GPT_LOG_COL_MAGENTA,
GPT_LOG_COL_CYAN,
GPT_LOG_COL_WHITE,
enum common_log_col : int {
COMMON_LOG_COL_DEFAULT = 0,
COMMON_LOG_COL_BOLD,
COMMON_LOG_COL_RED,
COMMON_LOG_COL_GREEN,
COMMON_LOG_COL_YELLOW,
COMMON_LOG_COL_BLUE,
COMMON_LOG_COL_MAGENTA,
COMMON_LOG_COL_CYAN,
COMMON_LOG_COL_WHITE,
};
// disable colors by default
......@@ -80,7 +80,7 @@ static std::vector<const char *> g_col = {
"",
};
struct gpt_log_entry {
struct common_log_entry {
enum ggml_log_level level;
bool prefix;
......@@ -97,7 +97,7 @@ struct gpt_log_entry {
if (!fcur) {
// stderr displays DBG messages only when their verbosity level is not higher than the threshold
// these messages will still be logged to a file
if (level == GGML_LOG_LEVEL_DEBUG && gpt_log_verbosity_thold < LOG_DEFAULT_DEBUG) {
if (level == GGML_LOG_LEVEL_DEBUG && common_log_verbosity_thold < LOG_DEFAULT_DEBUG) {
return;
}
......@@ -112,19 +112,19 @@ struct gpt_log_entry {
if (timestamp) {
// [M.s.ms.us]
fprintf(fcur, "%s%d.%02d.%03d.%03d%s ",
g_col[GPT_LOG_COL_BLUE],
g_col[COMMON_LOG_COL_BLUE],
(int) (timestamp / 1000000 / 60),
(int) (timestamp / 1000000 % 60),
(int) (timestamp / 1000 % 1000),
(int) (timestamp % 1000),
g_col[GPT_LOG_COL_DEFAULT]);
g_col[COMMON_LOG_COL_DEFAULT]);
}
switch (level) {
case GGML_LOG_LEVEL_INFO: fprintf(fcur, "%sI %s", g_col[GPT_LOG_COL_GREEN], g_col[GPT_LOG_COL_DEFAULT]); break;
case GGML_LOG_LEVEL_WARN: fprintf(fcur, "%sW %s", g_col[GPT_LOG_COL_MAGENTA], "" ); break;
case GGML_LOG_LEVEL_ERROR: fprintf(fcur, "%sE %s", g_col[GPT_LOG_COL_RED], "" ); break;
case GGML_LOG_LEVEL_DEBUG: fprintf(fcur, "%sD %s", g_col[GPT_LOG_COL_YELLOW], "" ); break;
case GGML_LOG_LEVEL_INFO: fprintf(fcur, "%sI %s", g_col[COMMON_LOG_COL_GREEN], g_col[COMMON_LOG_COL_DEFAULT]); break;
case GGML_LOG_LEVEL_WARN: fprintf(fcur, "%sW %s", g_col[COMMON_LOG_COL_MAGENTA], "" ); break;
case GGML_LOG_LEVEL_ERROR: fprintf(fcur, "%sE %s", g_col[COMMON_LOG_COL_RED], "" ); break;
case GGML_LOG_LEVEL_DEBUG: fprintf(fcur, "%sD %s", g_col[COMMON_LOG_COL_YELLOW], "" ); break;
default:
break;
}
......@@ -133,18 +133,18 @@ struct gpt_log_entry {
fprintf(fcur, "%s", msg.data());
if (level == GGML_LOG_LEVEL_WARN || level == GGML_LOG_LEVEL_ERROR || level == GGML_LOG_LEVEL_DEBUG) {
fprintf(fcur, "%s", g_col[GPT_LOG_COL_DEFAULT]);
fprintf(fcur, "%s", g_col[COMMON_LOG_COL_DEFAULT]);
}
fflush(fcur);
}
};
struct gpt_log {
struct common_log {
// default capacity - will be expanded if needed
gpt_log() : gpt_log(256) {}
common_log() : common_log(256) {}
gpt_log(size_t capacity) {
common_log(size_t capacity) {
file = nullptr;
prefix = false;
timestamps = false;
......@@ -163,7 +163,7 @@ struct gpt_log {
resume();
}
~gpt_log() {
~common_log() {
pause();
if (file) {
fclose(file);
......@@ -184,12 +184,12 @@ private:
int64_t t_start;
// ring buffer of entries
std::vector<gpt_log_entry> entries;
std::vector<common_log_entry> entries;
size_t head;
size_t tail;
// worker thread copies into this
gpt_log_entry cur;
common_log_entry cur;
public:
void add(enum ggml_log_level level, const char * fmt, va_list args) {
......@@ -245,7 +245,7 @@ public:
tail = (tail + 1) % entries.size();
if (tail == head) {
// expand the buffer
std::vector<gpt_log_entry> new_entries(2*entries.size());
std::vector<common_log_entry> new_entries(2*entries.size());
size_t new_tail = 0;
......@@ -346,15 +346,15 @@ public:
pause();
if (colors) {
g_col[GPT_LOG_COL_DEFAULT] = LOG_COL_DEFAULT;
g_col[GPT_LOG_COL_BOLD] = LOG_COL_BOLD;
g_col[GPT_LOG_COL_RED] = LOG_COL_RED;
g_col[GPT_LOG_COL_GREEN] = LOG_COL_GREEN;
g_col[GPT_LOG_COL_YELLOW] = LOG_COL_YELLOW;
g_col[GPT_LOG_COL_BLUE] = LOG_COL_BLUE;
g_col[GPT_LOG_COL_MAGENTA] = LOG_COL_MAGENTA;
g_col[GPT_LOG_COL_CYAN] = LOG_COL_CYAN;
g_col[GPT_LOG_COL_WHITE] = LOG_COL_WHITE;
g_col[COMMON_LOG_COL_DEFAULT] = LOG_COL_DEFAULT;
g_col[COMMON_LOG_COL_BOLD] = LOG_COL_BOLD;
g_col[COMMON_LOG_COL_RED] = LOG_COL_RED;
g_col[COMMON_LOG_COL_GREEN] = LOG_COL_GREEN;
g_col[COMMON_LOG_COL_YELLOW] = LOG_COL_YELLOW;
g_col[COMMON_LOG_COL_BLUE] = LOG_COL_BLUE;
g_col[COMMON_LOG_COL_MAGENTA] = LOG_COL_MAGENTA;
g_col[COMMON_LOG_COL_CYAN] = LOG_COL_CYAN;
g_col[COMMON_LOG_COL_WHITE] = LOG_COL_WHITE;
} else {
for (size_t i = 0; i < g_col.size(); i++) {
g_col[i] = "";
......@@ -381,47 +381,47 @@ public:
// public API
//
struct gpt_log * gpt_log_init() {
return new gpt_log;
struct common_log * common_log_init() {
return new common_log;
}
struct gpt_log * gpt_log_main() {
static struct gpt_log log;
struct common_log * common_log_main() {
static struct common_log log;
return &log;
}
void gpt_log_pause(struct gpt_log * log) {
void common_log_pause(struct common_log * log) {
log->pause();
}
void gpt_log_resume(struct gpt_log * log) {
void common_log_resume(struct common_log * log) {
log->resume();
}
void gpt_log_free(struct gpt_log * log) {
void common_log_free(struct common_log * log) {
delete log;
}
void gpt_log_add(struct gpt_log * log, enum ggml_log_level level, const char * fmt, ...) {
void common_log_add(struct common_log * log, enum ggml_log_level level, const char * fmt, ...) {
va_list args;
va_start(args, fmt);
log->add(level, fmt, args);
va_end(args);
}
void gpt_log_set_file(struct gpt_log * log, const char * file) {
void common_log_set_file(struct common_log * log, const char * file) {
log->set_file(file);
}
void gpt_log_set_colors(struct gpt_log * log, bool colors) {
void common_log_set_colors(struct common_log * log, bool colors) {
log->set_colors(colors);
}
void gpt_log_set_prefix(struct gpt_log * log, bool prefix) {
void common_log_set_prefix(struct common_log * log, bool prefix) {
log->set_prefix(prefix);
}
void gpt_log_set_timestamps(struct gpt_log * log, bool timestamps) {
void common_log_set_timestamps(struct common_log * log, bool timestamps) {
log->set_timestamps(timestamps);
}
llama/log.h
View file @ 527cc978
/**
* llama.cpp - commit 3f1ae2e32cde00c39b96be6d01c2997c29bae555 - do not edit this file
* llama.cpp - commit 40c6d79fb52f995f47507fedfeaae2ac05d9b35c - do not edit this file
*
* MIT License
*
......@@ -40,23 +40,23 @@
#define LOG_DEFAULT_LLAMA 0
// needed by the LOG_TMPL macro to avoid computing log arguments if the verbosity lower
// set via gpt_log_set_verbosity()
extern int gpt_log_verbosity_thold;
// set via common_log_set_verbosity()
extern int common_log_verbosity_thold;
void gpt_log_set_verbosity_thold(int verbosity); // not thread-safe
void common_log_set_verbosity_thold(int verbosity); // not thread-safe
// the gpt_log uses an internal worker thread to print/write log messages
// the common_log uses an internal worker thread to print/write log messages
// when the worker thread is paused, incoming log messages are discarded
struct gpt_log;
struct common_log;
struct gpt_log * gpt_log_init();
struct gpt_log * gpt_log_main(); // singleton, automatically destroys itself on exit
void gpt_log_pause (struct gpt_log * log); // pause the worker thread, not thread-safe
void gpt_log_resume(struct gpt_log * log); // resume the worker thread, not thread-safe
void gpt_log_free (struct gpt_log * log);
struct common_log * common_log_init();
struct common_log * common_log_main(); // singleton, automatically destroys itself on exit
void common_log_pause (struct common_log * log); // pause the worker thread, not thread-safe
void common_log_resume(struct common_log * log); // resume the worker thread, not thread-safe
void common_log_free (struct common_log * log);
LOG_ATTRIBUTE_FORMAT(3, 4)
void gpt_log_add(struct gpt_log * log, enum ggml_log_level level, const char * fmt, ...);
void common_log_add(struct common_log * log, enum ggml_log_level level, const char * fmt, ...);
// defaults: file = NULL, colors = false, prefix = false, timestamps = false
//
......@@ -80,10 +80,10 @@ void gpt_log_add(struct gpt_log * log, enum ggml_log_level level, const char * f
// D - debug (stderr, V = LOG_DEFAULT_DEBUG)
//
void gpt_log_set_file (struct gpt_log * log, const char * file); // not thread-safe
void gpt_log_set_colors (struct gpt_log * log, bool colors); // not thread-safe
void gpt_log_set_prefix (struct gpt_log * log, bool prefix); // whether to output prefix to each log
void gpt_log_set_timestamps(struct gpt_log * log, bool timestamps); // whether to output timestamps in the prefix
void common_log_set_file (struct common_log * log, const char * file); // not thread-safe
void common_log_set_colors (struct common_log * log, bool colors); // not thread-safe
void common_log_set_prefix (struct common_log * log, bool prefix); // whether to output prefix to each log
void common_log_set_timestamps(struct common_log * log, bool timestamps); // whether to output timestamps in the prefix
// helper macros for logging
// use these to avoid computing log arguments if the verbosity of the log is higher than the threshold
......@@ -92,13 +92,13 @@ void gpt_log_set_timestamps(struct gpt_log * log, bool timestamps); // w
//
// LOG_DBG("this is a debug message: %d\n", expensive_function());
//
// this will avoid calling expensive_function() if LOG_DEFAULT_DEBUG > gpt_log_verbosity_thold
// this will avoid calling expensive_function() if LOG_DEFAULT_DEBUG > common_log_verbosity_thold
//
#define LOG_TMPL(level, verbosity, ...) \
do { \
if ((verbosity) <= gpt_log_verbosity_thold) { \
gpt_log_add(gpt_log_main(), (level), __VA_ARGS__); \
if ((verbosity) <= common_log_verbosity_thold) { \
common_log_add(common_log_main(), (level), __VA_ARGS__); \
} \
} while (0)
......
llama/mllama.cpp
View file @ 527cc978
......@@ -3,6 +3,7 @@
#include "ggml-alloc.h"
#include "ggml-backend.h"
#include "ggml-cpu.h"
#include "ggml.h"
#ifdef GGML_USE_CUDA
......
llama/mmq.cpp 0 → 100644
View file @ 527cc978
/**
* llama.cpp - commit 40c6d79fb52f995f47507fedfeaae2ac05d9b35c - do not edit this file
*
* MIT License
*
* Copyright (c) 2023-2024 The ggml authors
*
* Permission is hereby granted, free of charge, to any person obtaining a copy
* of this software and associated documentation files (the "Software"), to deal
* in the Software without restriction, including without limitation the rights
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
* copies of the Software, and to permit persons to whom the Software is
* furnished to do so, subject to the following conditions:
*
* The above copyright notice and this permission notice shall be included in all
* copies or substantial portions of the Software.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
* SOFTWARE.
*/
#if defined(__GNUC__)
#pragma GCC diagnostic ignored "-Wpedantic"
#pragma GCC diagnostic ignored "-Wunused-local-typedefs"
#endif
#include "amx.h"
#include "mmq.h"
#include "ggml-impl.h"
#include "ggml-cpu-impl.h"
#include "ggml-cpu-quants.h"
#include "ggml-quants.h"
#include <algorithm>
#include <type_traits>
#if defined(__gnu_linux__)
#include <sys/syscall.h>
#include <unistd.h>
#endif
#if defined(_OPENMP)
#include <omp.h>
#endif
#if (defined(_WIN32) || defined(_WIN64))
#define RESTRICT __restrict
#else
#define RESTRICT __restrict__
#endif
#if (defined(_WIN32) || defined(_WIN64))
#define ALWAYS_INLINE __forceinline
#elif __has_attribute(always_inline) || defined(__GNUC__)
#define ALWAYS_INLINE __attribute__((__always_inline__)) inline
#else
#define ALWAYS_INLINE inline
#endif
#if defined(__AMX_INT8__) && defined(__AVX512VNNI__)
namespace {
// Forced unrolling
template <int n>
struct Unroll {
template <typename Func, typename... Args>
ALWAYS_INLINE void operator()(const Func& f, Args... args) const {
Unroll<n - 1>{}(f, args...);
f(std::integral_constant<int, n - 1>{}, args...);
}
};
template <>
struct Unroll<1> {
template <typename Func, typename... Args>
ALWAYS_INLINE void operator()(const Func& f, Args... args) const {
f(std::integral_constant<int, 0>{}, args...);
}
};
// type traits
template <typename T> struct PackedTypes {};
template <> struct PackedTypes<block_q4_0> { using type = int8_t; };
template <> struct PackedTypes<block_q4_1> { using type = uint8_t; };
template <> struct PackedTypes<block_q8_0> { using type = int8_t; };
template <typename T> using packed_B_type = typename PackedTypes<T>::type;
template <typename T>
struct do_compensate : std::integral_constant<bool,
std::is_same<T, block_q8_0>::value> {};
template <typename T>
struct do_unpack : std::integral_constant<bool,
std::is_same<T, block_q4_0>::value ||
std::is_same<T, block_q4_1>::value> {};
template <typename T>
struct is_type_qkk : std::integral_constant<bool,
std::is_same<T, block_q4_K>::value ||
std::is_same<T, block_q5_K>::value ||
std::is_same<T, block_q6_K>::value ||
std::is_same<T, block_iq4_xs>::value> {};
#define GGML_DISPATCH_FLOATING_TYPES(TYPE, ...) \
[&] { \
switch (TYPE) { \
case GGML_TYPE_F16: { \
using type = ggml_fp16_t; \
constexpr int blck_size = 16; \
return __VA_ARGS__(); \
} \
case GGML_TYPE_BF16: { \
using type = ggml_bf16_t; \
constexpr int blck_size = 32; \
return __VA_ARGS__(); \
} \
default: \
fprintf(stderr, "Unsupported floating data type\n"); \
} \
}()
#define GGML_DISPATCH_QTYPES(QT, ...) \
[&] { \
switch (QT) { \
case GGML_TYPE_Q4_0: { \
using type = block_q4_0; \
using vec_dot_type = block_q8_0; \
constexpr int blck_size = QK4_0; \
return __VA_ARGS__(); \
} \
case GGML_TYPE_Q4_1: { \
using type = block_q4_1; \
using vec_dot_type = block_q8_1; \
constexpr int blck_size = QK4_1; \
return __VA_ARGS__(); \
} \
case GGML_TYPE_Q8_0: { \
using type = block_q8_0; \
using vec_dot_type = block_q8_0; \
constexpr int blck_size = QK8_0; \
return __VA_ARGS__(); \
} \
case GGML_TYPE_Q4_K: { \
using type = block_q4_K; \
using vec_dot_type = block_q8_K; \
constexpr int blck_size = QK_K; \
return __VA_ARGS__(); \
} \
case GGML_TYPE_Q5_K: { \
using type = block_q5_K; \
using vec_dot_type = block_q8_K; \
constexpr int blck_size = QK_K; \
return __VA_ARGS__(); \
} \
case GGML_TYPE_Q6_K: { \
using type = block_q6_K; \
using vec_dot_type = block_q8_K; \
constexpr int blck_size = QK_K; \
return __VA_ARGS__(); \
} \
case GGML_TYPE_IQ4_XS: { \
using type = block_iq4_xs; \
using vec_dot_type = block_q8_K; \
constexpr int blck_size = QK_K; \
return __VA_ARGS__(); \
} \
default: \
fprintf(stderr, "Unsupported quantized data type: %d\n", int(TYPE)); \
} \
}()
#define GGML_DISPATCH_BOOL(BOOL_V, BOOL_NAME, ...) \
[&] { \
if (BOOL_V) { \
constexpr bool BOOL_NAME = true; \
return __VA_ARGS__(); \
} else { \
constexpr bool BOOL_NAME = false; \
return __VA_ARGS__(); \
} \
}()
// define amx tile config data structure
struct tile_config_t{
uint8_t palette_id = 0;
uint8_t start_row = 0;
uint8_t reserved_0[14] = {0};
uint16_t colsb[16] = {0};
uint8_t rows[16] = {0};
};
// Notes: amx tile config
//
// Typically, TMUL calculates A and B of size 16 x 64 containing INT8 values,
// and accumulate the result to a 16 x 16 matrix C containing INT32 values,
//
// As many GGUF quantized types as `block_size` of 32, so a 16-16-32 config is used
// instead of the normally used 16-16-64 config.
//
// Block A: {16, 32}, dtype = int8_t
// Block B: {16, 32}, dtype = uint8_t/int8_t
// Block C: {16, 16}, dtype = int32_t
//
// Block B needs to be prepacked to vnni format before feeding into TMUL:
// packed_B: from {n, k} to {k/vnni_blk, n, vnni_blck}, viewed in 2d, we get {8, 64}
//
// Therefore, we get tileconfig:
// A B C
// rows 16 8 16
// colsb 32 64 16
//
// For tile distribution, follow a 2-2-4 pattern, e.g. A used TMM2-TMM3, B used TMM0-TMM1,
// C used TMM4-TMM7:
// B TMM0 B TMM1
// A TMM2 C TMM4 C TMM6
// A TMM3 C TMM5 C TMM7
//
// Each `amx` kernel handles 4 blocks at a time: 2MB * 2NB, when m < 2 * BLOCK_M, unpack A
// will be needed.
//
// Here another commonly used pattern 1-3-3 is skipped, as it is mostly used when m <=16;
// and the sinlge batch gemm (m=1) has a special fast path with `avx512-vnni`.
//
// ref: https://www.intel.com/content/www/us/en/developer/articles/code-sample/
// advanced-matrix-extensions-intrinsics-functions.html
//
#define TC_CONFIG_TILE(i, r, cb) tc.rows[i] = r; tc.colsb[i] = cb
void ggml_tile_config_init(void) {
static thread_local bool is_first_time = true;
if (!is_first_time) {
return;
}
static thread_local tile_config_t tc;
tile_config_t current_tc;
_tile_storeconfig(&current_tc);
// load only when config changes
if (tc.palette_id == 0 || (memcmp(&current_tc.colsb, &tc.colsb, sizeof(uint16_t) * 8) != 0 &&
memcmp(&current_tc.rows, &tc.rows, sizeof(uint8_t) * 8) != 0)) {
tc.palette_id = 1;
tc.start_row = 0;
TC_CONFIG_TILE(TMM0, 8, 64);
TC_CONFIG_TILE(TMM1, 8, 64);
TC_CONFIG_TILE(TMM2, 16, 32);
TC_CONFIG_TILE(TMM3, 16, 32);
TC_CONFIG_TILE(TMM4, 16, 64);
TC_CONFIG_TILE(TMM5, 16, 64);
TC_CONFIG_TILE(TMM6, 16, 64);
TC_CONFIG_TILE(TMM7, 16, 64);
_tile_loadconfig(&tc);
}
is_first_time = false;
}
// we need an extra 16 * 4B (TILE_N * int32_t) for each NB/KB block for compensation.
// See the notes `s8s8 igemm compensation in avx512-vnni` for detail.
template <typename TB>
int get_tile_size() {
int tile_size = TILE_N * sizeof(TB);
if (do_compensate<TB>::value) {
tile_size += TILE_N * sizeof(int32_t);
}
if (std::is_same<TB, block_q4_K>::value ||
std::is_same<TB, block_q5_K>::value) {
tile_size += TILE_N * 4;
}
if (std::is_same<TB, block_iq4_xs>::value) {
tile_size += TILE_N * 2;
}
return tile_size;
}
template <typename TB, int BLOCK_K>
int get_row_size(int K) {
int KB = K / BLOCK_K;
int row_size = KB * sizeof(TB);
if (do_compensate<TB>::value) {
row_size += KB * sizeof(int32_t);
}
if (std::is_same<TB, block_q4_K>::value ||
std::is_same<TB, block_q5_K>::value) {
row_size += KB * 4;
}
if (std::is_same<TB, block_iq4_xs>::value) {
row_size += KB * 2;
}
return row_size;
}
// vectorized dtype conversion
inline float FP16_TO_FP32(ggml_half val) {
__m256i v = _mm256_setr_epi16(
val, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0);
__m512 o = _mm512_cvtph_ps(v);
return _mm512_cvtss_f32(o);
}
inline __m512 FP16_TO_FP32_VEC(ggml_half val) {
__m256i v = _mm256_set1_epi16(val);
return _mm512_cvtph_ps(v);
}
// horizontal reduce
inline float _mm512_reduce_max_ps(const __m512 x) {
__m512 v = x;
__m512 v1 = _mm512_shuffle_f32x4(v, v, 0x4E);
v = _mm512_max_ps(v, v1);
v1 = _mm512_shuffle_f32x4(v, v, 0xB1);
v = _mm512_max_ps(v, v1);
v1 = _mm512_shuffle_ps(v, v, 0x4E);
v = _mm512_max_ps(v, v1);
v1 = _mm512_shuffle_ps(v, v, 0xB1);
v = _mm512_max_ps(v, v1);
return _mm512_cvtss_f32(v);
}
// transpose utils
#define SHUFFLE_EPI32(a, b, mask) \
_mm256_castps_si256(_mm256_shuffle_ps(_mm256_castsi256_ps(a), _mm256_castsi256_ps(b), mask))
inline void transpose_8x8_32bit(__m256i * v, __m256i * v1) {
// unpacking and 32-bit elements
v1[0] = _mm256_unpacklo_epi32(v[0], v[1]);
v1[1] = _mm256_unpackhi_epi32(v[0], v[1]);
v1[2] = _mm256_unpacklo_epi32(v[2], v[3]);
v1[3] = _mm256_unpackhi_epi32(v[2], v[3]);
v1[4] = _mm256_unpacklo_epi32(v[4], v[5]);
v1[5] = _mm256_unpackhi_epi32(v[4], v[5]);
v1[6] = _mm256_unpacklo_epi32(v[6], v[7]);
v1[7] = _mm256_unpackhi_epi32(v[6], v[7]);
// shuffling the 32-bit elements
v[0] = SHUFFLE_EPI32(v1[0], v1[2], 0x44);
v[1] = SHUFFLE_EPI32(v1[0], v1[2], 0xee);
v[2] = SHUFFLE_EPI32(v1[4], v1[6], 0x44);
v[3] = SHUFFLE_EPI32(v1[4], v1[6], 0xee);
v[4] = SHUFFLE_EPI32(v1[1], v1[3], 0x44);
v[5] = SHUFFLE_EPI32(v1[1], v1[3], 0xee);
v[6] = SHUFFLE_EPI32(v1[5], v1[7], 0x44);
v[7] = SHUFFLE_EPI32(v1[5], v1[7], 0xee);
// shuffling 128-bit elements
v1[0] = _mm256_permute2f128_si256(v[2], v[0], 0x02);
v1[1] = _mm256_permute2f128_si256(v[3], v[1], 0x02);
v1[2] = _mm256_permute2f128_si256(v[6], v[4], 0x02);
v1[3] = _mm256_permute2f128_si256(v[7], v[5], 0x02);
v1[4] = _mm256_permute2f128_si256(v[2], v[0], 0x13);
v1[5] = _mm256_permute2f128_si256(v[3], v[1], 0x13);
v1[6] = _mm256_permute2f128_si256(v[6], v[4], 0x13);
v1[7] = _mm256_permute2f128_si256(v[7], v[5], 0x13);
}
inline void transpose_16x4_32bit(__m512i * r, __m512i * d) {
static const __m512i index1 = _mm512_set_epi32(
0x0f, 0x0b, 0x07, 0x03,
0x0e, 0x0a, 0x06, 0x02,
0x0d, 0x09, 0x05, 0x01,
0x0c, 0x08, 0x04, 0x00);
d[0] = _mm512_permutexvar_epi32(index1, r[0]);
d[1] = _mm512_permutexvar_epi32(index1, r[1]);
d[2] = _mm512_permutexvar_epi32(index1, r[2]);
d[3] = _mm512_permutexvar_epi32(index1, r[3]);
r[0] = _mm512_shuffle_i32x4(d[0], d[1], 0x44);
r[1] = _mm512_shuffle_i32x4(d[0], d[1], 0xee);
r[2] = _mm512_shuffle_i32x4(d[2], d[3], 0x44);
r[3] = _mm512_shuffle_i32x4(d[2], d[3], 0xee);
d[0] = _mm512_shuffle_i32x4(r[0], r[2], 0x88);
d[1] = _mm512_shuffle_i32x4(r[0], r[2], 0xdd);
d[2] = _mm512_shuffle_i32x4(r[1], r[3], 0x88);
d[3] = _mm512_shuffle_i32x4(r[1], r[3], 0xdd);
}
inline void transpose_16x16_32bit(__m512i * v) {
__m512i v1[16];
v1[0] = _mm512_unpacklo_epi32(v[0], v[1]);
v1[1] = _mm512_unpackhi_epi32(v[0], v[1]);
v1[2] = _mm512_unpacklo_epi32(v[2], v[3]);
v1[3] = _mm512_unpackhi_epi32(v[2], v[3]);
v1[4] = _mm512_unpacklo_epi32(v[4], v[5]);
v1[5] = _mm512_unpackhi_epi32(v[4], v[5]);
v1[6] = _mm512_unpacklo_epi32(v[6], v[7]);
v1[7] = _mm512_unpackhi_epi32(v[6], v[7]);
v1[8] = _mm512_unpacklo_epi32(v[8], v[9]);
v1[9] = _mm512_unpackhi_epi32(v[8], v[9]);
v1[10] = _mm512_unpacklo_epi32(v[10], v[11]);
v1[11] = _mm512_unpackhi_epi32(v[10], v[11]);
v1[12] = _mm512_unpacklo_epi32(v[12], v[13]);
v1[13] = _mm512_unpackhi_epi32(v[12], v[13]);
v1[14] = _mm512_unpacklo_epi32(v[14], v[15]);
v1[15] = _mm512_unpackhi_epi32(v[14], v[15]);
v[0] = _mm512_unpacklo_epi64(v1[0], v1[2]);
v[1] = _mm512_unpackhi_epi64(v1[0], v1[2]);
v[2] = _mm512_unpacklo_epi64(v1[1], v1[3]);
v[3] = _mm512_unpackhi_epi64(v1[1], v1[3]);
v[4] = _mm512_unpacklo_epi64(v1[4], v1[6]);
v[5] = _mm512_unpackhi_epi64(v1[4], v1[6]);
v[6] = _mm512_unpacklo_epi64(v1[5], v1[7]);
v[7] = _mm512_unpackhi_epi64(v1[5], v1[7]);
v[8] = _mm512_unpacklo_epi64(v1[8], v1[10]);
v[9] = _mm512_unpackhi_epi64(v1[8], v1[10]);
v[10] = _mm512_unpacklo_epi64(v1[9], v1[11]);
v[11] = _mm512_unpackhi_epi64(v1[9], v1[11]);
v[12] = _mm512_unpacklo_epi64(v1[12], v1[14]);
v[13] = _mm512_unpackhi_epi64(v1[12], v1[14]);
v[14] = _mm512_unpacklo_epi64(v1[13], v1[15]);
v[15] = _mm512_unpackhi_epi64(v1[13], v1[15]);
v1[0] = _mm512_shuffle_i32x4(v[0], v[4], 0x88);
v1[1] = _mm512_shuffle_i32x4(v[1], v[5], 0x88);
v1[2] = _mm512_shuffle_i32x4(v[2], v[6], 0x88);
v1[3] = _mm512_shuffle_i32x4(v[3], v[7], 0x88);
v1[4] = _mm512_shuffle_i32x4(v[0], v[4], 0xdd);
v1[5] = _mm512_shuffle_i32x4(v[1], v[5], 0xdd);
v1[6] = _mm512_shuffle_i32x4(v[2], v[6], 0xdd);
v1[7] = _mm512_shuffle_i32x4(v[3], v[7], 0xdd);
v1[8] = _mm512_shuffle_i32x4(v[8], v[12], 0x88);
v1[9] = _mm512_shuffle_i32x4(v[9], v[13], 0x88);
v1[10] = _mm512_shuffle_i32x4(v[10], v[14], 0x88);
v1[11] = _mm512_shuffle_i32x4(v[11], v[15], 0x88);
v1[12] = _mm512_shuffle_i32x4(v[8], v[12], 0xdd);
v1[13] = _mm512_shuffle_i32x4(v[9], v[13], 0xdd);
v1[14] = _mm512_shuffle_i32x4(v[10], v[14], 0xdd);
v1[15] = _mm512_shuffle_i32x4(v[11], v[15], 0xdd);
v[0] = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88);
v[1] = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88);
v[2] = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88);
v[3] = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88);
v[4] = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88);
v[5] = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88);
v[6] = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88);
v[7] = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88);
v[8] = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd);
v[9] = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd);
v[10] = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd);
v[11] = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd);
v[12] = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd);
v[13] = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd);
v[14] = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd);
v[15] = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd);
}
void quantize_row_q8_K_vnni(const float * RESTRICT x, void * RESTRICT vy, int64_t k) {
assert(k % QK_K == 0);
const int KB = k / QK_K;
constexpr int kVecs = QK_K / 16;
block_q8_K * y = reinterpret_cast<block_q8_K *>(vy);
// hold 16 float vecs from x
__m512 v[kVecs];
// hold the quants vecs
__m512i vq[kVecs / 4];
// hold the packed quants vecs
__m512i vq_packed[kVecs / 4];
const __m512 signBit = _mm512_set1_ps(-0.f);
for (int i = 0; i < KB; ++i) {
// Compute max(abs(e)) for the block
__m512 vamax = _mm512_set1_ps(0.f);
for (int j = 0; j < kVecs; ++j) {
v[j] = _mm512_loadu_ps(x); x += 16;
vamax = _mm512_max_ps(vamax, _mm512_andnot_ps(signBit, v[j]));
}
const float amax = _mm512_reduce_max_ps(vamax);
// Quantize these floats
const float iscale = 127.f / amax;
y[i].d = GGML_FP32_TO_FP16(1 / iscale);
const float id = ( amax != 0.0f ) ? iscale : 0.f;
const __m512 vscale = _mm512_set1_ps(id);
// Apply multiplier and round to nearest integer
for (int j = 0; j < kVecs; ++j) {
v[j] = _mm512_mul_ps(v[j], vscale);
v[j] = _mm512_roundscale_ps(v[j], (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
}
// Pack to epi8 vecs
for (int j = 0; j < kVecs / 4; ++j) {
__m128i q8_0 = _mm512_cvtepi32_epi8(_mm512_cvtps_epi32(v[j * 4 + 0]));
__m128i q8_1 = _mm512_cvtepi32_epi8(_mm512_cvtps_epi32(v[j * 4 + 1]));
__m128i q8_2 = _mm512_cvtepi32_epi8(_mm512_cvtps_epi32(v[j * 4 + 2]));
__m128i q8_3 = _mm512_cvtepi32_epi8(_mm512_cvtps_epi32(v[j * 4 + 3]));
__m256i q8_01 = _mm256_insertf128_si256(_mm256_castsi128_si256(q8_0), (q8_1), 1);
__m256i q8_23 = _mm256_insertf128_si256(_mm256_castsi128_si256(q8_2), (q8_3), 1);
vq[j] = _mm512_inserti32x8(_mm512_castsi256_si512(q8_01), q8_23, 1);
_mm512_storeu_si512((__m512i *)(y[i].qs + j * 64), vq[j]);
}
// Compute the bsums with vnni
transpose_16x4_32bit(vq, vq_packed);
const __m512i one = _mm512_set1_epi8(1);
__m512i sum = _mm512_setzero_si512();
for (int k = 0; k < 4; ++k) {
sum = _mm512_dpbusd_epi32(sum, one, vq_packed[k]);
}
_mm256_storeu_si256((__m256i *)(y[i].bsums), _mm512_cvtepi32_epi16(sum));
}
}
// quantize A from float to `vec_dot_type`
template <typename T>
inline void from_float(const float * x, char * vy, int64_t k);
template <>
inline void from_float<block_q8_0>(const float * x, char * vy, int64_t k) {
quantize_row_q8_0(x, (block_q8_0 *)vy, k);
}
template <>
inline void from_float<block_q8_1>(const float * x, char * vy, int64_t k) {
quantize_row_q8_1(x, (block_q8_1 *)vy, k);
}
template <>
inline void from_float<block_q8_K>(const float * x, char * vy, int64_t k) {
#if 1
// TODO: this is reference impl!
quantize_row_q8_K_ref(x, (block_q8_K *)vy, k);
#else
quantize_row_q8_K_vnni(x, vy, k);
#endif
}
// load A from memory to array when nrows can not fill in whole tile
void unpack_A(int8_t * RESTRICT tile, const block_q8_0 * RESTRICT A, int lda, int nr) {
assert(nr != TILE_M);
for (int m = 0; m < nr; ++m) {
const __m256i v = _mm256_loadu_si256((const __m256i *)(A[m * lda].qs));
_mm256_storeu_si256((__m256i *)(tile + m * TILE_K), v);
}
}
void unpack_A(int8_t * RESTRICT tile, const block_q8_1 * RESTRICT A, int lda, int nr) {
assert(nr != TILE_M);
for (int m = 0; m < nr; ++m) {
const __m256i v = _mm256_loadu_si256((const __m256i *)(A[m * lda].qs));
_mm256_storeu_si256((__m256i *)(tile + m * TILE_K), v);
}
}
template <typename TB>
void unpack_A(int8_t * RESTRICT tile, const block_q8_K * RESTRICT A, int lda, int k, int nr) {
assert(nr <= TILE_M);
for (int m = 0; m < nr; ++m) {
const __m256i v = _mm256_loadu_si256((const __m256i *)(A[m * lda].qs + k * 32));
_mm256_storeu_si256((__m256i *)(tile + m * TILE_K), v);
}
}
template <>
void unpack_A<block_q6_K>(int8_t * RESTRICT tile, const block_q8_K * RESTRICT A, int lda, int k, int nr) {
assert(nr <= TILE_M);
// zero padding k from 16 to 32, so that we don't have to re-config amx
const __m128i zero = _mm_setzero_si128();
for (int m = 0; m < nr; ++m) {
const __m128i v = _mm_loadu_si128((const __m128i *)(A[m * lda].qs + k * 16));
const __m256i r = _mm256_insertf128_si256(_mm256_castsi128_si256(v), zero, 1);
_mm256_storeu_si256((__m256i *)(tile + m * TILE_K), r);
}
}
#define MM256_SET_M128I(a, b) _mm256_insertf128_si256(_mm256_castsi128_si256(b), (a), 1)
inline __m256i bytes_from_nibbles_32(const uint8_t * rsi) {
const __m128i tmp = _mm_loadu_si128((const __m128i *)rsi);
const __m256i bytes = MM256_SET_M128I(_mm_srli_epi16(tmp, 4), tmp);
const __m256i lowMask = _mm256_set1_epi8(0xF);
return _mm256_and_si256(lowMask, bytes);
}
// used for block_q4_K
inline __m512i bytes_from_nibbles_64(const uint8_t * rsi) {
const __m256i tmp = _mm256_loadu_si256((const __m256i *)rsi);
const __m256i lowMask = _mm256_set1_epi8(0xF);
const __m256i q4l = _mm256_and_si256(tmp, lowMask);
const __m256i q4h = _mm256_and_si256(_mm256_srli_epi16(tmp, 4), lowMask);
return _mm512_inserti32x8(_mm512_castsi256_si512(q4l), q4h, 1);
}
// used for block_q5_K
inline __m512i bytes_from_nibbles_64(const uint8_t * qs, const uint8_t * qh, int k) {
const __m256i lowMask = _mm256_set1_epi8(0xF);
__m256i hmask = _mm256_set1_epi8(1);
hmask = _mm256_slli_epi16(hmask, k);
const __m256i q5bits = _mm256_loadu_si256((const __m256i *)qs);
const __m256i hbits = _mm256_loadu_si256((const __m256i *)qh);
const __m256i q5l_0 = _mm256_and_si256(q5bits, lowMask);
const __m256i q5h_0 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_and_si256(hbits, hmask), k + 0), 4);
const __m256i q5_0 = _mm256_add_epi8(q5l_0, q5h_0);
hmask = _mm256_slli_epi16(hmask, 1);
const __m256i q5l_1 = _mm256_and_si256(_mm256_srli_epi16(q5bits, 4), lowMask);
const __m256i q5h_1 = _mm256_slli_epi16(_mm256_srli_epi16(_mm256_and_si256(hbits, hmask), k + 1), 4);
const __m256i q5_1 = _mm256_add_epi8(q5l_1, q5h_1);
return _mm512_inserti32x8(_mm512_castsi256_si512(q5_0), q5_1, 1);
}
// used for block_q6_K
inline void bytes_from_nibbles_128(__m512i& r0, __m512i& r1, const uint8_t * qs, const uint8_t * qh) {
const __m256i m4 = _mm256_set1_epi8(0xF);
const __m256i m2 = _mm256_set1_epi8(0x3);
const __m256i q6bits1 = _mm256_loadu_si256((const __m256i *)qs);
const __m256i q6bits2 = _mm256_loadu_si256((const __m256i *)(qs + 32));
const __m256i q6bitsH = _mm256_loadu_si256((const __m256i *)qh);
const __m256i q6h_0 = _mm256_slli_epi16(_mm256_and_si256( q6bitsH, m2), 4);
const __m256i q6h_1 = _mm256_slli_epi16(_mm256_and_si256(_mm256_srli_epi16(q6bitsH, 2), m2), 4);
const __m256i q6h_2 = _mm256_slli_epi16(_mm256_and_si256(_mm256_srli_epi16(q6bitsH, 4), m2), 4);
const __m256i q6h_3 = _mm256_slli_epi16(_mm256_and_si256(_mm256_srli_epi16(q6bitsH, 6), m2), 4);
const __m256i q6_0 = _mm256_or_si256(_mm256_and_si256(q6bits1, m4), q6h_0);
const __m256i q6_1 = _mm256_or_si256(_mm256_and_si256(q6bits2, m4), q6h_1);
const __m256i q6_2 = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(q6bits1, 4), m4), q6h_2);
const __m256i q6_3 = _mm256_or_si256(_mm256_and_si256(_mm256_srli_epi16(q6bits2, 4), m4), q6h_3);
r0 = _mm512_inserti32x8(_mm512_castsi256_si512(q6_0), q6_1, 1);
r1 = _mm512_inserti32x8(_mm512_castsi256_si512(q6_2), q6_3, 1);
}
inline __m512i packNibbles(__m512i r0, __m512i r1) {
return _mm512_or_si512(r0, _mm512_slli_epi16(r1, 4));
}
template <typename TB>
inline void pack_qs(void * RESTRICT packed_B, const TB * RESTRICT B, int KB) {
int8_t tmp[8 * 64];
__m256i v[8], v2[8];
for (int n = 0; n < 8; ++n) {
v[n] = bytes_from_nibbles_32(B[n * KB].qs);
}
transpose_8x8_32bit(v, v2);
for (int n = 0; n < 8; ++n) {
_mm256_storeu_si256((__m256i *)(tmp + n * 64), v2[n]);
}
for (int n = 0; n < 8; ++n) {
v[n] = bytes_from_nibbles_32(B[(n + 8) * KB].qs);
}
transpose_8x8_32bit(v, v2);
for (int n = 0; n < 8; ++n) {
_mm256_storeu_si256((__m256i *)(tmp + n * 64 + 32), v2[n]);
}
// pack again with 128 to fully utilize vector length
for (int n = 0; n < 8; n += 2) {
__m512i r0 = _mm512_loadu_si512((const __m512i *)(tmp + n * 64));
__m512i r1 = _mm512_loadu_si512((const __m512i *)(tmp + n * 64 + 64));
__m512i r1r0 = packNibbles(r0, r1);
_mm512_storeu_si512((__m512i *)((char *)packed_B + n * 32), r1r0);
}
}
template <>
inline void pack_qs<block_q8_0>(void * RESTRICT packed_B, const block_q8_0 * RESTRICT B, int KB) {
__m256i v[8], v2[8];
for (int n = 0; n < 8; ++n) {
v[n] = _mm256_loadu_si256((const __m256i *)(B[n * KB].qs));
}
transpose_8x8_32bit(v, v2);
for (int n = 0; n < 8; ++n) {
_mm256_storeu_si256((__m256i *)((char *)packed_B + n * 64), v2[n]);
}
for (int n = 0; n < 8; ++n) {
v[n] = _mm256_loadu_si256((const __m256i *)(B[(n + 8) * KB].qs));
}
transpose_8x8_32bit(v, v2);
for (int n = 0; n < 8; ++n) {
_mm256_storeu_si256((__m256i *)((char *)packed_B + n * 64 + 32), v2[n]);
}
}
template <>
inline void pack_qs<block_q4_K>(void * RESTRICT packed_B, const block_q4_K * RESTRICT B, int KB) {
__m512i v[16];
// QK_K 256 with 8 groups, handle 2 groups at a time
char * pb = (char *)packed_B;
for (int k = 0; k < QK_K / 64; ++k) {
// pack 2 groups { n, g, k} to {g, k/4, 4n}
// e.g. {16, 2, 32} to {2, 8, 64}
for (int n = 0; n < TILE_N; ++n) {
v[n] = bytes_from_nibbles_64(B[n * KB].qs + k * 32);
}
transpose_16x16_32bit(v);
// pack again with 128 to fully utilize vector length
for (int n = 0; n < TILE_N; n += 2) {
_mm512_storeu_si512((__m512i *)pb, packNibbles(v[n], v[n + 1]));
pb += 64;
}
}
}
template <>
inline void pack_qs<block_q5_K>(void * RESTRICT packed_B, const block_q5_K * RESTRICT B, int KB) {
__m512i v[16];
const __m512i lowMask = _mm512_set1_epi8(0xF);
// QK_K 256 with 8 groups, handle 2 groups at a time
char * pb = (char *)packed_B;
char * ph = (char *)packed_B + (QK_K / 2) * TILE_N;
for (int k = 0; k < QK_K / 64; ++k) {
// pack 2 groups { n, g, k} to {g, k/4, 4n}
// e.g. {16, 2, 32} to {2, 8, 64}
for (int n = 0; n < TILE_N; ++n) {
v[n] = bytes_from_nibbles_64(B[n * KB].qs + k * 32, B[n * KB].qh, /* group */2 * k);
}
transpose_16x16_32bit(v);
// 1. pack lower 4bits with 2 groups
for (int n = 0; n < TILE_N; n += 2) {
// get lower 4 bits
const __m512i r0 = _mm512_and_si512(v[n], lowMask);
const __m512i r1 = _mm512_and_si512(v[n + 1], lowMask);
_mm512_storeu_si512((__m512i *)pb, packNibbles(r0, r1)); pb += 64;
}
// 2. pack higher 1bit with 2 groups
const __m512i hmask = _mm512_set1_epi8(0x10);
for (int g = 0; g < 2; ++g) {
__m512i hbits = _mm512_setzero_si512();
hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 8 + 0], hmask), 4));
hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 8 + 1], hmask), 3));
hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 8 + 2], hmask), 2));
hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 8 + 3], hmask), 1));
hbits = _mm512_add_epi8(hbits, _mm512_and_si512(v[g * 8 + 4], hmask) );
hbits = _mm512_add_epi8(hbits, _mm512_slli_epi16(_mm512_and_si512(v[g * 8 + 5], hmask), 1));
hbits = _mm512_add_epi8(hbits, _mm512_slli_epi16(_mm512_and_si512(v[g * 8 + 6], hmask), 2));
hbits = _mm512_add_epi8(hbits, _mm512_slli_epi16(_mm512_and_si512(v[g * 8 + 7], hmask), 3));
_mm512_storeu_si512((__m512i *)ph, hbits); ph += 64;
}
}
}
template <>
inline void pack_qs<block_q6_K>(void * RESTRICT packed_B, const block_q6_K * RESTRICT B, int KB) {
__m512i v[32];
const __m512i lowMask = _mm512_set1_epi8(0xF);
// QK_K 256 with 8 groups, handle 4 groups at a time
char * pb = (char *)packed_B;
char * ph = (char *)packed_B + (QK_K / 2) * TILE_N;
for (int k = 0; k < QK_K / 128; ++k) {
for (int n = 0; n < TILE_N; ++n) {
bytes_from_nibbles_128(v[n], v[n + 16], B[n * KB].ql + k * 64, B[n * KB].qh + k * 32);
}
// top half: group 0,1 or 4,5; bottom half: group 2,3 or 6,7
transpose_16x16_32bit(v);
transpose_16x16_32bit(v + 16);
// 1. pack lower 4bits with 4 groups
for (int n = 0; n < 32; n += 2) {
const __m512i r0 = _mm512_and_si512(v[n], lowMask);
const __m512i r1 = _mm512_and_si512(v[n + 1], lowMask);
_mm512_storeu_si512((__m512i *)pb, packNibbles(r0, r1)); pb += 64;
}
// 2. pack higher 2bit with 4 groups
const __m512i hmask = _mm512_set1_epi8(0x30);
for (int g = 0; g < 8; ++g) {
__m512i hbits = _mm512_setzero_si512();
hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 4 + 0], hmask), 4));
hbits = _mm512_add_epi8(hbits, _mm512_srli_epi16(_mm512_and_si512(v[g * 4 + 1], hmask), 2));
hbits = _mm512_add_epi8(hbits, _mm512_and_si512(v[g * 4 + 2], hmask) );
hbits = _mm512_add_epi8(hbits, _mm512_slli_epi16(_mm512_and_si512(v[g * 4 + 3], hmask), 2));
_mm512_storeu_si512((__m512i *)ph, hbits); ph += 64;
}
}
}
template <>
inline void pack_qs<block_iq4_xs>(void * RESTRICT packed_B, const block_iq4_xs * RESTRICT B, int KB) {
__m512i v[16];
char * pb = (char *)packed_B;
for (int k = 0; k < QK_K / 64; ++k) {
for (int n = 0; n < TILE_N; ++n) {
__m256i r0 = bytes_from_nibbles_32(B[n * KB].qs + k * 32 + 0);
__m256i r1 = bytes_from_nibbles_32(B[n * KB].qs + k * 32 + 16);
v[n] = _mm512_inserti32x8(_mm512_castsi256_si512(r0), r1, 1);
}
transpose_16x16_32bit(v);
// pack again with 128 to fully utilize vector length
for (int n = 0; n < TILE_N; n += 2) {
_mm512_storeu_si512((__m512i *)pb, packNibbles(v[n], v[n + 1]));
pb += 64;
}
}
}
// pack B to vnni formats in 4bits or 8 bits
void pack_B(void * RESTRICT packed_B, const block_q4_0 * RESTRICT B, int KB) {
pack_qs(packed_B, B, KB);
ggml_half * d0 = reinterpret_cast<ggml_half *>((char *)packed_B + TILE_N * TILE_K / 2);
for (int n = 0; n < TILE_N; ++n) {
d0[n] = B[n * KB].d;
}
}
void pack_B(void * RESTRICT packed_B, const block_q4_1 * RESTRICT B, int KB) {
pack_qs(packed_B, B, KB);
ggml_half * d0 = reinterpret_cast<ggml_half *>((char *)packed_B + TILE_N * TILE_K / 2);
ggml_half * m0 = d0 + TILE_N;
for (int n = 0; n < TILE_N; ++n) {
d0[n] = B[n * KB].d;
m0[n] = B[n * KB].m;
}
}
inline void s8s8_compensation(void * RESTRICT packed_B) {
// packed_B layout:
// quants {TILE_N, TILEK} int8_t
// d0 {TILE_N} ggml_half
// comp {TILE_N} int32_t
const int offset = TILE_N * TILE_K + TILE_N * sizeof(ggml_half);
__m512i vcomp = _mm512_setzero_si512();
const __m512i off = _mm512_set1_epi8(static_cast<char>(0x80));
for (int k = 0; k < 8; ++k) {
__m512i vb = _mm512_loadu_si512((const __m512i *)((const char *)packed_B + k * 64));
vcomp = _mm512_dpbusd_epi32(vcomp, off, vb);
}
_mm512_storeu_si512((__m512i *)((char *)(packed_B) + offset), vcomp);
}
void pack_B(void * RESTRICT packed_B, const block_q8_0 * RESTRICT B, int KB) {
pack_qs(packed_B, B, KB);
ggml_half * d0 = reinterpret_cast<ggml_half *>((char *)packed_B + TILE_N * TILE_K);
for (int n = 0; n < TILE_N; ++n) {
d0[n] = B[n * KB].d;
}
s8s8_compensation(packed_B);
}
// convert 8 * {min, scale} from int6 to int8
inline void unpack_mins_and_scales(const uint8_t * scales, uint32_t * utmp) {
const uint32_t kmask1 = 0x3f3f3f3f;
const uint32_t kmask2 = 0x0f0f0f0f;
const uint32_t kmask3 = 0x03030303;
memcpy(utmp, scales, 12);
utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4);
const uint32_t uaux = utmp[1] & kmask1;
utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4);
utmp[2] = uaux;
utmp[0] &= kmask1;
}
// packed_B layout:
// quants {8, TILE_N, 16} uint8
// scales {8, TILE_N} uint8
// mins {8, TILE_N} uint8
// d {TILE_N} ggml_half
// dmin {TILE_N} ggml_half
void pack_B(void * RESTRICT packed_B, const block_q4_K * RESTRICT B, int KB) {
pack_qs(packed_B, B, KB);
uint8_t * scales = reinterpret_cast<uint8_t *>((char *)packed_B + (QK_K / 2) * TILE_N);
uint8_t * mins = scales + 8 * TILE_N;
ggml_half * d = reinterpret_cast<ggml_half *>(mins + 8 * TILE_N);
ggml_half * dmin = d + TILE_N;
union {
uint32_t u32[4];
uint8_t u8[16];
} s;
for (int n = 0; n < TILE_N; ++n) {
unpack_mins_and_scales(B[n * KB].scales, s.u32);
for (int k = 0; k < 8; ++k) {
scales[k * TILE_N + n] = s.u8[k];
mins[(k >> 1) * TILE_N * 2 + n * 2 + (k & 0x1)] = s.u8[k + 8];
}
d[n] = B[n * KB].d;
dmin[n] = B[n * KB].dmin;
}
}
// packed_B layout:
// quants {8, TILE_N, 16} uint8
// qh {8, TILE_N, 4} uint8
// scales {8, TILE_N} uint8
// mins {8, TILE_N} uint8
// d {TILE_N} ggml_half
// dmin {TILE_N} ggml_half
void pack_B(void * RESTRICT packed_B, const block_q5_K * RESTRICT B, int KB) {
pack_qs(packed_B, B, KB);
uint8_t * scales = reinterpret_cast<uint8_t *>((char *)packed_B + (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N);
uint8_t * mins = scales + 8 * TILE_N;
ggml_half * d = reinterpret_cast<ggml_half *>(mins + 8 * TILE_N);
ggml_half * dmin = d + TILE_N;
union {
uint32_t u32[4];
uint8_t u8[16];
} s;
for (int n = 0; n < TILE_N; ++n) {
unpack_mins_and_scales(B[n * KB].scales, s.u32);
for (int k = 0; k < 8; ++k) {
scales[k * TILE_N + n] = s.u8[k];
mins[(k >> 1) * TILE_N * 2 + n * 2 + (k & 0x1)] = s.u8[k + 8];
}
d[n] = B[n * KB].d;
dmin[n] = B[n * KB].dmin;
}
}
// packed_B layout:
// quants {16, TILE_N, 8} uint8
// qh {16, TILE_N, 4} uint8
// scales {16, TILE_N} uint8
// d {TILE_N} ggml_half
void pack_B(void * RESTRICT packed_B, const block_q6_K * RESTRICT B, int KB) {
pack_qs(packed_B, B, KB);
uint8_t * scales = reinterpret_cast<uint8_t *>((char *)packed_B + (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N);
ggml_half * d = reinterpret_cast<ggml_half *>(scales + 16 * TILE_N);
for (int n = 0; n < TILE_N; ++n) {
const int8_t * ps = B[n * KB].scales;
for (int k = 0; k < 16; ++k) {
scales[k * TILE_N + n] = ps[k];
}
d[n] = B[n * KB].d;
}
}
// packed_B layout:
// quants {8, TILE_N, 16} uint8
// scales {8, TILE_N} int8
// d {TILE_N} ggml_half
void pack_B(void * RESTRICT packed_B, const block_iq4_xs * RESTRICT B, int KB) {
pack_qs(packed_B, B, KB);
int8_t * scales = reinterpret_cast<int8_t *>((char *)packed_B + (QK_K / 2) * TILE_N);
ggml_half * d = reinterpret_cast<ggml_half *>(scales + 8 * TILE_N);
// pack the scales
for (int n = 0; n < TILE_N; ++n) {
uint16_t sh = B[n * KB].scales_h;
for (int k = 0; k < 8; k += 2) {
const int16_t ls1 = ((B[n * KB].scales_l[k / 2] & 0xf) | ((sh << 4) & 0x30)) - 32;
const int16_t ls2 = ((B[n * KB].scales_l[k / 2] >> 4) | ((sh << 2) & 0x30)) - 32;
scales[(k + 0) * TILE_N + n] = ls1;
scales[(k + 1) * TILE_N + n] = ls2;
sh >>= 4;
}
d[n] = B[n * KB].d;
}
}
template<typename TB, typename packed_B_t = packed_B_type<TB>>
void unpack_B(packed_B_t * RESTRICT tile, const void * RESTRICT packed_B) {
GGML_UNUSED(tile);
GGML_UNUSED(packed_B);
}
template <>
void unpack_B<block_q4_0>(int8_t * RESTRICT tile, const void * RESTRICT packed_B) {
const __m512i off = _mm512_set1_epi8(8);
const __m512i lowMask = _mm512_set1_epi8(0xF);
for (int n = 0; n < 8; n += 2) {
__m512i bytes = _mm512_loadu_si512((const __m512i *)((const char *)packed_B + n * 32));
const __m512i r0 = _mm512_sub_epi8(_mm512_and_si512(bytes, lowMask), off);
const __m512i r1 = _mm512_sub_epi8(_mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask), off);
_mm512_storeu_si512((__m512i *)(tile + n * 64 + 0), r0);
_mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1);
}
}
template <>
void unpack_B<block_q4_1>(uint8_t * RESTRICT tile, const void * RESTRICT packed_B) {
const __m512i lowMask = _mm512_set1_epi8(0xF);
for (int n = 0; n < 8; n += 2) {
__m512i bytes = _mm512_loadu_si512((const __m512i *)((const char *)packed_B + n * 32));
const __m512i r0 = _mm512_and_si512(bytes, lowMask);
const __m512i r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
_mm512_storeu_si512((__m512i *)(tile + n * 64 + 0), r0);
_mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1);
}
}
// packed_B_t for QKK is int8_t
template <typename TB>
void unpack_B(int8_t * RESTRICT tile, const void * RESTRICT packed_B, int k) {
const int packed_B_group_size = QK_K / 2 * TILE_N / 8;
const char * packed_B_group = (const char *)packed_B + k * packed_B_group_size;
const __m512i lowMask = _mm512_set1_epi8(0xF);
for (int n = 0; n < 8; n += 2) {
__m512i bytes = _mm512_loadu_si512(packed_B_group + n * 32);
const __m512i r0 = _mm512_and_si512(bytes, lowMask);
const __m512i r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
_mm512_storeu_si512((__m512i *)(tile + n * 64 + 0), r0);
_mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1);
}
}
template <>
void unpack_B<block_q5_K>(int8_t * RESTRICT tile, const void * RESTRICT packed_B, int k) {
// lower 4bits, stride 256 bytes
const int packed_l4_group_size = QK_K / 2 * TILE_N / 8;
const char * pb = (const char *)packed_B + k * packed_l4_group_size;
// higher 1bit, stride 64 bytes
const int packed_h1_group_size = QK_K / 8 * TILE_N / 8;
const char * ph = (const char *)packed_B + (QK_K / 2) * TILE_N + k * packed_h1_group_size;
const __m512i hbits = _mm512_loadu_si512(ph);
const __m512i lowMask = _mm512_set1_epi8(0xF);
__m512i hmask0 = _mm512_set1_epi8(0x1);
__m512i hmask1 = _mm512_set1_epi8(0x2);
for (int n = 0; n < 8; n += 2) {
__m512i bytes = _mm512_loadu_si512(pb + n * 32);
__m512i r0 = _mm512_and_si512(bytes, lowMask);
__m512i r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
__m512i h0 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask0), n), 4);
__m512i h1 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask1), n + 1), 4);
hmask0 = _mm512_slli_epi16(hmask0, 2);
hmask1 = _mm512_slli_epi16(hmask1, 2);
r0 = _mm512_add_epi8(r0, h0);
r1 = _mm512_add_epi8(r1, h1);
_mm512_storeu_si512((__m512i *)(tile + n * 64 + 0), r0);
_mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1);
}
}
template <>
void unpack_B<block_q6_K>(int8_t * RESTRICT tile, const void * RESTRICT packed_B, int k) {
// lower 4bits, stride 128 bytes
const int packed_l4_group_size = QK_K / 2 * TILE_N / 16;
const char * pb = (const char *)packed_B + k * packed_l4_group_size;
// higher 2bits, stride 64 bytes
const int packed_h2_group_size = QK_K / 4 * TILE_N / 16;
const char * ph = (const char *)packed_B + (QK_K / 2) * TILE_N + k * packed_h2_group_size;
const __m512i hbits = _mm512_loadu_si512(ph);
const __m512i off = _mm512_set1_epi8(32);
const __m512i lowMask = _mm512_set1_epi8(0xF);
__m512i hmask0 = _mm512_set1_epi8(0x3); // 0011
__m512i hmask1 = _mm512_set1_epi8(0xC); // 1100
// notes: skip zero padding from row4 to row7 as we have done so in `unpack_A`
__m512i bytes = _mm512_loadu_si512(pb);
__m512i r0 = _mm512_and_si512(bytes, lowMask);
__m512i r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
__m512i h0 = _mm512_slli_epi16(_mm512_and_si512(hbits, hmask0), 4);
__m512i h1 = _mm512_slli_epi16(_mm512_and_si512(hbits, hmask1), 2);
_mm512_storeu_si512((__m512i *)(tile + 0), _mm512_sub_epi8(_mm512_add_epi8(r0, h0), off));
_mm512_storeu_si512((__m512i *)(tile + 64), _mm512_sub_epi8(_mm512_add_epi8(r1, h1), off));
hmask0 = _mm512_slli_epi16(hmask0, 4);
hmask1 = _mm512_slli_epi16(hmask1, 4);
bytes = _mm512_loadu_si512(pb + 64);
r0 = _mm512_and_si512(bytes, lowMask);
r1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
h0 = _mm512_and_si512(hbits, hmask0);
h1 = _mm512_srli_epi16(_mm512_and_si512(hbits, hmask1), 2);
_mm512_storeu_si512((__m512i *)(tile + 128), _mm512_sub_epi8(_mm512_add_epi8(r0, h0), off));
_mm512_storeu_si512((__m512i *)(tile + 192), _mm512_sub_epi8(_mm512_add_epi8(r1, h1), off));
}
template <>
void unpack_B<block_iq4_xs>(int8_t * RESTRICT tile, const void * RESTRICT packed_B, int k) {
static const __m512i values128 = _mm512_set_epi8(
113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127,
113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127,
113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127,
113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127
);
const int packed_B_group_size = QK_K / 2 * TILE_N / 8;
const char * pb = (const char *)packed_B + k * packed_B_group_size;
const __m512i lowMask = _mm512_set1_epi8(0xF);
for (int n = 0; n < 8; n += 2) {
__m512i bytes = _mm512_loadu_si512(pb + n * 32);
const __m512i r0 = _mm512_shuffle_epi8(values128, _mm512_and_si512(bytes, lowMask));
const __m512i r1 = _mm512_shuffle_epi8(values128, _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask));
_mm512_storeu_si512((__m512i *)(tile + n * 64 + 0), r0);
_mm512_storeu_si512((__m512i *)(tile + n * 64 + 64), r1);
}
}
template <typename TA, typename TB, bool is_acc>
struct acc_C {};
template <bool is_acc>
struct acc_C<block_q8_0, block_q4_0, is_acc> {
static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_0 * A, int lda, const void * packed_B, int nr) {
const int offset = TILE_N * TILE_K / 2;
const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)((const char *)packed_B + offset)));
for (int m = 0; m < nr; ++m) {
const __m512 vd1 = _mm512_set1_ps(GGML_FP16_TO_FP32(A[m * lda].d));
const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N));
__m512 vsum;
if (is_acc) {
vsum = _mm512_loadu_ps(C + m * ldc);
} else {
vsum = _mm512_set1_ps(0.f);
}
vsum = _mm512_fmadd_ps(vtile, _mm512_mul_ps(vd0, vd1), vsum);
_mm512_storeu_ps(C + m * ldc, vsum);
}
}
};
template <bool is_acc>
struct acc_C<block_q8_1, block_q4_1, is_acc> {
static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_1 * A, int lda, const void * packed_B, int nr) {
const int offset = TILE_N * TILE_K / 2;
const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)((const char *)packed_B + offset)));
const __m512 vm0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)((const char *)packed_B + offset + TILE_N * sizeof(ggml_half))));
for (int m = 0; m < nr; ++m) {
const __m512 vd1 = _mm512_set1_ps(GGML_FP16_TO_FP32(A[m * lda].d));
const __m512 vs1 = _mm512_set1_ps(GGML_FP16_TO_FP32(A[m * lda].s));
const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N));
__m512 vsum;
if (is_acc) {
vsum = _mm512_loadu_ps(C + m * ldc);
} else {
vsum = _mm512_set1_ps(0.f);
}
vsum = _mm512_fmadd_ps(vtile, _mm512_mul_ps(vd0, vd1), vsum);
vsum = _mm512_fmadd_ps(vm0, vs1, vsum);
_mm512_storeu_ps(C + m * ldc, vsum);
}
}
};
template <bool is_acc>
struct acc_C<block_q8_0, block_q8_0, is_acc> {
static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_0 * A, int lda, const void * packed_B, int nr) {
const int offset = TILE_N * TILE_K;
const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)((const char *)packed_B + offset)));
for (int m = 0; m < nr; ++m) {
const __m512 vd1 = _mm512_set1_ps(GGML_FP16_TO_FP32(A[m * lda].d));
const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N));
__m512 vsum;
if (is_acc) {
vsum = _mm512_loadu_ps(C + m * ldc);
} else {
vsum = _mm512_set1_ps(0.f);
}
vsum = _mm512_fmadd_ps(vtile, _mm512_mul_ps(vd0, vd1), vsum);
_mm512_storeu_ps(C + m * ldc, vsum);
}
}
};
template <bool is_acc>
struct acc_C<block_q8_K, block_q4_K, is_acc> {
static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_K * A, int lda, const void * packed_B, int nr) {
const uint8_t * scales = reinterpret_cast<const uint8_t *>((const char *)packed_B + (QK_K / 2) * TILE_N);
const uint8_t * mins = scales + 8 * TILE_N;
const ggml_half * d0 = reinterpret_cast<const ggml_half *>(mins + 8 * TILE_N);
const ggml_half * dmin = d0 + TILE_N;
const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)d0));
const __m512 vdmin = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)dmin));
for (int m = 0; m < nr; ++m) {
const float d1 = A[m * lda].d;
const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(d1), vd0);
const __m512 vdm = _mm512_mul_ps(_mm512_set1_ps(-d1), vdmin);
const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N));
__m512 vsum;
if (is_acc) {
vsum = _mm512_loadu_ps(C + m * ldc);
} else {
vsum = _mm512_set1_ps(0.f);
}
const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[m * lda].bsums);
const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1));
__m512i acc_m = _mm512_setzero_si512();
for (int k = 0; k < 4; ++k) {
__m512i vmask = _mm512_set1_epi32(k);
__m512i va = _mm512_permutexvar_epi32(vmask, _mm512_castsi128_si512(q8s));
__m512i vb = _mm512_cvtepi8_epi16(_mm256_loadu_si256((const __m256i *)(mins + k * 32)));
acc_m = _mm512_dpwssds_epi32(acc_m, va, vb);
}
vsum = _mm512_fmadd_ps(vtile, vd, vsum);
vsum = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc_m), vdm, vsum);
_mm512_storeu_ps(C + m * ldc, vsum);
}
}
};
template <bool is_acc>
struct acc_C<block_q8_K, block_q5_K, is_acc> {
static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_K * A, int lda, const void * packed_B, int nr) {
const uint8_t * scales = reinterpret_cast<const uint8_t *>((const char *)packed_B + (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N);
const uint8_t * mins = scales + 8 * TILE_N;
const ggml_half * d0 = reinterpret_cast<const ggml_half *>(mins + 8 * TILE_N);
const ggml_half * dmin = d0 + TILE_N;
const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)d0));
const __m512 vdmin = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)dmin));
for (int m = 0; m < nr; ++m) {
const float d1 = A[m * lda].d;
const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(d1), vd0);
const __m512 vdm = _mm512_mul_ps(_mm512_set1_ps(-d1), vdmin);
const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N));
__m512 vsum;
if (is_acc) {
vsum = _mm512_loadu_ps(C + m * ldc);
} else {
vsum = _mm512_set1_ps(0.f);
}
const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[m * lda].bsums);
const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1));
__m512i acc_m = _mm512_setzero_si512();
for (int k = 0; k < 4; ++k) {
__m512i vmask = _mm512_set1_epi32(k);
__m512i va = _mm512_permutexvar_epi32(vmask, _mm512_castsi128_si512(q8s));
__m512i vb = _mm512_cvtepi8_epi16(_mm256_loadu_si256((const __m256i *)(mins + k * 32)));
acc_m = _mm512_dpwssds_epi32(acc_m, va, vb);
}
vsum = _mm512_fmadd_ps(vtile, vd, vsum);
vsum = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc_m), vdm, vsum);
_mm512_storeu_ps(C + m * ldc, vsum);
}
}
};
template <bool is_acc>
struct acc_C<block_q8_K, block_q6_K, is_acc> {
static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_K * A, int lda, const void * packed_B, int nr) {
const uint8_t * scales = reinterpret_cast<const uint8_t *>((const char *)packed_B + (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N);
const ggml_half * d0 = reinterpret_cast<const ggml_half *>(scales + 16 * TILE_N);
const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)d0));
for (int m = 0; m < nr; ++m) {
const float d1 = A[m * lda].d;
const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(d1), vd0);
const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N));
__m512 vsum;
if (is_acc) {
vsum = _mm512_loadu_ps(C + m * ldc);
} else {
vsum = _mm512_set1_ps(0.f);
}
vsum = _mm512_fmadd_ps(vtile, vd, vsum);
_mm512_storeu_ps(C + m * ldc, vsum);
}
}
};
template <bool is_acc>
struct acc_C<block_q8_K, block_iq4_xs, is_acc> {
static void apply(float * RESTRICT C, int ldc, const int32_t * RESTRICT tile, const block_q8_K * A, int lda, const void * packed_B, int nr) {
const int8_t * scales = reinterpret_cast<const int8_t *>((const char *)packed_B + (QK_K / 2) * TILE_N);
const ggml_half * d0 = reinterpret_cast<const ggml_half *>(scales + 8 * TILE_N);
const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)d0));
for (int m = 0; m < nr; ++m) {
const float d1 = A[m * lda].d;
const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(d1), vd0);
const __m512 vtile = _mm512_cvtepi32_ps(_mm512_loadu_si512(tile + m * TILE_N));
__m512 vsum;
if (is_acc) {
vsum = _mm512_loadu_ps(C + m * ldc);
} else {
vsum = _mm512_set1_ps(0.f);
}
vsum = _mm512_fmadd_ps(vtile, vd, vsum);
_mm512_storeu_ps(C + m * ldc, vsum);
}
}
};
template <typename TB> constexpr int get_quants_size();
template <> constexpr int get_quants_size<block_q4_K>() { return (QK_K / 2) * TILE_N; }
template <> constexpr int get_quants_size<block_q5_K>() { return (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N; }
template <> constexpr int get_quants_size<block_q6_K>() { return (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N; }
template <> constexpr int get_quants_size<block_iq4_xs>() { return (QK_K / 2) * TILE_N; }
// used for QKK format
template <typename TB, bool is_acc,
typename std::enable_if<is_type_qkk<TB>::value, int>::type = 0>
inline void scale_C(const int32_t * RESTRICT tile, int32_t * RESTRICT sumi, const void * packed_B, int k, int nr) {
const uint8_t * scales = reinterpret_cast<const uint8_t *>((const char *)packed_B + get_quants_size<TB>());
const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(scales + k * TILE_N)));
for (int m = 0; m < nr; ++m) {
__m512i vsumi;
if (is_acc) {
vsumi = _mm512_loadu_si512(sumi + m * TILE_N);
} else {
vsumi = _mm512_setzero_si512();
}
__m512i vtile = _mm512_loadu_si512(tile + m * TILE_N);
vsumi = _mm512_add_epi32(vsumi, _mm512_mullo_epi32(vtile, vscale));
_mm512_storeu_si512((__m512i *)(sumi + m * TILE_N), vsumi);
}
}
template <typename TA, typename TB, typename TC, int BLOCK_M, int BLOCK_N, int BLOCK_K>
struct tinygemm_kernel_avx {
static void apply(int K, const TA * RESTRICT A, const TB * RESTRICT B, TC * RESTRICT C, int ldc) {
GGML_UNUSED(K);
GGML_UNUSED(A);
GGML_UNUSED(B);
GGML_UNUSED(C);
GGML_UNUSED(ldc);
}
};
template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
struct tinygemm_kernel_avx<float, ggml_fp16_t, float, BLOCK_M, BLOCK_N, BLOCK_K> {
static void apply(int K, const float * RESTRICT A, const ggml_fp16_t * RESTRICT B, float * RESTRICT C, int ldc) {
constexpr int ROWS = BLOCK_M;
constexpr int COLS = BLOCK_N;
assert(BLOCK_K == 16);
__m512 va;
__m512 vb[COLS];
__m512 vc[ROWS * COLS];
auto loadc = [&](auto idx) {
vc[idx] = _mm512_setzero_ps();
};
Unroll<ROWS * COLS>{}(loadc);
auto compute = [&](auto idx, auto k) {
constexpr int row = idx / COLS;
constexpr int col = idx % COLS;
if constexpr (col == 0) {
va = _mm512_loadu_ps(A + row * K + k);
}
if constexpr (row == 0) {
vb[col] = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(B + col * K + k)));
}
vc[idx] = _mm512_fmadd_ps(va, vb[col], vc[idx]);
};
for (int k = 0; k < K; k += 16) {
Unroll<ROWS * COLS>{}(compute, k);
}
auto storec = [&](auto idx) {
constexpr int row = idx / COLS;
constexpr int col = idx % COLS;
C[row * ldc + col] = _mm512_reduce_add_ps(vc[idx]);
};
Unroll<ROWS * COLS>{}(storec);
}
};
#define LAUNCH_TINYGEMM_KERNEL_AVX(MB_SIZE, NB_SIZE) \
tinygemm_kernel_avx<float, type, float, MB_SIZE, NB_SIZE, blck_size>::apply( \
K, (const float *)src1->data + mb_start * K, \
(const type *)src0->data + nb_start * K, \
(float *)dst->data + mb_start * ldc + nb_start, ldc);
// re-organize in the format {NB, KB, TILE_SIZE}:
#define PACKED_INDEX(n, k, KB, tile_size) (n * KB + k) * tile_size
template<typename TB, int BLOCK_K>
void convert_B_packed_format(void * RESTRICT packed_B, const TB * RESTRICT B, int N, int K, int n_threads) {
const int NB = N / TILE_N;
const int KB = K / BLOCK_K;
const int TILE_SIZE = get_tile_size<TB>();
// parallel on NB should be enough
parallel_for(n_threads, NB, [&](int begin, int end) {
for (int n = begin; n < end; ++n) {
for (int k = 0; k < KB; ++k) {
int n0 = n * TILE_N;
pack_B((char *)packed_B + PACKED_INDEX(n, k, KB, TILE_SIZE), &B[n0 * KB + k], KB);
}
}
});
}
template <typename TA, typename TB, typename TC, int BLOCK_M, int BLOCK_N, int BLOCK_K>
struct tinygemm_kernel_vnni {};
template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
struct tinygemm_kernel_vnni<block_q8_0, block_q4_0, float, BLOCK_M, BLOCK_N, BLOCK_K> {
static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) {
constexpr int COLS = BLOCK_N / 16;
const int TILE_SIZE = TILE_N * sizeof(block_q4_0);
const block_q8_0 * RESTRICT A = static_cast<const block_q8_0 *>(_A);
const char * RESTRICT B = static_cast<const char *>(_B);
__m512i va[8];
__m512 vc[COLS];
__m512 vd1;
// sum of offsets, shared across COLS
//
// avx512-vnni does not have `_mm512_dpbssd_epi32`,
// need to transfrom ss to us:
// a * (b - 8) is equavilent to b * a - 8 * a
// s u u u s u s
//
__m512i vcomp;
const __m512i off = _mm512_set1_epi8(8);
const __m512i lowMask = _mm512_set1_epi8(0xF);
auto loadc = [&](auto col) {
vc[col] = _mm512_setzero_ps();
};
Unroll<COLS>{}(loadc);
auto compute = [&](auto col, auto i) {
// load a and compute compensation
if constexpr (col == 0) {
const int32_t * a_ptr = reinterpret_cast<const int32_t *>(A[0 * KB + i].qs);
vcomp = _mm512_setzero_si512();
for (int k = 0; k < 8; ++k) {
va[k] = _mm512_set1_epi32(a_ptr[k]);
vcomp = _mm512_dpbusd_epi32(vcomp, off, va[k]);
}
vd1 = _mm512_set1_ps(GGML_FP16_TO_FP32(A[0 * KB + i].d));
}
// load b
__m512i vsum = _mm512_setzero_si512();
const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE);
for (int k = 0; k < 8; k += 2) {
__m512i bytes = _mm512_loadu_si512((const __m512i *)(b_ptr + k * 32));
__m512i vb0 = _mm512_and_si512(bytes, lowMask);
vsum = _mm512_dpbusd_epi32(vsum, vb0, va[k + 0]);
__m512i vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
vsum = _mm512_dpbusd_epi32(vsum, vb1, va[k + 1]);
}
const int offset = TILE_N * TILE_K / 2;
const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset)));
vsum = _mm512_sub_epi32(vsum, vcomp);
vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(vsum), _mm512_mul_ps(vd0, vd1), vc[col]);
};
for (int i = 0; i < KB; ++i) {
Unroll<COLS>{}(compute, i);
}
//store to C
auto storec = [&](auto col) {
_mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]);
};
Unroll<COLS>{}(storec);
}
};
template <int BLOCK_N, int BLOCK_K>
struct tinygemm_kernel_vnni<block_q8_1, block_q4_1, float, 1, BLOCK_N, BLOCK_K> {
static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) {
constexpr int COLS = BLOCK_N / 16;
const int TILE_SIZE = TILE_N * sizeof(block_q4_1);
const block_q8_1 * RESTRICT A = static_cast<const block_q8_1 *>(_A);
const char * RESTRICT B = static_cast<const char *>(_B);
__m512i va[8];
__m512i vb[8];
__m512 vc[COLS];
__m512 vd1, vs1;
const __m512i lowMask = _mm512_set1_epi8(0xF);
auto loadc = [&](auto col) {
vc[col] = _mm512_setzero_ps();
};
Unroll<COLS>{}(loadc);
auto compute = [&](auto col, auto i) {
// load a
if constexpr (col == 0) {
const int32_t * a_ptr = reinterpret_cast<const int32_t *>(A[0 * KB + i].qs);
for (int k = 0; k < 8; ++k) {
va[k] = _mm512_set1_epi32(a_ptr[k]);
}
vd1 = _mm512_set1_ps(GGML_FP16_TO_FP32(A[0 * KB + i].d));
vs1 = _mm512_set1_ps(GGML_FP16_TO_FP32(A[0 * KB + i].s));
}
// load b
const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE);
for (int k = 0; k < 8; k += 2) {
__m512i bytes = _mm512_loadu_si512((const __m512i *)(b_ptr + k * 32));
vb[k + 0] = _mm512_and_si512(bytes, lowMask);
vb[k + 1] = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
}
const int offset = TILE_N * TILE_K / 2;
const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset)));
const __m512 vm0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset + TILE_N * sizeof(ggml_half))));
__m512i vsum = _mm512_setzero_si512();
for (int k = 0; k < 8; ++k) {
vsum = _mm512_dpbusd_epi32(vsum, vb[k], va[k]);
}
vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(vsum), _mm512_mul_ps(vd0, vd1), vc[col]);
vc[col] = _mm512_fmadd_ps(vm0, vs1, vc[col]);
};
for (int i = 0; i < KB; ++i) {
Unroll<COLS>{}(compute, i);
}
//store to C
auto storec = [&](auto col) {
_mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]);
};
Unroll<COLS>{}(storec);
}
};
template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
struct tinygemm_kernel_vnni<block_q8_0, block_q8_0, float, BLOCK_M, BLOCK_N, BLOCK_K> {
static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) {
constexpr int COLS = BLOCK_N / 16;
const int TILE_SIZE = TILE_N * sizeof(block_q8_0) + TILE_N * sizeof(int32_t);
const block_q8_0 * RESTRICT A = static_cast<const block_q8_0 *>(_A);
const char * RESTRICT B = static_cast<const char *>(_B);
__m512i va[8];
__m512i vb[8];
__m512 vc[COLS];
__m512 vd1;
// Notes: s8s8 igemm compensation in avx512-vnni
// change s8s8 to u8s8 with compensate
// a * b = (a + 128) * b - 128 * b
// s s u s u s
//
// (128 * b is pre-computed when packing B to vnni formats)
//
const __m512i off = _mm512_set1_epi8(static_cast<char>(0x80));
auto loadc = [&](auto col) {
vc[col] = _mm512_setzero_ps();
};
Unroll<COLS>{}(loadc);
auto compute = [&](auto col, auto i) {
// load a and add offset 128
if constexpr (col == 0) {
const int32_t * a_ptr = reinterpret_cast<const int32_t *>(A[0 * KB + i].qs);
for (int k = 0; k < 8; ++k) {
va[k] = _mm512_set1_epi32(a_ptr[k]);
va[k] = _mm512_add_epi8(va[k], off);
}
vd1 = _mm512_set1_ps(GGML_FP16_TO_FP32(A[0 * KB + i].d));
}
// load b
const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE);
for (int k = 0; k < 8; ++k) {
vb[k] = _mm512_loadu_si512((const __m512i *)(b_ptr + k * 64));
}
const int offset = TILE_N * TILE_K;
const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset)));
const int offset2 = TILE_N * TILE_K + TILE_N * sizeof(ggml_half);
const __m512i vcomp = _mm512_loadu_si512((const __m512i *)(b_ptr + offset2));
__m512i vsum = _mm512_setzero_si512();
for (int k = 0; k < 8; ++k) {
vsum = _mm512_dpbusd_epi32(vsum, va[k], vb[k]);
}
vsum = _mm512_sub_epi32(vsum, vcomp);
vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(vsum), _mm512_mul_ps(vd0, vd1), vc[col]);
};
for (int i = 0; i < KB; ++i) {
Unroll<COLS>{}(compute, i);
}
//store to C
auto storec = [&](auto col) {
_mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]);
};
Unroll<COLS>{}(storec);
}
};
template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
struct tinygemm_kernel_vnni<block_q8_K, block_q4_K, float, BLOCK_M, BLOCK_N, BLOCK_K> {
static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) {
constexpr int COLS = BLOCK_N / 16;
const int TILE_SIZE = TILE_N * sizeof(block_q4_K) + TILE_N * 4;
const block_q8_K * RESTRICT A = static_cast<const block_q8_K *>(_A);
const char * RESTRICT B = static_cast<const char *>(_B);
// a.qs: 8 groups, 32 bytes each group (m256i)
__m512i va[8];
// a.bsum: 8 groups, 2 bytes each group (m128i)
__m512i va_bsum;
__m512 vc[COLS];
__m512 vd1;
// packed_B:
const int offset_scales = (QK_K / 2) * TILE_N;
const int offset_mins = (QK_K / 2) * TILE_N + 8 * TILE_N;
const int offset_d0 = (QK_K / 2) * TILE_N + 16 * TILE_N;
const int offset_dmin = (QK_K / 2) * TILE_N + 16 * TILE_N + TILE_N * sizeof(ggml_half);
const __m512i lowMask = _mm512_set1_epi8(0xF);
auto loadc = [&](auto col) {
vc[col] = _mm512_setzero_ps();
};
Unroll<COLS>{}(loadc);
// Notes: vnni formats in QK_K
// a) quants vnni format
// int8 {k/4, n, 4}, viewed as 2d {k/4, 4n}, k = 32
// from {16, 32} to {8, 64}
//
// b) min vnni format
// int16 {k/2, n, 2}, viewed as 2d {k/2, 2n}, k = 8
// from {16, 8} to {4, 32}
//
auto compute = [&](auto col, auto i) {
// load a
if constexpr (col == 0) {
for (int k_group = 0; k_group < QK_K / 32; ++k_group) {
va[k_group] = _mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)(A[0 * KB + i].qs + k_group * 32)));
}
const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[0 * KB + i].bsums);
const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1));
va_bsum = _mm512_castsi128_si512(q8s);
vd1 = _mm512_set1_ps(A[0 * KB + i].d);
}
// step 1: accumultate the quants
__m512i acc = _mm512_setzero_si512();
const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE);
const char * b_qs = b_ptr;
for (int k_group = 0; k_group < QK_K / 32; ++k_group) {
__m512i vsum = _mm512_setzero_si512();
for (int k = 0; k < 8; k += 2) {
__m512i va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(k + 0), va[k_group]);
__m512i va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(k + 1), va[k_group]);
__m512i bytes = _mm512_loadu_si512((const __m512i *)b_qs);
__m512i vb0 = _mm512_and_si512(bytes, lowMask);
vsum = _mm512_dpbusd_epi32(vsum, vb0, va0);
__m512i vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
vsum = _mm512_dpbusd_epi32(vsum, vb1, va1);
b_qs += 64;
}
// vacc += scale * (q8 @ q4)
const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(b_ptr + offset_scales + k_group * TILE_N)));
acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(vsum, vscale));
}
const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_d0)));
vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc), _mm512_mul_ps(vd0, vd1), vc[col]);
// step 2: accumulate the mins
__m512i acc_m = _mm512_setzero_si512();
for (int k = 0; k < 4; ++k) {
__m512i vmask = _mm512_set1_epi32(k);
__m512i va = _mm512_permutexvar_epi32(vmask, va_bsum);
__m512i vb = _mm512_cvtepi8_epi16(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_mins + k * 32)));
acc_m = _mm512_dpwssds_epi32(acc_m, va, vb);
}
const __m512 vdmin = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_dmin)));
vc[col] = _mm512_fnmadd_ps(_mm512_cvtepi32_ps(acc_m), _mm512_mul_ps(vdmin, vd1), vc[col]);
};
for (int i = 0; i < KB; ++i) {
Unroll<COLS>{}(compute, i);
}
//store to C
auto storec = [&](auto col) {
_mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]);
};
Unroll<COLS>{}(storec);
}
};
template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
struct tinygemm_kernel_vnni<block_q8_K, block_q5_K, float, BLOCK_M, BLOCK_N, BLOCK_K> {
static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) {
constexpr int COLS = BLOCK_N / 16;
const int TILE_SIZE = TILE_N * sizeof(block_q5_K) + TILE_N * 4;
const block_q8_K * RESTRICT A = static_cast<const block_q8_K *>(_A);
const char * RESTRICT B = static_cast<const char *>(_B);
// a.qs: 8 groups, 32 bytes each group (m256i)
__m512i va[8];
// a.bsum: 8 groups, 2 bytes each group (m128i)
__m512i va_bsum;
__m512 vc[COLS];
__m512 vd1;
// packed_B:
const int offset_qh = (QK_K / 2) * TILE_N;
const int offset_scales = (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N;
const int offset_mins = (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N + 8 * TILE_N;
const int offset_d0 = (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N + 16 * TILE_N;
const int offset_dmin = (QK_K / 2) * TILE_N + (QK_K / 8) * TILE_N + 16 * TILE_N + TILE_N * sizeof(ggml_half);
const __m512i lowMask = _mm512_set1_epi8(0xF);
auto loadc = [&](auto col) {
vc[col] = _mm512_setzero_ps();
};
Unroll<COLS>{}(loadc);
// Q5_K and Q4_K shares the same vnni formats, refer to notes above.
auto compute = [&](auto col, auto i) {
// load a
if constexpr (col == 0) {
for (int k_group = 0; k_group < QK_K / 32; ++k_group) {
va[k_group] = _mm512_castsi256_si512(_mm256_loadu_si256((const __m256i *)(A[0 * KB + i].qs + k_group * 32)));
}
const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[0 * KB + i].bsums);
const __m128i q8s = _mm_hadd_epi16(_mm256_extracti128_si256(q8sums, 0), _mm256_extracti128_si256(q8sums, 1));
va_bsum = _mm512_castsi128_si512(q8s);
vd1 = _mm512_set1_ps(A[0 * KB + i].d);
}
// step 1: accumultate the quants
__m512i acc = _mm512_setzero_si512();
const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE);
const char * b_qs = b_ptr;
const char * b_qh = b_ptr + offset_qh;
for (int k_group = 0; k_group < QK_K / 32; ++k_group) {
__m512i vsum = _mm512_setzero_si512();
__m512i hmask0 = _mm512_set1_epi8(0x1);
__m512i hmask1 = _mm512_set1_epi8(0x2);
__m512i hbits = _mm512_loadu_si512((const __m512i *)(b_qh + k_group * 64));
for (int k = 0; k < 8; k += 2) {
__m512i va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(k + 0), va[k_group]);
__m512i va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(k + 1), va[k_group]);
__m512i bytes = _mm512_loadu_si512((const __m512i *)b_qs);
__m512i vb0 = _mm512_and_si512(bytes, lowMask);
__m512i vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
__m512i vh0 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask0), k), 4);
__m512i vh1 = _mm512_slli_epi16(_mm512_srli_epi16(_mm512_and_si512(hbits, hmask1), k + 1), 4);
hmask0 = _mm512_slli_epi16(hmask0, 2);
hmask1 = _mm512_slli_epi16(hmask1, 2);
vb0 = _mm512_add_epi8(vb0, vh0);
vb1 = _mm512_add_epi8(vb1, vh1);
vsum = _mm512_dpbusd_epi32(vsum, vb0, va0);
vsum = _mm512_dpbusd_epi32(vsum, vb1, va1);
b_qs += 64;
}
// vacc += scale * (q8 @ q5)
const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(b_ptr + offset_scales + k_group * TILE_N)));
acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(vsum, vscale));
}
const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_d0)));
vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc), _mm512_mul_ps(vd0, vd1), vc[col]);
// step 2: accumulate the mins
__m512i acc_m = _mm512_setzero_si512();
for (int k = 0; k < 4; ++k) {
__m512i vmask = _mm512_set1_epi32(k);
__m512i va = _mm512_permutexvar_epi32(vmask, va_bsum);
__m512i vb = _mm512_cvtepi8_epi16(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_mins + k * 32)));
acc_m = _mm512_dpwssds_epi32(acc_m, va, vb);
}
const __m512 vdmin = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_dmin)));
vc[col] = _mm512_fnmadd_ps(_mm512_cvtepi32_ps(acc_m), _mm512_mul_ps(vdmin, vd1), vc[col]);
};
for (int i = 0; i < KB; ++i) {
Unroll<COLS>{}(compute, i);
}
//store to C
auto storec = [&](auto col) {
_mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]);
};
Unroll<COLS>{}(storec);
}
};
template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
struct tinygemm_kernel_vnni<block_q8_K, block_q6_K, float, BLOCK_M, BLOCK_N, BLOCK_K> {
static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) {
constexpr int COLS = BLOCK_N / 16;
const int TILE_SIZE = TILE_N * sizeof(block_q6_K);
const block_q8_K * RESTRICT A = static_cast<const block_q8_K *>(_A);
const char * RESTRICT B = static_cast<const char *>(_B);
// load the 256 bytes from A to 4 avx512 vectors
__m512i va[4];
__m512 vc[COLS];
__m512 vd1;
// packed_B:
const int offset_qh = (QK_K / 2) * TILE_N;
const int offset_scales = (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N;
const int offset_d0 = (QK_K / 2) * TILE_N + (QK_K / 4) * TILE_N + 16 * TILE_N;
// compensation
__m512i vcomp;
const __m512i m32s = _mm512_set1_epi32(32);
const __m512i lowMask = _mm512_set1_epi8(0xF);
auto loadc = [&](auto col) {
vc[col] = _mm512_setzero_ps();
};
Unroll<COLS>{}(loadc);
auto compute = [&](auto col, auto i) {
if constexpr (col == 0) {
// load a
va[0] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 0));
va[1] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 64));
va[2] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 128));
va[3] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 192));
const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[0 * KB + i].bsums);
vcomp = _mm512_mullo_epi32(_mm512_cvtepi16_epi32(q8sums), m32s);
vd1 = _mm512_set1_ps(A[0 * KB + i].d);
}
// accmulate the quants
__m512i acc = _mm512_setzero_si512();
const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE);
const char * b_qs = b_ptr;
const char * b_qh = b_ptr + offset_qh;
int mask = 0;
for (int k_group = 0; k_group < QK_K / 16; ++k_group) {
int r = k_group >> 2;
__m512i va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]);
__m512i va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]);
__m512i vsum = _mm512_setzero_si512();
__m512i hmask = _mm512_set1_epi8(0x3);
__m512i bytes = _mm512_loadu_si512(b_qs);
__m512i hbits = _mm512_loadu_si512(b_qh);
__m512i vb0 = _mm512_and_si512(bytes, lowMask);
__m512i vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
__m512i vh0 = _mm512_slli_epi16(_mm512_and_si512(hbits, hmask), 4);
__m512i vh1 = _mm512_slli_epi16(_mm512_and_si512(hbits, _mm512_slli_epi16(hmask, 2)), 2);
vb0 = _mm512_add_epi8(vb0, vh0);
vb1 = _mm512_add_epi8(vb1, vh1);
vsum = _mm512_dpbusd_epi32(vsum, vb0, va0);
vsum = _mm512_dpbusd_epi32(vsum, vb1, va1);
b_qs += 64;
va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]);
va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]);
bytes = _mm512_loadu_si512(b_qs);
vb0 = _mm512_and_si512(bytes, lowMask);
vb1 = _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask);
vh0 = _mm512_and_si512(hbits, _mm512_slli_epi16(hmask, 4));
vh1 = _mm512_srli_epi16(_mm512_and_si512(hbits, _mm512_slli_epi16(hmask, 6)), 2);
vb0 = _mm512_add_epi8(vb0, vh0);
vb1 = _mm512_add_epi8(vb1, vh1);
vsum = _mm512_dpbusd_epi32(vsum, vb0, va0);
vsum = _mm512_dpbusd_epi32(vsum, vb1, va1);
b_qs += 64;
b_qh += 64;
// B * A - 32 * A
__m512i vmask = _mm512_set1_epi32(k_group);
vsum = _mm512_sub_epi32(vsum, _mm512_permutexvar_epi32(vmask, vcomp));
// vacc += scale * (q8 @ q6)
const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(b_ptr + offset_scales + k_group * TILE_N)));
acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(vsum, vscale));
}
const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_d0)));
vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc), _mm512_mul_ps(vd0, vd1), vc[col]);
};
for (int i = 0; i < KB; ++i) {
Unroll<COLS>{}(compute, i);
}
//store to C
auto storec = [&](int col) {
_mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]);
};
Unroll<COLS>{}(storec);
}
};
template <int BLOCK_M, int BLOCK_N, int BLOCK_K>
struct tinygemm_kernel_vnni<block_q8_K, block_iq4_xs, float, BLOCK_M, BLOCK_N, BLOCK_K> {
static void apply(int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) {
constexpr int COLS = BLOCK_N / 16;
const int TILE_SIZE = TILE_N * sizeof(block_iq4_xs) + TILE_N * 2;
const block_q8_K * RESTRICT A = static_cast<const block_q8_K *>(_A);
const char * RESTRICT B = static_cast<const char *>(_B);
// load the 256 bytes from A to 4 avx512 vectors
__m512i va[4];
__m512 vc[COLS];
__m512 vd1;
// packed_B:
const int offset_scales = (QK_K / 2) * TILE_N ;
const int offset_d0 = (QK_K / 2) * TILE_N + 8 * TILE_N;
// compensation
__m512i vcomp;
const __m256i m128s = _mm256_set1_epi16(128);
const __m512i lowMask = _mm512_set1_epi8(0xF);
const __m512i values128 = _mm512_set_epi8(
113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127,
113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127,
113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127,
113, 89, 69, 53, 38, 25, 13, 1, -10, -22, -35, -49, -65, -83, -104, -127
);
const __m512i off = _mm512_set1_epi8(static_cast<char>(0x80));
const __m512i values256 = _mm512_add_epi8(values128, off);
auto loadc = [&](auto col) {
vc[col] = _mm512_setzero_ps();
};
Unroll<COLS>{}(loadc);
auto compute = [&](auto col, auto i) {
if constexpr (col == 0) {
// load a
va[0] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 0));
va[1] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 64));
va[2] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 128));
va[3] = _mm512_loadu_si512((const __m512i *)(A[0 * KB + i].qs + 192));
// compensation: 128 * A
const __m256i q8sums = _mm256_loadu_si256((const __m256i *)A[0 * KB + i].bsums);
vcomp = _mm512_castsi256_si512(_mm256_madd_epi16(q8sums, m128s));
vd1 = _mm512_set1_ps(A[0 * KB + i].d);
}
// accmulate the quants
__m512i acc = _mm512_setzero_si512();
const char * b_ptr = B + PACKED_INDEX(col, i, KB, TILE_SIZE);
const char * b_qs = b_ptr;
int mask = 0;
for (int k_group = 0; k_group < QK_K / 32; ++k_group) {
int r = k_group >> 1;
__m512i vmask = _mm512_set1_epi32(k_group);
__m512i vsum = _mm512_setzero_si512();
for (int k = 0; k < 8; k += 2) {
__m512i va0 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]);
__m512i va1 = _mm512_permutexvar_epi32(_mm512_set1_epi32(mask++), va[r]);
__m512i bytes = _mm512_loadu_si512(b_qs);
__m512i vb0 = _mm512_shuffle_epi8(values256, _mm512_and_si512(bytes, lowMask));
__m512i vb1 = _mm512_shuffle_epi8(values256, _mm512_and_si512(_mm512_srli_epi16(bytes, 4), lowMask));
vsum = _mm512_dpbusd_epi32(vsum, vb0, va0);
vsum = _mm512_dpbusd_epi32(vsum, vb1, va1);
b_qs += 64;
}
// (B + 128) * A - 128 * A
vsum = _mm512_sub_epi32(vsum, _mm512_permutexvar_epi32(vmask, vcomp));
// vacc += scale * (q8 @ q4)
const __m512i vscale = _mm512_cvtepi8_epi32(_mm_loadu_si128((const __m128i *)(b_ptr + offset_scales + k_group * TILE_N)));
acc = _mm512_add_epi32(acc, _mm512_mullo_epi32(vsum, vscale));
}
const __m512 vd0 = _mm512_cvtph_ps(_mm256_loadu_si256((const __m256i *)(b_ptr + offset_d0)));
vc[col] = _mm512_fmadd_ps(_mm512_cvtepi32_ps(acc), _mm512_mul_ps(vd0, vd1), vc[col]);
};
for (int i = 0; i < KB; ++i) {
Unroll<COLS>{}(compute, i);
}
//store to C
auto storec = [&](auto col) {
_mm512_storeu_ps((__m512i*)(C + 0 * ldc + col * 16), vc[col]);
};
Unroll<COLS>{}(storec);
}
};
#define LAUNCH_TINYGEMM_KERNEL_VNNI(NB_SIZE) \
tinygemm_kernel_vnni<vec_dot_type, type, float, 1, NB_SIZE, blck_size>::apply( \
KB, (const char *)wdata + 0 * row_size_A, \
(const char *)src0->data + PACKED_INDEX(nb * kTilesN, 0, KB, TILE_SIZE), \
(float *) dst->data + 0 * N + nb_start, ldc)
template <typename TA, typename TB, typename TC, int BLOCK_K,
typename std::enable_if<!is_type_qkk<TB>::value, int>::type = 0>
void tinygemm_kernel_amx(int M, int N, int KB, const void * RESTRICT _A, const void * RESTRICT _B, TC * RESTRICT C, int ldc) {
using packed_B_t = packed_B_type<TB>;
const int TILE_SIZE = get_tile_size<TB>();
const bool need_unpack = do_unpack<TB>::value;
GGML_ASSERT(M <= 2 * TILE_M && N == 2 * TILE_N);
const TA * RESTRICT A = static_cast<const TA *>(_A);
const char * RESTRICT B = static_cast<const char *>(_B);
const int m0 = std::min(M, TILE_M);
const int m1 = std::max(M - TILE_M, 0);
const int lda = KB * sizeof(TA);
//const int ldb = KB * sizeof(TB);
static thread_local packed_B_t Tile0[TILE_N * TILE_K];
static thread_local packed_B_t Tile1[TILE_N * TILE_K];
static thread_local int8_t Tile23[TILE_M * TILE_K];
static thread_local int32_t TileC0[TILE_M * TILE_N * 4];
static thread_local int32_t TileC1[TILE_M * TILE_N * 4];
// double buffering C to interleave avx512 and amx
int32_t * C_cur = TileC0;
int32_t * C_pre = TileC1;
auto Tile4 = [&](int32_t * base) { return base; };
auto Tile5 = [&](int32_t * base) { return base + TILE_M * TILE_N; };
auto Tile6 = [&](int32_t * base) { return base + 2 * TILE_M * TILE_N; };
auto Tile7 = [&](int32_t * base) { return base + 3 * TILE_M * TILE_N; };
if (M == 2 * TILE_M) {
// i = 0
const char * B_blk0 = B + PACKED_INDEX(0, 0, KB, TILE_SIZE);
const char * B_blk1 = B + PACKED_INDEX(1, 0, KB, TILE_SIZE);
if (need_unpack) {
unpack_B<TB>(Tile0, B_blk0);
_tile_loadd(TMM0, Tile0, TILE_N * VNNI_BLK);
} else {
_tile_loadd(TMM0, B_blk0, TILE_N * VNNI_BLK);
}
_tile_zero(TMM4);
_tile_loadd(TMM2, A[0].qs, lda);
_tile_dpbssd(TMM4, TMM2, TMM0);
_tile_stored(TMM4, Tile4(C_pre), TILE_N * sizeof(int32_t));
_tile_zero(TMM5);
_tile_loadd(TMM3, A[TILE_M * KB + 0].qs, lda);
_tile_dpbssd(TMM5, TMM3, TMM0);
_tile_stored(TMM5, Tile5(C_pre), TILE_N * sizeof(int32_t));
if (need_unpack) {
unpack_B<TB>(Tile1, B_blk0);
_tile_loadd(TMM1, Tile1, TILE_N * VNNI_BLK);
} else {
_tile_loadd(TMM1, B_blk1, TILE_N * VNNI_BLK);
}
_tile_zero(TMM6);
_tile_dpbssd(TMM6, TMM2, TMM1);
_tile_stored(TMM6, Tile6(C_pre), TILE_N * sizeof(int32_t));
_tile_zero(TMM7);
_tile_dpbssd(TMM7, TMM3, TMM1);
_tile_stored(TMM7, Tile7(C_pre), TILE_N * sizeof(int32_t));
for (int i = 1; i < KB; ++i) {
// index of previous iter
const int ii = i - 1;
const char * B_blk0 = B + PACKED_INDEX(0, i, KB, TILE_SIZE);
const char * B_blk1 = B + PACKED_INDEX(1, i, KB, TILE_SIZE);
GGML_DISPATCH_BOOL(ii > 0, is_acc, [&] {
if (need_unpack) {
unpack_B<TB>(Tile0, B_blk0);
_tile_loadd(TMM0, Tile0, TILE_N * VNNI_BLK);
} else {
_tile_loadd(TMM0, B_blk0, TILE_N * VNNI_BLK);
}
_tile_zero(TMM4);
_tile_loadd(TMM2, A[i].qs, lda);
acc_C<TA, TB, is_acc>::apply(C, ldc, Tile4(C_pre), &A[ii], KB, B + PACKED_INDEX(0, ii, KB, TILE_SIZE), TILE_M);
_tile_dpbssd(TMM4, TMM2, TMM0);
_tile_stored(TMM4, Tile4(C_cur), TILE_N * sizeof(int32_t));
_tile_zero(TMM5);
_tile_loadd(TMM3, A[TILE_M * KB + i].qs, lda);
acc_C<TA, TB, is_acc>::apply(C + TILE_M * ldc, ldc, Tile5(C_pre), &A[TILE_M * KB + ii], KB, B + PACKED_INDEX(0, ii, KB, TILE_SIZE), TILE_M);
_tile_dpbssd(TMM5, TMM3, TMM0);
_tile_stored(TMM5, Tile5(C_cur), TILE_N * sizeof(int32_t));
if (need_unpack) {
unpack_B<TB>(Tile1, B_blk1);
_tile_loadd(TMM1, Tile1, TILE_N * VNNI_BLK);
} else {
_tile_loadd(TMM1, B_blk1, TILE_N * VNNI_BLK);
}
_tile_zero(TMM6);
acc_C<TA, TB, is_acc>::apply(C + TILE_N, ldc, Tile6(C_pre), &A[ii], KB, B + PACKED_INDEX(1, ii, KB, TILE_SIZE), TILE_M);
_tile_dpbssd(TMM6, TMM2, TMM1);
_tile_stored(TMM6, Tile6(C_cur), TILE_N * sizeof(int32_t));
_tile_zero(TMM7);
acc_C<TA, TB, is_acc>::apply(C + TILE_M * ldc + TILE_N, ldc, Tile7(C_pre), &A[TILE_M * KB + ii], KB, B + PACKED_INDEX(1, ii, KB, TILE_SIZE), TILE_M);
_tile_dpbssd(TMM7, TMM3, TMM1);
_tile_stored(TMM7, Tile7(C_cur), TILE_N * sizeof(int32_t));
std::swap(C_cur, C_pre);
});
}
// final accumulation
{
int ii = KB - 1;
acc_C<TA, TB, true>::apply(C, ldc, Tile4(C_pre), &A[ii], KB, B + PACKED_INDEX(0, ii, KB, TILE_SIZE), TILE_M);
acc_C<TA, TB, true>::apply(C + TILE_M * ldc, ldc, Tile5(C_pre), &A[TILE_M * KB + ii], KB, B + PACKED_INDEX(0, ii, KB, TILE_SIZE), TILE_M);
acc_C<TA, TB, true>::apply(C + TILE_N, ldc, Tile6(C_pre), &A[ii], KB, B + PACKED_INDEX(1, ii, KB, TILE_SIZE), TILE_M);
acc_C<TA, TB, true>::apply(C + TILE_M * ldc + TILE_N, ldc, Tile7(C_pre), &A[TILE_M * KB + ii], KB, B + PACKED_INDEX(1, ii, KB, TILE_SIZE), TILE_M);
}
} else {
for (int i = 0; i < KB; ++i) {
_tile_zero(TMM4);
_tile_zero(TMM6);
if (m1 != 0) {
_tile_zero(TMM5);
_tile_zero(TMM7);
}
const char * B_blk0 = B + PACKED_INDEX(0, i, KB, TILE_SIZE);
const char * B_blk1 = B + PACKED_INDEX(1, i, KB, TILE_SIZE);
if (need_unpack) {
unpack_B<TB>(Tile0, B_blk0);
_tile_loadd(TMM0, Tile0, TILE_N * VNNI_BLK);
} else {
_tile_loadd(TMM0, B_blk0, TILE_N * VNNI_BLK);
}
if (need_unpack) {
unpack_B<TB>(Tile1, B_blk1);
_tile_loadd(TMM1, Tile1, TILE_N * VNNI_BLK);
} else {
_tile_loadd(TMM1, B_blk1, TILE_N * VNNI_BLK);
}
if (m0 == TILE_M) {
_tile_loadd(TMM2, A[i].qs, lda);
} else {
unpack_A(Tile23, &A[i], KB, m0);
_tile_loadd(TMM2, Tile23, TILE_K);
}
_tile_dpbssd(TMM4, TMM2, TMM0);
_tile_dpbssd(TMM6, TMM2, TMM1);
_tile_stored(TMM4, Tile4(C_cur), TILE_N * sizeof(int32_t));
_tile_stored(TMM6, Tile6(C_cur), TILE_N * sizeof(int32_t));
GGML_DISPATCH_BOOL(i > 0, is_acc, [&] {
acc_C<TA, TB, is_acc>::apply(C, ldc, Tile4(C_cur), &A[i], KB, B + PACKED_INDEX(0, i, KB, TILE_SIZE), m0);
acc_C<TA, TB, is_acc>::apply(C + TILE_N, ldc, Tile6(C_cur), &A[i], KB, B + PACKED_INDEX(1, i, KB, TILE_SIZE), m0);
});
if (m1 != 0) {
unpack_A(Tile23, &A[TILE_M * KB + i], KB, m1);
_tile_loadd(TMM3, Tile23, TILE_K);
_tile_dpbssd(TMM5, TMM3, TMM0);
_tile_dpbssd(TMM7, TMM3, TMM1);
_tile_stored(TMM5, Tile5(C_cur), TILE_N * sizeof(int32_t));
_tile_stored(TMM7, Tile7(C_cur), TILE_N * sizeof(int32_t));
GGML_DISPATCH_BOOL(i > 0, is_acc, [&] {
acc_C<TA, TB, is_acc>::apply(C + TILE_M * ldc, ldc, Tile5(C_cur), &A[TILE_M * KB + i], KB, B + PACKED_INDEX(0, i, KB, TILE_SIZE), m1);
acc_C<TA, TB, is_acc>::apply(C + TILE_M * ldc + TILE_N, ldc, Tile7(C_cur), &A[TILE_M * KB + i], KB, B + PACKED_INDEX(1, i, KB, TILE_SIZE), m1);
});
}
}
}
return;
}
template <typename TA, typename TB, typename TC, int BLOCK_K,
typename std::enable_if<is_type_qkk<TB>::value, int>::type = 0>
void tinygemm_kernel_amx(int M, int N, int KB, const void * RESTRICT _A, const void * RESTRICT _B, float * RESTRICT C, int ldc) {
static_assert(std::is_same<TA, block_q8_K>::value);
const int TILE_SIZE = get_tile_size<TB>();
GGML_ASSERT(M <= 2 * TILE_M && N == 2 * TILE_N);
const TA * RESTRICT A = static_cast<const TA *>(_A);
const char * RESTRICT B = static_cast<const char *>(_B);
const int m0 = std::min(M, TILE_M);
const int m1 = std::max(M - TILE_M, 0);
//const int lda = KB * sizeof(TA);
static thread_local int8_t Tile0[TILE_N * TILE_K];
static thread_local int8_t Tile1[TILE_N * TILE_K];
static thread_local int8_t Tile23[TILE_M * TILE_K];
// mat mul result for each group
static thread_local int32_t Tile4[TILE_M * TILE_N];
static thread_local int32_t Tile5[TILE_M * TILE_N];
static thread_local int32_t Tile6[TILE_M * TILE_N];
static thread_local int32_t Tile7[TILE_M * TILE_N];
// sum of each QK_K block, contains 8 groups, int32
static thread_local int32_t Sumi4[TILE_M * TILE_N];
static thread_local int32_t Sumi5[TILE_M * TILE_N];
static thread_local int32_t Sumi6[TILE_M * TILE_N];
static thread_local int32_t Sumi7[TILE_M * TILE_N];
const int k_group_size = std::is_same<TB, block_q6_K>::value ? 16 : 32;
for (int i = 0; i < KB; ++i) {
// step 1: accumulate the quants across 8 groups, each group with 32
for (int k = 0; k < QK_K / k_group_size; ++k) {
GGML_DISPATCH_BOOL(k > 0, is_acc, [&] {
_tile_zero(TMM4);
_tile_zero(TMM6);
unpack_B<TB>(Tile0, B + PACKED_INDEX(0, i, KB, TILE_SIZE), k);
_tile_loadd(TMM0, Tile0, TILE_N * VNNI_BLK);
unpack_B<TB>(Tile1, B + PACKED_INDEX(1, i, KB, TILE_SIZE), k);
_tile_loadd(TMM1, Tile1, TILE_N * VNNI_BLK);
unpack_A<TB>(Tile23, &A[i], KB, k, m0);
_tile_loadd(TMM2, Tile23, TILE_K);
_tile_dpbssd(TMM4, TMM2, TMM0);
_tile_dpbssd(TMM6, TMM2, TMM1);
_tile_stored(TMM4, Tile4, TILE_N * sizeof(int32_t));
_tile_stored(TMM6, Tile6, TILE_N * sizeof(int32_t));
scale_C<TB, is_acc>(Tile4, Sumi4, B + PACKED_INDEX(0, i, KB, TILE_SIZE), k, m0);
scale_C<TB, is_acc>(Tile6, Sumi6, B + PACKED_INDEX(1, i, KB, TILE_SIZE), k, m0);
if (m1 != 0) {
_tile_zero(TMM5);
_tile_zero(TMM7);
unpack_A<TB>(Tile23, &A[TILE_M * KB + i], KB, k, m1);
_tile_loadd(TMM3, Tile23, TILE_K);
_tile_dpbssd(TMM5, TMM3, TMM0);
_tile_dpbssd(TMM7, TMM3, TMM1);
_tile_stored(TMM5, Tile5, TILE_N * sizeof(int32_t));
_tile_stored(TMM7, Tile7, TILE_N * sizeof(int32_t));
scale_C<TB, is_acc>(Tile5, Sumi5, B + PACKED_INDEX(0, i, KB, TILE_SIZE), k, m1);
scale_C<TB, is_acc>(Tile7, Sumi7, B + PACKED_INDEX(1, i, KB, TILE_SIZE), k, m1);
}
});
}
// step 2: accmulate the mins
GGML_DISPATCH_BOOL(i > 0, is_acc, [&] {
acc_C<TA, TB, is_acc>::apply(C, ldc, Sumi4, &A[i], KB, B + PACKED_INDEX(0, i, KB, TILE_SIZE), m0);
acc_C<TA, TB, is_acc>::apply(C + TILE_N, ldc, Sumi6, &A[i], KB, B + PACKED_INDEX(1, i, KB, TILE_SIZE), m0);
if (m1 != 0) {
acc_C<TA, TB, is_acc>::apply(C + TILE_M * ldc, ldc, Sumi5, &A[TILE_M * KB + i], KB, B + PACKED_INDEX(0, i, KB, TILE_SIZE), m1);
acc_C<TA, TB, is_acc>::apply(C + TILE_M * ldc + TILE_N, ldc, Sumi7, &A[TILE_M * KB + i], KB, B + PACKED_INDEX(1, i, KB, TILE_SIZE), m1);
}
});
}
return;
}
} // anonymous namespace
// get the packed tensor size for quantized weights
size_t ggml_backend_amx_get_alloc_size(const struct ggml_tensor * tensor) {
const enum ggml_type TYPE = tensor->type;
const int K = tensor->ne[0]; // ne0: in_features
const int N = tensor->ne[1]; // ne1: out_features
auto get_tensor_size = [&] {
size_t row_size_B{0};
GGML_DISPATCH_QTYPES(TYPE, [&] {
row_size_B = get_row_size<type, blck_size>(K);
});
return N * row_size_B;
};
if (qtype_has_amx_kernels(TYPE)) {
return get_tensor_size();
} else {
// for f16, bf16 we don't do packing
return ggml_nbytes(tensor);
}
}
// pack weight to vnni format
void ggml_backend_amx_convert_weight(struct ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
GGML_ASSERT(offset == 0 && size == ggml_nbytes(tensor)); // only full tensor conversion is supported for now
const enum ggml_type TYPE = tensor->type;
const int K = tensor->ne[0]; // ne0: in_features
const int N = tensor->ne[1]; // ne1: out_features
#if defined(_OPENMP)
// the buffer ctx is not initialized when .set_tensor is called
int n_threads = omp_get_num_threads();
#else
int n_threads = 1;
#endif
GGML_DISPATCH_QTYPES(TYPE, [&] {
convert_B_packed_format<type, blck_size>((void *)((char *)tensor->data + offset), (const type *)data, N, K, n_threads);
});
}
size_t ggml_backend_amx_desired_wsize(const struct ggml_tensor * dst) {
struct ggml_tensor * src0 = dst->src[0];
const enum ggml_type TYPE = src0->type;
const bool is_floating_type = TYPE == GGML_TYPE_F16;
if (is_floating_type) {
return 0;
}
const int M = dst->ne[1];
const int K = src0->ne[0];
size_t desired_wsize = 0;
GGML_DISPATCH_QTYPES(TYPE, [&] {
const size_t row_size_A = K / blck_size * sizeof(vec_dot_type);
desired_wsize = M * row_size_A;
});
return desired_wsize;
}
// NB: mixed dtype gemm with Advanced Matrix Extensions (Intel AMX)
//
// src0: weight in shape of {N, K}, quantized
// src1: input in shape of {M, K}, float32
// dst: output in shape of {M, N}, float32
//
// the function performs: dst = src1 @ src0.T
//
void ggml_backend_amx_mul_mat(const ggml_compute_params * params, struct ggml_tensor * dst) {
struct ggml_tensor * src0 = dst->src[0];
struct ggml_tensor * src1 = dst->src[1];
const enum ggml_type TYPE = src0->type;
// f16 only has avx512 kernels for now,
// amx kernels will be added once 6th gen xeon is released.
const bool is_floating_type = TYPE == GGML_TYPE_F16;
const int M = dst->ne[1];
const int N = dst->ne[0];
const int K = src0->ne[0];
const int ldc = dst->nb[1] / dst->nb[0];
if (is_floating_type) {
constexpr int BLOCK_M = 4;
constexpr int BLOCK_N = 6;
const int MB = div_up(M, BLOCK_M);
const int NB = div_up(N, BLOCK_N);
parallel_for_ggml(params, MB * NB, [&](int begin, int end) {
GGML_DISPATCH_FLOATING_TYPES(TYPE, [&] {
for (int i = begin; i < end; ++i) {
int mb = i / NB;
int nb = i % NB;
int mb_start = mb * BLOCK_M;
int mb_size = std::min(BLOCK_M, M - mb_start);
int nb_start = nb * BLOCK_N;
int nb_size = std::min(BLOCK_N, N - nb_start);
switch (mb_size << 4 | nb_size) {
case 0x12: LAUNCH_TINYGEMM_KERNEL_AVX(1, 2); break;
case 0x14: LAUNCH_TINYGEMM_KERNEL_AVX(1, 4); break;
case 0x16: LAUNCH_TINYGEMM_KERNEL_AVX(1, 6); break;
case 0x22: LAUNCH_TINYGEMM_KERNEL_AVX(2, 2); break;
case 0x24: LAUNCH_TINYGEMM_KERNEL_AVX(2, 4); break;
case 0x26: LAUNCH_TINYGEMM_KERNEL_AVX(2, 6); break;
case 0x32: LAUNCH_TINYGEMM_KERNEL_AVX(3, 2); break;
case 0x34: LAUNCH_TINYGEMM_KERNEL_AVX(3, 4); break;
case 0x36: LAUNCH_TINYGEMM_KERNEL_AVX(3, 6); break;
case 0x42: LAUNCH_TINYGEMM_KERNEL_AVX(4, 2); break;
case 0x44: LAUNCH_TINYGEMM_KERNEL_AVX(4, 4); break;
case 0x46: LAUNCH_TINYGEMM_KERNEL_AVX(4, 6); break;
default: fprintf(stderr, "Unexpected block size!\n");
}
}
});
});
return;
}
// pointer to work space, used convert A from float to quantized type
void * wdata = params->wdata;
//TODO: performance improvement: merge quant A
if (params->ith == 0) {
GGML_DISPATCH_QTYPES(TYPE, [&] {
const size_t row_size_A = K / blck_size * sizeof(vec_dot_type);
const size_t desired_wsize = M * row_size_A;
if (params->wsize < desired_wsize) {
GGML_ABORT("insufficient work space size");
}
// Q4_0, Q4_1, Q8_0 handles 1 TILE_K per blck_size
// Q4_K, Q5_K, Q6_K, IQ4_XS handles 8 TILE_K per blck_size
GGML_ASSERT(TILE_K == blck_size || TILE_K * 8 == blck_size);
const float * A_data = static_cast<const float *>(src1->data);
for (int m = 0; m < M; ++m) {
from_float<vec_dot_type>(A_data + m * K, (char *)wdata + m * row_size_A, K);
}
});
}
ggml_barrier(params->threadpool);
if (M == 1) {
// MB = 1 and handle 8 tiles in each block
constexpr int kTilesN = 4;
constexpr int BLOCK_N = TILE_N * kTilesN;
const int NB = div_up(N, BLOCK_N);
parallel_for_ggml(params, NB, [&](int begin, int end) {
GGML_DISPATCH_QTYPES(TYPE, [&] {
const int KB = K / blck_size;
const int TILE_SIZE = get_tile_size<type>();
const int row_size_A = KB * sizeof(vec_dot_type);
for (int i = begin; i < end; ++i) {
int nb = i;
int nb_start = nb * BLOCK_N;
int nb_size = std::min(BLOCK_N, N - nb_start); // 32, 64, 96
switch (nb_size) {
//case 160: LAUNCH_TINYGEMM_KERNEL_VNNI(160); break;
case 128: LAUNCH_TINYGEMM_KERNEL_VNNI(128); break;
case 96: LAUNCH_TINYGEMM_KERNEL_VNNI(96); break;
case 64: LAUNCH_TINYGEMM_KERNEL_VNNI(64); break;
case 32: LAUNCH_TINYGEMM_KERNEL_VNNI(32); break;
default: fprintf(stderr, "Unexpected n block size!\n");
}
}
});
});
return;
}
// handle 4 tiles at a tile
constexpr int BLOCK_M = TILE_M * 2;
constexpr int BLOCK_N = TILE_N * 2;
const int MB = div_up(M, BLOCK_M);
const int NB = div_up(N, BLOCK_N);
parallel_for_ggml(params, MB * NB, [&](int begin, int end) {
// init tile config for each thread
ggml_tile_config_init();
GGML_DISPATCH_QTYPES(TYPE, [&] {
const int KB = K / blck_size;
const int TILE_SIZE = get_tile_size<type>();
const int row_size_A = KB * sizeof(vec_dot_type);
for (int i = begin; i < end; ++i) {
int mb = i / NB;
int nb = i % NB;
int mb_start = mb * BLOCK_M;
int mb_size = std::min(BLOCK_M, M - mb_start);
int nb_start = nb * BLOCK_N;
int nb_size = BLOCK_N;
tinygemm_kernel_amx<vec_dot_type, type, float, blck_size>(
mb_size, nb_size, KB,
(const char *)wdata + mb_start * row_size_A,
(const char *)src0->data + PACKED_INDEX(nb * 2, 0, KB, TILE_SIZE),
(float *) dst->data + mb_start * N + nb_start, ldc);
}
});
});
}
#endif // if defined(__AMX_INT8__) && defined(__AVX512VNNI__)
llama/mmq.h 0 → 100644
View file @ 527cc978
/**
* llama.cpp - commit 40c6d79fb52f995f47507fedfeaae2ac05d9b35c - do not edit this file
*
* MIT License
*
* Copyright (c) 2023-2024 The ggml authors
*
* Permission is hereby granted, free of charge, to any person obtaining a copy
* of this software and associated documentation files (the "Software"), to deal
* in the Software without restriction, including without limitation the rights
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
* copies of the Software, and to permit persons to whom the Software is
* furnished to do so, subject to the following conditions:
*
* The above copyright notice and this permission notice shall be included in all
* copies or substantial portions of the Software.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
* SOFTWARE.
*/
#pragma once
#include "common.h"
#ifdef __cplusplus
extern "C" {
#endif
size_t ggml_backend_amx_get_alloc_size(const struct ggml_tensor * tensor);
void ggml_backend_amx_convert_weight(struct ggml_tensor * tensor, const void * data, size_t offset, size_t size);
void ggml_backend_amx_mul_mat(const struct ggml_compute_params * params, struct ggml_tensor * dst);
#ifdef __cplusplus
}
#endif
llama/patches/0001-cuda.patch
View file @ 527cc978
......@@ -4,47 +4,33 @@ Date: Thu, 6 Jun 2024 23:55:47 -0700
Subject: [PATCH] cuda
---
ggml/include/ggml-cuda.h | 2 ++
ggml/src/ggml-backend.c | 5 +++++
ggml/src/ggml-cuda.cu | 6 ++++--
3 files changed, 11 insertions(+), 2 deletions(-)
ggml/src/ggml-backend.cpp | 5 +++++
ggml/src/ggml-cuda/ggml-cuda.cu | 4 ++++
2 files changed, 9 insertions(+)
diff --git a/ggml/include/ggml-cuda.h b/ggml/include/ggml-cuda.h
index 71bb6dcf..08be0895 100644
--- a/ggml/include/ggml-cuda.h
+++ b/ggml/include/ggml-cuda.h
@@ -34,6 +34,8 @@ GGML_API GGML_CALL ggml_backend_buffer_type_t ggml_backend_cuda_split_buffer_typ
// pinned host buffer for use with the CPU backend for faster copies between CPU and GPU
GGML_API GGML_CALL ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type(void);
+GGML_API GGML_CALL int ggml_backend_cuda_reg_devices();
+
GGML_API GGML_CALL int ggml_backend_cuda_get_device_count(void);
GGML_API GGML_CALL void ggml_backend_cuda_get_device_description(int device, char * description, size_t description_size);
GGML_API GGML_CALL void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total);
diff --git a/ggml/src/ggml-backend.c b/ggml/src/ggml-backend.c
index ba280e06..d5c3fe49 100644
--- a/ggml/src/ggml-backend.c
+++ b/ggml/src/ggml-backend.c
@@ -83,7 +83,12 @@ void ggml_backend_buffer_free(ggml_backend_buffer_t buffer) {
diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp
index fdb4b986..9b80fe07 100644
--- a/ggml/src/ggml-backend.cpp
+++ b/ggml/src/ggml-backend.cpp
@@ -106,7 +106,12 @@ void ggml_backend_buffer_free(ggml_backend_buffer_t buffer) {
if (buffer->iface.free_buffer != NULL) {
buffer->iface.free_buffer(buffer);
}
+
+// TODO: this needs to be freed in cuda and hipblas backends because
+// TODO: this needs to be freed in cuda and hip backends because
+// the cuda backend implementation compiled with msvc
+#if !defined(GGML_USE_CUDA) && !defined(GGML_USE_HIPBLAS)
free(buffer);
+#if !defined(GGML_USE_CUDA) && !defined(GGML_USE_HIP)
delete buffer;
+#endif
}
size_t ggml_backend_buffer_get_size(ggml_backend_buffer_t buffer) {
diff --git a/ggml/src/ggml-cuda.cu b/ggml/src/ggml-cuda.cu
index 6efdab14..809d6ab1 100644
--- a/ggml/src/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda.cu
@@ -469,6 +469,10 @@ GGML_CALL static bool ggml_backend_buffer_is_cuda(ggml_backend_buffer_t buffer)
GGML_CALL static void ggml_backend_cuda_buffer_free_buffer(ggml_backend_buffer_t buffer) {
diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index d6e4bfdd..52aec229 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -424,6 +424,10 @@ struct ggml_backend_cuda_buffer_context {
static void ggml_backend_cuda_buffer_free_buffer(ggml_backend_buffer_t buffer) {
ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context;
delete ctx;
+
......@@ -53,13 +39,4 @@ index 6efdab14..809d6ab1 100644
+ free(buffer);
}
GGML_CALL static void * ggml_backend_cuda_buffer_get_base(ggml_backend_buffer_t buffer) {
@@ -3204,8 +3208,6 @@ GGML_CALL static ggml_backend_t ggml_backend_reg_cuda_init(const char * params,
GGML_UNUSED(params);
}
-extern "C" GGML_CALL int ggml_backend_cuda_reg_devices();
-
GGML_CALL int ggml_backend_cuda_reg_devices() {
int device_count = ggml_backend_cuda_get_device_count();
//int device_count = 1; // DEBUG: some tools require delaying CUDA initialization
static bool ggml_backend_buffer_is_cuda(ggml_backend_buffer_t buffer) {
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment