bench.cpp

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <iomanip>
#include <iostream>
#include <numeric>
#include <string>
#include <vector>

#include "qwen.hpp"

// ── helpers ────────────────────────────────────────────────────────────────────

using Clock = std::chrono::steady_clock;

struct BenchResult {
    int    tokens;
    double total_ms;
    double tok_per_sec;
    double ms_per_tok;
};

static BenchResult make_result(int tokens, double total_ms) {
    return {
        tokens,
        total_ms,
        tokens > 0 ? (tokens * 1000.0 / total_ms) : 0.0,
        tokens > 0 ? (total_ms / tokens) : 0.0,
    };
}

struct BenchStats {
    int    tokens;
    double mean_tps;
    double stddev_tps;
    double best_tps;
};

static BenchStats compute_stats(int tokens, const std::vector<double>& samples) {
    double sum = 0.0;
    double best = 0.0;
    for (double s : samples) {
        sum += s;
        if (s > best) best = s;
    }
    double mean = sum / samples.size();
    double sq_sum = 0.0;
    for (double s : samples) {
        double d = s - mean;
        sq_sum += d * d;
    }
    // sample stddev (unbiased, matching llama-bench)
    double stddev = samples.size() > 1
        ? std::sqrt(sq_sum / (samples.size() - 1))
        : 0.0;
    return {tokens, mean, stddev, best};
}

static void print_stats(const char* label, const BenchStats& s) {
    std::printf("  %-28s %5d tokens  %8.2f ± %5.2f tok/s  (best: %8.2f tok/s)\n",
                label, s.tokens, s.mean_tps, s.stddev_tps, s.best_tps);
}

static void print_separator() {
    std::printf("  %s\n", std::string(78, '-').c_str());
}

// ── benchmark core ─────────────────────────────────────────────────────────────

static BenchResult bench_prefill(Qwen2Model& model,
                                 const std::vector<int32_t>& prompt,
                                 Qwen2Model::KVCache& cache) {
    model.reset_kv_cache(cache);
    auto t0 = Clock::now();
    model.forward_last_logits_cached(prompt, cache);
    auto t1 = Clock::now();
    double ms = std::chrono::duration<double, std::milli>(t1 - t0).count();
    return make_result(static_cast<int>(prompt.size()), ms);
}

static BenchResult bench_decode(Qwen2Model& model,
                                int decode_tokens,
                                Qwen2Model::KVCache& cache) {
    // Cache must already be prefilled.
    // We just need any valid token id to feed – use token 0.
    // We don't care about coherent output, just throughput.
    auto t0 = Clock::now();
    for (int i = 0; i < decode_tokens; i++) {
        model.forward_next_logits_cached(0, cache);
    }
    auto t1 = Clock::now();
    double ms = std::chrono::duration<double, std::milli>(t1 - t0).count();
    return make_result(decode_tokens, ms);
}

// ── main ───────────────────────────────────────────────────────────────────────

int main(int argc, char** argv) {
    try {
        std::string model_path = "qwen2.5-0.5b-instruct-fp16.gguf";
        int threads = 0;
        int decode_tokens = 64;
        int warmup_iters = 1;
        int bench_iters = 3;

        for (int i = 1; i < argc; i++) {
            std::string arg = argv[i];
            if (arg == "--model" && i + 1 < argc) { model_path = argv[++i]; continue; }
            if (arg == "--threads" && i + 1 < argc) { threads = std::stoi(argv[++i]); continue; }
            if (arg == "--decode-tokens" && i + 1 < argc) { decode_tokens = std::stoi(argv[++i]); continue; }
            if (arg == "--warmup" && i + 1 < argc) { warmup_iters = std::stoi(argv[++i]); continue; }
            if (arg == "--iters" && i + 1 < argc) { bench_iters = std::stoi(argv[++i]); continue; }
            if (arg == "--help" || arg == "-h") {
                std::printf("Usage: %s [--model PATH] [--threads N] [--decode-tokens N] [--warmup N] [--iters N]\n", argv[0]);
                return 0;
            }
            std::fprintf(stderr, "Unknown argument: %s\n", arg.c_str());
            return 1;
        }

        // ── load model ─────────────────────────────────────────────────────
        std::printf("Loading model: %s\n", model_path.c_str());
        auto load_t0 = Clock::now();
        Qwen2Model model(model_path, threads);
        auto load_t1 = Clock::now();
        double load_ms = std::chrono::duration<double, std::milli>(load_t1 - load_t0).count();

        const Qwen2Config& cfg = model.config();
        std::printf("\n");
        std::printf("=== STEEL Inference Benchmark ===\n");
        std::printf("\n");
        std::printf("Model config:\n");
        std::printf("  layers        = %lld\n", (long long)cfg.block_count);
        std::printf("  hidden_dim    = %lld\n", (long long)cfg.embedding_length);
        std::printf("  heads         = %lld\n", (long long)cfg.attention_head_count);
        std::printf("  kv_heads      = %lld\n", (long long)cfg.attention_head_count_kv);
        std::printf("  vocab         = %lld\n", (long long)cfg.vocab_size);
        std::printf("  ctx_length    = %lld\n", (long long)cfg.context_length);
        std::printf("  ffn_dim       = %lld\n", (long long)cfg.feed_forward_length);
        std::printf("  threads       = %d\n",   model.num_threads());
        std::printf("  model load    = %.1f ms\n", load_ms);
        std::printf("\n");
        std::printf("Benchmark settings:\n");
        std::printf("  warmup iters  = %d\n", warmup_iters);
        std::printf("  bench iters   = %d\n", bench_iters);
        std::printf("  decode tokens = %d\n", decode_tokens);
        std::printf("\n");

        // ── build prompt sets ──────────────────────────────────────────────
        // Use a realistic chat prompt + synthetic longer sequences.
        struct PromptSet {
            const char* label;
            std::vector<int32_t> tokens;
        };

        std::vector<PromptSet> prompts;

        // 1) Real chat prompt
        {
            std::string text = model.format_chat_prompt(
                "Explain the theory of relativity in simple terms.",
                "You are a helpful assistant.", true);
            auto ids = model.encode_text(text, true);
            prompts.push_back({"chat prompt", std::move(ids)});
        }

        // 2) Synthetic prompts of various lengths
        int synth_lengths[] = {32, 128, 256, 512};
        for (int len : synth_lengths) {
            // Fill with token id 1 (common token) repeated
            std::vector<int32_t> ids(static_cast<size_t>(len), 1);
            char label[64];
            std::snprintf(label, sizeof(label), "synthetic %d tok", len);
            prompts.push_back({label, std::move(ids)});
        }

        Qwen2Model::KVCache cache = model.make_kv_cache();

        // ── prefill benchmark ──────────────────────────────────────────────
        std::printf("--- Prefill Benchmark (tokens/s) ---\n\n");
        print_separator();
        std::printf("  %-28s %5s tokens  %18s tok/s  %16s tok/s\n",
                    "prompt", "", "mean ± stddev", "best");
        print_separator();

        for (auto& ps : prompts) {
            // warmup
            for (int w = 0; w < warmup_iters; w++) {
                model.reset_kv_cache(cache);
                model.forward_last_logits_cached(ps.tokens, cache);
            }

            // timed runs
            std::vector<double> samples;
            for (int r = 0; r < bench_iters; r++) {
                auto res = bench_prefill(model, ps.tokens, cache);
                samples.push_back(res.tok_per_sec);
            }

            auto stats = compute_stats(static_cast<int>(ps.tokens.size()), samples);
            char label[128];
            std::snprintf(label, sizeof(label), "%s (%d tok)", ps.label, (int)ps.tokens.size());
            print_stats(label, stats);
        }
        print_separator();

        // ── decode benchmark ───────────────────────────────────────────────
        std::printf("\n--- Decode Benchmark (tokens/s) ---\n\n");

        // Test decode at different KV cache fill levels
        struct DecodeScenario {
            const char* label;
            int prefill_len;
        };

        DecodeScenario scenarios[] = {
            {"after short prefill",   32},
            {"after medium prefill", 128},
            {"after long prefill",   256},
        };

        print_separator();
        std::printf("  %-28s %5s tokens  %18s tok/s  %16s tok/s\n",
                    "scenario", "", "mean ± stddev", "best");
        print_separator();

        for (auto& sc : scenarios) {
            // Prefill with synthetic tokens to set up KV cache
            std::vector<int32_t> prefill_tok(static_cast<size_t>(sc.prefill_len), 1);

            // warmup
            for (int w = 0; w < warmup_iters; w++) {
                model.reset_kv_cache(cache);
                model.forward_last_logits_cached(prefill_tok, cache);
                for (int d = 0; d < decode_tokens; d++) {
                    model.forward_next_logits_cached(0, cache);
                }
            }

            // timed runs
            std::vector<double> samples;
            for (int r = 0; r < bench_iters; r++) {
                model.reset_kv_cache(cache);
                model.forward_last_logits_cached(prefill_tok, cache);
                auto res = bench_decode(model, decode_tokens, cache);
                samples.push_back(res.tok_per_sec);
            }

            auto stats = compute_stats(decode_tokens, samples);
            char label[128];
            std::snprintf(label, sizeof(label), "%s (%d)", sc.label, sc.prefill_len);
            print_stats(label, stats);
        }
        print_separator();

        // ── end-to-end generation benchmark ────────────────────────────────
        std::printf("\n--- End-to-End Generation Benchmark ---\n\n");
        {
            std::string text = model.format_chat_prompt(
                "What is 2+2?", "", true);
            auto prompt_ids = model.encode_text(text, true);
            int prompt_len = static_cast<int>(prompt_ids.size());
            int gen_tokens = decode_tokens;

            // warmup
            for (int w = 0; w < warmup_iters; w++) {
                model.reset_kv_cache(cache);
                model.forward_last_logits_cached(prompt_ids, cache);
                for (int d = 0; d < gen_tokens; d++) {
                    model.forward_next_logits_cached(0, cache);
                }
            }

            // timed runs
            std::vector<double> ttft_samples, prefill_tps_samples, decode_tps_samples, overall_tps_samples;
            for (int r = 0; r < bench_iters; r++) {
                model.reset_kv_cache(cache);

                auto e2e_t0 = Clock::now();
                model.forward_last_logits_cached(prompt_ids, cache);
                auto prefill_done = Clock::now();
                for (int d = 0; d < gen_tokens; d++) {
                    model.forward_next_logits_cached(0, cache);
                }
                auto e2e_t1 = Clock::now();

                double prefill_ms = std::chrono::duration<double, std::milli>(prefill_done - e2e_t0).count();
                double decode_ms = std::chrono::duration<double, std::milli>(e2e_t1 - prefill_done).count();
                double total_ms = std::chrono::duration<double, std::milli>(e2e_t1 - e2e_t0).count();

                ttft_samples.push_back(prefill_ms);
                prefill_tps_samples.push_back(prompt_len * 1000.0 / prefill_ms);
                decode_tps_samples.push_back(gen_tokens * 1000.0 / decode_ms);
                overall_tps_samples.push_back((prompt_len + gen_tokens) * 1000.0 / total_ms);
            }

            auto pp_stats = compute_stats(prompt_len, prefill_tps_samples);
            auto tg_stats = compute_stats(gen_tokens, decode_tps_samples);
            auto oa_stats = compute_stats(prompt_len + gen_tokens, overall_tps_samples);

            // TTFT stats
            double ttft_sum = 0.0, ttft_best = 1e18;
            for (double t : ttft_samples) { ttft_sum += t; if (t < ttft_best) ttft_best = t; }
            double ttft_mean = ttft_sum / ttft_samples.size();
            double ttft_sq = 0.0;
            for (double t : ttft_samples) { double d = t - ttft_mean; ttft_sq += d * d; }
            double ttft_sd = ttft_samples.size() > 1 ? std::sqrt(ttft_sq / (ttft_samples.size() - 1)) : 0.0;

            std::printf("  prompt tokens     = %d\n", prompt_len);
            std::printf("  generated tokens  = %d\n", gen_tokens);
            std::printf("  TTFT              = %.2f ± %.2f ms  (best: %.2f ms)\n", ttft_mean, ttft_sd, ttft_best);
            print_stats("prefill", pp_stats);
            print_stats("decode", tg_stats);
            print_stats("overall", oa_stats);
        }

        std::printf("\nDone.\n");
        return 0;

    } catch (const std::exception& e) {
        std::fprintf(stderr, "error: %s\n", e.what());
        return 1;
    }
}