Skip to content

Commit b02a5a8

Browse files
ggerganovggerganov
committed
kv-cache : simplify SWA logic
ggml-ci
1 parent 5b4e000 commit b02a5a8

File tree

5 files changed

+96
-54
lines changed

5 files changed

+96
-54
lines changed

‎src/llama-graph.cpp‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -362,17 +362,17 @@ void llm_graph_input_attn_no_cache::set_input(const llama_ubatch * ubatch) {
362362

363363
void llm_graph_input_attn_kv_unified::set_input(const llama_ubatch * ubatch) {
364364
if (self_kq_mask) {
365-
kv_self->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn, false);
365+
kv_self->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
366366
}
367367
}
368368

369369
void llm_graph_input_attn_kv_unified_iswa::set_input(const llama_ubatch * ubatch) {
370370
if (self_kq_mask) {
371-
kv_self->get_kv_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn, false);
371+
kv_self->get_kv_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
372372
}
373373

374374
if (self_kq_mask_swa) {
375-
kv_self->get_kv_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn, true);
375+
kv_self->get_kv_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn);
376376
}
377377
}
378378

‎src/llama-hparams.h‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,9 @@ enum llama_expert_gating_func_type {
1515
};
1616

1717
enum llama_swa_type {
18-
LLAMA_SWA_TYPE_STANDARD = 0,
19-
LLAMA_SWA_TYPE_CHUNKED = 1,
18+
LLAMA_SWA_TYPE_NONE = 0,
19+
LLAMA_SWA_TYPE_STANDARD = 1,
20+
LLAMA_SWA_TYPE_CHUNKED = 2,
2021
};
2122

2223
struct llama_hparams_posnet {
@@ -100,7 +101,7 @@ struct llama_hparams {
100101
std::array<int, 4> rope_sections;
101102

102103
// Sliding Window Attention (SWA)
103-
llama_swa_type swa_type = LLAMA_SWA_TYPE_STANDARD;
104+
llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE;
104105

105106
uint32_t n_swa = 0; // the size of the sliding window (0 - no SWA)
106107
uint32_t n_swa_pattern = 1; // by default, all layers use non-sliding-window attention

‎src/llama-kv-cache.cpp‎

Lines changed: 57 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,9 @@ llama_kv_cache_unified::llama_kv_cache_unified(
3030
bool v_trans,
3131
bool offload,
3232
uint32_t kv_size,
33-
uint32_t padding) : model(model), hparams(model.hparams), v_trans(v_trans), padding(padding) {
33+
uint32_t padding,
34+
uint32_t n_swa,
35+
llama_swa_type swa_type) : model(model), hparams(model.hparams), v_trans(v_trans), padding(padding), n_swa(n_swa), swa_type(swa_type) {
3436
GGML_ASSERT(kv_size % padding == 0 && "kv_size must be a multiple of padding");
3537

3638
this->type_k = type_k;
@@ -594,8 +596,8 @@ ggml_tensor * llama_kv_cache_unified::get_v(ggml_context * ctx, int32_t il) cons
594596
// note: v->nb[1] > v->nb[2]
595597
return ggml_view_3d(ctx, v,
596598
n, hparams.n_head_kv(il), hparams.n_embd_head_v,
597-
ggml_element_size(v)*v->ne[1]*hparams.n_embd_head_v, // v->nb[1]
598-
ggml_element_size(v)*v->ne[1], // v->nb[2]
599+
ggml_row_size(v->type, v->ne[1]*hparams.n_embd_head_v), // v->nb[1]
600+
ggml_row_size(v->type, v->ne[1]), // v->nb[2]
599601
0);
600602
}
601603

@@ -640,7 +642,7 @@ ggml_tensor * llama_kv_cache_unified::cpy_v(ggml_context * ctx, ggml_tensor * v_
640642
return ggml_cpy(ctx, v_cur, v_view);
641643
}
642644

643-
void llama_kv_cache_unified::set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn, bool swa) const {
645+
void llama_kv_cache_unified::set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const {
644646
const int64_t n_tokens = ubatch->n_tokens;
645647
const int64_t n_seq_tokens = ubatch->n_seq_tokens;
646648
const int64_t n_seqs = ubatch->n_seqs;
@@ -667,41 +669,28 @@ void llama_kv_cache_unified::set_input_kq_mask(ggml_tensor * dst, const llama_ub
667669
const llama_seq_id seq_id = ubatch->seq_id[s][0];
668670

669671
for (int j = 0; j < n_seq_tokens; ++j) {
670-
const llama_pos pos = ubatch->pos[s*n_seq_tokens + j];
672+
const llama_pos p1 = ubatch->pos[s*n_seq_tokens + j];
671673

672674
for (int i = 0; i < n_kv; ++i) {
673-
float f;
674-
// mask the token if:
675-
if (!cells[i].has_seq_id(seq_id) // not the correct sequence
676-
|| (causal_attn && cells[i].pos > pos) // for causal, mask future tokens
677-
) {
678-
f = -INFINITY;
679-
} else {
680-
if (hparams.use_alibi) {
681-
f = -std::abs(cells[i].pos - pos);
682-
} else {
683-
f = 0.0f;
684-
}
685-
}
675+
const llama_pos p0 = cells[i].pos;
676+
677+
bool masked = false;
678+
679+
// mask the token if not the same sequence
680+
masked = masked || (!cells[i].has_seq_id(seq_id));
681+
682+
// mask future tokens
683+
masked = masked || (causal_attn && p0 > p1);
686684

687-
if (swa) {
688-
// may need to cut off old tokens for sliding window
689-
switch (hparams.swa_type) {
690-
case LLAMA_SWA_TYPE_STANDARD:
691-
{
692-
if (pos - cells[i].pos >= (int32_t) hparams.n_swa) {
693-
f = -INFINITY;
694-
}
695-
} break;
696-
case LLAMA_SWA_TYPE_CHUNKED:
697-
{
698-
const llama_pos pos_chunk_start = (pos / hparams.n_swa) * hparams.n_swa;
699-
700-
if (cells[i].pos < pos_chunk_start) {
701-
f = -INFINITY;
702-
}
703-
} break;
704-
}
685+
// apply SWA if any
686+
masked = masked || (is_masked_swa(p0, p1));
687+
688+
float f = 0.0f;
689+
690+
if (masked) {
691+
f = -INFINITY;
692+
} else if (hparams.use_alibi) {
693+
f = -std::abs(p0 - p1);
705694
}
706695

707696
data[h*(n_kv*n_tokens) + s*(n_kv*n_seq_tokens) + j*n_kv + i] = f;
@@ -1191,6 +1180,30 @@ uint32_t llama_kv_cache_unified::cell_max() const {
11911180
return 0;
11921181
}
11931182

1183+
bool llama_kv_cache_unified::is_masked_swa(llama_pos p0, llama_pos p1) const {
1184+
switch (swa_type) {
1185+
case LLAMA_SWA_TYPE_NONE:
1186+
{
1187+
} break;
1188+
case LLAMA_SWA_TYPE_STANDARD:
1189+
{
1190+
if (p1 - p0 >= (int32_t) n_swa) {
1191+
return true;
1192+
}
1193+
} break;
1194+
case LLAMA_SWA_TYPE_CHUNKED:
1195+
{
1196+
const llama_pos pos_chunk_start = (p1 / n_swa) * n_swa;
1197+
1198+
if (p0 < pos_chunk_start) {
1199+
return true;
1200+
}
1201+
} break;
1202+
}
1203+
1204+
return false;
1205+
}
1206+
11941207
void llama_kv_cache_unified::state_write(llama_io_write_i & io, llama_seq_id seq_id) const {
11951208
std::vector<std::pair<uint32_t, uint32_t>> cell_ranges; // ranges, from inclusive, to exclusive
11961209
uint32_t cell_count = 0;
@@ -1586,11 +1599,17 @@ llama_kv_cache_unified_iswa::llama_kv_cache_unified_iswa(
15861599

15871600
LLAMA_LOG_INFO("%s: creating non-SWA KV cache, size = %u cells\n", __func__, kv_size_base);
15881601

1589-
kv_base = std::make_unique<llama_kv_cache_unified>(model, std::move(filter_base), type_k, type_v, v_trans, offload, kv_size_base, padding);
1602+
kv_base = std::make_unique<llama_kv_cache_unified>(
1603+
model, std::move(filter_base), type_k, type_v,
1604+
v_trans, offload, kv_size_base, padding,
1605+
0, LLAMA_SWA_TYPE_NONE);
15901606

15911607
LLAMA_LOG_INFO("%s: creating SWA KV cache, size = %u cells\n", __func__, kv_size_swa);
15921608

1593-
kv_swa = std::make_unique<llama_kv_cache_unified>(model, std::move(filter_swa), type_k, type_v, v_trans, offload, kv_size_swa, padding);
1609+
kv_swa = std::make_unique<llama_kv_cache_unified>(
1610+
model, std::move(filter_swa), type_k, type_v,
1611+
v_trans, offload, kv_size_swa, padding,
1612+
hparams.n_swa, hparams.swa_type);
15941613
}
15951614

15961615
void llama_kv_cache_unified_iswa::clear() {

‎src/llama-kv-cache.h‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,9 @@ class llama_kv_cache_unified : public llama_kv_cache {
102102
bool v_trans,
103103
bool offload,
104104
uint32_t kv_size,
105-
uint32_t padding);
105+
uint32_t padding,
106+
uint32_t n_swa,
107+
llama_swa_type swa_type);
106108

107109
~llama_kv_cache_unified() = default;
108110

@@ -169,7 +171,7 @@ class llama_kv_cache_unified : public llama_kv_cache {
169171
ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, int32_t il) const;
170172
ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, int32_t il) const;
171173

172-
void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn, bool swa) const;
174+
void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const;
173175
void set_input_k_shift (ggml_tensor * dst) const;
174176
void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const;
175177

@@ -223,6 +225,11 @@ class llama_kv_cache_unified : public llama_kv_cache {
223225
ggml_type type_k = GGML_TYPE_F16;
224226
ggml_type type_v = GGML_TYPE_F16;
225227

228+
// SWA
229+
uint32_t n_swa = 0;
230+
231+
llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE;
232+
226233
std::vector<ggml_context_ptr> ctxs;
227234
std::vector<ggml_backend_buffer_ptr> bufs;
228235

@@ -264,6 +271,8 @@ class llama_kv_cache_unified : public llama_kv_cache {
264271
size_t size_k_bytes() const;
265272
size_t size_v_bytes() const;
266273

274+
bool is_masked_swa(llama_pos p0, llama_pos p1) const;
275+
267276
ggml_tensor * build_rope_shift(
268277
const llama_cparams & cparams,
269278
ggml_context * ctx,

‎src/llama-model.cpp‎

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -572,7 +572,7 @@ void llama_model::load_hparams(llama_model_loader & ml) {
572572
ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp);
573573
ml.get_key(LLM_KV_INTERLEAVE_MOE_LAYER_STEP, hparams.n_moe_layer_step);
574574

575-
hparams.swa_type = (llama_swa_type) LLAMA_SWA_TYPE_CHUNKED;
575+
hparams.swa_type = LLAMA_SWA_TYPE_CHUNKED;
576576
hparams.n_swa = 8192; // should this be a gguf kv? currently it's the same for Scout and Maverick
577577
hparams.n_swa_pattern = 4; // pattern: 3 chunked - 1 full
578578

@@ -858,18 +858,24 @@ void llama_model::load_hparams(llama_model_loader & ml) {
858858
// default value for Phi-3-mini-4k-instruct and Phi-3-medium-4k-instruct
859859
LLAMA_LOG_WARN("%s: assuming n_swa = 2047 for Phi-3-mini-4k-instruct and Phi-3-medium-4k-instruct\n", __func__);
860860

861+
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
862+
861863
hparams.n_swa = 2047;
862864
} else if (hparams.n_layer == 32 && hparams.n_head_kv(0) == 32 && hparams.n_ctx_train == 131072) {
863865
// default value for Phi-3-mini-128k-instruct
864-
LLAMA_LOG_WARN("%s: assuming n_swa = n_ctx_train for Phi-3-mini-128k-instruct\n", __func__);
866+
LLAMA_LOG_WARN("%s: assuming no SWA for Phi-3-mini-128k-instruct\n", __func__);
867+
868+
hparams.swa_type = LLAMA_SWA_TYPE_NONE;
865869

866-
hparams.n_swa = hparams.n_ctx_train;
870+
hparams.n_swa = hparams.n_ctx_train;
867871
hparams.n_swa_pattern = 1;
868872
} else if (hparams.n_layer == 40 && hparams.n_ctx_train == 131072) {
869873
// default value for Phi-3-medium-128k-instruct
870-
LLAMA_LOG_WARN("%s: assuming n_swa = n_ctx_train for Phi-3-medium-128k-instruct\n", __func__);
874+
LLAMA_LOG_WARN("%s: assuming no SWA for Phi-3-medium-128k-instruct\n", __func__);
875+
876+
hparams.swa_type = LLAMA_SWA_TYPE_NONE;
871877

872-
hparams.n_swa = hparams.n_ctx_train;
878+
hparams.n_swa = hparams.n_ctx_train;
873879
hparams.n_swa_pattern = 1;
874880
}
875881

@@ -879,9 +885,11 @@ void llama_model::load_hparams(llama_model_loader & ml) {
879885
}
880886

881887
if (hparams.n_swa > hparams.n_ctx_train) {
882-
LLAMA_LOG_WARN("%s: unexpected n_swa: %d >= %d, setting to 0\n", __func__, hparams.n_swa, hparams.n_ctx_train);
888+
LLAMA_LOG_WARN("%s: unexpected n_swa: %d >= %d, disabling SWA\n", __func__, hparams.n_swa, hparams.n_ctx_train);
883889

884-
hparams.n_swa = hparams.n_ctx_train;
890+
hparams.swa_type = LLAMA_SWA_TYPE_NONE;
891+
892+
hparams.n_swa = hparams.n_ctx_train;
885893
hparams.n_swa_pattern = 1;
886894
}
887895
} break;
@@ -952,6 +960,7 @@ void llama_model::load_hparams(llama_model_loader & ml) {
952960
} break;
953961
case LLM_ARCH_GEMMA2:
954962
{
963+
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
955964
hparams.n_swa = 4096; // default value of gemma 2
956965
hparams.n_swa_pattern = 2;
957966
hparams.attn_soft_cap = true;
@@ -970,6 +979,7 @@ void llama_model::load_hparams(llama_model_loader & ml) {
970979
} break;
971980
case LLM_ARCH_GEMMA3:
972981
{
982+
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
973983
hparams.n_swa_pattern = 6;
974984

975985
hparams.rope_freq_base_train_swa = 10000.0f;
@@ -1054,6 +1064,7 @@ void llama_model::load_hparams(llama_model_loader & ml) {
10541064
} break;
10551065
case LLM_ARCH_COHERE2:
10561066
{
1067+
hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
10571068
hparams.n_swa_pattern = 4;
10581069

10591070
ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
@@ -13228,7 +13239,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
1322813239
!cparams.flash_attn,
1322913240
cparams.offload_kqv,
1323013241
cparams.n_ctx,
13231-
padding);
13242+
padding,
13243+
hparams.n_swa,
13244+
hparams.swa_type);
1323213245
}
1323313246
}
1323413247
}

0 commit comments

Comments
 (0)