happyz
/
gemma.cpp
mirror of https://github.com/google/gemma.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Internal change. 

PiperOrigin-RevId: 765218260
This commit is contained in:
The gemma.cpp Authors 2025-05-30 09:18:08 -07:00 committed by Copybara-Service
parent 1e8642f8f4
commit 9c3e089b09
2 changed files with 44 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
gemma/configs.cc
View File

@ -572,7 +572,8 @@ ModelConfig::ModelConfig(const Model model, Type weight,
static Model FindModel(const std::string& specifier) { static Model FindModel(const std::string& specifier) {
Model found_model = Model::UNKNOWN; Model found_model = Model::UNKNOWN;
ForEachModel([&](Model model) { ForEachModel([&](Model model) {
const char* prefix = ModelPrefix(model); // Some model names are prefixes of other model names
const std::string prefix = std::string(ModelPrefix(model)) + "-";
if (specifier.rfind(prefix, 0) == 0) { // Starts with prefix. if (specifier.rfind(prefix, 0) == 0) { // Starts with prefix.
// We only expect one match. // We only expect one match.
HWY_ASSERT_M(found_model == Model::UNKNOWN, specifier.c_str()); HWY_ASSERT_M(found_model == Model::UNKNOWN, specifier.c_str());

19
gemma/gemma-inl.h
View File

@ -424,7 +424,8 @@ class GemmaAttention {
const size_t kHeadGroups = layer_config_.heads / layer_config_.kv_heads; const size_t kHeadGroups = layer_config_.heads / layer_config_.kv_heads;
// For each head (token, query), compute Q.K, softmax, and weighted V. // For each head (token, query), compute Q.K, softmax, and weighted V.
pool_.Run(0, layer_config_.heads * num_interleaved, pool_.Run(
0, layer_config_.heads * num_interleaved,
[&](uint64_t task, size_t /*thread*/) HWY_ATTR { [&](uint64_t task, size_t /*thread*/) HWY_ATTR {
const size_t head = task % layer_config_.heads; const size_t head = task % layer_config_.heads;
const size_t interleaved_idx = task / layer_config_.heads; const size_t interleaved_idx = task / layer_config_.heads;
@ -435,8 +436,7 @@ class GemmaAttention {
float* HWY_RESTRICT q = float* HWY_RESTRICT q =
activations_.q.Row(interleaved_idx) + head * q_stride_; activations_.q.Row(interleaved_idx) + head * q_stride_;
float* HWY_RESTRICT att = float* HWY_RESTRICT att = activations_.att.Row(interleaved_idx) +
activations_.att.Row(interleaved_idx) +
head * activations_.seq_len; head * activations_.seq_len;
float* HWY_RESTRICT att_out = float* HWY_RESTRICT att_out =
activations_.att_out.Row(interleaved_idx) + head * qkv_dim; activations_.att_out.Row(interleaved_idx) + head * qkv_dim;
@ -446,12 +446,10 @@ class GemmaAttention {
KVCache& kv_cache = kv_caches_[query_idx]; KVCache& kv_cache = kv_caches_[query_idx];
const size_t kv_head_offset = const size_t kv_head_offset =
layer_ * cache_layer_size_ + head_offset; layer_ * cache_layer_size_ + head_offset;
MatPtrT<float> k("k_view", MatPtrT<float> k("k_view", Extents2D(kv_cache.seq_len, qkv_dim));
Extents2D(kv_cache.seq_len, qkv_dim));
k.SetPtr(kv_cache.kv_cache.get() + kv_head_offset, k.SetPtr(kv_cache.kv_cache.get() + kv_head_offset,
/*stride=*/cache_pos_size_); /*stride=*/cache_pos_size_);
MatPtrT<float> v("v_view", MatPtrT<float> v("v_view", Extents2D(kv_cache.seq_len, qkv_dim));
Extents2D(kv_cache.seq_len, qkv_dim));
v.SetPtr(kv_cache.kv_cache.get() + kv_head_offset + qkv_dim, v.SetPtr(kv_cache.kv_cache.get() + kv_head_offset + qkv_dim,
/*stride=*/cache_pos_size_); /*stride=*/cache_pos_size_);
@ -466,8 +464,8 @@ class GemmaAttention {
last_pos = prefix_end - 1; last_pos = prefix_end - 1;
} }
SingleDotSoftmaxWeightedSum(q, k, v, att, att_out, query_scale, SingleDotSoftmaxWeightedSum(q, k, v, att, att_out, query_scale, pos,
pos, start_pos, last_pos); start_pos, last_pos);
}); });
} }
@ -1510,8 +1508,7 @@ void GenerateSingleT(const ModelStore& model,
} }
template <typename T> template <typename T>
void GenerateBatchT(const ModelStore& model, void GenerateBatchT(const ModelStore& model, const ModelWeightsPtrs<T>& weights,
const ModelWeightsPtrs<T>& weights,
const RuntimeConfig& runtime_config, const RuntimeConfig& runtime_config,
const QueriesPromptTokens& queries_prompt, const QueriesPromptTokens& queries_prompt,
const QueriesPos& queries_pos, const QueriesPos& queries_pos,