internal change

PiperOrigin-RevId: 718824952
This commit is contained in:
Apoorv Reddy 2025-01-23 05:28:51 -08:00 committed by Copybara-Service
parent a60b564b88
commit ce807a31a1
1 changed files with 37 additions and 35 deletions

View File

@ -28,7 +28,6 @@
#include "gemma/configs.h" #include "gemma/configs.h"
#include "gemma/gemma.h" #include "gemma/gemma.h"
#include "gemma/weights.h" #include "gemma/weights.h"
// Placeholder for internal test4, do not remove
#include "paligemma/image.h" #include "paligemma/image.h"
#include "util/allocator.h" #include "util/allocator.h"
#include "util/basics.h" #include "util/basics.h"
@ -1312,10 +1311,12 @@ void GenerateT(const ModelWeightsStorage& model, Activations& activations,
0.0f); 0.0f);
} }
{
const size_t vocab_size = model.Config().vocab_size; const size_t vocab_size = model.Config().vocab_size;
const double gen_start = hwy::platform::Now(); const double gen_start = hwy::platform::Now();
for (size_t gen = 0; gen < max_generated_tokens; ++gen) { for (size_t gen = 0; gen < max_generated_tokens; ++gen) {
// Decode generates one token per query and increments queries_mutable_pos. // Decode generates one token per query and increments
// queries_mutable_pos.
Transformer(QueriesToken(gen_tokens.data(), num_queries), Transformer(QueriesToken(gen_tokens.data(), num_queries),
queries_mutable_pos, queries_prefix_end, weights, activations, queries_mutable_pos, queries_prefix_end, weights, activations,
div_seq_len, kv_caches, runtime_config.layers_output, div_seq_len, kv_caches, runtime_config.layers_output,
@ -1334,7 +1335,8 @@ void GenerateT(const ModelWeightsStorage& model, Activations& activations,
PROFILER_ZONE("Gen.Softcap+Sample+Stream"); PROFILER_ZONE("Gen.Softcap+Sample+Stream");
for (size_t query_idx = 0; query_idx < num_queries; ++query_idx) { for (size_t query_idx = 0; query_idx < num_queries; ++query_idx) {
float* HWY_RESTRICT logits = activations.logits.Batch(query_idx); float* HWY_RESTRICT logits = activations.logits.Batch(query_idx);
MaybeLogitsSoftCap(weights.weights_config.final_cap, logits, vocab_size); MaybeLogitsSoftCap(weights.weights_config.final_cap, logits,
vocab_size);
const TokenAndProb tp = sample_token(logits, vocab_size); const TokenAndProb tp = sample_token(logits, vocab_size);
timing_info.NotifyGenerated(prefill_start, gen_start); timing_info.NotifyGenerated(prefill_start, gen_start);
@ -1346,8 +1348,8 @@ void GenerateT(const ModelWeightsStorage& model, Activations& activations,
} }
if (all_queries_eos) break; if (all_queries_eos) break;
} // foreach token to generate } // foreach token to generate
timing_info.NotifyGenerateDone(gen_start); timing_info.NotifyGenerateDone(gen_start);
}
} }
template <typename T> template <typename T>