perplexity: honor larger batch size for KLD processing

tests/test-perplexity-plumbing.cpp

@@ -45,12 +45,18 @@ int main(int argc, char ** argv) {
 
     const std::string root = argv[1];
     const std::string perplexity = read_file(root + "/tools/perplexity/perplexity.cpp");
+    const std::string ppl = slice_between(perplexity,
+            "static results_perplexity perplexity(llama_context * ctx, const common_params & params, const int32_t n_ctx)",
+            "static bool decode_helper");
     const std::string kl = slice_between(perplexity,
             "static bool kl_divergence(llama_context * ctx, const common_params & params)",
             "if (kld.count < 100) return true;");
 
     ok &= expect(perplexity.find("static int ppl_max_logits_rows(int n_vocab, const common_params & params)") != std::string::npos,
         "perplexity must cap full-vocab logits rows to avoid multi-GiB output buffers");
+    ok &= expect(ppl.find("const int max_logits_rows = params.logits_file.empty()") != std::string::npos &&
+                 ppl.find(": std::max(1, std::min(n_ctx, params.n_batch))") != std::string::npos,
+        "logits baseline generation must bypass the ordinary perplexity row cap");
     ok &= expect(perplexity.find("logits_stream.write(\"_logits_\", 8)") != std::string::npos,
         "perplexity must write upstream-compatible v1 logits baselines");
     ok &= expect(perplexity.find("kld_logits::") == std::string::npos,
@@ -61,14 +67,14 @@ int main(int argc, char ** argv) {
         "KLD must retain the upstream v1 probability cutoff");
     ok &= expect(perplexity.find("if (!kl_divergence(ctx, params))") != std::string::npos,
         "perplexity KL failures must propagate to a nonzero process exit");
-    ok &= expect(kl.find("const int max_logits_rows = ppl_max_logits_rows(n_vocab, params)") != std::string::npos,
-        "KL divergence must use the bounded logits-row cap");
-    ok &= expect(kl.find("const int n_batch = std::max(1, std::min(n_ctx_i, std::min(params.n_batch, max_logits_rows)))") != std::string::npos,
-        "KL divergence batch size must be bounded by max_logits_rows");
+    ok &= expect(kl.find("ppl_max_logits_rows") == std::string::npos,
+        "KL divergence must not silently cap the requested batch by logits memory");
+    ok &= expect(kl.find("const int n_batch = std::max(1, std::min(n_ctx_i, params.n_batch))") != std::string::npos,
+        "KL divergence batch size must honor the requested batch size");
     ok &= expect(kl.find("llama_batch_init(n_batch, 0, 1)") != std::string::npos,
         "KL divergence batch allocation must match bounded n_batch");
-    ok &= expect(kl.find("std::vector<uint16_t> log_probs_uint16(size_t(max_logits_rows) * nv)") != std::string::npos,
-        "KL divergence base-logit buffer must be bounded by max_logits_rows");
+    ok &= expect(kl.find("std::vector<uint16_t> log_probs_uint16(size_t(n_batch) * nv)") != std::string::npos,
+        "KL divergence base-logit buffer must match the decode batch size");
     ok &= expect(kl.find("const int logits_first = std::max(first, pos_start)") != std::string::npos &&
                  kl.find("const int logits_end   = std::min(n_ctx_i - 1, pos_start + batch_size)") != std::string::npos,
         "KL divergence must process only the logits rows produced by the current decode slice");

tools/perplexity/README.md

@@ -22,6 +22,17 @@ and finally the `--kl-divergence` argument to indicate that the program should c
 This is a measure of how similar the FP16 and the quantized logit distributions are with a value of 0 indicating that the distribution are the same.
 The uncertainty on the mean KL divergence is calculated by assuming the KL divergence per token follows a Gaussian distribution.
 
+KL-divergence baselines and candidates are processed in blocks controlled by
+`--batch-size`; `--ubatch-size` controls how those blocks are split internally.
+The F32 logits output requires `vocabulary size * batch size * 4` bytes of host
+memory. KL-divergence evaluation also keeps a compressed baseline block that
+requires roughly half as much memory. Use matching `--batch-size` and
+`--ubatch-size` settings when generating a baseline and evaluating a candidate.
+
+For example, a model with 248320 vocabulary entries and a batch size of 2048
+uses 1940 MiB for the F32 logits block and about 970 MiB for the compressed
+baseline block.
+
 In addition to the KL divergence the following statistics are calculated with `--kl-divergence`:
 
 * Ratio of mean FP16 PPL and quantized PPL. Uncertainty is estimated on logits, then propagated. The logarithm of this metric is also calculated and printed, it is 0 if the logit distributions are the same.

tools/perplexity/perplexity.cpp

@@ -518,7 +518,9 @@ static results_perplexity perplexity(llama_context * ctx, const common_params &
     const int n_chunk_max = tokens.size() / n_ctx;
 
     const int n_vocab = llama_vocab_n_tokens(vocab);
-    const int max_logits_rows = ppl_max_logits_rows(n_vocab, params);
+    const int max_logits_rows = params.logits_file.empty()
+            ? ppl_max_logits_rows(n_vocab, params)
+            : std::max(1, std::min(n_ctx, params.n_batch));
     const int n_seq_ctx = std::max(1, params.n_ctx / n_ctx);
     const int n_seq = params.logits_file.empty() ? std::min(n_seq_ctx, max_logits_rows) : 1;
     const int n_batch = std::max(1, std::min(n_ctx, std::min(
@@ -1796,8 +1798,7 @@ static bool kl_divergence(llama_context * ctx, const common_params & params) {
     }
 
     const int n_ctx_i = static_cast<int>(n_ctx);
-    const int max_logits_rows = ppl_max_logits_rows(n_vocab, params);
-    const int n_batch = std::max(1, std::min(n_ctx_i, std::min(params.n_batch, max_logits_rows)));
+    const int n_batch = std::max(1, std::min(n_ctx_i, params.n_batch));
     const int num_batches = (n_ctx_i + n_batch - 1) / n_batch;
     const int n_seq = 1;
     const int nv = 2*((n_vocab + 1)/2) + 4;
@@ -1806,7 +1807,7 @@ static bool kl_divergence(llama_context * ctx, const common_params & params) {
 
     llama_batch batch = llama_batch_init(n_batch, 0, 1);
 
-    std::vector<uint16_t> log_probs_uint16(size_t(max_logits_rows) * nv);
+    std::vector<uint16_t> log_probs_uint16(size_t(n_batch) * nv);
     std::vector<float>    kld_values(size_t(n_ctx - 1 - n_ctx/2)*n_chunk);
     std::vector<float> p_diff_values(size_t(n_ctx - 1 - n_ctx/2)*n_chunk);