Fixed wrong stride. Use v_trans to detect use of FA.

Author: jukofyork

src/llama-graph.cpp

@@ -1432,7 +1432,7 @@ ggml_tensor * llm_graph_context::build_attn(
 
         v_cur = ggml_reshape_2d(ctx0, v_cur, n_embd_v_gqa, n_tokens);
 
-        // note: MLA flash attention now uses the last 512 elements of K in place of V
+        // note: MLA with flash attention now uses the last 512 elements of K in place of V
         if (v_trans || !v_mla) {
             ggml_tensor * v_cache_view = nullptr;
 
@@ -1488,11 +1488,11 @@ ggml_tensor * llm_graph_context::build_attn(
                 ggml_row_size(kv_self->v_l[il]->type, n_embd_head_v),
                 0);
     } else {
-        // note: MLA flash attention now uses the last 512 elements of K in place of V
+        // note: MLA with flash attention now uses the last 512 elements of K in place of V
         v = ggml_view_3d(ctx0, kv_self->k_l[il],
                 n_embd_head_v, n_kv, n_head_kv,
-                ggml_row_size(kv_self->v_l[il]->type, n_embd_v_gqa),
-                ggml_row_size(kv_self->v_l[il]->type, n_embd_head_v),
+                ggml_row_size(kv_self->v_l[il]->type, n_embd_k_gqa),
+                ggml_row_size(kv_self->v_l[il]->type, n_embd_head_k),
                 n_embd_head_k-n_embd_head_v); // offset by n_rot elements
     }
 

src/llama-kv-cache.cpp

@@ -33,7 +33,7 @@ llama_kv_cache_unified::llama_kv_cache_unified(
     const int32_t n_layer = hparams.n_layer;
 
     has_shift = false;
-    can_shift = !(model.arch == LLM_ARCH_DEEPSEEK2 && cparams.flash_attn);
+    can_shift = model.arch != LLM_ARCH_DEEPSEEK2 || v_trans; // TODO: allow context shifting for MLA with flash attention
 
     LLAMA_LOG_INFO("%s: kv_size = %d, type_k = '%s', type_v = '%s', n_layer = %d, can_shift = %d, padding = %d\n",
             __func__, kv_size, ggml_type_name(type_k), ggml_type_name(type_v), n_layer, can_shift, padding);
@@ -100,8 +100,9 @@ llama_kv_cache_unified::llama_kv_cache_unified(
             throw std::runtime_error("failed to create ggml context for kv cache");
         }
 
+        // note: MLA with flash attention now uses the last 512 elements of K in place of V
         ggml_tensor * k = ggml_new_tensor_1d(ctx, type_k, n_embd_k_gqa*kv_size);
-        ggml_tensor * v = ggml_new_tensor_1d(ctx, type_v, model.arch == LLM_ARCH_DEEPSEEK2 && cparams.flash_att ? 0 : n_embd_v_gqa*kv_size);
+        ggml_tensor * v = ggml_new_tensor_1d(ctx, type_v, model.arch != LLM_ARCH_DEEPSEEK2 || v_trans ? n_embd_v_gqa*kv_size : 0);
         ggml_format_name(k, "cache_k_l%d", i);
         ggml_format_name(v, "cache_v_l%d", i);
         k_l.push_back(k);