@@ -4889,6 +4889,7 @@ static void llm_load_hparams(
4889
4889
} break;
4890
4890
case LLM_ARCH_PHI3:
4891
4891
{
4892
+ ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
4892
4893
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
4893
4894
4894
4895
switch (hparams.n_layer) {
@@ -10748,7 +10749,7 @@ struct llm_build_context {
10748
10749
struct ggml_tensor * inp_pos = build_inp_pos();
10749
10750
10750
10751
// KQ_mask (mask for 1 head, it will be broadcasted to all heads)
10751
- struct ggml_tensor * KQ_mask = build_inp_KQ_mask ();
10752
+ struct ggml_tensor * KQ_mask_swa = build_inp_KQ_mask_swa ();
10752
10753
10753
10754
for (int il = 0; il < n_layer; ++il) {
10754
10755
auto residual = inpL;
@@ -10806,7 +10807,7 @@ struct llm_build_context {
10806
10807
10807
10808
cur = llm_build_kv(ctx0, lctx, kv_self, gf,
10808
10809
model.layers[il].wo, model.layers[il].bo,
10809
- Kcur, Vcur, Qcur, KQ_mask , n_tokens, kv_head, n_kv, 1.0f, cb, il);
10810
+ Kcur, Vcur, Qcur, KQ_mask_swa , n_tokens, kv_head, n_kv, 1.0f, cb, il);
10810
10811
}
10811
10812
10812
10813
if (il == n_layer - 1) {
@@ -14013,18 +14014,23 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) {
14013
14014
"causal attention is not supported by this model"
14014
14015
);
14015
14016
14016
- if (lctx.inp_KQ_mask) {
14017
+ if (lctx.inp_KQ_mask || lctx.inp_KQ_mask_swa ) {
14017
14018
// NOTE: hparams.causal_attn indicates the model is capable of generation and uses the kv cache.
14018
14019
if (cparams.causal_attn && !lctx.is_encoding) {
14019
14020
const int64_t n_kv = kv_self.n;
14020
14021
const int64_t n_tokens = batch.n_tokens;
14021
14022
14022
- GGML_ASSERT(ggml_backend_buffer_is_host(lctx.inp_KQ_mask->buffer));
14023
14023
14024
- float * data = (float *) lctx.inp_KQ_mask->data ;
14024
+ float * data = nullptr ;
14025
14025
float * data_swa = nullptr;
14026
14026
14027
+ if (lctx.inp_KQ_mask) {
14028
+ GGML_ASSERT(ggml_backend_buffer_is_host(lctx.inp_KQ_mask->buffer));
14029
+ data = (float *) lctx.inp_KQ_mask->data;
14030
+ }
14031
+
14027
14032
if (lctx.inp_KQ_mask_swa) {
14033
+ GGML_ASSERT(ggml_backend_buffer_is_host(lctx.inp_KQ_mask_swa->buffer));
14028
14034
data_swa = (float *) lctx.inp_KQ_mask_swa->data;
14029
14035
}
14030
14036
@@ -14047,7 +14053,10 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) {
14047
14053
f = 0.0f;
14048
14054
}
14049
14055
}
14050
- data[h*(n_kv*n_tokens) + j*n_kv + i] = f;
14056
+
14057
+ if (data) {
14058
+ data[h*(n_kv*n_tokens) + j*n_kv + i] = f;
14059
+ }
14051
14060
14052
14061
// may need to cut off old tokens for sliding window
14053
14062
if (data_swa) {
@@ -14059,9 +14068,19 @@ static void llama_set_inputs(llama_context & lctx, const llama_batch & batch) {
14059
14068
}
14060
14069
}
14061
14070
14062
- for (int i = n_tokens; i < GGML_PAD(n_tokens, GGML_KQ_MASK_PAD); ++i) {
14063
- for (int j = 0; j < n_kv; ++j) {
14064
- data[h*(n_kv*n_tokens) + i*n_kv + j] = -INFINITY;
14071
+ if (data) {
14072
+ for (int i = n_tokens; i < GGML_PAD(n_tokens, GGML_KQ_MASK_PAD); ++i) {
14073
+ for (int j = 0; j < n_kv; ++j) {
14074
+ data[h*(n_kv*n_tokens) + i*n_kv + j] = -INFINITY;
14075
+ }
14076
+ }
14077
+ }
14078
+
14079
+ if (data_swa) {
14080
+ for (int i = n_tokens; i < GGML_PAD(n_tokens, GGML_KQ_MASK_PAD); ++i) {
14081
+ for (int j = 0; j < n_kv; ++j) {
14082
+ data_swa[h*(n_kv*n_tokens) + i*n_kv + j] = -INFINITY;
14083
+ }
14065
14084
}
14066
14085
}
14067
14086
}
0 commit comments