From a860922827c560a42d46bb4c3c8f2dda98b4ead4 Mon Sep 17 00:00:00 2001
From: Erik Schultheis <erik.schultheis@aalto.fi>
Date: Sun, 4 May 2025 22:10:38 +0200
Subject: [PATCH] set stream for attention softmax

---
 llmc/attention.cuh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llmc/attention.cuh b/llmc/attention.cuh
index f6294a2..639ab75 100644
--- a/llmc/attention.cuh
+++ b/llmc/attention.cuh
@@ -263,7 +263,7 @@ void attention_backward(floatX* dinp, floatX* dqkvr, floatX* datt, floatX* scrat
     matmul_cublaslt(dv, scratch, att, nullptr, HS, T, T, stream, false, true, B * NH, T * HS, T * T, T * HS);
     const float scale = 1.0f / sqrtf((float)HS);
     // backward into preatt. this is an in-place operation; datt turns into dpreatt here
-    softmax_autoregressive_backward_inplace_kernel<<<dim3(T / 4, B * NH), 256>>>(datt, att, B, T, C, scale);
+    softmax_autoregressive_backward_inplace_kernel<<<dim3(T / 4, B * NH), 256, 0, stream>>>(datt, att, B, T, C, scale);
     const floatX* dpreatt = datt;
     // backward into q
     matmul_cublaslt(dq, k, dpreatt, nullptr, HS, T, T, stream, false, false, B * NH, T * HS, T * T, T * HS);