diff --git a/CMakeLists.txt b/CMakeLists.txt index 643d7d10cc1..88c334bf41f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -180,11 +180,21 @@ endif () if (FA3_ENABLED AND ${CMAKE_CUDA_COMPILER_VERSION} GREATER_EQUAL 12.0) # BF16 source files file(GLOB FA3_BF16_GEN_SRCS - "hopper/instantiations/flash_fwd_hdim64_bf16*_sm90.cu" - "hopper/instantiations/flash_fwd_hdim96_bf16*_sm90.cu" - "hopper/instantiations/flash_fwd_hdim128_bf16*_sm90.cu" - "hopper/instantiations/flash_fwd_hdim192_bf16*_sm90.cu" - "hopper/instantiations/flash_fwd_hdim256_bf16*_sm90.cu") + "hopper/instantiations/flash_fwd_hdim64_bf16_*packgqa_sm90.cu" + "hopper/instantiations/flash_fwd_hdim96_bf16_*packgqa_sm90.cu" + "hopper/instantiations/flash_fwd_hdim128_bf16_*packgqa_sm90.cu" + "hopper/instantiations/flash_fwd_hdim192_bf16_*packgqa_sm90.cu" + "hopper/instantiations/flash_fwd_hdim256_bf16_*packgqa_sm90.cu" + "hopper/instantiations/flash_fwd_hdim64_bf16_split*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim96_bf16_split*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim128_bf16_split*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim192_bf16_split*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim256_bf16_split*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim64_bf16_paged*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim96_bf16_paged*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim128_bf16_paged*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim192_bf16_paged*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim256_bf16_paged*_sm90.cu") # Add these for hdim diff cases file(GLOB FA3_BF16_GEN_SRCS_ "hopper/instantiations/flash_fwd_hdim64_256_bf16*_sm90.cu" @@ -194,11 +204,22 @@ if (FA3_ENABLED AND ${CMAKE_CUDA_COMPILER_VERSION} GREATER_EQUAL 12.0) # FP16 source files file(GLOB FA3_FP16_GEN_SRCS - "hopper/instantiations/flash_fwd_hdim64_fp16*_sm90.cu" - "hopper/instantiations/flash_fwd_hdim96_fp16*_sm90.cu" - "hopper/instantiations/flash_fwd_hdim128_fp16*_sm90.cu" - "hopper/instantiations/flash_fwd_hdim192_fp16*_sm90.cu" - "hopper/instantiations/flash_fwd_hdim256_fp16*_sm90.cu") + "hopper/instantiations/flash_fwd_hdim64_fp16_*packgqa_sm90.cu" + "hopper/instantiations/flash_fwd_hdim96_fp16_*packgqa_sm90.cu" + "hopper/instantiations/flash_fwd_hdim128_fp16_*packgqa_sm90.cu" + "hopper/instantiations/flash_fwd_hdim192_fp16_*packgqa_sm90.cu" + "hopper/instantiations/flash_fwd_hdim256_fp16_*packgqa_sm90.cu" + "hopper/instantiations/flash_fwd_hdim64_fp16_split*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim96_fp16_split*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim128_fp16_split*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim192_fp16_split*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim256_fp16_split*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim64_fp16_paged*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim96_fp16_paged*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim128_fp16_paged*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim192_fp16_paged*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim256_fp16_paged*_sm90.cu" + ) # Add these for hdim diff cases file(GLOB FA3_FP16_GEN_SRCS_ "hopper/instantiations/flash_fwd_hdim64_256_fp16*_sm90.cu" @@ -208,11 +229,21 @@ if (FA3_ENABLED AND ${CMAKE_CUDA_COMPILER_VERSION} GREATER_EQUAL 12.0) # FP8 source files file(GLOB FA3_FP8_GEN_SRCS - "hopper/instantiations/flash_fwd_hdim64_e4m3*_sm90.cu" - "hopper/instantiations/flash_fwd_hdim96_e4m3*_sm90.cu" - "hopper/instantiations/flash_fwd_hdim128_e4m3*_sm90.cu" - "hopper/instantiations/flash_fwd_hdim192_e4m3*_sm90.cu" - "hopper/instantiations/flash_fwd_hdim256_e4m3*_sm90.cu") + "hopper/instantiations/flash_fwd_hdim64_e4m3_*packgqa_sm90.cu" + "hopper/instantiations/flash_fwd_hdim96_e4m3_*packgqa_sm90.cu" + "hopper/instantiations/flash_fwd_hdim128_e4m3_*packgqa_sm90.cu" + "hopper/instantiations/flash_fwd_hdim192_e4m3_*packgqa_sm90.cu" + "hopper/instantiations/flash_fwd_hdim256_e4m3_*packgqa_sm90.cu" + "hopper/instantiations/flash_fwd_hdim64_e4m3_split*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim96_e4m3_split*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim128_e4m3_split*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim192_e4m3_split*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim256_e4m3_split*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim64_e4m3_paged*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim96_e4m3_paged*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim128_e4m3_paged*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim192_e4m3_paged*_sm90.cu" + "hopper/instantiations/flash_fwd_hdim256_e4m3_paged*_sm90.cu") # Add these for hdim diff cases (192 only) file(GLOB FA3_FP8_GEN_SRCS_ "hopper/instantiations/flash_fwd_hdim192_128_e4m3*_sm90.cu") @@ -261,11 +292,12 @@ if (FA3_ENABLED AND ${CMAKE_CUDA_COMPILER_VERSION} GREATER_EQUAL 12.0) FLASHATTENTION_DISABLE_BACKWARD FLASHATTENTION_DISABLE_DROPOUT # FLASHATTENTION_DISABLE_ALIBI - FLASHATTENTION_DISABLE_SOFTCAP - # FLASHATTENTION_DISABLE_UNEVEN_K + # FLASHATTENTION_DISABLE_SOFTCAP + FLASHATTENTION_DISABLE_UNEVEN_K # FLASHATTENTION_DISABLE_LOCAL FLASHATTENTION_DISABLE_PYBIND FLASHATTENTION_VARLEN_ONLY # Custom flag to save on binary size + FLASHATTENTION_PACKGQA_ONLY # Custom flag to save on binary size FLASHATTENTION_DISABLE_CLUSTER # disabled for varlen in any case FLASHATTENTION_DISABLE_SM8x # FLASHATTENTION_DISABLE_HDIMDIFF64 diff --git a/hopper/flash.h b/hopper/flash.h index dcc912249a0..e5c6d6238f6 100644 --- a/hopper/flash.h +++ b/hopper/flash.h @@ -152,9 +152,16 @@ struct Flash_fwd_params : public Qkv_params { int * __restrict__ tile_count_semaphore; // int * __restrict__ num_m_blocks_ptr; + int * __restrict__ prepare_seqlen_q_ptr; // int * __restrict__ num_n_blocks_ptr; int * __restrict__ num_splits_dynamic_ptr; + int * __restrict__ varlen_batch_idx_ptr; // virtual -> actual + int * __restrict__ num_nheads_in_l2_ptr; bool skip_scheduler_metadata_computation; + bool varlen_sort_batches; + int tile_count_semaphore_offset; + bool head_swizzle; + bool prepare_varlen_pdl; int arch; int num_sm; diff --git a/hopper/flash_api.cpp b/hopper/flash_api.cpp index 0d6a0f3a127..9e0baffcd2a 100644 --- a/hopper/flash_api.cpp +++ b/hopper/flash_api.cpp @@ -59,6 +59,8 @@ namespace pybind11::detail { #define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")") #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") +#define PREPARE_VARLEN_MAX_BATCHES_1CTA 992 + namespace { inline at::cuda::CUDAGuard make_cuda_guard_from_tensor(const at::Tensor& t) { return at::cuda::CUDAGuard(static_cast(t.get_device())); @@ -272,7 +274,14 @@ void run_mha_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { PAGEDKV_SWITCH(params.page_table && !params.pagedkv_tma, PagedKVNonTMA, [&] { PACKGQA_SWITCH(params.pack_gqa, PackGQA_, [&] { // Always enable PackGQA for Sm8x or PagedKVNonTMA or Split to reduce compilation + #ifdef FLASHATTENTION_PACKGQA_ONLY + // Always enable PackGQA except for hdimdiff cases to reduce compilation + static constexpr bool PackGQA = true; + static constexpr bool PackGQA_HDIMDIFF = PackGQA_ || Arch < 90 || PagedKVNonTMA || Split; + #else static constexpr bool PackGQA = PackGQA_ || Arch < 90 || PagedKVNonTMA || Split; + static constexpr bool PackGQA_HDIMDIFF = PackGQA_ || Arch < 90 || PagedKVNonTMA || Split; + #endif SOFTCAP_SWITCH(params.softcap > 0.0, Has_softcap, [&] { if (!params.is_e4m3) { if (params.is_bf16) { @@ -280,9 +289,9 @@ void run_mha_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { if (params.d <= 64) { #ifndef FLASHATTENTION_DISABLE_HDIMDIFF64 if (params.dv > 256 && Arch == 90) { - return run_mha_fwd_(params, stream); + return run_mha_fwd_(params, stream); } else if (params.dv > 64 && Arch == 90) { - return run_mha_fwd_(params, stream); + return run_mha_fwd_(params, stream); } else { return run_mha_fwd_(params, stream); } @@ -301,7 +310,7 @@ void run_mha_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { if (params.d <= 192) { #ifndef FLASHATTENTION_DISABLE_HDIMDIFF192 if (params.dv <= 128 && Arch == 90) { - return run_mha_fwd_(params, stream); + return run_mha_fwd_(params, stream); } else { return run_mha_fwd_(params, stream); } @@ -319,9 +328,9 @@ void run_mha_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { if (params.d <= 64) { #ifndef FLASHATTENTION_DISABLE_HDIMDIFF64 if (params.dv > 256 && Arch == 90) { - return run_mha_fwd_(params, stream); + return run_mha_fwd_(params, stream); } else if (params.dv > 64 && Arch == 90) { - return run_mha_fwd_(params, stream); + return run_mha_fwd_(params, stream); } else { return run_mha_fwd_(params, stream); } @@ -340,7 +349,7 @@ void run_mha_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { if (params.d <= 192) { #ifndef FLASHATTENTION_DISABLE_HDIMDIFF192 if (params.dv <= 128 && Arch == 90) { - return run_mha_fwd_(params, stream); + return run_mha_fwd_(params, stream); } else { return run_mha_fwd_(params, stream); } @@ -371,7 +380,7 @@ void run_mha_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { if (params.d <= 192) { #ifndef FLASHATTENTION_DISABLE_HDIMDIFF192 if (params.dv <= 128 && Arch == 90) { - return run_mha_fwd_<90, cutlass::float_e4m3_t, 192, 128, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream); + return run_mha_fwd_<90, cutlass::float_e4m3_t, 192, 128, Split, PagedKVNonTMA, Has_softcap, PackGQA_HDIMDIFF>(params, stream); } else { return run_mha_fwd_<90, cutlass::float_e4m3_t, 192, 192, Split, PagedKVNonTMA, Has_softcap, PackGQA>(params, stream); } @@ -438,15 +447,15 @@ inline bool get_pack_gqa(Flash_fwd_params const& params) { // Always enable PackGQA for Sm8x or PagedKVNonTMA or Split to reduce compilation and binary size. // Has little effect on speed. if (params.arch < 90 || (params.page_table && !params.pagedkv_tma) || params.num_splits > 1) { return true; } + #ifdef FLASHATTENTION_DISABLE_PACKGQA + return false; + #else // Always enable PackGQA for special case of hdim = 64, qheads/kvheads = 8, local attention // TODO: investigate more cases where PackGQA improves perf due to better tile quantization bool const packgqa_override = params.arch >= 90 && (params.h / params.h_k) == 8 && params.is_local && params.d == 64 && (params.dv == params.d); if (packgqa_override) { return true; } - #ifdef FLASHATTENTION_DISABLE_PACKGQA - return false; - #else // params.page_table must already be set if (params.h == params.h_k) { return false; } // This needs to match the kernel configs @@ -620,9 +629,10 @@ mha_fwd_get_scheduler_metadata( params.page_size = page_size.has_value() ? page_size.value() : 1; params.page_table = !page_size.has_value() ? nullptr : reinterpret_cast(1); - // disable dynamic split if given explicit instructions to not split - bool const use_dynamic_split = params.b <= 992 && num_splits != 1; - params.num_splits_dynamic_ptr = !use_dynamic_split ? nullptr : reinterpret_cast(1); + bool const use_prepare_varlen = true; + params.prepare_varlen_pdl = use_prepare_varlen && params.b <= PREPARE_VARLEN_MAX_BATCHES_1CTA; + // set to use in split heuristic + params.num_splits_dynamic_ptr = reinterpret_cast(1); params.pagedkv_tma = get_pagedkv_tma(params); params.num_splits = num_splits <= 0 ? get_num_splits(params) : num_splits; @@ -630,6 +640,11 @@ mha_fwd_get_scheduler_metadata( // Always enable PackGQA for Split params.pack_gqa |= params.num_splits > 1; // printf("Num splits (metadata) = %d.\n", params.num_splits); + #ifdef FLASHATTENTION_PACKGQA_ONLY + params.pack_gqa |= params.d == params.dv; + #endif + + bool const use_dynamic_split = params.b <= PREPARE_VARLEN_MAX_BATCHES_1CTA && params.num_splits > 1; bool is_varlen = true; @@ -641,18 +656,38 @@ mha_fwd_get_scheduler_metadata( // This needs to be set after get_num_splits at::Tensor tile_count_semaphore; // Contains the semaphore and optionally num_splits_dynamic bool const scheduler_needs_semaphore = params.arch >= 90 || params.num_splits > 1; - if (scheduler_needs_semaphore || use_dynamic_split) { - tile_count_semaphore = torch::empty({int(scheduler_needs_semaphore) + int(use_dynamic_split) * params.b}, opts.dtype(torch::kInt32)); + auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; }; + // params.varlen_sort_batches = !params.is_local; // Use this value for Sort in scheduler template + params.varlen_sort_batches = false; + params.head_swizzle = params.is_causal || params.is_local; // Use this value for LPT in scheduler template + if (scheduler_needs_semaphore || use_prepare_varlen) { + int b_rounded = round_multiple(params.b, 4); // for 16 byte alignment of pointers + int num_prepare_batch_vectors = use_prepare_varlen ? 1 : 0; + if (use_dynamic_split) { num_prepare_batch_vectors += 1; } + if (params.varlen_sort_batches) { num_prepare_batch_vectors += 1; } + if (params.head_swizzle) { num_prepare_batch_vectors += 1; } + int sort_offset = b_rounded * (use_dynamic_split ? 2 : 1); + int head_swizzle_offset = b_rounded * (num_prepare_batch_vectors - 1); + int tile_count_semaphore_offset = b_rounded * num_prepare_batch_vectors; + // printf("(Metadata) num prepare batch vectors = %d.\n", num_prepare_batch_vectors); + tile_count_semaphore = torch::empty( + {int(scheduler_needs_semaphore) + tile_count_semaphore_offset}, + opts.dtype(torch::kInt32)); + // ORDER: {prepare_seqlen_q, num_splits_dynamic, varlen_batch_idx, num_nheads_in_l2} + params.prepare_seqlen_q_ptr = use_prepare_varlen ? tile_count_semaphore.data_ptr() : nullptr; + params.num_splits_dynamic_ptr = use_prepare_varlen && use_dynamic_split ? tile_count_semaphore.data_ptr() + b_rounded : nullptr; + params.varlen_batch_idx_ptr = use_prepare_varlen && params.varlen_sort_batches ? tile_count_semaphore.data_ptr() + sort_offset : nullptr; + // params.num_n_blocks_ptr = use_prepare_varlen && params.head_swizzle ? tile_count_semaphore.data_ptr() + head_swizzle_offset : nullptr; + params.num_nheads_in_l2_ptr = use_prepare_varlen && params.head_swizzle ? tile_count_semaphore.data_ptr() + head_swizzle_offset : nullptr; if (scheduler_needs_semaphore) { - if (!use_dynamic_split) { tile_count_semaphore.zero_(); } // If varlen we'll manually do the zero-ing - params.tile_count_semaphore = tile_count_semaphore.data_ptr(); + if (!use_prepare_varlen) { tile_count_semaphore.zero_(); } // If varlen we'll manually do the zero-ing + params.tile_count_semaphore = tile_count_semaphore.data_ptr() + tile_count_semaphore_offset; } else { params.tile_count_semaphore = nullptr; } - params.num_splits_dynamic_ptr = use_dynamic_split ? tile_count_semaphore.data_ptr() + 1 : nullptr; } - if (params.num_splits_dynamic_ptr) { + if (use_prepare_varlen) { auto kBlockMN_kernel_args_sm90 = tile_size_fwd_sm90(params.d_rounded, params.dv_rounded, params.is_causal, params.is_local, params.is_e4m3 ? 1 : 2 /*element_size*/, false /*v_colmajor*/, params.page_table && !params.pagedkv_tma, params.softcap > 0.f, use_one_mma_wg(params)); auto kBlockMN_kernel_args_sm8x = tile_size_fwd_sm8x(params.arch == 86 || params.arch == 89, params.d_rounded, params.dv_rounded, params.is_causal, params.is_local, params.is_e4m3 ? 1 : 2 /*element_size*/, params.page_table, is_varlen && params.num_splits > 1, params.softcap > 0.f, params.knew_ptr); int const kBlockM = params.arch >= 90 ? std::get<0>(kBlockMN_kernel_args_sm90) : std::get<0>(kBlockMN_kernel_args_sm8x); @@ -990,10 +1025,10 @@ mha_fwd(at::Tensor &q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seq } } - // 992 = 32 * 31 is the max supported batch in prepare_varlen_num_blocks kernel - bool const use_dynamic_split = is_varlen && params.b <= 992 && num_splits != 1; + bool const use_prepare_varlen = is_varlen; + params.prepare_varlen_pdl = use_prepare_varlen && params.b <= PREPARE_VARLEN_MAX_BATCHES_1CTA; // Temporarily set num_splits_dynamic_ptr to 1 since get_num_splits checks it - params.num_splits_dynamic_ptr = !use_dynamic_split ? nullptr : reinterpret_cast(1); + params.num_splits_dynamic_ptr = !use_prepare_varlen ? nullptr : reinterpret_cast(1); params.pagedkv_tma = get_pagedkv_tma(params); params.num_splits = num_splits <= 0 ? get_num_splits(params) : num_splits; @@ -1001,6 +1036,13 @@ mha_fwd(at::Tensor &q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seq params.pack_gqa = pack_gqa_.has_value() ? pack_gqa_.value() : get_pack_gqa(params); // Always enable PackGQA for Split params.pack_gqa |= (params.num_splits > 1); + #ifdef FLASHATTENTION_PACKGQA_ONLY + params.pack_gqa |= params.d == params.dv; + #endif + + bool const use_dynamic_split = use_prepare_varlen && params.b <= PREPARE_VARLEN_MAX_BATCHES_1CTA && params.num_splits > 1; + // disable split for varlen and >992 batches for now + if (use_prepare_varlen && params.b > PREPARE_VARLEN_MAX_BATCHES_1CTA) { params.num_splits = 1; } // This needs to be set after get_num_splits at::Tensor tile_count_semaphore; // Contains the semaphore and optionally num_splits_dynamic @@ -1008,8 +1050,20 @@ mha_fwd(at::Tensor &q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seq bool const scheduler_needs_semaphore = params.arch >= 90 ? (((params.is_causal || params.is_local) && (params.num_splits == 1)) || is_varlen) : ((params.is_causal && !is_varlen) || (is_varlen && params.num_splits > 1)); - if (scheduler_needs_semaphore || use_dynamic_split) { - int metadata_size = int(scheduler_needs_semaphore) + int(use_dynamic_split) * params.b; + // params.varlen_sort_batches = !params.is_local; // Use this value for Sort in scheduler template + params.varlen_sort_batches = false; + params.head_swizzle = params.is_causal || params.is_local; // Use this value for LPT in scheduler template + if (scheduler_needs_semaphore || use_prepare_varlen) { + int b_rounded = round_multiple(params.b, 4); // for 16 byte alignment of pointers + int num_prepare_batch_vectors = use_prepare_varlen ? 1 : 0; + if (use_dynamic_split) { num_prepare_batch_vectors += 1; } + if (params.varlen_sort_batches) { num_prepare_batch_vectors += 1; } + if (params.head_swizzle) { num_prepare_batch_vectors += 1; } + int sort_offset = b_rounded * (use_dynamic_split ? 2 : 1); + int head_swizzle_offset = b_rounded * (num_prepare_batch_vectors - 1); + int tile_count_semaphore_offset = b_rounded * num_prepare_batch_vectors; + int metadata_size = int(scheduler_needs_semaphore) + tile_count_semaphore_offset; + // printf("Num prepare batch vectors = %d, metadata_size = %d.\n", num_prepare_batch_vectors, metadata_size); params.skip_scheduler_metadata_computation = scheduler_metadata_.has_value(); if (scheduler_metadata_.has_value()) { at::Tensor scheduler_metadata = scheduler_metadata_.value(); @@ -1021,11 +1075,17 @@ mha_fwd(at::Tensor &q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seq } else { tile_count_semaphore = torch::empty({metadata_size}, opts.dtype(torch::kInt32)); } - if (scheduler_needs_semaphore && !use_dynamic_split) { + if (scheduler_needs_semaphore && !use_prepare_varlen) { tile_count_semaphore.zero_(); // If varlen we'll manually do the zero-ing } - params.tile_count_semaphore = scheduler_needs_semaphore ? tile_count_semaphore.data_ptr() : nullptr; - params.num_splits_dynamic_ptr = use_dynamic_split ? tile_count_semaphore.data_ptr() + 1 : nullptr; + // ORDER: {prepare_seqlen_q, num_splits_dynamic, varlen_batch_idx, num_nheads_in_l2} + params.prepare_seqlen_q_ptr = use_prepare_varlen ? tile_count_semaphore.data_ptr() : nullptr; + params.num_splits_dynamic_ptr = use_prepare_varlen && use_dynamic_split ? tile_count_semaphore.data_ptr() + b_rounded : nullptr; + params.varlen_batch_idx_ptr = use_prepare_varlen && params.varlen_sort_batches ? tile_count_semaphore.data_ptr() + sort_offset : nullptr; + // params.num_n_blocks_ptr = use_prepare_varlen && params.head_swizzle ? tile_count_semaphore.data_ptr() + head_swizzle_offset : nullptr; + params.num_nheads_in_l2_ptr = use_prepare_varlen && params.head_swizzle ? tile_count_semaphore.data_ptr() + head_swizzle_offset : nullptr; + params.tile_count_semaphore = scheduler_needs_semaphore ? tile_count_semaphore.data_ptr() + tile_count_semaphore_offset : nullptr; + params.tile_count_semaphore_offset = tile_count_semaphore_offset; // might need to zero out semaphore later } if (q_v_.has_value()) { @@ -1214,7 +1274,7 @@ mha_fwd(at::Tensor &q, // (b, s_q, h, d) or (total_q, h, d) if there is cu_seq run_mha_fwd_combine(params, stream, true /*enable_pdl*/); } else if (scheduler_needs_semaphore && params.skip_scheduler_metadata_computation) { // need to zero out the semaphore in this case - tile_count_semaphore.index({torch::indexing::Slice(0, 1)}).zero_(); + tile_count_semaphore.index({torch::indexing::Slice(params.tile_count_semaphore_offset, params.tile_count_semaphore_offset + 1)}).zero_(); } } else if (total_q > 0 && num_heads_k > 0) { // If seqlen_k == 0, then we have an empty tensor. We need to set the output to 0. diff --git a/hopper/flash_fwd_combine_kernel.h b/hopper/flash_fwd_combine_kernel.h index 826ccef6f88..cb3f6dbebcb 100644 --- a/hopper/flash_fwd_combine_kernel.h +++ b/hopper/flash_fwd_combine_kernel.h @@ -153,6 +153,7 @@ class FlashAttnFwdCombine { int const* const cu_seqlens = nullptr; int const* const seqused = nullptr; int const* const num_splits_dynamic_ptr = nullptr; + int const* const varlen_batch_idx_ptr = nullptr; int* const semaphore_to_reset = nullptr; }; @@ -173,6 +174,7 @@ class FlashAttnFwdCombine { int const* const cu_seqlens = nullptr; int const* const seqused = nullptr; int const* const num_splits_dynamic_ptr = nullptr; + int const* const varlen_batch_idx_ptr = nullptr; int* const semaphore_to_reset = nullptr; }; @@ -197,6 +199,7 @@ class FlashAttnFwdCombine { args.cu_seqlens, args.seqused, args.num_splits_dynamic_ptr, + args.varlen_batch_idx_ptr, args.semaphore_to_reset }; } @@ -206,9 +209,12 @@ class FlashAttnFwdCombine { int seqlen_q; int total_q; int num_heads; + int num_heads_kv; int dv; + bool pack_gqa; int const* cu_seqlens_q; int const* seqused_q; + int const* prepare_seqlen_q_ptr; }; struct StaticTileScheduler { @@ -253,10 +259,13 @@ class FlashAttnFwdCombine { }; struct Params { - int b; - int num_heads; + int const b; + int const num_heads; + int const num_heads_kv; + bool const pack_gqa; int const* const cu_seqlens_q; int const* const seqused_q; + int const* const prepare_seqlen_q_ptr; SchedulingAlgo algo; }; @@ -281,8 +290,11 @@ class FlashAttnFwdCombine { return { args.b, args.num_heads, + args.num_heads_kv, + args.pack_gqa, args.cu_seqlens_q, args.seqused_q, + args.prepare_seqlen_q_ptr, choose_scheduling_algo(args) }; } @@ -309,7 +321,6 @@ class FlashAttnFwdCombine { } CUTE_DEVICE BlockCoord get_block_coord_linearized_m_and_batch(Params const& params) { - int num_heads = params.num_heads; int curr_tile_id = blockIdx.x; // Scan through the batches find the batch that contains the current @@ -332,8 +343,13 @@ class FlashAttnFwdCombine { auto get_num_m_blocks = [&](int bidb) { if (bidb >= params.b) return 0; - flash::SeqlenInfo seqlen_info{bidb, 0, params.cu_seqlens_q, params.seqused_q}; - return cute::ceil_div(seqlen_info.seqlen * num_heads, Int{}()); + if (params.prepare_seqlen_q_ptr) { + int length = params.prepare_seqlen_q_ptr[bidb] * (!params.pack_gqa ? params.num_heads : params.num_heads_kv); + return cute::ceil_div(length, Int{}); + } else { + flash::SeqlenInfo seqlen_info{bidb, 0, params.cu_seqlens_q, params.seqused_q}; + return cute::ceil_div(seqlen_info.seqlen * params.num_heads, Int{}); + } }; // Cumulative number of blocks for the next 31 batches @@ -408,6 +424,11 @@ class FlashAttnFwdCombine { SharedStorage& shared_storage = *reinterpret_cast(smem_buf); TileScheduler tile_scheduler{shared_storage}; + if (params.semaphore_to_reset && threadIdx.x == 0 && blockIdx.x == gridDim.x - 1 && blockIdx.y == gridDim.y - 1 && blockIdx.z == gridDim.z - 1) { + cutlass::arch::wait_on_dependent_grids(); + *params.semaphore_to_reset = 0; + } + Tensor sLSE = make_tensor(make_smem_ptr(shared_storage.smem_lse_partial.data()), SmemLayoutLSE{}); Tensor sMaxValidSplit = make_tensor(make_smem_ptr(shared_storage.smem_max_valid_split.data()), Shape>{}); Tensor sO = make_tensor(make_smem_ptr(shared_storage.smem_o_partial.data()), SmemLayoutO{}); @@ -418,14 +439,10 @@ class FlashAttnFwdCombine { int const m_block = block_coord.block_m; int const k_block = block_coord.block_k; - int const batch = block_coord.bidb; - - if (params.semaphore_to_reset && threadIdx.x == 0 && blockIdx.x == gridDim.x - 1 && blockIdx.y == gridDim.y - 1 && blockIdx.z == gridDim.z - 1) { - cutlass::arch::wait_on_dependent_grids(); - *params.semaphore_to_reset = 0; - } - - if (batch >= params.b) { return; } + int const maybe_virtual_batch = block_coord.bidb; + if (maybe_virtual_batch >= params.b) { return; } + int const batch = params.varlen_batch_idx_ptr ? params.varlen_batch_idx_ptr[maybe_virtual_batch] : maybe_virtual_batch; + flash::SeqlenInfo seqlen_info{batch, size<0>(params.shape_LSE_partial), params.cu_seqlens, params.seqused}; int const offset = seqlen_info.offset; int const seqlen = seqlen_info.seqlen; @@ -433,7 +450,7 @@ class FlashAttnFwdCombine { if (m_block >= cute::ceil_div(max_idx, Int{})) { return; } - int const num_splits = params.num_splits_dynamic_ptr ? params.num_splits_dynamic_ptr[batch] : get<1>(params.shape_LSE_partial); + int const num_splits = params.num_splits_dynamic_ptr ? params.num_splits_dynamic_ptr[maybe_virtual_batch] : get<1>(params.shape_LSE_partial); if (num_splits <= 1) { return; } cutlass::FastDivmod seqlen_divmod_dynamic(seqlen); diff --git a/hopper/flash_fwd_combine_launch_template.h b/hopper/flash_fwd_combine_launch_template.h index c99efdadfeb..cbaa40f34b3 100644 --- a/hopper/flash_fwd_combine_launch_template.h +++ b/hopper/flash_fwd_combine_launch_template.h @@ -36,12 +36,12 @@ void run_flash_fwd_combine(Flash_fwd_params ¶ms, cudaStream_t stream, bool e {params.o_row_stride, _1{}, params.o_head_stride, !Varlen ? params.o_batch_stride : 0}, // stride_O static_cast(params.softmax_lse_ptr), {_1{}, !Varlen ? params.seqlen_q : params.total_q, !Varlen ? params.h * params.seqlen_q : 0}, // stride_LSE - params.cu_seqlens_q, params.seqused_q, params.num_splits_dynamic_ptr, params.tile_count_semaphore + params.cu_seqlens_q, params.seqused_q, params.num_splits_dynamic_ptr, params.varlen_batch_idx_ptr, params.tile_count_semaphore }; typename CombineKernel::SchedulerArguments scheduler_args { - params.b, params.seqlen_q, params.total_q, params.h, params.dv, - params.cu_seqlens_q, params.seqused_q + params.b, params.seqlen_q, params.total_q, params.h, params.h_k, params.dv, params.pack_gqa, + params.cu_seqlens_q, params.seqused_q, params.prepare_seqlen_q_ptr }; typename CombineKernel::Params kernel_params = { diff --git a/hopper/flash_fwd_kernel_sm90.h b/hopper/flash_fwd_kernel_sm90.h index fd44f33889f..df416386ea4 100644 --- a/hopper/flash_fwd_kernel_sm90.h +++ b/hopper/flash_fwd_kernel_sm90.h @@ -81,8 +81,10 @@ class FlashAttnFwdSm90 { /// Register requirement for Load and Math WGs // If we use cp.async to load K and V, we need more registers for the producer WG. - static constexpr uint32_t LoadRegisterRequirement = NumMmaWarpGroups == 1 ? 56 : (NumMmaWarpGroups == 2 ? (Use_TMA_KV ? 24 : 40) : 32); - static constexpr uint32_t MmaRegisterRequirement = NumMmaWarpGroups == 1 ? 256 : (NumMmaWarpGroups == 2 ? (Use_TMA_KV ? 240 : 232) : 160); + // If we use varlen, we would like more registers for the varlen dynamic persistent scheduler in producer WG, but can only donate sometimes. + static constexpr bool ExtraProducerRegsForVarlen = Varlen && !Is_FP8 && (get<1>(TileShape_MNK_PV{}) == 128 || get<1>(TileShape_MNK_PV{}) == 192); + static constexpr uint32_t LoadRegisterRequirement = NumMmaWarpGroups == 1 ? 56 : (NumMmaWarpGroups == 2 ? (Use_TMA_KV && !ExtraProducerRegsForVarlen ? 24 : 40) : 32); + static constexpr uint32_t MmaRegisterRequirement = NumMmaWarpGroups == 1 ? 256 : (NumMmaWarpGroups == 2 ? (Use_TMA_KV && !ExtraProducerRegsForVarlen ? 240 : 232) : 160); // If you want to print from the producer warp, you'd need to increase the number of registers // Otherwise you'll get CUDA error. // static constexpr uint32_t LoadRegisterRequirement = 40; diff --git a/hopper/flash_fwd_launch_template.h b/hopper/flash_fwd_launch_template.h index 73a128bc986..e8db360b144 100644 --- a/hopper/flash_fwd_launch_template.h +++ b/hopper/flash_fwd_launch_template.h @@ -59,8 +59,11 @@ void run_flash_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { using CollectiveEpilogue = flash::CollectiveEpilogueFwd; static constexpr int NumProducerThreads = Arch >= 90 ? CollectiveMainloop::NumProducerThreads : CollectiveMainloop::NumMmaThreads; + static constexpr bool LPT = Is_causal || Is_local; + // static constexpr bool Sort = !Is_local; + static constexpr bool Sort = false; using SchedulerPersistent = std::conditional_t= 90 /*WarpSpecialized*/>, + flash::VarlenDynamicPersistentTileScheduler= 90 /*WarpSpecialized*/, LPT, Sort, true /*Prepared*/>, std::conditional_t, flash::DynamicPersistentTileScheduler= 90 /*WarpSpecialized*/> @@ -155,14 +158,16 @@ void run_flash_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { params.seqlen_q, params.seqlen_k, params.d, params.dv, sizeof(Element), params.tile_count_semaphore, params.cu_seqlens_q, params.seqused_q, - // params.num_m_blocks_ptr, params.num_splits_dynamic_ptr, + params.prepare_seqlen_q_ptr, + params.varlen_batch_idx_ptr, + params.num_nheads_in_l2_ptr, params.cp_world_size, params.cp_rank, }; - if (Varlen && params.num_splits_dynamic_ptr && !params.skip_scheduler_metadata_computation) { - prepare_varlen_num_blocks(params, stream, PackGQA, kBlockM, kBlockN, Arch >= 90 /*enable_pdl*/); + if (Varlen && !params.skip_scheduler_metadata_computation) { + prepare_varlen_num_blocks(params, stream, PackGQA, kBlockM, kBlockN, Arch >= 90 && params.prepare_varlen_pdl /*enable_pdl*/); CHECK_CUDA_KERNEL_LAUNCH(); } @@ -195,7 +200,7 @@ void run_flash_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { } // kernel<<>>(kernel_params); cutlass::kernel_launch(grid_dims, block_dims, smem_size, stream, kernel_params, - Arch >= 90 && Varlen && params.num_splits_dynamic_ptr && !params.skip_scheduler_metadata_computation /*launch_with_pdl*/); + Arch >= 90 && Varlen && !params.skip_scheduler_metadata_computation && params.prepare_varlen_pdl /*launch_with_pdl*/); } CHECK_CUDA_KERNEL_LAUNCH(); } @@ -211,7 +216,7 @@ void run_mha_fwd_(Flash_fwd_params ¶ms, cudaStream_t stream) { VARLEN_SWITCH(params.cu_seqlens_q || params.cu_seqlens_k || params.seqused_q || params.seqused_k || params.leftpad_k, Varlen, [&] { BOOL_SWITCH(use_one_mma_wg(params), Use_one_mma_wg_, [&] { // Avoid over compiliation by making sure this only get set if it is actually used, i.e. we currently only support one mma wg for 64/128 head dim and hopper - static constexpr bool Use_one_mma_wg = Use_one_mma_wg_ && Arch >= 90 && (kHeadDim == 128 || kHeadDim == 64) && (kHeadDimV == kHeadDim); + static constexpr bool Use_one_mma_wg = Use_one_mma_wg_ && Arch >= 90 && (kHeadDim == 128 || kHeadDim == 64) && (kHeadDimV == kHeadDim) && !Is_local; // Only needed here to decide if we should use cluster static constexpr int kBlockM = Arch >= 90 ? std::get<0>(tile_size_fwd_sm90(kHeadDim, kHeadDimV, Is_causal, Is_local, sizeof(T) /*element_size*/, V_colmajor, PagedKVNonTMA, Has_softcap, Use_one_mma_wg)) : 128; static constexpr bool Enable_cluster = Arch == 90 && (sizeof(T) == 2 ? (kHeadDim >= 128) : (kHeadDim == 192)) && !Is_causal && !Is_local && !Split && !PagedKVNonTMA && !Varlen && !Use_one_mma_wg; @@ -223,8 +228,9 @@ void run_mha_fwd_(Flash_fwd_params ¶ms, cudaStream_t stream) { static constexpr int ClusterM = Enable_cluster && Use_cluster ? 2 : 1; int const qhead_per_khead = !PackGQA ? 1 : cutlass::ceil_div(params.h, params.h_k); PACK_GQA_BLOCK_SWITCH(qhead_per_khead, kBlockH_, [&] { - // TODO: look at pack gqa tma for hdim diff - static constexpr int kBlockH = !PackGQA || Arch < 90 || (kHeadDim != kHeadDimV) ? 1 : kBlockH_; + // Non-unary values of kBlockH can improve GQA perf for specific ratios (4, 8, 16) by enabling TMA for loading Q + // Disable for hdim diff, fp16, 1 mma wg or split to shrink build + static constexpr int kBlockH = !PackGQA || Arch < 90 || (kHeadDim != kHeadDimV) || cute::is_same_v || Use_one_mma_wg || Split ? 1 : kBlockH_; run_flash_fwd(params, stream); }); }); diff --git a/hopper/flash_prepare_scheduler.cu b/hopper/flash_prepare_scheduler.cu index 7093fff32b6..bc0a51cf0ab 100644 --- a/hopper/flash_prepare_scheduler.cu +++ b/hopper/flash_prepare_scheduler.cu @@ -2,6 +2,7 @@ * Copyright (c) 2024, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao. ******************************************************************************/ +#include #include "cutlass/fast_math.h" #include "cutlass/barrier.h" #include "cutlass/arch/barrier.h" @@ -10,8 +11,35 @@ #include "flash.h" +#include "static_switch.h" + namespace flash { +// Sort in descending order +template +struct PrepareSortOp +{ + __device__ __forceinline__ bool operator()(T const & lhs, T const & rhs) + { + return lhs > rhs; + } +}; + +template <> +struct PrepareSortOp { + __device__ __forceinline__ bool operator()(int2 const & lhs, int2 const & rhs) const { + return lhs.x > rhs.x; + } +}; + +template <> +struct PrepareSortOp { + __device__ __forceinline__ bool operator()(int4 const & lhs, int4 const & rhs) const { + return lhs.x > rhs.x; + } +}; + +template __global__ void prepare_varlen_num_blocks_kernel( int seqlen_q_static, int seqlen_k_static, int seqlen_k_new_static, int const* const cu_seqlens_q, int const* const cu_seqlens_k, int const* const cu_seqlens_k_new, @@ -19,16 +47,28 @@ __global__ void prepare_varlen_num_blocks_kernel( int num_batch, int num_head, int qhead_per_khead, int num_sm, int num_splits_static, cutlass::FastDivmod blockm_divmod, cutlass::FastDivmod blockn_divmod, int* const tile_count_semaphore, - // int* const num_m_blocks_ptr, + int* const prepare_seqlen_q_ptr, int* const num_splits_dynamic_ptr, - bool enable_pdl) { + int* const varlen_batch_idx_ptr, + // int* const num_n_blocks_ptr, + int* const num_nheads_in_l2_ptr, + bool enable_pdl, + bool is_causal, + bool packgqa, + int max_kvblocks_in_l2) { static constexpr int kNumBatchPerWarp = cutlass::NumThreadsPerWarp - 1; static constexpr int kSmemSize = 1; - // Assume that there's only one block in the grid + static constexpr int BLOCK_DIM_X = NumWarps * 32; + static constexpr int ITEMS_PER_THREAD = 1; + static_assert(BLOCK_DIM_X * ITEMS_PER_THREAD == NumWarps * 32); + using BlockMergeSort = cub::BlockMergeSort; + __shared__ int total_blocks_smem[kSmemSize]; - // There's only 1 block in the grid, so might as well start launching the main attn kernel + // Allocate shared memory for BlockMergeSort operations + __shared__ typename BlockMergeSort::TempStorage temp_storage; + if (enable_pdl) { cutlass::arch::launch_dependent_grids(); } if (threadIdx.x < kSmemSize) { total_blocks_smem[threadIdx.x] = 0; } @@ -38,8 +78,7 @@ __global__ void prepare_varlen_num_blocks_kernel( int lane = threadIdx.x % cutlass::NumThreadsPerWarp; - auto get_num_m_blocks = [&](int bidb_start) { - int batch_idx = lane + bidb_start; + auto get_num_m_blocks_and_seqlen = [&](int batch_idx) { int seqlen; if (seqused_q) { seqlen = batch_idx < num_batch ? seqused_q[batch_idx] : 0; @@ -50,13 +89,12 @@ __global__ void prepare_varlen_num_blocks_kernel( } else { seqlen = seqlen_q_static; } - seqlen *= qhead_per_khead; return batch_idx < num_batch && lane < kNumBatchPerWarp - ? blockm_divmod.div(seqlen + blockm_divmod.divisor - 1) : 0; + ? cute::make_tuple(blockm_divmod.div(seqlen * (packgqa ? qhead_per_khead : 1) + blockm_divmod.divisor - 1), seqlen) + : cute::make_tuple(0, 0); }; - auto get_num_n_blocks = [&](int bidb_start) { - int batch_idx = lane + bidb_start; + auto get_num_n_blocks = [&](int batch_idx) { int leftpad_k = batch_idx < num_batch && leftpad_k_ptr != nullptr ? leftpad_k_ptr[batch_idx] : 0; int seqlen; if (seqused_k) { @@ -83,42 +121,130 @@ __global__ void prepare_varlen_num_blocks_kernel( }; int warp_idx = threadIdx.x / cutlass::NumThreadsPerWarp; - int bidb_start = kNumBatchPerWarp * warp_idx; - int num_m_blocks = get_num_m_blocks(bidb_start); - int num_n_blocks = get_num_n_blocks(bidb_start); - - int total_blocks = num_m_blocks * num_n_blocks; - // Warp sum - #pragma unroll - for (int i = cutlass::NumThreadsPerWarp / 2; i >= 1; i /= 2) { - total_blocks += __shfl_down_sync(0xffffffff, total_blocks, i); + int batch_cta_idx_offset = int(blockIdx.x) * 992; + int bidb_start = batch_cta_idx_offset + kNumBatchPerWarp * warp_idx; + int batch_idx = lane + bidb_start; + // int num_m_blocks = get_num_m_blocks(batch_idx); + auto seqlen_q_info = get_num_m_blocks_and_seqlen(batch_idx); + int num_m_blocks = cute::get<0>(seqlen_q_info); + int seqlen_q = cute::get<1>(seqlen_q_info); + int num_n_blocks = get_num_n_blocks(batch_idx); + + auto get_nheads_in_l2 = [&](int n_blocks) { + int nheads_in_l2 = n_blocks * 16 <= max_kvblocks_in_l2 ? 16 + : n_blocks * 8 <= max_kvblocks_in_l2 ? 8 + : n_blocks * 4 <= max_kvblocks_in_l2 ? 4 + : n_blocks * 2 <= max_kvblocks_in_l2 ? 2 + : 1; + if(!packgqa) { nheads_in_l2 *= qhead_per_khead; } + return min(nheads_in_l2, num_head); + }; + + int num_splits_dynamic; + if (int(gridDim.x) > 1 || num_splits_static == 1) { + // set num splits for all batches to 1 (note that user expects num_splits_static to mean upper bound on splits) + // for batch size > 992, we expect GPU occupancy to not be an issue except in degenerate cases (e.g., most are zero-length) + num_splits_dynamic = 1; + } else { + int total_blocks = num_m_blocks * num_n_blocks; + // Warp sum + #pragma unroll + for (int i = cutlass::NumThreadsPerWarp / 2; i >= 1; i /= 2) { + total_blocks += __shfl_down_sync(0xffffffff, total_blocks, i); + } + if (lane == 0) { atomicAdd(total_blocks_smem, total_blocks); } + __syncthreads(); + total_blocks = total_blocks_smem[0]; + // 10% margin + int blocks_per_sm = static_cast(ceilf(float(total_blocks) * 1.1f * float(num_head) / float(num_sm))); + // blocks_per_sm = std::max(1, blocks_per_sm); // 1 is the minimum number of blocks per SM + num_splits_dynamic = std::max(std::min((num_n_blocks + blocks_per_sm - 1) / blocks_per_sm, num_splits_static), 1); + // num_n_blocks per work tile for the batch + num_n_blocks = cutlass::ceil_div(num_n_blocks, num_splits_dynamic); } - if (lane == 0) { atomicAdd(total_blocks_smem, total_blocks); } - __syncthreads(); - total_blocks = total_blocks_smem[0]; - // 10% margin - int blocks_per_sm = static_cast(ceilf(float(total_blocks) * 1.1f * float(num_head) / float(num_sm))); - // blocks_per_sm = std::max(1, blocks_per_sm); // 1 is the minimum number of blocks per SM - int num_splits_dynamic = std::max(std::min((num_n_blocks + blocks_per_sm - 1) / blocks_per_sm, num_splits_static), 1); - if (bidb_start + lane < num_batch && lane < kNumBatchPerWarp) { - num_splits_dynamic_ptr[bidb_start + lane] = num_splits_dynamic; - // printf("idx = %d, num_m_blocks = %d, num_n_blocks = %d, num_split_static = %d, num_splits_dynamic = %d\n", bidb_start + lane, num_m_blocks_ptr[bidb_start + lane], num_n_blocks, num_splits_static, num_splits_dynamic); + + if constexpr (Sort) { + if(lane == kNumBatchPerWarp || batch_idx >= num_batch) { + num_n_blocks = INT_MIN; // sort last + } + else if (is_causal) { + // sort by middle member to process + num_n_blocks = num_n_blocks * blockn_divmod.divisor - (seqlen_q / 2); + // sort by shortest member to process + // num_n_blocks = num_n_blocks * blockn_divmod.divisor - seqlen_q; + } + int4 batch_coords[ITEMS_PER_THREAD]; // 1 item per thread + batch_coords[0] = make_int4(num_n_blocks, seqlen_q, num_splits_dynamic, batch_idx); + + // Sort batches by num_n_blocks in descending order + BlockMergeSort(temp_storage).Sort(batch_coords, PrepareSortOp()); + + if (is_causal) { + // reset value to num_n_blocks + batch_coords[0].x = blockn_divmod.div(batch_coords[0].x + (batch_coords[0].y / 2)); + // batch_coords[0].x = blockn_divmod.div(batch_coords[0].x + batch_coords[0].y); + } + + // When sorting, we re-index some metadata by 'virtual batch index' + // and also store the vbidx -> bidx mapping. + // 1. num_nheads_in_l2_ptr: virtual_batch_idx -> num_nheads_in_l2[batch_idx] + // 2. num_splits_dynamic_ptr: virtual_batch_idx -> num_splits[batch_idx] + // 3. prepare_seqlen_q_ptr: virtual_batch_idx -> seqlen_q[batch_idx] * (packgqa ? qhead_per_khead : 1) + // 4. varlen_batch_idx_ptr: virtual_batch_idx -> batch_idx + batch_idx = batch_cta_idx_offset + threadIdx.x; + if (batch_idx < num_batch && threadIdx.x < 992) { + if(num_nheads_in_l2_ptr) { num_nheads_in_l2_ptr[batch_idx] = get_nheads_in_l2(max(batch_coords[0].x, 1)); } + prepare_seqlen_q_ptr[batch_idx] = batch_coords[0].y * (packgqa ? qhead_per_khead : 1); + if(num_splits_dynamic_ptr) { num_splits_dynamic_ptr[batch_idx] = batch_coords[0].z; } + varlen_batch_idx_ptr[batch_idx] = batch_coords[0].w; + } + } else { + if (batch_idx < num_batch && lane < kNumBatchPerWarp) { + prepare_seqlen_q_ptr[batch_idx] = seqlen_q * (packgqa ? qhead_per_khead : 1); + if(num_splits_dynamic_ptr) { num_splits_dynamic_ptr[batch_idx] = num_splits_dynamic; } + if(num_nheads_in_l2_ptr) { num_nheads_in_l2_ptr[batch_idx] = get_nheads_in_l2(max(num_n_blocks, 1)); } + // printf("idx = %d, num_m_blocks = %d, num_n_blocks = %d, num_split_static = %d, num_splits_dynamic = %d\n", bidb_start + lane, num_m_blocks_ptr[bidb_start + lane], num_n_blocks, num_splits_static, num_splits_dynamic); + } } + } } // flash void prepare_varlen_num_blocks(Flash_fwd_params ¶ms, cudaStream_t stream, bool packgqa, int blockM, int blockN, bool enable_pdl) { - // Only support batch <= 992 (32 warps, each with 31 batches) - int qhead_per_khead = !packgqa ? 1 : cutlass::ceil_div(params.h, params.h_k); - flash::prepare_varlen_num_blocks_kernel<<<1 /*grid*/, 1024 /*block*/, 0, stream>>>( - params.seqlen_q, params.seqlen_k, params.seqlen_knew, - params.cu_seqlens_q, params.cu_seqlens_k, params.cu_seqlens_knew, - params.seqused_q, params.seqused_k, params.leftpad_k, - params.b, !packgqa ? params.h : params.h_k, qhead_per_khead, params.num_sm, params.num_splits, - cutlass::FastDivmod(blockM), cutlass::FastDivmod(blockN), - params.tile_count_semaphore, - // params.num_m_blocks_ptr, - params.num_splits_dynamic_ptr, enable_pdl); + int qhead_per_khead = cutlass::ceil_div(params.h, params.h_k); + int num_warps = cutlass::ceil_div(params.b, 31); // warp switch will cap this at 32 + int num_ctas = cutlass::ceil_div(params.b, 31 * 32); + // int const size_l2 = 50 * 1024 * 1024; // 50 MB + int const size_l2_divisor = qhead_per_khead == 1 ? 1 + : qhead_per_khead <= 2 ? 2 + : qhead_per_khead <= 4 ? 4 + : qhead_per_khead <= 8 ? 8 + : 16; + int const size_l2 = (32 * 1024 * 1024) / size_l2_divisor; // experimental + int const element_size = params.is_e4m3 ? 1 : 2; + int const size_one_kvblock = blockN * (params.d + params.dv) * element_size; + // printf("block size = %d, element size = %d, headdim = %d, headdim_v = %d, size 1 kblock = %d.\n", blockN, element_size, params.d, params.dv, size_one_kvblock); + int const max_kvblocks_in_l2 = size_l2 / size_one_kvblock; + BOOL_SWITCH(params.varlen_sort_batches, Sort, [&] { + NUM_WARP_SWITCH(num_warps, NumWarps, [&] { + flash::prepare_varlen_num_blocks_kernel<<>>( + params.seqlen_q, params.seqlen_k, params.seqlen_knew, + params.cu_seqlens_q, params.cu_seqlens_k, params.cu_seqlens_knew, + params.seqused_q, params.seqused_k, params.leftpad_k, + params.b, !packgqa ? params.h : params.h_k, qhead_per_khead, params.num_sm, params.num_splits, + cutlass::FastDivmod(blockM), cutlass::FastDivmod(blockN), + params.tile_count_semaphore, + params.prepare_seqlen_q_ptr, + params.num_splits_dynamic_ptr, + params.varlen_batch_idx_ptr, + // params.num_n_blocks_ptr, + params.num_nheads_in_l2_ptr, + enable_pdl, + params.is_causal, + packgqa, + max_kvblocks_in_l2); + }); + }); } diff --git a/hopper/heuristics.h b/hopper/heuristics.h index 621dc16bc34..f8f02710bd7 100644 --- a/hopper/heuristics.h +++ b/hopper/heuristics.h @@ -9,7 +9,7 @@ inline bool use_one_mma_wg(Flash_fwd_params const& params) { // assume pack_gqa for seqlen calculation - return params.arch >= 90 && (params.d == 128 || params.d == 64) && + return params.arch >= 90 && (params.d == 128 || params.d == 64) && !params.is_local && params.seqlen_q * (params.h / params.h_k) <= 64; }; diff --git a/hopper/setup.py b/hopper/setup.py index 7d1776f565d..887b6339023 100644 --- a/hopper/setup.py +++ b/hopper/setup.py @@ -68,7 +68,9 @@ DISABLE_HDIMDIFF64 = os.getenv("FLASH_ATTENTION_DISABLE_HDIMDIFF64", "FALSE") == "TRUE" DISABLE_HDIMDIFF192 = os.getenv("FLASH_ATTENTION_DISABLE_HDIMDIFF192", "FALSE") == "TRUE" -# DISABLE_BACKWARD = True +PACKGQA_ONLY = os.getenv("FLASH_ATTENTION_PACKGQA_ONLY", "FALSE") == "TRUE" + +DISABLE_BACKWARD = True # DISABLE_SPLIT = True # DISABLE_PAGEDKV = True # DISABLE_APPENDKV = True @@ -78,17 +80,19 @@ # DISABLE_FP16 = True # DISABLE_FP8 = True # DISABLE_VARLEN = True -# DISABLE_CLUSTER = True +DISABLE_CLUSTER = True # DISABLE_HDIM64 = True # DISABLE_HDIM96 = True # DISABLE_HDIM128 = True # DISABLE_HDIM192 = True # DISABLE_HDIM256 = True -# DISABLE_SM8x = True +DISABLE_SM8x = True -# DISABLE_HDIMDIFF64 = True +DISABLE_HDIMDIFF64 = True # DISABLE_HDIMDIFF192 = True +PACKGQA_ONLY = True + # HACK: we monkey patch pytorch's _write_ninja_file to pass # "-gencode arch=compute_sm90a,code=sm_90a" to files ending in '_sm90.cu', # and pass "-gencode arch=compute_sm80,code=sm_80" to files ending in '_sm80.cu' @@ -494,6 +498,8 @@ def nvcc_threads_args(): + (["-DFLASHATTENTION_ENABLE_VCOLMAJOR"] if ENABLE_VCOLMAJOR else []) + (["-DFLASHATTENTION_DISABLE_HDIMDIFF64"] if DISABLE_HDIMDIFF64 else []) + (["-DFLASHATTENTION_DISABLE_HDIMDIFF192"] if DISABLE_HDIMDIFF192 else []) + + (["-DFLASHATTENTION_PACKGQA_ONLY"] if PACKGQA_ONLY else []) + + (["-DFLASHATTENTION_VARLEN_ONLY"]) ) DTYPE_FWD_SM80 = ["bf16"] + (["fp16"] if not DISABLE_FP16 else []) @@ -539,10 +545,15 @@ def nvcc_threads_args(): # We already always hard-code PackGQA=true for Sm8x sources_fwd_sm80 = [f"instantiations/flash_fwd_hdim{hdim}_{dtype}{paged}{split}{softcap}_sm80.cu" for hdim, dtype, split, paged, softcap in itertools.product(HEAD_DIMENSIONS_FWD_SM80, DTYPE_FWD_SM80, SPLIT, PAGEDKV, SOFTCAP_ALL)] - # We already always hard-code PackGQA=true for Sm9x if PagedKV or Split + # We already always hard-code PackGQA=true for Sm9x + # "not (packgqa and (paged or split))" is default setting since PackGQA=true if PagedKV or Split + # We then add "(packgqa or paged or split)" to enable PackGQA always sources_fwd_sm90 = [f"instantiations/flash_fwd_hdim{hdim}_{dtype}{paged}{split}{softcap}{packgqa}_sm90.cu" for hdim, dtype, split, paged, softcap, packgqa in itertools.product(HEAD_DIMENSIONS_FWD, DTYPE_FWD_SM90, SPLIT, PAGEDKV, SOFTCAP, PACKGQA) - if not (packgqa and (paged or split))] + if not (packgqa and (paged or split)) and (not PACKGQA_ONLY or packgqa or paged or split)] + # sources_fwd_sm90 = [f"instantiations/flash_fwd_hdim{hdim}_{dtype}{paged}{split}{softcap}{packgqa}_sm90.cu" + # for hdim, dtype, split, paged, softcap, packgqa in itertools.product(HEAD_DIMENSIONS_FWD, DTYPE_FWD_SM90, SPLIT, PAGEDKV, SOFTCAP, PACKGQA) + # if not (packgqa and (paged or split))] if not DISABLE_HDIMDIFF64: sources_fwd_sm90 += [f"instantiations/flash_fwd_hdim{hdim}_{dtype}{paged}{split}{softcap}{packgqa}_sm90.cu" for hdim, dtype, split, paged, softcap, packgqa in itertools.product(HEAD_DIMENSIONS_DIFF64_FWD, HALF_DTYPE_FWD_SM90, SPLIT, PAGEDKV, SOFTCAP, PACKGQA) diff --git a/hopper/static_switch.h b/hopper/static_switch.h index 45bf0cf771e..4d2f3dec0fb 100644 --- a/hopper/static_switch.h +++ b/hopper/static_switch.h @@ -210,12 +210,41 @@ #else #define PACK_GQA_BLOCK_SWITCH(QHEADS_PER_KHEADS, BLOCK_H, ...) \ [&] { \ - if (QHEADS_PER_KHEADS == 8) { \ + if (QHEADS_PER_KHEADS == 16) { \ + constexpr static int BLOCK_H = 16; \ + return __VA_ARGS__(); \ + } else if (QHEADS_PER_KHEADS == 8) { \ constexpr static int BLOCK_H = 8; \ return __VA_ARGS__(); \ + } else if (QHEADS_PER_KHEADS == 4) { \ + constexpr static int BLOCK_H = 4; \ + return __VA_ARGS__(); \ } else { \ constexpr static int BLOCK_H = 1; \ return __VA_ARGS__(); \ } \ }() -#endif \ No newline at end of file +#endif + +#define NUM_WARP_SWITCH(VALUE, CONST_NAME, ...) \ + [&] { \ + if (VALUE <= 1) { \ + constexpr static int CONST_NAME = 1; \ + return __VA_ARGS__(); \ + } else if (VALUE <= 2) { \ + constexpr static int CONST_NAME = 2; \ + return __VA_ARGS__(); \ + } else if (VALUE <= 4) { \ + constexpr static int CONST_NAME = 4; \ + return __VA_ARGS__(); \ + } else if (VALUE <= 8) { \ + constexpr static int CONST_NAME = 8; \ + return __VA_ARGS__(); \ + } else if (VALUE <= 16) { \ + constexpr static int CONST_NAME = 16; \ + return __VA_ARGS__(); \ + } else { \ + constexpr static int CONST_NAME = 32; \ + return __VA_ARGS__(); \ + } \ + }() diff --git a/hopper/test_flash_attn.py b/hopper/test_flash_attn.py index 0b1f58ce70b..cbb7f70089b 100644 --- a/hopper/test_flash_attn.py +++ b/hopper/test_flash_attn.py @@ -398,12 +398,9 @@ def test_flash_attn_varlen_output( device = "cuda" # set seed torch.random.manual_seed(seqlen_q + seqlen_k + d + int(causal) * 2 + int(local)) - # batch_size = 40 - # nheads = 16 + # batch_size = 1024 batch_size = 9 if seqlen_q <= 2048 else 2 nheads = 6 - # batch_size = 2 - # nheads = 1 nheads_kv = nheads if mha_type == "mha" else (2 if mha_type == "gqa" else 1) dtype_ref = torch.bfloat16 if dtype == torch.float8_e4m3fn else dtype if d == 192 and not DISABLE_HDIMDIFF192: @@ -546,6 +543,8 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): k_descale=k_descale, v_descale=v_descale, window_size=window_size, softcap=softcap, + num_splits=num_splits, + pack_gqa=pack_gqa, s_aux=s_aux, ) print("Pack GQA =",pack_gqa) @@ -648,10 +647,10 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): @pytest.mark.parametrize("mha_type", ["mha", "mqa", "gqa"]) # @pytest.mark.parametrize("mha_type", ["mha"]) @pytest.mark.parametrize("new_kv", [False] + ([True] if not DISABLE_APPENDKV else [])) -# @pytest.mark.parametrize("new_kv", [False]) +# @pytest.mark.parametrize("new_kv", [True]) @pytest.mark.parametrize("causal,local", [(False, False), (True, False)] + ([(False, True)] if not DISABLE_LOCAL else [])) # @pytest.mark.parametrize("causal,local", [(False, False), (True, False)]) -# @pytest.mark.parametrize("causal,local", [(False, True)]) +# @pytest.mark.parametrize("causal,local", [(True, False)]) @pytest.mark.parametrize("seqlen_new_eq_seqlen_q", [True, False] if not DISABLE_APPENDKV else [True]) # @pytest.mark.parametrize("seqlen_new_eq_seqlen_q", [True]) @pytest.mark.parametrize("has_rotary_seqlens", [False, True]) @@ -661,7 +660,7 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): @pytest.mark.parametrize("rotary_fraction", [0.0, 0.5, 1.0] if (not DISABLE_APPENDKV) and (apply_rotary_emb is not None) else [0.0]) # @pytest.mark.parametrize("rotary_fraction", [0.0]) @pytest.mark.parametrize("page_size", [None] + ([1, 4, 128] if not DISABLE_PAGEDKV else [])) -# @pytest.mark.parametrize("page_size", [4]) +# @pytest.mark.parametrize("page_size", [None]) @pytest.mark.parametrize("has_leftpad", [False, True]) # @pytest.mark.parametrize("has_leftpad", [False]) @pytest.mark.parametrize("has_batch_idx", [False, True]) @@ -673,7 +672,7 @@ def _gen_unused_masks(padding_mask, add_unused, max_seq_len, bs, device): # @pytest.mark.parametrize('d', [32, 40, 64, 80, 96, 128, 160, 192]) # @pytest.mark.parametrize('d', [56, 80]) # @pytest.mark.parametrize("d", COMPILED_HDIMS) -@pytest.mark.parametrize("d", [64]) +@pytest.mark.parametrize("d", [128]) # @pytest.mark.parametrize("d", [192]) # @pytest.mark.parametrize("test_sink", [False, True]) @pytest.mark.parametrize("test_sink", [False]) @@ -928,7 +927,6 @@ def test_flash_attn_kvcache( v_cache_saved = v_cache.clone() if page_size is None else v_cache_paged.clone() num_splits_vals = [1, 3, 0] if not DISABLE_SPLIT else [1] precompute_metadata_vals = [False, True] - # precompute_metadata_vals = [False] for num_splits, precompute_metadata in itertools.product(num_splits_vals, precompute_metadata_vals): print("Num splits = ",num_splits) print("Precompute metadata = ",precompute_metadata) @@ -936,7 +934,10 @@ def test_flash_attn_kvcache( if precompute_metadata: # WARNING: seqlen_k is not max_seqlen_k if using page table, so we can't expect this to make sense? scheduler_metadata = get_scheduler_metadata( - batch_size, max_seqlen_q if varlen_q else seqlen_q, seqlen_k, nheads, nheads_k, d, + batch_size, + max_seqlen_q if varlen_q else seqlen_q, + seqlen_k if page_size is None else page_table.shape[1] * page_size, + nheads, nheads_k, d, cache_seqlens, q.dtype, headdim_v=dv, cu_seqlens_q=cu_seqlens_q, cu_seqlens_k_new=cu_seqlens_k_new, cache_leftpad=cache_leftpad, max_seqlen_k_new=seqlen_new, page_size=page_size, diff --git a/hopper/tile_scheduler.hpp b/hopper/tile_scheduler.hpp index c147725dd8c..1109de2b8a4 100644 --- a/hopper/tile_scheduler.hpp +++ b/hopper/tile_scheduler.hpp @@ -24,8 +24,11 @@ struct TileSchedulerArguments { int* const tile_count_semaphore = nullptr; int const* const cu_seqlens = nullptr; int const* const seqused = nullptr; - // int const* const num_m_blocks_ptr = nullptr; int const* const num_splits_dynamic_ptr = nullptr; + int const* const prepare_seqlen_q_ptr = nullptr; + int const* const varlen_batch_idx_ptr = nullptr; + // int const* const num_n_blocks_ptr = nullptr; + int const* const num_nheads_in_l2_ptr = nullptr; // CP (Context Parallelism) parameters int const cp_world_size = 1; int const cp_rank = 0; @@ -361,7 +364,8 @@ class DynamicPersistentTileScheduler { }; -template +template class VarlenDynamicPersistentTileScheduler { static_assert(WarpSpecialized || NumProducerThreads == NumMmaThreads); @@ -380,13 +384,17 @@ class VarlenDynamicPersistentTileScheduler { int num_head, num_batch; int const qhead_per_khead; int const seqlen; + // int const max_kvblocks_in_l2; cutlass::FastDivmod head_divmod; cutlass::FastDivmod nsplits_divmod; int* const tile_count_semaphore; int const* const cu_seqlens; int const* const seqused; - // int* const num_m_blocks_ptr; int const* const num_splits_dynamic_ptr; + int const* const prepare_seqlen_q_ptr; + int const* const varlen_batch_idx_ptr; + // int const* const num_n_blocks_ptr; + int const* const num_nheads_in_l2_ptr; }; static Params @@ -396,13 +404,20 @@ class VarlenDynamicPersistentTileScheduler { assert(args.tile_count_semaphore != nullptr); assert(args.num_head < (1 << 16)); // We use the top 16 bits to store num_splits & split_idx assert(!Split || args.num_splits < (1 << 8)); // We use the top 8 bits to store num_splits + // int const size_l2 = 50 * 1024 * 1024; // 50 MB + // int const size_one_kvblock = kBlockN * (args.headdim + args.headdim_v) * args.element_size; + // int max_kvblocks_in_l2 = size_l2 / size_one_kvblock; return {args.num_head, args.num_batch, args.qhead_per_khead, args.seqlen, + // max_kvblocks_in_l2, cutlass::FastDivmod(args.num_head), cutlass::FastDivmod(!Split ? 1 : args.num_splits), args.tile_count_semaphore, args.cu_seqlens, args.seqused, - // args.num_m_blocks_ptr, - args.num_splits_dynamic_ptr}; + args.num_splits_dynamic_ptr, + args.prepare_seqlen_q_ptr, + args.varlen_batch_idx_ptr, + // aras.num_n_blocks_ptr, + args.num_nheads_in_l2_ptr}; } static dim3 @@ -423,8 +438,15 @@ class VarlenDynamicPersistentTileScheduler { CUTLASS_DEVICE cute::tuple get_block_coord(Params const& params) const { + auto get_actual_batch = [&](int virtual_batch) { + if constexpr(Prepared && Sort) { + return params.varlen_batch_idx_ptr[virtual_batch]; + } else { + return virtual_batch; + } + }; if constexpr (!Split) { - return {block, bidh, bidb, 0 /*split_idx*/}; + return {block, bidh, get_actual_batch(bidb), 0 /*split_idx*/}; } else { // the top 8 bits of bidh store num_splits and the next 8 bits store split_idx // reinterpret_cast to uint32_t to make sure we're not doing sign extension when we shift @@ -438,7 +460,7 @@ class VarlenDynamicPersistentTileScheduler { // if (threadIdx.x == 128) { // printf("blockIdx.x = %d, bidb = %d, bidh = %d, bidh_actual = %d, split_idx = %d\n", blockIdx.x, bidb, bidh, bidh_actual, split_idx); // } - return {block, bidh_actual, bidb, split_idx}; + return {block, bidh_actual, get_actual_batch(bidb), split_idx}; } } }; @@ -452,31 +474,38 @@ class VarlenDynamicPersistentTileScheduler { int lane = threadIdx.x % cutlass::NumThreadsPerWarp; auto get_num_m_blocks = [&] (int bidb_start) { int batch_idx = lane + bidb_start; - int seqlen = params.seqlen * (!PackGQA ? 1 : params.qhead_per_khead); - if (seqlen > kBlock) { - if (params.seqused) { - seqlen = batch_idx < params.num_batch ? params.seqused[batch_idx] : 0; - } else if (params.cu_seqlens) { - int cur_cu_seqlen = batch_idx <= params.num_batch ? params.cu_seqlens[batch_idx] : 0; - int next_cu_seqlen = __shfl_down_sync(0xffffffff, cur_cu_seqlen, 1); - seqlen = next_cu_seqlen - cur_cu_seqlen; - } else { - seqlen = params.seqlen; + if constexpr (Prepared) { + return batch_idx < params.num_batch && lane < cutlass::NumThreadsPerWarp - 1 + ? cute::ceil_div(params.prepare_seqlen_q_ptr[batch_idx], kBlockM) : 0; + } else { + int seqlen = params.seqlen * (!PackGQA ? 1 : params.qhead_per_khead); + if (seqlen > kBlockM) { + if (params.seqused) { + seqlen = batch_idx < params.num_batch ? params.seqused[batch_idx] : 0; + } else if (params.cu_seqlens) { + int cur_cu_seqlen = batch_idx <= params.num_batch ? params.cu_seqlens[batch_idx] : 0; + int next_cu_seqlen = __shfl_down_sync(0xffffffff, cur_cu_seqlen, 1); + seqlen = next_cu_seqlen - cur_cu_seqlen; + } else { + seqlen = params.seqlen; + } + if constexpr (PackGQA) { seqlen *= params.qhead_per_khead; } } - if constexpr (PackGQA) { seqlen *= params.qhead_per_khead; } + return batch_idx < params.num_batch && lane < cutlass::NumThreadsPerWarp - 1 + ? cute::ceil_div(seqlen, kBlockM) : 0; } - return batch_idx < params.num_batch && lane < cutlass::NumThreadsPerWarp - 1 - ? cute::ceil_div(seqlen, kBlock) : 0; - // ? params.num_m_blocks_ptr[batch_idx] : 0; }; auto get_num_splits = [&] (int bidb_start) { int batch_idx = lane + bidb_start; - return batch_idx < params.num_batch && lane < cutlass::NumThreadsPerWarp - 1 - ? (!Split ? 1 : (params.num_splits_dynamic_ptr - ? params.num_splits_dynamic_ptr[batch_idx] - : params.nsplits_divmod.divisor)) - : 0; + bool is_valid = batch_idx < params.num_batch && lane < cutlass::NumThreadsPerWarp - 1; + if constexpr (!Split) { + return is_valid ? 1 : 0; + } else if constexpr(Prepared) { + return is_valid ? params.num_splits_dynamic_ptr[batch_idx] : 0; + } else { + return is_valid ? params.nsplits_divmod.divisor : 0; + } }; int num_m_blocks = get_num_m_blocks(current_work.bidb); // Different for each lane @@ -487,12 +516,14 @@ class VarlenDynamicPersistentTileScheduler { // Total number of blocks for the next 31 batches int m_blocks_in_group = __shfl_sync(0xffffffff, num_m_blocks_cumulative, cutlass::NumThreadsPerWarp - 1); // Only the lower 16 bits are the actual bidh - int current_bidh = !Split ? current_work.bidh : (current_work.bidh & 0x0000FFFF); - int group_end_tile = current_work.tile_idx - current_work.block - current_bidh * __shfl_sync(0xffffffff, num_split_m_blocks, 0 /*lane*/) + m_blocks_in_group * params.num_head; // Same for all lanes - if constexpr (Split) { - int current_split_idx = (current_work.bidh & 0x00FF0000) >> 16; - group_end_tile -= current_split_idx * __shfl_sync(0xffffffff, num_m_blocks, 0 /*lane*/); - } + // int current_bidh = !Split ? current_work.bidh : (current_work.bidh & 0x0000FFFF); + // int group_end_tile = current_work.tile_idx - current_work.block - current_bidh * __shfl_sync(0xffffffff, num_split_m_blocks, 0 /*lane*/) + m_blocks_in_group * params.num_head; // Same for all lanes + // if constexpr (Split) { + // int current_split_idx = (current_work.bidh & 0x00FF0000) >> 16; + // group_end_tile -= current_split_idx * __shfl_sync(0xffffffff, num_m_blocks, 0 /*lane*/); + // } + // NEW: current_work.tile_idx holds group_start_tile for starting batch + int group_end_tile = current_work.tile_idx + m_blocks_in_group * params.num_head; // Same for all lanes int bidb = current_work.bidb; // if (blockIdx.x <= 9 && threadIdx.x == 0) { // printf("Before while, blockIdx.x = %d, threadIdx.x = %d, bidb = %d, num_m_blocks = %d, next_tile_idx = %d, cur tile_idx = %d, cur block = %d, cur bidh = %d, num_split_m_blocks = %d, group_end_tile = %d, m_blocks_in_group = %d\n", blockIdx.x, threadIdx.x, current_work.bidb, num_m_blocks, next_tile_idx, current_work.tile_idx, current_work.block, current_bidh, num_split_m_blocks, group_end_tile, m_blocks_in_group); @@ -524,27 +555,81 @@ class VarlenDynamicPersistentTileScheduler { bidb += batch_idx_in_group; num_m_blocks = __shfl_sync(0xffffffff, num_m_blocks, batch_idx_in_group); if constexpr (Split) { num_splits = __shfl_sync(0xffffffff, num_splits, batch_idx_in_group); } - int mh_block = next_tile_idx - group_start_tile - (batch_idx_in_group == 0 ? 0 : __shfl_sync(0xffffffff, num_m_blocks_cumulative, batch_idx_in_group - 1)) * params.num_head; - int bidh = mh_block / num_m_blocks; - int block = mh_block - bidh * num_m_blocks; - if constexpr (Split) { - int bidh_actual = bidh / num_splits; - int split_idx = bidh - bidh_actual * num_splits; - // TODO: idk why this gives wrong answer nondeterministically - // int bidh_actual, split_idx; - // split_idx = params.head_divmod.divmod(bidh_actual, bidh); - // Use the top 8 bits to store num_splits and the next 8 bits to store split_idx - // reinterpret_cast to uint32_t to make sure we're not doing sign extension when we shift - uint32_t bidh_packed = reinterpret_cast(bidh_actual) + (reinterpret_cast(split_idx) << 16) + (reinterpret_cast(num_splits) << 24); - // if (threadIdx.x == 0) { - // printf("blockIdx.x = %d, group_start_tiled = %d, bidb = %d, batch_idx_in_group = %d, mh_block = %d, num_m_blocks = %d, bidh = %d, bidh_actual = %d, split_idx = %d, num_splits = %d, bidh_packed = %d\n", blockIdx.x, group_start_tile, bidb, batch_idx_in_group, mh_block, num_m_blocks, bidh, bidh_actual, split_idx, num_splits, bidh_packed); + group_start_tile += (batch_idx_in_group == 0 ? 0 : __shfl_sync(0xffffffff, num_m_blocks_cumulative, batch_idx_in_group - 1)) * params.num_head; + int mh_block = next_tile_idx - group_start_tile; + int block, bidh; + if constexpr (LPT) { + if (!Split || num_splits == 1) { + // NOTE: code for computing nheads_in_l2 directly left as reference + // int num_n_blocks = params.num_n_blocks_ptr ? params.num_n_blocks_ptr[bidb] : num_m_blocks; + // auto find_log2_floor = [&](int n) { return 31 - cutlass::clz(n); }; + // int nheads_in_l2 = params.max_kvblocks_in_l2 < num_n_blocks + // ? 1 : 1 << find_log2_floor(params.max_kvblocks_in_l2 / num_n_blocks); + // if constexpr (!PackGQA) { nheads_in_l2 *= params.qhead_per_khead; } + // nheads_in_l2 = min(nheads_in_l2, params.num_head); + auto get_nheads_in_l2 = [&](int batch_idx) { + if constexpr(Prepared) { + return params.num_nheads_in_l2_ptr[batch_idx]; + } else { + return !PackGQA ? params.qhead_per_khead : 1; + } + }; + int nheads_in_l2 = get_nheads_in_l2(bidb); + int mh_in_l2 = nheads_in_l2 * num_m_blocks; + int section_idx = mh_block / mh_in_l2; + int l2_mod = mh_block - section_idx * mh_in_l2; + // tail section + int nheads_remainder = params.num_head - section_idx * nheads_in_l2; + int nheads_in_this_section = nheads_in_l2 <= nheads_remainder ? nheads_in_l2 : nheads_remainder; + block = l2_mod / nheads_in_this_section; + int bidh_residual = l2_mod - block * nheads_in_this_section; + bidh = section_idx * nheads_in_l2 + bidh_residual; + if constexpr(Split) { + // remember to set num_splits = 1 in work tile + uint32_t bidh_packed = reinterpret_cast(bidh) + (reinterpret_cast(num_splits) << 24); + bidh = reinterpret_cast(bidh_packed); + } + } else { + // NOTE: leave traverse heads first version for reference + // block = params.head_divmod.divmod(bidh, mh_block); + // if constexpr (Split) { + // int split_idx = block / num_m_blocks; + // block = block - split_idx * num_m_blocks; + // uint32_t bidh_packed = reinterpret_cast(bidh) + (reinterpret_cast(split_idx) << 16) + (reinterpret_cast(num_splits) << 24); + // bidh = reinterpret_cast(bidh_packed); + // } + bidh = mh_block / num_m_blocks; + block = mh_block - bidh * num_m_blocks; + if constexpr (Split) { + int bidh_actual = bidh / num_splits; + int split_idx = bidh - bidh_actual * num_splits; + uint32_t bidh_packed = reinterpret_cast(bidh_actual) + (reinterpret_cast(split_idx) << 16) + (reinterpret_cast(num_splits) << 24); + bidh = reinterpret_cast(bidh_packed); + } + } + block = num_m_blocks - 1 - block; + } else { + bidh = mh_block / num_m_blocks; + block = mh_block - bidh * num_m_blocks; + if constexpr (Split) { + int bidh_actual = bidh / num_splits; + int split_idx = bidh - bidh_actual * num_splits; + // TODO: idk why this gives wrong answer nondeterministically + // int bidh_actual, split_idx; + // split_idx = params.head_divmod.divmod(bidh_actual, bidh); + // Use the top 8 bits to store num_splits and the next 8 bits to store split_idx + // reinterpret_cast to uint32_t to make sure we're not doing sign extension when we shift + uint32_t bidh_packed = reinterpret_cast(bidh_actual) + (reinterpret_cast(split_idx) << 16) + (reinterpret_cast(num_splits) << 24); + // if (threadIdx.x == 0) { + // printf("blockIdx.x = %d, group_start_tiled = %d, bidb = %d, batch_idx_in_group = %d, mh_block = %d, num_m_blocks = %d, bidh = %d, bidh_actual = %d, split_idx = %d, num_splits = %d, bidh_packed = %d\n", blockIdx.x, group_start_tile, bidb, batch_idx_in_group, mh_block, num_m_blocks, bidh, bidh_actual, split_idx, num_splits, bidh_packed); + // } + bidh = reinterpret_cast(bidh_packed); + } + // if (blockIdx.x <= 9 && threadIdx.x == 0) { + // printf("Before returning, blockIdx.x = %d, threadIdx.x = %d, group_start_tile = %d, batch_idx_in_group = %d, bidb = %d, num_m_blocks = %d, next_tile_idx = %d, group_end_tile = %d, m_blocks_in_group = %d, mh_block = %d, bidh = %d, block = %d\n", blockIdx.x, threadIdx.x, group_start_tile, batch_idx_in_group, bidb, num_m_blocks, next_tile_idx, group_end_tile, m_blocks_in_group, mh_block, bidh, block); // } - bidh = reinterpret_cast(bidh_packed); } - // if (blockIdx.x <= 9 && threadIdx.x == 0) { - // printf("Before returning, blockIdx.x = %d, threadIdx.x = %d, group_start_tile = %d, batch_idx_in_group = %d, bidb = %d, num_m_blocks = %d, next_tile_idx = %d, group_end_tile = %d, m_blocks_in_group = %d, mh_block = %d, bidh = %d, block = %d\n", blockIdx.x, threadIdx.x, group_start_tile, batch_idx_in_group, bidb, num_m_blocks, next_tile_idx, group_end_tile, m_blocks_in_group, mh_block, bidh, block); - // } - return {next_tile_idx, block, bidh, bidb}; + return {group_start_tile, block, bidh, bidb}; } template diff --git a/hopper/tile_size.h b/hopper/tile_size.h index cede4f62326..9b9aef5a704 100644 --- a/hopper/tile_size.h +++ b/hopper/tile_size.h @@ -41,12 +41,12 @@ constexpr std::tuple tile_size_fwd_sm90( if (use_one_mma_wg) { return {64, is_causal || is_local || paged_kv_non_TMA ? 128 : 176, true, true}; } else { - return {128, is_causal || is_local || paged_kv_non_TMA ? 128 : 176, true, true}; + return {128, is_causal || is_local || paged_kv_non_TMA ? 128 : 160, true, true}; } // {128, 192, false, false} and {192, 128, false, true} are quite good too // 128 x 192 hits the limit of smem if MmaPV_is_RS, 128 x 144 hits the limit if !MmaPV_is_RS } else if (headdim <= 192) { - return {128, paged_kv_non_TMA || is_local ? 96 : (headdim_v <= 128 ? 128 : 112), true, true}; // 128 x 112 hits the limit of smem + return {128, paged_kv_non_TMA || is_local ? 96 : (headdim_v <= 128 ? 128 : 96), true, true}; // 128 x 112 hits the limit of smem } else { return {128, is_local ? 64 : 80, true, true}; // 128 x 80 hits the limit of smem }