Skip to content

Commit a410713

Browse files
authored
llama : add guard for K/V rotation input when buffer is unallocated (ggml-org#25215)
llm_graph_input_attn_kv::set_input and llm_graph_input_attn_kv_iswa::set_input call set_input_k_rot / set_input_v_rot whenever the rotation tensor pointer is non-null, but the tensor's buffer can be unallocated (NULL) when a graph only stores K/V without attending -- e.g. DFlash speculative decoding's KV-injection pass. set_input_k_rot then calls ggml_backend_buffer_is_host() on a NULL buffer and aborts with GGML_ASSERT(buffer). Guard the four k_rot/v_rot inputs with the same "&& ->buffer" check that the adjacent kq_mask inputs already use in these two functions. When the buffer is unallocated there is no data to upload, so skipping is correct. Fixes ggml-org#25191 Signed-off-by: liminfei-amd <91481003+liminfei-amd@users.noreply.github.com>
1 parent 6658925 commit a410713

1 file changed

Lines changed: 6 additions & 6 deletions

File tree

src/llama-graph.cpp

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -494,11 +494,11 @@ void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) {
494494
mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn);
495495
}
496496

497-
if (self_k_rot) {
497+
if (self_k_rot && self_k_rot->buffer) {
498498
mctx->set_input_k_rot(self_k_rot);
499499
}
500500

501-
if (self_v_rot) {
501+
if (self_v_rot && self_v_rot->buffer) {
502502
mctx->set_input_v_rot(self_v_rot);
503503
}
504504
}
@@ -592,19 +592,19 @@ void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) {
592592
mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn);
593593
}
594594

595-
if (self_k_rot) {
595+
if (self_k_rot && self_k_rot->buffer) {
596596
mctx->get_base()->set_input_k_rot(self_k_rot);
597597
}
598598

599-
if (self_v_rot) {
599+
if (self_v_rot && self_v_rot->buffer) {
600600
mctx->get_base()->set_input_v_rot(self_v_rot);
601601
}
602602

603-
if (self_k_rot_swa) {
603+
if (self_k_rot_swa && self_k_rot_swa->buffer) {
604604
mctx->get_swa()->set_input_k_rot(self_k_rot_swa);
605605
}
606606

607-
if (self_v_rot_swa) {
607+
if (self_v_rot_swa && self_v_rot_swa->buffer) {
608608
mctx->get_swa()->set_input_v_rot(self_v_rot_swa);
609609
}
610610
}

0 commit comments

Comments
 (0)