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

7x compile time speedup: shard gemma.cc 

Use overloaded functions defined in gemma/instantiations.
Also split out activations.h.

PiperOrigin-RevId: 649053122
This commit is contained in:
Jan Wassenberg 2024-07-03 06:31:37 -07:00 committed by Copybara-Service
parent a40165dea2
commit c7c3daa624
28 changed files with 1553 additions and 985 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
BUILD.bazel
View File

@ -99,6 +99,7 @@ cc_library(
srcs = ["gemma/tokenizer.cc"], srcs = ["gemma/tokenizer.cc"],
hdrs = ["gemma/tokenizer.h"], hdrs = ["gemma/tokenizer.h"],
deps = [ deps = [
":common",
"//compression:io", "//compression:io",
"@hwy//:hwy", "@hwy//:hwy",
"@hwy//:nanobenchmark", # timer "@hwy//:nanobenchmark", # timer
@ -121,8 +122,27 @@ cc_library(
name = "gemma_lib", name = "gemma_lib",
srcs = [ srcs = [
"gemma/gemma.cc", "gemma/gemma.cc",
"gemma/instantiations/27b_bf16.cc",
"gemma/instantiations/27b_f32.cc",
"gemma/instantiations/27b_sfp.cc",
"gemma/instantiations/2b_bf16.cc",
"gemma/instantiations/2b_f32.cc",
"gemma/instantiations/2b_sfp.cc",
"gemma/instantiations/7b_bf16.cc",
"gemma/instantiations/7b_f32.cc",
"gemma/instantiations/7b_sfp.cc",
"gemma/instantiations/9b_bf16.cc",
"gemma/instantiations/9b_f32.cc",
"gemma/instantiations/9b_sfp.cc",
"gemma/instantiations/tiny_bf16.cc",
"gemma/instantiations/tiny_f32.cc",
"gemma/instantiations/tiny_sfp.cc",
"gemma/instantiations/gr2b_bf16.cc",
"gemma/instantiations/gr2b_f32.cc",
"gemma/instantiations/gr2b_sfp.cc",
], ],
hdrs = [ hdrs = [
"gemma/activations.h",
"gemma/gemma.h", "gemma/gemma.h",
], ],
exec_properties = { exec_properties = {
@ -130,7 +150,7 @@ cc_library(
"mem": "28g", "mem": "28g",
}, },
textual_hdrs = [ textual_hdrs = [
# Placeholder for internal file1, do not remove, "gemma/gemma-inl.h",
# Placeholder for internal file2, do not remove, # Placeholder for internal file2, do not remove,
], ],
deps = [ deps = [

20
CMakeLists.txt
View File

@ -63,13 +63,33 @@ set(SOURCES
backprop/optimizer.h backprop/optimizer.h
evals/cross_entropy.cc evals/cross_entropy.cc
evals/cross_entropy.h evals/cross_entropy.h
gemma/activations.h
gemma/benchmark_helper.cc gemma/benchmark_helper.cc
gemma/benchmark_helper.h gemma/benchmark_helper.h
gemma/common.cc gemma/common.cc
gemma/common.h gemma/common.h
gemma/configs.h gemma/configs.h
gemma/gemma-inl.h
gemma/gemma.cc gemma/gemma.cc
gemma/gemma.h gemma/gemma.h
gemma/instantiations/27b_bf16.cc
gemma/instantiations/27b_f32.cc
gemma/instantiations/27b_sfp.cc
gemma/instantiations/2b_bf16.cc
gemma/instantiations/2b_f32.cc
gemma/instantiations/2b_sfp.cc
gemma/instantiations/7b_bf16.cc
gemma/instantiations/7b_f32.cc
gemma/instantiations/7b_sfp.cc
gemma/instantiations/9b_bf16.cc
gemma/instantiations/9b_f32.cc
gemma/instantiations/9b_sfp.cc
gemma/instantiations/gr2b_bf16.cc
gemma/instantiations/gr2b_f32.cc
gemma/instantiations/gr2b_sfp.cc
gemma/instantiations/tiny_bf16.cc
gemma/instantiations/tiny_f32.cc
gemma/instantiations/tiny_sfp.cc
gemma/kv_cache.cc gemma/kv_cache.cc
gemma/kv_cache.h gemma/kv_cache.h
gemma/ops.h gemma/ops.h

93
gemma/activations.h Normal file
View File

@ -0,0 +1,93 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef THIRD_PARTY_GEMMA_CPP_GEMMA_ACTIVATIONS_H_
#define THIRD_PARTY_GEMMA_CPP_GEMMA_ACTIVATIONS_H_
#include <stddef.h>
#include <array>
#include "gemma/common.h" // AllocateSizeof
#include "hwy/base.h" // hwy::bfloat16_t
namespace gcpp {
// Must be aligned.
template <class TConfig, size_t kBatchSize>
struct Activations {
static constexpr size_t kModelDim = TConfig::kModelDim;
static constexpr size_t kQKVDim = TConfig::kQKVDim;
static constexpr size_t kHeads = TConfig::kHeads;
static constexpr size_t kKVHeads = TConfig::kKVHeads;
static constexpr bool kIsMHA = kHeads == kKVHeads; // Multi-Head Attention
// Stride between subsequent queries. Each of Q, K, V are of length kQKVDim,
// but for MHA we store them as Q,K,V, Q,K,V, .. instead of Q..Q, K..K, V..V.
static constexpr size_t kQStride = kQKVDim * (kIsMHA ? 3 : 1);
std::array<float, kBatchSize * kModelDim> x; // input
std::array<float, kBatchSize * kModelDim> pre_att_rms_out;
std::array<float, kBatchSize * kHeads * kQStride> q; // query vector
std::array<float, kBatchSize * kHeads * TConfig::kSeqLen>
att; // attention vector
std::array<float, kBatchSize * kHeads * kQKVDim> att_out; // attention output
std::array<float, kHeads * kBatchSize * kModelDim>
att_post1; // attention output after linear transformation, per head
std::array<float, kBatchSize * kModelDim>
att_post2; // accumulation of attention outputs over heads
std::array<hwy::bfloat16_t, kBatchSize * kModelDim> bf_pre_ffw_rms_out;
std::array<float, kBatchSize * TConfig::kFFHiddenDim * 2> ffw_hidden;
// For FFW MatMul.
std::array<float, kBatchSize * TConfig::kFFHiddenDim> C1;
std::array<float, kBatchSize * TConfig::kFFHiddenDim> C2;
std::array<float, kBatchSize * kModelDim> ffw_out;
std::array<float, kBatchSize * TConfig::kVocabSize> logits;
// For bf16/f32 vectors * bf16 matrix: faster to unpack once beforehand, into
// per-thread storage.
std::array<float, kModelDim * kMaxThreads> even_odd;
// Griffin layer internal activations
static constexpr size_t kGriffinDim =
TConfig::kGriffinLayers > 0 ? kModelDim : 0;
std::array<float, kBatchSize * kGriffinDim> griffin_x;
std::array<float, kBatchSize * kGriffinDim> griffin_y;
std::array<float, kBatchSize * kGriffinDim> griffin_gate_x;
std::array<float, kBatchSize * kGriffinDim> griffin_multiplier;
};
template <typename TConfig>
struct AllocateState {
void operator()(ByteStorageT& prefill, ByteStorageT& decode) const {
// When batching queries, the prefill batch size is reduced by a factor
// of kBatchedQueryBatchSize
prefill =
AllocateSizeof<Activations<TConfig, kMinAdjustedPrefillBatchSize *
kBatchedQueryBatchSize>>();
decode = AllocateSizeof<
Activations<TConfig, kDecodeBatchSize * kBatchedQueryBatchSize>>();
}
};
template <class TConfig, size_t kBatchSize>
Activations<TConfig, kBatchSize>& GetActivations(const ByteStorageT& state_u8) {
return *reinterpret_cast<Activations<TConfig, kBatchSize>*>(state_u8.get());
}
} // namespace gcpp
#endif // THIRD_PARTY_GEMMA_CPP_GEMMA_ACTIVATIONS_H_

14
gemma/common.cc
View File

@ -53,7 +53,8 @@ constexpr ModelTraining kModelTraining[] = {
ModelTraining::GEMMA_IT, // Gemma Tiny ModelTraining::GEMMA_IT, // Gemma Tiny
}; };
constexpr size_t kNumModelFlags = std::end(kModelFlags) - std::begin(kModelFlags); constexpr size_t kNumModelFlags =
std::end(kModelFlags) - std::begin(kModelFlags);
static_assert(kNumModelFlags == static_assert(kNumModelFlags ==
std::end(kModelTypes) - std::begin(kModelTypes)); std::end(kModelTypes) - std::begin(kModelTypes));
static_assert(kNumModelFlags == static_assert(kNumModelFlags ==
@ -123,4 +124,15 @@ const char* ParseType(const std::string& type_string, Type& type) {
return kErrorMessageBuffer; return kErrorMessageBuffer;
} }
void Wrap(const ModelInfo& info, size_t pos, std::string& prompt) {
// Instruction-tuned models are trained to expect control tokens.
if (info.training == ModelTraining::GEMMA_IT) {
// Prepend "<end_of_turn>" if this is a multi-turn dialogue continuation.
const std::string start = (pos == 0)
? "<start_of_turn>user\n"
: "<end_of_turn>\n<start_of_turn>user\n";
prompt = start + prompt + "<end_of_turn>\n<start_of_turn>model\n";
}
}
} // namespace gcpp } // namespace gcpp

46
gemma/common.h
View File

@ -42,7 +42,8 @@ constexpr size_t kBatchedQueryBatchSize = 16;
constexpr size_t kMinAdjustedPrefillBatchSize = constexpr size_t kMinAdjustedPrefillBatchSize =
HWY_MAX((size_t)1, kPrefillBatchSize / kBatchedQueryBatchSize); HWY_MAX((size_t)1, kPrefillBatchSize / kBatchedQueryBatchSize);
// Model variants: see configs.h for details. // Model variants: see configs.h for details. When adding a new one, also
// update GEMMA_FOREACH* and Call* below, and add instantiations/*.cc.
enum class Model { enum class Model {
GEMMA_2B, GEMMA_2B,
GEMMA_7B, GEMMA_7B,
@ -55,16 +56,29 @@ enum class Model {
// Instruction-tuned models require extra 'turn structure' tokens in prompts. // Instruction-tuned models require extra 'turn structure' tokens in prompts.
enum class ModelTraining { GEMMA_IT, GEMMA_PT }; enum class ModelTraining { GEMMA_IT, GEMMA_PT };
// Tensor types for loading weights. // Tensor types for loading weights. When adding a new one, also
// update GEMMA_FOREACH* and Call* below, and add instantiations/*.cc.
enum class Type { kF32, kBF16, kSFP }; enum class Type { kF32, kBF16, kSFP };
// TODO(janwas): merge with parser/ToString. // TODO(janwas): merge with functions below.
struct ModelInfo { struct ModelInfo {
Model model; Model model;
ModelTraining training; ModelTraining training;
Type weight; Type weight;
}; };
// Returns error string or nullptr if OK.
// Thread-hostile.
const char* ParseModelTypeAndTraining(const std::string& model_flag,
Model& model, ModelTraining& training);
const char* ParseType(const std::string& type_string, Type& type);
// Inverse of ParseModelTypeAndTraining.
const char* ModelString(Model model, ModelTraining training);
const char* StringFromType(Type type);
void Wrap(const ModelInfo& info, size_t pos, std::string& prompt);
// Returns the return value of FuncT<Config*<TWeight>>().operator()(args), where // Returns the return value of FuncT<Config*<TWeight>>().operator()(args), where
// Config* is selected via `model`. Typically called by CallForModelAndWeight, // Config* is selected via `model`. Typically called by CallForModelAndWeight,
// but can also be called directly when FuncT does not actually use TWeight. // but can also be called directly when FuncT does not actually use TWeight.
@ -122,6 +136,20 @@ decltype(auto) CallForModelAndWeight(Model model, Type weight,
} }
} }
#define GEMMA_FOREACH_WEIGHT(X, CONFIGT) \
X(CONFIGT, float) \
X(CONFIGT, hwy::bfloat16_t) \
X(CONFIGT, SfpStream)
#define GEMMA_FOREACH_CONFIG_AND_WEIGHT(X) \
GEMMA_FOREACH_WEIGHT(X, ConfigGemmaTiny) \
GEMMA_FOREACH_WEIGHT(X, ConfigGemma2B) \
GEMMA_FOREACH_WEIGHT(X, ConfigGemma7B) \
GEMMA_FOREACH_WEIGHT(X, ConfigGemma9B) \
GEMMA_FOREACH_WEIGHT(X, ConfigGemma27B) \
GEMMA_FOREACH_WEIGHT(X, ConfigGriffin2B) \
static_assert(true, "Allow trailing ;")
// Used by GEMMA_EXPORT_AND_DISPATCH. For a given TWEIGHT (e.g. float), // Used by GEMMA_EXPORT_AND_DISPATCH. For a given TWEIGHT (e.g. float),
// calls FUNC<ConfigT<TWEIGHT>> where ConfigT is chosen via MODEL enum. // calls FUNC<ConfigT<TWEIGHT>> where ConfigT is chosen via MODEL enum.
#define GEMMA_DISPATCH_MODEL(MODEL, TWEIGHT, FUNC, ARGS) \ #define GEMMA_DISPATCH_MODEL(MODEL, TWEIGHT, FUNC, ARGS) \
@ -163,6 +191,8 @@ decltype(auto) CallForModelAndWeight(Model model, Type weight,
// Like CallForModelAndWeight, but for SIMD function templates. This is a macro // Like CallForModelAndWeight, but for SIMD function templates. This is a macro
// because it boils down to N_SSE4::FUNC, which would not work if FUNC was a // because it boils down to N_SSE4::FUNC, which would not work if FUNC was a
// normal function argument. MODEL and WEIGHT are enums. // normal function argument. MODEL and WEIGHT are enums.
// For gemma.cc, we use overloaded extern functions for faster builds. However,
// this is still used in compress_weights because its compile time is OK.
#define GEMMA_EXPORT_AND_DISPATCH(MODEL, WEIGHT, FUNC, ARGS) \ #define GEMMA_EXPORT_AND_DISPATCH(MODEL, WEIGHT, FUNC, ARGS) \
switch (WEIGHT) { \ switch (WEIGHT) { \
case Type::kF32: \ case Type::kF32: \
@ -178,16 +208,6 @@ decltype(auto) CallForModelAndWeight(Model model, Type weight,
HWY_ABORT("Weight type %d unknown.", static_cast<int>(WEIGHT)); \ HWY_ABORT("Weight type %d unknown.", static_cast<int>(WEIGHT)); \
} }
// Returns error string or nullptr if OK.
// Thread-hostile.
const char* ParseModelTypeAndTraining(const std::string& model_flag,
Model& model, ModelTraining& training);
const char* ParseType(const std::string& type_string, Type& type);
// Inverse of ParseModelTypeAndTraining.
const char* ModelString(Model model, ModelTraining training);
const char* StringFromType(Type type);
// ---------------------------------------------------------------------------- // ----------------------------------------------------------------------------
// //

906
gemma/gemma-inl.h Normal file
View File

@ -0,0 +1,906 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// SIMD functions for Gemma/Griffin transformers.
// Include guard (still compiled once per target)
#if defined(THIRD_PARTY_GEMMA_CPP_GEMMA_GEMMA_INL_H_) == \
defined(HWY_TARGET_TOGGLE)
#ifdef THIRD_PARTY_GEMMA_CPP_GEMMA_GEMMA_INL_H_
#undef THIRD_PARTY_GEMMA_CPP_GEMMA_GEMMA_INL_H_
#else
#define THIRD_PARTY_GEMMA_CPP_GEMMA_GEMMA_INL_H_
#endif
#include <stddef.h>
#include <stdio.h>
#include <string.h> // memcpy
#include <algorithm>
#include <string>
#include <vector>
#include "gemma/activations.h"
#include "gemma/common.h"
#include "gemma/gemma.h"
#include "gemma/ops.h"
#include "gemma/weights.h"
// Placeholder for internal test4, do not remove
#include "hwy/aligned_allocator.h"
#include "hwy/base.h"
#include "hwy/contrib/matvec/matvec-inl.h"
#include "hwy/contrib/thread_pool/thread_pool.h"
#include "hwy/highway.h"
#include "hwy/profiler.h"
#include "hwy/timer.h"
#ifndef GEMMA_CONFIG
#if HWY_IDE
// Provide a definition so the IDE does not complain.
#define GEMMA_CONFIG ConfigGemmaTiny<float>
#else
#error "Only include from instantiations/*.cc, which must define GEMMA_CONFIG"
#endif // HWY_IDE
#endif // GEMMA_CONFIG
HWY_BEFORE_NAMESPACE();
namespace gcpp {
namespace HWY_NAMESPACE {
template <class TConfig, size_t kBatchSize, size_t kQueryBatchSize>
HWY_NOINLINE void GriffinRecurrent(
size_t batch_start, size_t num_tokens, size_t num_queries, size_t layer,
Activations<TConfig, kBatchSize * kQueryBatchSize>& activations,
const CompressedLayer<TConfig>* layer_weights,
const std::vector<KVCache*>& kv_caches, hwy::ThreadPool& pool) {
PROFILER_ZONE("Gen.Griffin");
static_assert(kQueryBatchSize == 1,
"Griffin does not support batched queries.");
HWY_DASSERT(num_queries == 1); // TODO: add batch query support for Griffin.
KVCache& kv_cache = *kv_caches[0];
namespace hn = hwy::HWY_NAMESPACE;
using D = hn::ScalableTag<float>;
HWY_DASSERT(num_tokens <= kBatchSize);
static constexpr size_t kModelDim =
gcpp::Activations<TConfig, kBatchSize * kQueryBatchSize>::kModelDim;
static constexpr size_t kConv1dWidth = TConfig::kConv1dWidth;
static constexpr size_t kHeads = TConfig::kHeads;
// X / Y linear layers.
for (size_t batch_idx = 0; batch_idx < num_tokens; ++batch_idx) {
const size_t batch_offset = batch_idx * kModelDim;
float* HWY_RESTRICT y = activations.griffin_y.data() + batch_offset;
float* HWY_RESTRICT x = activations.griffin_x.data() + batch_offset;
TwoMatVecAdd<kModelDim, kModelDim>(
layer_weights->griffin.linear_x_w, layer_weights->griffin.linear_y_w, 0,
activations.pre_att_rms_out.data() + batch_offset,
/*add0=*/layer_weights->griffin.linear_x_biases.data(),
/*add1=*/layer_weights->griffin.linear_y_biases.data(), /*out0=*/x,
/*out1=*/y, pool);
Gelu(y, kModelDim);
}
// Conv1D.
for (size_t batch_idx = 0; batch_idx < num_tokens; ++batch_idx) {
const size_t batch_offset = batch_idx * kModelDim;
const size_t pos = batch_start + batch_idx;
float* HWY_RESTRICT x = activations.griffin_x.data() + batch_offset;
HWY_FULL(float) df;
HWY_DASSERT(kModelDim % hn::Lanes(df) == 0);
const size_t layer_offset = layer * kModelDim * (kConv1dWidth - 1);
// cache[i] = input at time t-i.
float* HWY_RESTRICT cache[HWY_MAX(kConv1dWidth, 1)];
cache[0] = x;
for (size_t i = 1; i < kConv1dWidth; i++) {
cache[i] =
kv_cache.conv1d_cache.get() + layer_offset +
((pos + kConv1dWidth - 1 - i) % (kConv1dWidth - 1)) * kModelDim;
}
for (size_t i = 0; i < kModelDim; i += hn::Lanes(df)) {
auto xv = hn::Load(df, x + i);
auto accum0 =
hn::Load(df, layer_weights->griffin.conv_biases.data() + i);
auto accum1 = hn::Zero(df);
static_assert(kConv1dWidth % 2 == 0, "Conv width must be even");
for (size_t l = 0; 2 * l < kConv1dWidth; l++) {
auto wv0 = hn::Load(df, layer_weights->griffin.conv_w.data() +
(kConv1dWidth - 1 - 2 * l) * kModelDim + i);
auto wv1 = hn::Load(df, layer_weights->griffin.conv_w.data() +
(kConv1dWidth - 2 - 2 * l) * kModelDim + i);
accum0 = hn::MulAdd(wv0, hn::Load(df, cache[l * 2] + i), accum0);
accum1 = hn::MulAdd(wv1, hn::Load(df, cache[l * 2 + 1] + i), accum1);
}
hn::Store(hn::Add(accum0, accum1), df, x + i);
hn::Store(xv, df, cache[HWY_MAX(kConv1dWidth, 1) - 1] + i);
}
}
// RGLRU
for (size_t batch_idx = 0; batch_idx < num_tokens; ++batch_idx) {
const size_t batch_offset = batch_idx * kModelDim;
const size_t pos = batch_start + batch_idx;
float* HWY_RESTRICT y = activations.griffin_y.data() + batch_offset;
float* HWY_RESTRICT x = activations.griffin_x.data() + batch_offset;
float* HWY_RESTRICT gate_x =
activations.griffin_gate_x.data() + batch_offset;
float* HWY_RESTRICT a =
activations.griffin_multiplier.data() + batch_offset;
float* HWY_RESTRICT rnn_state =
kv_cache.rglru_cache.get() + layer * kModelDim;
pool.Run(0, kHeads, [&](const uint64_t head, size_t /*thread*/) HWY_ATTR {
constexpr size_t kHeadDim = kModelDim / kHeads;
constexpr size_t kMatrixSize = kHeadDim * kHeadDim;
size_t head_offset = head * kHeadDim;
TwoOfsMatVecAddLoop<kHeadDim, kHeadDim>(
layer_weights->griffin.gate_w, kMatrixSize * head,
kMatrixSize * (kHeads + head), x + head_offset,
/*add0=*/layer_weights->griffin.gate_biases.data() + head_offset,
/*add1=*/layer_weights->griffin.gate_biases.data() + kModelDim +
head_offset,
/*out0=*/gate_x + head_offset, /*out1=*/a + head_offset);
Sigmoid(gate_x + head_offset, kHeadDim);
Sigmoid(a + head_offset, kHeadDim);
const auto fn_mul = [](D d, hn::Vec<D> x, hn::Vec<D> gate_x)
HWY_ATTR { return hn::Mul(x, gate_x); };
hn::Transform1(D(), a + head_offset, kHeadDim,
layer_weights->griffin.a.data() + head_offset, fn_mul);
hn::Transform1(D(), x + head_offset, kHeadDim, gate_x + head_offset,
fn_mul);
// RNN scan
HWY_FULL(float) df;
HWY_DASSERT(kHeadDim % hn::Lanes(df) == 0);
for (size_t i = 0; i < kHeadDim; i += hn::Lanes(df)) {
auto log_a = hn::Load(df, a + head_offset + i);
auto gated_x = hn::Load(df, x + head_offset + i);
auto rnn = hn::Load(df, rnn_state + head_offset + i);
auto a = hn::Exp(df, log_a);
auto x_multiplier = hn::Sqrt(hn::NegMulAdd(a, a, hn::Set(df, 1.0f)));
if (pos == 0) {
x_multiplier = hn::Set(df, 1.0f);
}
auto new_x = hn::MulAdd(x_multiplier, gated_x, hn::Mul(a, rnn));
hn::Store(new_x, df, rnn_state + head_offset + i);
// Join branches.
auto yv = hn::Load(df, y + head_offset + i);
auto pre_out = hn::Mul(yv, new_x);
hn::Store(pre_out, df, x + head_offset + i);
}
});
}
// Final linear layer.
for (size_t batch_idx = 0; batch_idx < num_tokens; ++batch_idx) {
const size_t batch_offset = batch_idx * kModelDim;
float* HWY_RESTRICT x = activations.griffin_x.data() + batch_offset;
float* out_ptr = activations.att_post2.data() + batch_idx * kModelDim;
MatVecAdd<kModelDim, kModelDim>(
layer_weights->griffin.linear_out_w, 0, x,
layer_weights->griffin.linear_out_biases.data(),
activations.even_odd.data(), out_ptr, pool);
}
}
template <class TConfig, size_t kBatchSize, size_t kQueryBatchSize>
HWY_NOINLINE void Attention(
size_t batch_and_query_start, size_t num_tokens, size_t num_queries,
size_t layer,
Activations<TConfig, kBatchSize * kQueryBatchSize>& activations,
const CompressedLayer<TConfig>* layer_weights,
const std::vector<KVCache*>& kv_caches,
hwy::ThreadPool& pool) {
PROFILER_ZONE("Gen.Attention");
HWY_DASSERT(num_tokens <= kBatchSize);
HWY_DASSERT(num_queries <= kQueryBatchSize);
HWY_DASSERT(batch_and_query_start % num_queries == 0);
using TActivations = Activations<TConfig, kBatchSize * kQueryBatchSize>;
constexpr size_t kQKVDim = TActivations::kQKVDim;
constexpr size_t kQStride = TActivations::kQStride;
constexpr size_t kCachePosSize = CachePosSize<TConfig>()();
constexpr size_t kCacheLayerSize = CacheLayerSize<TConfig>()();
constexpr size_t kModelDim = TActivations::kModelDim;
constexpr size_t kHeads = TConfig::kHeads;
constexpr size_t kKVHeads = TConfig::kKVHeads;
constexpr size_t kSeqLen = TConfig::kSeqLen;
GEMMA_CONSTEXPR_SQRT const float kQueryScale =
1.0f / Sqrt(static_cast<float>(kQKVDim));
constexpr bool kIsMHA = TActivations::kIsMHA; // Multi-Head Attention
const size_t batch_start = batch_and_query_start / num_queries;
const size_t num_tokens_and_queries = num_tokens * num_queries;
// If MHA, this also computes KV, which we copy to the KV cache below.
static_assert(!kIsMHA || TConfig::kInterleaveQKV); // MHA => interleaved
MatMul_4x4_Batch<kModelDim, kHeads * kQStride>(
num_tokens_and_queries, activations.pre_att_rms_out.data(),
layer_weights->qkv_einsum_w.data(), activations.q.data(), pool);
for (size_t batch_and_query_idx = 0;
batch_and_query_idx < num_tokens_and_queries; ++batch_and_query_idx) {
const float* x = activations.pre_att_rms_out.data() + batch_and_query_idx
* kModelDim;
const size_t query_idx = batch_and_query_idx % num_queries;
const size_t batch_idx = batch_and_query_idx / num_queries;
KVCache& kv_cache = *kv_caches[query_idx];
// QKV projections:
if constexpr (!kIsMHA) {
const size_t pos = batch_start + batch_idx;
const size_t cache_pos = pos % (kSeqLen + kPrefillBatchSize);
const size_t kv_offset =
cache_pos * kCachePosSize + layer * kCacheLayerSize;
float* HWY_RESTRICT kv = kv_cache.kv_cache.get() + kv_offset;
// TODO: requires MatMul support for offsets.
MatVec<kKVHeads * kQKVDim * 2, kModelDim>(
layer_weights->qkv_einsum_w, kHeads * kQKVDim * kModelDim, x,
activations.even_odd.data(), kv, pool);
}
}
// Positional encodings for kv:
pool.Run(
0, kKVHeads * num_tokens_and_queries,
[&](uint64_t task, size_t thread) HWY_ATTR {
const size_t head = task % kKVHeads;
const size_t batch_and_query_idx = task / kKVHeads;
const size_t query_idx = batch_and_query_idx % num_queries;
const size_t batch_idx = batch_and_query_idx / num_queries;
const size_t pos = batch_start + batch_idx;
const size_t cache_pos = pos % (kSeqLen + kPrefillBatchSize);
const size_t kv_offset = cache_pos * kCachePosSize +
layer * kCacheLayerSize + head * kQKVDim * 2;
KVCache& kv_cache = *kv_caches[query_idx];
float* HWY_RESTRICT kv = kv_cache.kv_cache.get() + kv_offset;
if constexpr (kIsMHA) {
// For MHA, copy kv into the KV cache from scratch space (see above).
const float* HWY_RESTRICT q =
activations.q.data() + (batch_and_query_idx * kHeads
+ head) * kQStride;
// Skip past the Q part of `q`, and copy KV to `kv`.
memcpy(kv, q + kQKVDim, 2 * kQKVDim * sizeof(float));
}
Rope(kv, TConfig::kUseHalfRope ? kQKVDim / 2 : kQKVDim, pos);
});
static_assert((kHeads % kKVHeads) == 0,
"query heads must be a multiple of key-value heads");
constexpr size_t kGroupHeads = kHeads / kKVHeads;
pool.Run(0, kHeads * num_tokens_and_queries,
[&](uint64_t task, size_t thread) HWY_ATTR {
const size_t head = task % kHeads;
const size_t batch_and_query_idx = task / kHeads;
const size_t query_idx = batch_and_query_idx % num_queries;
const size_t batch_idx = batch_and_query_idx / num_queries;
const size_t head_offset = (head / kGroupHeads) * kQKVDim * 2;
KVCache& kv_cache = *kv_caches[query_idx];
float* HWY_RESTRICT q =
activations.q.data() + (batch_and_query_idx * kHeads + head) * kQStride;
const size_t pos = batch_start + batch_idx;
// Calculate scores
float* HWY_RESTRICT head_att =
activations.att.data() + head * kSeqLen
+ batch_and_query_idx * kHeads * kSeqLen;
Rope(q, TConfig::kUseHalfRope ? kQKVDim / 2 : kQKVDim, pos);
MulByConst(kQueryScale, q, kQKVDim);
// Compute Q dot K scores
const size_t start_pos =
pos - std::min(TConfig::kAttentionWindowSizes[layer] - 1, pos);
for (size_t pos2 = start_pos; pos2 <= pos; ++pos2) {
const size_t cache_pos = pos2 % (kSeqLen + kPrefillBatchSize);
const size_t kv_offset =
cache_pos * kCachePosSize + layer * kCacheLayerSize + head_offset;
const float* HWY_RESTRICT k2 = kv_cache.kv_cache.get() + kv_offset;
const float score = Dot(q, k2, kQKVDim);
head_att[pos2 % kSeqLen] = score;
}
const size_t head_att_len = std::min(pos + 1, kSeqLen);
if constexpr (TConfig::kAttCap > 0.0f) {
LogitsSoftCap(TConfig::kAttCap, head_att, head_att_len);
}
Softmax(head_att, head_att_len);
// Weighted summation
float* HWY_RESTRICT att_out = activations.att_out.data() + head * kQKVDim +
batch_and_query_idx * kHeads * kQKVDim;
hwy::ZeroBytes(att_out, kQKVDim * sizeof(*att_out));
for (size_t pos2 = start_pos; pos2 <= pos; ++pos2) {
const size_t cache_pos = pos2 % (kSeqLen + kPrefillBatchSize);
const size_t kv_offset =
cache_pos * kCachePosSize + layer * kCacheLayerSize + head_offset;
float* HWY_RESTRICT v2 = kv_cache.kv_cache.get() + kv_offset + kQKVDim;
MulByConstAndAdd(head_att[pos2 % kSeqLen], v2, att_out, kQKVDim);
}
});
for (size_t batch_and_query_idx = 0;
batch_and_query_idx < num_tokens_and_queries; ++batch_and_query_idx) {
// TODO(szabadka) Use a single MatVecAdd like in GriffinRecurrent() after
// rearranging the weights.
float* HWY_RESTRICT att_out =
activations.att_out.data() + batch_and_query_idx * kHeads * kQKVDim;
float* HWY_RESTRICT layer_out =
activations.att_post2.data() + batch_and_query_idx * kModelDim;
MatVecT</*kAdd=*/TConfig::kSoftmaxAttnOutputBiases, kModelDim, kQKVDim>(
layer_weights->attn_vec_einsum_w, 0, att_out,
layer_weights->attention_output_biases.data(),
activations.even_odd.data(), layer_out, pool);
for (size_t head = 1; head < kHeads; ++head) {
// TODO(patrickms): Check this calculation
float* HWY_RESTRICT head_out =
activations.att_post1.data() +
head * kBatchSize * kQueryBatchSize * kModelDim;
// TODO: requires MatMul support for offsets.
MatVec<kModelDim, kQKVDim>(
layer_weights->attn_vec_einsum_w, head * kModelDim * kQKVDim,
att_out + head * kQKVDim,
activations.even_odd.data(), head_out, pool);
AddFrom(head_out, layer_out, kModelDim);
}
}
}
template <class TConfig, size_t kBatchSize>
HWY_NOINLINE void FFW(Activations<TConfig, kBatchSize>& activations,
size_t num_tokens,
const CompressedLayer<TConfig>* layer_weights,
hwy::ThreadPool& pool) {
HWY_DASSERT(num_tokens <= kBatchSize);
constexpr size_t kModelDim = TConfig::kModelDim;
constexpr size_t kFFHiddenDim = TConfig::kFFHiddenDim;
float* HWY_RESTRICT even_odd = activations.even_odd.data();
// TODO: MatMul does not yet support adding another matrix to the result.
if constexpr (!TConfig::kFFBiases) {
PROFILER_ZONE("Gen.FFW.GatedGELU");
// MatMul expects col-major B, which is what we have: kModelDim consecutive
// elements in memory, repeated kFFHiddenDim times.
const auto b1 = layer_weights->gating_einsum_w.data();
constexpr size_t kColsA = kModelDim;
constexpr size_t kColsB = kFFHiddenDim;
const auto b2 = b1 + kColsA * kColsB;
auto A = activations.bf_pre_ffw_rms_out.data();
// Will go through GELU.
MatMul_4x4_Batch<kColsA, kColsB>(num_tokens, A, b1, activations.C1.data(),
pool);
// What to multiply by.
MatMul_4x4_Batch<kColsA, kColsB>(num_tokens, A, b2, activations.C2.data(),
pool);
// Gelu and multiply by gate.
namespace hn = hwy::HWY_NAMESPACE;
using DF = hn::ScalableTag<float>;
using VF = hn::Vec<DF>;
hn::Transform1(DF(), activations.C1.data(), kFFHiddenDim * num_tokens,
activations.C2.data(), [](DF df, VF v, VF mul) HWY_ATTR {
return hn::Mul(mul, Gelu(df, v));
});
MatMul_4x4_Batch<kFFHiddenDim, kModelDim>(num_tokens, activations.C1.data(),
layer_weights->linear_w.data(),
activations.ffw_out.data(), pool);
} else { // TConfig::kFFBiases == true
for (size_t batch_idx = 0; batch_idx < num_tokens; ++batch_idx) {
const size_t hidden_offset = batch_idx * kFFHiddenDim * 2;
const hwy::bfloat16_t* HWY_RESTRICT vec =
activations.bf_pre_ffw_rms_out.data() + batch_idx * kModelDim;
float* HWY_RESTRICT out = activations.ffw_hidden.data() + hidden_offset;
float* HWY_RESTRICT out_mul = out + kFFHiddenDim;
PROFILER_ZONE("Gen.FFW.GatedGELU");
// Same matrix, first and second half of rows. Could fuse into one MatVec.
MatVecT</*kAdd=*/true, kFFHiddenDim, kModelDim>(
layer_weights->gating_einsum_w, kFFHiddenDim * kModelDim, vec,
layer_weights->ffw_gating_biases.data() + kFFHiddenDim, even_odd,
out_mul, pool);
// Gate, will go through the nonlinearity.
MatVecT</*kAdd=*/true, kFFHiddenDim, kModelDim>(
layer_weights->gating_einsum_w, 0, vec,
layer_weights->ffw_gating_biases.data(), even_odd, out, pool);
namespace hn = hwy::HWY_NAMESPACE;
using DF = hn::ScalableTag<float>;
using VF = hn::Vec<DF>;
hn::Transform1(DF(), out, kFFHiddenDim, out_mul,
[](DF df, VF v, VF mul)
HWY_ATTR { return hn::Mul(mul, Gelu(df, v)); });
MatVecT</*kAdd=*/true, kModelDim, kFFHiddenDim>(
layer_weights->linear_w, 0,
activations.ffw_hidden.data() + hidden_offset,
layer_weights->ffw_output_biases.data(), even_odd,
activations.ffw_out.data() + batch_idx * kModelDim, pool);
}
}
}
template <class TConfig, size_t kBatchSize>
HWY_NOINLINE void EmbedToken(int token, size_t token_idx, size_t pos,
const CompressedWeights<TConfig>& weights,
Activations<TConfig, kBatchSize>& activations) {
constexpr size_t kModelDim = TConfig::kModelDim;
GEMMA_CONSTEXPR_EMBSCALING const float kEmbScaling =
EmbeddingScaling<TConfig>();
HWY_DASSERT(token >= 0);
HWY_DASSERT(token < TConfig::kVocabSize);
Decompress(weights.embedder_input_embedding, token * kModelDim,
activations.x.data() + token_idx * kModelDim, kModelDim);
MulByConst(kEmbScaling, activations.x.data() + token_idx * kModelDim,
kModelDim);
if constexpr (TConfig::kAbsolutePE) {
AddAbsolutePositionalEmbeddings(
activations.x.data() + token_idx * kModelDim, kModelDim,
pos + token_idx);
};
}
template <class TConfig, size_t kBatchSize, size_t kQueryBatchSize>
HWY_NOINLINE void TransformerLayer(
size_t num_tokens, size_t num_queries, size_t pos, size_t layer,
const CompressedLayer<TConfig>* layer_weights,
Activations<TConfig, kBatchSize * kQueryBatchSize>& activations,
const std::vector<KVCache*>& kv_caches, hwy::ThreadPool& pool) {
constexpr size_t kModelDim = TConfig::kModelDim;
const size_t num_tokens_and_queries = num_tokens * num_queries;
auto type = TConfig::kLayerConfig[layer];
size_t layer_of_type =
NumLayersOfTypeBefore(TConfig::kLayerConfig, type, layer);
RMSNormBatched<kBatchSize * kQueryBatchSize>(
num_tokens_and_queries, activations.x.data(),
layer_weights->pre_attention_norm_scale.data(),
activations.pre_att_rms_out.data(), kModelDim);
if (type == LayerAttentionType::kGemma) {
Attention<TConfig, kBatchSize, kQueryBatchSize>(
pos, num_tokens, num_queries, layer_of_type, activations,
layer_weights, kv_caches, pool);
} else {
// This Griffin layers should never exist unless the model is a Griffin
// model. This conditional prevents the compiler from generating code for
// this branch when building a non-Griffin model, since we have static
// asserts about the query batch size for Griffin layers.
if constexpr (TConfig::kGriffinLayers > 0) {
GriffinRecurrent<TConfig, kBatchSize, kQueryBatchSize>(
pos, num_tokens, num_queries, layer_of_type, activations,
layer_weights, kv_caches, pool);
}
}
if (TConfig::kPostNormScale) {
RMSNormInplaceBatched<kBatchSize * kQueryBatchSize>(
num_tokens_and_queries,
layer_weights->post_attention_norm_scale.data(),
activations.att_post2.data(), kModelDim);
}
AddFromBatched<kBatchSize * kQueryBatchSize>(num_tokens_and_queries,
activations.att_post2.data(),
activations.x.data(), kModelDim);
RMSNormBatched<kBatchSize * kQueryBatchSize>(
num_tokens_and_queries, activations.x.data(),
layer_weights->pre_ffw_norm_scale.data(),
activations.bf_pre_ffw_rms_out.data(), kModelDim);
FFW<TConfig, kBatchSize * kQueryBatchSize>(
activations, num_tokens_and_queries, layer_weights, pool);
if (TConfig::kPostNormScale) {
RMSNormInplaceBatched<kBatchSize * kQueryBatchSize>(
num_tokens_and_queries, layer_weights->post_ffw_norm_scale.data(),
activations.ffw_out.data(), kModelDim);
}
AddFromBatched<kBatchSize * kQueryBatchSize>(
num_tokens_and_queries, activations.ffw_out.data(),
activations.x.data(), kModelDim);
}
template <class TConfig, size_t kBatchSize, size_t kQueryBatchSize>
HWY_NOINLINE void Prefill(
const int* tokens, size_t num_tokens, size_t num_queries, size_t pos,
const CompressedWeights<TConfig>& weights,
Activations<TConfig, kBatchSize * kQueryBatchSize>& activations,
const std::vector<KVCache*>& kv_caches, hwy::ThreadPool& pool) {
HWY_DASSERT(num_queries <= kQueryBatchSize);
const size_t minibatch_size = std::min(num_tokens, kBatchSize);
PROFILER_ZONE("Gen.Prefill\\Att\\FFW");
// TODO(patrickms): Try to hoist pool.Run out of the loop.
for (size_t i = 0; i < num_tokens; i += minibatch_size) {
const size_t offset = i * num_queries;
const size_t current_token_count = std::min(
minibatch_size, num_tokens - i);
pool.Run(0, current_token_count * num_queries,
[&](const uint64_t token_idx, size_t /*thread*/) HWY_ATTR {
EmbedToken<TConfig, kBatchSize * kQueryBatchSize>(
tokens[token_idx + offset], token_idx, pos + offset,
weights, activations);
});
for (size_t layer = 0; layer < TConfig::kLayers; ++layer) {
const auto* layer_weights = weights.GetLayer(layer);
TransformerLayer<TConfig, kBatchSize, kQueryBatchSize>(
current_token_count, num_queries, pos + offset , layer, layer_weights,
activations, kv_caches, pool);
}
}
}
// Compute the transformer for a batch of input tokens. During generation,
// we usually have num_tokens == 1 (and also kBatchSize == 1).
template <class TConfig, size_t kBatchSize, size_t kQueryBatchSize>
HWY_NOINLINE void Transformer(
const int* tokens, size_t num_tokens, size_t num_queries, size_t pos,
const CompressedWeights<TConfig>& weights,
Activations<TConfig, kBatchSize * kQueryBatchSize>& activations,
const std::vector<KVCache*>& kv_caches,
hwy::ThreadPool& pool,
const LayersOutputFunc& layers_output) {
HWY_ASSERT(num_tokens <= kBatchSize);
const size_t num_tokens_and_queries = num_tokens * num_queries;
if (layers_output) {
for (size_t token_idx = 0; token_idx < num_tokens_and_queries;
++token_idx) {
float token_f = tokens[token_idx];
layers_output(pos + token_idx, "Tokens", &token_f, 1);
}
}
constexpr size_t kModelDim = TConfig::kModelDim;
for (size_t token_idx = 0; token_idx < num_tokens_and_queries; ++token_idx) {
EmbedToken<TConfig, kBatchSize * kQueryBatchSize>(
tokens[token_idx], token_idx, pos, weights, activations);
}
for (size_t layer = 0; layer < TConfig::kLayers; ++layer) {
const CompressedLayer<TConfig>* layer_weights = weights.GetLayer(layer);
TransformerLayer<TConfig, kBatchSize, kQueryBatchSize>(
num_tokens, num_queries, pos, layer, layer_weights,
activations, kv_caches, pool);
if (layers_output) {
const std::string block_name = "blocks." + std::to_string(layer);
for (size_t token_idx = 0; token_idx < num_tokens_and_queries;
++token_idx) {
layers_output(pos + token_idx, block_name,
activations.x.data() + token_idx * kModelDim, kModelDim);
}
}
}
RMSNormInplaceBatched<kBatchSize * kQueryBatchSize>(
num_tokens * num_queries, weights.final_norm_scale.data(),
activations.x.data(), kModelDim);
if (layers_output) {
for (size_t token_idx = 0; token_idx < num_tokens_and_queries;
++token_idx) {
layers_output(pos + token_idx, "final_norm",
activations.x.data() + token_idx * kModelDim, kModelDim);
}
}
}
template <class TConfig>
void RangeChecks(size_t& max_tokens, size_t& max_generated_tokens,
size_t& prompt_size) {
if (!TConfig::kUseLocalAttention) {
if (max_tokens > TConfig::kSeqLen) {
fprintf(stderr, "WARNING: max_tokens %zu > kSeqLen %d, truncating.\n",
max_tokens, TConfig::kSeqLen);
max_tokens = static_cast<size_t>(TConfig::kSeqLen);
}
}
if (max_generated_tokens > max_tokens) {
fprintf(stderr,
"WARNING: max_generated_tokens %zu > max_tokens %zu, truncating.\n",
max_generated_tokens, max_tokens);
max_generated_tokens = max_tokens - 1;
}
if (!TConfig::kUseLocalAttention) {
if (prompt_size + max_generated_tokens > max_tokens) {
fprintf(stderr,
"WARNING: prompt_size %zu + max_generated_tokens %zu > "
"max_tokens %zu, truncating to ",
prompt_size, max_generated_tokens, max_tokens);
prompt_size = std::min(prompt_size, max_tokens - max_generated_tokens);
fprintf(stderr, "%zu\n", prompt_size);
}
}
HWY_ASSERT(prompt_size > 0);
}
// Placeholder for internal test3, do not remove
template <class TConfig, size_t kQueryBatchSize>
void GenerateT(const ByteStorageT& weights_u8, const ByteStorageT& prefill_u8,
const ByteStorageT& decode_u8,
const RuntimeConfig& runtime_config,
const hwy::Span<const hwy::Span<int>>& prompts, size_t pos,
const size_t query_index_offset,
const std::vector<KVCache*>& kv_caches, hwy::ThreadPool& pool,
TimingInfo& timing_info) {
constexpr size_t kAdjustedPrefillBatchSize =
std::max((size_t)1, kPrefillBatchSize / kQueryBatchSize);
static_assert(kAdjustedPrefillBatchSize >= kMinAdjustedPrefillBatchSize);
const size_t num_queries = prompts.size();
HWY_DASSERT(num_queries <= kQueryBatchSize);
pos *= num_queries; // position in (num_queries) interleaved token sequence.
const CompressedWeights<TConfig>& weights =
*reinterpret_cast<const CompressedWeights<TConfig>*>(weights_u8.get());
auto& prefill_activations =
GetActivations<TConfig,
kAdjustedPrefillBatchSize * kQueryBatchSize>(prefill_u8);
auto& activations = GetActivations<TConfig, kQueryBatchSize>(decode_u8);
size_t min_prompt_size = (size_t)-1;
size_t max_prompt_size = 0;
for (int i=0; i < prompts.size(); ++i) {
min_prompt_size = std::min(min_prompt_size, prompts[i].size());
max_prompt_size = std::max(max_prompt_size, prompts[i].size());
}
std::vector<int> prompt;
prompt.reserve(max_prompt_size * prompts.size());
for (int i = 0; i < max_prompt_size; ++i) {
for (int j=0; j < prompts.size(); ++j) {
if (i < prompts[j].size()) {
prompt.push_back(prompts[j][i]);
} else {
prompt.push_back(0);
}
}
}
constexpr size_t kVocabSize = TConfig::kVocabSize;
size_t max_tokens = runtime_config.max_tokens;
size_t max_generated_tokens = runtime_config.max_generated_tokens;
RangeChecks<TConfig>(max_tokens, max_generated_tokens, max_prompt_size);
if (pos >= max_tokens) {
fprintf(stderr, "Warning: pos %zu >= max_tokens %zu, aborting.\n", pos,
max_tokens);
return;
}
// If no sample_func is provided, we use top-k sampling.
const SampleFunc sample_token =
runtime_config.sample_func
? runtime_config.sample_func
: [&](const float* logits, size_t vocab_size) -> int {
return SampleTopK<TConfig::kTopK>(logits, vocab_size, *runtime_config.gen,
runtime_config.temperature,
runtime_config.accept_token);
};
std::vector<bool> reached_eos(num_queries);
std::fill(reached_eos.begin(), reached_eos.end(), false);
// pos indexes the KV cache. In the first turn of a chat, pos = 0.
//
// After the first turn, pos gets passed in with > 0 corresponding to the
// current token position in the KV cache.
//
// pos_offset keeps track of the relative position within the turn, starting
// at 0 each turn. During prefill, pos_offset corresponds to the index into
// the prompt vector.
//
// In single-turn (non-chat) usage, pos and pos_offset start at 0 and are
// always equal.
size_t pos_offset = 0; // offset relative to pos
// Used to keep track of how many tokens are processed per prompt,
// so that we know when to start generating tokens.
size_t single_prompt_pos_offset = 0;
const double prefill_start = hwy::platform::Now();
// Prefill stops before prompt_size - 1 since the last prompt token is the
// first input token for generation.
while (single_prompt_pos_offset < min_prompt_size - 1) {
const size_t batch_size = std::min(
kPrefillBatchSize, min_prompt_size - 1 - single_prompt_pos_offset);
const size_t batch_and_query_size = batch_size * num_queries;
HWY_DASSERT(batch_size <= kPrefillBatchSize);
HWY_DASSERT(single_prompt_pos_offset + batch_size <= min_prompt_size - 1);
HWY_DASSERT(pos_offset + batch_size <= (min_prompt_size - 1) * num_queries);
const int* batch_tokens = prompt.data() + pos_offset;
Prefill<TConfig, kAdjustedPrefillBatchSize, kQueryBatchSize>(
batch_tokens, batch_size, num_queries, pos, weights,
prefill_activations, kv_caches, pool);
for (size_t idx = 0; idx < batch_size; ++idx) {
bool all_tokens_eos = true;
for (size_t query_idx = 0; query_idx < num_queries; ++query_idx) {
if (reached_eos[query_idx]) continue;
if (runtime_config.StreamToken(
query_idx + query_index_offset, single_prompt_pos_offset,
batch_tokens[idx * num_queries + query_idx], 0.0f)) {
all_tokens_eos = false;
} else {
reached_eos[query_idx] = true;
}
}
if (all_tokens_eos) {
return;
}
}
pos += batch_and_query_size;
pos_offset += batch_and_query_size;
single_prompt_pos_offset += batch_size;
}
timing_info.prefill_tok_sec =
static_cast<double>(pos_offset) / (hwy::platform::Now() - prefill_start);
// Start generation.
const double gen_start = hwy::platform::Now();
HWY_DASSERT(single_prompt_pos_offset == min_prompt_size - 1);
size_t pos_gen_start = pos_offset;
int token = prompt.at(pos_offset);
std::vector<int>::const_iterator first = prompt.begin() + pos_offset;
std::vector<int>::const_iterator last = first + num_queries;
std::vector<int> gen_tokens(first, last);
// The loop below is not yet prepared for decode batch size > 1.
HWY_ASSERT(kDecodeBatchSize == 1);
bool all_tokens_eos = true;
for (size_t i=0; i < num_queries; ++i) {
if (reached_eos[i]) continue;
if (runtime_config.StreamToken(i + query_index_offset,
single_prompt_pos_offset, gen_tokens[i],
0.0f)) {
all_tokens_eos = false;
} else {
reached_eos[i] = true;
}
}
if (all_tokens_eos) {
return;
}
for (size_t generate_pos = 0;
generate_pos < max_tokens && generate_pos < max_generated_tokens;
++single_prompt_pos_offset, ++generate_pos) {
Transformer<TConfig, kDecodeBatchSize, kQueryBatchSize>(
gen_tokens.data(), kDecodeBatchSize, num_queries, pos, weights,
activations, kv_caches, pool, runtime_config.layers_output);
float token_logit = 0.0f;
// The condition below is always true if we are doing Prefill above.
// We keep it here for clarity so that the code is correct even if Prefill
// is disabled.
bool all_tokens_eos = true;
float* x = activations.x.data();
float* logits = activations.logits.data();
for (size_t i = 0; i < num_queries; ++i, ++pos, ++pos_offset,
x += TConfig::kModelDim, logits += kVocabSize) {
const size_t prompt_size = prompts[i].size();
const bool is_generating_phase =
(single_prompt_pos_offset >= prompt_size - 1);
if (is_generating_phase) {
PROFILER_ZONE("Gen.Embedding");
// Compute logits from last layer activations.
MatVec<kVocabSize, TConfig::kModelDim>(
weights.embedder_input_embedding, 0, x, activations.even_odd.data(),
logits, pool);
if constexpr (TConfig::kFinalCap > 0.0f) {
LogitsSoftCap(TConfig::kFinalCap, activations.logits.data(),
kVocabSize);
}
// Barrier: must have all logits so we can subtract max.
Softmax(logits, kVocabSize);
token = sample_token(logits, kVocabSize);
token_logit = logits[token];
if (generate_pos == 0) {
timing_info.time_to_first_token = hwy::platform::Now() - gen_start;
}
} else {
// We would take this branch if we were not doing Prefill but would
// process the tokens of the prompt one at a time.
token = prompt.at(pos_offset);
token_logit = 0.0f;
}
if (!reached_eos[i]) {
if (!runtime_config.StreamToken(i + query_index_offset,
single_prompt_pos_offset + 1, token,
token_logit)) {
token = runtime_config.eos_id;
}
if (token != runtime_config.eos_id) {
all_tokens_eos = false;
} else {
reached_eos[i] = true;
}
}
gen_tokens[i] = token;
}
if (all_tokens_eos) {
break;
}
}
timing_info.gen_tok_sec = static_cast<double>(pos_offset - pos_gen_start) /
(hwy::platform::Now() - gen_start);
}
template <class TConfig>
void GenerateSingleT(const ByteStorageT& weights_u8,
const ByteStorageT& prefill_u8,
const ByteStorageT& decode_u8,
const RuntimeConfig& runtime_config,
const std::vector<int>& prompt, size_t pos,
KVCache& kv_cache, hwy::ThreadPool& pool,
TimingInfo& timing_info) {
// TODO: the input should also be span, not a vector.
const hwy::Span<int> prompt_span(const_cast<int*>(prompt.data()),
prompt.size());
const hwy::Span<const hwy::Span<int>> prompts(&prompt_span, 1);
// TODO: also span of kv_cache.
std::vector<KVCache*> kv_caches = {&kv_cache};
const size_t query_index_offset = 0;
GenerateT<TConfig, /*kQueryBatchSize=*/1>(
weights_u8, prefill_u8, decode_u8, runtime_config, prompts, pos,
query_index_offset, kv_caches, pool, timing_info);
}
template <class TConfig>
void GenerateBatchT(const ByteStorageT& weights_u8,
const ByteStorageT& prefill_u8,
const ByteStorageT& decode_u8,
const RuntimeConfig& runtime_config,
const hwy::Span<const hwy::Span<int>>& prompts,
size_t pos, const std::vector<KVCache*>& kv_caches,
hwy::ThreadPool& pool,
TimingInfo& timing_info) {
// Disable query batching for Griffin models.
constexpr size_t kQueryBatchSize =
(TConfig::kGriffinLayers > 0) ? 1 : kBatchedQueryBatchSize;
for (size_t i = 0; i < prompts.size(); i += kQueryBatchSize) {
const size_t num_queries = std::min(prompts.size() - i, kQueryBatchSize);
const hwy::Span<const hwy::Span<int>> current_prompts(
prompts.data() + i, num_queries);
GenerateT<TConfig, kQueryBatchSize>(weights_u8, prefill_u8, decode_u8,
runtime_config, current_prompts,
pos, i, kv_caches, pool, timing_info);
}
}
} // namespace HWY_NAMESPACE
#if HWY_ONCE
// These are extern functions defined by instantiations/*.cc, which include this
// 'header' after defining GEMMA_CONFIG, which is for function overloading.
void GenerateSingle( // NOLINT(misc-definitions-in-headers)
GEMMA_CONFIG, const ByteStorageT& weights_u8,
const ByteStorageT& prefill_u8, const ByteStorageT& decode_u8,
const RuntimeConfig& runtime_config, const std::vector<int>& prompt,
size_t pos, KVCache& kv_cache, hwy::ThreadPool& pool,
TimingInfo& timing_info) {
HWY_EXPORT_AND_DYNAMIC_DISPATCH_T(GenerateSingleT<GEMMA_CONFIG>)
(weights_u8, prefill_u8, decode_u8, runtime_config, prompt, pos, kv_cache,
pool, timing_info);
}
void GenerateBatch( // NOLINT(misc-definitions-in-headers)
GEMMA_CONFIG, const ByteStorageT& weights_u8,
const ByteStorageT& prefill_u8, const ByteStorageT& decode_u8,
const RuntimeConfig& runtime_config,
const hwy::Span<const hwy::Span<int>>& prompts, size_t pos,
const std::vector<KVCache*>& kv_caches, hwy::ThreadPool& pool,
TimingInfo& timing_info) {
HWY_EXPORT_AND_DYNAMIC_DISPATCH_T(GenerateBatchT<GEMMA_CONFIG>)
(weights_u8, prefill_u8, decode_u8, runtime_config, prompts, pos, kv_caches,
pool, timing_info);
}
#endif // HWY_ONCE
} // namespace gcpp
HWY_AFTER_NAMESPACE();
#endif // THIRD_PARTY_GEMMA_CPP_GEMMA_GEMMA_INL_H_

1028
gemma/gemma.cc
View File

File diff suppressed because it is too large Load Diff

7
gemma/gemma.h
View File

@ -56,6 +56,13 @@ using LayersOutputFunc =
std::function<void(int, const std::string&, const float*, size_t)>; std::function<void(int, const std::string&, const float*, size_t)>;
struct RuntimeConfig { struct RuntimeConfig {
bool StreamToken(size_t query_idx, size_t pos, int token, float prob) const {
if (batch_stream_token) {
return batch_stream_token(query_idx, pos, token, prob);
}
return stream_token(token, prob);
}
size_t max_tokens; size_t max_tokens;
size_t max_generated_tokens; size_t max_generated_tokens;
float temperature; float temperature;

21
gemma/instantiations/27b_bf16.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/27b_bf16.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGemma27B<hwy::bfloat16_t>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/27b_f32.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/27b_f32.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGemma27B<float>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/27b_sfp.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/27b_sfp.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGemma27B<SfpStream>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/2b_bf16.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/2b_bf16.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGemma2B<hwy::bfloat16_t>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/2b_f32.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/2b_f32.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGemma2B<float>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/2b_sfp.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/2b_sfp.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGemma2B<SfpStream>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/7b_bf16.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/7b_bf16.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGemma7B<hwy::bfloat16_t>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/7b_f32.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/7b_f32.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGemma7B<float>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/7b_sfp.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/7b_sfp.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGemma7B<SfpStream>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/9b_bf16.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/9b_bf16.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGemma9B<hwy::bfloat16_t>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/9b_f32.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/9b_f32.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGemma9B<float>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/9b_sfp.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/9b_sfp.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGemma9B<SfpStream>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/gr2b_bf16.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/gr2b_bf16.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGriffin2B<hwy::bfloat16_t>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/gr2b_f32.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/gr2b_f32.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGriffin2B<float>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/gr2b_sfp.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/gr2b_sfp.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGriffin2B<SfpStream>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/tiny_bf16.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/tiny_bf16.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGemmaTiny<hwy::bfloat16_t>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/tiny_f32.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/tiny_f32.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGemmaTiny<float>
#include "gemma/gemma-inl.h"

21
gemma/instantiations/tiny_sfp.cc Normal file
View File

@ -0,0 +1,21 @@
// Copyright 2024 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE \
"gemma/instantiations/tiny_sfp.cc"
#include "hwy/foreach_target.h" // IWYU pragma: keep
#define GEMMA_CONFIG ConfigGemmaTiny<SfpStream>
#include "gemma/gemma-inl.h"

17
gemma/tokenizer.cc
View File

@ -22,7 +22,8 @@
#include <vector> #include <vector>
#include "compression/io.h" // Path #include "compression/io.h" // Path
#include "hwy/base.h" #include "gemma/common.h" // Wrap
#include "hwy/base.h" // HWY_ASSERT
#include "hwy/profiler.h" #include "hwy/profiler.h"
// copybara:import_next_line:sentencepiece // copybara:import_next_line:sentencepiece
#include "src/sentencepiece_processor.h" #include "src/sentencepiece_processor.h"
@ -95,4 +96,18 @@ bool GemmaTokenizer::Decode(const std::vector<int>& ids,
return impl_->Decode(ids, detokenized); return impl_->Decode(ids, detokenized);
} }
std::vector<int> WrapAndTokenize(const GemmaTokenizer& tokenizer,
const ModelInfo& info, size_t pos,
std::string& prompt) {
Wrap(info, pos, prompt);
std::vector<int> tokens;
HWY_ASSERT(tokenizer.Encode(prompt, &tokens));
// Both pre-trained and instruction-tuned require BOS as first token.
if (pos == 0) {
tokens.insert(tokens.begin(), BOS_ID);
}
return tokens;
}
} // namespace gcpp } // namespace gcpp

7
gemma/tokenizer.h
View File

@ -16,11 +16,14 @@
#ifndef THIRD_PARTY_GEMMA_CPP_GEMMA_TOKENIZER_H_ #ifndef THIRD_PARTY_GEMMA_CPP_GEMMA_TOKENIZER_H_
#define THIRD_PARTY_GEMMA_CPP_GEMMA_TOKENIZER_H_ #define THIRD_PARTY_GEMMA_CPP_GEMMA_TOKENIZER_H_
#include <stddef.h>
#include <memory> #include <memory>
#include <string> #include <string>
#include <vector> #include <vector>
#include "compression/io.h" // Path #include "compression/io.h" // Path
#include "gemma/common.h" // ModelInfo
namespace gcpp { namespace gcpp {
@ -47,6 +50,10 @@ class GemmaTokenizer {
std::unique_ptr<Impl> impl_; std::unique_ptr<Impl> impl_;
}; };
std::vector<int> WrapAndTokenize(const GemmaTokenizer& tokenizer,
const ModelInfo& info, size_t pos,
std::string& prompt);
} // namespace gcpp } // namespace gcpp
#endif // THIRD_PARTY_GEMMA_CPP_GEMMA_TOKENIZER_H_ #endif // THIRD_PARTY_GEMMA_CPP_GEMMA_TOKENIZER_H_