bug : Vulkan multi-stream attention problem

[ggerganov]

While investigating #14893 (reply in thread), it appears that there is a problem with parallel streams on the Vulkan backend when LLAMA_SET_ROWS=1 is enabled. The problem occurs only when -fa is disabled.

Here is a basic repro:

make -j && LLAMA_SET_ROWS=1 ./bin/llama-parallel -hf ggml-org/Qwen2.5-Coder-3B-Q8_0-GGUF -np 2 -ns 32 -s 1 -c 4096 -ngl 99 --top-k 1

Garbage output in second request

Input:    What is the meaning of life?
Response: The meaning of life is a philosophical question that has been debated for centuries. There is no one definitive answer, but some people believe that the meaning of life is to find happiness and fulfillment in one's own life. Others believe that the meaning of life is to contribute to society and make a positive impact on the world. Ultimately, the meaning of life is up to each individual to define for themselves.

0.05.487.441 I Client   0, seq    2, junk =    0, prompt = 284, started decoding ...
0.07.614.732 I Client   1, seq   1/ 32, prompt  284 t, response  128 t, time  5.29 s, speed 77.83 t/s, cache miss 0  

Input:    What is the meaning of life?
Response: The  _,  _,SID  _,  _,SID
 _,
 _,
 _,SID   _,
 _,
 _,
 _,
 _,SIDdata  _,SID  _, =  _,  _,  _,  _,  _,SID =  _,  _,SID`  _,
 _,  _,


 _,SID  _,  _,  _,  _,  _,SIDim  _,SID =  _, _,  _,
 _,SID  _,
 _,  _,\
 _,  _,  _,
 _,
 _,SID
 _,  _,SID
 _,SID
  _,ID  _,
 _,\  _,SID =

The problem is not observed if:

• Remove LLAMA_SET_ROWS=1
• Keep LLAMA_SET_ROWS=1 and add -fa

This makes me think there is bug in the SOFT_MAX implementation for batched cases (#14449).

Edit: investigating further, it seems that the problem is not in the SOFT_MAX, because when I keep it on the CPU, the problem persists:

diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index 75b58c26f..3f3efad7f 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -10863,7 +10863,9 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
         case GGML_OP_PAD:
         case GGML_OP_ROLL:
         case GGML_OP_DIAG_MASK_INF:
+            return true;
         case GGML_OP_SOFT_MAX:
+            return false;
         case GGML_OP_SOFT_MAX_BACK:
         case GGML_OP_ARGSORT:
         case GGML_OP_SUM:

Looking at the batched matrix multiplications now across dim 3.

[ggerganov]

cc @jeffbolznv @0cc4m

[ggerganov]

I have narrowed it down to the KQV multiplication. The following patch fixes the issue:

diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index 491a26b63..9ad5cbb5f 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -1234,6 +1234,7 @@ ggml_tensor * llm_graph_context::build_attn_mha(
             v = ggml_cont(ctx0, ggml_transpose(ctx0, v));
         }
 
+        v = ggml_cont(ctx0, v);
         ggml_tensor * kqv = ggml_mul_mat(ctx0, v, kq);
 
         // for MLA with the absorption optimization, we need to "decompress" from MQA back to MHA

So it seems some case of 4d inputs with permuted, non-contiguous a matrix multiplication does not seem to be handled correctly. Will try to add a test-backend-ops test to reproduce it.

[jeffbolznv]

I can reproduce this with llama-parallel. Disabling the ggml_vk_mul_mat_vec_nc_f16_f32 path fixes it. @ggerganov let me know if you're close to having a backend test, otherwise I can try to do it.

[ggerganov]

Sorry, I got distracted and haven't looked since my earlier comment. If you have repro - go for it, as I will be busy for a bit.

[jeffbolznv]

I think this path is not handling ne[3] > 1. I'll see if it's easy to add, otherwise I'll route it through another path.