mtmd: use RAII for setting and resetting non-causal attention (#25723)

* mtmd: use RAII for setting and resetting non-causal attention

* mtmd: drop dependency on <optional>

* mtmd: shorten class and variable names
This commit is contained in:
Niklas Wenzel
2026-07-22 18:10:03 +02:00
committed by GitHub
parent 61328e6a91
commit e0833bf686

View File

@@ -238,6 +238,29 @@ struct decode_embd_batch {
}
};
// Helper class to set non-causal attention via RAII
class scope_non_causal {
public:
scope_non_causal(llama_context * context, bool enabled) : context_(context), enabled_(enabled) {
if (enabled_) {
// TODO @ngxson : need to make sure only one image is processed at a time, and n_ubatch must be enough to hold the image
llama_set_causal_attn(context_, false);
}
}
~scope_non_causal() {
if (enabled_) {
llama_set_causal_attn(context_, true);
}
}
scope_non_causal(const scope_non_causal &) = delete;
scope_non_causal & operator=(const scope_non_causal &) = delete;
private:
llama_context * context_;
bool enabled_;
};
// Helper function for decoding an image whose embeddings have already been calculated
int32_t mtmd_helper_decode_image_chunk(
mtmd_context * ctx,
@@ -288,10 +311,7 @@ int32_t mtmd_helper_decode_image_chunk(
}
const bool use_non_causal = mtmd_decode_use_non_causal(ctx, chunk);
if (use_non_causal) {
llama_set_causal_attn(lctx, false);
// TODO @ngxson : need to make sure only one image is processed at a time, and n_ubatch must be enough to hold the image
}
const scope_non_causal non_causal(lctx, use_non_causal);
while (i_batch < n_img_batches) { // split into batches
int pos_offset = i_batch*n_batch;
@@ -304,9 +324,6 @@ int32_t mtmd_helper_decode_image_chunk(
int32_t ret = llama_decode(lctx, batch_embd_view);
if (ret != 0) {
LOG_ERR("failed to decode %s\n", name);
if (use_non_causal) {
llama_set_causal_attn(lctx, true);
}
return ret;
}
@@ -314,9 +331,6 @@ int32_t mtmd_helper_decode_image_chunk(
ret = callback(batch_embd_view, user_data);
if (ret != 0) {
LOG_ERR("post-decode callback failed\n");
if (use_non_causal) {
llama_set_causal_attn(lctx, true);
}
return ret;
}
}
@@ -329,9 +343,6 @@ int32_t mtmd_helper_decode_image_chunk(
n_past += mtmd_input_chunk_get_n_pos(chunk);
*new_n_past = n_past;
if (use_non_causal) {
llama_set_causal_attn(lctx, true);
}
return 0;
}