diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 4c86e43c1..dc41d6690 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -494,11 +494,11 @@ void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) { mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); } - if (self_k_rot) { + if (self_k_rot && self_k_rot->buffer) { mctx->set_input_k_rot(self_k_rot); } - if (self_v_rot) { + if (self_v_rot && self_v_rot->buffer) { mctx->set_input_v_rot(self_v_rot); } } @@ -592,19 +592,19 @@ void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) { mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn); } - if (self_k_rot) { + if (self_k_rot && self_k_rot->buffer) { mctx->get_base()->set_input_k_rot(self_k_rot); } - if (self_v_rot) { + if (self_v_rot && self_v_rot->buffer) { mctx->get_base()->set_input_v_rot(self_v_rot); } - if (self_k_rot_swa) { + if (self_k_rot_swa && self_k_rot_swa->buffer) { mctx->get_swa()->set_input_k_rot(self_k_rot_swa); } - if (self_v_rot_swa) { + if (self_v_rot_swa && self_v_rot_swa->buffer) { mctx->get_swa()->set_input_v_rot(self_v_rot_swa); } }