Separate profiler output for image token generation

PiperOrigin-RevId: 880895239
This commit is contained in:
Jan Wassenberg 2026-03-09 09:26:29 -07:00 committed by Copybara-Service
parent bea8b1cdbd
commit 70cb9cf1c2
1 changed files with 4 additions and 0 deletions

View File

@ -740,6 +740,10 @@ void GenerateImageTokensT(const ModelConfig& config,
// Weights are for the full PaliGemma model, not just the ViT part.
PrefillVit(config, weights, prefill_runtime_config, image, image_tokens,
prefill_activations, env);
// No-op if the profiler is disabled. Printing now ensures that the
// `PrintResults` after prefill does not include the image token part.
env.ctx.profiler.PrintResults();
}
// NOLINTNEXTLINE(google-readability-namespace-comments)