Minor fix in compute_attn_1rowblock_splitkv (#900)

3e9414f1 · ljss · GitHub · 36587c01 · 3e9414f1
Unverified Commit 3e9414f1 authored Mar 28, 2024 by ljss Committed by GitHub Mar 27, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

csrc/flash_attn/src/flash_fwd_kernel.h csrc/flash_attn/src/flash_fwd_kernel.h +2 -2

No files found.
--- a/csrc/flash_attn/src/flash_fwd_kernel.h
+++ b/csrc/flash_attn/src/flash_fwd_kernel.h
@@ -491,8 +491,8 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params &params, cons

    using GmemTiledCopyO = std::conditional_t<
        !Split,
-        typename Kernel_traits::GmemTiledCopyOaccum,
-        typename Kernel_traits::GmemTiledCopyO
+        typename Kernel_traits::GmemTiledCopyO,
+        typename Kernel_traits::GmemTiledCopyOaccum
    >;
    using ElementO = std::conditional_t<!Split, Element, ElementAccum>;