Minor fix in compute_attn_1rowblock_splitkv (#900)

0c07b792 · ljss · skrider · 79681482 · 0c07b792
Commit 0c07b792 authored Mar 28, 2024 by ljss Committed by skrider Mar 28, 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
@@ -492,8 +492,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>;