Skip to content

kv-cache : add SWA support #13194

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

Merged
merged 16 commits into from
May 20, 2025
+1,414 −638
Merged

kv-cache : add SWA support #13194

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
c011e4e
kv-cache : prepare for SWA
ggerganov Apr 28, 2025
85f5fc5
kv-cache : initial iSWA implementation
ggerganov May 11, 2025
b9ce306
kv-cache : rework error recovery logic
ggerganov May 12, 2025
a4aafa5
models : fix Phi-3 SWA parameters
ggerganov May 12, 2025
c7d8175
model : adjust Granite to rope factor changes
ggerganov May 14, 2025
554b4d0
server : check if context can do shifts
ggerganov May 11, 2025
4a258ff
iswa : for now, always enable shifts (experiment)
ggerganov May 11, 2025
e743246
kv-cache : simplify SWA logic
ggerganov May 15, 2025
6390125
kv-cache : apply defrag when we fail to find slots for the batch
ggerganov May 15, 2025
86c526a
llama : update docs about llama_decode
ggerganov May 15, 2025
0073157
kv-cache : update warning logs when no space for the batch is available
ggerganov May 17, 2025
12ee6db
llama : add llama_kv_self_seq_pos_min()
ggerganov May 17, 2025
ca52e19
kv-cache : keep track of partial SWA computes and print warnings
ggerganov May 17, 2025
84742ef
server : disallow use cases involving partial SWA context
ggerganov May 17, 2025
c699abc
llama : add param to control SWA cache size
ggerganov May 18, 2025
5db1110
minor : clean-up
ggerganov May 20, 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
Next Next commit
models : fix Phi-3 SWA parameters 
ggml-ci
  • Loading branch information
@ggerganov
ggerganov committed May 17, 2025
commit a4aafa5374975253e6d6e2a519251e91849224be
5 changes: 2 additions & 3 deletions src/llama-graph.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1316,9 +1316,8 @@ llm_graph_input_attn_kv_unified_iswa * llm_graph_context::build_attn_inp_kv_unif
inp->self_kq_mask_cnv = cparams.flash_attn ? ggml_cast(ctx0, inp->self_kq_mask, GGML_TYPE_F16) : inp->self_kq_mask;
}

{
GGML_ASSERT(hparams.n_swa_pattern > 1 && "Use llama_kv_cache_unified for non-SWA");
GGML_ASSERT(hparams.n_swa > 0 && "Use llama_kv_cache_unified for non-SWA");
if (hparams.n_swa_pattern > 1) {
GGML_ASSERT(hparams.n_swa > 0 && "Use llama_kv_cache_unified for non-SWA");

const auto n_kv = kv_self->get_kv_swa()->get_n();

Expand Down
6 changes: 3 additions & 3 deletions src/llama-graph.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -256,10 +256,10 @@ class llm_graph_input_attn_kv_unified : public llm_graph_input_i {

void set_input(const llama_ubatch * ubatch) override;

ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }
ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; }

ggml_tensor * self_kq_mask = nullptr; // F32 [n_kv, n_batch]
ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch]
ggml_tensor * self_kq_mask = nullptr; // F32 [n_kv, n_batch]
ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch]

const llama_hparams & hparams;
const llama_cparams & cparams;
Expand Down
22 changes: 18 additions & 4 deletions src/llama-model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -856,20 +856,34 @@ void llama_model::load_hparams(llama_model_loader & ml) {
// for backward compatibility ; see: https://github.com/ggerganov/llama.cpp/pull/8931
if ((hparams.n_layer == 32 || hparams.n_layer == 40) && hparams.n_ctx_train == 4096) {
// default value for Phi-3-mini-4k-instruct and Phi-3-medium-4k-instruct
LLAMA_LOG_WARN("%s: assuming n_swa = 2047 for Phi-3-mini-4k-instruct and Phi-3-medium-4k-instruct\n", __func__);

hparams.n_swa = 2047;
} else if (hparams.n_layer == 32 && hparams.n_head_kv(0) == 32 && hparams.n_ctx_train == 131072) {
// default value for Phi-3-mini-128k-instruct
// note: this seems incorrect because the window is bigger than the train context?
hparams.n_swa = 262144;
LLAMA_LOG_WARN("%s: assuming n_swa = n_ctx_train for Phi-3-mini-128k-instruct\n", __func__);

hparams.n_swa = hparams.n_ctx_train;
hparams.n_swa_pattern = 1;
} else if (hparams.n_layer == 40 && hparams.n_ctx_train == 131072) {
// default value for Phi-3-medium-128k-instruct
// note: this seems incorrect because the window is equal to the train context?
hparams.n_swa = 131072;
LLAMA_LOG_WARN("%s: assuming n_swa = n_ctx_train for Phi-3-medium-128k-instruct\n", __func__);

hparams.n_swa = hparams.n_ctx_train;
hparams.n_swa_pattern = 1;
}

bool found_swa = ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false);
if (!found_swa && hparams.n_swa == 0) {
throw std::runtime_error("invalid value for sliding_window");
}

if (hparams.n_swa > hparams.n_ctx_train) {
LLAMA_LOG_WARN("%s: unexpected n_swa: %d >= %d, setting to 0\n", __func__, hparams.n_swa, hparams.n_ctx_train);

hparams.n_swa = hparams.n_ctx_train;
hparams.n_swa_pattern = 1;
}
} break;
case LLM_ARCH_PHIMOE:
{
Expand Down