fixed dropout random number write in gridwise attn fwd

3724ab55 · guangzlu · 36d2871b · 3724ab55
Commit 3724ab55 authored Mar 08, 2023 by guangzlu
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

include/ck/tensor_operation/gpu/grid/gridwise_batched_multihead_attention_forward_xdl_cshuffle.hpp ...wise_batched_multihead_attention_forward_xdl_cshuffle.hpp +2 -2

No files found.
--- a/include/ck/tensor_operation/gpu/grid/gridwise_batched_multihead_attention_forward_xdl_cshuffle.hpp
+++ b/include/ck/tensor_operation/gpu/grid/gridwise_batched_multihead_attention_forward_xdl_cshuffle.hpp
@@ -891,8 +891,8 @@ struct GridwiseBatchedMultiheadAttentionForward_Xdl_CShuffle
                     n3, // NInputNum
                     n4>,
            Sequence<0, 1, 2, 3, 4, 5, 6, 7, 8, 9>,
-            9,  // DstVectorDim
-            n4, // DstScalarPerVector
+            9, // DstVectorDim
+            1, // DstScalarPerVector
            InMemoryDataOperationEnum::Set,
            1, // DstScalarStrideInVector
            true>{z_grid_desc_m0_n0_m1_n1_m2_n2_m3_n3_n4_n5,