Repository: google/gemma.cpp
File: gemma/run.cc (C++)

Code before:
    }

    if (inference.prompt.empty() && prompt_string.empty()) {
      std::cout << "Use '%q' to quit.\n";
      continue;
    }

    // Set up runtime config.
    TimingInfo timing_info = {.verbosity = inference.verbosity};
    RuntimeConfig runtime_config = {.gen = &gen,
                                    .verbosity = inference.verbosity,
                                    .stream_token = stream_token,
                                    .use_spinning = threading.spin};
    inference.CopyTo(runtime_config);
    std::vector<int> prompt;
    size_t prefix_end = 0;
    if (have_image) {
      prompt = WrapAndTokenize(gemma.Tokenizer(), gemma.ChatTemplate(),
                               config.wrapping, abs_pos, prompt_string,
                               image_tokens.Rows());
      runtime_config.image_tokens = &image_tokens;
      prompt_size = prompt.size();
      if (config.wrapping == PromptWrapping::PALIGEMMA) {
        // The end of the prefix for prefix-LM style attention in Paligemma.
        // See Figure 2 of https://arxiv.org/abs/2407.07726.
        prefix_end = prompt_size;
        // We need to look at all the tokens for the prefix.
        runtime_config.prefill_tbatch_size = prompt_size;
      }
    } else {
      prompt = WrapAndTokenize(gemma.Tokenizer(), gemma.ChatTemplate(),
                               config.wrapping, abs_pos, prompt_string);
      prompt_size = prompt.size();
    }

    if constexpr (kVerboseLogTokens) {
      for (int i = 0; i < prompt_size; ++i) {
        fprintf(stderr, "DDD TOKEN %3d: %6d\n", i, prompt[i]);
      }
    }

    // Generate until EOS or max_generated_tokens.
    if (inference.verbosity >= 1) {
      std::cerr << "\n[ Reading prompt ] " << std::flush;
    }
    gemma.Generate(runtime_config, prompt, abs_pos, prefix_end, kv_cache,
                   timing_info);
    std::cout << "\n\n";

    // Break the loop if in non-interactive mode
    if (!inference.prompt.empty()) {

Review comment (performance):
The prefix_end change makes sense, thanks.
But the prefill_tbatch_size change can result in huge batch sizes and should not be necessary. Results ought to be the same regardless of batch size. Can we revert that?

What I can imagine is that we should still call activations.SetBatchSize. This sets the (active) number of rows in each tensor, which is initially the batch size passed to the Activations ctor.

Code after:
    }

    if (inference.prompt.empty() && prompt_string.empty()) {
      std::cout << "Use '%q' to quit.\n";
      continue;
    }

    // Set up runtime config.
    TimingInfo timing_info = {.verbosity = inference.verbosity};
    RuntimeConfig runtime_config = {.gen = &gen,
                                    .verbosity = inference.verbosity,
                                    .stream_token = stream_token,
                                    .use_spinning = threading.spin};
    inference.CopyTo(runtime_config);
    std::vector<int> prompt;
    size_t prefix_end = 0;
    if (have_image) {
      prompt = WrapAndTokenize(gemma.Tokenizer(), gemma.ChatTemplate(),
                               config.wrapping, abs_pos, prompt_string,
                               image_tokens.Rows());
      runtime_config.image_tokens = &image_tokens;
      prompt_size = prompt.size();
      if (config.wrapping == PromptWrapping::PALIGEMMA) {
        // The end of the prefix for prefix-LM style attention in Paligemma.
        // See Figure 2 of https://arxiv.org/abs/2407.07726.
        prefix_end = prompt_size;
        // We need to look at all the tokens for the prefix.
        // NOTE: Online softmax is on the roadmap, after which this requirement
        // can be lifted.
        runtime_config.prefill_tbatch_size = prompt_size;
      }
    } else {
      prompt = WrapAndTokenize(gemma.Tokenizer(), gemma.ChatTemplate(),
                               config.wrapping, abs_pos, prompt_string);
      prompt_size = prompt.size();
    }

    if constexpr (kVerboseLogTokens) {
      for (int i = 0; i < prompt_size; ++i) {
        fprintf(stderr, "DDD TOKEN %3d: %6d\n", i, prompt[i]);
      }
    }

    // Generate until EOS or max_generated_tokens.
    if (inference.verbosity >= 1) {
      std::cerr << "\n[ Reading prompt ] " << std::flush;
    }
    gemma.Generate(runtime_config, prompt, abs_pos, prefix_end, kv_cache,
                   timing_info);
    std::cout << "\n\n";