diff --git a/common/arg.cpp b/common/arg.cpp index c21598e7687f..4c803db56fdb 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1328,7 +1328,50 @@ common_params_context common_params_parser_init(common_params & params, llama_ex [](common_params & params, bool value) { params.kv_unified = value; } - ).set_env("LLAMA_ARG_KV_UNIFIED").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_BATCHED, LLAMA_EXAMPLE_BENCH, LLAMA_EXAMPLE_PARALLEL})); + ).set_env("LLAMA_ARG_KV_UNIFIED").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_BATCHED, LLAMA_EXAMPLE_BENCH, LLAMA_EXAMPLE_PARALLEL, LLAMA_EXAMPLE_PAGED})); + add_opt(common_arg( + {"-kvp", "--kv-paged"}, + {"-no-kvp", "--no-kv-paged"}, + "use paged KV buffer shared across all sequences (default: disabled)", + [](common_params & params, bool value) { + params.kv_paged = value; + } + ).set_env("LLAMA_ARG_KV_PAGED").set_examples({LLAMA_EXAMPLE_PAGED})); + add_opt(common_arg( + {"-ncpub", "--n-cpu-blocks"}, "N", + "number of physical CPU blocks for paged KV cache (default: 1)", + [](common_params & params, int value) { + params.n_cpu_blocks = value; + } + ).set_examples({LLAMA_EXAMPLE_PAGED})); + add_opt(common_arg( + {"-ngpub", "--n-gpu-blocks"}, "N", + "number of physical GPU blocks for paged KV cache (default: 1)", + [](common_params & params, int value) { + params.n_gpu_blocks = value; + } + ).set_examples({LLAMA_EXAMPLE_PAGED})); + add_opt(common_arg( + {"-kvbls", "--kv-block-size"}, "N", + "fixed number of tokens for a given paged block (default: 16)", + [](common_params & params, int value) { + if (value <= 0 || (value & (value - 1)) != 0) { + throw std::invalid_argument("--kv-block-size must be a positive power of 2"); + } + params.block_size = value; + } + ).set_examples({LLAMA_EXAMPLE_PAGED})); + add_opt(common_arg( + {"--kv-paged-watermark"}, "N", + "fraction of blocks reserved before processing new requests (default: 0.05, range [0.0, 1.0))", + [](common_params & params, const std::string & value) { + float potential_watermark = std::stof(value); + if (potential_watermark < 0.0f || potential_watermark >= 1.0f) { + throw std::invalid_argument("--kv-paged-watermark must be in range [0.0, 1.0)"); + } + params.kv_paged_watermark = potential_watermark; + } + ).set_examples({LLAMA_EXAMPLE_PAGED})); add_opt(common_arg( {"--cache-idle-slots"}, {"--no-cache-idle-slots"}, @@ -2152,7 +2195,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex [](common_params & params, int value) { params.n_sequences = value; } - ).set_examples({LLAMA_EXAMPLE_PARALLEL})); + ).set_examples({LLAMA_EXAMPLE_PARALLEL, LLAMA_EXAMPLE_PAGED})); add_opt(common_arg( {"-cb", "--cont-batching"}, {"-nocb", "--no-cont-batching"}, diff --git a/common/common.cpp b/common/common.cpp index 793b8fee7b84..7fe525f36c21 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1141,6 +1141,56 @@ struct common_init_result::impl { std::vector samplers_seq_config; }; +static void common_fit_paged_kv_blocks(common_params& params, const llama_model * model) { + GGML_ASSERT(model && "model must be loaded before fitting paged KV blocks."); + ggml_backend_dev_t dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_GPU); + if (!dev) { + LOG_WRN("%s: no GPU device found, cannot fit paged KV blocks.\n", __func__); + return; + } + + size_t free_vram = 0; + size_t total_vram = 0; + ggml_backend_dev_memory(dev, &free_vram, &total_vram); + + const uint32_t n_heads_kv = llama_model_n_head_kv(model); + const uint32_t n_layers = llama_model_n_layer(model); + const uint32_t head_dim = llama_model_n_embd(model) / llama_model_n_head(model); + const uint32_t block_size = params.block_size; + + const size_t bytes_per_block = (size_t)2 * head_dim * n_heads_kv * block_size * n_layers * ggml_type_size(GGML_TYPE_F16); + + const size_t margin = params.fit_params_target.empty() + ? (size_t)(total_vram * 0.05f) + : (size_t)params.fit_params_target[0]; + + if (free_vram <= margin) { + LOG_ERR("%s: not enough free VRAM for paged KV blocks. " + "free_vram=%.1f MiB <= margin=%.1f MiB. " + "Try reducing --margin or offloading fewer layers to GPU.\n", + __func__, free_vram / 1024.0f / 1024.0f, margin / 1024.0f / 1024.0f); + return; // leave params.n_gpu_blocks at its existing value + } + + const size_t available = (free_vram > margin) ? free_vram - margin : 0; + + if (bytes_per_block == 0 || available < bytes_per_block) { + LOG_ERR("%s: available VRAM (%.1f MiB) is less than one block (%.1f MiB). " + "Try increasing n_gpu_blocks manually or reducing block_size.\n", + __func__, available / 1024.0f / 1024.0f, bytes_per_block / 1024.0f / 1024.0f); + return; + } + + const uint32_t n_gpu_blocks = (uint32_t)(available / bytes_per_block); + const uint32_t n_cpu_blocks = (uint32_t)(n_gpu_blocks * params.cpu_to_gpu_blocks_ratio); + + LOG_INF("%s: free_vram=%0.1f MiB, bytes_per_block=%ld, n_gpu_blocks=%d, n_cpu_blocks=%d\n", + __func__, free_vram / 1024.0f / 1024.0f, bytes_per_block, n_gpu_blocks, n_cpu_blocks); + + params.n_gpu_blocks = n_gpu_blocks; + params.n_cpu_blocks = n_cpu_blocks; +} + common_init_result::common_init_result(common_params & params) : pimpl(new impl{}) { auto mparams = common_model_params_to_llama(params); @@ -1163,6 +1213,12 @@ common_init_result::common_init_result(common_params & params) : pimpl->model.reset(model); + if (params.fit_params && params.kv_paged) { + LOG_INF("%s: fitting KV paged params to device memory\n", __func__); + common_fit_paged_kv_blocks(params, pimpl->model.get()); + cparams = common_context_params_to_llama(params); // re-derive this params to reflect changes + } + const llama_vocab * vocab = llama_model_get_vocab(model); // load and optionally apply lora adapters @@ -1488,32 +1544,37 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { struct llama_context_params common_context_params_to_llama(const common_params & params) { auto cparams = llama_context_default_params(); - cparams.n_ctx = params.n_ctx; - cparams.n_seq_max = params.n_parallel; - cparams.n_batch = params.n_batch; - cparams.n_ubatch = params.n_ubatch; - cparams.n_threads = params.cpuparams.n_threads; - cparams.n_threads_batch = params.cpuparams_batch.n_threads == -1 ? - params.cpuparams.n_threads : params.cpuparams_batch.n_threads; - cparams.embeddings = params.embedding; - cparams.rope_scaling_type = params.rope_scaling_type; - cparams.rope_freq_base = params.rope_freq_base; - cparams.rope_freq_scale = params.rope_freq_scale; - cparams.yarn_ext_factor = params.yarn_ext_factor; - cparams.yarn_attn_factor = params.yarn_attn_factor; - cparams.yarn_beta_fast = params.yarn_beta_fast; - cparams.yarn_beta_slow = params.yarn_beta_slow; - cparams.yarn_orig_ctx = params.yarn_orig_ctx; - cparams.pooling_type = params.pooling_type; - cparams.attention_type = params.attention_type; - cparams.flash_attn_type = params.flash_attn_type; - cparams.cb_eval = params.cb_eval; - cparams.cb_eval_user_data = params.cb_eval_user_data; - cparams.offload_kqv = !params.no_kv_offload; - cparams.no_perf = params.no_perf; - cparams.op_offload = !params.no_op_offload; - cparams.swa_full = params.swa_full; - cparams.kv_unified = params.kv_unified; + cparams.n_ctx = params.n_ctx; + cparams.n_seq_max = params.n_parallel; + cparams.n_batch = params.n_batch; + cparams.n_ubatch = params.n_ubatch; + cparams.n_threads = params.cpuparams.n_threads; + cparams.n_threads_batch = params.cpuparams_batch.n_threads == -1 ? + params.cpuparams.n_threads : params.cpuparams_batch.n_threads; + cparams.embeddings = params.embedding; + cparams.rope_scaling_type = params.rope_scaling_type; + cparams.rope_freq_base = params.rope_freq_base; + cparams.rope_freq_scale = params.rope_freq_scale; + cparams.yarn_ext_factor = params.yarn_ext_factor; + cparams.yarn_attn_factor = params.yarn_attn_factor; + cparams.yarn_beta_fast = params.yarn_beta_fast; + cparams.yarn_beta_slow = params.yarn_beta_slow; + cparams.yarn_orig_ctx = params.yarn_orig_ctx; + cparams.pooling_type = params.pooling_type; + cparams.attention_type = params.attention_type; + cparams.flash_attn_type = params.flash_attn_type; + cparams.cb_eval = params.cb_eval; + cparams.cb_eval_user_data = params.cb_eval_user_data; + cparams.offload_kqv = !params.no_kv_offload; + cparams.no_perf = params.no_perf; + cparams.op_offload = !params.no_op_offload; + cparams.swa_full = params.swa_full; + cparams.kv_unified = params.kv_unified; + cparams.kv_paged = params.kv_paged; + cparams.block_size = params.block_size; + cparams.n_gpu_blocks = params.n_gpu_blocks; + cparams.n_cpu_blocks = params.n_cpu_blocks; + cparams.kv_paged_watermark = params.kv_paged_watermark; cparams.type_k = params.cache_type_k; cparams.type_v = params.cache_type_v; diff --git a/common/common.h b/common/common.h index a564b3b8c2b4..e31de4afb9dc 100644 --- a/common/common.h +++ b/common/common.h @@ -95,7 +95,7 @@ enum llama_example { LLAMA_EXAMPLE_FIT_PARAMS, LLAMA_EXAMPLE_RESULTS, LLAMA_EXAMPLE_EXPORT_GRAPH_OPS, - + LLAMA_EXAMPLE_PAGED, LLAMA_EXAMPLE_COUNT, }; @@ -539,6 +539,15 @@ struct common_params { bool ctx_shift = false; // context shift on infinite text generation bool swa_full = false; // use full-size SWA cache (https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055) bool kv_unified = false; // enable unified KV cache + bool kv_paged = false; // enable paged KV cache + + int32_t block_size = 16; + + uint32_t n_gpu_blocks = 1; + uint32_t n_cpu_blocks = 1; + + float cpu_to_gpu_blocks_ratio = 0.25; + float kv_paged_watermark = 0.05; // percentage bool input_prefix_bos = false; // prefix BOS to user inputs, preceding input_prefix bool use_mmap = true; // enable mmap to use filesystem cache diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index a29dc707c3dc..6bbd13e73e3c 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -35,6 +35,8 @@ else() add_subdirectory(gen-docs) add_subdirectory(training) add_subdirectory(diffusion) + add_subdirectory(paged) + add_subdirectory(continuous-batch) if (NOT GGML_BACKEND_DL) add_subdirectory(convert-llama2c-to-ggml) # these examples use the backends directly and cannot be built with dynamic loading diff --git a/examples/continuous-batch/CMakeLists.txt b/examples/continuous-batch/CMakeLists.txt new file mode 100644 index 000000000000..de99c8b50ed9 --- /dev/null +++ b/examples/continuous-batch/CMakeLists.txt @@ -0,0 +1,5 @@ +set(TARGET llama-continuous-batch) +add_executable(${TARGET} continuous-batch.cpp) +install(TARGETS ${TARGET} RUNTIME) +target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_17) diff --git a/examples/continuous-batch/continuous-batch.cpp b/examples/continuous-batch/continuous-batch.cpp new file mode 100644 index 000000000000..fd6d2c3269d2 --- /dev/null +++ b/examples/continuous-batch/continuous-batch.cpp @@ -0,0 +1,368 @@ +#include "arg.h" +#include "common.h" +#include "llama.h" +#include "log.h" +#include "sampling.h" + +#include +#include +#include +#include +#include +#include +#include + +static std::vector k_prompts = { + "What is the tallest mountain in the world?", + "Who was the first person to win two Nobel Prizes?", + "Which country invented paper?", + "What organ is primarily responsible for pumping blood throughout the body?", + "Which planet is known for its prominent ring system?", + "Who directed the movie 'Inception'?", + "What is the freezing point of water in Fahrenheit?", + "Which animal is known to have the longest lifespan?", + "What language has the most native speakers worldwide?", + "What is the capital city of Canada?", + "Who is credited with inventing the World Wide Web?", + "Which metal is liquid at room temperature?", + "What is the term for an animal that eats both plants and meat?", + "Who painted 'The Starry Night'?", + "What gas do humans exhale that plants use for photosynthesis?", + "What year did World War II end?", + "Which continent has the most countries?", + "Who wrote the novel 'Frankenstein'?", + "What does DNA stand for?", + "What is the main ingredient in traditional Japanese miso soup?" +}; +static const size_t k_n_prompts = 20; + +enum class seq_status { PREFILL, DECODE, DONE }; + +struct sequence_state { + int32_t seq_id = -1; + std::string prompt = ""; + int32_t n_prompt = 0; + int32_t n_past = 0; + int32_t n_decoded = 0; + seq_status status = seq_status::PREFILL; + int64_t t_arrival_us = 0; + int64_t t_first_token_us = 0; + int64_t t_finished_us = 0; + + std::vector prompt_tokens; + std::vector output_tokens; + + common_sampler * sampler = nullptr; +}; + +struct request_result { + int32_t request_id = -1; + int32_t n_prompt = 0; + int32_t n_decoded = 0; + float ttft_ms = 0.f; // time to first token + float tpot_ms = 0.f; // avg time per output token (excl. first) + float tps = 0.f; // generation tokens/s + float e2e_ms = 0.f; // arrival to last token + std::string response = ""; +}; + +int main(int argc, char ** argv) { + std::setlocale(LC_NUMERIC, "C"); + srand(1234); + + common_params params; + + // Mirror paged main.cpp defaults exactly for a fair comparison + params.kv_unified = true; + params.kv_paged = false; + params.warmup = false; + + // Example set of parameters that work + // params.n_sequences = 100; + // params.n_parallel = params.n_sequences; + // params.n_predict = 150; + // params.n_batch = 1024; + // params.n_ubatch = params.n_batch; + + // Using PAGED examples args because I don't want to create a new one for unified + if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_PAGED)) { + return 1; + } + + if (params.n_sequences != params.n_parallel) { + LOG_INF("%s: n_sequences (%d) needs to be equal to n_parallel (%d)\n", __func__, params.n_sequences, + params.n_parallel); + return 1; + } + + common_init(); + llama_backend_init(); + llama_numa_init(params.numa); + + auto llama_init = common_init_from_params(params); + auto * model = llama_init->model(); + auto * ctx = llama_init->context(); + + if (!model || !ctx) { + LOG_ERR("%s: failed to load model or create context\n", __func__); + return 1; + } + + const llama_vocab * vocab = llama_model_get_vocab(model); + const int32_t n_sequences = params.n_sequences; + const int32_t n_predict = params.n_predict; + const int32_t n_ctx_seq = llama_n_ctx_seq(ctx); + + LOG_INF("%s: Start continuous batching loop. n_seq=%d, n_predict=%d, n_batch=%d\n", __func__, n_sequences, + n_predict, params.n_batch); + + // Initialize sequences + std::vector seqs(n_sequences); + for (int i = 0; i < n_sequences; ++i) { + sequence_state & s = seqs[i]; + s.seq_id = i; + s.prompt = k_prompts[i % k_n_prompts]; + s.t_arrival_us = ggml_time_us(); + s.status = seq_status::PREFILL; + s.sampler = common_sampler_init(model, params.sampling); + + s.prompt_tokens = common_tokenize(ctx, s.prompt, /*add_bos=*/true); + s.n_prompt = (int32_t) s.prompt_tokens.size(); + + if (s.n_prompt > n_ctx_seq) { + LOG_WRN("%s: prompt for seq %d too long (%d > %d), truncating\n", __func__, i, s.n_prompt, n_ctx_seq); + s.prompt_tokens.resize(n_ctx_seq); + s.n_prompt = n_ctx_seq; + } + + LOG_INF("%s: seq %d: \"%s\" (%d tokens)\n", __func__, i, s.prompt.c_str(), s.n_prompt); + } + + // Continuous batching inference loop + const int64_t t_start_us = ggml_time_us(); + int32_t n_finished = 0; + std::vector results; + std::unordered_map accumulated_responses; + + while (n_finished < n_sequences) { + llama_batch batch = llama_batch_init(params.n_batch, 0, 1); + batch.n_tokens = 0; + + std::unordered_map logit_pos; // seq_id mapped to batch index + std::unordered_set decode_submitted; // sequences with decode token in pass 1 + int32_t tokens_in_batch = 0; // used for token budget + + // Pass 1 (decode has higher priority): one decode token per decode sequence + for (int i = 0; i < n_sequences; ++i) { + sequence_state & s = seqs[i]; + if (s.status != seq_status::DECODE) { + continue; + } + if (tokens_in_batch >= params.n_batch) { + break; + } + + const llama_token last_tok = s.output_tokens.empty() ? s.prompt_tokens.back() : s.output_tokens.back(); + + batch.token[batch.n_tokens] = last_tok; + batch.pos[batch.n_tokens] = s.n_past; + batch.n_seq_id[batch.n_tokens] = 1; + batch.seq_id[batch.n_tokens][0] = s.seq_id; + batch.logits[batch.n_tokens] = 1; + + logit_pos[s.seq_id] = batch.n_tokens; + decode_submitted.insert(s.seq_id); + batch.n_tokens++; + tokens_in_batch++; + } + + // Pass 2: prefill tokens within remaining budget + for (int i = 0; i < n_sequences; ++i) { + sequence_state & s = seqs[i]; + if (s.status != seq_status::PREFILL) { + continue; + } + + const int32_t submitted = s.n_past; + const int32_t remaining = s.n_prompt - submitted; + if (remaining <= 0) { + s.status = seq_status::DECODE; + continue; + } + + // Budget system to avoid running out of VRAM for decodes + const int32_t budget = params.n_batch - tokens_in_batch; + const int32_t to_add = std::min(remaining, budget); + if (to_add <= 0) { + break; + } + + for (int32_t tok_pos = 0; tok_pos < to_add; ++tok_pos) { + const bool is_last = (tok_pos == (to_add - 1)) && (submitted + to_add == s.n_prompt); + + batch.token[batch.n_tokens] = s.prompt_tokens[submitted + tok_pos]; + batch.pos[batch.n_tokens] = submitted + tok_pos; + batch.n_seq_id[batch.n_tokens] = 1; + batch.seq_id[batch.n_tokens][0] = s.seq_id; + batch.logits[batch.n_tokens] = is_last ? 1 : 0; + + if (is_last) { + logit_pos[s.seq_id] = batch.n_tokens; + } + + batch.n_tokens++; + tokens_in_batch++; + } + + s.n_past += to_add; + if (s.n_past >= s.n_prompt) { + s.status = seq_status::DECODE; + } + } + + if (batch.n_tokens == 0) { + llama_batch_free(batch); + break; + } + + LOG_INF("batch.n_tokens=%d, n_batch=%d, n_sequences=%d, n_parallel=%d\n", batch.n_tokens, params.n_batch, + params.n_sequences, params.n_parallel); + GGML_ASSERT(batch.n_tokens <= params.n_batch && "batch exceeds n_batch"); + + if (llama_decode(ctx, batch) != 0) { + LOG_ERR("%s: llama_decode failed\n", __func__); + llama_batch_free(batch); + break; + } + llama_synchronize(ctx); + + // Sampling the decoded tokens + for (auto & [seq_id, logit_idx] : logit_pos) { + sequence_state & s = seqs[seq_id]; + if (s.status == seq_status::DONE) { + continue; + } + + llama_token next_tok = common_sampler_sample(s.sampler, ctx, logit_idx); + common_sampler_accept(s.sampler, next_tok, /*accept_grammar=*/true); + accumulated_responses[seq_id] += common_token_to_piece(ctx, next_tok); + + if (s.n_decoded == 0) { + s.t_first_token_us = ggml_time_us(); + } + + s.output_tokens.push_back(next_tok); + s.n_decoded++; + if (decode_submitted.count(seq_id)) { + s.n_past++; + } + + // TTFT will only be relevant for first generated token the rest will be 0.0 + LOG_INF("[Request %d] token=%s, decoded=%d, ttft=%.1fms\n", seq_id, + common_token_to_piece(ctx, next_tok).c_str(), s.n_decoded, + s.n_decoded == 1 ? (s.t_first_token_us - s.t_arrival_us) / 1000.0f : 0.0f); + + const bool is_eog = llama_vocab_is_eog(vocab, next_tok); + const bool hit_limit = s.n_decoded >= n_predict; + const bool ctx_full = s.n_past >= n_ctx_seq; + + if (is_eog || hit_limit || ctx_full) { + s.t_finished_us = ggml_time_us(); + s.status = seq_status::DONE; + n_finished++; + + llama_memory_seq_rm(llama_get_memory(ctx), s.seq_id, -1, -1); + + request_result r; + r.request_id = s.seq_id; + r.n_prompt = s.n_prompt; + r.n_decoded = s.n_decoded; + r.ttft_ms = (s.t_first_token_us - s.t_arrival_us) / 1000.0f; + r.e2e_ms = (s.t_finished_us - s.t_arrival_us) / 1000.0f; + r.tps = s.n_decoded > 0 && s.t_finished_us > s.t_first_token_us ? + s.n_decoded / ((s.t_finished_us - s.t_first_token_us) / 1e6f) : + 0.0f; + r.tpot_ms = s.n_decoded > 1 && s.t_finished_us > s.t_first_token_us ? + (s.t_finished_us - s.t_first_token_us) / 1000.0f / (s.n_decoded - 1) : + 0.0f; + r.response = accumulated_responses.count(s.seq_id) ? accumulated_responses[s.seq_id] : ""; + results.push_back(r); + + LOG_INF("[Request %d] finished. decoded=%d tokens, tps=%.1f\n", s.seq_id, s.n_decoded, + s.n_decoded / ((s.t_finished_us - s.t_arrival_us) / 1e6f)); + + if (ctx_full && !is_eog && !hit_limit) { + LOG_WRN("%s: seq %d hit context limit\n", __func__, seq_id); + } + accumulated_responses.erase(s.seq_id); + } + } + + llama_batch_free(batch); + } + + // Per-request output + std::sort(results.begin(), results.end(), + [](const request_result & a, const request_result & b) { return a.request_id < b.request_id; }); + + LOG_INF("%s: OUTPUTS:\n", __func__); + for (const auto & r : results) { + LOG_INF("Request Id: %d:%s<\\s>\n", r.request_id, r.response.c_str()); + } + + // Summary metrics + const float elapsed_s = (ggml_time_us() - t_start_us) / 1e6f; + + int32_t total_prompt_tokens = 0; + int32_t total_decoded_tokens = 0; + float sum_ttft_ms = 0.f; + float sum_tpot_ms = 0.f; // avg time per output token + float sum_e2e_ms = 0.f; + float min_ttft_ms = FLT_MAX; + float max_ttft_ms = 0.f; + float min_tps = FLT_MAX; + float max_tps = 0.f; + + for (const auto & r : results) { + total_prompt_tokens += r.n_prompt; + total_decoded_tokens += r.n_decoded; + sum_ttft_ms += r.ttft_ms; + sum_tpot_ms += r.tpot_ms; + sum_e2e_ms += r.e2e_ms; + min_ttft_ms = std::min(min_ttft_ms, r.ttft_ms); + max_ttft_ms = std::max(max_ttft_ms, r.ttft_ms); + min_tps = std::min(min_tps, r.tps); + max_tps = std::max(max_tps, r.tps); + } + + const int32_t n_results = (int32_t) results.size(); + const float avg_ttft_ms = n_results > 0 ? sum_ttft_ms / n_results : 0.f; + const float avg_tpot_ms = n_results > 0 ? sum_tpot_ms / n_results : 0.f; + const float avg_e2e_ms = n_results > 0 ? sum_e2e_ms / n_results : 0.f; + const float agg_tps = elapsed_s > 0 ? total_decoded_tokens / elapsed_s : 0.f; + + LOG_INF("\n"); + LOG_INF("=== Unified KV Cache Summary ===\n"); + LOG_INF(" n_sequences : %d\n", n_sequences); + LOG_INF(" n_predict : %d\n", n_predict); + LOG_INF(" n_batch : %d\n", params.n_batch); + LOG_INF(" total elapsed : %.2f s\n", elapsed_s); + LOG_INF(" total prompt tokens : %d\n", total_prompt_tokens); + LOG_INF(" total decoded tokens : %d\n", total_decoded_tokens); + LOG_INF(" aggregate tps : %.2f tokens/s\n", agg_tps); + LOG_INF(" --- per-request latency ---\n"); + LOG_INF(" ttft avg / min / max : %.1f / %.1f / %.1f ms\n", avg_ttft_ms, min_ttft_ms, max_ttft_ms); + LOG_INF(" tpot avg : %.1f ms/token\n", avg_tpot_ms); + LOG_INF(" e2e avg : %.1f ms\n", avg_e2e_ms); + LOG_INF(" tps min / max : %.1f / %.1f tokens/s\n", min_tps, max_tps); + LOG_INF("================================\n"); + + // Clean-up + for (auto & s : seqs) { + if (s.sampler) { + common_sampler_free(s.sampler); + } + } + llama_backend_free(); + return 0; +} diff --git a/examples/paged/CMakeLists.txt b/examples/paged/CMakeLists.txt new file mode 100644 index 000000000000..9ed615a0e353 --- /dev/null +++ b/examples/paged/CMakeLists.txt @@ -0,0 +1,5 @@ +set(TARGET llama-paged) +add_executable(${TARGET} paged.cpp) +install(TARGETS ${TARGET} RUNTIME) +target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_17) diff --git a/examples/paged/README.md b/examples/paged/README.md new file mode 100644 index 000000000000..4bdc56772522 --- /dev/null +++ b/examples/paged/README.md @@ -0,0 +1,31 @@ +# llama.cpp/example/paged + +Minimal end-to-end demo of the paged KV cache and continuous-batching scheduler. Queues N requests up front, then runs a loop of `llama_paged_scheduler_prepare_batch` → `llama_decode` → `llama_paged_scheduler_update` until every request finishes. + +## Example + +Run 10 sequences in parallel with greedy sampling, up to 50 generated tokens per sequence on a single GPU: + +```bash +llama-paged -m model.gguf -kvp -ngl 100 -sm none -mg 0 \ + -ns 10 -np 10 -n 50 -b 512 -ub 512 \ + -ngpub 500 -ncpub 100 +``` + +- `-kvp` / `--kv-paged`: enable the paged KV cache (required) +- `-ngpub N`: number of GPU KV blocks to allocate (auto-fitted from free VRAM if omitted) +- `-ncpub N`: number of CPU KV blocks for swap-out (optional) +- `-ns N`: number of sequences to queue +- `-np N`: maximum sequences decoded in parallel (must equal `-ns` in this example) + +## Phase 1 restrictions +The paged path currently requires: +- Single device (fully loaded on one GPU or CPU). Use `-sm none -mg ` to pin to one GPU on multi-GPU machines. +- Full offload of the model (`-ngl` must cover all layers). +- `-b == -ub` (batch and ubatch sizes must match). + +Note: SWA architectures (gemma3, llama4, etc.) are not yet supported. + +## Output +- Prints all the generated outputs at the end. +- Prints performance timings (general and per-request). diff --git a/examples/paged/paged.cpp b/examples/paged/paged.cpp new file mode 100644 index 000000000000..a7a4b16cb145 --- /dev/null +++ b/examples/paged/paged.cpp @@ -0,0 +1,286 @@ +#include "arg.h" +#include "common.h" +#include "llama.h" +#include "log.h" +#include "sampling.h" + +#include +#include +#include +#include +#include +#include + +static std::vector k_prompts = { + "What is the tallest mountain in the world?", + "Who was the first person to win two Nobel Prizes?", + "Which country invented paper?", + "What organ is primarily responsible for pumping blood throughout the body?", + "Which planet is known for its prominent ring system?", + "Who directed the movie 'Inception'?", + "What is the freezing point of water in Fahrenheit?", + "Which animal is known to have the longest lifespan?", + "What language has the most native speakers worldwide?", + "What is the capital city of Canada?", + "Who is credited with inventing the World Wide Web?", + "Which metal is liquid at room temperature?", + "What is the term for an animal that eats both plants and meat?", + "Who painted 'The Starry Night'?", + "What gas do humans exhale that plants use for photosynthesis?", + "What year did World War II end?", + "Which continent has the most countries?", + "Who wrote the novel 'Frankenstein'?", + "What does DNA stand for?", + "What is the main ingredient in traditional Japanese miso soup?" +}; +static const size_t k_n_prompts = 20; + +struct request_result { + int32_t request_id = -1; + int32_t n_prompt = 0; + int32_t n_decoded = 0; + float ttft_ms = 0.f; // time to first token + float tpot_ms = 0.f; // avg time per output token (excluding the first token) + float tps = 0.f; // generation tokens/s + float e2e_ms = 0.f; // time from arrival to last token + std::string response = ""; +}; + +static void add_request_from_pool(struct llama_paged_scheduler * scheduler, + llama_context * ctx, + size_t pool_index, + int seq_id) { + const std::string input_prompt = k_prompts[pool_index % k_n_prompts]; + std::vector tokens = common_tokenize(ctx, input_prompt, true); + + bool success = llama_paged_scheduler_add_request(scheduler, tokens.data(), tokens.size(), seq_id); + if (!success) { + LOG_ERR("Failed to add request %ld from pool\n", pool_index); + } + LOG_INF("%s: Successfully added request %ld: %s\n", __func__, pool_index, input_prompt.c_str()); +} + +struct callback_context { + llama_context * ctx; +}; + +static void log_output(int32_t req_id, const llama_token * /*tokens*/, int32_t n_tokens, void * user_data) { + if (!user_data) { + printf("\n--- [Request %d n_tokens] ---\n%d\n---------------------------\n", req_id, n_tokens); + return; + } +} + +int main(int argc, char ** argv) { + std::setlocale(LC_NUMERIC, "C"); + srand(1234); + + common_params params; + params.warmup = false; // we do not support warmup for paged KV cache yet + + // Example parameters that work + // params.n_sequences = 100; + // params.n_parallel = params.n_sequences; + // params.n_predict = 150; + // params.n_batch = 1024; + // params.n_ubatch = params.n_batch; + + if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_PAGED)) { + return 1; + } + + if (params.n_sequences != params.n_parallel) { + LOG_INF("%s: n_sequences (%d) needs to be equal to n_parallel (%d)\n", __func__, params.n_sequences, + params.n_parallel); + return 1; + } + + common_init(); + llama_backend_init(); + llama_numa_init(params.numa); + + auto llama_init = common_init_from_params(params); + auto * model = llama_init->model(); + auto * ctx = llama_init->context(); + + if (!model || !ctx) { + LOG_ERR("%s: failed to load model or create context\n", __func__); + return 1; + } + + const llama_vocab * vocab = llama_model_get_vocab(model); + + LOG_INF("%s: Loaded model and created context\n", __func__); + + struct llama_paged_scheduler * scheduler = llama_paged_scheduler_init(ctx); + if (!scheduler) { + LOG_ERR("%s: Failed to initialize scheduler.\n", __func__); + return 1; + } + + callback_context cb_ctx = { ctx }; + llama_paged_scheduler_set_on_finish(scheduler, log_output, &cb_ctx); + + std::unordered_map samplers; + for (int i = 0; i < params.n_sequences; ++i) { + add_request_from_pool(scheduler, ctx, (size_t) i, i); + samplers[i] = common_sampler_init(model, params.sampling); + } + + std::vector results; + std::unordered_map accumulated_responses; + + llama_batch batch = {}; + LOG_INF( + "%s: Start continuous batching loop. n_seq=%d, n_predict=%d, " + "n_gpu_blocks=%d, n_cpu_blocks=%d\n", + __func__, params.n_sequences, params.n_predict, params.n_gpu_blocks, params.n_cpu_blocks); + + const int64_t t_start_us = ggml_time_us(); + + while (true) { + const bool success = llama_paged_scheduler_prepare_batch(scheduler, &batch); + if (!success || batch.n_tokens == 0) { + break; + } + + LOG_DBG("prepared batch: batch.n_tokens=%d, n_batch=%d, n_sequences=%d, n_parallel=%d\n", batch.n_tokens, + params.n_batch, params.n_sequences, params.n_parallel); + GGML_ASSERT(batch.n_tokens <= params.n_batch && "batch exceeds n_batch"); + + if (llama_decode(ctx, batch) != 0) { + LOG_INF("%s: llama_decode failed\n", __func__); + break; + } + llama_synchronize(ctx); + + std::vector sampled_tokens; + std::vector stop_flags; + + const llama_paged_batch_info * info = llama_paged_scheduler_get_batch_info(scheduler); + GGML_ASSERT(info != nullptr && "llama_paged_batch_info is nullptr."); + for (int i = 0; i < info->n_seq; ++i) { + int32_t request_id = batch.seq_id[info->batch_offsets[i]][0]; + + auto it = samplers.find(request_id); + if (it == samplers.end()) { + samplers[request_id] = common_sampler_init(model, params.sampling); + it = samplers.find(request_id); + } + common_sampler * sampler = it->second; + + int32_t last_token_in_batch = info->batch_offsets[i] + info->batch_lens[i] - 1; + llama_token next_token = common_sampler_sample(sampler, ctx, last_token_in_batch); + common_sampler_accept(sampler, next_token, /*accept_grammar=*/true); + sampled_tokens.push_back(next_token); + accumulated_responses[request_id] += common_token_to_piece(ctx, next_token); + + llama_paged_seq_state state = {}; + llama_paged_scheduler_get_seq_state(scheduler, request_id, &state); + + LOG_DBG("[Request %d] token=%s, decoded=%d, ttft=%.1fms\n", request_id, + common_token_to_piece(ctx, next_token).c_str(), state.n_decoded, + state.n_decoded == 1 ? (state.t_first_token_us - state.t_arrival_us) / 1000.0f : 0.0f); + + bool stop = llama_vocab_is_eog(vocab, next_token) || state.n_decoded >= params.n_predict; + stop_flags.push_back(stop ? 1 : 0); + + if (stop) { + const int64_t t_finished = ggml_time_us(); // microsecs + + request_result r; + r.request_id = request_id; + r.n_prompt = state.n_prompt; + r.n_decoded = state.n_decoded; + r.ttft_ms = (state.t_first_token_us - state.t_arrival_us) / 1000.0f; + r.e2e_ms = (t_finished - state.t_arrival_us) / 1000.0f; + r.tps = state.n_decoded > 0 && t_finished > state.t_first_token_us ? + state.n_decoded / ((t_finished - state.t_first_token_us) / 1e6f) : + 0.0f; + r.tpot_ms = state.n_decoded > 1 && t_finished > state.t_first_token_us ? + (t_finished - state.t_first_token_us) / 1000.0f / (state.n_decoded - 1) : + 0.0f; + r.response = accumulated_responses.count(request_id) ? accumulated_responses[request_id] : ""; + results.push_back(r); + + LOG_DBG("[Request %d] finished. decoded=%d tokens, tps=%.1f\n", request_id, state.n_decoded, + state.n_decoded / (r.e2e_ms / 1000.0f)); + + accumulated_responses.erase(request_id); + common_sampler_free(samplers[request_id]); + samplers.erase(request_id); + } + } + + llama_paged_scheduler_update(scheduler, &batch, sampled_tokens.data(), stop_flags.data()); + } + + LOG_INF("%s: Finished paged example.\n", __func__); + + // For easy visualization of output + std::sort(results.begin(), results.end(), + [](const request_result & a, const request_result & b) { return a.request_id < b.request_id; }); + + LOG_INF("%s: Paged KV cache outputs:\n", __func__); + for (const auto & res : results) { + LOG_INF("Request Id: %d:%s<\\s>\n", res.request_id, res.response.c_str()); + } + + // Summary + const float elapsed_s = (ggml_time_us() - t_start_us) / 1e6f; + + int32_t total_prompt_tokens = 0; + int32_t total_decoded_tokens = 0; + float sum_ttft_ms = 0.f; + float sum_tpot_ms = 0.f; + float sum_e2e_ms = 0.f; + float min_ttft_ms = FLT_MAX; + float max_ttft_ms = 0.f; + float min_tps = FLT_MAX; + float max_tps = 0.f; + + for (const auto & r : results) { + total_prompt_tokens += r.n_prompt; + total_decoded_tokens += r.n_decoded; + sum_ttft_ms += r.ttft_ms; + sum_tpot_ms += r.tpot_ms; + sum_e2e_ms += r.e2e_ms; + min_ttft_ms = std::min(min_ttft_ms, r.ttft_ms); + max_ttft_ms = std::max(max_ttft_ms, r.ttft_ms); + min_tps = std::min(min_tps, r.tps); + max_tps = std::max(max_tps, r.tps); + } + + const int32_t n_results = (int32_t) results.size(); + const float avg_ttft_ms = n_results > 0 ? sum_ttft_ms / n_results : 0.f; + const float avg_tpot_ms = n_results > 0 ? sum_tpot_ms / n_results : 0.f; + const float avg_e2e_ms = n_results > 0 ? sum_e2e_ms / n_results : 0.f; + const float agg_tps = elapsed_s > 0 ? total_decoded_tokens / elapsed_s : 0.f; + + LOG_INF("\n"); + LOG_INF("=== Paged KV Cache Summary ===\n"); + LOG_INF(" n_sequences : %d\n", params.n_sequences); + LOG_INF(" n_predict : %d\n", params.n_predict); + LOG_INF(" n_batch : %d\n", params.n_batch); + LOG_INF(" n_gpu_blocks : %d\n", params.n_gpu_blocks); + LOG_INF(" n_cpu_blocks : %d\n", params.n_cpu_blocks); + LOG_INF(" total elapsed : %.2f s\n", elapsed_s); + LOG_INF(" total prompt tokens : %d\n", total_prompt_tokens); + LOG_INF(" total decoded tokens : %d\n", total_decoded_tokens); + LOG_INF(" aggregate tps : %.2f tokens/s\n", agg_tps); + LOG_INF(" --- per-request latency ---\n"); + LOG_INF(" ttft avg / min / max : %.1f / %.1f / %.1f ms\n", avg_ttft_ms, min_ttft_ms, max_ttft_ms); + LOG_INF(" tpot avg : %.1f ms/token\n", avg_tpot_ms); + LOG_INF(" e2e avg : %.1f ms\n", avg_e2e_ms); + LOG_INF(" tps min / max : %.1f / %.1f tokens/s\n", min_tps, max_tps); + LOG_INF("==============================\n"); + + // Clean-up + for (auto & [rid, s] : samplers) { + common_sampler_free(s); + } + llama_paged_scheduler_free(scheduler); + llama_backend_free(); + + return 0; +} diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 703e37831361..4844060155f9 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -578,6 +578,8 @@ extern "C" { GGML_OP_GLU, GGML_OP_COUNT, + + GGML_OP_PAGED_ATTN, }; enum ggml_unary_op { @@ -2824,6 +2826,21 @@ extern "C" { GGML_API void ggml_threadpool_params_init (struct ggml_threadpool_params * p, int n_threads); GGML_API bool ggml_threadpool_params_match (const struct ggml_threadpool_params * p0, const struct ggml_threadpool_params * p1); + GGML_API struct ggml_tensor * ggml_paged_attn(struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k_new, + struct ggml_tensor * v_new, + struct ggml_tensor * k_cache, + struct ggml_tensor * v_cache, + struct ggml_tensor * block_table, + struct ggml_tensor * write_slots, + struct ggml_tensor * context_lens, + struct ggml_tensor * batch_offsets, + struct ggml_tensor * batch_lens, + float scale, + int block_size, + int max_blocks); + #ifdef __cplusplus } #endif diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 2b3eb5b5ce65..c38bf2f19097 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2097,6 +2097,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { // nop } break; + case GGML_OP_PAGED_ATTN: + { + ggml_compute_forward_paged_attn(params, tensor); + } break; case GGML_OP_COUNT: { GGML_ABORT("fatal error"); @@ -2427,6 +2431,10 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { { GGML_ABORT("fatal error"); } + case GGML_OP_PAGED_ATTN: + { + n_tasks = 0; + } break; default: { fprintf(stderr, "%s: op not implemented: ", __func__); diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index a9bc21da6f0f..52b85fa643cd 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -11212,3 +11212,192 @@ void ggml_compute_forward_opt_step_sgd(const ggml_compute_params * params, ggml_ } } } + +void ggml_compute_forward_paged_attn(const ggml_compute_params * params, ggml_tensor * dst) { + // Single threaded reference + if (params->ith != 0) { + return; + } + + static bool log_warning = false; + if (!log_warning) { + log_warning = true; + GGML_LOG_WARN( + "%s: running CPU reference implementation of paged attention. This is for correctness validation only and " + "is not optimized. See docs/paged-attention.md for details.\n", + __func__); + } + + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k_new = dst->src[1]; + const ggml_tensor * v_new = dst->src[2]; + const ggml_tensor * kv_cache = dst->src[3]; // for reads + ggml_tensor * kv_cache_mut = dst->src[3]; // for writes + const ggml_tensor * block_table = dst->src[5]; + const ggml_tensor * write_slots = dst->src[6]; + const ggml_tensor * ctx_lens = dst->src[7]; + const ggml_tensor * batch_offsets = dst->src[8]; + const ggml_tensor * batch_lens = dst->src[9]; + + const float * op_params_f = (const float *) (dst->op_params); + const float scale = op_params_f[0]; + const int block_size = ((const int32_t *) (op_params_f + 1))[0]; + const int max_blocks = ((const int32_t *) (op_params_f + 2))[0]; + + const int head_dim = q->ne[0]; + const int n_heads = q->ne[1]; + const int n_seq = batch_lens->ne[0]; + const int n_heads_kv = k_new->ne[1]; + + GGML_ASSERT(block_size != 0 && "block_size cannot be 0."); + GGML_ASSERT(n_heads != 0 && "n_head cannot be 0."); + GGML_ASSERT(n_heads_kv != 0 && "n_head_kv cannot be 0."); + GGML_ASSERT(n_heads % n_heads_kv == 0 && "n_heads must be divisible by n_head_kv."); + + const size_t stride_token = kv_cache->nb[1] / sizeof(ggml_fp16_t); + const size_t stride_head = kv_cache->nb[2] / sizeof(ggml_fp16_t); + const size_t stride_block = kv_cache->nb[3] / sizeof(ggml_fp16_t); + + // Accessing tensors via backend API to make the CPU reference implementation backend agnostic + std::vector q_host(ggml_nelements(q)); + std::vector k_host(ggml_nelements(k_new)); + std::vector v_host(ggml_nelements(v_new)); + std::vector block_table_host(ggml_nelements(block_table)); + std::vector slots_host(ggml_nelements(write_slots)); + std::vector ctx_lens_host(ggml_nelements(ctx_lens)); + std::vector batch_offsets_host(ggml_nelements(batch_offsets)); + std::vector batch_lens_host(ggml_nelements(batch_lens)); + std::vector out_host(ggml_nelements(dst)); + + ggml_backend_tensor_get(q, q_host.data(), 0, ggml_nbytes(q)); + ggml_backend_tensor_get(k_new, k_host.data(), 0, ggml_nbytes(k_new)); + ggml_backend_tensor_get(v_new, v_host.data(), 0, ggml_nbytes(v_new)); + ggml_backend_tensor_get(block_table, block_table_host.data(), 0, ggml_nbytes(block_table)); + ggml_backend_tensor_get(write_slots, slots_host.data(), 0, ggml_nbytes(write_slots)); + ggml_backend_tensor_get(ctx_lens, ctx_lens_host.data(), 0, ggml_nbytes(ctx_lens)); + ggml_backend_tensor_get(batch_offsets, batch_offsets_host.data(), 0, ggml_nbytes(batch_offsets)); + ggml_backend_tensor_get(batch_lens, batch_lens_host.data(), 0, ggml_nbytes(batch_lens)); + + const float * q_data = q_host.data(); + const float * k_data = k_host.data(); + const float * v_data = v_host.data(); + const int32_t * block_table_data = block_table_host.data(); + const int32_t * slots_data = slots_host.data(); + const int32_t * ctx_lens_data = ctx_lens_host.data(); + const int32_t * batch_offsets_data = batch_offsets_host.data(); + const int32_t * batch_lens_data = batch_lens_host.data(); + float * out_data = out_host.data(); // use host buffer throughout + + // We use staging buffers for KV cache access to make it agnostic to where + // the KV cache is allocated. + const size_t head_bytes = (size_t) head_dim * sizeof(ggml_fp16_t); + std::vector staging_k(head_dim); + std::vector staging_v(head_dim); + std::vector staging_write_k(head_dim); + std::vector staging_write_v(head_dim); + + // Write to KV cache + for (int seq = 0; seq < n_seq; ++seq) { + const int seq_start = batch_offsets_data[seq]; + const int num_tokens = batch_lens_data[seq]; + + for (int i = 0; i < num_tokens; ++i) { + const int token_batch_idx = seq_start + i; + const int target_slot = slots_data[token_batch_idx]; + const int block_id = target_slot / block_size; + const int token_in_block = target_slot % block_size; + + for (int h_id = 0; h_id < n_heads_kv; ++h_id) { + const size_t k_cache_byte_offset = ((size_t) block_id * stride_block + (size_t) h_id * stride_head + + (size_t) token_in_block * stride_token) * + sizeof(ggml_fp16_t); + const size_t v_cache_byte_offset = + ((size_t) block_id * stride_block + (size_t) (n_heads_kv + h_id) * stride_head + + (size_t) token_in_block * stride_token) * + sizeof(ggml_fp16_t); + const size_t input_offset = (size_t) token_batch_idx * n_heads_kv * head_dim + (size_t) h_id * head_dim; + + for (int d_id = 0; d_id < head_dim; ++d_id) { + staging_write_k[d_id] = GGML_FP32_TO_FP16(k_data[input_offset + d_id]); + staging_write_v[d_id] = GGML_FP32_TO_FP16(v_data[input_offset + d_id]); + } + ggml_backend_tensor_set((ggml_tensor *) kv_cache_mut, staging_write_k.data(), k_cache_byte_offset, + head_bytes); + ggml_backend_tensor_set((ggml_tensor *) kv_cache_mut, staging_write_v.data(), v_cache_byte_offset, + head_bytes); + } + } + } + + // Decode + for (int seq = 0; seq < n_seq; ++seq) { + const int seq_start = batch_offsets_data[seq]; + const int num_new_tokens = batch_lens_data[seq]; + const int ctx_len = ctx_lens_data[seq]; + + for (int i = 0; i < num_new_tokens; ++i) { + const int token_batch_idx = seq_start + i; + const int q_pos = (ctx_len - num_new_tokens) + i; + const int num_blocks = (q_pos / block_size) + 1; + + for (int h_id = 0; h_id < n_heads; ++h_id) { + const int kv_h = h_id / (n_heads / n_heads_kv); + + const float * q_vec = q_data + token_batch_idx * n_heads * head_dim + h_id * head_dim; + + float qk_max = -FLT_MAX; + float exp_sum = 0.0f; + std::vector acc(head_dim, 0.0f); + std::fill(acc.begin(), acc.end(), 0.0f); + + for (int bid = 0; bid < num_blocks; ++bid) { + const int physical_block = block_table_data[seq * max_blocks + bid]; + const int start_token = bid * block_size; + const int end_token = + ((start_token + block_size) < (q_pos + 1)) ? start_token + block_size : q_pos + 1; + + for (int tok = start_token; tok < end_token; ++tok) { + const int token_in_block = tok % block_size; + const size_t k_byte_offset = + ((size_t) physical_block * stride_block + (size_t) kv_h * stride_head + + (size_t) token_in_block * stride_token) * + sizeof(ggml_fp16_t); + const size_t v_byte_offset = + ((size_t) physical_block * stride_block + (size_t) (n_heads_kv + kv_h) * stride_head + + (size_t) token_in_block * stride_token) * + sizeof(ggml_fp16_t); + + // Fetch K and V from cache (this might involve device to host transfers) + ggml_backend_tensor_get(kv_cache, staging_k.data(), k_byte_offset, head_bytes); + ggml_backend_tensor_get(kv_cache, staging_v.data(), v_byte_offset, head_bytes); + + // QK dot product + float qk = 0.0f; + for (int d_id = 0; d_id < head_dim; ++d_id) { + qk += q_vec[d_id] * GGML_FP16_TO_FP32(staging_k[d_id]); + } + qk *= scale; + + // Online softmax update + const float qk_max_new = fmaxf(qk_max, qk); + const float exp_old = expf(qk_max - qk_max_new); + const float exp_new = expf(qk - qk_max_new); + + exp_sum = exp_sum * exp_old + exp_new; + for (int d_id = 0; d_id < head_dim; ++d_id) { + acc[d_id] = acc[d_id] * exp_old + exp_new * GGML_FP16_TO_FP32(staging_v[d_id]); + } + qk_max = qk_max_new; + } + } + // Write output + const size_t out_idx = (size_t) token_batch_idx * n_heads * head_dim + (size_t) h_id * head_dim; + for (int d_id = 0; d_id < head_dim; ++d_id) { + out_data[out_idx + d_id] = acc[d_id] / (exp_sum + 1e-6f); + } + } + } + } + // Write output back to dst (might involve host to device transfer) + ggml_backend_tensor_set(dst, out_host.data(), 0, ggml_nbytes(dst)); +} diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 3fa1443abc48..643f22255634 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -112,6 +112,7 @@ void ggml_compute_forward_cross_entropy_loss_back(const struct ggml_compute_para void ggml_compute_forward_opt_step_adamw(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_mul_mat(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_opt_step_sgd(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_paged_attn(const struct ggml_compute_params * params, struct ggml_tensor * dst); #ifdef __cplusplus } #endif diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index fbe0fa06242c..a05476881b70 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -61,6 +61,7 @@ #include "ggml-cuda/tri.cuh" #include "ggml-cuda/cumsum.cuh" #include "ggml-cuda/fill.cuh" +#include "ggml-cuda/pagedattn.cuh" #include "ggml.h" #include @@ -2952,6 +2953,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_FILL: ggml_cuda_op_fill(ctx, dst); break; + case GGML_OP_PAGED_ATTN: + ggml_cuda_op_paged_attn(ctx, dst); + break; default: return false; } @@ -5212,8 +5216,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_TRI: case GGML_OP_DIAG: case GGML_OP_SOLVE_TRI: + case GGML_OP_PAGED_ATTN: return true; - default: return false; } diff --git a/ggml/src/ggml-cuda/pagedattn.cu b/ggml/src/ggml-cuda/pagedattn.cu new file mode 100644 index 000000000000..c759b1797a4e --- /dev/null +++ b/ggml/src/ggml-cuda/pagedattn.cu @@ -0,0 +1,208 @@ +#include "pagedattn.cuh" + +__device__ __forceinline__ float block_reduce_sum_full(float val, float * __restrict__ smem, int tid, int head_dim) { + const int lane = tid & 31; + const int warp_id = tid >> 5; + const int n_warps = (head_dim + 31) >> 5; + + // warp-level reduce + for (int offset = 16; offset > 0; offset >>= 1) { + val += __shfl_down_sync(0xffffffffu, val, offset); + } + // Each warp's lane 0 depositits partial sum into shared memory + if (lane == 0) { + smem[warp_id] = val; + } + __syncthreads(); + + // First warp reduces the per-warp partial + float warp_val = (tid < n_warps) ? smem[tid] : 0.0f; + if (warp_id == 0) { + for (int offset = 16; offset > 0; offset >>= 1) { + warp_val += __shfl_down_sync(0xffffffffu, warp_val, offset); + } + if (lane == 0) { + smem[0] = warp_val; + } + } + __syncthreads(); + return smem[0]; // this will be identical in every thread +} + +__global__ void paged_attention_write_kernel(const float * __restrict__ k_new, // [batch_size, n_heads_kv, head_dim] + const float * __restrict__ v_new, // [batch_size, n_heads_kv, head_dim] + half * __restrict__ kv_cache, // The paged cache + const int * __restrict__ write_slots, // Global slot index for each token + const int * __restrict__ batch_offsets, + const int * __restrict__ batch_lens, + const size_t stride_token, // Elements between tokens in a block (nb1) + const size_t stride_head, // Elements between heads (nb2) + const size_t stride_block, // Elements between physical blocks (nb3) + const int n_heads_kv, + const int block_size) { + const int head_idx = blockIdx.x; // 0 to n_heads_kv - 1 + const int seq_idx = blockIdx.y; + const int tid = threadIdx.x; // 0 to head_dim - 1 + const int head_dim = blockDim.x; + + const int seq_start = batch_offsets[seq_idx]; + const int num_tokens = batch_lens[seq_idx]; + + for (int i = 0; i < num_tokens; ++i) { + const int token_batch_idx = seq_start + i; + const int target_slot = write_slots[token_batch_idx]; + + // Map slot to block and internal offset + const int block_id = target_slot / block_size; + const int token_in_block = target_slot % block_size; + + // K is at head_idx, V is at n_heads_kv + head_idx + const size_t k_cache_idx = (size_t) block_id * stride_block + (size_t) head_idx * stride_head + + (size_t) token_in_block * stride_token + tid; + const size_t v_cache_idx = (size_t) block_id * stride_block + (size_t) (n_heads_kv + head_idx) * stride_head + + (size_t) token_in_block * stride_token + tid; + + // Input offset: [token][head][dim] + const size_t input_off = (size_t) token_batch_idx * n_heads_kv * head_dim + (size_t) head_idx * head_dim + tid; + + kv_cache[k_cache_idx] = __float2half(k_new[input_off]); + kv_cache[v_cache_idx] = __float2half(v_new[input_off]); + } +} + +__global__ void paged_attention_decode_kernel(const float * __restrict__ q, + const half * __restrict__ kv_cache, + const int * __restrict__ block_table, + const int * __restrict__ context_lens, + const int * __restrict__ batch_offsets, + const int * __restrict__ batch_lens, + const size_t stride_token, + const size_t stride_head, + const size_t stride_block, + const int n_heads_kv, + const int block_size, + const int max_blocks, + const float scale, + float * __restrict__ out) { + extern __shared__ float smem[]; + + const int head_idx = blockIdx.x; + const int seq_idx = blockIdx.y; + const int tid = threadIdx.x; + + const int n_heads = gridDim.x; + const int head_dim = blockDim.x; + + const int kv_head_idx = head_idx / (n_heads / n_heads_kv); + + const int seq_start = batch_offsets[seq_idx]; + const int num_new_tokens = batch_lens[seq_idx]; + + for (int i = 0; i < num_new_tokens; i++) { + const int token_batch_idx = seq_start + i; + + float q_val = q[(size_t) token_batch_idx * n_heads * head_dim + (size_t) head_idx * head_dim + tid] * scale; + + float qk_max = -FLT_MAX; + float exp_sum = 0.0f; + float acc = 0.0f; + + const int ctx_len = context_lens[seq_idx]; + const int q_pos = (ctx_len - num_new_tokens) + i; + const int num_blocks = (q_pos / block_size) + 1; + + for (int bid = 0; bid < num_blocks; bid++) { + const int physical_block = block_table[seq_idx * max_blocks + bid]; + const int start_token = bid * block_size; + const int end_token = min(start_token + block_size, q_pos + 1); + + for (int token = start_token; token < end_token; ++token) { + const int token_in_block = token % block_size; + + const size_t k_idx = + tid + token_in_block * stride_token + kv_head_idx * stride_head + physical_block * stride_block; + + const size_t v_idx = tid + token_in_block * stride_token + (n_heads_kv + kv_head_idx) * stride_head + + physical_block * stride_block; + + float k_val = __half2float(kv_cache[k_idx]); + float v_val = __half2float(kv_cache[v_idx]); + + // Calculate full dot product and return the same scalar in every thread + const float qk = block_reduce_sum_full(q_val * k_val, smem, tid, head_dim); + + // Online softmax update + const float qk_max_new = fmaxf(qk_max, qk); + const float exp_old = __expf(qk_max - qk_max_new); + const float exp_new = __expf(qk - qk_max_new); + + exp_sum = exp_sum * exp_old + exp_new; + acc = acc * exp_old + exp_new * v_val; + qk_max = qk_max_new; + } + } + + const int out_idx = (size_t) token_batch_idx * n_heads * head_dim + (size_t) head_idx * head_dim + tid; + + out[out_idx] = acc / (exp_sum + 1e-6f); + } +} + +void ggml_cuda_op_paged_attn(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k_new = dst->src[1]; + const ggml_tensor * v_new = dst->src[2]; + const ggml_tensor * kv_cache = dst->src[3]; // KV interleaved layout + const ggml_tensor * block_table = dst->src[5]; + const ggml_tensor * write_slots = dst->src[6]; + const ggml_tensor * context_lens = dst->src[7]; + const ggml_tensor * batch_offsets = dst->src[8]; + const ggml_tensor * batch_lens = dst->src[9]; + + const float * op_params_f = (const float *) (dst->op_params); + const float scale = op_params_f[0]; + const int block_size = ((const int32_t *) (op_params_f + 1))[0]; + const int max_blocks = ((const int32_t *) (op_params_f + 2))[0]; + + const int head_dim = q->ne[0]; + const int n_heads = q->ne[1]; + const int n_seq = batch_lens->ne[0]; + const int n_heads_kv = k_new->ne[1]; + + GGML_ASSERT(n_heads != 0 && "n_head cannot be 0."); + GGML_ASSERT(n_heads_kv != 0 && "n_heads_kv cannot be 0."); + GGML_ASSERT(head_dim <= 1024 && "head_dim exceeds maximum supported (1024)"); + GGML_ASSERT(n_heads % n_heads_kv == 0 && "n_heads must be divisible by n_heads_kv"); + + // Extracting strides + const size_t stride_token = kv_cache->nb[1] / sizeof(half); + const size_t stride_head = kv_cache->nb[2] / sizeof(half); + const size_t stride_block = kv_cache->nb[3] / sizeof(half); + + dim3 block_dims(head_dim); // one thread per dimension of head + dim3 grid_dims(n_heads, n_seq); // one block per head per sequence + + // Write kernel - Grid (n_heads_kv, n_seq), Block (head_dim) + paged_attention_write_kernel<<>>( + (const float *) k_new->data, (const float *) v_new->data, (half *) kv_cache->data, + (const int *) write_slots->data, (const int *) batch_offsets->data, (const int *) batch_lens->data, + stride_token, stride_head, stride_block, n_heads_kv, block_size); + + // Shared memory + const size_t n_warps = ((size_t) head_dim + 31) / 32; + const size_t smem_bytes = n_warps * sizeof(float); + + // Manually request extended shared memory if needed (>48 KB) + // https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/compute-capabilities.html + if (smem_bytes > 48 * 1024) { + GGML_ASSERT(smem_bytes <= 96 * 1024 && "smem exceeds 96KB limit"); + CUDA_CHECK(cudaFuncSetAttribute(paged_attention_decode_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + (int) smem_bytes)); + } + + // Read kernel - Grid (n_heads_kv, n_seq), Block (head_dim) + paged_attention_decode_kernel<<>>( + (const float *) q->data, (const half *) kv_cache->data, (const int *) block_table->data, + (const int *) context_lens->data, (const int *) batch_offsets->data, (const int *) batch_lens->data, + stride_token, stride_head, stride_block, n_heads_kv, block_size, max_blocks, scale, (float *) dst->data); +} diff --git a/ggml/src/ggml-cuda/pagedattn.cuh b/ggml/src/ggml-cuda/pagedattn.cuh new file mode 100644 index 000000000000..2345cde0559d --- /dev/null +++ b/ggml/src/ggml-cuda/pagedattn.cuh @@ -0,0 +1,30 @@ +#include "common.cuh" + +__global__ void paged_attention_write_kernel(const float * k_new, // [batch_size, n_heads_kv, head_dim] + const float * v_new, // [batch_size, n_heads_kv, head_dim] + half * kv_cache, // The paged cache + const int * write_slots, // Global slot index for each token + const int * batch_offsets, + const int * batch_lens, + const size_t stride_token, // Elements between tokens in a block (nb1) + const size_t stride_head, // Elements between heads (nb2) + const size_t stride_block, // Elements between physical blocks (nb3) + const int n_heads_kv, + const int block_size); + +__global__ void paged_attention_decode_kernel(const float * __restrict__ q, + const half * __restrict__ kv_cache, + const int * __restrict__ block_table, + const int * __restrict__ context_lens, + const int * __restrict__ batch_offsets, + const int * __restrict__ batch_lens, + const size_t stride_token, + const size_t stride_head, + const size_t stride_block, + const int n_heads_kv, + const int block_size, + const int max_blocks, + const float scale, + float * __restrict__ out); + +void ggml_cuda_op_paged_attn(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 54d3eae3e4da..df63b8034996 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -7758,3 +7758,42 @@ bool ggml_threadpool_params_match(const struct ggml_threadpool_params * p0, cons if (p0->strict_cpu != p1->strict_cpu ) return false; return memcmp(p0->cpumask, p1->cpumask, GGML_MAX_N_THREADS) == 0; } + +struct ggml_tensor * ggml_paged_attn( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k_new, + struct ggml_tensor * v_new, + struct ggml_tensor * k_cache, + struct ggml_tensor * v_cache, + struct ggml_tensor * block_table, + struct ggml_tensor * write_slots, + struct ggml_tensor * context_lens, + struct ggml_tensor * batch_offsets, + struct ggml_tensor * batch_lens, + float scale, + int block_size, + int max_blocks) { + + struct ggml_tensor * result = ggml_new_tensor(ctx, q->type, ggml_n_dims(q), q->ne); + result->op = GGML_OP_PAGED_ATTN; + result->src[0] = q; + result->src[1] = k_new; + result->src[2] = v_new; + result->src[3] = k_cache; + result->src[4] = v_cache; + result->src[5] = block_table; + result->src[6] = write_slots; + result->src[7] = context_lens; + result->src[8] = batch_offsets; + result->src[9] = batch_lens; + + // Storing hyperparams directly in op_params + float * op_params_f = (float *)result->op_params; + op_params_f[0] = scale; + int32_t * op_params_i = (int32_t *)(op_params_f + 1); + op_params_i[0] = block_size; + op_params_i[1] = max_blocks; + + return result; +} diff --git a/include/llama.h b/include/llama.h index eb8698140978..37634d7c886c 100644 --- a/include/llama.h +++ b/include/llama.h @@ -63,6 +63,8 @@ extern "C" { struct llama_context; struct llama_sampler; + struct llama_paged_scheduler; + typedef struct llama_memory_i * llama_memory_t; typedef int32_t llama_pos; @@ -243,6 +245,22 @@ extern "C" { int8_t * logits; // TODO: rename this to "output" } llama_batch; + // CPU-side metadata produced by the paged scheduler and consumed by + // llama_kv_cache_paged_context during graph build. These arrays are owned + // by the scheduler and must remain valid until the next scheduler step + // clears them. + typedef struct llama_paged_batch_info { + int32_t n_blocks_per_seq = 0; + int32_t n_seq = 0; + int32_t n_tokens = 0; + + int32_t * write_slots = NULL; // [n_tokens] + int32_t * block_table = NULL; // [n_seq * n_blocks_per_seq] + int32_t * context_lens = NULL; // [n_seq] + int32_t * batch_offsets = NULL; // [n_seq] + int32_t * batch_lens = NULL; // [n_seq] + } llama_paged_batch_info; + enum llama_model_kv_override_type { LLAMA_KV_OVERRIDE_TYPE_INT, LLAMA_KV_OVERRIDE_TYPE_FLOAT, @@ -375,6 +393,14 @@ extern "C" { // try to disable when n_seq_max > 1 for improved performance when the sequences do not share a large prefix // ref: https://github.com/ggml-org/llama.cpp/pull/14363 + // Paged KV cache (experimental) + // Opt-in block-indexing KV cache with continuous-batching scheduling + bool kv_paged; // enable paged KV cache + uint32_t block_size; // tokens per physical KV block + uint32_t n_gpu_blocks; // GPU block pool size + uint32_t n_cpu_blocks; // CPU block pool size for swap-out + float kv_paged_watermark; // percentage of GPU blocks reserved as safety margin [0, 0.1) + // [EXPERIMENTAL] // backend sampler chain configuration (make sure the caller keeps the sampler chains alive) // note: the samplers must be sampler chains (i.e. use llama_sampler_chain_init) @@ -1558,6 +1584,53 @@ extern "C" { ggml_opt_epoch_callback callback_train, ggml_opt_epoch_callback callback_eval); + // + // Paged inference + // + struct llama_paged_seq_state { + int32_t request_id; + int32_t n_prompt; + int32_t n_decoded; + int32_t n_past; + int64_t t_arrival_us; + int64_t t_first_token_us; + }; + + typedef void (*llama_paged_on_finish_cb)(int32_t request_id, + const llama_token * tokens, + int32_t n_tokens, + void * user_data); + LLAMA_API struct llama_paged_scheduler * llama_paged_scheduler_init(struct llama_context * ctx); + LLAMA_API void llama_paged_scheduler_free(struct llama_paged_scheduler * sched); + + // Queueing and stepping. + LLAMA_API bool llama_paged_scheduler_add_request(struct llama_paged_scheduler * sched, + const llama_token * tokens, + int32_t n_tokens, + int32_t request_id); + + LLAMA_API bool llama_paged_scheduler_prepare_batch(struct llama_paged_scheduler * sched, + struct llama_batch * batch); + + LLAMA_API void llama_paged_scheduler_update(struct llama_paged_scheduler * sched, + struct llama_batch * batch, + const llama_token * tokens, + const int8_t * stop_flags); + + // Introspection. + LLAMA_API bool llama_paged_scheduler_get_seq_state(struct llama_paged_scheduler * sched, + int32_t request_id, + struct llama_paged_seq_state * out_state); + + // Returns the current batch's paged routing metadata. Valid until the next + // call to llama_paged_scheduler_prepare_batch on this scheduler. Do not free. + LLAMA_API const struct llama_paged_batch_info * llama_paged_scheduler_get_batch_info( + const struct llama_paged_scheduler * sched); + + // Optional finish callback. + LLAMA_API void llama_paged_scheduler_set_on_finish(struct llama_paged_scheduler * sched, + llama_paged_on_finish_cb cb, + void * user_data); #ifdef __cplusplus } #endif diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 7b1fcfca0ada..72400c5c852b 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -24,6 +24,10 @@ add_library(llama llama-io.cpp llama-kv-cache.cpp llama-kv-cache-iswa.cpp + llama-kv-cache-paged.cpp + llama-block-manager.cpp + llama-paged-scheduler.cpp + llama-paged-scheduler-impl.cpp llama-memory.cpp llama-memory-hybrid.cpp llama-memory-hybrid-iswa.cpp diff --git a/src/llama-block-manager.cpp b/src/llama-block-manager.cpp new file mode 100644 index 000000000000..4816da7a7a91 --- /dev/null +++ b/src/llama-block-manager.cpp @@ -0,0 +1,109 @@ +#include "llama-block-manager.h" + +#include "llama-impl.h" + +#include + +void llama_block_manager::init(uint32_t n_gpu, uint32_t n_cpu, float watermark) { + LLAMA_LOG_INFO("%s: Block manager initialized: n_free_gpu_blocks=%d, n_free_cpu_blocks=%d\n", __func__, n_gpu, + n_cpu); + total_num_gpu_blocks = n_gpu; + total_num_cpu_blocks = n_cpu; + + watermark_gpu_safety_num_blocks = std::ceil(total_num_gpu_blocks * watermark); + watermark_cpu_safety_num_blocks = std::ceil(total_num_cpu_blocks * watermark); + + gpu_registry.resize(n_gpu); + for (uint32_t i = 0; i < n_gpu; ++i) { + gpu_registry[i].id = i; + gpu_registry[i].is_gpu = true; + free_gpu_ids.push_back(i); + } + + cpu_registry.resize(n_cpu); + for (uint32_t i = 0; i < n_cpu; ++i) { + cpu_registry[i].id = i + total_num_gpu_blocks; + cpu_registry[i].is_gpu = false; + free_cpu_ids.push_back(i + total_num_gpu_blocks); + } +} + +size_t llama_block_manager::n_free_gpu_blocks() const { + return free_gpu_ids.size(); +} + +size_t llama_block_manager::n_free_cpu_blocks() const { + return free_cpu_ids.size(); +} + +bool llama_block_manager::has_free_gpu_blocks(uint32_t num_requested_blocks) const { + size_t curr_free_gpus = free_gpu_ids.size(); + if (curr_free_gpus < watermark_gpu_safety_num_blocks) { + return false; + } + return num_requested_blocks <= (curr_free_gpus - watermark_gpu_safety_num_blocks); +} + +bool llama_block_manager::has_free_cpu_blocks(uint32_t num_requested_blocks) const { + size_t curr_free_cpus = free_cpu_ids.size(); + if (curr_free_cpus < watermark_cpu_safety_num_blocks) { + return false; + } + return num_requested_blocks <= (curr_free_cpus - watermark_cpu_safety_num_blocks); +} + +llama_block_manager::physical_block_ids llama_block_manager::checkout_gpu_blocks(uint32_t num_blocks) { + physical_block_ids new_ids = {}; + if (num_blocks > free_gpu_ids.size()) { + return new_ids; + } + + new_ids.insert(new_ids.end(), std::make_move_iterator(free_gpu_ids.end() - num_blocks), + std::make_move_iterator(free_gpu_ids.end())); + free_gpu_ids.erase(free_gpu_ids.end() - num_blocks, free_gpu_ids.end()); + + for (const uint32_t & id : new_ids) { + gpu_registry[id].ref_count += 1; + } + return new_ids; +} + +llama_block_manager::physical_block_ids llama_block_manager::checkout_cpu_blocks(uint32_t num_blocks) { + physical_block_ids new_ids = {}; + if (num_blocks > free_cpu_ids.size()) { + return new_ids; + } + + new_ids.insert(new_ids.end(), std::make_move_iterator(free_cpu_ids.end() - num_blocks), + std::make_move_iterator(free_cpu_ids.end())); + free_cpu_ids.erase(free_cpu_ids.end() - num_blocks, free_cpu_ids.end()); + + for (const uint32_t & id : new_ids) { + cpu_registry[id - total_num_gpu_blocks].ref_count += 1; + } + return new_ids; +} + +void llama_block_manager::release_gpu_blocks(const physical_block_ids & freed_blocks_ids) { + for (const uint32_t & id : freed_blocks_ids) { + gpu_registry[id].ref_count -= 1; + if (gpu_registry[id].ref_count <= 0) { + gpu_registry[id].ref_count = 0; + free_gpu_ids.push_back(id); + } + } +} + +void llama_block_manager::release_cpu_blocks(const physical_block_ids & freed_blocks_ids) { + for (const uint32_t & id : freed_blocks_ids) { + cpu_registry[id - total_num_gpu_blocks].ref_count -= 1; + if (cpu_registry[id - total_num_gpu_blocks].ref_count <= 0) { + cpu_registry[id - total_num_gpu_blocks].ref_count = 0; + free_cpu_ids.push_back(id); + } + } +} + +bool llama_block_manager::is_gpu(uint32_t block_id) const { + return block_id < total_num_gpu_blocks; +} diff --git a/src/llama-block-manager.h b/src/llama-block-manager.h new file mode 100644 index 000000000000..603bca7fb451 --- /dev/null +++ b/src/llama-block-manager.h @@ -0,0 +1,45 @@ +#pragma once + +#include +#include +#include + +class llama_block_manager { + struct physical_block { + uint32_t id = 0; + uint32_t ref_count = 0; + bool is_gpu = false; + }; + + using physical_block_bool = std::vector; + using physical_block_ids = std::vector; + + physical_block_bool gpu_registry; + physical_block_bool cpu_registry; + + physical_block_ids free_gpu_ids; // [0 to total_num_gpu_blocks - 1] + physical_block_ids free_cpu_ids; // [total_num_gpu_blocks, total_num_gpu_blocks + total_num_cpu_blocks] + + uint32_t watermark_gpu_safety_num_blocks; + uint32_t watermark_cpu_safety_num_blocks; + + uint32_t total_num_gpu_blocks; + uint32_t total_num_cpu_blocks; + + public: + void init(uint32_t n_gpu, uint32_t n_cpu, float watermark); + + size_t n_free_gpu_blocks() const; + size_t n_free_cpu_blocks() const; + + bool has_free_gpu_blocks(uint32_t num_requested_blocks) const; + bool has_free_cpu_blocks(uint32_t num_requested_blocks) const; + + physical_block_ids checkout_gpu_blocks(uint32_t num_blocks); + physical_block_ids checkout_cpu_blocks(uint32_t num_blocks); + + void release_gpu_blocks(const physical_block_ids & freed_blocks); + void release_cpu_blocks(const physical_block_ids & freed_blocks); + + bool is_gpu(uint32_t block) const; +}; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 8126249e1436..4081f06e1ac4 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -21,6 +21,36 @@ // llama_context // +namespace { +// Returns empty string on success, or an error message describing why +// paged KV cache cannot be used with the current model placement. +std::string validate_paged_kv_placement(const llama_model & model) { + std::set devs_used; + for (uint32_t il = 0; il < model.hparams.n_layer; ++il) { + ggml_backend_dev_t dev = model.dev_layer(il); + if (dev != nullptr) { + devs_used.insert(dev); + } + } + + if (devs_used.size() > 1) { + std::string dev_list; + for (auto * d : devs_used) { + if (!dev_list.empty()) dev_list += ", "; + dev_list += ggml_backend_dev_name(d); + } + return format( + "paged KV cache (--kv-paged) currently requires all model layers " + "to live on a single device, but the model is split across %zu " + "devices (%s). Re-run with `-sm none -mg ` to pin the " + "model to one device, or omit --kv-paged.", + devs_used.size(), dev_list.c_str()); + } + + return {}; +} +} // namespace + llama_context::llama_context( const llama_model & model, llama_context_params params) : @@ -164,6 +194,11 @@ llama_context::llama_context( cparams.op_offload = params.op_offload; cparams.kv_unified = params.kv_unified; + cparams.kv_paged = params.kv_paged; + cparams.block_size = params.block_size; + cparams.n_gpu_blocks = params.n_gpu_blocks; + cparams.n_cpu_blocks = params.n_cpu_blocks; + cparams.kv_paged_watermark = params.kv_paged_watermark; // initialized later cparams.pipeline_parallel = false; @@ -280,7 +315,28 @@ llama_context::llama_context( /*.swa_full =*/ params.swa_full, }; - memory.reset(model.create_memory(params_mem, cparams)); + ggml_backend_t gpu_handle = nullptr; + if (cparams.kv_paged) { + std::string err = validate_paged_kv_placement(model); + if (!err.empty()) { + LLAMA_LOG_ERROR("%s: %s\n", __func__, err.c_str()); + throw std::runtime_error(err); + } + // Find the first non-CPU backend to use as the primary GPU/Compute backend + for (auto & b : backends) { + if (ggml_backend_get_device(b.get()) && + ggml_backend_dev_type(ggml_backend_get_device(b.get())) != GGML_BACKEND_DEVICE_TYPE_CPU) { + gpu_handle = b.get(); + break; + } + } + // Fallback: If no GPU, use CPU for both + if (!gpu_handle) { + gpu_handle = backend_cpu; + } + } + + memory.reset(model.create_memory(params_mem, cparams, gpu_handle, backend_cpu)); } // init backends @@ -695,6 +751,10 @@ uint32_t llama_context::n_seq_max() const { return cparams.n_seq_max; } +uint32_t llama_context::block_size() const { + return cparams.block_size; +} + uint32_t llama_context::n_threads() const { return cparams.n_threads; } @@ -2916,6 +2976,11 @@ llama_context_params llama_context_default_params() { /*.op_offload =*/ true, /*.swa_full =*/ true, /*.kv_unified =*/ false, + /*.kv_paged =*/ false, + /*.block_size =*/ 16, + /*.n_gpu_blocks =*/ 0, + /*.n_cpu_blocks =*/ 0, + /*.kv_paged_watermark =*/ 0.05, /*.sampler =*/ nullptr, /*.n_sampler =*/ 0, }; diff --git a/src/llama-context.h b/src/llama-context.h index 53c705eaffce..715cc8df5faf 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -46,11 +46,12 @@ struct llama_context { ggml_backend_sched_t get_sched() const; - uint32_t n_ctx() const; - uint32_t n_ctx_seq() const; - uint32_t n_batch() const; - uint32_t n_ubatch() const; - uint32_t n_seq_max() const; + uint32_t n_ctx() const; + uint32_t n_ctx_seq() const; + uint32_t n_batch() const; + uint32_t n_ubatch() const; + uint32_t n_seq_max() const; + uint32_t block_size() const; uint32_t n_threads() const; uint32_t n_threads_batch() const; diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 9d359474132f..a01b6472bfc9 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -15,6 +15,11 @@ struct llama_cparams { int32_t n_threads; // number of threads to use for generation int32_t n_threads_batch; // number of threads to use for batch processing + uint32_t block_size; + uint32_t n_gpu_blocks; + uint32_t n_cpu_blocks; + float kv_paged_watermark; + float rope_freq_base; float rope_freq_scale; @@ -38,6 +43,7 @@ struct llama_cparams { bool warmup; bool op_offload; bool kv_unified; + bool kv_paged; bool pipeline_parallel; enum llama_pooling_type pooling_type; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 2ff23f87cf44..0fdf5d37d8a7 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -7,6 +7,7 @@ #include "llama-kv-cache.h" #include "llama-kv-cache-iswa.h" +#include "llama-kv-cache-paged.h" #include "llama-memory-hybrid.h" #include "llama-memory-hybrid-iswa.h" #include "llama-memory-recurrent.h" @@ -541,6 +542,34 @@ bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) { return res; } +void llm_graph_input_attn_kv_paged::set_input(const llama_ubatch* ubatch) { + GGML_ASSERT(ubatch != nullptr); + + if (paged_write_slots) { + ggml_backend_tensor_set(paged_write_slots, mctx->get_write_slots(), 0, ggml_nbytes(paged_write_slots)); + last_n_tokens = paged_write_slots->ne[0]; + } + if (paged_block_table) { + ggml_backend_tensor_set(paged_block_table, mctx->get_block_table(), 0, ggml_nbytes(paged_block_table)); + } + if (paged_context_lens) { + ggml_backend_tensor_set(paged_context_lens, mctx->get_context_lens(), 0, ggml_nbytes(paged_context_lens)); + } + if (paged_batch_offsets) { + ggml_backend_tensor_set(paged_batch_offsets, mctx->get_batch_offsets(), 0, ggml_nbytes(paged_batch_offsets)); + } + if (paged_batch_lens) { + ggml_backend_tensor_set(paged_batch_lens, mctx->get_batch_lens(), 0, ggml_nbytes(paged_batch_lens)); + } +} + +bool llm_graph_input_attn_kv_paged::can_reuse(const llm_graph_params & /*params*/) { + // In paged KV cache, we can 'never' re-use the graph. Because we have write_slots + // which encode the complete physical memory mapping for a specific batch at a specific + // step. + return false; +} + void llm_graph_input_attn_cross::set_input(const llama_ubatch * ubatch) { GGML_ASSERT(cross_kq_mask); @@ -1929,6 +1958,36 @@ ggml_tensor * llm_graph_context::build_pos_bias(ggml_tensor * pos_bucket, ggml_t return pos_bias; } +ggml_tensor * llm_graph_context::build_attn_mha_paged( + ggml_tensor * q, // [n_embd_head, n_head, n_tokens] + ggml_tensor * k_cur, // [n_embd_head, n_head_kv, n_tokens] + ggml_tensor * v_cur, // [n_embd_head, n_head_kv, n_tokens] + ggml_tensor * k_cache, // master K buffer + ggml_tensor * v_cache, // master V buffer + ggml_tensor * block_table, // [max_blocks, batch_size] + ggml_tensor * write_slots, // [n_tokens] + ggml_tensor * context_lens, // [batch_size] + ggml_tensor * batch_offsets, // [batch_size] + ggml_tensor * batch_lens, // [batch_size] + float kq_scale, + int block_size, + int max_blocks) const { + + // Paged attention kernel (write) assumes dense layout [n_tokens. n_heads_kv, head_dim]. + // Architectures like (Falcon, GPT-2, etc.) produce KV as views into a fused QKV tensor + // We force contiguity before passing to kernel. + // This can be optimized in phase 2. + k_cur = ggml_cont(ctx0, k_cur); + v_cur = ggml_cont(ctx0, v_cur); + q = ggml_cont(ctx0, q); + + ggml_tensor * cur = ggml_paged_attn(ctx0, + q, k_cur, v_cur, k_cache, v_cache, + block_table, write_slots, context_lens, batch_offsets, batch_lens, + kq_scale, block_size, max_blocks); + return cur; +} + ggml_tensor * llm_graph_context::build_attn_mha( ggml_tensor * q, ggml_tensor * k, @@ -2246,6 +2305,48 @@ ggml_tensor * llm_graph_context::build_attn( return cur; } +llm_graph_input_i * llm_graph_context::build_attn_inp_kv_auto() const { + if (cparams.kv_paged) { + return build_attn_inp_kv_paged(); + } + return build_attn_inp_kv(); +} + +ggml_tensor * llm_graph_context::build_attn( + llm_graph_input_i * inp, + ggml_tensor * wo, + ggml_tensor * wo_b, + ggml_tensor * wo_s, + ggml_tensor * q_cur, + ggml_tensor * k_cur, + ggml_tensor * v_cur, + ggml_tensor * kq_b, + ggml_tensor * sinks, + ggml_tensor * v_mla, + float kq_scale, + int il) const { + if (auto * no_cache = dynamic_cast(inp)) { + return build_attn(no_cache, wo, wo_b, wo_s, q_cur, k_cur, v_cur, kq_b, sinks, v_mla, kq_scale, il); + } + if (auto * kv = dynamic_cast(inp)) { + return build_attn(kv, wo, wo_b, wo_s, q_cur, k_cur, v_cur, kq_b, sinks, v_mla, kq_scale, il); + } + if (auto * iswa = dynamic_cast(inp)) { + return build_attn(iswa, wo, wo_b, wo_s, q_cur, k_cur, v_cur, kq_b, sinks, v_mla, kq_scale, il); + } + if (auto * paged = dynamic_cast(inp)) { + return build_attn(paged, wo, wo_b, wo_s, q_cur, k_cur, v_cur, kq_b, sinks, v_mla, kq_scale, il); + } + if (auto * k = dynamic_cast(inp)) { + return build_attn(k, wo, wo_b, wo_s, q_cur, k_cur, v_cur, kq_b, sinks, v_mla, kq_scale, il); + } + if (auto * cross = dynamic_cast(inp)) { + return build_attn(cross, wo, wo_b, wo_s, q_cur, k_cur, v_cur, kq_b, sinks, v_mla, kq_scale, il); + } + GGML_ASSERT(false && "unknown attention input type"); + return nullptr; +} + static std::unique_ptr build_attn_inp_k_impl( ggml_context * ctx0, const llama_ubatch & ubatch, @@ -2420,6 +2521,55 @@ ggml_tensor * llm_graph_context::build_attn( return cur; } +ggml_tensor * llm_graph_context::build_attn( + llm_graph_input_attn_kv_paged * inp, + ggml_tensor * wo, + ggml_tensor * wo_b, + ggml_tensor * wo_s, + ggml_tensor * q_cur, + ggml_tensor * k_cur, + ggml_tensor * v_cur, + ggml_tensor * /*kq_b*/, + ggml_tensor * /*sinks*/, + ggml_tensor * /*v_mla*/, + float kq_scale, + int il) const { + GGML_ASSERT(inp && "attn input is nullptr."); + auto * paged_mctx = inp->mctx; + GGML_ASSERT(paged_mctx && "llama_kv_cache_context_paged is nullptr."); + + const int32_t max_blocks = paged_mctx->get_max_blocks(); + ggml_tensor * k_physical = paged_mctx->get_k(il); // interleaved KV + ggml_tensor * v_physical = paged_mctx->get_v(il); // interleaved KV + + ggml_tensor * cur = build_attn_mha_paged( + q_cur, k_cur, v_cur, k_physical, v_physical, + inp->paged_block_table, + inp->paged_write_slots, + inp->paged_context_lens, + inp->paged_batch_offsets, + inp->paged_batch_lens, + kq_scale, cparams.block_size, max_blocks); + cb(cur, "kqv_out", il); + + // Reshape to [n_embd, n_tokens] (just a view) + cur = ggml_reshape_2d(ctx0, cur, hparams.n_embd, n_tokens); + + if (wo) { + cur = build_lora_mm(wo, cur, wo_s); + } + + if (wo_b) { + //cb(cur, "kqv_wo", il); + } + + if (wo_b) { + cur = ggml_add(ctx0, cur, wo_b); + } + + return cur; +} + llm_graph_input_attn_cross * llm_graph_context::build_attn_inp_cross() const { auto inp = std::make_unique(cross); @@ -2476,6 +2626,31 @@ ggml_tensor * llm_graph_context::build_attn( return cur; } +llm_graph_input_attn_kv_paged * llm_graph_context::build_attn_inp_kv_paged() const { + const auto * mctx_paged = static_cast(mctx); + + auto inp = std::make_unique(hparams, cparams, mctx_paged); + + const int32_t n_tokens = mctx_paged->get_n_tokens(); + const int32_t batch_size = mctx_paged->get_batch_size(); + const int32_t max_blocks = mctx_paged->get_max_blocks(); + + // Create the GGML descriptors + inp->paged_write_slots = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + inp->paged_block_table = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, max_blocks, batch_size); + inp->paged_context_lens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, batch_size); + inp->paged_batch_offsets = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, batch_size); + inp->paged_batch_lens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, batch_size); + + ggml_set_input(inp->paged_write_slots); + ggml_set_input(inp->paged_block_table); + ggml_set_input(inp->paged_context_lens); + ggml_set_input(inp->paged_batch_offsets); + ggml_set_input(inp->paged_batch_lens); + + return (llm_graph_input_attn_kv_paged *) res->add_input(std::move(inp)); +} + // TODO: maybe separate the inner implementation into a separate function // like with the non-sliding window equivalent // once sliding-window hybrid caches are a thing. diff --git a/src/llama-graph.h b/src/llama-graph.h index 5cb1756c6a97..0b236cd0ec5a 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -23,6 +23,7 @@ struct llama_memory_context_i; class llama_kv_cache_context; class llama_kv_cache_iswa_context; +class llama_kv_cache_paged_context; class llama_memory_recurrent_context; class llama_memory_hybrid_context; class llama_memory_hybrid_iswa_context; @@ -401,6 +402,36 @@ class llm_graph_input_attn_kv_iswa : public llm_graph_input_i { const llama_kv_cache_iswa_context * mctx; }; +class llm_graph_input_attn_kv_paged : public llm_graph_input_i { +public: + llm_graph_input_attn_kv_paged( + const llama_hparams & hparams, + const llama_cparams & cparams, + const llama_kv_cache_paged_context * mctx) : + hparams(hparams), + cparams(cparams), + mctx(mctx) { + } + ~llm_graph_input_attn_kv_paged() = default; + + void set_input(const llama_ubatch * ubatch) override; + bool can_reuse(const llm_graph_params & params) override; + + // The tensors the attention kernel will actually use + ggml_tensor * paged_write_slots = nullptr; + ggml_tensor * paged_block_table = nullptr; + ggml_tensor * paged_context_lens = nullptr; + ggml_tensor * paged_batch_offsets = nullptr; + ggml_tensor * paged_batch_lens = nullptr; + + const llama_hparams hparams; + const llama_cparams cparams; + + int32_t last_n_tokens; + + const llama_kv_cache_paged_context * mctx; +}; + class llm_graph_input_attn_cross : public llm_graph_input_i { public: llm_graph_input_attn_cross(const llama_cross * cross) : cross(cross) {} @@ -892,6 +923,20 @@ struct llm_graph_context { // // attention // + ggml_tensor * build_attn_mha_paged( + ggml_tensor * q, // [n_embd_head, n_head, n_tokens] + ggml_tensor * k_cur, // [n_embd_head, n_head_kv, n_tokens] + ggml_tensor * v_cur, // [n_embd_head, n_head_kv, n_tokens] + ggml_tensor * k_cache, // master K buffer + ggml_tensor * v_cache, // master V buffer + ggml_tensor * block_table, // [max_blocks, batch_size] + ggml_tensor * write_slots, // [n_tokens] + ggml_tensor * context_lens, // [batch_size] + ggml_tensor * batch_offsets, // [batch_size] + ggml_tensor * batch_lens, // [batch_size] + float kq_scale, + int block_size, + int max_blocks) const; ggml_tensor * build_attn_mha( ggml_tensor * q, // [n_embd_head_q, n_head_q, n_tokens] @@ -904,6 +949,22 @@ struct llm_graph_context { float kq_scale, int il) const; + llm_graph_input_i * build_attn_inp_kv_auto() const; + + ggml_tensor * build_attn( + llm_graph_input_i * inp, + ggml_tensor * wo, + ggml_tensor * wo_b, + ggml_tensor * wo_s, + ggml_tensor * q_cur, + ggml_tensor * k_cur, + ggml_tensor * v_cur, + ggml_tensor * kq_b, + ggml_tensor * sinks, + ggml_tensor * v_mla, + float kq_scale, + int il) const; + llm_graph_input_attn_no_cache * build_attn_inp_no_cache() const; ggml_tensor * build_attn( @@ -969,6 +1030,22 @@ struct llm_graph_context { float kq_scale, int il) const; + llm_graph_input_attn_kv_paged * build_attn_inp_kv_paged() const; + + ggml_tensor * build_attn( + llm_graph_input_attn_kv_paged * inp, + ggml_tensor * wo, + ggml_tensor * wo_b, + ggml_tensor * wo_s, + ggml_tensor * q_cur, // [n_embd_head_q, n_head_q, n_tokens] + ggml_tensor * k_cur, // [n_embd_head_k, n_head_k, n_tokens] optional + ggml_tensor * v_cur, // [n_embd_head_v, n_head_v, n_tokens] optional + ggml_tensor * kq_b, + ggml_tensor * sinks, // [n_head_q] + ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] + float kq_scale, + int il) const; + llm_graph_input_attn_cross * build_attn_inp_cross() const; ggml_tensor * build_attn( diff --git a/src/llama-kv-cache-paged.cpp b/src/llama-kv-cache-paged.cpp new file mode 100644 index 000000000000..8ddd8d4955b5 --- /dev/null +++ b/src/llama-kv-cache-paged.cpp @@ -0,0 +1,443 @@ +#include "llama-kv-cache-paged.h" + +#include "llama-impl.h" + +// +// llama_kv_cache_paged +// + +llama_kv_cache_paged::llama_kv_cache_paged(uint32_t head_dim, + uint32_t n_heads_kv, + uint32_t block_size, + uint32_t n_layers, + uint32_t n_ubatch, + uint32_t n_seq_max) : + kv_type(GGML_TYPE_F16), + head_dim(head_dim), + n_heads_kv(n_heads_kv), + block_size(block_size), + n_layers(n_layers), + n_ubatch(n_ubatch), + n_seq_max(n_seq_max), + num_gpu_blocks(0), + num_cpu_blocks(0), + gpu_backend(nullptr), + cpu_backend(nullptr) {} + +void llama_kv_cache_paged::init(ggml_backend_t backend_gpu, + ggml_backend_t backend_cpu, + enum ggml_type type, + uint32_t n_gpu_blocks, + uint32_t n_cpu_blocks, + float watermark) { + GGML_ASSERT(backend_cpu && "backend_cpu is nullptr"); + GGML_ASSERT(backend_gpu && "backend_gpu is nullptr"); + const ggml_backend_dev_t dev = ggml_backend_get_device(backend_gpu); + if (!dev || ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) { + LLAMA_LOG_WARN( + "%s: no GPU device found, allocating KV block pool on CPU. " + "This is valid for testing but it will be slow.\n", + __func__); + } + + GGML_ASSERT(n_gpu_blocks && "n_gpu_blocks need to be greater than 0."); + GGML_ASSERT(n_cpu_blocks && "n_cpu_blocks need to be greater than 0."); + + LLAMA_LOG_INFO( + "%s: initializing paged KV cache. n_gpu_blocks=%d, n_cpu_blocks=%d, block_size=%d, watermark=%0.2f\n", __func__, + n_gpu_blocks, n_cpu_blocks, block_size, watermark); + num_gpu_blocks = n_gpu_blocks; + num_cpu_blocks = n_cpu_blocks; + kv_type = type; + gpu_backend = backend_gpu; + cpu_backend = backend_cpu; + block_bytes = 2 * block_size * n_heads_kv * head_dim * ggml_type_size(kv_type); + + // Set up GPU context and tensor + // Interleaved shape: [num_blocks, 2, n_heads_kv, block_size, head_dim] (5D) + struct ggml_init_params gpu_params; + gpu_params.mem_size = ggml_tensor_overhead() * 5 * n_layers; + gpu_params.mem_buffer = NULL; + gpu_params.no_alloc = true; + + struct ggml_context * ctx_gpu = ggml_init(gpu_params); + + for (uint32_t il = 0; il < n_layers; ++il) { + // Since GGML_MAX_DIMS is set to 4, we flatten the layout to be 4D: [num_blocks, 2 * n_heads_kv, block_size, head_dim] + ggml_tensor * kv_layer_gpu = + ggml_new_tensor_4d(ctx_gpu, type, head_dim, block_size, 2 * n_heads_kv, n_gpu_blocks); + kv_gpu_layers.push_back(kv_layer_gpu); + } + + // Allocate on GPU backend + ggml_backend_buffer_t buf_gpu = ggml_backend_alloc_ctx_tensors(ctx_gpu, backend_gpu); + GGML_ASSERT(buf_gpu && "Failed to allocate GPU KV cache buffer"); + ggml_backend_buffer_clear(buf_gpu, 0); // zero out the cache + for (uint32_t il = 0; il < n_layers; ++il) { + GGML_ASSERT(kv_gpu_layers[il]->buffer && "GPU layer tensor has null buffer"); + } + + // For non CUDA backends, we would split views to allow for standard ggml operators to work out of the box + + // Set up CPU context and tensor (for swapping) + struct ggml_init_params cpu_params; + cpu_params.mem_size = ggml_tensor_overhead() * 5 * n_layers; + cpu_params.mem_buffer = NULL; + cpu_params.no_alloc = true; + struct ggml_context * ctx_cpu = ggml_init(cpu_params); + for (uint32_t il = 0; il < n_layers; ++il) { + ggml_tensor * kv_layer_cpu = + ggml_new_tensor_4d(ctx_cpu, type, head_dim, block_size, 2 * n_heads_kv, n_cpu_blocks); + kv_cpu_layers.push_back(kv_layer_cpu); + } + + // Allocate on the CPU backend (using pinned memory for faster PCIe transfer) + ggml_backend_buffer_t buf_cpu = ggml_backend_alloc_ctx_tensors(ctx_cpu, backend_cpu); + GGML_ASSERT(buf_cpu && "Failed to allocate CPU KV cache buffer"); + ggml_backend_buffer_clear(buf_cpu, 0); // zero out the cache + for (uint32_t il = 0; il < n_layers; ++il) { + GGML_ASSERT(kv_cpu_layers[il]->buffer && "CPU layer tensor has null buffer"); + } + + // Setting up our block accountant + block_manager.init(n_gpu_blocks, n_cpu_blocks, watermark); +} + +bool llama_kv_cache_paged::allocate(int32_t num_tokens, llama_sequence_group & group) { + uint32_t curr_block_count = group.block_table.size(); + uint32_t total_num_tokens = group.n_prompt + group.n_decoded + num_tokens; + uint32_t num_requested_blocks = std::ceil((float) total_num_tokens / block_size) - curr_block_count; + LLAMA_LOG_DEBUG("%s: curr_block_count=%d, total_num_tokens=%d, num_requested_blocks=%d\n", __func__, + curr_block_count, total_num_tokens, num_requested_blocks); + + if (num_requested_blocks == 0) { + return true; + } + + if (!block_manager.has_free_gpu_blocks(num_requested_blocks)) { + LLAMA_LOG_DEBUG("%s: insufficient GPU blocks. Requested: %d.\n", __func__, num_requested_blocks); + return false; + } + + llama_block_ids new_ids = block_manager.checkout_gpu_blocks(num_requested_blocks); + concat_block_ids(group.block_table, new_ids); + LLAMA_LOG_DEBUG("%s: successfully allocated %d.\n", __func__, num_requested_blocks); + return true; +} + +void llama_kv_cache_paged::free_blocks(llama_sequence_group & group) { + if (group.block_table.empty()) { + return; + } + + llama_block_ids blocks_to_free_gpu; + llama_block_ids blocks_to_free_cpu; + + for (uint32_t block_id : group.block_table) { + if (block_manager.is_gpu(block_id)) { + blocks_to_free_gpu.push_back(block_id); + } else { + blocks_to_free_cpu.push_back(block_id); + } + } + + if (!blocks_to_free_gpu.empty()) { + block_manager.release_gpu_blocks(blocks_to_free_gpu); + } + if (!blocks_to_free_cpu.empty()) { + block_manager.release_cpu_blocks(blocks_to_free_cpu); + } + + group.block_table.clear(); + seq_rm(group.request_id, llama_pos{}, llama_pos{}); +} + +void llama_kv_cache_paged::do_block_copy(const llama_block_ids & src_ids, + const llama_block_ids & new_ids, + bool to_gpu) { + const uint32_t num_blocks = src_ids.size(); + LLAMA_LOG_DEBUG("%s: num_blocks_size=%d, new_ids_size=%ld\n", __func__, num_blocks, new_ids.size()); + GGML_ASSERT(num_blocks == new_ids.size() && "src_ids and new_ids do not have the same size."); + + const auto & src_layers = to_gpu ? kv_cpu_layers : kv_gpu_layers; + const auto & dst_layers = to_gpu ? kv_gpu_layers : kv_cpu_layers; + + GGML_ASSERT(src_layers.size() == n_layers && "src layer count mismatch."); + GGML_ASSERT(dst_layers.size() == n_layers && "src layer count mismatch."); + + // Buffer on HOST to faciliate block data transfer + // Note: an optimization would be to use views and async copies. Beware of + // memory overhead heurisitcs. + std::vector staging(block_bytes); + + for (uint32_t il = 0; il < n_layers; ++il) { + struct ggml_tensor * src_main = src_layers[il]; + struct ggml_tensor * dst_main = dst_layers[il]; + + for (uint32_t i = 0; i < num_blocks; ++i) { + const uint32_t src_global = src_ids[i]; + const uint32_t dst_global = new_ids[i]; + + // GPU and CPu blocks may differ (usually CPU < GPU) + // We substract the diffence to calculate where the local starts before we calculate offsets + const uint32_t src_local = to_gpu ? src_global - num_gpu_blocks : src_global; + const uint32_t dst_local = to_gpu ? dst_global : dst_global - num_gpu_blocks; + + const size_t src_offset = (size_t) src_local * block_bytes; + const size_t dst_offset = (size_t) dst_local * block_bytes; + + // Put src tensor into HOST staging buffer + ggml_backend_tensor_get(src_main, staging.data(), src_offset, block_bytes); + // Put tensor from HOST staging into dst tensor + ggml_backend_tensor_set(dst_main, staging.data(), dst_offset, block_bytes); + } + } +} + +bool llama_kv_cache_paged::swap_in(llama_sequence_group & group) { + const uint32_t num_blocks = group.block_table.size(); + if (num_blocks == 0) { + return true; + } + + // A potential optimization to reduce thrashing is to have a heuristic to check if + // if we can continue decoding after swap_in. + if (!block_manager.has_free_gpu_blocks(num_blocks)) { + return false; + } + + llama_block_ids new_ids = block_manager.checkout_gpu_blocks(num_blocks); + do_block_copy(group.block_table, new_ids, /*to_gpu=*/true); + + free_blocks(group); + group.block_table = new_ids; + return true; +} + +bool llama_kv_cache_paged::swap_out(llama_sequence_group & group) { + const uint32_t num_blocks = group.block_table.size(); + if (num_blocks == 0) { + return true; + } + + if (!block_manager.has_free_cpu_blocks(num_blocks)) { + return false; + } + + llama_block_ids new_ids = block_manager.checkout_cpu_blocks(num_blocks); + do_block_copy(group.block_table, new_ids, /*to_gpu=*/false); + + free_blocks(group); + group.block_table = new_ids; + return true; +} + +void llama_kv_cache_paged::set_paged_batch_info(const llama_paged_batch_info * info) { + last_paged_info = info; +} + +uint32_t llama_kv_cache_paged::get_num_gpu_blocks() const { + return num_gpu_blocks; +} + +void llama_kv_cache_paged::concat_block_ids(llama_block_ids & to_block_table, + const llama_block_ids & from_block_table) { + to_block_table.insert(to_block_table.end(), from_block_table.begin(), from_block_table.end()); +} + +// llama_memory_i + +llama_memory_context_ptr llama_kv_cache_paged::init_batch(llama_batch_allocr & balloc, + uint32_t n_ubatch, + bool /*embd_all*/) { + do { + balloc.split_reset(); + + std::vector ubatches; + while (true) { + auto ubatch = balloc.split_simple(n_ubatch); + if (ubatch.n_tokens == 0) { + break; + } + ubatches.push_back(std::move(ubatch)); + } + + // Failed to find a suitable split + if (balloc.get_n_used() < balloc.get_n_tokens()) { + break; + } + + auto ctx = std::make_unique(this, std::move(ubatches)); + + // Do not use balloc's internal batch. It does not carry any paged metadata. + GGML_ASSERT(last_paged_info && "no paged batch info set before init_batch was called."); + ctx->set_batch_data(*last_paged_info); + return ctx; + } while (false); + + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); +} + +// Used by llama_context scheduler to dry-run +llama_memory_context_ptr llama_kv_cache_paged::init_full() { + LLAMA_LOG_DEBUG("%s: reserving graph for n_ubatch=%d, n_seq_max=%d, num_gpu_blocks=%d\n", __func__, n_ubatch, + n_seq_max, num_gpu_blocks); + + // Create a "dummy" ubatch that represents the maximum capacity + // of the system to let the scheduler reserve enough space for metadata. + llama_ubatch ubatch = {}; + ubatch.n_tokens = n_ubatch; // maximum tokens + ubatch.n_seqs = n_seq_max; // maximum sequences + ubatch.n_pos = 1; + + std::vector ubatches = { ubatch }; + + auto ctx = std::make_unique(this, ubatches); + + ctx->set_batch_size(n_seq_max); // maximum possible sequences + ctx->set_n_tokens(n_ubatch); // representative token count + ctx->set_max_blocks(num_gpu_blocks); // every block could theoretically belong to one seq + + return ctx; +} + +llama_memory_context_ptr llama_kv_cache_paged::init_update(llama_context * /*lctx*/, bool /*optimize*/) { + std::vector dummy_ubatch = {}; + auto ctx = std::make_unique(this, dummy_ubatch); + // TODO maybe confirm block counts or clean up stale pointers + return ctx; +} + +struct ggml_tensor * llama_kv_cache_paged::get_kv_tensor(int layer_idx) const { + return kv_gpu_layers[layer_idx]; +} + +void llama_kv_cache_paged::clear(bool /*data*/) { + sequence_positions.clear(); +} + +bool llama_kv_cache_paged::seq_rm(llama_seq_id seq_id, llama_pos /*p0*/, llama_pos /*p1*/) { + sequence_positions.erase(seq_id); + return true; +} + +llama_pos llama_kv_cache_paged::seq_pos_min(llama_seq_id seq_id) const { + auto it = sequence_positions.find(seq_id); + return (it != sequence_positions.end()) ? it->second.min : -1; +} + +llama_pos llama_kv_cache_paged::seq_pos_max(llama_seq_id seq_id) const { + auto it = sequence_positions.find(seq_id); + return (it != sequence_positions.end()) ? it->second.max : -1; +} + +std::map llama_kv_cache_paged::memory_breakdown() const { + std::map breakdown; + const size_t n_gpu_kvs = kv_gpu_layers.size(); + const size_t n_cpu_kvs = kv_cpu_layers.size(); + + for (size_t il = 0; il < n_layers; ++il) { + auto * kv_gpu = (il < n_gpu_kvs) ? kv_gpu_layers[il] : nullptr; + if (kv_gpu) { + breakdown[ggml_backend_buffer_get_type(kv_gpu->buffer)] = ggml_nbytes(kv_gpu); + } + auto * kv_cpu = (il < n_cpu_kvs) ? kv_cpu_layers[il] : nullptr; + if (kv_cpu) { + breakdown[ggml_backend_buffer_get_type(kv_cpu->buffer)] = ggml_nbytes(kv_cpu); + } + } + return breakdown; +} + +void llama_kv_cache_paged::set_seq_min_pos(llama_seq_id seq_id, llama_pos new_min) { + sequence_positions[seq_id].min = new_min; +} + +void llama_kv_cache_paged::set_seq_max_pos(llama_seq_id seq_id, llama_pos new_max) { + sequence_positions[seq_id].max = new_max; +} + +// llama_kv_cache_paged_context + +void llama_kv_cache_paged_context::set_batch_data(const llama_paged_batch_info & info) { + paged_write_slots = info.write_slots; + paged_block_table = info.block_table; + paged_context_lens = info.context_lens; + paged_batch_offsets = info.batch_offsets; + paged_batch_lens = info.batch_lens; + n_tokens = info.n_tokens; + max_blocks = info.n_blocks_per_seq; + batch_size = info.n_seq; +} + +bool llama_kv_cache_paged_context::next() { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + if (++i_cur >= ubatches.size()) { + return false; + } + return true; +} + +bool llama_kv_cache_paged_context::apply() { + // Nothing to do for paged KV cache, return true to allow for execution + return true; +} + +const llama_ubatch & llama_kv_cache_paged_context::get_ubatch() const { + assert(status == LLAMA_MEMORY_STATUS_SUCCESS); + return ubatches[i_cur]; +} + +struct ggml_tensor * llama_kv_cache_paged_context::get_k(int layer_idx) const { + GGML_ASSERT(manager && "manager has not been initialized."); + return manager->get_kv_tensor(layer_idx); +} + +struct ggml_tensor * llama_kv_cache_paged_context::get_v(int layer_idx) const { + GGML_ASSERT(manager && "manager has not been initialized."); + return manager->get_kv_tensor(layer_idx); +} + +int32_t llama_kv_cache_paged_context::get_n_tokens() const { + return n_tokens; +} + +int32_t llama_kv_cache_paged_context::get_batch_size() const { + return batch_size; +} + +int32_t llama_kv_cache_paged_context::get_max_blocks() const { + return max_blocks; +} + +int32_t * llama_kv_cache_paged_context::get_write_slots() const { + return paged_write_slots; +} + +int32_t * llama_kv_cache_paged_context::get_block_table() const { + return paged_block_table; +} + +int32_t * llama_kv_cache_paged_context::get_context_lens() const { + return paged_context_lens; +} + +int32_t * llama_kv_cache_paged_context::get_batch_offsets() const { + return paged_batch_offsets; +} + +int32_t * llama_kv_cache_paged_context::get_batch_lens() const { + return paged_batch_lens; +} + +void llama_kv_cache_paged_context::set_n_tokens(int32_t new_n_tokens) { + n_tokens = new_n_tokens; +} + +void llama_kv_cache_paged_context::set_batch_size(int32_t new_batch_size) { + batch_size = new_batch_size; +} + +void llama_kv_cache_paged_context::set_max_blocks(int32_t new_max_blocks) { + max_blocks = new_max_blocks; +} diff --git a/src/llama-kv-cache-paged.h b/src/llama-kv-cache-paged.h new file mode 100644 index 000000000000..9db497bcd4f8 --- /dev/null +++ b/src/llama-kv-cache-paged.h @@ -0,0 +1,186 @@ +#pragma once + +#include "llama-batch.h" +#include "llama-block-manager.h" +#include "llama-graph.h" +#include "llama-memory.h" +#include "llama-sequence-group.h" + +#include + +// +// llama_kv_cache_paged +// + +class llama_kv_cache_paged : public llama_memory_i { + public: + llama_kv_cache_paged(uint32_t head_dim, + uint32_t n_head_kv, + uint32_t block_size, + uint32_t n_layers, + uint32_t n_ubatch, + uint32_t n_seq_max); + + void init(ggml_backend_t backend_gpu, + ggml_backend_t backend_cpu, + enum ggml_type type, + uint32_t n_gpu_blocks, + uint32_t n_cpu_blocks, + float watermark); // percentage + + bool allocate(int32_t num_tokens, llama_sequence_group & group); + void free_blocks(llama_sequence_group & group); + bool swap_in(llama_sequence_group & group); + bool swap_out(llama_sequence_group & group); + + void set_paged_batch_info(const llama_paged_batch_info * info); + uint32_t get_num_gpu_blocks() const; + + // + // llama_memory_i + // + llama_memory_context_ptr init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) override; + + llama_memory_context_ptr init_full() override; + llama_memory_context_ptr init_update(llama_context * lctx, bool optimize) override; + + struct ggml_tensor * get_kv_tensor(int layer_idx) const; + + bool get_can_shift() const override { return false; } + + void clear(bool data) override; + + bool seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; + + void seq_cp(llama_seq_id /*seq_id_src*/, + llama_seq_id /*seq_id_dst*/, + llama_pos /*p0*/, + llama_pos /*p1*/) override { /* implement later CoW mechanism */ + } + + void seq_keep(llama_seq_id /*seq_id*/) override {} + + void seq_add(llama_seq_id /*seq_id*/, llama_pos /*p0*/, llama_pos /*p1*/, llama_pos /*shift*/) override {} + + void seq_div(llama_seq_id /*seq_id*/, llama_pos /*p0*/, llama_pos /*p1*/, int /*d*/) override {} + + llama_pos seq_pos_min(llama_seq_id seq_id) const override; + llama_pos seq_pos_max(llama_seq_id seq_id) const override; + + std::map memory_breakdown() const override; + + // state write/load + void state_write(llama_io_write_i & /*io*/, + llama_seq_id /*seq_id*/ = -1, + llama_state_seq_flags /*flags*/ = 0) const override {} + + void state_read(llama_io_read_i & /*io*/, + llama_seq_id /*seq_id*/ = -1, + llama_state_seq_flags /*flags*/ = 0) override {} + + // + // Helpers to llama_memory_i + // + void set_seq_min_pos(llama_seq_id seq_id, llama_pos new_min); + void set_seq_max_pos(llama_seq_id seq_id, llama_pos new_max); + + private: + void concat_block_ids(llama_block_ids & to_block_table, const llama_block_ids & from_block_table); + void do_block_copy(const llama_block_ids & src_ids, const llama_block_ids & new_ids, bool to_gpu); + + // Master physical buffer + // For CUDA: memory is interleaved + // For other backends: we treat the exact same memory buffer as two virtual views + std::vector kv_gpu_layers; + std::vector kv_cpu_layers; + + enum ggml_type kv_type; + + llama_block_manager block_manager; + + // Non-owning pointer to the batch currently being processed. + // Lifetime: set by the scheduler at the end of step(), cleared at the + // start of the next step() (before the batch's paged_* arrays are freed). + // The ordering in llama_paged_scheduler_impl::clear_batch is load-bearing; + // do not reorder without updating init_batch's contract. + const llama_paged_batch_info * last_paged_info = nullptr; + + const uint32_t head_dim; + const uint32_t n_heads_kv; + const uint32_t block_size; + const uint32_t n_layers; + const uint32_t n_ubatch; + const uint32_t n_seq_max; + uint32_t num_gpu_blocks; + uint32_t num_cpu_blocks; + uint32_t block_bytes; + + ggml_backend_t gpu_backend; + ggml_backend_t cpu_backend; + + struct seq_range { + llama_pos min = -1; + llama_pos max = -1; + }; + + std::unordered_map sequence_positions; +}; + +class llama_kv_cache_paged_context : public llama_memory_context_i { + public: + llama_kv_cache_paged_context(llama_kv_cache_paged * parent, const std::vector & in_ubatch) : + manager(parent), + ubatches(in_ubatch) { + i_cur = 0; + } + + llama_kv_cache_paged_context(llama_memory_status status) : status(status) {} + + void set_batch_data(const llama_paged_batch_info & info); + int32_t get_n_tokens() const; + int32_t get_batch_size() const; + int32_t get_max_blocks() const; + + int32_t * get_write_slots() const; + int32_t * get_block_table() const; + int32_t * get_context_lens() const; + int32_t * get_batch_offsets() const; + int32_t * get_batch_lens() const; + + void set_n_tokens(int32_t new_n_tokens); + void set_batch_size(int32_t new_batch_size); + void set_max_blocks(int32_t new_max_blocks); + + struct ggml_tensor * get_k(int layer_idx) const; + struct ggml_tensor * get_v(int layer_idx) const; + + // + // llama_memory_context_i + // + bool next() override; + bool apply() override; + const llama_ubatch & get_ubatch() const override; + + llama_memory_status get_status() const override { return status; } + + private: + const llama_kv_cache_paged * manager; + + // + // batch processing context + // + std::vector ubatches; + size_t i_cur = 0; // index of ubatch to process + + int32_t * paged_write_slots = nullptr; // [n_tokens] + int32_t * paged_block_table = nullptr; // [batch_size, max_blocks] + int32_t * paged_context_lens = nullptr; // [batch_size] + int32_t * paged_batch_offsets = nullptr; // [batch_size] + int32_t * paged_batch_lens = nullptr; // [batch_size] + + int32_t n_tokens = 0; + int32_t batch_size = 0; + int32_t max_blocks = 0; + + llama_memory_status status = LLAMA_MEMORY_STATUS_SUCCESS; +}; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 9e2a13cbd43e..adbdbf0604f7 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -10,6 +10,7 @@ #include "llama-kv-cache.h" #include "llama-kv-cache-iswa.h" +#include "llama-kv-cache-paged.h" #include "llama-memory-hybrid.h" #include "llama-memory-hybrid-iswa.h" #include "llama-memory-recurrent.h" @@ -8421,7 +8422,8 @@ ggml_tensor * llama_model::get_rope_factors(const llama_cparams & cparams, int i return layers[il].rope_short; } -llama_memory_i * llama_model::create_memory(const llama_memory_params & params, const llama_cparams & cparams) const { +llama_memory_i * llama_model::create_memory(const llama_memory_params & params, const llama_cparams & cparams, + ggml_backend_t backend_gpu, ggml_backend_t backend_cpu) const { llama_memory_i * res; switch (arch) { @@ -8544,21 +8546,50 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, reuse); } else { GGML_ASSERT(!hparams.is_swa_any()); - - res = new llama_kv_cache( - *this, + if (cparams.kv_paged) { + GGML_ASSERT(!cparams.kv_unified && "conflicting parameters: kv_unified cannot be used with kv_paged."); + GGML_ASSERT(cparams.n_ubatch == cparams.n_batch && "kv_paged requires n_ubatch == n_batch."); + LLAMA_LOG_INFO("%s: Detected kv_paged=%d, creating llama_kv_cache_paged.\n", __func__, cparams.kv_paged); + const uint32_t head_dim = hparams.n_embd_head_v(); + const uint32_t n_head = hparams.n_head_kv(); + const uint32_t n_layers = hparams.n_layer; + const uint32_t block_size = cparams.block_size; + + const uint32_t n_gpu_blocks = cparams.n_gpu_blocks; + const uint32_t n_cpu_blocks = cparams.n_cpu_blocks; + const float watermark = cparams.kv_paged_watermark; + + const uint32_t n_ubatch = cparams.n_ubatch; + const uint32_t n_seq_max = cparams.n_seq_max; + + auto * paged_cache = new llama_kv_cache_paged(head_dim, n_head, block_size, n_layers, n_ubatch, n_seq_max); + GGML_ASSERT(paged_cache && "unable to create paged KV cache."); + + paged_cache->init( + backend_gpu, + backend_cpu, params.type_k, - params.type_v, - !cparams.flash_attn, - cparams.offload_kqv, - cparams.kv_unified, - cparams.n_ctx_seq, - cparams.n_seq_max, - 1, - hparams.n_swa, - hparams.swa_type, - nullptr, - nullptr); + n_gpu_blocks, + n_cpu_blocks, + watermark); + + res = paged_cache; + } else { + res = new llama_kv_cache( + *this, + params.type_k, + params.type_v, + !cparams.flash_attn, + cparams.offload_kqv, + cparams.kv_unified, + cparams.n_ctx_seq, + cparams.n_seq_max, + 1, + hparams.n_swa, + hparams.swa_type, + nullptr, + nullptr); + } } } } diff --git a/src/llama-model.h b/src/llama-model.h index 5f101bd63745..429117e749b1 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -621,7 +621,8 @@ struct llama_model { ggml_tensor * get_rope_factors(const llama_cparams & cparams, int il) const; // TODO: move this to new llm_arch_model_i interface - llama_memory_i * create_memory(const llama_memory_params & params, const llama_cparams & cparams) const; + llama_memory_i * create_memory(const llama_memory_params & params, const llama_cparams & cparams, + ggml_backend_t backend_gpu, ggml_backend_t backend_cpu) const; // TODO: move this to new llm_arch_model_i interface ggml_cgraph * build_graph(const llm_graph_params & params) const; diff --git a/src/llama-paged-scheduler-impl.cpp b/src/llama-paged-scheduler-impl.cpp new file mode 100644 index 000000000000..206ca6b64fae --- /dev/null +++ b/src/llama-paged-scheduler-impl.cpp @@ -0,0 +1,491 @@ +#include "llama-paged-scheduler-impl.h" + +#include "llama-impl.h" + +llama_paged_scheduler_impl::llama_paged_scheduler_impl(uint32_t n_ctx, + uint32_t block_sz, + int32_t n_batch, + llama_kv_cache_paged * kv_manager) : + n_seq_max_ctx(n_ctx), + block_size(block_sz), + n_batch(n_batch), + kv_cache_manager(kv_manager) {} + +bool llama_paged_scheduler_impl::check_deadlock(uint32_t n_candidates, uint32_t n_swapped, uint32_t n_waiting) const { + if (n_candidates == 0 && (n_swapped > 0 || n_waiting > 0)) { + LLAMA_LOG_ERROR( + "%s: Scheduler deadlock detected. " + "%d sequence(s) are swapped out and %d are waiting, " + "but there are not enough free GPU blocks to make progress. " + "Hint: increase n_gpu_blocks (currently %d) or reduce n_sequences.\n", + __func__, n_swapped, n_waiting, kv_cache_manager->get_num_gpu_blocks()); + return true; + } + return false; +} + +bool llama_paged_scheduler_impl::check_livelock(uint32_t n_swapped, uint32_t prev_n_swapped) { + // swapped count is non-zero and not decreasing + if (n_swapped > 0 && n_swapped >= prev_n_swapped) { + n_livelock_steps++; + if (n_livelock_steps >= max_livelock_steps) { + LLAMA_LOG_ERROR( + "%s: Livelock detected. Swapped count has been " + "non-decreasing for %d steps (currently %d swapped). " + "Increase n_gpu_blocks (currently %d) or reduce " + "n_sequences.\n", + __func__, n_livelock_steps, n_swapped, kv_cache_manager->get_num_gpu_blocks()); + return true; + } + } else { + // Swapped count decreased — swap-ins are happening, reset counter + n_livelock_steps = 0; + } + return false; +} + +int32_t llama_paged_scheduler_impl::get_curr_decode_tokens() const { + return running.size(); +} + +llama_scheduler_status llama_paged_scheduler_impl::step(llama_batch & batch) { + // Free previous inference batches + clear_batch(batch); + + llama_sequence_group_raw_list candidates; + process_running_list(candidates); + process_swapped_list(candidates); + + const int32_t remaining = n_batch - get_curr_decode_tokens(); + process_waiting_list(candidates, remaining); + + const uint32_t n_running = running.size(); + const uint32_t n_swapped = swapped.size(); + const uint32_t n_waiting = waiting.size(); + const uint32_t n_candidates = candidates.size(); + + LLAMA_LOG_INFO("%s: Scheduler status: running=%d, swapped=%d, waiting=%d, candidates=%d\n", __func__, n_running, + n_swapped, n_waiting, n_candidates); + + const bool deadlock = check_deadlock(n_candidates, n_swapped, n_waiting); + const bool livelock = check_livelock(n_swapped, prev_n_swapped); // updates n_livelock_steps + if (deadlock || livelock) { + return llama_scheduler_status::DEADLOCK; + } + + prev_n_swapped = n_swapped; + populate_batch_from(candidates, batch); + kv_cache_manager->set_paged_batch_info(&curr_info); + return llama_scheduler_status::OK; +} + +bool llama_paged_scheduler_impl::queue_request(llama_sequence_group group) { + // Rejecting any requests that exceeds max context for a seq + if (group.n_prompt >= n_seq_max_ctx) { + LLAMA_LOG_ERROR("%s: request %d exceeds max context (%d > %d).\n", __func__, group.request_id, group.n_prompt, + n_seq_max_ctx); + return false; + } + + auto group_ptr = std::make_unique(std::move(group)); + + id_to_group[group_ptr->request_id] = group_ptr.get(); + + set_waiting(std::move(group_ptr)); + return true; +} + +void llama_paged_scheduler_impl::insert_sorted_by_arrival_time(llama_sequence_group_ptr new_group_ptr, + llama_sequence_group_list & list) { + GGML_ASSERT(new_group_ptr && "New group cannot be sorted because it's nullptr."); + auto it = + std::lower_bound(list.begin(), list.end(), new_group_ptr->t_arrival_time, + [](const llama_sequence_group_ptr & group, int64_t time) { + GGML_ASSERT(group && "group cannot be checked for arrival time because it's nullptr."); + return group->t_arrival_time < time; + }); + list.insert(it, std::move(new_group_ptr)); +} + +void llama_paged_scheduler_impl::set_running(llama_sequence_group_ptr group_ptr) { + GGML_ASSERT(group_ptr && group_ptr->status != llama_sequence_group_status::RUNNING && + "Request is already running."); + group_ptr->status = llama_sequence_group_status::RUNNING; + insert_sorted_by_arrival_time(std::move(group_ptr), running); +} + +void llama_paged_scheduler_impl::set_swapped(llama_sequence_group_ptr group_ptr) { + GGML_ASSERT(group_ptr && group_ptr->status != llama_sequence_group_status::SWAPPED && + "Request is already swapped."); + group_ptr->status = llama_sequence_group_status::SWAPPED; + insert_sorted_by_arrival_time(std::move(group_ptr), swapped); +} + +void llama_paged_scheduler_impl::set_waiting(llama_sequence_group_ptr group_ptr) { + GGML_ASSERT(group_ptr && group_ptr->status != llama_sequence_group_status::WAITING && + "Request is already waiting."); + group_ptr->status = llama_sequence_group_status::WAITING; + insert_sorted_by_arrival_time(std::move(group_ptr), waiting); +} + +void llama_paged_scheduler_impl::finish(llama_sequence_group & group) { + GGML_ASSERT(kv_cache_manager && "kv_cache_manager is nullptr."); + GGML_ASSERT(group.status == llama_sequence_group_status::FINISHED && "Request was not marked as finished."); + // We prioritize user CB, otherwise we log by default + if (on_finish_cb) { + // TODO perhaps just have the callback take sequence_group and user_data + on_finish_cb(group.request_id, group.logical_seq.data(), (int32_t) group.logical_seq.size(), + on_finish_user_data); + } else { + LLAMA_LOG_DEBUG("%s: Request: %d generated %d tokens.\n", __func__, group.request_id, group.n_decoded); + } + kv_cache_manager->free_blocks(group); + group.status = llama_sequence_group_status::FINISHED; + id_to_group.erase(group.request_id); +} + +// Try to swap a running sequence out to CPU. +// if the CPU pool is full, fall back to recomputation by resetting the sequence's decode state +// and sending it back to the waiting queue. +// +// Takes ownership of group_ptr. +// On return, the sequence is either in the swapped list (CPU pool had room) +// or the waiting list (recomputed). +void llama_paged_scheduler_impl::swap_out_or_recompute(llama_sequence_group_ptr group_ptr) { + GGML_ASSERT(group_ptr && "group_ptr is nullptr"); + GGML_ASSERT(kv_cache_manager && "kv_cache_manager is nullptr"); + + const int32_t rid = group_ptr->request_id; + + const bool swap_ok = kv_cache_manager->swap_out(*group_ptr); + if (swap_ok) { + LLAMA_LOG_DEBUG("%s: (swapped_out) request_id=%d was swapped out to make room.\n", __func__, rid); + set_swapped(std::move(group_ptr)); + return; + } + + // There was not enough CPU memory to swap the request (recomputation) + kv_cache_manager->free_blocks(*group_ptr); + kv_cache_manager->seq_rm(rid, -1, -1); + group_ptr->n_past = 0; + group_ptr->n_decoded = 0; + group_ptr->logical_seq.resize(group_ptr->n_prompt); + + LLAMA_LOG_DEBUG("%s: (recomputation) request_id=%d was sent for recomputation.\n", __func__, rid); + set_waiting(std::move(group_ptr)); +} + +void llama_paged_scheduler_impl::evict() { + GGML_ASSERT(kv_cache_manager && "kv_cache_manager is nullptr."); + LLAMA_LOG_DEBUG("%s: Eviction requested...\n", __func__); + if (running.empty()) { + return; + } + + llama_sequence_group_ptr most_recent_request = std::move(running.back()); + running.pop_back(); + GGML_ASSERT(most_recent_request && "request selected for eviction is nullptr."); + + swap_out_or_recompute(std::move(most_recent_request)); +} + +void llama_paged_scheduler_impl::process_running_list(llama_sequence_group_raw_list & candidates) { + GGML_ASSERT(kv_cache_manager && "kv_cache_manager is nullptr."); + + llama_sequence_group_list::iterator it = running.begin(); + while (it != running.end()) { + llama_sequence_group * group = it->get(); + GGML_ASSERT(group && "group is nullptr."); + + if (group->status == llama_sequence_group_status::FINISHED) { + finish(*group); + it = running.erase(it); + continue; + } + + // Dynamically allocate more blocks to decode the request + uint32_t current_capacity = group->block_table.size() * block_size; + uint32_t required_capacity = group->n_past + 1; + LLAMA_LOG_DEBUG( + "%s: (running) request_id=%d: current_capacity (tokens)=%d toks, required capacity (tokens) = %d toks\n", + __func__, group->request_id, current_capacity, required_capacity); + if (required_capacity >= current_capacity) { + LLAMA_LOG_DEBUG("%s: (running_pending) request_id=%d: requires a new block to decode.\n", __func__, + group->request_id); + bool success = kv_cache_manager->allocate(1, *group); // decode phase + if (!success) { + if (running.size() > 1) { + const bool curr_is_back = (std::next(it) == running.end()); + // Evict pops the back of the list (most recent request) + evict(); + + // Evict might have removed the current group from running + if (curr_is_back) { + // Current group was evicted + it = running.end(); + continue; + } + + // Try allocating again after eviction + success = kv_cache_manager->allocate(1, *group); + } + + if (!success) { + // If allocate failed, it means we must evict the current request + llama_sequence_group_ptr self = std::move(*it); + it = running.erase(it); + + swap_out_or_recompute(std::move(self)); + continue; + } + } + LLAMA_LOG_DEBUG("%s: (running_restored) request_id=%d: found a new block to continue decoding.\n", __func__, + group->request_id); + } + + // A request is a candidate if we there is still room for generation without adding blocks + // or if there was enough GPU memory to allocate another physical block. + candidates.push_back(group); + ++it; + } +} + +void llama_paged_scheduler_impl::process_swapped_list(llama_sequence_group_raw_list & candidates) { + GGML_ASSERT(kv_cache_manager && "kv_cache_manager is nullptr."); + llama_sequence_group_list::iterator it = swapped.begin(); + while (it != swapped.end()) { + llama_sequence_group * group = it->get(); + GGML_ASSERT(group && "the group to swap is nullptr."); + const bool success = kv_cache_manager->swap_in(*group); + if (!success) { + // We respect FCFS, so we stop here to prevent a younger swapped request from jumping ahead. + break; + } + candidates.push_back(group); + llama_sequence_group_ptr group_ptr = std::move(*it); + LLAMA_LOG_DEBUG("%s: (swapped_in) request_id=%d back in for processing.\n", __func__, group_ptr->request_id); + set_running(std::move(group_ptr)); + it = swapped.erase(it); + } +} + +void llama_paged_scheduler_impl::process_waiting_list(llama_sequence_group_raw_list & candidates, + int32_t remaining_token_budget) { + GGML_ASSERT(kv_cache_manager && "kv_cache_manager is nullptr."); + llama_sequence_group_list::iterator it = waiting.begin(); + size_t count = 0; + while (it != waiting.end()) { + llama_sequence_group * group = it->get(); + GGML_ASSERT(group && "the waiting group is nullptr."); + + const int32_t tokens_needed = group->n_prompt + 1; + if (tokens_needed > remaining_token_budget) { + break; + } + + ++count; + // When prefilling, we want to always guarantee at least one decode to avoid thrashing + const bool success = kv_cache_manager->allocate(tokens_needed, *group); + if (!success) { + // We respect FCFS, so we stop here to prevent a younger waiting request from jumping ahead. + break; + } + candidates.push_back(group); + remaining_token_budget -= tokens_needed; + llama_sequence_group_ptr group_ptr = std::move(*it); + LLAMA_LOG_DEBUG("%s: (start) request_id=%d sent for processing.\n", __func__, group_ptr->request_id); + set_running(std::move(group_ptr)); + it = waiting.erase(it); + } + if (count > 0) { + LLAMA_LOG_DEBUG("%s: Started %ld waiting requests\n", __func__, count); + } +} + +int32_t llama_paged_scheduler_impl::calculate_global_slot_index(int32_t token_pos, + std::vector & block_table) { + GGML_ASSERT(block_size && "block_size needs to be greater than 0"); + const int32_t block_table_id = token_pos / block_size; + const int32_t offset = token_pos % block_size; + + const size_t block_table_size = block_table.size(); + if ((size_t) block_table_id >= block_table_size) { + LLAMA_LOG_ERROR("%s: block_table_id=%d is OOB for pos=%d. Block table size=%ld.\n", __func__, block_table_id, + token_pos, block_table_size); + LLAMA_LOG_ERROR("%s: block_table_contents: [ ", __func__); + for (size_t id = 0; id < block_table_size; ++id) { + LLAMA_LOG_ERROR("%d ", block_table[id]); + if (id == block_table_size - 1) { + LLAMA_LOG_ERROR("]\n"); + } + } + GGML_ASSERT(false && "block_table_id OOB"); + } + const int32_t block_id = block_table.at(block_table_id); + + return (block_id * block_size) + offset; +} + +void llama_paged_scheduler_impl::clear_batch(llama_batch & batch) { + LLAMA_LOG_DEBUG("%s: clearing batch.", __func__); + // Invalidate last scheduled batch info before freeing the arrays + // (MUST be called before the delete[]). + kv_cache_manager->set_paged_batch_info(nullptr); + + delete[] curr_info.write_slots; + delete[] curr_info.block_table; + delete[] curr_info.context_lens; + delete[] curr_info.batch_offsets; + delete[] curr_info.batch_lens; + curr_info = {}; // reset to defaults + + if (batch.n_tokens == 0) { + return; + } + + llama_batch_free(batch); + batch.n_tokens = 0; +} + +void llama_paged_scheduler_impl::populate_batch_from(const llama_sequence_group_raw_list & candidates, + llama_batch & batch) { + if (candidates.empty()) { + LLAMA_LOG_DEBUG("%s: No candidates for this step.\n", __func__); + batch.n_tokens = 0; + return; + } + int32_t total_tokens = 0; + int32_t batch_size = candidates.size(); + int32_t max_blocks = 0; + + LLAMA_LOG_DEBUG("%s: Creating batch from candidates (%d requests). n_batch=%d\n", __func__, batch_size, n_batch); + + // Calculating required sizes + for (const auto & group : candidates) { + GGML_ASSERT(group && "candidate request is nullptr."); + total_tokens += (group->n_decoded > 0) ? 1 : group->n_prompt; + max_blocks = std::max(max_blocks, (int32_t) group->block_table.size()); + } + + GGML_ASSERT(total_tokens <= (int32_t) n_batch && "total_tokens exceeds n_batch — token budget logic is broken"); + + // Initialize the batch (assumed it was cleared before) + batch = llama_batch_init(total_tokens, 0, 1); + GGML_ASSERT(batch.token != nullptr && "llama_batch_init failed to allocate tokens."); + + batch.n_tokens = total_tokens; + + curr_info.n_seq = batch_size; + curr_info.n_tokens = total_tokens; + curr_info.n_blocks_per_seq = max_blocks; + + curr_info.write_slots = new int32_t[total_tokens]; + curr_info.block_table = new int32_t[batch_size * max_blocks]; + curr_info.context_lens = new int32_t[batch_size]; + curr_info.batch_offsets = new int32_t[batch_size]; + curr_info.batch_lens = new int32_t[batch_size]; + LLAMA_LOG_DEBUG("%s: created llama_batch: n_seq=%d, n_tokens=%d, n_blocks_per_seq=%d\n", __func__, curr_info.n_seq, + batch.n_tokens, curr_info.n_blocks_per_seq); + + int32_t token_offset = 0; + for (int seq_id = 0; seq_id < batch_size; ++seq_id) { + llama_sequence_group * group = candidates[seq_id]; + GGML_ASSERT(group && "Make sure the candidates are not nullptr."); + + const bool is_prefill = group->n_decoded == 0; + const int32_t new_tokens = is_prefill ? group->n_prompt : 1; + + if (is_prefill) { + GGML_ASSERT(group->logical_seq.size() >= (size_t) new_tokens && "logical_seq too small for prefill"); + } else { + GGML_ASSERT(!group->logical_seq.empty() && "logical_seq empty during decode"); + } + + for (int token_idx = 0; token_idx < new_tokens; ++token_idx) { + int32_t batch_start_id = token_offset + token_idx; + + batch.token[batch_start_id] = is_prefill ? group->logical_seq[token_idx] : group->logical_seq.back(); + batch.pos[batch_start_id] = group->n_past + token_idx; // n_past starts at 0 + + batch.n_seq_id[batch_start_id] = 1; + batch.seq_id[batch_start_id][0] = group->request_id; + + batch.logits[batch_start_id] = (token_idx == (new_tokens - 1)); // only the last token + + int32_t token_pos = group->n_past + token_idx; + curr_info.write_slots[batch_start_id] = calculate_global_slot_index(token_pos, group->block_table); + LLAMA_LOG_DEBUG("%s: llama_batch seq_id: %d (req_id %d) token %d: pos: %d, global_slot_idx=%d\n", __func__, + seq_id, group->request_id, token_idx, token_pos, curr_info.write_slots[batch_start_id]); + } + + // Populate block table (1D): [batch_size * max_blocks] + const int32_t curr_block_table_size = group->block_table.size(); + for (int block = 0; block < max_blocks; ++block) { + int flattened_id = (seq_id * max_blocks) + block; // row-major + bool need_padding = block >= curr_block_table_size; + curr_info.block_table[flattened_id] = need_padding ? -1 : group->block_table[block]; + } + + curr_info.context_lens[seq_id] = group->n_past + new_tokens; + curr_info.batch_offsets[seq_id] = token_offset; + curr_info.batch_lens[seq_id] = new_tokens; + token_offset += new_tokens; + } +} + +// new_tokens contain 1 token per sequence in the batch +void llama_paged_scheduler_impl::update(const llama_batch & batch, + const std::vector & new_tokens, + const int8_t * stop_flags) { + GGML_ASSERT((int32_t) new_tokens.size() >= curr_info.n_seq && "new_tokens size does not match with batch size."); + GGML_ASSERT(stop_flags != nullptr && "stop_flags can't be null"); + + for (int i = 0; i < curr_info.n_seq; ++i) { + int32_t token_offset = curr_info.batch_offsets[i]; + int32_t request_id = batch.seq_id[token_offset][0]; + + auto it = id_to_group.find(request_id); + if (it == id_to_group.end()) { + LLAMA_LOG_WARN("%s: request_id %d not found in scheduler, skipping\n", __func__, request_id); + continue; + } + + llama_sequence_group * group = it->second; + GGML_ASSERT(group && "group is nullptr."); + + // TTFT + if (group->n_decoded == 0) { + group->t_first_token_us = ggml_time_us(); + } + + // Setting token ranges + llama_pos range_min = kv_cache_manager->seq_pos_min(group->request_id); + if (range_min == -1) { + kv_cache_manager->set_seq_min_pos(group->request_id, batch.pos[token_offset]); + } + int32_t last_token_in_batch_idx = token_offset + curr_info.batch_lens[i] - 1; + kv_cache_manager->set_seq_max_pos(group->request_id, batch.pos[last_token_in_batch_idx]); + + group->n_past += curr_info.batch_lens[i]; + group->n_decoded += curr_info.batch_lens[i]; + group->logical_seq.push_back(new_tokens[i]); + + // Default stop flags are n_seq_max + if (stop_flags[i] || group->n_past >= n_seq_max_ctx) { + group->status = llama_sequence_group_status::FINISHED; + } + } +} + +void llama_paged_scheduler_impl::set_on_finish(llama_paged_on_finish_cb cb, void * user_data) { + on_finish_cb = cb; + on_finish_user_data = user_data; +} + +llama_sequence_group * llama_paged_scheduler_impl::get_group_from_id(int32_t request_id) const { + return id_to_group.count(request_id) ? id_to_group.at(request_id) : nullptr; +} + +const llama_paged_batch_info * llama_paged_scheduler_impl::get_curr_batch_info() const { + return &curr_info; +} diff --git a/src/llama-paged-scheduler-impl.h b/src/llama-paged-scheduler-impl.h new file mode 100644 index 000000000000..0e340c1ec37d --- /dev/null +++ b/src/llama-paged-scheduler-impl.h @@ -0,0 +1,70 @@ +#pragma once + +#include "llama-kv-cache-paged.h" + +#include +#include + +enum class llama_scheduler_status { + OK, + DEADLOCK, // cannot make progress +}; + +class llama_paged_scheduler_impl { + public: + llama_paged_scheduler_impl(uint32_t n_ctx, uint32_t block_sz, int32_t n_batch, llama_kv_cache_paged * kv_manager); + + llama_scheduler_status step(llama_batch & batch); + bool queue_request(llama_sequence_group group); + void update(const llama_batch & batch, const std::vector & new_tokens, const int8_t * stop_flags); + void set_on_finish(llama_paged_on_finish_cb cb, void * user_data); + llama_sequence_group * get_group_from_id(int32_t request_id) const; + const llama_paged_batch_info * get_curr_batch_info() const; + + private: + void insert_sorted_by_arrival_time(llama_sequence_group_ptr new_group, llama_sequence_group_list & list); + + bool check_deadlock(uint32_t n_candidates, uint32_t n_swapped, uint32_t n_waiting) const; + bool check_livelock(uint32_t n_swapped, uint32_t prev_n_swapped); + + void set_running(llama_sequence_group_ptr group); + void set_swapped(llama_sequence_group_ptr group); + void set_waiting(llama_sequence_group_ptr group); + + void finish(llama_sequence_group & group); + + int32_t get_curr_decode_tokens() const; + + void evict(); + void process_running_list(llama_sequence_group_raw_list & candidates); + void process_swapped_list(llama_sequence_group_raw_list & candidates); + void process_waiting_list(llama_sequence_group_raw_list & candidates, int32_t remaining_token_bugdet); + + void swap_out_or_recompute(llama_sequence_group_ptr group_ptr); + + int32_t calculate_global_slot_index(int32_t token_pos, std::vector & block_table); + + void clear_batch(llama_batch & batch); + void populate_batch_from(const llama_sequence_group_raw_list & candidates, llama_batch & batch); + + llama_sequence_group_list running; + llama_sequence_group_list swapped; + llama_sequence_group_list waiting; + + // Used for fast lookups + std::unordered_map id_to_group; + + const uint32_t n_seq_max_ctx; + const uint32_t block_size; + const int32_t n_batch; + llama_kv_cache_paged * kv_cache_manager = nullptr; + llama_paged_batch_info curr_info; + + uint32_t n_livelock_steps = 0; + uint32_t prev_n_swapped = 0; + uint32_t max_livelock_steps = 20; + + // Callback for output tracking + llama_paged_on_finish_cb on_finish_cb = nullptr; + void * on_finish_user_data = nullptr; +}; diff --git a/src/llama-paged-scheduler.cpp b/src/llama-paged-scheduler.cpp new file mode 100644 index 000000000000..c9212e38d2c8 --- /dev/null +++ b/src/llama-paged-scheduler.cpp @@ -0,0 +1,140 @@ +#include "ggml.h" +#include "llama-context.h" +#include "llama-impl.h" +#include "llama-paged-scheduler-impl.h" + +struct llama_paged_scheduler { + llama_paged_scheduler_impl impl; + + llama_paged_scheduler(uint32_t n_ctx, uint32_t block_sz, uint32_t n_batch, llama_kv_cache_paged * kv_manager) : + impl(n_ctx, block_sz, n_batch, kv_manager) {} +}; + +LLAMA_API struct llama_paged_scheduler * llama_paged_scheduler_init(struct llama_context * ctx) { + if (!ctx) { + return nullptr; + } + + // Get the paged kv cache + auto * paged_kv = dynamic_cast(ctx->get_memory()); + if (!paged_kv) { + LLAMA_LOG_ERROR( + "%s: context does not have a paged KV cache. " + "Make sure to pass --kv-paged (-kvp) and use a " + "supported architecture. SWA architectures (gemma3, llama4, etc.) " + "are not yet supported.\n", + __func__); + return nullptr; + } + + // Extract params + const uint32_t n_ctx = ctx->n_ctx(); + const uint32_t block_sz = ctx->block_size(); + const uint32_t n_batch = ctx->n_batch(); + GGML_ASSERT(n_batch == ctx->n_ubatch() && "kv_paged requires n_batch == n_ubatch."); + + try { + return new llama_paged_scheduler(n_ctx, block_sz, n_batch, paged_kv); + } catch (const std::exception & e) { + LLAMA_LOG_ERROR("%s: Error when creating llama_paged_scheduler: %s\n", __func__, e.what()); + return nullptr; + } +} + +LLAMA_API void llama_paged_scheduler_free(struct llama_paged_scheduler * sched) { + if (sched) { + delete sched; + } +} + +LLAMA_API bool llama_paged_scheduler_prepare_batch(struct llama_paged_scheduler * sched, struct llama_batch * batch) { + if (!sched || !batch) { + return false; + } + + llama_scheduler_status status; + try { + status = sched->impl.step(*batch); + } catch (const std::exception & e) { + LLAMA_LOG_ERROR("%s: %s\n", __func__, e.what()); + return false; + } + + if (status == llama_scheduler_status::DEADLOCK) { + LLAMA_LOG_ERROR("%s: Deadlock detected.\n", __func__); + return false; + } + + return true; +} + +LLAMA_API bool llama_paged_scheduler_add_request(struct llama_paged_scheduler * sched, + const llama_token * tokens, + int32_t n_tokens, + int32_t request_id) { + if (!sched || !tokens) { + return false; + } + + llama_sequence_group group; + group.request_id = request_id; + group.n_prompt = n_tokens; + group.n_decoded = 0; + group.n_past = 0; + for (int i = 0; i < n_tokens; ++i) { + group.logical_seq.push_back(tokens[i]); + } + group.t_arrival_time = ggml_time_us(); // int64_t milliseconds + + return sched->impl.queue_request(group); +} + +LLAMA_API void llama_paged_scheduler_update(struct llama_paged_scheduler * sched, + struct llama_batch * batch, + const llama_token * tokens, + const int8_t * stop_flags) { + if (!sched || !batch || !tokens || !stop_flags) { + return; + } + + const auto * info = sched->impl.get_curr_batch_info(); + GGML_ASSERT(info != nullptr && "no batch info was set."); + std::vector tokens_vec(tokens, tokens + info->n_seq); + sched->impl.update(*batch, tokens_vec, stop_flags); +} + +LLAMA_API void llama_paged_scheduler_set_on_finish(struct llama_paged_scheduler * sched, + llama_paged_on_finish_cb cb, + void * user_data) { + sched->impl.set_on_finish(cb, user_data); +} + +LLAMA_API bool llama_paged_scheduler_get_seq_state(struct llama_paged_scheduler * sched, + int32_t request_id, + struct llama_paged_seq_state * out_state) { + if (!sched || !out_state) { + return false; + } + + llama_sequence_group * group = sched->impl.get_group_from_id(request_id); + if (group == nullptr) { + LLAMA_LOG_ERROR("%s: request_id=%d does not exist.", __func__, request_id); + return false; + } + + out_state->request_id = group->request_id; + out_state->n_prompt = group->n_prompt; + out_state->n_decoded = group->n_decoded; + out_state->n_past = group->n_past; + out_state->t_arrival_us = group->t_arrival_time; + out_state->t_first_token_us = group->t_first_token_us; + return true; +} + +LLAMA_API const struct llama_paged_batch_info * llama_paged_scheduler_get_batch_info( + const struct llama_paged_scheduler * sched) { + if (!sched) { + return nullptr; + } + return sched->impl.get_curr_batch_info(); +} diff --git a/src/llama-sequence-group.h b/src/llama-sequence-group.h new file mode 100644 index 000000000000..c34129adc5bf --- /dev/null +++ b/src/llama-sequence-group.h @@ -0,0 +1,31 @@ +#pragma once + +#include "llama.h" + +#include +#include +#include +#include + +enum class llama_sequence_group_status { PENDING, WAITING, RUNNING, SWAPPED, FINISHED }; + +using llama_block_ids = std::vector; + +struct llama_sequence_group { + int32_t request_id = -1; + llama_sequence_group_status status = llama_sequence_group_status::PENDING; + + int64_t t_arrival_time = 0; + int64_t t_first_token_us = 0; + + uint32_t n_prompt = 0; + uint32_t n_decoded = 0; + uint32_t n_past = 0; + + std::vector logical_seq; + llama_block_ids block_table; +}; + +using llama_sequence_group_raw_list = std::vector; +using llama_sequence_group_ptr = std::unique_ptr; +using llama_sequence_group_list = std::list; diff --git a/src/models/command-r.cpp b/src/models/command-r.cpp index 067961caa086..7663da25805b 100644 --- a/src/models/command-r.cpp +++ b/src/models/command-r.cpp @@ -18,7 +18,7 @@ llm_build_command_r::llm_build_command_r(const llama_model & model, const llm_gr // inp_pos - contains the positions ggml_tensor * inp_pos = build_inp_pos(); - auto * inp_attn = build_attn_inp_kv(); + auto * inp_attn = build_attn_inp_kv_auto(); ggml_tensor * inp_out_ids = build_inp_out_ids(); diff --git a/src/models/falcon.cpp b/src/models/falcon.cpp index 2f65fa56e1fd..7c10ac5dd5b7 100644 --- a/src/models/falcon.cpp +++ b/src/models/falcon.cpp @@ -14,7 +14,7 @@ llm_build_falcon::llm_build_falcon(const llama_model & model, const llm_graph_pa // inp_pos - contains the positions ggml_tensor * inp_pos = build_inp_pos(); - auto * inp_attn = build_attn_inp_kv(); + auto * inp_attn = build_attn_inp_kv_auto(); ggml_tensor * inp_out_ids = build_inp_out_ids(); diff --git a/src/models/gemma.cpp b/src/models/gemma.cpp index 09d2ff8bae7f..80bb2f1eef0b 100644 --- a/src/models/gemma.cpp +++ b/src/models/gemma.cpp @@ -14,7 +14,7 @@ llm_build_gemma::llm_build_gemma(const llama_model & model, const llm_graph_para // inp_pos - contains the positions ggml_tensor * inp_pos = build_inp_pos(); - auto * inp_attn = build_attn_inp_kv(); + auto * inp_attn = build_attn_inp_kv_auto(); ggml_tensor * inp_out_ids = build_inp_out_ids(); diff --git a/src/models/gemma3.cpp b/src/models/gemma3.cpp index 0da4af21c173..a6b7c892ed2b 100644 --- a/src/models/gemma3.cpp +++ b/src/models/gemma3.cpp @@ -17,13 +17,13 @@ llm_build_gemma3::llm_build_gemma3(const llama_model & model, const llm_gr ggml_tensor * inp_pos = build_inp_pos(); // TODO: is causal == true correct? might need some changes - using inp_attn_type = std::conditional_t; - inp_attn_type * inp_attn = nullptr; + // using inp_attn_type = std::conditional_t; + llm_graph_input_i * inp_attn = nullptr; if constexpr (iswa) { inp_attn = build_attn_inp_kv_iswa(); } else { - inp_attn = build_attn_inp_kv(); + inp_attn = build_attn_inp_kv_auto(); } ggml_tensor * inp_out_ids = build_inp_out_ids(); diff --git a/src/models/internlm2.cpp b/src/models/internlm2.cpp index 83be2ca0aee6..682110a0beb0 100644 --- a/src/models/internlm2.cpp +++ b/src/models/internlm2.cpp @@ -14,7 +14,7 @@ llm_build_internlm2::llm_build_internlm2(const llama_model & model, const llm_gr // inp_pos - contains the positions ggml_tensor * inp_pos = build_inp_pos(); - auto * inp_attn = build_attn_inp_kv(); + auto * inp_attn = build_attn_inp_kv_auto(); ggml_tensor * inp_out_ids = build_inp_out_ids(); diff --git a/src/models/llama.cpp b/src/models/llama.cpp index 8d478dc67475..2db3be51b897 100644 --- a/src/models/llama.cpp +++ b/src/models/llama.cpp @@ -15,13 +15,11 @@ llm_build_llama::llm_build_llama(const llama_model & model, const llm_gra // inp_pos - contains the positions ggml_tensor * inp_pos = build_inp_pos(); - using inp_attn_type = std::conditional_t; - - inp_attn_type * inp_attn = nullptr; + llm_graph_input_i * inp_attn = nullptr; if constexpr (embed) { inp_attn = build_attn_inp_no_cache(); } else { - inp_attn = build_attn_inp_kv(); + inp_attn = build_attn_inp_kv_auto(); } const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale; diff --git a/src/models/mistral3.cpp b/src/models/mistral3.cpp index b5ae72a2ee13..d7083df75408 100644 --- a/src/models/mistral3.cpp +++ b/src/models/mistral3.cpp @@ -20,7 +20,7 @@ llm_build_mistral3::llm_build_mistral3(const llama_model & model, const llm_grap inp_attn_scale = build_inp_attn_scale(); } - auto * inp_attn = build_attn_inp_kv(); + auto * inp_attn = build_attn_inp_kv_auto(); const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale; diff --git a/src/models/phi3.cpp b/src/models/phi3.cpp index 39af285d3c52..93498c59e2f0 100644 --- a/src/models/phi3.cpp +++ b/src/models/phi3.cpp @@ -14,13 +14,12 @@ llm_build_phi3::llm_build_phi3(const llama_model & model, const llm_graph_ // inp_pos - contains the positions ggml_tensor * inp_pos = build_inp_pos(); - using inp_attn_type = std::conditional_t; - inp_attn_type * inp_attn = nullptr; + llm_graph_input_i * inp_attn = nullptr; if constexpr (iswa) { inp_attn = build_attn_inp_kv_iswa(); } else { - inp_attn = build_attn_inp_kv(); + inp_attn = build_attn_inp_kv_auto(); } ggml_tensor * inp_out_ids = build_inp_out_ids(); diff --git a/src/models/qwen2.cpp b/src/models/qwen2.cpp index 2892dd750876..ae6504d1accf 100644 --- a/src/models/qwen2.cpp +++ b/src/models/qwen2.cpp @@ -14,7 +14,7 @@ llm_build_qwen2::llm_build_qwen2(const llama_model & model, const llm_graph_para // inp_pos - contains the positions ggml_tensor * inp_pos = build_inp_pos(); - auto * inp_attn = build_attn_inp_kv(); + auto * inp_attn = build_attn_inp_kv_auto(); ggml_tensor * inp_out_ids = build_inp_out_ids(); diff --git a/src/models/qwen3.cpp b/src/models/qwen3.cpp index 883dd5f9a905..8eb2aab7976e 100644 --- a/src/models/qwen3.cpp +++ b/src/models/qwen3.cpp @@ -14,7 +14,7 @@ llm_build_qwen3::llm_build_qwen3(const llama_model & model, const llm_graph_para // inp_pos - contains the positions ggml_tensor * inp_pos = build_inp_pos(); - auto * inp_attn = build_attn_inp_kv(); + auto * inp_attn = build_attn_inp_kv_auto(); ggml_tensor * inp_out_ids = build_inp_out_ids(); diff --git a/src/models/starcoder2.cpp b/src/models/starcoder2.cpp index b6d4d5aac1ab..51332fb5a482 100644 --- a/src/models/starcoder2.cpp +++ b/src/models/starcoder2.cpp @@ -14,7 +14,7 @@ llm_build_starcoder2::llm_build_starcoder2(const llama_model & model, const llm_ // inp_pos - contains the positions ggml_tensor * inp_pos = build_inp_pos(); - auto * inp_attn = build_attn_inp_kv(); + auto * inp_attn = build_attn_inp_kv_auto(); ggml_tensor * inp_out_ids = build_inp_out_ids(); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index edb585b9f655..a9e820217525 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -232,6 +232,8 @@ set_tests_properties(test-download-model PROPERTIES FIXTURES_SETUP test-download llama_build_and_test(test-thread-safety.cpp ARGS -m "${MODEL_DEST}" -ngl 99 -p "The meaning of life is" -n 128 -c 256 -ub 32 -np 4 -t 2) set_tests_properties(test-thread-safety PROPERTIES FIXTURES_REQUIRED test-download-model) +llama_build_and_test(test-paged-kv-e2e.cpp ARGS -m "${MODEL_DEST}" -ngl 99) +set_tests_properties(test-paged-kv-e2e PROPERTIES FIXTURES_REQUIRED test-download-model) llama_build_and_test(test-arg-parser.cpp) @@ -300,3 +302,13 @@ if (TARGET gguf-model-data) target_link_libraries(export-graph-ops PRIVATE gguf-model-data) target_compile_definitions(export-graph-ops PRIVATE LLAMA_HF_FETCH) endif() + +# Special-case: test-paged-kv needs access to internal headers in src +add_executable(test-paged-kv test-paged-kv.cpp get-model.cpp) +target_link_libraries(test-paged-kv PRIVATE llama-common llama) +target_include_directories(test-paged-kv PRIVATE ${PROJECT_SOURCE_DIR}/src) +if (LLAMA_TESTS_INSTALL) + install(TARGETS test-paged-kv RUNTIME) +endif() +add_test(NAME test-paged-kv COMMAND $) +set_property(TEST test-paged-kv PROPERTY LABELS main) diff --git a/tests/test-paged-kv-e2e.cpp b/tests/test-paged-kv-e2e.cpp new file mode 100644 index 000000000000..5a352e3ffaa5 --- /dev/null +++ b/tests/test-paged-kv-e2e.cpp @@ -0,0 +1,279 @@ +// tests/test-paged-kv-e2e.cpp +// +// End-to-end equivalence test for paged KV cache. +// We compare top-K agreement rather than raw logit values because the paged +// attention path uses a custom CUDA kernel with online softmax, while the +// unified path uses standard ggml attention with two-pass softmax. The two +// produce mathematically equivalent results but with different F16 +// accumulation order, drifts on the order of 0.05-0.5 in raw +// logit values is expected and not a correctness issue. Top-K set agreement +// is robust to this drift while still catching the real issues (e.g. cross-device +// reads, layout corruption, MQA broadcast bugs) which produce wildly +// different distributions. +// +// Also samples N_COMPARE tokens greedy as a secondary 'cheap' check. + +#include "arg.h" +#include "common.h" +#include "llama.h" +#include "sampling.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#define EXPECT_TRUE(x) \ + do { \ + if (!(x)) { \ + fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, #x); \ + throw std::runtime_error("FAILED assertion."); \ + } \ + } while (0) + +static constexpr const char * TEST_PROMPT = "Once upon a time there was a lovely"; +static constexpr int N_PREDICT = 16; +static constexpr int N_COMPARE = 4; // token-equivalence window +static constexpr int TOP_K = 5; +static constexpr int MIN_TOP_K_OVERLAP = 4; // at least 4 of top-5 must match + +// Result of running one path: prefill-final logits + sampled token sequence. +struct path_result { + std::vector prefill_logits; // [n_vocab] + std::vector tokens; // [N_PREDICT] + int n_vocab = 0; +}; + +static path_result run_non_paged(const std::string & model_path) { + common_params params; + params.model.path = model_path; + params.n_ctx = 256; + params.n_batch = 64; + params.n_ubatch = 64; + params.n_predict = N_PREDICT; + params.sampling.temp = 0.0f; // greedy + params.warmup = false; + params.kv_paged = false; + + auto init = common_init_from_params(params); + llama_model * model = init->model(); + llama_context * ctx = init->context(); + EXPECT_TRUE(model != nullptr); + EXPECT_TRUE(ctx != nullptr); + + const llama_vocab * vocab = llama_model_get_vocab(model); + const int n_vocab = llama_vocab_n_tokens(vocab); + + std::vector prompt_tokens = common_tokenize(ctx, TEST_PROMPT, true); + EXPECT_TRUE(!prompt_tokens.empty()); + + // Prefill + llama_batch batch = llama_batch_get_one(prompt_tokens.data(), prompt_tokens.size()); + EXPECT_TRUE(llama_decode(ctx, batch) == 0); + + // Capture prefill-final logits BEFORE any further decode steps overwrite them. + path_result result; + result.n_vocab = n_vocab; + { + const float * raw = llama_get_logits_ith(ctx, -1); // last logit + EXPECT_TRUE(raw != nullptr); + result.prefill_logits.assign(raw, raw + n_vocab); + } + + // Sample N_PREDICT tokens for the secondary token-equivalence check. + common_sampler * smpl = common_sampler_init(model, params.sampling); + EXPECT_TRUE(smpl != nullptr); + + llama_token cur = -1; + for (int i = 0; i < N_PREDICT; ++i) { + cur = common_sampler_sample(smpl, ctx, -1); + common_sampler_accept(smpl, cur, true); + result.tokens.push_back(cur); + if (llama_vocab_is_eog(vocab, cur)) { + break; + } + + llama_batch step = llama_batch_get_one(&cur, 1); + EXPECT_TRUE(llama_decode(ctx, step) == 0); + } + + common_sampler_free(smpl); + return result; +} + +static path_result run_paged(const std::string & model_path) { + common_params params; + params.model.path = model_path; + params.n_ctx = 256; + params.n_batch = 64; + params.n_ubatch = 64; + params.n_predict = N_PREDICT; + params.sampling.temp = 0.0f; // greedy + params.warmup = false; + params.kv_paged = true; + params.n_gpu_blocks = 64; + params.n_cpu_blocks = 16; + params.n_sequences = 1; + params.n_parallel = 1; + + auto init = common_init_from_params(params); + llama_model * model = init->model(); + llama_context * ctx = init->context(); + EXPECT_TRUE(model != nullptr); + EXPECT_TRUE(ctx != nullptr); + + const llama_vocab * vocab = llama_model_get_vocab(model); + const int n_vocab = llama_vocab_n_tokens(vocab); + + llama_paged_scheduler * sched = llama_paged_scheduler_init(ctx); + EXPECT_TRUE(sched != nullptr); + + std::vector prompt_tokens = common_tokenize(ctx, TEST_PROMPT, true); + EXPECT_TRUE(!prompt_tokens.empty()); + + bool ok = llama_paged_scheduler_add_request(sched, prompt_tokens.data(), prompt_tokens.size(), 0); + EXPECT_TRUE(ok); + + common_sampler * smpl = common_sampler_init(model, params.sampling); + EXPECT_TRUE(smpl != nullptr); + + path_result result; + result.n_vocab = n_vocab; + bool captured_prefill_logits = false; + llama_batch batch = {}; + + while ((int) result.tokens.size() < N_PREDICT) { + bool prepared = llama_paged_scheduler_prepare_batch(sched, &batch); + EXPECT_TRUE(prepared); + if (batch.n_tokens == 0) { + break; + } + + EXPECT_TRUE(llama_decode(ctx, batch) == 0); + llama_synchronize(ctx); + + const llama_paged_batch_info * info = llama_paged_scheduler_get_batch_info(sched); + EXPECT_TRUE(info != nullptr && info->n_seq == 1); + + const int32_t last_idx = info->batch_offsets[0] + info->batch_lens[0] - 1; + + // First decode is the prefill — capture its final logits before + // sampling anything else. + if (!captured_prefill_logits) { + const float * raw = llama_get_logits_ith(ctx, last_idx); + EXPECT_TRUE(raw != nullptr); + result.prefill_logits.assign(raw, raw + n_vocab); + captured_prefill_logits = true; + } + + llama_token next = common_sampler_sample(smpl, ctx, last_idx); + common_sampler_accept(smpl, next, true); + result.tokens.push_back(next); + + bool stop = llama_vocab_is_eog(vocab, next) || (int) result.tokens.size() >= N_PREDICT; + int8_t stop_flag = stop ? 1 : 0; + llama_paged_scheduler_update(sched, &batch, &next, &stop_flag); + if (stop) { + break; + } + } + + common_sampler_free(smpl); + llama_paged_scheduler_free(sched); + return result; +} + +static void compare_results(const path_result & ref, const path_result & paged) { + auto top_k = [](const std::vector & l, int k) { + std::vector idx(l.size()); + std::iota(idx.begin(), idx.end(), 0); + std::partial_sort(idx.begin(), idx.begin() + k, idx.end(), [&l](int a, int b) { return l[a] > l[b]; }); + idx.resize(k); + return idx; + }; + + const auto top_ref = top_k(ref.prefill_logits, TOP_K); + const auto top_paged = top_k(paged.prefill_logits, TOP_K); + + // Argmax must match: the most-confident next token should be identical. + EXPECT_TRUE(top_ref[0] == top_paged[0]); + + // Top-K set overlap: at least MIN_TOP_K_OVERLAP of the K most likely + // tokens must appear in both distributions. + std::set ref_set(top_ref.begin(), top_ref.end()); + int overlap = 0; + for (int t : top_paged) { + if (ref_set.count(t)) { + overlap++; + } + } + + fprintf(stderr, "test-paged-kv-e2e: top-%d argmax match: ref=%d paged=%d\n", TOP_K, top_ref[0], top_paged[0]); + fprintf(stderr, "test-paged-kv-e2e: top-%d set overlap: %d/%d (require >= %d)\n", TOP_K, overlap, TOP_K, + MIN_TOP_K_OVERLAP); + + if (overlap < MIN_TOP_K_OVERLAP) { + fprintf(stderr, + "FAIL: top-%d distributions diverge too much. Only %d of %d most-likely " + "tokens match between ref and paged. Real correctness issue likely.\n", + TOP_K, overlap, TOP_K); + fprintf(stderr, " ref: "); + for (int t : top_ref) { + fprintf(stderr, "%d(%.3f) ", t, ref.prefill_logits[t]); + } + fprintf(stderr, "\n paged: "); + for (int t : top_paged) { + fprintf(stderr, "%d(%.3f) ", t, paged.prefill_logits[t]); + } + fprintf(stderr, "\n"); + throw std::runtime_error("FAILED test."); + } + + // Token-level secondary check: first N_COMPARE tokens must match. + // We don't compare beyond N_COMPARE because greedy sampling on small + // models is sensitive to argmax tiebreakers, and minor floating-point + // accumulation differences between the paged and non-paged paths can + // flip individual tokens after a handful of decode steps. + EXPECT_TRUE((int) ref.tokens.size() >= N_COMPARE); + EXPECT_TRUE((int) paged.tokens.size() >= N_COMPARE); + for (int i = 0; i < N_COMPARE; ++i) { + if (ref.tokens[i] != paged.tokens[i]) { + fprintf(stderr, "FAIL: token %d differs in the equivalence window: ref=%d paged=%d\n", i, ref.tokens[i], + paged.tokens[i]); + throw std::runtime_error("FAILED test."); + } + } + fprintf(stderr, "test-paged-kv-e2e: PASSED\n"); +} + +int main(int argc, char ** argv) { + common_params params; + if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_PAGED)) { + fprintf(stderr, "usage: %s -m \n", argv[0]); + return 1; + } + if (params.model.path.empty()) { + fprintf(stderr, "skip: no --model provided\n"); + return 0; + } + + common_init(); + llama_backend_init(); + + fprintf(stderr, "test-paged-kv-e2e: running non-paged reference\n"); + path_result ref = run_non_paged(params.model.path); + fprintf(stderr, " got %zu tokens, %d-vocab logits\n", ref.tokens.size(), ref.n_vocab); + + fprintf(stderr, "test-paged-kv-e2e: running paged path\n"); + path_result paged = run_paged(params.model.path); + fprintf(stderr, " got %zu tokens, %d-vocab logits\n", paged.tokens.size(), paged.n_vocab); + + compare_results(ref, paged); + + llama_backend_free(); + return 0; +} diff --git a/tests/test-paged-kv.cpp b/tests/test-paged-kv.cpp new file mode 100644 index 000000000000..907f7e4c2d7b --- /dev/null +++ b/tests/test-paged-kv.cpp @@ -0,0 +1,417 @@ +#include "ggml-backend.h" +#include "llama-block-manager.h" +#include "llama-kv-cache-paged.h" +#include "llama-paged-scheduler-impl.h" + +#include +#include +#include + +#define TEST(name) static void name() +#define RUN(name) \ + do { \ + fprintf(stderr, " running %-40s ", #name); \ + fflush(stderr); \ + name(); \ + fprintf(stderr, "OK\n"); \ + } while (0) + +#define EXPECT_EQ(a, b) \ + do { \ + auto _a = (a); \ + auto _b = (b); \ + if (_a != _b) { \ + fprintf(stderr, "\n FAIL %s:%d: expected %s == %s, got %lld vs %lld\n", __FILE__, __LINE__, #a, #b, \ + (long long) _a, (long long) _b); \ + } \ + } while (0) + +#define EXPECT_TRUE(x) \ + do { \ + if (!(x)) { \ + fprintf(stderr, "\n FAIL %s:%d: %s\n", __FILE__, __LINE__, #x); \ + std::abort(); \ + } \ + } while (0) +#define EXPECT_FALSE(x) EXPECT_TRUE(!(x)) + +// Testing block_manager main functionality + +TEST(test_block_manager_leak_simple) { + const uint32_t n_gpu_blocks = 16; + const uint32_t n_cpu_blocks = 8; + const float watermark = 0.0f; + + llama_block_manager block_manager; + block_manager.init(n_gpu_blocks, n_cpu_blocks, watermark); + + EXPECT_EQ(block_manager.n_free_gpu_blocks(), n_gpu_blocks); + EXPECT_EQ(block_manager.n_free_cpu_blocks(), n_cpu_blocks); + + auto gpu_ids = block_manager.checkout_gpu_blocks(10); + EXPECT_EQ(gpu_ids.size(), 10u); + EXPECT_EQ(block_manager.n_free_gpu_blocks(), (n_gpu_blocks - 10u)); + + auto cpu_ids = block_manager.checkout_cpu_blocks(5); + EXPECT_EQ(cpu_ids.size(), 5u); + EXPECT_EQ(block_manager.n_free_cpu_blocks(), (n_cpu_blocks - 5u)); + + block_manager.release_gpu_blocks(gpu_ids); + block_manager.release_cpu_blocks(cpu_ids); + + EXPECT_EQ(block_manager.n_free_gpu_blocks(), n_gpu_blocks); + EXPECT_EQ(block_manager.n_free_cpu_blocks(), n_cpu_blocks); +} + +// Testing that we always return to full (stress test) +TEST(test_block_manager_leak_repeated) { + const uint32_t block_size = 16; + const uint32_t n_gpu_blocks = 64; + const uint32_t n_cpu_blocks = 32; + const float watermark = 0.0f; + const int n_iter = 1000; + + llama_block_manager block_manager; + block_manager.init(n_gpu_blocks, n_cpu_blocks, watermark); + + for (int iter = 0; iter < n_iter; ++iter) { + const uint32_t n = (iter % block_size) + 1; + auto ids = block_manager.checkout_gpu_blocks(n); + EXPECT_EQ(ids.size(), (size_t) n); + block_manager.release_gpu_blocks(ids); + EXPECT_EQ(block_manager.n_free_gpu_blocks(), n_gpu_blocks); + } +} + +// Attempting to check-out more blocks than available. It should return empty. +TEST(test_block_manager_checkout_too_many) { + const uint32_t n_gpu_blocks = 8; + const uint32_t n_cpu_blocks = 4; + const float watermark = 0.0f; + + llama_block_manager block_manager; + block_manager.init(n_gpu_blocks, n_cpu_blocks, watermark); + + auto ids = block_manager.checkout_gpu_blocks(9); + EXPECT_EQ(ids.size(), 0u); + EXPECT_EQ(block_manager.n_free_gpu_blocks(), n_gpu_blocks); + + auto ids2 = block_manager.checkout_cpu_blocks(5); + EXPECT_EQ(ids2.size(), 0u); + EXPECT_EQ(block_manager.n_free_cpu_blocks(), n_cpu_blocks); +} + +TEST(test_block_manager_watermark) { + // watermark=0.2 means 2 blocks reserved as safety (always consider max gpu blocks) + const uint32_t n_gpu_blocks = 10; + const uint32_t n_cpu_blocks = 10; + const float watermark = 0.2f; + + llama_block_manager block_manager; + block_manager.init(n_gpu_blocks, n_cpu_blocks, watermark); + + // 10 free, safety=2 + EXPECT_TRUE(block_manager.has_free_gpu_blocks(8)); + EXPECT_FALSE(block_manager.has_free_gpu_blocks(9)); + EXPECT_FALSE(block_manager.has_free_gpu_blocks(10)); + + auto ids = block_manager.checkout_gpu_blocks(5); + EXPECT_EQ(ids.size(), 5u); + // 5 free, safety=2 + EXPECT_TRUE(block_manager.has_free_gpu_blocks(3)); + EXPECT_FALSE(block_manager.has_free_gpu_blocks(4)); + + block_manager.release_gpu_blocks(ids); + EXPECT_TRUE(block_manager.has_free_gpu_blocks(8)); +} + +TEST(test_block_manager_watermark_zero) { + // watermark=0 means the entire pool is requestable. + const uint32_t n_gpu_blocks = 10; + const uint32_t n_cpu_blocks = 10; + const float watermark = 0.0f; + + llama_block_manager block_manager; + block_manager.init(n_gpu_blocks, n_cpu_blocks, watermark); + EXPECT_TRUE(block_manager.has_free_gpu_blocks(10)); + EXPECT_FALSE(block_manager.has_free_gpu_blocks(11)); +} + +TEST(test_block_manager_gpu_cpu_disjoint) { + // Just making sure CPU and GPU blocks are disjoint + const uint32_t n_gpu_blocks = 8; + const uint32_t n_cpu_blocks = 4; + const float watermark = 0.0f; + + llama_block_manager block_manager; + block_manager.init(n_gpu_blocks, n_cpu_blocks, watermark); + + auto gpu_ids = block_manager.checkout_gpu_blocks(8); + auto cpu_ids = block_manager.checkout_cpu_blocks(4); + + for (auto id : gpu_ids) { + EXPECT_TRUE(block_manager.is_gpu(id)); + } + for (auto id : cpu_ids) { + EXPECT_FALSE(block_manager.is_gpu(id)); + } + + block_manager.release_gpu_blocks(gpu_ids); + block_manager.release_cpu_blocks(cpu_ids); +} + +// Testing llama_kv_cache_paged book-keeping + +// kv_cache_paged with arbitrary shape (we only care about the bookkeeping) +static llama_kv_cache_paged make_kv() { + return llama_kv_cache_paged( + /*head_dim=*/64, + /*n_heads_kv=*/4, + /*block_size=*/16, + /*n_layers=*/2, + /*n_ubatch=*/32, + /*n_seq_max=*/8); +} + +TEST(test_seq_pos_default_unknown) { + auto kv = make_kv(); + EXPECT_EQ(kv.seq_pos_min(0), -1); + EXPECT_EQ(kv.seq_pos_max(0), -1); + EXPECT_EQ(kv.seq_pos_min(42), -1); +} + +TEST(test_seq_pos_set_and_get) { + auto kv = make_kv(); + kv.set_seq_min_pos(0, 5); + kv.set_seq_max_pos(0, 17); + EXPECT_EQ(kv.seq_pos_min(0), 5); + EXPECT_EQ(kv.seq_pos_max(0), 17); + + // Independent per seq_id. + kv.set_seq_min_pos(1, 100); + EXPECT_EQ(kv.seq_pos_min(1), 100); + EXPECT_EQ(kv.seq_pos_min(0), 5); // unchanged +} + +TEST(test_seq_pos_overwrite) { + auto kv = make_kv(); + kv.set_seq_min_pos(0, 5); + kv.set_seq_min_pos(0, 9); + EXPECT_EQ(kv.seq_pos_min(0), 9); +} + +TEST(test_seq_pos_seq_rm_removes) { + auto kv = make_kv(); + kv.set_seq_min_pos(0, 5); + kv.set_seq_max_pos(0, 17); + EXPECT_EQ(kv.seq_pos_min(0), 5); + + kv.seq_rm(0, 0, 0); + EXPECT_EQ(kv.seq_pos_min(0), -1); + EXPECT_EQ(kv.seq_pos_max(0), -1); +} + +TEST(test_seq_pos_clear_removes_all) { + auto kv = make_kv(); + kv.set_seq_min_pos(0, 5); + kv.set_seq_min_pos(1, 7); + kv.set_seq_min_pos(2, 9); + + kv.clear(/*data=*/false); + EXPECT_EQ(kv.seq_pos_min(0), -1); + EXPECT_EQ(kv.seq_pos_min(1), -1); + EXPECT_EQ(kv.seq_pos_min(2), -1); +} + +TEST(test_free_blocks_releases_to_pool) { + // Initialize a KV cache with a real CPU backend (used as both "GPU" and CPU). + // This is fine for testing + ggml_backend_t backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + EXPECT_TRUE(backend != nullptr); + + const uint32_t n_gpu_blocks = 16; + const uint32_t n_cpu_blocks = 8; + const float watermark = 0.0f; + + auto kv = make_kv(); + kv.init(/*backend_gpu=*/backend, + /*backend_cpu=*/backend, GGML_TYPE_F16, n_gpu_blocks, n_cpu_blocks, watermark); + + // allocate() pulls from the GPU block pool. After releasing, the count + // must return to the initial value. + const uint32_t n_gpu_initial = kv.get_num_gpu_blocks(); + EXPECT_EQ(n_gpu_initial, n_gpu_blocks); + + llama_sequence_group group; + group.request_id = 0; + group.n_prompt = 32; // 2 blocks + group.n_decoded = 0; + + bool success = kv.allocate(/*num_tokens=*/0, group); + EXPECT_TRUE(success); + EXPECT_EQ(group.block_table.size(), 2u); + + // Allocating a second sequence further reduces the pool. + llama_sequence_group group2; + group2.request_id = 1; + group2.n_prompt = 48; // 3 blocks + group2.n_decoded = 0; + + success = kv.allocate(0, group2); + EXPECT_TRUE(success); + EXPECT_EQ(group2.block_table.size(), 3u); + + // Free the first sequence: we release 2 blocks + kv.free_blocks(group); + EXPECT_EQ(group.block_table.size(), 0u); + + llama_sequence_group group3; + group3.request_id = 2; + group3.n_prompt = 32; // 2 blocks + group3.n_decoded = 0; + success = kv.allocate(0, group3); + EXPECT_TRUE(success); + EXPECT_EQ(group3.block_table.size(), 2u); + + // All blocks released + kv.free_blocks(group2); + kv.free_blocks(group3); + + llama_sequence_group group_full; + group_full.request_id = 99; + group_full.n_prompt = n_gpu_blocks * 16; // exactly n_gpu_blocks worth + group_full.n_decoded = 0; + success = kv.allocate(0, group_full); + EXPECT_TRUE(success); + EXPECT_EQ(group_full.block_table.size(), 16u); + kv.free_blocks(group_full); + + ggml_backend_free(backend); +} + +// Testing scheduler + +// For easy testing and clean-up. +// KV cache paged + scheduler (must free CPU backend) +struct paged_test_fixture { + ggml_backend_t backend = nullptr; + std::unique_ptr kv; + std::unique_ptr sched; + + ~paged_test_fixture() { + if (backend) { + ggml_backend_free(backend); + } + } + + // rule of 5 + paged_test_fixture() = default; + paged_test_fixture(paged_test_fixture &&) = default; + paged_test_fixture & operator=(paged_test_fixture &&) = default; + paged_test_fixture(const paged_test_fixture &) = delete; + paged_test_fixture & operator=(const paged_test_fixture &) = delete; +}; + +static paged_test_fixture make_fixture(uint32_t n_ctx = 128, + uint32_t block_size = 16, + uint32_t n_batch = 64, + uint32_t n_gpu_blocks = 4, + uint32_t n_cpu_blocks = 2) { + paged_test_fixture fixture; + fixture.backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + EXPECT_TRUE(fixture.backend != nullptr); + + fixture.kv = std::unique_ptr(new llama_kv_cache_paged( + /*head_dim=*/64u, + /*n_heads_kv=*/4u, + /*block_size=*/block_size, + /*n_layers=*/2u, + /*n_ubatch=*/n_batch, + /*n_seq_max=*/8u)); + fixture.kv->init(fixture.backend, fixture.backend, GGML_TYPE_F16, n_gpu_blocks, n_cpu_blocks, /*watermark=*/0.0f); + + fixture.sched = std::unique_ptr( + new llama_paged_scheduler_impl(n_ctx, block_size, n_batch, fixture.kv.get())); + return fixture; +} + +static llama_sequence_group make_group(int32_t request_id, uint32_t n_prompt) { + llama_sequence_group group; + group.request_id = request_id; + group.n_prompt = n_prompt; + group.n_decoded = 0; + group.n_past = 0; + group.t_arrival_time = request_id; // control ordering based on request_id + group.logical_seq.assign(n_prompt, /*dummy token=*/1); + return group; +} + +TEST(test_scheduler_no_deadlock_on_empty) { + // No requests queued means no deadlock + auto fixture = make_fixture(); + llama_batch batch = {}; + llama_scheduler_status status = fixture.sched->step(batch); + EXPECT_TRUE(status == llama_scheduler_status::OK); + EXPECT_EQ(batch.n_tokens, 0); +} + +TEST(test_scheduler_deadlock_oversize_waiting_request) { + // There are 2 blocks, the waiting request needs 3. + // Attempt to process prefill for this requets will fail and the request will remain in waiting (deadlock). + auto fixture = make_fixture(128, 16, 64, /*n_gpu_blocks=*/2, /*n_cpu_blocks=*/1); + auto group = make_group(0, /*n_prompts=*/32); + + bool queued = fixture.sched->queue_request(group); + EXPECT_TRUE(queued); + + llama_batch batch = {}; + llama_scheduler_status status = fixture.sched->step(batch); + EXPECT_TRUE(status == llama_scheduler_status::DEADLOCK); + + // Next steps will continue remained deadlocked + status = fixture.sched->step(batch); + EXPECT_TRUE(status == llama_scheduler_status::DEADLOCK); +} + +TEST(test_scheduler_rejects_oversized_prompt) { + auto fixture = make_fixture(/*n_ctx=*/64, /*block_size=*/16, /*n_batch=*/128, + /*n_gpu_blocks=*/32, /*n_cpu_blocks=*/8); + + bool queued = fixture.sched->queue_request(make_group(/*id=*/0, /*n_prompt=*/64)); + EXPECT_FALSE(queued); + + queued = fixture.sched->queue_request(make_group(/*id=*/1, /*n_prompt=*/128)); + EXPECT_FALSE(queued); + + // A request just under the limit is accepted. + queued = fixture.sched->queue_request(make_group(/*id=*/2, /*n_prompt=*/63)); + EXPECT_TRUE(queued); +} + +int main(int /*argc*/, char ** /*argv*/) { + fprintf(stderr, "test-paged-kv: block_manager\n"); + RUN(test_block_manager_leak_simple); + RUN(test_block_manager_leak_repeated); + RUN(test_block_manager_checkout_too_many); + RUN(test_block_manager_watermark); + RUN(test_block_manager_watermark_zero); + RUN(test_block_manager_gpu_cpu_disjoint); + + fprintf(stderr, "test-paged-kv: llama_kv_cache_paged seq_pos\n"); + RUN(test_seq_pos_default_unknown); + RUN(test_seq_pos_set_and_get); + RUN(test_seq_pos_overwrite); + RUN(test_seq_pos_seq_rm_removes); + RUN(test_seq_pos_clear_removes_all); + + fprintf(stderr, "test-paged-kv: llama_kv_cache_paged free_blocks\n"); + RUN(test_free_blocks_releases_to_pool); + + fprintf(stderr, "test-paged-kv: llama_kv_cache_paged scheduler\n"); + RUN(test_scheduler_no_deadlock_on_empty); + RUN(test_scheduler_deadlock_oversize_waiting_request); + RUN(test_scheduler_rejects_oversized_prompt); + + fprintf(stderr, "test-paged-kv: ALL PASSED\n"); + return 0; +}