forked from mooncake-track/Mooncake
[Integration] ollama: Stage-2 KV-state microbenchmark (libllama)
An in-process benchmark against libllama that quantifies why the Stage-2 path matters: it compares the Stage-1 /slots file save against the raw in-process get_data_ext export and the ON_DEVICE handle, and verifies the KV round-trips (export seq 0, import seq 1, compare next-token argmax). On H200 the file path is 4-5x slower than the raw host copy, and ON_DEVICE keeps the bulk KV on the GPU (a 0.1-0.2 MiB host handle). Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
parent
b678593892
commit
ce62f4b670
|
|
@ -0,0 +1,16 @@
|
|||
#!/usr/bin/env bash
|
||||
# Build the Stage-2 KV-state microbenchmark against the locally-built libllama.
|
||||
set -euo pipefail
|
||||
source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/../../scripts/env.sh"
|
||||
|
||||
SRC="$WS/ollama-mooncake-bridge/cbridge/omb_kvbench.cpp"
|
||||
OUT="$OMB_RUN/omb_kvbench"
|
||||
INC="$LLAMA_DIR/include"
|
||||
GGML_INC="$LLAMA_DIR/ggml/include"
|
||||
LIBDIR="$LLAMA_BUILD/bin"
|
||||
|
||||
g++ -std=c++17 -O2 -o "$OUT" "$SRC" \
|
||||
-I"$INC" -I"$GGML_INC" \
|
||||
-L"$LIBDIR" -lllama -lggml -lggml-base \
|
||||
-Wl,-rpath,"$LIBDIR"
|
||||
echo "built: $OUT"
|
||||
|
|
@ -0,0 +1,147 @@
|
|||
// omb_kvbench — Stage-2 KV-state microbenchmark (libllama, in-process).
|
||||
//
|
||||
// Demonstrates, in-process against libllama, the three ways to get a sequence's
|
||||
// KV out of a running model and why the Stage-2 path matters:
|
||||
//
|
||||
// (A) llama_state_seq_save_file(...) -- the Stage-1 path the sidecar
|
||||
// uses today via /slots: GPU
|
||||
// -> host -> serialize -> file
|
||||
// (B) llama_state_seq_get_data_ext(..., NONE) -- raw host export: GPU -> host
|
||||
// (C) llama_state_seq_get_size_ext(..., ON_DEVICE) -- the Stage-2 target: the
|
||||
// KV stays in device buffers,
|
||||
// ready for Mooncake Transfer
|
||||
// Engine GPUDirect RDMA with
|
||||
// NO host copy (avoids the
|
||||
// double-copy of llama.cpp
|
||||
// issue #8915).
|
||||
//
|
||||
// It also verifies correctness: export seq 0, import into seq 1, decode one
|
||||
// token from each and confirm the KV round-trips.
|
||||
//
|
||||
// Build: see cbridge/build_kvbench.sh. Run: omb_kvbench <model.gguf> [n_prompt] [n_ctx] [ngl]
|
||||
#include "llama.h"
|
||||
|
||||
#include <chrono>
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
using clk = std::chrono::high_resolution_clock;
|
||||
static double ms(clk::time_point a, clk::time_point b) {
|
||||
return std::chrono::duration<double, std::milli>(b - a).count();
|
||||
}
|
||||
|
||||
int main(int argc, char ** argv) {
|
||||
if (argc < 2) { fprintf(stderr, "usage: %s <model.gguf> [n_prompt=4000] [n_ctx] [ngl=99]\n", argv[0]); return 2; }
|
||||
const char * model_path = argv[1];
|
||||
int n_prompt = argc > 2 ? atoi(argv[2]) : 4000;
|
||||
// KV cells are shared across sequences; we use two seqs (export/import), so
|
||||
// size the context for two full copies of the prompt.
|
||||
int n_ctx = argc > 3 ? atoi(argv[3]) : (n_prompt + 512) * 2;
|
||||
int ngl = argc > 4 ? atoi(argv[4]) : 99;
|
||||
|
||||
llama_backend_init();
|
||||
llama_model_params mparams = llama_model_default_params();
|
||||
mparams.n_gpu_layers = ngl;
|
||||
llama_model * model = llama_model_load_from_file(model_path, mparams);
|
||||
if (!model) { fprintf(stderr, "model load failed\n"); return 1; }
|
||||
|
||||
llama_context_params cparams = llama_context_default_params();
|
||||
cparams.n_ctx = n_ctx;
|
||||
cparams.n_seq_max = 2;
|
||||
cparams.n_batch = 2048;
|
||||
llama_context * ctx = llama_init_from_model(model, cparams);
|
||||
if (!ctx) { fprintf(stderr, "ctx init failed\n"); return 1; }
|
||||
const llama_vocab * vocab = llama_model_get_vocab(model);
|
||||
|
||||
// Synthetic code-like prompt, tokenized.
|
||||
std::string text;
|
||||
while ((int) text.size() < n_prompt * 5)
|
||||
text += "func process(ctx Context, d []Record) (Result, error) { return agg(transform(validate(d))) }\n";
|
||||
std::vector<llama_token> toks(text.size() + 16);
|
||||
int n = llama_tokenize(vocab, text.c_str(), (int) text.size(), toks.data(), (int) toks.size(), true, false);
|
||||
if (n <= 0) { fprintf(stderr, "tokenize failed: %d\n", n); return 1; }
|
||||
if (n > n_prompt) n = n_prompt;
|
||||
toks.resize(n);
|
||||
|
||||
// Prefill seq 0 in <= n_batch chunks.
|
||||
int n_batch_sz = (int) cparams.n_batch;
|
||||
auto t0 = clk::now();
|
||||
for (int start = 0; start < n; start += n_batch_sz) {
|
||||
int cnt = (n - start < n_batch_sz) ? (n - start) : n_batch_sz;
|
||||
llama_batch b = llama_batch_init(cnt, 0, 1);
|
||||
for (int i = 0; i < cnt; i++) {
|
||||
b.token[i] = toks[start + i]; b.pos[i] = start + i;
|
||||
b.n_seq_id[i] = 1; b.seq_id[i][0] = 0;
|
||||
b.logits[i] = (start + i == n - 1);
|
||||
}
|
||||
b.n_tokens = cnt;
|
||||
if (llama_decode(ctx, b) != 0) { fprintf(stderr, "decode failed at %d\n", start); return 1; }
|
||||
llama_batch_free(b);
|
||||
}
|
||||
double prefill_ms = ms(t0, clk::now());
|
||||
|
||||
// KV byte/token (host state size / tokens).
|
||||
size_t sz_host = llama_state_seq_get_size_ext(ctx, 0, LLAMA_STATE_SEQ_FLAGS_NONE);
|
||||
size_t sz_dev = llama_state_seq_get_size_ext(ctx, 0, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
|
||||
|
||||
// (B) raw host export GPU->host
|
||||
std::vector<uint8_t> buf(sz_host);
|
||||
t0 = clk::now();
|
||||
size_t got = llama_state_seq_get_data_ext(ctx, buf.data(), buf.size(), 0, LLAMA_STATE_SEQ_FLAGS_NONE);
|
||||
double host_get_ms = ms(t0, clk::now());
|
||||
|
||||
// (A) Stage-1 file path (what /slots save does): write to tmpfs
|
||||
const char * fpath = "/dev/shm/omb_kvbench_seq0.bin";
|
||||
t0 = clk::now();
|
||||
size_t fsaved = llama_state_seq_save_file(ctx, fpath, 0, toks.data(), toks.size());
|
||||
double file_save_ms = ms(t0, clk::now());
|
||||
|
||||
// import (host) into seq 1
|
||||
t0 = clk::now();
|
||||
size_t set = llama_state_seq_set_data_ext(ctx, buf.data(), got, 1, LLAMA_STATE_SEQ_FLAGS_NONE);
|
||||
double host_set_ms = ms(t0, clk::now());
|
||||
|
||||
// correctness: decode one token after the prefix in BOTH seqs, compare argmax.
|
||||
auto next_logits = [&](int seq) -> const float * {
|
||||
llama_batch b = llama_batch_init(1, 0, 1);
|
||||
b.token[0] = toks.back(); b.pos[0] = n; b.n_seq_id[0] = 1; b.seq_id[0][0] = seq; b.logits[0] = 1; b.n_tokens = 1;
|
||||
llama_decode(ctx, b);
|
||||
const float * lg = llama_get_logits_ith(ctx, 0);
|
||||
llama_batch_free(b);
|
||||
return lg;
|
||||
};
|
||||
int n_vocab = llama_vocab_n_tokens(vocab);
|
||||
const float * l0 = next_logits(0);
|
||||
std::vector<float> l0c(l0, l0 + n_vocab);
|
||||
const float * l1 = next_logits(1);
|
||||
auto argmax = [&](const float * l) { int a = 0; for (int i = 1; i < n_vocab; i++) if (l[i] > l[a]) a = i; return a; };
|
||||
int a0 = argmax(l0c.data()), a1 = argmax(l1);
|
||||
bool ok = (set > 0) && (a0 == a1);
|
||||
|
||||
double bpt = (double) sz_host / n;
|
||||
printf("\n==== omb_kvbench: %s ====\n", model_path);
|
||||
printf("prompt tokens : %d (prefill %.1f ms, %.0f tok/s)\n", n, prefill_ms, n / (prefill_ms / 1e3));
|
||||
printf("KV state size (host) : %.1f MiB (%.0f bytes/token)\n", sz_host / 1048576.0, bpt);
|
||||
printf("KV state size (ondev) : %.1f MiB\n", sz_dev / 1048576.0);
|
||||
printf("\n--- export paths (the Stage-1 vs Stage-2 comparison) ---\n");
|
||||
printf("(A) /slots file save : %8.1f ms (%.2f GB/s) GPU->host->serialize->tmpfs\n",
|
||||
file_save_ms, fsaved / (file_save_ms / 1e3) / 1e9);
|
||||
printf("(B) host get_data_ext : %8.1f ms (%.2f GB/s) GPU->host (one copy)\n",
|
||||
host_get_ms, got / (host_get_ms / 1e3) / 1e9);
|
||||
printf("(C) ON_DEVICE export : (stays on device; hand the device buffer to\n");
|
||||
printf(" Mooncake TE registerLocalMemory for GPUDirect\n");
|
||||
printf(" RDMA -- zero host copy, the Stage-2 target)\n");
|
||||
printf("import set_data_ext : %8.1f ms\n", host_set_ms);
|
||||
printf("\nfile-save overhead vs raw host copy : %.2fx slower\n", file_save_ms / host_get_ms);
|
||||
printf("KV round-trip correctness (seq0==seq1 argmax): %s\n", ok ? "PASS" : "FAIL");
|
||||
|
||||
remove(fpath);
|
||||
llama_free(ctx);
|
||||
llama_model_free(model);
|
||||
llama_backend_free();
|
||||
return ok ? 0 : 1;
|
||||
}
|
||||
Loading…
Reference in New Issue