Adding stream to 1 kernel.

exllamav2/exllamav2_ext/cuda/cache.cu

@@ -97,7 +97,8 @@ void array_fp16_to_fp8_cuda(const half* pIn, unsigned char *pOut, int stride, in
  gridDim.x = DIVIDE((max - min) / 8, THREADS);
  gridDim.y = height;
 
- fp16_to_fp8_kernel<<<gridDim, blockDim>>>(pIn, pOut, stride, height, min, max);
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ fp16_to_fp8_kernel<<<gridDim, blockDim, 0, stream>>>(pIn, pOut, stride, height, min, max);
  // cuda_check( cudaPeekAtLastError() );
 }
 
@@ -113,7 +114,8 @@ void array_fp8_to_fp16_cuda(const unsigned char* pIn, half* pOut, int stride, in
  gridDim.x = DIVIDE((max - min) / 8, THREADS);
  gridDim.y = height;
 
- fp8_to_fp16_kernel<<<gridDim, blockDim>>>(pIn, pOut, stride, height, min, max);
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ fp8_to_fp16_kernel<<<gridDim, blockDim, 0, stream>>>(pIn, pOut, stride, height, min, max);
  // cuda_check( cudaPeekAtLastError() );
 }
 

exllamav2/exllamav2_ext/cuda/h_add.cu

@@ -137,14 +137,15 @@ void cuda_vector_add_
  int width
 )
 {
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  if (width % 8 == 0)
  {
  dim3 blockDim, gridDim;
  blockDim.x = min(NUM_THREADS_INT4, width / 8);
  gridDim.x = DIVIDE(width, NUM_EL_INT4);
  gridDim.y = DIVIDE(height, NUM_THREADS_Y_INT4);
 
- cuda_vector_add_int4_kernel<<<gridDim, blockDim>>>(dest, source, height, width);
+ cuda_vector_add_int4_kernel<<<gridDim, blockDim, 0, stream>>>(dest, source, height, width);
  }
  else
  {
@@ -153,7 +154,7 @@ void cuda_vector_add_
  gridDim.x = DIVIDE(width, NUM_THREADS_X * 2);
  gridDim.y = DIVIDE(height, NUM_THREADS_Y);
 
- cuda_vector_add_kernel<<<gridDim, blockDim>>>(dest, source, height, width);
+ cuda_vector_add_kernel<<<gridDim, blockDim, 0, stream>>>(dest, source, height, width);
  }
 }
 
@@ -165,14 +166,15 @@ void cuda_vector_set_
  int width
 )
 {
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  if (width % 8 == 0)
  {
  dim3 blockDim, gridDim;
  blockDim.x = min(NUM_THREADS_INT4, width / 8);
  gridDim.x = DIVIDE(width, NUM_EL_INT4);
  gridDim.y = DIVIDE(height, NUM_THREADS_Y_INT4);
 
- cuda_vector_set_int4_kernel<<<gridDim, blockDim>>>(dest, source, height, width);
+ cuda_vector_set_int4_kernel<<<gridDim, blockDim, 0, stream>>>(dest, source, height, width);
  }
  else
  {
@@ -181,7 +183,7 @@ void cuda_vector_set_
  gridDim.x = DIVIDE(width, NUM_THREADS_X * 2);
  gridDim.y = DIVIDE(height, NUM_THREADS_Y);
 
- cuda_vector_set_kernel<<<gridDim, blockDim>>>(dest, source, height, width);
+ cuda_vector_set_kernel<<<gridDim, blockDim, 0, stream>>>(dest, source, height, width);
  }
 }
 

exllamav2/exllamav2_ext/cuda/h_gemm.cu

@@ -220,6 +220,7 @@ void h_gemm_cuda
  const float beta
 )
 {
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  if ((beta == 1.0f || beta == 0.0f) && (alpha == 1.0f))
  {
  bool clear = (beta == 0.0f);
@@ -241,7 +242,7 @@ void h_gemm_cuda
 // DBGI3(blockDim.x, blockDim.y, blockDim.z);
 // DBGI3(gridDim.x, gridDim.y, gridDim.z);
 
- h_gemm_tall_kernel<<<gridDim, blockDim>>>(size_m, size_n, size_k, a, b, c, clear);
+ h_gemm_tall_kernel<<<gridDim, blockDim, 0, stream>>>(size_m, size_n, size_k, a, b, c, clear);
  cuda_check( cudaPeekAtLastError() );
  return;
  }
@@ -261,7 +262,7 @@ void h_gemm_cuda
 // DBGI3(blockDim.x, blockDim.y, blockDim.z);
 // DBGI3(gridDim.x, gridDim.y, gridDim.z);
 
- h_gemm_wide_kernel<<<gridDim, blockDim>>>(size_m, size_n, size_k, a, b, c, clear);
+ h_gemm_wide_kernel<<<gridDim, blockDim, 0, stream>>>(size_m, size_n, size_k, a, b, c, clear);
  cuda_check( cudaPeekAtLastError() );
  return;
  }
@@ -271,4 +272,4 @@ void h_gemm_cuda
 // DBGI3(size_m, size_n, size_k);
  cuda_check( cudaPeekAtLastError() );
 
-}
+}

exllamav2/exllamav2_ext/cuda/head_norm.cu

@@ -114,5 +114,6 @@ void head_norm_cuda
 
  float r_dim = 1.0f / (float) head_dim;
 
- head_norm_kernel<<<gridDim, blockDim>>>(x, w, b, y, epsilon, r_dim, rows, num_heads, head_dim);
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ head_norm_kernel<<<gridDim, blockDim, 0, stream>>>(x, w, b, y, epsilon, r_dim, rows, num_heads, head_dim);
 }

exllamav2/exllamav2_ext/cuda/layer_norm.cu

@@ -204,5 +204,6 @@ void layer_norm_cuda
 
  int blocks_per_warp = DIVIDE(dim, NUM_THREADS * 2);
  fp_layer_norm_kernel kernel = pick_layer_norm_kernel(blocks_per_warp);
- kernel<<<gridDim, blockDim>>>(x, w, b, y, epsilon, r_dim, rows, dim, add_residual);
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ kernel<<<gridDim, blockDim, 0, stream>>>(x, w, b, y, epsilon, r_dim, rows, dim, add_residual);
 }

exllamav2/exllamav2_ext/cuda/pack_tensor.cu

@@ -47,7 +47,8 @@ void pack_rows_4_cuda
  dim3 threads(BLOCKSIZE_X, BLOCKSIZE_Y);
  dim3 blocks(DIVIDE(out_columns, BLOCKSIZE_X), DIVIDE(rows, BLOCKSIZE_Y));
 
- pack_rows_4_kernel<<<blocks, threads>>>(input, output, rows, out_columns);
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ pack_rows_4_kernel<<<blocks, threads, 0, stream>>>(input, output, rows, out_columns);
 }
 
 // Pack rows:
@@ -93,7 +94,8 @@ void pack_rows_6_cuda
  dim3 threads(BLOCKSIZE_X, BLOCKSIZE_Y);
  dim3 blocks(DIVIDE(out_columns, BLOCKSIZE_X), DIVIDE(rows, BLOCKSIZE_Y));
 
- pack_rows_6_kernel<<<blocks, threads>>>(input, output, rows, out_columns);
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ pack_rows_6_kernel<<<blocks, threads, 0, stream>>>(input, output, rows, out_columns);
 }
 
 // Pack columns

exllamav2/exllamav2_ext/cuda/q_gemm.cu

@@ -104,7 +104,8 @@ void gemm_half_q_half_cuda_part
 
  // Launch kernel
 
- kernel<<<gridDim, blockDim>>>
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ kernel<<<gridDim, blockDim, 0, stream>>>
  (
  a,
  b->cuda_q_weight,
@@ -165,7 +166,8 @@ void gemm_half_q_half_cuda_part
 // print_global_mem(r_weights, 1, 1, 1);
 // DBGI(r_weights_stride);
 
- kernel<<<gridDim, blockDim>>>
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ kernel<<<gridDim, blockDim, 0, stream>>>
  (
  a,
  b->cuda_q_weight,

exllamav2/exllamav2_ext/cuda/q_matrix.cu

@@ -178,7 +178,8 @@ QMatrix::QMatrix
  gridDim.x = DIVIDE(width, THREADS_X);
  gridDim.y = 1;
 
- shuffle_kernel<<<gridDim, blockDim>>>(cuda_q_weight, height, width, rows_8, rows_6, rows_5, rows_4, rows_3, rows_2);
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ shuffle_kernel<<<gridDim, blockDim, 0, stream>>>(cuda_q_weight, height, width, rows_8, rows_6, rows_5, rows_4, rows_3, rows_2);
 }
 
 QMatrix::~QMatrix()
@@ -491,10 +492,11 @@ void QMatrix::reconstruct(half* out, int row_a, int row_b)
 
  gridDim.y = DIVIDE(row_b - row_a, BLOCK_KN_SIZE);
 
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  if (!is_gptq)
  {
  gridDim.x = DIVIDE(width, BLOCK_KN_SIZE);
- reconstruct_kernel<<<gridDim, blockDim>>>
+ reconstruct_kernel<<<gridDim, blockDim, 0, stream>>>
  (
  cuda_q_weight,
  cuda_q_perm,
@@ -519,7 +521,7 @@ void QMatrix::reconstruct(half* out, int row_a, int row_b)
  else
  {
  gridDim.x = DIVIDE(width, BLOCK_KN_SIZE * 4);
- reconstruct_gptq_kernel<<<gridDim, blockDim>>>
+ reconstruct_gptq_kernel<<<gridDim, blockDim, 0, stream>>>
  (
  cuda_q_weight,
  cuda_q_perm,
@@ -640,7 +642,8 @@ bool QMatrix::make_sequential(const uint32_t* cpu_g_idx)
  gridDim.x = DIVIDE(width, THREADS_X);
  gridDim.y = height / 8;
 
- make_sequential_kernel<<<gridDim, blockDim>>>
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ make_sequential_kernel<<<gridDim, blockDim, 0, stream>>>
  (
  cuda_q_weight,
  cuda_new_qweight,
@@ -722,7 +725,8 @@ void matrix_fp8_to_fp16_cuda
  dim3 blockDim, gridDim;
  blockDim.x = THREADS_F;
  gridDim.x = numel / (BLOCKSIZE_F * THREADS_F);
- matrix_fp8_to_fp16_kernel<<<gridDim, blockDim>>>(in_ptr, out_ptr);
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ matrix_fp8_to_fp16_kernel<<<gridDim, blockDim, 0, stream>>>(in_ptr, out_ptr);
 }
 
 void matrix_fp16_to_fp8_cuda
@@ -738,7 +742,8 @@ void matrix_fp16_to_fp8_cuda
  dim3 blockDim, gridDim;
  blockDim.x = THREADS_F;
  gridDim.x = numel / (BLOCKSIZE_F * THREADS_F);
- matrix_fp16_to_fp8_kernel<<<gridDim, blockDim>>>(in_ptr, out_ptr);
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ matrix_fp16_to_fp8_kernel<<<gridDim, blockDim, 0, stream>>>(in_ptr, out_ptr);
 }
 
 // Q4/FP16 convert funcs

exllamav2/exllamav2_ext/cuda/q_mlp.cu

@@ -96,6 +96,7 @@ void QMLP::forward_
 
  // Up proj with gate
 
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  if (gate)
  {
  gemm_half_q_half_cuda(cublas_handle, norm_state, gate, temp_a, rows, intermediate_size, columns, true, temp_dq);
@@ -105,7 +106,7 @@ void QMLP::forward_
  apply_loras_cuda(cublas_handle, up_proj_lora, loras, up, norm_state, temp_b, lora_temp, rows);
 
  fp_act_mul_kernel kernel = pick_act_mul_kernel(use_half2, false, act_gelu);
- kernel<<<gridDim, blockDim>>>(temp_a, temp_b, rows, intermediate_size, NULL, 0);
+ kernel<<<gridDim, blockDim, 0, stream>>>(temp_a, temp_b, rows, intermediate_size, NULL, 0);
  }
 
  // Up proj without gate
@@ -117,7 +118,7 @@ void QMLP::forward_
  apply_loras_cuda(cublas_handle, up_proj_lora, loras, up, norm_state, temp_a, lora_temp, rows);
 
  fp_act_kernel kernel = pick_act_kernel(use_half2, false, act_gelu);
- kernel<<<gridDim, blockDim>>>(temp_a, rows, intermediate_size, NULL, 0);
+ kernel<<<gridDim, blockDim, 0, stream>>>(temp_a, rows, intermediate_size, NULL, 0);
  }
 
  // Down proj without post_layernorm
@@ -244,12 +245,13 @@ void QMoEMLP::forward_
  blockDim.y = 1;
  gridDim.x = 1;
  gridDim.y = DIVIDE(rows, WARPS);
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  if (num_experts == 4)
- softmax4_topk_norm_kernel<<<gridDim, blockDim>>>(temp_logits, rows, num_experts_per_token);
+ softmax4_topk_norm_kernel<<<gridDim, blockDim, 0, stream>>>(temp_logits, rows, num_experts_per_token);
  else if (num_experts == 8)
- softmax8_topk_norm_kernel<<<gridDim, blockDim>>>(temp_logits, rows, num_experts_per_token);
+ softmax8_topk_norm_kernel<<<gridDim, blockDim, 0, stream>>>(temp_logits, rows, num_experts_per_token);
  else if (num_experts == 16)
- softmax16_topk_norm_kernel<<<gridDim, blockDim>>>(temp_logits, rows, num_experts_per_token);
+ softmax16_topk_norm_kernel<<<gridDim, blockDim, 0, stream>>>(temp_logits, rows, num_experts_per_token);
 
  // For small no. rows, execute all kernels but pass the routing weights. Rows with a weight of zero will skip dot
  // product accum and kernels launched with only zero-weights will exit prematurely.
@@ -271,7 +273,7 @@ void QMoEMLP::forward_
  blockDim.y = THREADS_Y;
  gridDim.x = DIVIDE(intermediate_size, THREADS_X) / (use_half2 ? 2 : 1);
  gridDim.y = DIVIDE(rows, THREADS_Y);
- kernel<<<gridDim, blockDim>>>(temp_a, temp_b, rows, intermediate_size, temp_logits + i, num_experts);
+ kernel<<<gridDim, blockDim, 0, stream>>>(temp_a, temp_b, rows, intermediate_size, temp_logits + i, num_experts);
 
  gemm_half_q_half_cuda(cublas_handle, temp_a, w2[i], x, rows, columns, intermediate_size, false, temp_dq, true, temp_logits + i, num_experts, true);
 

exllamav2/exllamav2_ext/cuda/quantize.cu

@@ -66,7 +66,8 @@ void quantize_rtn_cuda
  dim3 threads(BLOCKSIZE_X, 1);
  dim3 blocks(DIVIDE(columns, BLOCKSIZE_X), 1);
 
- quantize_rtn_kernel<<<blocks, threads>>>
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ quantize_rtn_kernel<<<blocks, threads, 0, stream>>>
  (
  weights,
  scale,
@@ -151,7 +152,8 @@ void fused_quantize_adjust_cuda
  dim3 threads(BLOCKSIZE_X, 1);
  dim3 blocks(DIVIDE(columns, BLOCKSIZE_X), 1);
 
- fused_quantize_adjust_kernel<<<blocks, threads>>>
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ fused_quantize_adjust_kernel<<<blocks, threads, 0, stream>>>
  (
  weights,
  quant,
@@ -232,7 +234,8 @@ void quantize_cuda
 // DBGI2(rows, columns);
 // DBGF2(qzero, maxq);
 
- quantize_kernel<<<blocks, threads>>>
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ quantize_kernel<<<blocks, threads, 0, stream>>>
  (
  input,
  output,
@@ -281,7 +284,8 @@ void adjust_error_row_cuda
  dim3 threads(BLOCKSIZE_X, 1);
  dim3 blocks(DIVIDE(columns, BLOCKSIZE_X), 1);
 
- adjust_error_row_kernel<<<blocks, threads>>>(hessian_inv, error, weights, quant, c, columns, hcolumns);
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ adjust_error_row_kernel<<<blocks, threads, 0, stream>>>(hessian_inv, error, weights, quant, c, columns, hcolumns);
 }
 
 __global__ void quantize_err_kernel
@@ -353,7 +357,8 @@ void quantize_err_cuda
 // DBGI2(rows, columns);
 // DBGF2(qzero, maxq);
 
- quantize_err_kernel<<<blocks, threads>>>
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ quantize_err_kernel<<<blocks, threads, 0, stream>>>
  (
  input,
  output,
@@ -414,5 +419,6 @@ void vv_mul_sub_cuda
  gridDim.y = DIVIDE(x_size, BLOCKSIZE_Y);
  gridDim.z = 1;
 
- vv_mul_sub_kernel<<<gridDim, blockDim>>>(x, y, z, x_size, y_size);
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ vv_mul_sub_kernel<<<gridDim, blockDim, 0, stream>>>(x, y, z, x_size, y_size);
 }

exllamav2/exllamav2_ext/cuda/rms_norm.cu

@@ -220,5 +220,6 @@ void rms_norm_cuda
 
  int blocks_per_warp = DIVIDE(dim, NUM_THREADS * 2);
  fp_rms_norm_kernel kernel = pick_rms_norm_kernel(blocks_per_warp);
- kernel<<<gridDim, blockDim>>>(x, w, y, epsilon, r_dim, rows, dim, add_residual, input_fp32, output_fp32);
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ kernel<<<gridDim, blockDim, 0, stream>>>(x, w, y, epsilon, r_dim, rows, dim, add_residual, input_fp32, output_fp32);
 }

exllamav2/exllamav2_ext/cuda/rope.cu

@@ -195,7 +195,8 @@ void rope_cuda
  gridDim.y = DIVIDE(rows_per_batch, threads_y);
  gridDim.z = batch_size;
 
- rope_cuda_kernel<<<gridDim, blockDim>>>
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ rope_cuda_kernel<<<gridDim, blockDim, 0, stream>>>
  (
  x,
  sin,
@@ -240,7 +241,8 @@ void rope_cuda_qk
  gridDim.y = DIVIDE(rows_per_batch, threads_y);
  gridDim.z = batch_size;
 
- rope_cuda_qk_kernel<<<gridDim, blockDim>>>
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ rope_cuda_qk_kernel<<<gridDim, blockDim, 0, stream>>>
  (
  x_q,
  x_k,

exllamav2/exllamav2_ext/cuda/softcap.cu

@@ -33,7 +33,8 @@ void softcap_cuda_
  blockDim.x = NUM_THREADS;
  gridDim.x = DIVIDE(numel, NUM_THREADS);
 
- cuda_softcap_kernel<<<gridDim, blockDim>>>(x, numel, scale);
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ cuda_softcap_kernel<<<gridDim, blockDim, 0, stream>>>(x, numel, scale);
 }
 
 // TODO: Profile
@@ -73,6 +74,7 @@ void h_softcap_cuda_
  blockDim.x = NUM_THREADS;
  gridDim.x = DIVIDE(numel / 2, NUM_THREADS);
 
- h_cuda_softcap_kernel<<<gridDim, blockDim>>>(x, numel, scale);
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+ h_cuda_softcap_kernel<<<gridDim, blockDim, 0, stream>>>(x, numel, scale);
 }