ggml : add GPU support for Mamba models

[ggerganov]

Recently, initial Mamba support (CPU-only) has been introduced in #5328 by @compilade

In order to support running these models efficiently on the GPU, we seem to be lacking kernel implementations for the following 2 ops:

• GGML_OP_SSM_CONV
• GGML_OP_SSM_SCAN

Creating this issue to keep track of this and give more visibility of this feature. Help with implementing the missing kernels for CUDA and Metal (and other backends potentially) is welcome. We can also discuss if anything else is required to better support this architecture in llama.cpp

[jploski]

I tried to add these both operations (in a very naive form, just the boilerplate and copy/paste of the CPU implementation, replacing memcpy by for loops, and using 1 block with WARP_SIZE==32 threads) here:

jploski@677ad0a

While it works in the sense that the implementation generates plausible-looking output, it differs from the CPU version. The initial prompt seems to be processed incorrectly, with the first generated token already looking off, and the generation not influenced by the content of the prompt. Furthermore, changing batch size (the -b parameter) affects the generation (it shouldn't). So maybe something else needs to be improved besides the new ops implementation.

Sample correct output from CPU version:

/mnt/seagate/dalai/llama.cpp.bak/bin/main -m /mnt/f2fs/mamba/tinystories.gguf -t 1 -p 'Sara and Ben are playing in the snow.' --temp 0.0 -c 1024 --keep -1 --repeat_penalty 1.3 -n 1024
Sara and Ben are playing in the snow. They make a big snowman with a hat, a scarf and a carrot nose. They put on their warm clothes and go inside to get some cold chocolate.
"Look at my snowman!" Sara says. "He is so cool! He can talk like us."
Ben looks at his snowman and smiles. He likes the hat, the scarf and the carrot nose too. He wants to make a new friend.
Sara nods and takes her hat from the closet. She puts it on Ben's head and says, "Hello, Mr. Snowman! Do you want to play with me?"
Ben looks at Sara and smiles. He likes his snowman. He says, "Yes, Mr. Snowman! I like your hat, scarf and carrot nose too."
They hug each other and play together in the snow. They make a big snowman, a blue snowman, a red snowman and a yellow snowman. They are happy to have fun with their new friend. [end of text]

llama_print_timings:        load time =     299.08 ms
llama_print_timings:      sample time =     167.65 ms /   206 runs   (    0.81 ms per token,  1228.74 tokens per second)
llama_print_timings: prompt eval time =     123.71 ms /    11 tokens (   11.25 ms per token,    88.92 tokens per second)
llama_print_timings:        eval time =    5023.01 ms /   205 runs   (   24.50 ms per token,    40.81 tokens per second)
llama_print_timings:       total time =    5394.42 ms /   216 tokens

Sample incorrect output from the GPU version:

Sara and Ben are playing in the snow. not too sure to eat it.swed wasmometimes been
"I am so happy when a room is gone," was were, into them witholidated. They had done something for their family.
They were out at the park andsled down together best best of amidt- covari.
The sun was in his eye not long before they were been to one. He wasmnoted to each other and this is where he wasered been been been been been been been been been done.
He had been been beby, so to all who have been been there been been been been been been taken out of them best part onlys in the long- march with us.
The end. [end of text]

llama_print_timings:        load time =    1359.73 ms
llama_print_timings:      sample time =     117.06 ms /   147 runs   (    0.80 ms per token,  1255.71 tokens per second)
llama_print_timings: prompt eval time =      52.04 ms /    11 tokens (    4.73 ms per token,   211.40 tokens per second)
llama_print_timings:        eval time =     911.71 ms /   146 runs   (    6.24 ms per token,   160.14 tokens per second)
llama_print_timings:       total time =    1135.02 ms /   157 tokens

The GPU output looks especially mangled in the example above, but with a sinlge-token prompt consisting of a space it almost looks "normal" for this model, which makes me suspect that maybe it's something unrelated to the new CUDA kernels:

 avorable, a little girl was walking in the park. She saw something strange and she wanted to take a look. She walked closer and saw that it was a big, round, green frog! The frog hopped around and said "Hello!"
The girl was so surprised but also excited. She asked the frog if he could help her. The frog smiled and said "Of course I can help you!" He took out some food from his bag and gave it to the girl.
"Thank you very much," she said with a big smile on her face. The frog then hopped away, but he was still there when she got home.
The girl was so happy that she had helped the frog. She thanked him again and went back to playing in the park. [end of text]

(The same also happens when I reduce the CUDA kernels to use just a single thread rather than 32.)

[jploski]

Another observation: with one offloaded layer (-ngl 1) I get the same output as for -ngl 0 (= CPU). But as the number of offloaded layers is increased, the output changes each time.

[slaren]

You can add tests for the ops in test-backend-ops to compare the GPU and CPU implementations.

[jploski]

You can add tests for the ops in test-backend-ops to compare the GPU and CPU implementations.

Thanks. I added a test for each of the new ops here: 35f2f86

The CPU vs. GPU comparison ("test-backend-ops test -o SSM_CONV", "test-backend-ops test -o SSM_SCAN") passes. Maybe because of non-representative input. I set "sq", whatever it means, to zero because initializing it to any other value triggered an assertion on the CPU backend. (I should mention that I have little clue about how the Mamba algorithm works in theory or practice; I just attempted a blind port FWIW.)

[compilade]

@jploski Thank you for doing this!

I set "sq", whatever it means, to zero because initializing it to any other value triggered an assertion on the CPU backend

Zero is correct, sq is the seq_id of the tokens corresponding to each part of the input. It was added for simultaneous sequence processing.

But note that this is being changed in #7531. sq no longer exists there, because equal-length simultaneous sequences allow directly using the fourth tensor dimension to separate the states for each sequence, which simplifies SSM_SCAN and SSM_CONV a bit.

https://github.com/ggerganov/llama.cpp/blob/61200ef29fc0e76f264ada583b77e9228120779f/ggml.c#L16349-L16418

The CPU vs. GPU comparison ("test-backend-ops test -o SSM_CONV", "test-backend-ops test -o SSM_SCAN") passes. Maybe because of non-representative input.

Hmm. If both result in the same output with random input (for all tensors apart from sq), maybe the problem isn't with the operators, but how the data gets in and out of them in practice?

I wonder if using ctx_layer instead of ctx_split for the tensors that get passed to ggml_ssm_conv and ggml_ssm_scan would help?

diff --git a/llama.cpp b/llama.cpp
index 841be1de..0c710f8d 100644
--- a/llama.cpp
+++ b/llama.cpp
@@ -6099,7 +6099,7 @@ static bool llm_load_tensors(
 
                         layer.ssm_in = ml.create_tensor(ctx_split, tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, 2*d_inner});
 
-                        layer.ssm_conv1d = ml.create_tensor(ctx_split, tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner});
+                        layer.ssm_conv1d = ml.create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner});
                         layer.ssm_conv1d_b = ml.create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_CONV1D, "bias", i), {d_inner});
 
                         layer.ssm_x = ml.create_tensor(ctx_split, tn(LLM_TENSOR_SSM_X, "weight", i), {d_inner, dt_rank + 2*d_state});
@@ -6108,7 +6108,7 @@ static bool llm_load_tensors(
                         layer.ssm_dt_b = ml.create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner});
 
                         // no "weight" suffix for these
-                        layer.ssm_a = ml.create_tensor(ctx_split, tn(LLM_TENSOR_SSM_A, i), {d_state, d_inner});
+                        layer.ssm_a = ml.create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_A, i), {d_state, d_inner});
                         layer.ssm_d = ml.create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_D, i), {d_inner});
 
                         // out_proj

I have no idea what exactly this would do, but I guess it's worth a try?

[jploski]

I have no idea what exactly this would do, but I guess it's worth a try?

Replacing ctx_split by ctx_layer did not change anything.

Maybe it's good to wait with further tests until #7531 is merged (there's slight hope that the observed issue might go away...)

[jploski]

I have no idea what exactly this would do, but I guess it's worth a try?

Replacing ctx_split by ctx_layer did not change anything.

Maybe it's good to wait with further tests until #7531 is merged (there's slight hope that the observed issue might go away...)

So I was impatient, and applied my changes to #7531 here:

697fab6

The result is as before. The CPU output is the same as before (meaning that #7531 did not break anything.) The CPU vs. GPU backend tests for ssm_conv and ssm_scan with random data pass, but the actual generation derails (even faster than before, namely already with -ngl 1). The differences feel like there's some unintended "latent space exploration" going on, but with more GPU layers they get worse.

CPU:

Sara and Ben are playing in the snow. They make a big snowman with a hat, a scarf and a carrot nose. They put on their warm clothes and go inside to get some cold chocolate.
"Look at my snowman!" Sara says. "He is so cool! He can talk like us."
Ben looks at his snowman and smiles. He likes the hat, the scarf and the carrot nose too. He wants to make a new friend.

GPU (-ngl 1):

Sara and Ben are playing in the snow. They make a big snowman with a hat, a scarf and a carrot nose. They put on their warm clothes and go inside to have some hot chocolate.
"Look at my snowman!" Sara says. "He is so cool! He can do anything."
Ben looks around his room and sees that the snowman has no hat or scarf. He feels sad for him.

[slaren]

ctx_split should only be used for matrices. Unless using CUDA with -sm row, it is the same as ctx_layer. If the results look reasonable and the tests pass, they might not be anything wrong, the results from different backends are always slightly different due to small floating point differences. perplexity is good for verifying that.