pass stream for initialize (#5003)

Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/2016

Pull Request resolved: #5003

This doesn't actually do anything. But the debug logs it printed would be better.

Reviewed By: sryap

Differential Revision: D84372042

fbshipit-source-id: a0bd1f894005144026de55dac3fce11132677c0e

Author: henrylhtsang

fbgemm_gpu/experimental/gen_ai/src/attention/cuda/cutlass_blackwell_fmha/blackwell_fmha_utils.hpp

@@ -121,7 +121,10 @@ static void launch_fmha_op(const typename Operation::Arguments& arguments) {
       "This kernel is not supported. Last CUDA error is: ",
       cudaGetErrorString(cudaGetLastError()));
 
-  status = op.initialize(arguments, workspace.mutable_data_ptr());
+  status = op.initialize(
+      arguments,
+      workspace.mutable_data_ptr(),
+      at::cuda::getCurrentCUDAStream());
   TORCH_CHECK(
       status == cutlass::Status::kSuccess,
       "Failed to initialize the CUTLASS kernel. Last CUDA error is: ",

fbgemm_gpu/experimental/gen_ai/src/attention/cuda/cutlass_blackwell_fmha/device/fmha.hpp

@@ -205,7 +205,7 @@ class FMHA {
   /// Supplied params struct must be construct by calling Kernel::to_underling_arguments()
   static Status
   run(Params& params, cudaStream_t stream = nullptr) {
-    CUTLASS_TRACE_HOST("FMHA::run()");
+    CUTLASS_TRACE_HOST("FMHA::run(), stream: " << (stream ? "non-null" : "null"));
     dim3 const block = Kernel::get_block_shape();
     dim3 const grid = get_grid_shape(params);
 

fbgemm_gpu/experimental/gen_ai/src/attention/cuda/cutlass_blackwell_fmha/device/fmha_device_bwd.hpp

@@ -331,7 +331,7 @@ class Sm100FmhaBwd {
   /// Supplied params struct must be construct by calling Kernel::to_underling_arguments()
   static Status
   run(Params& params, cudaStream_t stream = nullptr) {
-    CUTLASS_TRACE_HOST("FmhaDeviceBwd::run()");
+    CUTLASS_TRACE_HOST("FmhaDeviceBwd::run(), stream: " << (stream ? "non-null" : "null"));
 
     Status result = Status::kSuccess;
     result = params.op_sum_OdO.run(stream);