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

Minor fix to profiler zone and add comment 

PiperOrigin-RevId: 681350546
This commit is contained in:
Jan Wassenberg 2024-10-02 01:37:08 -07:00 committed by Copybara-Service
parent dc2e5f1505
commit 96d2ab7d31
2 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
gemma/gemma-inl.h
View File

@ -1175,8 +1175,8 @@ SampleFunc ChooseSampleFunc(const RuntimeConfig& runtime_config) {
// Fast path for top-1 with no accept_token.
if (kTopK == 1 && !runtime_config.accept_token) {
PROFILER_ZONE("Gen.Sample Top1");
return [](float* logits, size_t vocab_size) -> TokenAndProb {
PROFILER_ZONE("Gen.Sample Top1");
return Top1OfSoftmax(logits, vocab_size);
};
}

1
ops/ops-inl.h
View File

@ -701,6 +701,7 @@ static HWY_INLINE TokenAndProb ArgmaxAndMax(const float* HWY_RESTRICT x,
// with normalized probabilities. Only equivalent to `Softmax` + `sample_func`
// if `kTopK` == 1. This is worthwhile because `num` is typically `kVocabSize`
// == 256K, and this avoids writing and then scanning again for the max.
// However, this is not enough to make parallelization worthwhile.
static HWY_MAYBE_UNUSED TokenAndProb Top1OfSoftmax(float* HWY_RESTRICT x,
const size_t num) {
namespace hn = hwy::HWY_NAMESPACE;