nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

ggml-cuda: bind cublas handle to the backend stream in DSA attention (#2347) 

The DSA attention kernel used the shared cublas handle without binding
it to the backend's stream, so its Q.K / P.V GEMMs ran on a different
stream than the gather and softmax kernels. The softmax could then read
the score buffer before the GEMM wrote it, picking up stale (NaN) values.
This commit is contained in:
Yap Sok Ann 2026-08-24 23:48:12 +07:00 committed by GitHub
parent c574620b12
commit 26113d1dd3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
ggml/src/ggml-cuda/dsa_attn.cu
View File

@ -344,6 +344,7 @@ bool ggml_cuda_dsa_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst)
(const float *)((const char *)Q->data + first*Q->nb[1]), q16.get());
}
CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(), ctx.stream()));
CUBLAS_CHECK(cublasHgemmStridedBatched(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N,
indexer->ne[0], Q->ne[2], Q->ne[0],
&alpha, k16.get(), K->ne[0], K->ne[0]*indexer->ne[0],