8000 feat: Hybrid unified/recurrent cache by gabe-l-hart · Pull Request #13276 · ggml-org/llama.cpp · GitHub
[go: up one dir, main page]
Skip to content

feat: Hybrid unified/recurrent cache #13276

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 21 commits into from
Closed

feat: Hybrid unified/recurrent cache #13276

Changes from 1 commit
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
0e9f0e0
tests: Initial unit tests for memory hierarchy
gabe-l-hart May 20, 2025
b656613
build: Add build step for test-memory on non-windows builds
gabe-l-hart May 20, 2025
9a15f27
fix(tests): Fix constructors in tests for signature changes after rebase
gabe-l-hart May 27, 2025
d74d76a
tests(wip): More robust test for unified cache
gabe-l-hart May 23, 2025
114f2ce
feat: Add can_seq_rm API to llama_kv_cache API
gabe-l-hart May 27, 2025
13efcf3
feat: Move layer_filter_cb up to llama_kv_cache
gabe-l-hart May 20, 2025
6625248
feat: Add layer filter to recurrent cache
gabe-l-hart May 20, 2025
76771e8
feat: Initial implementation of llama_kv_cache_hybrid
gabe-l-hart May 20, 2025
f031fb8
feat: Add llama_model_is_hybrid API call
gabe-l-hart May 9, 2025
298e147
feat: Add c++ side constants for attention layer indices hparam
gabe-l-hart May 9, 2025
f472373
feat: Add support for distinguishing recurrent vs non-recurrent layer…
gabe-l-hart May 9, 10000 2025
404783b
feat: Auto-fill hparams.recurrent_layer_arr based on whether the mode…
gabe-l-hart May 9, 2025
53ef2d4
feat: Instantiate hybrid cache for hybrid models (currently none)
gabe-l-hart May 20, 2025
e6ff93a
refactor: rename *_is_hybrid -> *_is_hybrid_recurrent
gabe-l-hart May 20, 2025
226955b
fix: Fix indexing into k_l for recurrent cache with filter
gabe-l-hart May 20, 2025
1fb08da
fix: Use per-layer sizing everywhere in kv caches
gabe-l-hart May 14, 2025
04fe482
fix: Remove unused kv cache methods after rebase
gabe-l-hart May 23, 2025
db9a618
fix(tests): Fix constructors in tests for signature changes after rebase
gabe-l-hart May 23, 2025
8aee2e7
feat: Add split_equal to init(...) signature
gabe-l-hart May 27, 2025
a4cc4aa
fix: Overhaul hybrid cache for refactor part3 (::init interface)
gabe-l-hart May 27, 2025
58994f6
tests(wip): Comment out broken test for now and fix other constructor…
gabe-l-hart May 27, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
tests(wip): Comment out broken test for now and fix other constructor… 
… signatures

Branch: HybridCache

Signed-off-by: Gabe Goodhart <ghart@us.ibm.com>
  • Loading branch information
@gabe-l-hart
gabe-l-hart committed May 27, 2025
commit 58994f653086b19be3d8d988f7176106bb60cdec
86 changes: 43 additions & 43 deletions tests/test-memory.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,48 +158,48 @@ static void test_llama_kv_cache_unified_single_seq() {
/* swa_type */ LLAMA_SWA_TYPE_NONE
);

// Create the micro batch with a single 3-token sequence
llama_batch batch1 = _make_batch({{101, 1, 102}}, {{42}});
llama_sbatch sbatch1 = cache.sbatch_init(batch1, false);
llama_ubatch ubatch1 = cache.ubatch_next(sbatch1, 4, false);

// Find a slot for a new sequence
GGML_ASSERT(cache.find_slot(ubatch1));

// Cache the k/v for a single layer in this slot
ggml_context * ctx = ggml_init({10240, NULL, false});
ggml_tensor * k1 = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, model->hparams.n_embd_k_gqa(0));
ggml_tensor * v1 = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, model->hparams.n_embd_v_gqa(0));
ggml_tensor * k1_view = cache.cpy_k(ctx, k1, 0);
ggml_tensor * v1_view = cache.cpy_v(ctx, v1, 0);
GGML_ASSERT(is_source_tensor(k1_view, k1));
GGML_ASSERT(is_source_tensor(v1_view, v1));

// Create a second batch with different tokens and find a slot for it
llama_batch batch2 = _make_batch({{1, 2, 3, 4}}, {{5}});
llama_sbatch sbatch2 = cache.sbatch_init(batch2, false);
llama_ubatch ubatch2 = cache.ubatch_next(sbatch2, 4, false);
GGML_ASSERT(cache.find_slot(ubatch2));

// Add some different tensors
ggml_tensor * k2 = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, model->hparams.n_embd_k_gqa(0));
ggml_tensor * v2 = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, model->hparams.n_embd_v_gqa(0));
ggml_tensor * k2_view = cache.cpy_k(ctx, k2, 0);
ggml_tensor * v2_view = cache.cpy_v(ctx, v2, 0);
GGML_ASSERT(is_source_tensor(k2_view, k2));
GGML_ASSERT(is_source_tensor(v2_view, v2));

// Make sure first batch's k/v aren't cache hit
GGML_ASSERT(!is_source_tensor(k2_view, k1));
GGML_ASSERT(!is_source_tensor(v2_view, v1));

// Re-find the slot for the first batch and make sure they cache hit
GGML_ASSERT(cache.find_slot(ubatch1));

// Clean up
llama_batch_free(batch1);
llama_batch_free(batch2);
ggml_free(ctx);
// // Create the micro batch with a single 3-token sequence
// llama_batch batch1 = _make_batch({{101, 1, 102}}, {{42}});
// llama_sbatch sbatch1 = cache.sbatch_init(batch1, false);
// llama_ubatch ubatch1 = cache.ubatch_next(sbatch1, 4, false);

// // Find a slot for a new sequence
// GGML_ASSERT(cache.find_slot(ubatch1));

// // Cache the k/v for a single layer in this slot
// ggml_context * ctx = ggml_init({10240, NULL, false});
// ggml_tensor * k1 = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, model->hparams.n_embd_k_gqa(0));
// ggml_tensor * v1 = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, model->hparams.n_embd_v_gqa(0));
// ggml_tensor * k1_view = cache.cpy_k(ctx, k1, 0);
// ggml_tensor * v1_view = cache.cpy_v(ctx, v1, 0);
// GGML_ASSERT(is_source_tensor(k1_view, k1));
// GGML_ASSERT(is_source_tensor(v1_view, v1));

// // Create a second batch with different tokens and find a slot for it
// llama_batch batch2 = _make_batch({{1, 2, 3, 4}}, {{5}});
// llama_sbatch sbatch2 = cache.sbatch_init(batch2, false);
// llama_ubatch ubatch2 = cache.ubatch_next(sbatch2, 4, false);
// GGML_ASSERT(cache.find_slot(ubatch2));

// // Add some different tensors
// ggml_tensor * k2 = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, model->hparams.n_embd_k_gqa(0));
// ggml_tensor * v2 = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, model->hparams.n_embd_v_gqa(0));
// ggml_tensor * k2_view = cache.cpy_k(ctx, k2, 0);
// ggml_tensor * v2_view = cache.cpy_v(ctx, v2, 0);
// GGML_ASSERT(is_source_tensor(k2_view, k2));
// GGML_ASSERT(is_source_tensor(v2_view, v2));

// // Make sure first batch's k/v aren't cache hit
// GGML_ASSERT(!is_source_tensor(k2_view, k1));
// GGML_ASSERT(!is_source_tensor(v2_view, v1));

// // Re-find the slot for the first batch and make sure they cache hit
// GGML_ASSERT(cache.find_slot(ubatch1));

// // Clean up
// llama_batch_free(batch1);
// llama_batch_free(batch2);
// ggml_free(ctx);
}

/*- Recurrent Cache ----------------------------------------------------------*/
Expand Down Expand Up @@ -280,7 +280,7 @@ static void test_llama_kv_cache_hybrid_constructor() {
children.emplace_back(std::move(u_cache), std::vector<size_t>{1, 3});
children.emplace_back(std::move(r_cache), std::vector<size_t>{0, 2});

llama_kv_cache_hybrid cache(model->hparams, std::move(children));
llama_kv_cache_hybrid cache(std::move(children));

GGML_ASSERT(cache.get_child_cache<llama_kv_cache_unified>() == u_cache_ptr);
GGML_ASSERT(cache.get_child_cache<llama_kv_cache_recurrent>() == r_cache_ptr);
Expand Down
Loading
0