flash_bwd_launch_template.h 16.8 KB
Newer Older
1
2
3
/******************************************************************************
 * Copyright (c) 2024, Tri Dao.
 ******************************************************************************/
Tri Dao's avatar
Tri Dao committed
4
5
6
7
8
9
10

#pragma once

#include <ATen/cuda/CUDAContext.h>

#include "static_switch.h"
#include "flash.h"
11
#include "flash_bwd_preprocess_kernel.h"
Tri Dao's avatar
Tri Dao committed
12
13
#include "flash_bwd_kernel.h"

14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
// Determine if the architecture supports FLASH and define a macro to handle parameter modifiers
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
#define ARCH_SUPPORTS_FLASH
#define KERNEL_PARAM_MODIFIER __grid_constant__
#else
#define KERNEL_PARAM_MODIFIER
#endif

// Define a macro for unsupported architecture handling to centralize the error message
#define FLASH_UNSUPPORTED_ARCH printf("FATAL: FlashAttention requires building with sm version sm80-sm90, but was built for < 8.0!");

// Use a macro to clean up kernel definitions
#define DEFINE_FLASH_BACKWARD_KERNEL(kernelName, ...) \
template<typename Kernel_traits, __VA_ARGS__> \
__global__ void kernelName(KERNEL_PARAM_MODIFIER const Flash_bwd_params params)

DEFINE_FLASH_BACKWARD_KERNEL(flash_bwd_dq_dk_dv_loop_kernel, bool Is_dropout, bool Is_causal, bool Has_alibi, bool Is_even_M, bool Is_even_K) {
    #if defined(ARCH_SUPPORTS_FLASH)
       flash::compute_dq_dk_dv<Kernel_traits, Is_dropout, Is_causal, Has_alibi, Is_even_M, Is_even_K>(params);
    #else
        FLASH_UNSUPPORTED_ARCH
    #endif
}

Phil Wang's avatar
Phil Wang committed
38
DEFINE_FLASH_BACKWARD_KERNEL(flash_bwd_dq_dk_dv_loop_seqk_parallel_kernel, bool Is_dropout, bool Is_causal, bool Is_local, bool Has_alibi, bool Is_even_MN, bool Is_even_K, bool Is_softcap) {
39
40
    #if defined(ARCH_SUPPORTS_FLASH)
        static_assert(!(Is_causal && Is_local));  // If Is_local is true, Is_causal should be false
Phil Wang's avatar
Phil Wang committed
41
        flash::compute_dq_dk_dv_seqk_parallel<Kernel_traits, Is_dropout, Is_causal, Is_local, Has_alibi, Is_even_MN, Is_even_K, Is_softcap>(params);
42
43
44
45
46
47
    #else
        FLASH_UNSUPPORTED_ARCH
    #endif
}


Tri Dao's avatar
Tri Dao committed
48
template<bool Clear_dQaccum=true, typename Kernel_traits>
49
__global__ void flash_bwd_dot_do_o_kernel(const Flash_bwd_params params) {
Tri Dao's avatar
Tri Dao committed
50
51
52
53
    flash::compute_dot_do_o<Clear_dQaccum, Kernel_traits>(params);
}

template<typename Kernel_traits>
54
__global__ void flash_bwd_clear_dkvaccum_kernel(const Flash_bwd_params params) {
Tri Dao's avatar
Tri Dao committed
55
56
57
58
    flash::clear_dKVaccum<Kernel_traits>(params);
}

template<typename Kernel_traits>
59
__global__ void flash_bwd_convert_dq_kernel(const Flash_bwd_params params, const int nsplits) {
60
    flash::convert_dQ<Kernel_traits>(params, nsplits);
Tri Dao's avatar
Tri Dao committed
61
62
63
}

template<typename Kernel_traits>
64
__global__ void flash_bwd_convert_dkv_kernel(const Flash_bwd_params params) {
Tri Dao's avatar
Tri Dao committed
65
66
67
68
    flash::convert_dKV<Kernel_traits>(params);
}

template<typename Kernel_traits, bool Is_dropout>
69
void run_flash_bwd_seqk_parallel(Flash_bwd_params &params, cudaStream_t stream) {
Tri Dao's avatar
Tri Dao committed
70
71
72
    const int num_m_block = (params.seqlen_q + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM;
    dim3 grid_m(num_m_block, params.b, params.h);
    const int num_n_block = (params.seqlen_k + Kernel_traits::kBlockN - 1) / Kernel_traits::kBlockN;
73
74
75
76
77
78
    int gridDimx = num_n_block;
    if (params.deterministic) {
        auto dprops = at::cuda::getCurrentDeviceProperties();
        gridDimx = (dprops->multiProcessorCount + params.b * params.h - 1) / (params.b * params.h);
    }
    dim3 grid_n(gridDimx, params.b, params.h);
Tri Dao's avatar
Tri Dao committed
79

80
81
82
83
84
    if (!params.deterministic) {
        flash_bwd_dot_do_o_kernel<true, Kernel_traits><<<grid_m, Kernel_traits::kNThreads, 0, stream>>>(params);
    } else {
        flash_bwd_dot_do_o_kernel<false, Kernel_traits><<<grid_m, Kernel_traits::kNThreads, 0, stream>>>(params);
    }
Tri Dao's avatar
Tri Dao committed
85
86
    C10_CUDA_KERNEL_LAUNCH_CHECK();

87
88
89
    // We want to specialize to is_even_MN and not just is_even_M, since in the case where N is not
    // a multiple of kBlockN, we'll need to apply mask in the loop.
    const bool is_even_MN = params.cu_seqlens_q == nullptr && params.cu_seqlens_k == nullptr && params.seqlen_q % Kernel_traits::kBlockM == 0 && params.seqlen_k % Kernel_traits::kBlockN == 0;
Tri Dao's avatar
Tri Dao committed
90
91
92
    const bool is_even_K = params.d == Kernel_traits::kHeadDim;
    constexpr int smem_size_dq_dk_dv = Kernel_traits::kSmemSize1colblock;
    // printf("smem_size_dq_dk_dv = %d\n", smem_size_dq_dk_dv);
93
    BOOL_SWITCH(params.is_causal, Is_causal, [&] {
94
        BOOL_SWITCH(is_even_MN, IsEvenMNConst, [&] {
95
96
97
            EVENK_SWITCH(is_even_K, IsEvenKConst, [&] {
                LOCAL_SWITCH((params.window_size_left >= 0 || params.window_size_right >= 0) && !params.is_causal, Is_local, [&] {
                    ALIBI_SWITCH(params.alibi_slopes_ptr != nullptr, Has_alibi, [&] {
Phil Wang's avatar
Phil Wang committed
98
99
100
101
                        SOFTCAP_SWITCH(params.softcap > 0.0, Is_softcap, [&] {
                            // If not IsEvenKConst, we also set IsEvenMNConst to false to reduce number of templates.
                            // If head dim > 128, set IsEvenMNConst to false to reduce number of templates
                            // If Is_local, set Is_causal to false
Tri Dao's avatar
Tri Dao committed
102
                            auto kernel = &flash_bwd_dq_dk_dv_loop_seqk_parallel_kernel<Kernel_traits, Is_dropout && !Is_softcap, Is_causal, Is_local && !Is_causal, Has_alibi, IsEvenMNConst && IsEvenKConst && !Is_local && Kernel_traits::kHeadDim <= 128, IsEvenKConst, Is_softcap>;
Phil Wang's avatar
Phil Wang committed
103
104
105
106
107
108
109
110
                            // auto kernel = &flash_bwd_dq_dk_dv_loop_seqk_parallel_kernel<Kernel_traits, false, Is_causal, false, false, true, true>;
                            if (smem_size_dq_dk_dv >= 48 * 1024)  {
                                C10_CUDA_CHECK(cudaFuncSetAttribute(
                                    kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size_dq_dk_dv));
                            }
                            kernel<<<grid_n, Kernel_traits::kNThreads, smem_size_dq_dk_dv, stream>>>(params);
                            C10_CUDA_KERNEL_LAUNCH_CHECK();
                        });
111
                    });
Tri Dao's avatar
Tri Dao committed
112
                });
Tri Dao's avatar
Tri Dao committed
113
114
115
116
117
            });
        });
    });

    auto kernel_dq = &flash_bwd_convert_dq_kernel<Kernel_traits>;
118
    if (Kernel_traits::kSmemdQSize >= 48 * 1024)  {
Tri Dao's avatar
Tri Dao committed
119
120
121
        C10_CUDA_CHECK(cudaFuncSetAttribute(
            kernel_dq, cudaFuncAttributeMaxDynamicSharedMemorySize, Kernel_traits::kSmemdQSize));
    }
122
    kernel_dq<<<grid_m, Kernel_traits::kNThreads, Kernel_traits::kSmemdQSize, stream>>>(params, !params.deterministic ? 1 : gridDimx);
Tri Dao's avatar
Tri Dao committed
123
124
125
126
    C10_CUDA_KERNEL_LAUNCH_CHECK();
}

template<typename Kernel_traits, bool Is_dropout>
127
void run_flash_bwd(Flash_bwd_params &params, cudaStream_t stream) {
128
#ifndef FLASHATTENTION_DISABLE_BACKWARD
129
    run_flash_bwd_seqk_parallel<Kernel_traits, Is_dropout>(params, stream);
130
#endif
Tri Dao's avatar
Tri Dao committed
131
132
133
}

template<typename T>
134
void run_mha_bwd_hdim32(Flash_bwd_params &params, cudaStream_t stream) {
135
    constexpr static int Headdim = 32;
Tri Dao's avatar
Tri Dao committed
136
137
138
139
140
    int device;
    cudaGetDevice(&device);
    int max_smem_per_block;
    cudaError status_ = cudaDeviceGetAttribute(
        &max_smem_per_block, cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
Driss Guessous's avatar
Driss Guessous committed
141
142
143
    if (status_ != cudaSuccess) {
      C10_CUDA_CHECK(status_);
    }
144
    DROPOUT_SWITCH(params.p_dropout < 1.f, Is_dropout, [&] {
Tri Dao's avatar
Tri Dao committed
145
146
        if (max_smem_per_block >= 2 * ((3 * 128 + 2 * 128) * Headdim + 2 * 128 * 128)) { // 104 KB
            if constexpr(!Is_dropout) {  // We can afford more registers to keep V in registers
147
                run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 128, 8, 4, 4, 4, true, false, T>, Is_dropout>(params, stream);
Tri Dao's avatar
Tri Dao committed
148
            } else {
149
                run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 128, 8, 4, 4, 4, false, false, T>, Is_dropout>(params, stream);
Tri Dao's avatar
Tri Dao committed
150
151
            }
        } else {  // 96 KB
152
            run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 128, 8, 4, 4, 4, true, false, T>, Is_dropout>(params, stream);
Tri Dao's avatar
Tri Dao committed
153
154
155
156
157
        }
    });
}

template<typename T>
158
void run_mha_bwd_hdim64(Flash_bwd_params &params, cudaStream_t stream) {
159
    constexpr static int Headdim = 64;
Tri Dao's avatar
Tri Dao committed
160
161
162
163
164
    int device;
    cudaGetDevice(&device);
    int max_smem_per_block;
    cudaError status_ = cudaDeviceGetAttribute(
        &max_smem_per_block, cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
Driss Guessous's avatar
Driss Guessous committed
165
166
167
    if (status_ != cudaSuccess) {
      C10_CUDA_CHECK(status_);
    }
Tri Dao's avatar
Tri Dao committed
168
    // printf("max_smem_per_block = %d\n", max_smem_per_block);
169
    DROPOUT_SWITCH(params.p_dropout < 1.f, Is_dropout, [&] {
Tri Dao's avatar
Tri Dao committed
170
        // Changing AtomLayoutMdQ from 2 to 4 takes the same time
171
172
173
174
        // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 2, false, false, T>>(params, stream);
        // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 2, true, false, T>>(params, stream);
        // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 128, 8, 2, 4, 4, false, false, T>>(params, stream);
        // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 64, 8, 4, 2, 4, false, false, T>, Is_dropout>(params, stream);
Tri Dao's avatar
Tri Dao committed
175
176
        // This is slightly faster. We want to split M more so we need fewer registers to store LSE.
        if (max_smem_per_block >= 144 * 1024) {
177
            run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 128, 8, 4, 4, 4, false, false, T>, Is_dropout>(params, stream);
Tri Dao's avatar
Tri Dao committed
178
            // This has a lot of register spilling
179
            // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 128, 8, 4, 4, 4, true, false, T>, Is_dropout>(params, stream);
Tri Dao's avatar
Tri Dao committed
180
181
        } else {
            // if (params.h == params.h_k) {
182
183
184
185
                // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 4, false, false, T>, Is_dropout>(params, stream);
                run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 4, true, false, T>, Is_dropout>(params, stream);
                // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 64, 8, 4, 2, 4, false, false, T>, Is_dropout>(params, stream);
                // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 64, 8, 4, 2, 4, true, false, T>, Is_dropout>(params, stream);
Tri Dao's avatar
Tri Dao committed
186
187
188
189
            // } else {
            // }
        }
    });
190
191
192
193
    // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 64, 8, 4, 2, 4, true, false, T>>(params, stream);
    // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 4, 2, 2, 2, true, false, T>>(params, stream);
    // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 32, 128, 4, 1, 4, 1, false, false, T>>(params, stream);
    // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 16, 128, 4, 1, 4, 1, false, false, T>>(params, stream);
Tri Dao's avatar
Tri Dao committed
194
    // M=128, N=64 is quite slow, I think because we need to read/write dQaccum twice as many times
195
196
197
    // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 64, 8, 2, 2, 2, false, T>>(params, stream);
    // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 64, 8, false, T>>(params, stream);
    // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 4, false, T>>(params, stream);
Tri Dao's avatar
Tri Dao committed
198

199
    // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 64, 4, 4, 2, 4, false, false, T>>(params, stream);
Tri Dao's avatar
Tri Dao committed
200
201
202
}

template<typename T>
203
void run_mha_bwd_hdim96(Flash_bwd_params &params, cudaStream_t stream) {
204
    constexpr static int Headdim = 96;
Tri Dao's avatar
Tri Dao committed
205
206
207
208
209
    int device;
    cudaGetDevice(&device);
    int max_smem_per_block;
    cudaError status_ = cudaDeviceGetAttribute(
        &max_smem_per_block, cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
Driss Guessous's avatar
Driss Guessous committed
210
211
212
    if (status_ != cudaSuccess) {
      C10_CUDA_CHECK(status_);
    }
Tri Dao's avatar
Tri Dao committed
213
    // printf("max_smem_per_block = %d\n", max_smem_per_block);
214
    DROPOUT_SWITCH(params.p_dropout < 1.f, Is_dropout, [&] {
215
216
        if (max_smem_per_block >= 116 * 1024) {
            if constexpr(!Is_dropout) {  // 92KB
217
                run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 4, true, false, T>, Is_dropout>(params, stream);
218
219
            } else {  // 116 KB
                // This is faster for dropout since we don't have many registers to spare
220
                run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 4, false, false, T>, Is_dropout>(params, stream);
Tri Dao's avatar
Tri Dao committed
221
            }
222
        } else {
223
            run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 4, true, false, T>, Is_dropout>(params, stream);
224
        }
Tri Dao's avatar
Tri Dao committed
225
226
227
228
    });
}

template<typename T>
229
void run_mha_bwd_hdim128(Flash_bwd_params &params, cudaStream_t stream) {
230
    constexpr static int Headdim = 128;
Tri Dao's avatar
Tri Dao committed
231
232
233
234
235
    int device;
    cudaGetDevice(&device);
    int max_smem_per_block;
    cudaError status_ = cudaDeviceGetAttribute(
        &max_smem_per_block, cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
Driss Guessous's avatar
Driss Guessous committed
236
237
238
    if (status_ != cudaSuccess) {
      C10_CUDA_CHECK(status_);
    }
Tri Dao's avatar
Tri Dao committed
239
    // printf("max_smem_per_block = %d\n", max_smem_per_block);
240
    DROPOUT_SWITCH(params.p_dropout < 1.f, Is_dropout, [&] {
241
        // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 32, 128, 8, 2, 2, 2, false, false, T>>(params, stream);
242
243
        // This is faster, in the case of sequence-parallel bwd (where we need fewer registers).
        // Out of these three, the 2nd one is slightly faster (2% faster than the first). Idk why.
244
        // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 2, 2, false, false, T>>(params, stream);
245
        if (max_smem_per_block >= 144 * 1024) {
246
247
248
249
250
251
            run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 2, false, false, T>, Is_dropout>(params, stream);
            // run_flash_bwd_seqk_parallel<Flash_bwd_kernel_traits<Headdim, 128, 128, 8, 4, 4, 4, false, false, T>, Is_dropout>(params, stream);
            // run_flash_bwd_seqk_parallel<Flash_bwd_kernel_traits<Headdim, 128, 128, 8, 4, 4, 4, false, true, T>, Is_dropout>(params, stream);
            // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 2, true, false, T>, Is_dropout>(params, stream);
            // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 64, 8, 4, 2, 2, false, false, T>, Is_dropout>(params, stream);
            // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 64, 8, 4, 2, 2, true, false, T>, Is_dropout>(params, stream);
252
        } else {
253
254
            // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 4, 2, 2, false, false, T>, Is_dropout>(params, stream);
            run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 4, 2, 2, true, false, T>, Is_dropout>(params, stream);
255
        }
256
        // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 128, 8, 2, 4, 4, false, false, T>>(params, stream);
Tri Dao's avatar
Tri Dao committed
257

258
        // run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 128, 64, 8, 4, 4, 4, false, false, T>>(params, stream);
Tri Dao's avatar
Tri Dao committed
259
260
261
262
    });
}

template<typename T>
263
void run_mha_bwd_hdim160(Flash_bwd_params &params, cudaStream_t stream) {
264
    constexpr static int Headdim = 160;
Tri Dao's avatar
Tri Dao committed
265
266
267
268
269
    int device;
    cudaGetDevice(&device);
    int max_smem_per_block;
    cudaError status_ = cudaDeviceGetAttribute(
        &max_smem_per_block, cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
Driss Guessous's avatar
Driss Guessous committed
270
271
272
    if (status_ != cudaSuccess) {
      C10_CUDA_CHECK(status_);
    }
273
    DROPOUT_SWITCH(params.p_dropout < 1.f, Is_dropout, [&] {
Tri Dao's avatar
Tri Dao committed
274
        if (max_smem_per_block >= 116 * 1024) {
275
            run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 4, 4, 4, false, false, T>, Is_dropout>(params, stream);
Tri Dao's avatar
Tri Dao committed
276
        } else {
277
            run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 4, 4, 4, false, true, T>, Is_dropout>(params, stream);
Tri Dao's avatar
Tri Dao committed
278
279
280
281
282
        }
    });
}

template<typename T>
283
void run_mha_bwd_hdim192(Flash_bwd_params &params, cudaStream_t stream) {
284
    constexpr static int Headdim = 192;
Tri Dao's avatar
Tri Dao committed
285
286
287
288
289
    int device;
    cudaGetDevice(&device);
    int max_smem_per_block;
    cudaError status_ = cudaDeviceGetAttribute(
        &max_smem_per_block, cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
Driss Guessous's avatar
Driss Guessous committed
290
291
292
    if (status_ != cudaSuccess) {
      C10_CUDA_CHECK(status_);
    }
293
    DROPOUT_SWITCH(params.p_dropout < 1.f, Is_dropout, [&] {
Tri Dao's avatar
Tri Dao committed
294
        if (max_smem_per_block >= 136 * 1024) {
295
            run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 4, 2, 2, false, false, T>, Is_dropout>(params, stream);
Tri Dao's avatar
Tri Dao committed
296
        } else {
297
            run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 4, 2, 2, true, true, T>, Is_dropout>(params, stream);
Tri Dao's avatar
Tri Dao committed
298
299
300
301
302
        }
    });
}

template<typename T>
303
void run_mha_bwd_hdim256(Flash_bwd_params &params, cudaStream_t stream) {
304
    constexpr static int Headdim = 256;
Tri Dao's avatar
Tri Dao committed
305
306
307
308
309
    int device;
    cudaGetDevice(&device);
    int max_smem_per_block;
    cudaError status_ = cudaDeviceGetAttribute(
        &max_smem_per_block, cudaDevAttrMaxSharedMemoryPerBlockOptin, device);
Driss Guessous's avatar
Driss Guessous committed
310
311
312
    if (status_ != cudaSuccess) {
      C10_CUDA_CHECK(status_);
    }
313
    DROPOUT_SWITCH(params.p_dropout < 1.f, Is_dropout, [&] {
Tri Dao's avatar
Tri Dao committed
314
        if (max_smem_per_block >= 176 * 1024) {  // H100
315
            run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 4, 2, 2, false, false, T>, Is_dropout>(params, stream);
316
        } else if (max_smem_per_block >= 144 * 1024) {  // A100, we don't do double buffering to save smem
317
            run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 64, 8, 4, 2, 2, false, true, T>, Is_dropout>(params, stream);
318
319
320
321
        } else { // sm86 and sm89, max smem is 99 KB. Only works without dropout. V in regs and no double buffering.
            if constexpr (!Is_dropout) {
                run_flash_bwd<Flash_bwd_kernel_traits<Headdim, 64, 32, 8, 4, 1, 2, true, true, T>, false>(params, stream);
            }
Tri Dao's avatar
Tri Dao committed
322
323
324
        }
    });
}