fixed a bug in device_batched_multihead_attention_forward_xdl_cshuffle.hpp

ee8f18f3 · guangzlu · 0db65bfb · ee8f18f3
Commit ee8f18f3 authored Mar 16, 2023 by guangzlu
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

include/ck/tensor_operation/gpu/device/impl/device_batched_multihead_attention_forward_xdl_cshuffle.hpp ...vice_batched_multihead_attention_forward_xdl_cshuffle.hpp +2 -2

No files found.
--- a/include/ck/tensor_operation/gpu/device/impl/device_batched_multihead_attention_forward_xdl_cshuffle.hpp
+++ b/include/ck/tensor_operation/gpu/device/impl/device_batched_multihead_attention_forward_xdl_cshuffle.hpp
@@ -106,8 +106,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,