From ce62f4b6703e1ac31efffcd03750dfad8e9e52d7 Mon Sep 17 00:00:00 2001 From: zbtrs2 Date: Tue, 23 Jun 2026 19:10:16 +0800 Subject: [PATCH] [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 --- .../cbridge/build_kvbench.sh | 16 ++ .../cbridge/omb_kvbench.cpp | 147 ++++++++++++++++++ 2 files changed, 163 insertions(+) create mode 100644 mooncake-integration/ollama/ollama-mooncake-bridge/cbridge/build_kvbench.sh create mode 100644 mooncake-integration/ollama/ollama-mooncake-bridge/cbridge/omb_kvbench.cpp diff --git a/mooncake-integration/ollama/ollama-mooncake-bridge/cbridge/build_kvbench.sh b/mooncake-integration/ollama/ollama-mooncake-bridge/cbridge/build_kvbench.sh new file mode 100644 index 00000000..bfa457ea --- /dev/null +++ b/mooncake-integration/ollama/ollama-mooncake-bridge/cbridge/build_kvbench.sh @@ -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" diff --git a/mooncake-integration/ollama/ollama-mooncake-bridge/cbridge/omb_kvbench.cpp b/mooncake-integration/ollama/ollama-mooncake-bridge/cbridge/omb_kvbench.cpp new file mode 100644 index 00000000..0a5c91df --- /dev/null +++ b/mooncake-integration/ollama/ollama-mooncake-bridge/cbridge/omb_kvbench.cpp @@ -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 [n_prompt] [n_ctx] [ngl] +#include "llama.h" + +#include +#include +#include +#include +#include +#include +#include + +using clk = std::chrono::high_resolution_clock; +static double ms(clk::time_point a, clk::time_point b) { + return std::chrono::duration(b - a).count(); +} + +int main(int argc, char ** argv) { + if (argc < 2) { fprintf(stderr, "usage: %s [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 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 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 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; +}