Repository: ollama/ollama
Pull request: fix multiple image inputs

* zero out inp_embd when setting cross attention states
* set cross attention if _any_ inputs contains image tokens, not just the new inputs

Changed files:
filename: llama/llama.cpp
patch: @@ -2699,7 +2699,7 @@ struct llama_hparams {
         GGML_ABORT("fatal error");
     }
 
-    bool cross_attention_layer(uint32_t il) const {
+    bool cross_attention_layers(uint32_t il) const {
         return std::find(cross_attn_layers.begin(), cross_attn_layers.end(), il) != cross_attn_layers.end();
     }
 };
@@ -2731,6 +2731,7 @@ struct llama_cparams {
     bool offload_kqv;
     bool flash_attn;
     bool no_perf;
+    bool cross_attn = false;
 
     enum llama_pooling_type pooling_type;
 
@@ -3546,7 +3547,6 @@ struct llama_context {
 
     // TODO (jmorganca): this should most likely be passed in as part of a batch
     // and not set on the context for all batches.
-    bool cross_attn_state = false;
     struct ggml_tensor * inp_cross_attn_state; // F32 [4, n_embd, 1061]
 };
 
@@ -3783,7 +3783,7 @@ static bool llama_kv_cache_init(
 
     for (int i = 0; i < (int) n_layer; i++) {
         // for cross attention layers
-        if (model.arch == LLM_ARCH_MLLAMA && hparams.cross_attention_layer(i)) {
+        if (model.arch == LLM_ARCH_MLLAMA && hparams.cross_attention_layers(i)) {
             struct ggml_context * ctx = offload ? ctx_map.at(model.buft_layer[i].buft) : cache.ctxs.front();
             ggml_tensor * k = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hparams.n_embd_head_k, 6404, hparams.n_head_kv(i));
             ggml_tensor * v = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hparams.n_embd_head_v, 6404, hparams.n_head_kv(i));
@@ -7390,7 +7390,7 @@ static bool llm_load_tensors(
 
                         auto & layer = model.layers[i];
 
-                        if (hparams.cross_attention_layer(i)) {
+                        if (hparams.cross_attention_layers(i)) {
                             layer.cross_attn_k_norm = ml.create_tensor(ctx_split, tn(LLM_TENSOR_CROSS_ATTN_K_NORM,   "weight", i), {128});
                             layer.cross_attn_k_proj = ml.create_tensor(ctx_split, tn(LLM_TENSOR_CROSS_ATTN_K_PROJ,   "weight", i), {n_embd, 1024});
                             layer.cross_attn_o_proj = ml.create_tensor(ctx_split, tn(LLM_TENSOR_CROSS_ATTN_O_PROJ,   "weight", i), {n_embd, n_embd});
@@ -9369,11 +9369,10 @@ static struct ggml_tensor * llm_build_inp_cross_attn_state(
          const llm_build_cb & cb) {
     const int64_t n_embd = hparams.n_embd;
 
-    struct ggml_tensor * inpCAS;
-    lctx.inp_cross_attn_state = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1601, 4);
-    cb(lctx.inp_cross_attn_state, "inp_cross_attn_state", -1);
-    ggml_set_input(lctx.inp_cross_attn_state);
-    inpCAS = lctx.inp_cross_attn_state;
+    struct ggml_tensor * inpCAS = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1601, 4);
+    cb(inpCAS, "inp_cross_attn_state", -1);
+    ggml_set_input(inpCAS);
+    lctx.inp_cross_attn_state = inpCAS;
 
     return inpCAS;
 }
@@ -10980,8 +10979,8 @@ struct llm_build_context {
                     LLM_NORM_RMS, cb, il);
             cb(cur, "attn_norm", il);
 
-            if (hparams.cross_attention_layer(il)) {
-                if (!batch.embd && !lctx.cross_attn_state) {
+            if (hparams.cross_attention_layers(il)) {
+                if (!batch.embd && !cparams.cross_attn) {
                     continue;
                 }
 
@@ -10992,44 +10991,28 @@ struct llm_build_context {
                 Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
                 cb(Qcur, "Qcur", il);
 
-                Qcur = ggml_permute(ctx0, Qcur, 0, 2, 1, 3);
-                cb(Qcur, "Qcur", il);
-
-                // TODO: is this required?
-                Qcur = ggml_cont(ctx0, Qcur);
+                Qcur = ggml_cont(ctx0, ggml_permute(ctx0, Qcur, 0, 2, 1, 3));
                 cb(Qcur, "Qcur", il);
 
                 Qcur = llm_build_norm(ctx0, Qcur, hparams, model.layers[il].cross_attn_q_norm, NULL, LLM_NORM_RMS, cb, il);
                 cb(Qcur, "Qcur", il);
 
-                struct ggml_tensor * Kcur;
+                struct ggml_tensor * Kcur, * Vcur;
                 if (batch.embd) {
-                    lctx.cross_attn_state = true;
-
                     Kcur = ggml_mul_mat(ctx0, model.layers[il].cross_attn_k_proj, inpCAS);
                     cb(Kcur, "Kcur", il);
 
                     Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, 6404);
                     cb(Kcur, "Kcur", il);
 
-                    Kcur = ggml_permute(ctx0, Kcur, 0, 2, 1, 3);
-                    cb(Kcur, "Kcur", il);
-
-                    // TODO: is this required?
-                    Kcur = ggml_cont(ctx0, Kcur);
+                    Kcur = ggml_cont(ctx0, ggml_permute(ctx0, Kcur, 0, 2, 1, 3));
                     cb(Kcur, "Kcur", il);
 
                     Kcur = llm_build_norm(ctx0, Kcur, hparams, model.layers[il].cross_attn_k_norm, NULL, LLM_NORM_RMS, cb, il);
                     cb(Kcur, "Kcur", il);
 
                     ggml_build_forward_expand(gf, ggml_cpy(ctx0, Kcur, kv_self.k_l[il]));
-                } else {
-                    Kcur = ggml_view_tensor(ctx0, kv_self.k_l[il]);
-                    cb(Kcur, "Kcur (view)", il);
-                }
 
-                struct ggml_tensor * Vcur;
-                if (batch.embd) {
                     Vcur = ggml_mul_mat(ctx0, model.layers[il].cross_attn_v_proj, inpCAS);
                     cb(Vcur, "Vcur", il);
 
@@ -11041,18 +11024,18 @@ struct llm_build_context {
 
                     ggml_build_forward_expand(gf, ggml_cpy(ctx0, Vcur, kv_self.v_l[il]));
                 } else {
+                    Kcur = ggml_view_tensor(ctx0, kv_self.k_l[il]);
+                    cb(Kcur, "Kcur (view)", il);
+
                     Vcur = ggml_view_tensor(ctx0, kv_self.v_l[il]);
                     cb(Vcur, "Vcur (view)", il);
                 }
 
                 struct ggml_tensor * kq = ggml_mul_mat(ctx0, Kcur, Qcur);
                 cb(kq, "kq", il);
 
-                kq = ggml_scale_inplace(ctx0, kq, 1.0f/sqrtf(float(n_embd_head)));
-                cb(kq, "kq_scaled", il);
-
                 // TODO: apply causal masks
-                struct ggml_tensor * kq_soft_max = ggml_soft_max_inplace(ctx0, kq);
+                struct ggml_tensor * kq_soft_max = ggml_soft_max_ext(ctx0, kq, nullptr, 1.f/sqrtf(float(n_embd_head)), hparams.f_max_alibi_bias);
                 cb(kq_soft_max, "kq_soft_max", il);
 
                 Vcur = ggml_cont(ctx0, ggml_transpose(ctx0, Vcur));
@@ -11142,8 +11125,8 @@ struct llm_build_context {
                 cb(Kcur, "Kcur", il);
 
                 cur = llm_build_kv(ctx0, lctx, kv_self, gf,
-                        model.layers[il].wo, model.layers[il].bo,
-                        Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, 1.0f/sqrtf(float(n_embd_head)), cb, il);
+                    model.layers[il].wo, model.layers[il].bo,
+                    Kcur, Vcur, Qcur, KQ_mask, n_tokens, kv_head, n_kv, 1.0f/sqrtf(float(n_embd_head)), cb, il);
 
 
                 if (il == n_layer - 1) {
@@ -17200,9 +17183,13 @@ static void llama_set_inputs(llama_context & lctx, const llama_ubatch & batch) {
     }
 
     if (batch.embd) {
-        if (lctx.inp_cross_attn_state &&
-            lctx.inp_cross_attn_state->buffer) {
-            ggml_backend_tensor_set(lctx.inp_cross_attn_state, batch.embd, 0, batch.n_tokens*batch.n_embd*ggml_element_size(lctx.inp_cross_attn_state));
+        if (lctx.inp_cross_attn_state && lctx.inp_cross_attn_state->buffer) {
+            ggml_backend_tensor_set(lctx.inp_cross_attn_state, batch.embd, 0, ggml_nbytes(lctx.inp_cross_attn_state));
+            // zero out inp_embd since it's not used
+            float * inp_embd_data = (float *)lctx.inp_embd->data;
+            for (int i = 0; i < ggml_nelements(lctx.inp_embd); ++i) {
+                inp_embd_data[i] = 0.0f;
+            }
         } else {
             ggml_backend_tensor_set(lctx.inp_embd, batch.embd, 0, batch.n_tokens*batch.n_embd*ggml_element_size(lctx.inp_embd));
         }
@@ -19787,10 +19774,6 @@ struct llama_model * llama_load_model_from_file(
     return model;
 }
 
-void llama_set_cross_attn_state(struct llama_context * ctx, bool cross_attn_state) {
-    ctx->cross_attn_state = cross_attn_state;
-}
-
 void llama_free_model(struct llama_model * model) {
     delete model;
 }
@@ -21678,6 +21661,10 @@ void llama_set_causal_attn(struct llama_context * ctx, bool causal_attn) {
     ctx->cparams.causal_attn = causal_attn;
 }
 
+void llama_set_cross_attention(struct llama_context * ctx, bool cross_attention) {
+    ctx->cparams.cross_attn = cross_attention;
+}
+
 struct llama_batch llama_batch_get_one(
              llama_token * tokens,
                  int32_t   n_tokens,

filename: llama/llama.go
patch: @@ -521,8 +521,8 @@ func (m *MllamaContext) EmbedSize(llamaContext *Context) int {
 }
 
 // This really needs to be set on a batch instead
-func (c *Context) MllamaSetCrossAttn(state bool) {
-	C.llama_set_cross_attn_state(c.c, (C.bool)(state))
+func (c *Context) SetCrossAttention(state bool) {
+	C.llama_set_cross_attention(c.c, C.bool(state))
 }
 
 // sampling

filename: llama/llama.h
patch: @@ -452,7 +452,7 @@ extern "C" {
 
     // TODO (jmorganca): this should most likely be passed in as part of a batch
     // and not set on the context for all batches.
-    LLAMA_API void llama_set_cross_attn_state(struct llama_context * ctx, bool cross_attn_state);
+    LLAMA_API void llama_set_cross_attention(struct llama_context * ctx, bool cross_attn_state);
 
     // Frees all allocated memory
     LLAMA_API void llama_free(struct llama_context * ctx);

filename: llama/runner/runner.go
patch: @@ -296,8 +296,8 @@ func flushPending(seq *Sequence) bool {
 
 func (s *Server) removeSequence(seqIndex int, reason string) {
 	seq := s.seqs[seqIndex]
+	s.lc.SetCrossAttention(false)
 
-	s.lc.MllamaSetCrossAttn(false)
 	flushPending(seq)
 	seq.doneReason = reason
 	close(seq.responses)
@@ -615,18 +615,19 @@ func (s *Server) completion(w http.ResponseWriter, r *http.Request) {
 	s.mu.Lock()
 	for i, sq := range s.seqs {
 		if sq == nil {
+			for _, input := range seq.inputs {
+				if input.embed != nil {
+					s.lc.SetCrossAttention(true)
+					break
+				}
+			}
+
 			seq.cache, seq.inputs, seq.numPast, err = s.cache.LoadCacheSlot(seq.inputs, req.CachePrompt)
 			if err != nil {
 				s.mu.Unlock()
 				http.Error(w, fmt.Sprintf("Failed to load cache: %v", err), http.StatusInternalServerError)
 				return
 			}
-			for _, input := range seq.cache.Inputs {
-				if input.embed != nil {
-					s.lc.MllamaSetCrossAttn(true)
-					break
-				}
-			}
 
 			s.seqs[i] = seq
 			s.cond.Signal()