nikola/llama_cpp
1
0
Fork 0
Code Issues Pull Requests Actions 23 Packages Projects Releases Wiki Activity

CUDA: Fix ssm_scan_f32 data-races (#24360)

Browse Source 
* Add missing syncthreads before resuing cub_temp_storage

__syncthreads() is required before being allowed to resue TempStorage
smem:
https://nvidia.github.io/cccl/unstable/cub/api/classcub_1_1BlockLoad.html#_CPPv4I0EN3cub9BlockLoad4LoadEv20RandomAccessIteratorRA14ItemsPerThread_1Ti

* Add one more missing __syncthreads

Could also double-buffer, but alternative is to simply ensure all
threads have read smem* before writing to it again in the next loop
iteration

* Remove unused smem from ssm_scan_f32
This commit is contained in:
Oliver Simons
2026-06-10 14:27:08 +02:00
committed by GitHub
parent 039e20a2db
commit fb83cc9a07
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-cuda/ssm-scan.cu
+3 -2
View File
@@ -67,6 +67,7 @@ __global__ void __launch_bounds__(splitD, 1)
__shared__ CubTempStorage cub_temp_storage; __shared__ CubTempStorage cub_temp_storage;
BlockLoad(cub_temp_storage.load_temp).Load(A_block, regA); BlockLoad(cub_temp_storage.load_temp).Load(A_block, regA);
__syncthreads();
BlockLoad(cub_temp_storage.load_temp).Load(s0_block, regs0); BlockLoad(cub_temp_storage.load_temp).Load(s0_block, regs0);
#else #else
const int stride_s0 = src0_nb2 / sizeof(float); const int stride_s0 = src0_nb2 / sizeof(float);
@@ -105,6 +106,7 @@ __global__ void __launch_bounds__(splitD, 1)
regs0[n] = state; regs0[n] = state;
} }
y_block[i * stride_y + threadIdx.x] = sumf; y_block[i * stride_y + threadIdx.x] = sumf;
__syncthreads();
} }
#ifdef USE_CUB #ifdef USE_CUB
@@ -249,9 +251,8 @@ static void ssm_scan_f32_cuda(const float * src0, const float * src1, const floa
GGML_ASSERT(head_dim == 1); GGML_ASSERT(head_dim == 1);
GGML_ASSERT(n_group == 1); GGML_ASSERT(n_group == 1);
const dim3 blocks(n_seq, (n_head + threads - 1) / threads, 1); const dim3 blocks(n_seq, (n_head + threads - 1) / threads, 1);
const int smem_size = (threads * (d_state + 1) * 2) * sizeof(float);
if (d_state == 16) { if (d_state == 16) {
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks, threads, smem_size, stream); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(blocks, threads, 0, stream);
switch (n_tok) switch (n_tok)
{ {
case 1: case 1: