added missing stream argument for repkv_backward

Author: ngc92

llmc/repkv.cuh

@@ -50,7 +50,7 @@ __global__ void repkv_forward_kernel1(floatX* replicated_qkv,
 
 __global__ void repkv_backward_kernel1(floatX* dinp, const floatX* dout,
                                 int B, int N, int NH, int replicate_factor, int HD) {
-    // we have a single tensor dout of shapae of (B, N 3 * NH * HD)
+    // we have a single tensor dout of shape of (B, N 3 * NH * HD)
     // we want to reduce sum (for K and V) into  (B, N, (NH + 2*(NH/replicate_factor)) * HD)
     int idx = blockIdx.x * blockDim.x + threadIdx.x;
     if (idx >= B * N * 3 * NH * HD) { return;}
@@ -111,11 +111,11 @@ void repkv_forward(floatX* out, const floatX* inp, int B, int T, int NH, int NH_
 }
 
 void repkv_backward(floatX* dinp, const floatX* dout,
-                    const int B, const int T, const int NH, const int NH_KV, const int d) {
+                    const int B, const int T, const int NH, const int NH_KV, const int d, cudaStream_t stream) {
     const int block_size = 128;
     int total_threads = B * T * (3 * NH) * d;
     int num_blocks = CEIL_DIV(total_threads, block_size);
     int replicate_factor = NH / NH_KV;
-    repkv_backward_kernel1<<<num_blocks, block_size>>>(dinp, dout, B, T, NH, replicate_factor, d);
+    repkv_backward_kernel1<<<num_blocks, block_size, 0, stream>>>(dinp, dout, B, T, NH, replicate_factor, d);
     cudaCheck(cudaGetLastError());
 }

train_llama3.cu

@@ -922,7 +922,7 @@ void llama3_backward_and_reduce(LLama3 *model, int* inputs, const int* targets,
         floatX* buffer_b = l_fch_pre_gelu;        // this is B x T x 4C, so even larger than what we need
         attention_backward(dl_bt4c, buffer_b, scratchX, buffer_a, dl_btc, l_qkvr, l_att, B, T, C, NH, main_stream);
         // backward repkv (use scratchX as gradient buffer here)
-        repkv_backward(dl_bt4c2, dl_bt4c, B, T, NH, n_kv_head, hd);
+        repkv_backward(dl_bt4c2, dl_bt4c, B, T, NH, n_kv_head, hd, main_stream);
         #endif
         // backward rope (this can be done in-place)
         rope_backward_inplace(dl_bt4c2, dl_bt4c2, model->freqs_cis, B, T, NH, n_kv_head, hd, main_stream);