llama : fix long/short rope factor selection by sequence length

Author: o7si

src/llama-batch.cpp

@@ -9,6 +9,16 @@
 #include <algorithm>
 #include <sstream>
 
+int32_t llama_ubatch::pos_max() const {
+    int32_t res = 0;
+    if (pos) {
+        for (uint32_t i = 0; i < n_tokens; ++i) {
+            res = std::max(res, (int32_t) pos[i]);
+        }
+    }
+    return res;
+}
+
 llama_batch_allocr::llama_batch_allocr(uint32_t n_pos_per_embd) : n_pos_per_embd(n_pos_per_embd) {
     const char * LLAMA_BATCH_DEBUG = getenv("LLAMA_BATCH_DEBUG");
     debug = LLAMA_BATCH_DEBUG ? atoi(LLAMA_BATCH_DEBUG) : 0;

src/llama-batch.h

@@ -27,6 +27,9 @@ struct llama_ubatch {
         return n_pos >= 3;
     }
 
+    // largest token position in the ubatch (max over all sequences)
+    int32_t pos_max() const;
+
     uint32_t b_equal_seqs; // note: this is a boolean, but we use an int32_t for alignment
                            //       otherwise address sanitizer complains
     // TODO: whole_seqs for embeddings?

src/llama-context.cpp

@@ -2400,6 +2400,7 @@ llm_graph_params llama_context::graph_params(
         /*.hparams     =*/ model.hparams,
         /*.cparams     =*/ cparams,
         /*.ubatch      =*/ ubatch,
+        /*.pos_max     =*/ ubatch.pos_max(),
         /*.gtype       =*/ gtype,
         /*.sched       =*/ sched.get(),
         /*.backend_cpu =*/ backend_cpu,

src/llama-graph.cpp

@@ -1054,6 +1054,7 @@ llm_graph_context::llm_graph_context(const llm_graph_params & params) :
     n_tokens         (ubatch.n_tokens),
     n_outputs        (params.n_outputs),
     n_ctx_orig       (cparams.n_ctx_orig_yarn),
+    pos_max          (params.pos_max),
     pooling_type     (cparams.pooling_type),
     rope_type        (hparams.rope_type),
     sched            (params.sched),

src/llama-graph.h

@@ -593,6 +593,8 @@ struct llm_graph_params {
 
     llama_ubatch ubatch; // note: intentionally make a copy
 
+    int32_t pos_max;
+
     llm_graph_type gtype;
 
     ggml_backend_sched_t sched;
@@ -660,6 +662,11 @@ struct llm_graph_params {
             return false;
         }
 
+        if ((pos_max       >= (int32_t) hparams.n_ctx_orig_yarn) !=
+            (other.pos_max >= (int32_t) hparams.n_ctx_orig_yarn)) {
+            return false;
+        }
+
         if (n_outputs != other.n_outputs) {
             return false;
         }
@@ -813,6 +820,8 @@ struct llm_graph_context {
     const int64_t n_outputs;
     const int32_t n_ctx_orig; // yarn
 
+    const int32_t pos_max;
+
     const enum llama_pooling_type pooling_type;
     const enum llama_rope_type    rope_type;
 

src/llama-kv-cache.cpp

@@ -1918,6 +1918,13 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co
 
     const auto & cparams = lctx->get_cparams();
 
+    int32_t pos_max = 0;
+    for (uint32_t s = 0; s < n_stream; ++s) {
+        for (llama_seq_id seq_id = 0; seq_id < LLAMA_MAX_SEQ; ++seq_id) {
+            pos_max = std::max(pos_max, (int32_t) v_cells[s].seq_pos_max(seq_id));
+        }
+    }
+
     for (const auto & layer : layers) {
         const uint32_t il = layer.il;
 
@@ -1931,7 +1938,7 @@ ggml_cgraph * llama_kv_cache::build_graph_shift(llm_graph_result * res, llama_co
         const float freq_base_l  = model.get_rope_freq_base (cparams, il);
         const float freq_scale_l = model.get_rope_freq_scale(cparams, il);
 
-        ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
+        ggml_tensor * rope_factors = model.get_rope_factors(cparams, pos_max, il);
 
         ggml_tensor * k =
             ggml_view_3d(ctx, layer.k,

src/llama-model.cpp

@@ -1985,15 +1985,15 @@ float llama_model::get_rope_freq_scale(const llama_cparams & cparams, int il) co
     return hparams.is_swa(il) ? hparams.rope_freq_scale_train_swa : cparams.rope_freq_scale;
 }
 
-ggml_tensor * llama_model::get_rope_factors(const llama_cparams & cparams, int il) const {
-    const uint32_t n_ctx_seq = cparams.n_ctx_seq;
+ggml_tensor * llama_model::get_rope_factors(const llama_cparams & cparams, int32_t pos_max, int il) const {
+    GGML_UNUSED(cparams);
 
-    // choose long/short freq factors based on the context size
     if (layers[il].rope_freqs != nullptr) {
         return layers[il].rope_freqs;
     }
 
-    if (n_ctx_seq > hparams.n_ctx_orig_yarn) {
+    // long factors once positions reach the original training length
+    if (pos_max >= (int32_t) hparams.n_ctx_orig_yarn) {
         return layers[il].rope_long;
     }
 

src/llama-model.h

@@ -638,7 +638,7 @@ struct llama_model {
     float get_rope_freq_base (const llama_cparams & cparams, int il) const;
     float get_rope_freq_scale(const llama_cparams & cparams, int il) const;
 
-    ggml_tensor * get_rope_factors(const llama_cparams & cparams, int il) const;
+    ggml_tensor * get_rope_factors(const llama_cparams & cparams, int32_t pos_max, int il) const;
 
     llama_memory_i * create_memory(const llama_memory_params & params, const llama_cparams & cparams) const;
 

src/models/apertus.cpp

@@ -84,7 +84,7 @@ llama_model_apertus::graph::graph(const llama_model & model, const llm_graph_par
 
         // self-attention
         {
-            ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
+            ggml_tensor * rope_factors = model.get_rope_factors(cparams, pos_max, il);
 
             // compute Q and K and RoPE them
             auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,

src/models/arcee.cpp

@@ -77,7 +77,7 @@ llama_model_arcee::graph::graph(const llama_model & model, const llm_graph_param
         // self-attention
         {
             // rope freq factors for llama3; may return nullptr for llama2 and other models
-            ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
+            ggml_tensor * rope_factors = model.get_rope_factors(cparams, pos_max, il);
 
             // compute Q and K and RoPE them
             auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,