Merge branch 'mha-train-develop' into mha-train-ldsbypass

include/ck/tensor_operation/gpu/device/impl/device_batched_mha_fwd_xdl_cshuffle_v1.hpp

@@ -113,8 +113,8 @@ __global__ void
                 p_b_grid + b_batch_offset,
                 p_b1_grid + b1_batch_offset,
                 p_c_grid + c_batch_offset,
-                nullptr ? nullptr : p_z_grid + z_batch_offset,
-                nullptr ? nullptr : p_lse_grid + lse_batch_offset,
+                p_z_grid == nullptr ? nullptr : p_z_grid + z_batch_offset,
+                p_lse_grid == nullptr ? nullptr : p_lse_grid + lse_batch_offset,
                 p_shared,
                 a_element_op,
                 b_element_op,
@@ -142,8 +142,8 @@ __global__ void
             p_b_grid + b_batch_offset,
             p_b1_grid + b1_batch_offset,
             p_c_grid + c_batch_offset,
-            nullptr ? nullptr : p_z_grid + z_batch_offset,
-            nullptr ? nullptr : p_lse_grid + lse_batch_offset,
+            p_z_grid == nullptr ? nullptr : p_z_grid + z_batch_offset,
+            p_lse_grid == nullptr ? nullptr : p_lse_grid + lse_batch_offset,
             p_shared,
             a_element_op,
             b_element_op,

include/ck/tensor_operation/gpu/device/impl/device_batched_mha_fwd_xdl_cshuffle_v2.hpp

@@ -117,8 +117,8 @@ __global__ void
                 p_b_grid + b_batch_offset,
                 p_b1_grid + b1_batch_offset,
                 p_c_grid + c_batch_offset,
-                nullptr ? nullptr : p_z_grid + z_batch_offset,
-                nullptr ? nullptr : p_lse_grid + lse_batch_offset,
+                p_z_grid == nullptr ? nullptr : p_z_grid + z_batch_offset,
+                p_lse_grid == nullptr ? nullptr : p_lse_grid + lse_batch_offset,
                 p_shared,
                 a_element_op,
                 b_element_op,
@@ -148,8 +148,8 @@ __global__ void
             p_b_grid + b_batch_offset,
             p_b1_grid + b1_batch_offset,
             p_c_grid + c_batch_offset,
-            nullptr ? nullptr : p_z_grid + z_batch_offset,
-            nullptr ? nullptr : p_lse_grid + lse_batch_offset,
+            p_z_grid == nullptr ? nullptr : p_z_grid + z_batch_offset,
+            p_lse_grid == nullptr ? nullptr : p_lse_grid + lse_batch_offset,
             p_shared,
             a_element_op,
             b_element_op,

include/ck/tensor_operation/gpu/grid/gridwise_batched_mha_fwd_xdl_cshuffle_v1.hpp

@@ -1172,7 +1172,7 @@ struct GridwiseBatchedMultiheadAttentionForward_Xdl_CShuffle_V1
             static_for<0, MXdlPerWave, 1>{}(
                 [&](auto I) { lse_thread_buf(I) = running_max(I) + math::log(running_sum(I)); });
 
-            if(get_warp_local_1d_id() < AccM2)
+            if(get_lane_local_1d_id() < AccM2)
             {
                 static_for<0, MXdlPerWave, 1>{}([&](auto I) {
                     // copy from VGPR to Global

include/ck/tensor_operation/gpu/grid/gridwise_batched_mha_fwd_xdl_cshuffle_v2.hpp

@@ -1350,7 +1350,7 @@ struct GridwiseBatchedMultiheadAttentionForward_Xdl_CShuffle_V2
             static_for<0, MXdlPerWave, 1>{}(
                 [&](auto I) { lse_thread_buf(I) = running_max(I) + math::log(running_sum(I)); });
 
-            if(get_warp_local_1d_id() < AccM2)
+            if(get_lane_local_1d_id() < AccM2)
             {
                 static_for<0, MXdlPerWave, 1>{}([&](auto I) {
                     // copy from VGPR to Global

include/ck/utility/get_id.hpp

@@ -19,6 +19,8 @@ __device__ index_t get_thread_global_1d_id() { return blockIdx.x * blockDim.x +
 
 __device__ index_t get_warp_local_1d_id() { return threadIdx.x / get_warp_size(); }
 
+__device__ index_t get_lane_local_1d_id() { return threadIdx.x % get_warp_size(); }
+
 __device__ index_t get_block_1d_id() { return blockIdx.x; }
 
 __device__ index_t get_grid_size() { return gridDim.x; }