cuda: refactored ssm_scan and use CUB

[Your-Cheese]

I modified the structure of the CUDA kernel for the ssm scan such parallelization is performed per thread across the channel dimension (D). This allows A and the initial state (s0) to be loaded into registers and reused across the sequence (L) and SSM state dimensions (N). Additionally, B and C can be loaded into shared memory since blocks process the same timestep in parallel. I also added another CUDA kernel specifically for a sequence length of 1 (recurrent mode) in order to reduce the number of registers used by removing the loop over the sequence dimension.

I'm unsure about optimizing the number of threads per block or the minimum number of blocks per multiprocessor in the launch bounds, however, so I left them as is.

Benchmarks

I got the following results with the following test cases added to test-backend-ops.cpp.

for (int seq_len : {1, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384 }) {
        for (int batch_size : { 1, 2, 4, 8, }) {
            test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 16, 1536, seq_len, batch_size));
        }
    }

./bin/test-backend-ops perf -o SSM_SCAN -b CUDA0

Hardware: Intel i7-13700K, Nvidia RTX 3090
Raw output:
cpu.txt
original_cuda.txt
improved_cuda.txt
improved_cuda_no_cub.txt

llama-bench

./bin/llama-bench -m falcon-mamba-7B-BF16.gguf 

Original:

  Device 0: NVIDIA GeForce RTX 3090, compute capability 8.6, VMM: yes
| model                          |       size |     params | backend    | ngl |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | --------------: | -------------------: |
| mamba ?B BF16                  |  14.10 GiB |     7.27 B | CUDA       |  99 |           pp512 |       2624.38 ± 2.44 |
| mamba ?B BF16                  |  14.10 GiB |     7.27 B | CUDA       |  99 |           tg128 |         39.67 ± 0.22 |

Improved:

  Device 0: NVIDIA GeForce RTX 3090, compute capability 8.6, VMM: yes
| model                          |       size |     params | backend    | ngl |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | --------------: | -------------------: |
| mamba ?B BF16                  |  14.10 GiB |     7.27 B | CUDA       |  99 |           pp512 |      2748.06 ± 11.23 |
| mamba ?B BF16                  |  14.10 GiB |     7.27 B | CUDA       |  99 |           tg128 |         40.64 ± 0.11 |

Improved (No CUB):

  Device 0: NVIDIA GeForce RTX 3090, compute capability 8.6, VMM: yes
| model                          |       size |     params | backend    | ngl |            test |                  t/s |
| ------------------------------ | ---------: | ---------: | ---------- | --: | --------------: | -------------------: |
| mamba ?B BF16                  |  14.10 GiB |     7.27 B | CUDA       |  99 |           pp512 |       2723.28 ± 6.63 |
| mamba ?B BF16                  |  14.10 GiB |     7.27 B | CUDA       |  99 |           tg128 |         39.68 ± 0.15 |

[JohannesGaessler]

Consider deduplicating the code by adding an optional template parameter for L. Ignore the template parameter if it's 0, otherwise use it instead of the runtime parameter (add #pragma unroll to the loops over L). Adding additional template specializations for L <= 8 would likely also improve performance. You can look at softmax.cu for an example.

If you are not doing this already, my recommendation for optimizing CUDA performance would be to first use NVIDIA NSight Systems to identify which kernels take up a large percentage of the total runtime (and are thus worth optimizing). Then you can use NVIDIA NSight Compute to get a detailed breakdown of a specific kernel and to identify bottlenecks. For this kernel I assume the bottleneck is I/0.

[JohannesGaessler]

In GPU code there can be performance issues if you cast to char *, do pointer arithmetic, and then cast back to float *. But since this is only done once here it should be fine and in my experience this mostly affects the HIP port for AMD anyways.

[JohannesGaessler]

In CUDA there are 64k registers per SM and each thread can at most use 255 registers. So with 128 threads the occupancy limit in terms of registers is 4 and telling the compiler to limit register usage in order to fit 2 blocks effectively tells it to just use as many registers as it wants. You could maybe change the args to (splitD, 1) to make this a little clearer but I think it's also fine as-is.

[Your-Cheese]

I could just remove it if it's not doing anything then, so it would be (splitD) only.

[JohannesGaessler]

No, this does in fact do something. The compiler is by default very conservative with how many registers it uses because this avoids the worst-performing cases but it also leaves potential performance on the table. If you explicitly tell the compiler to use as many registers as it wants the performance can be better (for this kernel it probably doesn't matter anyways).

[Your-Cheese]

Oh, I see that's why the register count used was 64 if I removed it. It does seem to make a small difference in performance. I'll change it to 1 since there doesn't seem to be a difference from 2 in the generated assembly.

[JohannesGaessler]

The memory access pattern here is inefficient though I also wouldn't know how to improve it.

[Your-Cheese]

Does the problem lie in that the loads aren't coalesced? Wouldn't using a coalesced loading pattern require the data to be in a different layout?

[JohannesGaessler]

Yes, the problem is the uncoalesced I/O. If you could somehow re-write the kernel to make the loads coalesced or change the memory pattern the previous kernel puts out the performance would likely be better. (I did not try to analyze whether something like this is possible.)

[JohannesGaessler]

The memory access pattern here is also inefficient.

[Your-Cheese]

Consider deduplicating the code by adding an optional template parameter for L. Ignore the template parameter if it's 0, otherwise use it instead of the runtime parameter (add #pragma unroll to the loops over L). Adding additional template specializations for L <= 8 would likely also improve performance. You can look at softmax.cu for an example.

If you are not doing this already, my recommendation for optimizing CUDA performance would be to first use NVIDIA NSight Systems to identify which kernels take up a large percentage of the total runtime (and are thus worth optimizing). Then you can use NVIDIA NSight Compute to get a detailed breakdown of a specific kernel and to identify bottlenecks. For this kernel I assume the bottleneck is I/0.

Here's what I saw when profiling with NSight Systems, and I think the speedup is visible, though it doesn't make too much of an impact overall since the ssm_scan only occupies around 10% of the overall execution time anyways.

It's more significant in prompt processing since the scan takes up a larger portion of time there, while it doesn't make much of a difference on text generation.

Original:

Improved:

I tried using the template parameter for L to deduplicate the functions, but it seems to have reduced the number of registers being used alongside slightly reducing the performance. Is there something I'm missing in how I'm using a template for L?

With 1 min block per MP:

WIthout min blocks per MP: