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

Allow changing k parameter of `SampleTopK` as a compiler flag

This commit is contained in:
RangerUFO 2024-03-13 13:55:37 +08:00
parent 5fa2eb1a86
commit 83ec42954f
1 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
configs.h
View File

@ -23,11 +23,17 @@
#define GEMMA_MAX_SEQLEN 4096
#endif // !GEMMA_MAX_SEQLEN
// Allow changing k parameter of `SampleTopK` as a compiler flag
#ifndef GEMMA_TOPK
#define GEMMA_TOPK 1
#endif // !GEMMA_TOPK
#include <stddef.h>
namespace gcpp {
static constexpr size_t kSeqLen = GEMMA_MAX_SEQLEN;
static constexpr size_t kTopK = GEMMA_TOPK;
struct ConfigGemma7B {
static constexpr int kSeqLen = gcpp::kSeqLen;
@ -38,7 +44,7 @@ struct ConfigGemma7B {
static constexpr int kHeads = 16;
static constexpr int kKVHeads = 16; // standard MHA
static constexpr int kQKVDim = 256; // query size == key size == value size
static constexpr int kTopK = 1;
static constexpr int kTopK = gcpp::kTopK;
};
struct ConfigGemma2B {
@ -50,7 +56,7 @@ struct ConfigGemma2B {
static constexpr int kHeads = 8;
static constexpr int kKVHeads = 8; // TODO(austinvhuang): add MQA support
static constexpr int kQKVDim = 256; // query size == key size == value size
static constexpr int kTopK = 1;
static constexpr int kTopK = gcpp::kTopK;
};
} // namespace gcpp